First pass

Author: Claire Pajot

nucleus/__init__.py

@@ -87,6 +87,7 @@
     CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
+    CategoryPrediction,
 )
 from .scene import Frame, LidarScene
 from .slice import Slice
@@ -742,6 +743,7 @@ def predict(
                 PolygonPrediction,
                 CuboidPrediction,
                 SegmentationPrediction,
+                CategoryPrediction,
             ]
         ],
         model_run_id: Optional[str] = None,
@@ -752,7 +754,7 @@ def predict(
     ):
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
-        :param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
+        :param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction, CategoryPrediction]],
         :param update: bool
         :return:
         {
@@ -914,7 +916,7 @@ def predictions_ref_id(self, model_run_id: str, ref_id: str):
         :param reference_id: reference_id of a dataset item.
         :return:
         {
-            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction, CategoryPrediction]],
         }
         """
         return self.make_request(
@@ -939,7 +941,7 @@ def predictions_iloc(self, model_run_id: str, i: int):
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
         :return:
         {
-            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction, CategoryPrediction]],
         }
         """
         return self.make_request(

nucleus/payload_constructor.py

@@ -14,6 +14,7 @@
     CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
+    CategoryPrediction,
 )
 from .constants import (
     ANNOTATION_UPDATE_KEY,
@@ -91,7 +92,12 @@ def construct_segmentation_payload(
 
 def construct_box_predictions_payload(
     box_predictions: List[
-        Union[BoxPrediction, PolygonPrediction, CuboidPrediction]
+        Union[
+            BoxPrediction,
+            PolygonPrediction,
+            CuboidPrediction,
+            CategoryPrediction,
+        ]
     ],
     update: bool,
 ) -> dict:

nucleus/prediction.py

@@ -1,6 +1,7 @@
 from typing import Dict, Optional, List
 from .annotation import (
     BoxAnnotation,
+    CategoryAnnotation,
     Point,
     PolygonAnnotation,
     Segment,
@@ -13,10 +14,12 @@
     BOX_TYPE,
     CUBOID_TYPE,
     POLYGON_TYPE,
+    CATEGORY_TYPE,
     REFERENCE_ID_KEY,
     METADATA_KEY,
     GEOMETRY_KEY,
     LABEL_KEY,
+    TAXONOMY_NAME_KEY,
     TYPE_KEY,
     X_KEY,
     Y_KEY,
@@ -40,6 +43,8 @@ def from_json(payload: dict):
         return PolygonPrediction.from_json(payload)
     elif payload.get(TYPE_KEY, None) == CUBOID_TYPE:
         return CuboidPrediction.from_json(payload)
+    elif payload.get(TYPE_KEY, None) == CATEGORY_TYPE:
+        return CategoryPrediction.from_json(payload)
     else:
         return SegmentationPrediction.from_json(payload)
 
@@ -207,3 +212,43 @@ def from_json(cls, payload: dict):
             metadata=payload.get(METADATA_KEY, {}),
             class_pdf=payload.get(CLASS_PDF_KEY, None),
         )
+
+
+class CategoryPrediction(CategoryAnnotation):
+    def __init__(
+        self,
+        label: str,
+        taxonomy_name: str,
+        reference_id: str,
+        confidence: Optional[float] = None,
+        metadata: Optional[Dict] = None,
+        class_pdf: Optional[Dict] = None,
+    ):
+        super().__init__(
+            label=label,
+            taxonomy_name=taxonomy_name,
+            reference_id=reference_id,
+            metadata=metadata,
+        )
+        self.confidence = confidence
+        self.class_pdf = class_pdf
+
+    def to_payload(self) -> dict:
+        payload = super().to_payload()
+        if self.confidence is not None:
+            payload[CONFIDENCE_KEY] = self.confidence
+        if self.class_pdf is not None:
+            payload[CLASS_PDF_KEY] = self.class_pdf
+
+        return payload
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        return cls(
+            label=payload.get(LABEL_KEY, 0),
+            taxonomy_name=payload.get(TAXONOMY_NAME_KEY, None),
+            reference_id=payload[REFERENCE_ID_KEY],
+            confidence=payload.get(CONFIDENCE_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
+            class_pdf=payload.get(CLASS_PDF_KEY, None),
+        )

nucleus/utils.py

@@ -34,6 +34,7 @@
 from .dataset_item import DatasetItem
 from .prediction import (
     BoxPrediction,
+    CategoryPrediction,
     CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
@@ -50,6 +51,7 @@ def format_prediction_response(
             BoxPrediction,
             PolygonPrediction,
             CuboidPrediction,
+            CategoryPrediction,
             SegmentationPrediction,
         ]
     ],
@@ -65,12 +67,14 @@ def format_prediction_response(
             Type[BoxPrediction],
             Type[PolygonPrediction],
             Type[CuboidPrediction],
+            Type[CategoryPrediction],
             Type[SegmentationPrediction],
         ],
     ] = {
         BOX_TYPE: BoxPrediction,
         POLYGON_TYPE: PolygonPrediction,
         CUBOID_TYPE: CuboidPrediction,
+        CATEGORY_TYPE: CategoryPrediction,
         SEGMENTATION_TYPE: SegmentationPrediction,
     }
     for type_key in annotation_payload:

tests/helpers.py

@@ -206,6 +206,11 @@ def reference_id_from_url(url):
     for polygon_annotation in TEST_POLYGON_ANNOTATIONS
 }
 
+TEST_CATEGORY_MODEL_PDF = {
+    category_annotation["label"]: 1 / len(TEST_CATEGORY_ANNOTATIONS)
+    for category_annotation in TEST_CATEGORY_ANNOTATIONS
+}
+
 TEST_BOX_PREDICTIONS = [
     {
         **TEST_BOX_ANNOTATIONS[i],
@@ -234,6 +239,20 @@ def reference_id_from_url(url):
     for i in range(len(TEST_POLYGON_ANNOTATIONS))
 ]
 
+TEST_CATEGORY_PREDICTIONS = [
+    {
+        **TEST_CATEGORY_ANNOTATIONS[i],
+        "confidence": 0.10 * i,
+        "class_pdf": TEST_CATEGORY_MODEL_PDF,
+    }
+    if i != 0
+    else {
+        **TEST_CATEGORY_ANNOTATIONS[i],
+        "confidence": 0.10 * i,
+    }
+    for i in range(len(TEST_CATEGORY_ANNOTATIONS))
+]
+
 TEST_INDEX_EMBEDDINGS_FILE = "https://raw.githubusercontent.com/scaleapi/nucleus-python-client/master/tests/testdata/pytest_embeddings_payload.json"
 
 
@@ -339,3 +358,12 @@ def assert_polygon_prediction_matches_dict(
         prediction_instance, prediction_dict
     )
     assert prediction_instance.confidence == prediction_dict["confidence"]
+
+
+def assert_category_prediction_matches_dict(
+    prediction_instance, prediction_dict
+):
+    assert_category_annotation_matches_dict(
+        prediction_instance, prediction_dict
+    )
+    assert prediction_instance.confidence == prediction_dict["confidence"]