adjust refloc payload format

Author: sasha-scale

nucleus/constants.py

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

nucleus/dataset.py

@@ -1,6 +1,10 @@
 from typing import List, Dict, Any, Optional, Union
 from .dataset_item import DatasetItem
-from .annotation import BoxAnnotation, PolygonAnnotation
+from .annotation import (
+    BoxAnnotation,
+    PolygonAnnotation,
+    SegmentationAnnotation,
+)
 from .constants import (
     DATASET_NAME_KEY,
     DATASET_MODEL_RUNS_KEY,
@@ -12,9 +16,13 @@
     ITEM_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     ANNOTATIONS_KEY,
+    BOX_TYPE,
+    POLYGON_TYPE,
+    SEGMENTATION_TYPE,
 )
 from .payload_constructor import construct_model_run_creation_payload
 
+
 class Dataset:
     """
     Nucleus Dataset. You can append images with metadata to your dataset,
@@ -170,6 +178,7 @@ 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:
@@ -227,17 +236,34 @@ def list_autotags(self):
 
     def _format_dataset_item_response(self, response: dict) -> dict:
         item = response.get(ITEM_KEY, None)
-        annotation_payload = response.get(ANNOTATIONS_KEY, [])
+        annotation_payload = response.get(ANNOTATIONS_KEY, {})
         if not item or not annotation_payload:
             # An error occured
             return response
-        annotations = [
-            BoxAnnotation.from_json(ann)
-            if ann["type"] == "box"
-            else PolygonAnnotation.from_json(ann)
-            for ann in annotation_payload
-        ]
+
+        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(ann)
+                for ann in annotation_payload[SEGMENTATION_TYPE]
+            ]
+        # annotations = [
+        #     BoxAnnotation.from_json(ann)
+        #     if ann["type"] == "box"
+        #     else PolygonAnnotation.from_json(ann)
+        #     for ann in annotation_payload
+        # ]
         return {
             ITEM_KEY: DatasetItem.from_json(item),
-            ANNOTATIONS_KEY: annotations,
+            ANNOTATIONS_KEY: annotation_response,
         }

nucleus/model.py

@@ -1,6 +1,10 @@
 from typing import List, Optional, Dict, Union
 from .dataset import Dataset
-from .prediction import BoxPrediction, PolygonPrediction
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .model_run import ModelRun
 from .constants import (
     NAME_KEY,
@@ -25,7 +29,7 @@ def __init__(
         self._client = client
 
     def __repr__(self):
-        return f'Model(model_id={self.id}, name={self.name}, reference_id={self.reference_id}, metadata={self.metadata}, client={self._client})'
+        return f"Model(model_id={self.id}, name={self.name}, reference_id={self.reference_id}, metadata={self.metadata}, client={self._client})"
 
     def __eq__(self, other):
         return self.id == other.id