[Nucleus] Scene categorization support (#343)

Author: Anthony Krivonos

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.14.16](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.15) - 2022-08-12
+
+### Added
+- Scene cateogorization support
+
 ## [0.14.15](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.15) - 2022-08-11
 
 ### Removed

nucleus/__init__.py

@@ -31,6 +31,8 @@
     "PolygonAnnotation",
     "PolygonPrediction",
     "Quaternion",
+    "SceneCategoryAnnotation",
+    "SceneCategoryPrediction",
     "Segment",
     "SegmentationAnnotation",
     "SegmentationPrediction",
@@ -61,6 +63,7 @@
     Point,
     Point3D,
     PolygonAnnotation,
+    SceneCategoryAnnotation,
     Segment,
     SegmentationAnnotation,
 )
@@ -132,6 +135,7 @@
     KeypointsPrediction,
     LinePrediction,
     PolygonPrediction,
+    SceneCategoryPrediction,
     SegmentationPrediction,
 )
 from .quaternion import Quaternion

nucleus/annotation.py

@@ -41,6 +41,9 @@
     Z_KEY,
 )
 
+# TODO: refactor to reduce this file to under 1000 lines.
+# pylint: disable=C0302
+
 
 class Annotation:
     """Internal base class, not to be used directly.
@@ -893,6 +896,53 @@ def to_payload(self) -> dict:
         return payload
 
 
+@dataclass
+class SceneCategoryAnnotation(Annotation):
+    """A scene category annotation.
+
+    ::
+
+        from nucleus import CategoryAnnotation
+
+        category = SceneCategoryAnnotation(
+            label="running",
+            reference_id="scene_1",
+            taxonomy_name="action",
+        )
+
+    Parameters:
+        label (str): The label for this annotation.
+        reference_id (str): User-defined ID of the scene to which to apply this annotation.
+        taxonomy_name (Optional[str]): The name of the taxonomy this annotation conforms to.
+          See :meth:`Dataset.add_taxonomy`.
+    """
+
+    label: str
+    reference_id: str
+    taxonomy_name: Optional[str] = None
+    # todo(546247): add metadata support when required
+    metadata: Optional[Dict] = field(default_factory=dict)
+
+    @classmethod
+    def from_json(cls, payload: dict):
+        return cls(
+            label=payload[LABEL_KEY],
+            reference_id=payload[REFERENCE_ID_KEY],
+            taxonomy_name=payload.get(TAXONOMY_NAME_KEY, None),
+        )
+
+    def to_payload(self) -> dict:
+        payload = {
+            LABEL_KEY: self.label,
+            TYPE_KEY: CATEGORY_TYPE,
+            GEOMETRY_KEY: {},
+            REFERENCE_ID_KEY: self.reference_id,
+        }
+        if self.taxonomy_name is not None:
+            payload[TAXONOMY_NAME_KEY] = self.taxonomy_name
+        return payload
+
+
 @dataclass
 class AnnotationList:
     """Wrapper class separating a list of annotations by type."""
@@ -910,6 +960,9 @@ class AnnotationList:
     multi_category_annotations: List[MultiCategoryAnnotation] = field(
         default_factory=list
     )
+    scene_category_annotations: List[SceneCategoryAnnotation] = field(
+        default_factory=list
+    )
     segmentation_annotations: List[SegmentationAnnotation] = field(
         default_factory=list
     )
@@ -934,6 +987,8 @@ def add_annotations(self, annotations: List[Annotation]):
                 self.category_annotations.append(annotation)
             elif isinstance(annotation, MultiCategoryAnnotation):
                 self.multi_category_annotations.append(annotation)
+            elif isinstance(annotation, SceneCategoryAnnotation):
+                self.scene_category_annotations.append(annotation)
             else:
                 assert isinstance(
                     annotation, SegmentationAnnotation
@@ -952,6 +1007,7 @@ def __len__(self):
             + len(self.cuboid_annotations)
             + len(self.category_annotations)
             + len(self.multi_category_annotations)
+            + len(self.scene_category_annotations)
             + len(self.segmentation_annotations)
         )
 

nucleus/dataset.py

@@ -10,6 +10,7 @@
     CategoryPrediction,
     CuboidPrediction,
     PolygonPrediction,
+    SceneCategoryPrediction,
     SegmentationPrediction,
     from_json,
 )
@@ -345,8 +346,8 @@ def annotate(
 
         Nucleus supports :class:`Box<BoxAnnotation>`, :class:`Polygon<PolygonAnnotation>`,
         :class:`Cuboid<CuboidAnnotation>`, :class:`Segmentation<SegmentationAnnotation>`,
-        and :class:`Category<CategoryAnnotation>` annotations. Cuboid annotations
-        can only be uploaded to a :class:`pointcloud DatasetItem<LidarScene>`.
+        :class:`Category<CategoryAnnotation>`, and :class:`Category<SceneCategoryAnnotation>` annotations.
+        Cuboid annotations can only be uploaded to a :class:`pointcloud DatasetItem<LidarScene>`.
 
         When uploading an annotation, you need to specify which item you are
         annotating via the reference_id you provided when uploading the image
@@ -1474,6 +1475,7 @@ def upload_predictions(
                 CuboidPrediction,
                 SegmentationPrediction,
                 CategoryPrediction,
+                SceneCategoryPrediction,
             ]
         ],
         update: bool = False,
@@ -1492,7 +1494,7 @@ def upload_predictions(
 
         Nucleus supports :class:`Box<BoxPrediction>`, :class:`Polygon<PolygonPrediction>`,
         :class:`Cuboid<CuboidPrediction>`, :class:`Segmentation<SegmentationPrediction>`,
-        and :class:`Category<CategoryPrediction>` predictions. Cuboid predictions
+        :class:`Category<CategoryPrediction>`, and :class:`Category<SceneCategoryPrediction>` predictions. Cuboid predictions
         can only be uploaded to a :class:`pointcloud DatasetItem<LidarScene>`.
 
         When uploading an prediction, you need to specify which item you are
@@ -1519,6 +1521,7 @@ def upload_predictions(
                 :class:`CuboidPrediction`, \
                 :class:`SegmentationPrediction`, \
                 :class:`CategoryPrediction` \
+                :class:`SceneCategoryPrediction` \
             ]]): List of prediction objects to upload.
             update: Whether or not to overwrite metadata or ignore on reference ID
               collision. Default is False.

nucleus/metrics/filtering.py

@@ -36,6 +36,7 @@
     LinePrediction,
     PolygonPrediction,
     PredictionList,
+    SceneCategoryPrediction,
     SegmentationPrediction,
 )
 
@@ -90,6 +91,7 @@ class FilterType(str, enum.Enum):
     CuboidPrediction,
     LinePrediction,
     PolygonPrediction,
+    SceneCategoryPrediction,
     SegmentationPrediction,
 ]
 

nucleus/payload_constructor.py

@@ -6,6 +6,7 @@
     CuboidAnnotation,
     MultiCategoryAnnotation,
     PolygonAnnotation,
+    SceneCategoryAnnotation,
     SegmentationAnnotation,
 )
 from .constants import (
@@ -32,6 +33,7 @@
     CategoryPrediction,
     CuboidPrediction,
     PolygonPrediction,
+    SceneCategoryPrediction,
     SegmentationPrediction,
 )
 from .scene import LidarScene, VideoScene
@@ -69,6 +71,7 @@ def construct_annotation_payload(
             CuboidAnnotation,
             CategoryAnnotation,
             MultiCategoryAnnotation,
+            SceneCategoryAnnotation,
             SegmentationAnnotation,
         ]
     ],
@@ -112,6 +115,7 @@ def construct_box_predictions_payload(
             PolygonPrediction,
             CuboidPrediction,
             CategoryPrediction,
+            SceneCategoryPrediction,
         ]
     ],
     update: bool,

nucleus/prediction.py

@@ -16,6 +16,7 @@
     Point,
     Point3D,
     PolygonAnnotation,
+    SceneCategoryAnnotation,
     Segment,
     SegmentationAnnotation,
 )
@@ -571,13 +572,25 @@ def from_json(cls, payload: dict):
         )
 
 
+class SceneCategoryPrediction(SceneCategoryAnnotation):
+    """A prediction of a category for a scene.
+
+    Parameters:
+        label: The label for this annotation (e.g. action, subject, scenario).
+        reference_id: The reference ID of the scene you wish to apply this annotation to.
+        taxonomy_name: The name of the taxonomy this annotation conforms to.
+          See :meth:`Dataset.add_taxonomy`.
+    """
+
+
 Prediction = Union[
     BoxPrediction,
     LinePrediction,
     PolygonPrediction,
     KeypointsPrediction,
     CuboidPrediction,
     CategoryPrediction,
+    SceneCategoryPrediction,
     SegmentationPrediction,
 ]
 
@@ -596,6 +609,9 @@ class PredictionList:
     category_predictions: List[CategoryPrediction] = field(
         default_factory=list
     )
+    scene_category_predictions: List[SceneCategoryPrediction] = field(
+        default_factory=list
+    )
     segmentation_predictions: List[SegmentationPrediction] = field(
         default_factory=list
     )
@@ -617,6 +633,8 @@ def add_predictions(self, predictions: List[Prediction]):
                 self.cuboid_predictions.append(prediction)
             elif isinstance(prediction, CategoryPrediction):
                 self.category_predictions.append(prediction)
+            elif isinstance(prediction, SceneCategoryPrediction):
+                self.scene_category_predictions.append(prediction)
             else:
                 assert isinstance(
                     prediction, SegmentationPrediction
@@ -631,5 +649,6 @@ def __len__(self):
             + len(self.keypoints_predictions)
             + len(self.cuboid_predictions)
             + len(self.category_predictions)
+            + len(self.scene_category_predictions)
             + len(self.segmentation_predictions)
         )

nucleus/utils.py

@@ -52,6 +52,7 @@
     KeypointsPrediction,
     LinePrediction,
     PolygonPrediction,
+    SceneCategoryPrediction,
     SegmentationPrediction,
 )
 from .scene import LidarScene, VideoScene
@@ -112,6 +113,7 @@ def format_prediction_response(
             KeypointsPrediction,
             CuboidPrediction,
             CategoryPrediction,
+            SceneCategoryPrediction,
             SegmentationPrediction,
         ]
     ],
@@ -138,6 +140,7 @@ def format_prediction_response(
             Type[CuboidPrediction],
             Type[CategoryPrediction],
             Type[KeypointsPrediction],
+            Type[SceneCategoryPrediction],
             Type[SegmentationPrediction],
         ],
     ] = {

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.14.15"
+version = "0.14.16"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]