address PR comments

Author: sasha-scale

nucleus/annotation.py

@@ -22,6 +22,54 @@
 )
 
 
+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)
+
+    # def from_json(cls, payload: dict):
+    #     if payload.get(TYPE_KEY, None) == BOX_TYPE:
+    #         geometry = payload.get(GEOMETRY_KEY, {})
+    #         return BoxAnnotation(
+    #             label=payload.get(LABEL_KEY, 0),
+    #             x=geometry.get(X_KEY, 0),
+    #             y=geometry.get(Y_KEY, 0),
+    #             width=geometry.get(WIDTH_KEY, 0),
+    #             height=geometry.get(HEIGHT_KEY, 0),
+    #             reference_id=payload.get(REFERENCE_ID_KEY, None),
+    #             item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+    #             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+    #             metadata=payload.get(METADATA_KEY, {}),
+    #         )
+    #     elif payload.get(TYPE_KEY, None) == POLYGON_TYPE:
+    #         geometry = payload.get(GEOMETRY_KEY, {})
+    #         return PolygonAnnotation(
+    #             label=payload.get(LABEL_KEY, 0),
+    #             vertices=geometry.get(VERTICES_KEY, []),
+    #             reference_id=payload.get(REFERENCE_ID_KEY, None),
+    #             item_id=payload.get(DATASET_ITEM_ID_KEY, None),
+    #             annotation_id=payload.get(ANNOTATION_ID_KEY, None),
+    #             metadata=payload.get(METADATA_KEY, {}),
+    #         )
+    #     else:
+    #         return SegmentationAnnotation(
+    #             mask_url=payload.get(MASK_URL_KEY),
+    #             annotations=[
+    #                 Segment.from_json(ann)
+    #                 for ann in payload.get(ANNOTATIONS_KEY, [])
+    #             ],
+    #             reference_id=payload.get(REFERENCE_ID_KEY, None),
+    #             item_id=payload.get(ITEM_ID_KEY, None),
+    #         )
+
+
 class Segment:
     def __init__(
         self, label: str, index: int, metadata: Optional[dict] = None
@@ -51,14 +99,17 @@ def to_payload(self) -> dict:
         return payload
 
 
-class SegmentationAnnotation:
+class SegmentationAnnotation(Annotation):
     def __init__(
         self,
         mask_url: str,
         annotations: List[Segment],
         reference_id: Optional[str] = None,
         item_id: Optional[str] = None,
     ):
+        super().__init__()
+        if not mask_url:
+            raise Exception("You must specify a mask_url.")
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
@@ -74,7 +125,7 @@ def __str__(self):
     @classmethod
     def from_json(cls, payload: dict):
         return cls(
-            mask_url=payload.get(MASK_URL_KEY, None),
+            mask_url=payload.get(MASK_URL_KEY),
             annotations=[
                 Segment.from_json(ann)
                 for ann in payload.get(ANNOTATIONS_KEY, [])
@@ -101,7 +152,7 @@ class AnnotationTypes(Enum):
 
 
 # TODO: Add base annotation class to reduce repeated code here
-class BoxAnnotation:
+class BoxAnnotation(Annotation):
     # pylint: disable=too-many-instance-attributes
     def __init__(
         self,
@@ -115,6 +166,7 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
+        super().__init__()
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
@@ -164,7 +216,7 @@ def __str__(self):
 
 
 # TODO: Add Generic type for 2D point
-class PolygonAnnotation:
+class PolygonAnnotation(Annotation):
     def __init__(
         self,
         label: str,
@@ -174,6 +226,7 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
+        super().__init__()
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."

nucleus/constants.py

@@ -48,6 +48,7 @@
 BOX_TYPE = "box"
 POLYGON_TYPE = "polygon"
 SEGMENTATION_TYPE = "segmentation"
+ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
 GEOMETRY_KEY = "geometry"
 AUTOTAGS_KEY = "autotags"
 MASK_URL_KEY = "mask_url"

nucleus/dataset.py

@@ -1,6 +1,7 @@
 from typing import List, Dict, Any, Optional, Union
 from .dataset_item import DatasetItem
 from .annotation import (
+    Annotation,
     BoxAnnotation,
     PolygonAnnotation,
     SegmentationAnnotation,
@@ -19,6 +20,7 @@
     BOX_TYPE,
     POLYGON_TYPE,
     SEGMENTATION_TYPE,
+    ANNOTATION_TYPES,
 )
 from .payload_constructor import construct_model_run_creation_payload
 
@@ -178,7 +180,6 @@ def refloc(self, reference_id: str) -> dict:
         }
         """
         response = self._client.dataitem_ref_id(self.id, reference_id)
-        print("RESPONSE" + str(response))
         return self._format_dataset_item_response(response)
 
     def loc(self, dataset_item_id: str) -> dict:
@@ -242,22 +243,12 @@ def _format_dataset_item_response(self, response: dict) -> dict:
             return response
 
         annotation_response = {}
-        if BOX_TYPE in annotation_payload:
-            annotation_response[BOX_TYPE] = [
-                BoxAnnotation.from_json(ann)
-                for ann in annotation_payload[BOX_TYPE]
-            ]
-        if POLYGON_TYPE in annotation_payload:
-            annotation_response[POLYGON_TYPE] = [
-                PolygonAnnotation.from_json(ann)
-                for ann in annotation_payload[POLYGON_TYPE]
-            ]
-        if SEGMENTATION_TYPE in annotation_payload:
-            annotation_response[
-                SEGMENTATION_TYPE
-            ] = SegmentationAnnotation.from_json(
-                annotation_payload[SEGMENTATION_TYPE]
-            )
+        for annotation_type in ANNOTATION_TYPES:
+            if annotation_type in annotation_payload:
+                annotation_response[annotation_type] = [
+                    Annotation.from_json(ann)
+                    for ann in annotation_payload[annotation_type]
+                ]
         return {
             ITEM_KEY: DatasetItem.from_json(item),
             ANNOTATIONS_KEY: annotation_response,

nucleus/model_run.py

@@ -127,7 +127,6 @@ def _format_prediction_response(
         self, response: dict
     ) -> Union[dict, List[Union[BoxPrediction, PolygonPrediction]]]:
         annotations = response.get(ANNOTATIONS_KEY, None)
-        print(annotations)
         if annotations:
             annotation_response = {}
             if BOX_TYPE in annotations: