Merge pull request #365 from Labelbox/ms/in-line-validation

in-line validation

Author: msokoloff1

labelbox/schema/bulk_import_request.py

@@ -781,33 +781,60 @@ def is_valid_location(cls, v):
         return v
 
 
-class MaskFeatures(BaseModel):
-    instanceURI: str
-    colorRGB: Union[List[int], Tuple[int, int, int]]
+class RLEMaskFeatures(BaseModel):
+    counts: List[int]
+    size: List[int]
 
+    @validator('counts')
+    def validate_counts(cls, counts):
+        if not all([count >= 0 for count in counts]):
+            raise ValueError(
+                "Found negative value for counts. They should all be zero or positive"
+            )
+        return counts
 
-class NDMask(NDBaseTool):
-    ontology_type: Literal["superpixel"] = "superpixel"
-    mask: MaskFeatures = pydantic.Field(determinant=True)
+    @validator('size')
+    def validate_size(cls, size):
+        if len(size) != 2:
+            raise ValueError(
+                f"Mask `size` should have two ints representing height and with. Found : {size}"
+            )
+        if not all([count > 0 for count in size]):
+            raise ValueError(
+                f"Mask `size` should be a postitive int. Found : {size}")
+        return size
 
-    @validator('mask')
-    def is_valid_mask(cls, v):
-        if isinstance(v, BaseModel):
-            v = v.dict()
 
-        colors = v['colorRGB']
+class PNGMaskFeatures(BaseModel):
+    # base64 encoded png bytes
+    png: str
+
+
+class URIMaskFeatures(BaseModel):
+    instanceURI: str
+    colorRGB: Union[List[int], Tuple[int, int, int]]
+
+    @validator('colorRGB')
+    def validate_color(cls, colorRGB):
         #Does the dtype matter? Can it be a float?
-        if not isinstance(colors, (tuple, list)):
+        if not isinstance(colorRGB, (tuple, list)):
             raise ValueError(
-                f"Received color that is not a list or tuple. Found : {colors}")
-        elif len(colors) != 3:
+                f"Received color that is not a list or tuple. Found : {colorRGB}"
+            )
+        elif len(colorRGB) != 3:
             raise ValueError(
-                f"Must provide RGB values for segmentation colors. Found : {colors}"
+                f"Must provide RGB values for segmentation colors. Found : {colorRGB}"
             )
-        elif not all([0 <= color <= 255 for color in colors]):
+        elif not all([0 <= color <= 255 for color in colorRGB]):
             raise ValueError(
-                f"All rgb colors must be between 0 and 255. Found : {colors}")
-        return v
+                f"All rgb colors must be between 0 and 255. Found : {colorRGB}")
+        return colorRGB
+
+
+class NDMask(NDBaseTool):
+    ontology_type: Literal["superpixel"] = "superpixel"
+    mask: Union[URIMaskFeatures, PNGMaskFeatures,
+                RLEMaskFeatures] = pydantic.Field(determinant=True)
 
 
 #A union with custom construction logic to improve error messages

tests/integration/annotation_import/conftest.py

@@ -232,6 +232,33 @@ def segmentation_inference(prediction_id_mapping):
     return segmentation
 
 
+@pytest.fixture
+def segmentation_inference_rle(prediction_id_mapping):
+    segmentation = prediction_id_mapping['superpixel'].copy()
+    segmentation.update({
+        'uuid': str(uuid.uuid4()),
+        'mask': {
+            'size': [10, 10],
+            'counts': [1, 0, 10, 100]
+        }
+    })
+    del segmentation['tool']
+    return segmentation
+
+
+@pytest.fixture
+def segmentation_inference_png(prediction_id_mapping):
+    segmentation = prediction_id_mapping['superpixel'].copy()
+    segmentation.update({
+        'uuid': str(uuid.uuid4()),
+        'mask': {
+            'png': "somedata",
+        }
+    })
+    del segmentation['tool']
+    return segmentation
+
+
 @pytest.fixture
 def checklist_inference(prediction_id_mapping):
     checklist = prediction_id_mapping['checklist'].copy()
@@ -283,10 +310,12 @@ def model_run_predictions(polygon_inference, rectangle_inference,
 # also used for label imports
 @pytest.fixture
 def object_predictions(polygon_inference, rectangle_inference, line_inference,
-                       entity_inference, segmentation_inference):
+                       entity_inference, segmentation_inference,
+                       segmentation_inference_rle, segmentation_inference_png):
     return [
         polygon_inference, rectangle_inference, line_inference,
-        entity_inference, segmentation_inference
+        entity_inference, segmentation_inference, segmentation_inference_rle,
+        segmentation_inference_png
     ]
 
 

tests/integration/annotation_import/test_ndjson_validation.py

@@ -28,7 +28,9 @@ def test_subclassification_construction(rectangle_inference):
                           (fixture_ref('rectangle_inference'), NDRectangle),
                           (fixture_ref('line_inference'), NDPolyline),
                           (fixture_ref('entity_inference'), NDTextEntity),
-                          (fixture_ref('segmentation_inference'), NDMask)])
+                          (fixture_ref('segmentation_inference'), NDMask),
+                          (fixture_ref('segmentation_inference_rle'), NDMask),
+                          (fixture_ref('segmentation_inference_png'), NDMask)])
 def test_tool_construction(inference, expected_type):
     assert isinstance(NDTool.build(inference), expected_type)
 
@@ -131,6 +133,14 @@ def test_incorrect_mask(segmentation_inference, configured_project):
     with pytest.raises(MALValidationError):
         _validate_ndjson([seg], configured_project)
 
+    seg['mask'] = {'counts': [0], 'size': [0, 1]}
+    with pytest.raises(MALValidationError):
+        _validate_ndjson([seg], configured_project)
+
+    seg['mask'] = {'counts': [-1], 'size': [1, 1]}
+    with pytest.raises(MALValidationError):
+        _validate_ndjson([seg], configured_project)
+
 
 def test_all_validate_json(configured_project, predictions):
     #Predictions contains one of each type of prediction.