Merge pull request #40 from scaleapi/sasha/annotation_id_fix

Add annotation_id field to segmentation annotations

Author: sasha-scale

nucleus/annotation.py

@@ -26,10 +26,8 @@ class Annotation:
     @classmethod
     def from_json(cls, payload: dict):
         if payload.get(TYPE_KEY, None) == BOX_TYPE:
-            geometry = payload.get(GEOMETRY_KEY, {})
             return BoxAnnotation.from_json(payload)
         elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
-            geometry = payload.get(GEOMETRY_KEY, {})
             return PolygonAnnotation.from_json(payload)
         else:
             return SegmentationAnnotation.from_json(payload)
@@ -71,6 +69,7 @@ def __init__(
         annotations: List[Segment],
         reference_id: Optional[str] = None,
         item_id: Optional[str] = None,
+        annotation_id: Optional[str] = None,
     ):
         super().__init__()
         if not mask_url:
@@ -83,6 +82,7 @@ def __init__(
         self.annotations = annotations
         self.reference_id = reference_id
         self.item_id = item_id
+        self.annotation_id = annotation_id
 
     def __str__(self):
         return str(self.to_payload())
@@ -97,12 +97,14 @@ def from_json(cls, payload: dict):
             ],
             reference_id=payload.get(REFERENCE_ID_KEY, None),
             item_id=payload.get(ITEM_ID_KEY, None),
+            annotation_id=payload.get(ANNOTATION_ID_KEY, None),
         )
 
     def to_payload(self) -> dict:
         payload = {
             MASK_URL_KEY: self.mask_url,
             ANNOTATIONS_KEY: [ann.to_payload() for ann in self.annotations],
+            ANNOTATION_ID_KEY: self.annotation_id,
         }
         if self.reference_id:
             payload[REFERENCE_ID_KEY] = self.reference_id

nucleus/payload_constructor.py

@@ -49,8 +49,8 @@ def construct_annotation_payload(
 
 
 def construct_segmentation_payload(
-    annotation_items: List[
-        Union[SegmentationAnnotation, SegmentationPrediction]
+    annotation_items: Union[
+        List[SegmentationAnnotation], List[SegmentationPrediction]
     ],
     update: bool,
 ) -> dict:

nucleus/prediction.py

@@ -25,6 +25,8 @@
 
 
 class SegmentationPrediction(SegmentationAnnotation):
+    # No need to define init or to_payload methods because
+    # we default to functions defined in the parent class
     @classmethod
     def from_json(cls, payload: dict):
         return cls(
@@ -37,10 +39,6 @@ def from_json(cls, payload: dict):
             item_id=payload.get(ITEM_ID_KEY, None),
         )
 
-    def to_payload(self) -> dict:
-        payload = super().to_payload()
-        return payload
-
 
 class BoxPrediction(BoxAnnotation):
     def __init__(

tests/helpers.py

@@ -94,10 +94,11 @@ def reference_id_from_url(url):
     for i in range(len(TEST_IMG_URLS))
 ]
 
-TEST_MASK_URL = "https://scale-temp.s3.amazonaws.com/scale-select/nucleus/mscoco_semseg_masks_uint8/000000000285.png"
+TEST_MASK_URL = "https://scale-ml.s3.amazonaws.com/tmp/nucleus/mscoco_semseg_masks/000000000285.png"
 TEST_SEGMENTATION_ANNOTATIONS = [
     {
         "reference_id": reference_id_from_url(TEST_IMG_URLS[i]),
+        "annotation_id": f"[Pytest] Segmentation Annotation Id{i}",
         "mask_url": get_signed_url(TEST_MASK_URL),
         "annotations": [
             {"label": "bear", "index": 2},
@@ -151,6 +152,9 @@ def assert_segmentation_annotation_matches_dict(
     annotation_instance, annotation_dict
 ):
     assert annotation_instance.mask_url == annotation_dict["mask_url"]
+    assert (
+        annotation_instance.annotation_id == annotation_dict["annotation_id"]
+    )
     # Cannot guarantee segments are in same order
     assert len(annotation_instance.annotations) == len(
         annotation_dict["annotations"]