add support for image bytes as a source for raster video mask import (#1277)

Author: apollonin

Makefile

@@ -1,6 +1,6 @@
 
 build-image:
-	docker build -t local/labelbox-python:test .
+	docker build --quiet -t local/labelbox-python:test .
 
 test-local: build-image
 

labelbox/data/annotation_types/video.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import List, Optional, Tuple
 
-from pydantic import BaseModel, validator
+from pydantic import BaseModel, validator, root_validator
 from labelbox.data.annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation
 
 from labelbox.data.annotation_types.annotation import ClassificationAnnotation, ObjectAnnotation
@@ -81,14 +81,24 @@ class DICOMObjectAnnotation(VideoObjectAnnotation):
         keyframe (bool): Whether or not this annotation was a human generated or interpolated annotation
         segment_id (Optional[Int]): Index of video segment this annotation belongs to
         classifications (List[ClassificationAnnotation]) = []
-        extra (Dict[str, Any])        
+        extra (Dict[str, Any])
     """
     group_key: GroupKey
 
 
 class MaskFrame(_CamelCaseMixin, BaseModel):
     index: int
-    instance_uri: str
+    instance_uri: Optional[str] = None
+    im_bytes: Optional[bytes] = None
+
+    @root_validator()
+    def validate_args(cls, values):
+        im_bytes = values.get("im_bytes")
+        instance_uri = values.get("instance_uri")
+
+        if im_bytes == instance_uri == None:
+            raise ValueError("One of `instance_uri`, `im_bytes` required.")
+        return values
 
     @validator("instance_uri")
     def validate_uri(cls, v):

labelbox/data/serialization/ndjson/objects.py

@@ -501,13 +501,23 @@ class NDVideoMasks(NDJsonBase, ConfidenceMixin):
     masks: NDVideoMasksFramesInstances
 
     def to_common(self) -> VideoMaskAnnotation:
+        for mask_frame in self.masks.frames:
+            if mask_frame.im_bytes:
+                mask_frame.im_bytes = base64.b64decode(
+                    mask_frame.im_bytes.encode('utf-8'))
+
         return VideoMaskAnnotation(
             frames=self.masks.frames,
             instances=self.masks.instances,
         )
 
     @classmethod
     def from_common(cls, annotation, data):
+        for mask_frame in annotation.frames:
+            if mask_frame.im_bytes:
+                mask_frame.im_bytes = base64.b64encode(
+                    mask_frame.im_bytes).decode('utf-8')
+
         return cls(
             data_row=DataRow(id=data.uid, global_key=data.global_key),
             masks=NDVideoMasksFramesInstances(frames=annotation.frames,

tests/data/annotation_types/test_video.py

@@ -6,6 +6,7 @@ def test_mask_frame():
                                     instance_uri="http://path/to/frame.png")
     assert mask_frame.dict(by_alias=True) == {
         'index': 1,
+        'imBytes': None,
         'instanceURI': 'http://path/to/frame.png'
     }
 

tests/data/serialization/ndjson/test_dicom.py

@@ -1,5 +1,6 @@
 from copy import copy
 import pytest
+import base64
 import labelbox.types as lb_types
 from labelbox.data.serialization import NDJsonConverter
 from labelbox.data.serialization.ndjson.objects import NDDicomSegments, NDDicomSegment, NDDicomLine
@@ -89,9 +90,11 @@
     'masks': {
         'frames': [{
             'index': 1,
+            'imBytes': None,
             'instanceURI': instance_uri_1
         }, {
             'index': 5,
+            'imBytes': None,
             'instanceURI': instance_uri_5
         }],
         'instances': [
@@ -157,9 +160,12 @@
     video_mask_label_with_global_key
 ]
 ndjsons = [
-    polyline_annotation_ndjson, polyline_annotation_ndjson_with_global_key,
-    dicom_mask_annotation_ndjson, dicom_mask_annotation_ndjson_with_global_key,
-    video_mask_annotation_ndjson, video_mask_annotation_ndjson_with_global_key
+    polyline_annotation_ndjson,
+    polyline_annotation_ndjson_with_global_key,
+    dicom_mask_annotation_ndjson,
+    dicom_mask_annotation_ndjson_with_global_key,
+    video_mask_annotation_ndjson,
+    video_mask_annotation_ndjson_with_global_key,
 ]
 labels_ndjsons = list(zip(labels, ndjsons))
 

tests/integration/annotation_import/test_data_types.py

@@ -65,6 +65,7 @@ def get_annotation_comparison_dicts_from_labels(labels):
         if 'masks' in annotation:
             for frame in annotation['masks']['frames']:
                 frame.pop('instanceURI')
+                frame.pop('imBytes')
             for instance in annotation['masks']['instances']:
                 instance.pop('colorRGB')
     return labels_ndjson