Merge pull request #35 from scaleapi/sasha/segmentation

Sasha/segmentation

Author: sasha-scale

nucleus/__init__.py

@@ -69,8 +69,17 @@
 
 from .dataset import Dataset
 from .dataset_item import DatasetItem
-from .annotation import BoxAnnotation, PolygonAnnotation
-from .prediction import BoxPrediction, PolygonPrediction
+from .annotation import (
+    BoxAnnotation,
+    PolygonAnnotation,
+    SegmentationAnnotation,
+    Segment,
+)
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .model_run import ModelRun
 from .slice import Slice
 from .upload_response import UploadResponse
@@ -79,6 +88,7 @@
     construct_annotation_payload,
     construct_model_creation_payload,
     construct_box_predictions_payload,
+    construct_segmentation_payload,
 )
 from .constants import (
     NUCLEUS_ENDPOINT,
@@ -465,7 +475,9 @@ def exception_handler(request, exception):
     def annotate_dataset(
         self,
         dataset_id: str,
-        annotations: List[Union[BoxAnnotation, PolygonAnnotation]],
+        annotations: List[
+            Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
+        ],
         update: bool,
         batch_size: int = 100,
     ):
@@ -477,9 +489,26 @@ def annotate_dataset(
         :return: {"dataset_id: str, "annotations_processed": int}
         """
 
+        # Split payload into segmentations and Box/Polygon
+        segmentations = [
+            ann
+            for ann in annotations
+            if isinstance(ann, SegmentationAnnotation)
+        ]
+        other_annotations = [
+            ann
+            for ann in annotations
+            if not isinstance(ann, SegmentationAnnotation)
+        ]
+
         batches = [
-            annotations[i : i + batch_size]
-            for i in range(0, len(annotations), batch_size)
+            other_annotations[i : i + batch_size]
+            for i in range(0, len(other_annotations), batch_size)
+        ]
+
+        semseg_batches = [
+            segmentations[i : i + batch_size]
+            for i in range(0, len(segmentations), batch_size)
         ]
 
         agg_response = {
@@ -488,22 +517,42 @@ def annotate_dataset(
             ANNOTATIONS_IGNORED_KEY: 0,
         }
 
+        total_batches = len(batches) + len(semseg_batches)
+
         tqdm_batches = self.tqdm_bar(batches)
 
-        for batch in tqdm_batches:
-            payload = construct_annotation_payload(batch, update)
-            response = self._make_request(
-                payload, f"dataset/{dataset_id}/annotate"
-            )
-            if STATUS_CODE_KEY in response:
-                agg_response[ERRORS_KEY] = response
-            else:
-                agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
-                    ANNOTATIONS_PROCESSED_KEY
-                ]
-                agg_response[ANNOTATIONS_IGNORED_KEY] += response[
-                    ANNOTATIONS_IGNORED_KEY
-                ]
+        with self.tqdm_bar(total=total_batches) as pbar:
+            for batch in tqdm_batches:
+                payload = construct_annotation_payload(batch, update)
+                response = self._make_request(
+                    payload, f"dataset/{dataset_id}/annotate"
+                )
+                pbar.update(1)
+                if STATUS_CODE_KEY in response:
+                    agg_response[ERRORS_KEY] = response
+                else:
+                    agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
+                        ANNOTATIONS_PROCESSED_KEY
+                    ]
+                    agg_response[ANNOTATIONS_IGNORED_KEY] += response[
+                        ANNOTATIONS_IGNORED_KEY
+                    ]
+
+            for s_batch in semseg_batches:
+                payload = construct_segmentation_payload(s_batch, update)
+                response = self._make_request(
+                    payload, f"dataset/{dataset_id}/annotate_segmentation"
+                )
+                pbar.update(1)
+                if STATUS_CODE_KEY in response:
+                    agg_response[ERRORS_KEY] = response
+                else:
+                    agg_response[ANNOTATIONS_PROCESSED_KEY] += response[
+                        ANNOTATIONS_PROCESSED_KEY
+                    ]
+                    agg_response[ANNOTATIONS_IGNORED_KEY] += response[
+                        ANNOTATIONS_IGNORED_KEY
+                    ]
 
         return agg_response
 
@@ -579,7 +628,9 @@ def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
     def predict(
         self,
         model_run_id: str,
-        annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        annotations: List[
+            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+        ],
         update: bool,
         batch_size: int = 100,
     ):
@@ -595,9 +646,26 @@ def predict(
             "predictions_ignored": int,
         }
         """
+        segmentations = [
+            ann
+            for ann in annotations
+            if isinstance(ann, SegmentationPrediction)
+        ]
+
+        other_predictions = [
+            ann
+            for ann in annotations
+            if not isinstance(ann, SegmentationPrediction)
+        ]
+
+        s_batches = [
+            segmentations[i : i + batch_size]
+            for i in range(0, len(segmentations), batch_size)
+        ]
+
         batches = [
-            annotations[i : i + batch_size]
-            for i in range(0, len(annotations), batch_size)
+            other_predictions[i : i + batch_size]
+            for i in range(0, len(other_predictions), batch_size)
         ]
 
         agg_response = {
@@ -610,7 +678,7 @@ def predict(
 
         for batch in tqdm_batches:
             batch_payload = construct_box_predictions_payload(
-                annotations,
+                batch,
                 update,
             )
             response = self._make_request(
@@ -626,8 +694,23 @@ def predict(
                     PREDICTIONS_IGNORED_KEY
                 ]
 
+        for s_batch in s_batches:
+            payload = construct_segmentation_payload(s_batch, update)
+            response = self._make_request(
+                payload, f"modelRun/{model_run_id}/predict_segmentation"
+            )
+            # pbar.update(1)
+            if STATUS_CODE_KEY in response:
+                agg_response[ERRORS_KEY] = response
+            else:
+                agg_response[PREDICTIONS_PROCESSED_KEY] += response[
+                    PREDICTIONS_PROCESSED_KEY
+                ]
+                agg_response[PREDICTIONS_IGNORED_KEY] += response[
+                    PREDICTIONS_IGNORED_KEY
+                ]
+
         return agg_response
-        # return self._make_request(payload, f"modelRun/{model_run_id}/predict")
 
     def commit_model_run(
         self, model_run_id: str, payload: Optional[dict] = None

nucleus/annotation.py

@@ -15,9 +15,86 @@
     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
+    ):
+        self.label = label
+        self.index = index
+        self.metadata = metadata
+
+    def __str__(self):
+        return str(self.to_payload())
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        return cls(
+            label=payload.get(LABEL_KEY, ""),
+            index=payload.get(INDEX_KEY, None),
+            metadata=payload.get(METADATA_KEY, None),
+        )
+
+    def to_payload(self) -> dict:
+        payload = {
+            LABEL_KEY: self.label,
+            INDEX_KEY: self.index,
+        }
+        if self.metadata is not None:
+            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,
+    ):
+        if bool(reference_id) == bool(item_id):
+            raise Exception(
+                "You must specify either a reference_id or an item_id for an annotation."
+            )
+        self.mask_url = mask_url
+        self.annotations = annotations
+        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=[
+                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),
+        )
+
+    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):
     BOX = BOX_TYPE
     POLYGON = POLYGON_TYPE
@@ -131,4 +208,4 @@ def to_payload(self) -> dict:
         }
 
     def __str__(self):
-        return str(self.to_payload())
+        return str(self.to_payload())

nucleus/constants.py

@@ -1,4 +1,5 @@
-NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
+# NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
+NUCLEUS_ENDPOINT = "http://localhost:3000/v1/nucleus"
 ITEMS_KEY = "items"
 ITEM_KEY = "item"
 REFERENCE_ID_KEY = "reference_id"
@@ -26,6 +27,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"
@@ -48,3 +50,6 @@
 POLYGON_TYPE = "polygon"
 GEOMETRY_KEY = "geometry"
 AUTOTAGS_KEY = "autotags"
+MASK_URL_KEY = "mask_url"
+INDEX_KEY = "index"
+SEGMENTATIONS_KEY = "segmentations"

nucleus/model.py

@@ -37,7 +37,9 @@ def create_run(
         self,
         name: str,
         dataset: Dataset,
-        predictions: List[Union[BoxPrediction, PolygonPrediction]],
+        predictions: List[
+            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+        ],
         metadata: Optional[Dict] = None,
     ) -> ModelRun:
         payload: dict = {

nucleus/model_run.py

@@ -1,6 +1,10 @@
 from typing import Optional, List, Dict, Any, Union
 from .constants import ANNOTATIONS_KEY, DEFAULT_ANNOTATION_UPDATE_MODE
-from .prediction import BoxPrediction, PolygonPrediction
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .payload_constructor import construct_box_predictions_payload
 
 
@@ -62,7 +66,9 @@ def commit(self, payload: Optional[dict] = None) -> dict:
 
     def predict(
         self,
-        annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        annotations: List[
+            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+        ],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
     ) -> dict:
         """

nucleus/payload_constructor.py

@@ -1,7 +1,15 @@
 from typing import List, Optional, Dict, Union
 from .dataset_item import DatasetItem
-from .annotation import BoxAnnotation, PolygonAnnotation
-from .prediction import BoxPrediction, PolygonPrediction
+from .annotation import (
+    BoxAnnotation,
+    PolygonAnnotation,
+    SegmentationAnnotation,
+)
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .constants import (
     ANNOTATION_UPDATE_KEY,
     NAME_KEY,
@@ -11,6 +19,7 @@
     ITEMS_KEY,
     FORCE_KEY,
     MODEL_ID_KEY,
+    SEGMENTATIONS_KEY,
 )
 
 
@@ -39,6 +48,19 @@ def construct_annotation_payload(
     return {ANNOTATIONS_KEY: annotations, ANNOTATION_UPDATE_KEY: update}
 
 
+def construct_segmentation_payload(
+    annotation_items: List[
+        Union[SegmentationAnnotation, SegmentationPrediction]
+    ],
+    update: bool,
+) -> dict:
+    annotations = []
+    for annotation_item in annotation_items:
+        annotations.append(annotation_item.to_payload())
+
+    return {SEGMENTATIONS_KEY: annotations, ANNOTATION_UPDATE_KEY: update}
+
+
 def construct_box_predictions_payload(
     box_predictions: List[Union[BoxPrediction, PolygonPrediction]],
     update: bool,