Working encode test

Author: feynmanliang

labelbox/lbx.py

@@ -24,23 +24,47 @@
 bytes 0-1 -> R | G value of pixel (superpixel layer #)
 bytes 2-3 -> # consecutive occurences
 """
+from io import BytesIO
 import itertools
 import struct
 
 import numpy as np
 from PIL import Image
 
+_BACKGROUND_RGBA = [255, 255, 255, 255]
+_HEADER_LENGTH = 6 * 4
+
 
 def encode(image):
-    """Converts a `PIL.Image` to a `io.BytesIO` with LBX encoded data.
+    """Converts a RGB `PIL.Image` to a `io.BytesIO` with LBX encoded data.
 
     Args:
         image (`PIL.Image`): The image to encode.
 
     Returns:
         A `io.BytesIO` containing the LBX encoded image.
     """
-    return image
+    im = image.convert('RGBA')
+    pixel_words = np.array(im).reshape(-1, 4)
+
+    colormap = list(filter(lambda x: not np.all(x == _BACKGROUND_RGBA),
+            np.unique(pixel_words, axis=0)))
+
+
+    input_byte_len = len(np.array(im).flat)
+    buff = BytesIO(bytes([0] * (len(colormap) * 4 + input_byte_len)))
+
+    offset = _HEADER_LENGTH  # make room for header
+    for color in colormap:
+        struct.pack_into('<BBBB', buff.getbuffer(), offset, *color)
+        offset += 4
+
+    # write header
+    struct.pack_into('<iiiiii', buff.getbuffer(), 0,
+            1, im.width, im.height, input_byte_len, len(colormap),
+            offset - len(colormap) - _HEADER_LENGTH)
+
+    return buff
 
 
 def decode(lbx):
@@ -53,13 +77,13 @@ def decode(lbx):
         A `PIL.Image` of the decoded data.
     """
     version, width, height, byteLength, numColors, numBlocks = \
-            map(lambda x: x[0], struct.iter_unpack('<i', lbx.read(24)))
+            map(lambda x: x[0], struct.iter_unpack('<i', lbx.read(_HEADER_LENGTH)))
     colormap = np.array(
             list(_grouper(
                 map(lambda x: x[0],
                     struct.iter_unpack('<B', lbx.read(4 * numColors))),
                 4)) +
-            [[0, 0, 0, 255]]
+            [_BACKGROUND_RGBA]
             )
 
     image_data = np.zeros((width * height, 4), dtype='uint8')

tests/test_lbx.py

@@ -3,6 +3,7 @@
 import labelbox.lbx as lbx
 from PIL import Image
 import pytest
+import struct
 
 @pytest.fixture
 def im_png(datadir):
@@ -20,12 +21,12 @@ def test_lbx_decode(lbx_sample):
 
 def test_lbx_encode(im_png):
     lbx_encoded = lbx.encode(im_png)
-    version, width, height = map(lambda x: x[0], struct.iter_unpack('<i', lbx.read(12)))
+    version, width, height = map(lambda x: x[0], struct.iter_unpack('<i', lbx_encoded.read(12)))
     assert version == 1
     assert width == 800
     assert height == 600
 
 
-# def test_identity(im_png):
-#     assert np.all(np.array(lbx.decode(lbx.encode(im_png))) == np.array(im_png))
+def test_identity(im_png):
+    assert np.all(np.array(lbx.decode(lbx.encode(im_png))) == np.array(im_png))