uses png if no signed url provided

Author: Matt Sokoloff

labelbox/data/serialization/ndjson/label.py

@@ -42,7 +42,6 @@ def _generate_annotations(
                                                         NDConfusionMatrixMetric,
                                                         NDScalarMetric]]]
     ) -> Generator[Label, None, None]:
-
         for data_row_id, annotations in grouped_annotations.items():
             annots = []
             for annotation in annotations:

labelbox/data/serialization/ndjson/objects.py

@@ -1,6 +1,11 @@
+from ast import Bytes
+from io import BytesIO
 from typing import Any, Dict, List, Tuple, Union
+import base64
+import numpy as np
 
 from pydantic import BaseModel
+from PIL import Image
 
 from ...annotation_types.data import ImageData, TextData, MaskData
 from ...annotation_types.ner import TextEntity
@@ -113,28 +118,48 @@ def from_common(cls, rectangle: Rectangle,
                    classifications=classifications)
 
 
-class _Mask(BaseModel):
+class _URIMask(BaseModel):
     instanceURI: str
     colorRGB: Tuple[int, int, int]
 
 
+class _PNGMask(BaseModel):
+    png: str
+
+
 class NDMask(NDBaseObject):
-    mask: _Mask
+    mask: Union[_URIMask, _PNGMask]
 
     def to_common(self) -> Mask:
-        return Mask(mask=MaskData(url=self.mask.instanceURI),
-                    color=self.mask.colorRGB)
+        if isinstance(self.mask, _URIMask):
+            return Mask(mask=MaskData(url=self.mask.instanceURI),
+                        color=self.mask.colorRGB)
+        else:
+            encoded_image_bytes = self.mask.png.encode('utf-8')
+            image_bytes = base64.b64decode(encoded_image_bytes)
+            image = np.array(Image.open(BytesIO(image_bytes)))
+            if np.max(image) > 1:
+                raise ValueError(
+                    f"Expected binary mask. Found max value of {np.max(image)}")
+            # Color is 1,1,1 because it is a binary array and we are just stacking it into 3 channels
+            return Mask(mask=MaskData.from_2D_arr(image), color=(1, 1, 1))
 
     @classmethod
     def from_common(cls, mask: Mask,
                     classifications: List[ClassificationAnnotation],
                     feature_schema_id: Cuid, extra: Dict[str, Any],
                     data: Union[ImageData, TextData]) -> "NDMask":
-        if mask.mask.url is None:
-            raise ValueError(
-                "Mask does not have a url. Use `LabelGenerator.add_url_to_masks`, `LabelList.add_url_to_masks`, or `Label.add_url_to_masks`."
-            )
-        return cls(mask=_Mask(instanceURI=mask.mask.url, colorRGB=mask.color),
+
+        if mask.mask.url is not None:
+            lbv1_mask = _URIMask(instanceURI=mask.mask.url, colorRGB=mask.color)
+        else:
+            binary = np.all(mask.mask.value == mask.color, axis=-1)
+            im_bytes = BytesIO()
+            Image.fromarray(binary, 'L').save(im_bytes, format="PNG")
+            lbv1_mask = _PNGMask(
+                png=base64.b64encode(im_bytes.getvalue()).decode('utf-8'))
+
+        return cls(mask=lbv1_mask,
                    dataRow=DataRow(id=data.uid),
                    schema_id=feature_schema_id,
                    uuid=extra.get('uuid'),

tests/data/serialization/ndjson/test_image.py

@@ -1,6 +1,9 @@
 import json
+import numpy as np
+import cv2
 
 from labelbox.data.serialization.ndjson.converter import NDJsonConverter
+from labelbox.data.annotation_types import Mask, Label, ObjectAnnotation, ImageData, MaskData
 
 
 def round_dict(data):
@@ -25,3 +28,59 @@ def test_image():
     for r in res:
         r.pop('classifications', None)
     assert [round_dict(x) for x in res] == [round_dict(x) for x in data]
+
+
+def test_mask():
+    data = [{
+        "uuid": "b862c586-8614-483c-b5e6-82810f70cac0",
+        "schemaId": "ckrazcueb16og0z6609jj7y3y",
+        "dataRow": {
+            "id": "ckrazctum0z8a0ybc0b0o0g0v"
+        },
+        "mask": {
+            "png":
+                "iVBORw0KGgoAAAANSUhEUgAAAAoAAAAKCAAAAACoWZBhAAAAMklEQVR4nD3MuQ3AQADDMOqQ/Vd2ijytaSiZLAcYuyLEYYYl9cvrlGftTHvsYl+u/3EDv0QLI8Z7FlwAAAAASUVORK5CYII="
+        }
+    }, {
+        "uuid": "751fc725-f7b6-48ed-89b0-dd7d94d08af6",
+        "schemaId": "ckrazcuec16ok0z66f956apb7",
+        "dataRow": {
+            "id": "ckrazctum0z8a0ybc0b0o0g0v"
+        },
+        "mask": {
+            "instanceURI":
+                "https://storage.labelbox.com/ckqcx1czn06830y61gh9v02cs%2F3e729327-f038-f66c-186e-45e921ef9717-1?Expires=1626806874672&KeyName=labelbox-assets-key-3&Signature=YsUOGKrsqmAZ68vT9BlPJOaRyLY",
+            "colorRGB": [255, 0, 0]
+        }
+    }]
+    res = NDJsonConverter.deserialize(data).as_list()
+    res = list(NDJsonConverter.serialize(res))
+    for r in res:
+        r.pop('classifications', None)
+
+    assert [round_dict(x) for x in res] == [round_dict(x) for x in data]
+
+
+def test_mask_from_arr():
+    mask_arr = np.round(np.zeros((32, 32))).astype(np.uint8)
+    mask_arr = cv2.rectangle(mask_arr, (5, 5), (10, 10), (1, 1), -1)
+
+    label = Label(annotations=[
+        ObjectAnnotation(feature_schema_id="1" * 25,
+                         value=Mask(mask=MaskData.from_2D_arr(arr=mask_arr),
+                                    color=(1, 1, 1)))
+    ],
+                  data=ImageData(uid="0" * 25))
+    res = next(NDJsonConverter.serialize([label]))
+    res.pop('uuid')
+    assert res == {
+        "classifications": [],
+        "schemaId": "1" * 25,
+        "dataRow": {
+            "id": "0" * 25
+        },
+        "mask": {
+            "png":
+                "iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAAAAABWESUoAAAAHklEQVR4nGNgGAKAEYn8j00BEyETBoOCUTAKhhwAAJW+AQwvpePVAAAAAElFTkSuQmCC"
+        }
+    }