Merge pull request #36 from scaleapi/sasha/segmentation_prediction

Sasha/segmentation prediction

Author: sasha-scale

nucleus/__init__.py

@@ -73,8 +73,13 @@
     BoxAnnotation,
     PolygonAnnotation,
     SegmentationAnnotation,
+    Segment,
+)
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
 )
-from .prediction import BoxPrediction, PolygonPrediction
 from .model_run import ModelRun
 from .slice import Slice
 from .upload_response import UploadResponse
@@ -623,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,
     ):
@@ -639,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 = {
@@ -670,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/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"

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

@@ -5,7 +5,11 @@
     PolygonAnnotation,
     SegmentationAnnotation,
 )
-from .prediction import BoxPrediction, PolygonPrediction
+from .prediction import (
+    BoxPrediction,
+    PolygonPrediction,
+    SegmentationPrediction,
+)
 from .constants import (
     ANNOTATION_UPDATE_KEY,
     NAME_KEY,
@@ -45,7 +49,9 @@ def construct_annotation_payload(
 
 
 def construct_segmentation_payload(
-    annotation_items: List[SegmentationAnnotation],
+    annotation_items: List[
+        Union[SegmentationAnnotation, SegmentationPrediction]
+    ],
     update: bool,
 ) -> dict:
     annotations = []

nucleus/prediction.py

@@ -1,5 +1,10 @@
 from typing import Dict, Optional, List, Any
-from .annotation import BoxAnnotation, PolygonAnnotation
+from .annotation import (
+    BoxAnnotation,
+    PolygonAnnotation,
+    Segment,
+    SegmentationAnnotation,
+)
 from .constants import (
     ANNOTATION_ID_KEY,
     DATASET_ITEM_ID_KEY,
@@ -13,9 +18,30 @@
     HEIGHT_KEY,
     CONFIDENCE_KEY,
     VERTICES_KEY,
+    ANNOTATIONS_KEY,
+    ITEM_ID_KEY,
+    MASK_URL_KEY,
 )
 
 
+class SegmentationPrediction(SegmentationAnnotation):
+    @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 = super().to_payload()
+        return payload
+
+
 class BoxPrediction(BoxAnnotation):
     def __init__(
         self,
@@ -31,7 +57,15 @@ def __init__(
         metadata: Optional[Dict] = None,
     ):
         super().__init__(
-            label, x, y, width, height, reference_id, item_id, annotation_id, metadata
+            label,
+            x,
+            y,
+            width,
+            height,
+            reference_id,
+            item_id,
+            annotation_id,
+            metadata,
         )
         self.confidence = confidence
 
@@ -73,7 +107,9 @@ def __init__(
         annotation_id: Optional[str] = None,
         metadata: Optional[Dict] = None,
     ):
-        super().__init__(label, vertices, reference_id, item_id, annotation_id, metadata)
+        super().__init__(
+            label, vertices, reference_id, item_id, annotation_id, metadata
+        )
         self.confidence = confidence
 
     def to_payload(self) -> dict:

tests/helpers.py

@@ -107,6 +107,8 @@ def reference_id_from_url(url):
     for i in range(len(TEST_IMG_URLS))
 ]
 
+TEST_SEGMENTATION_PREDICTIONS = TEST_SEGMENTATION_ANNOTATIONS
+
 TEST_BOX_PREDICTIONS = [
     {**TEST_BOX_ANNOTATIONS[i], "confidence": 0.10 * i}
     for i in range(len(TEST_BOX_ANNOTATIONS))