flesh out Segmentation Class

Author: sasha-scale

nucleus/annotation.py

@@ -14,20 +14,31 @@
     LABEL_KEY,
     TYPE_KEY,
     VERTICES_KEY,
+    ITEM_ID_KEY,
+    MASK_URL_KEY,
+    INDEX_KEY,
+    ANNOTATIONS_KEY,
 )
 
+
 class Segment:
-    def __init__(self, label: str, index: int, metadata: Optional[dict] = None):
+    def __init__(
+        self, label: str, index: int, metadata: Optional[dict] = None
+    ):
         self.label = label
         self.index = index
         self.metadata = metadata
 
     def __str__(self):
-        return self.to_payload()
+        return str(self.to_payload())
 
     @classmethod
     def from_json(cls, payload: dict):
-        return cls(payload.get(LABEL_KEY), payload.get(INDEX_KEY), payload.get(METADATA_KEY, None))
+        return cls(
+            label=payload.get(LABEL_KEY),
+            index=payload.get(INDEX_KEY),
+            metadata=payload.get(METADATA_KEY, None),
+        )
 
     def to_payload(self) -> dict:
         payload = {
@@ -38,8 +49,15 @@ def to_payload(self) -> dict:
             payload[METADATA_KEY] = self.metadata
         return payload
 
+
 class SegmentationAnnotation:
-    def __init__(self, mask_url: str, annotations: List[Segment], reference_id: Optional[str] = None, item_id: Optional[str] = None):
+    def __init__(
+        self,
+        mask_url: str,
+        annotations: List[Segment],
+        reference_id: Optional[str] = None,
+        item_id: Optional[str] = None,
+    ):
         if bool(reference_id) == bool(item_id):
             raise Exception(
                 "You must specify either a reference_id or an item_id for an annotation."
@@ -49,17 +67,28 @@ def __init__(self, mask_url: str, annotations: List[Segment], reference_id: Opti
         self.reference_id = reference_id
         self.item_id = item_id
 
-
     def __str__(self):
+        return str(self.to_payload())
 
     @classmethod
     def from_json(cls, payload: dict):
+        return cls(
+            mask_url=payload[MASK_URL_KEY],
+            annotations=payload[ANNOTATIONS_KEY],
+            reference_id=payload.get(REFERENCE_ID_KEY, None),
+            item_id=payload.get(ITEM_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],
         }
+        if self.reference_id:
+            payload[REFERENCE_ID_KEY] = self.reference_id
+        else:
+            payload[ITEM_ID_KEY] = self.item_id
+        return payload
 
 
 class AnnotationTypes(Enum):

nucleus/constants.py

@@ -21,6 +21,7 @@
 MODEL_RUN_ID_KEY = "model_run_id"
 MODEL_ID_KEY = "model_id"
 DATASET_ITEM_ID_KEY = "dataset_item_id"
+ITEM_ID_KEY = "item_id"
 DATASET_ITEM_IDS_KEY = "dataset_item_ids"
 SLICE_ID_KEY = "slice_id"
 DATASET_NAME_KEY = "name"