Merge pull request #84 from scaleapi/upload_scene_integration_test

Upload scene integration test

Author: sasha-scale

nucleus/__init__.py

@@ -615,21 +615,23 @@ def annotate_dataset(
         self,
         dataset_id: str,
         annotations: List[
-            Union[BoxAnnotation, PolygonAnnotation, SegmentationAnnotation]
+            Union[
+                BoxAnnotation,
+                PolygonAnnotation,
+                CuboidAnnotation,
+                SegmentationAnnotation,
+            ]
         ],
         update: bool,
         batch_size: int = 5000,
     ):
         """
         Uploads ground truth annotations for a given dataset.
         :param dataset_id: id of the dataset
-        :param annotations: List[Union[BoxAnnotation, PolygonAnnotation]]
+        :param annotations: List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]]
         :param update: whether to update or ignore conflicting annotations
         :return: {"dataset_id: str, "annotations_processed": int}
         """
-        if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
-            raise NotImplementedError("Cuboid annotations not yet supported")
-
         # Split payload into segmentations and Box/Polygon
         segmentations = [
             ann
@@ -770,14 +772,19 @@ def predict(
         self,
         model_run_id: str,
         annotations: List[
-            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
         ],
         update: bool,
         batch_size: int = 5000,
     ):
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
-        :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        :param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         :param update: bool
         :return:
         {
@@ -787,9 +794,6 @@ def predict(
             "predictions_ignored": int,
         }
         """
-        if any((isinstance(ann, CuboidPrediction) for ann in annotations)):
-            raise NotImplementedError("Cuboid predictions not yet supported")
-
         segmentations = [
             ann
             for ann in annotations
@@ -941,7 +945,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[BoxPrediction],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         return self.make_request(
@@ -966,7 +970,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[BoxPrediction],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         return self.make_request(
@@ -995,7 +999,7 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
         :param dataset_item_id: dataset_item_id of a dataset item.
         :return:
         {
-            "annotations": List[BoxPrediction],
+            "annotations": List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         return self.make_request(

nucleus/annotation.py

@@ -299,16 +299,26 @@ def to_payload(self) -> dict:
             LABEL_KEY: self.label,
             TYPE_KEY: CUBOID_TYPE,
             GEOMETRY_KEY: {
-                POSITION_KEY: self.position,
-                DIMENSIONS_KEY: self.dimensions,
+                POSITION_KEY: self.position.to_payload(),
+                DIMENSIONS_KEY: self.dimensions.to_payload(),
                 YAW_KEY: self.yaw,
             },
             REFERENCE_ID_KEY: self.reference_id,
+            ITEM_ID_KEY: self.item_id,
             ANNOTATION_ID_KEY: self.annotation_id,
             METADATA_KEY: self.metadata,
         }
 
 
+def check_all_frame_paths_remote(frames: List[str]):
+    for frame_url in frames:
+        if is_local_path(frame_url):
+            raise ValueError(
+                f"All paths must be remote, but {frame_url} is either "
+                "local, or a remote URL type that is not supported."
+            )
+
+
 def check_all_mask_paths_remote(
     annotations: Sequence[Union[Annotation]],
 ):

nucleus/constants.py

@@ -30,6 +30,7 @@
 ERROR_CODES = "error_codes"
 ERROR_ITEMS = "upload_errors"
 ERROR_PAYLOAD = "error_payload"
+FRAMES = "frames"
 GEOMETRY_KEY = "geometry"
 HEIGHT_KEY = "height"
 IGNORED_ITEMS = "ignored_items"
@@ -62,6 +63,7 @@
 REFERENCE_IDS_KEY = "reference_ids"
 REFERENCE_ID_KEY = "reference_id"
 REQUEST_ID_KEY = "requestId"
+SCENES = "scenes"
 SEGMENTATIONS_KEY = "segmentations"
 SLICE_ID_KEY = "slice_id"
 STATUS_CODE_KEY = "status_code"
@@ -70,6 +72,7 @@
 TYPE_KEY = "type"
 UPDATED_ITEMS = "updated_items"
 UPDATE_KEY = "update"
+URL = "url"
 VERTICES_KEY = "vertices"
 WIDTH_KEY = "width"
 YAW_KEY = "yaw"

nucleus/dataset.py

@@ -12,8 +12,8 @@
 
 from .annotation import (
     Annotation,
-    CuboidAnnotation,
     check_all_mask_paths_remote,
+    check_all_frame_paths_remote,
 )
 from .constants import (
     DATASET_ITEM_IDS_KEY,
@@ -23,10 +23,13 @@
     DATASET_SLICES_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     EXPORTED_ROWS,
+    FRAMES,
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    SCENES,
     UPDATE_KEY,
+    URL,
 )
 from .dataset_item import (
     DatasetItem,
@@ -167,9 +170,6 @@ def annotate(
             "ignored_items": int,
         }
         """
-        if any((isinstance(ann, CuboidAnnotation) for ann in annotations)):
-            raise NotImplementedError("Cuboid annotations not yet supported")
-
         check_all_mask_paths_remote(annotations)
 
         if asynchronous:
@@ -248,14 +248,53 @@ def append(
             batch_size=batch_size,
         )
 
+    def upload_scenes(
+        self,
+        payload: dict,
+        update: Optional[bool] = False,
+        asynchronous: bool = False,
+    ) -> Union[dict, AsyncJob]:
+        """
+        Uploads scenes with given frames to the dataset
+
+        Parameters:
+        :param payload: dictionary containing scenes to be uploaded
+        :param update: if True, overwrite scene on collision
+        :param aynchronous: if True, return a job object representing asynchronous ingestion job
+        :return:
+        {
+            'dataset_id': str,
+            'new_scenes': int,
+        }
+        """
+        if asynchronous:
+            for scene in payload[SCENES]:
+                for frame in scene[FRAMES]:
+                    check_all_frame_paths_remote(frame[URL])
+            request_id = serialize_and_write_to_presigned_url(
+                [payload], self.id, self._client
+            )
+            response = self._client.make_request(
+                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                route=f"{self.id}/upload_scenes?async=1",
+            )
+            return AsyncJob.from_json(response, self._client)
+
+        # TODO: create client method for sync scene upload
+        response = self._client.make_request(
+            payload=payload,
+            route=f"{self.id}/upload_scenes",
+        )
+        return response
+
     def iloc(self, i: int) -> dict:
         """
         Returns Dataset Item Info By Dataset Item Number.
         :param i: absolute number of dataset item for the given dataset.
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
         }
         """
         response = self._client.dataitem_iloc(self.id, i)
@@ -268,7 +307,7 @@ def refloc(self, reference_id: str) -> dict:
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
         }
         """
         response = self._client.dataitem_ref_id(self.id, reference_id)
@@ -281,7 +320,7 @@ def loc(self, dataset_item_id: str) -> dict:
         :return:
         {
             "item": DatasetItem,
-            "annotations": List[Union[BoxAnnotation, PolygonAnnotation]],
+            "annotations": List[Union[BoxAnnotation, PolygonAnnotation, CuboidAnnotation, SegmentationAnnotation]],
         }
         """
         response = self._client.dataitem_loc(self.id, dataset_item_id)

nucleus/model.py

@@ -2,6 +2,7 @@
 from .dataset import Dataset
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
@@ -57,7 +58,12 @@ def create_run(
         name: str,
         dataset: Dataset,
         predictions: List[
-            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
         ],
         metadata: Optional[Dict] = None,
         asynchronous: bool = False,

nucleus/model_run.py

@@ -7,6 +7,7 @@
 from .constants import (
     ANNOTATIONS_KEY,
     BOX_TYPE,
+    CUBOID_TYPE,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     POLYGON_TYPE,
     REQUEST_ID_KEY,
@@ -15,6 +16,7 @@
 )
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
@@ -89,14 +91,19 @@ def commit(self, payload: Optional[dict] = None) -> dict:
     def predict(
         self,
         annotations: List[
-            Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
         ],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
         asynchronous: bool = False,
     ) -> Union[dict, AsyncJob]:
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
-        :param annotations: List[Union[BoxPrediction, PolygonPrediction]],
+        :param annotations: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         :return:
         {
             "model_run_id": str,
@@ -122,7 +129,7 @@ def iloc(self, i: int):
         """
         Returns Model Run Info For Dataset Item by its number.
         :param i: absolute number of Dataset Item for a dataset corresponding to the model run.
-        :return: List[Union[BoxPrediction, PolygonPrediction]],
+        :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         }
         """
         response = self._client.predictions_iloc(self.model_run_id, i)
@@ -132,7 +139,7 @@ def refloc(self, reference_id: str):
         """
         Returns Model Run Info For Dataset Item by its reference_id.
         :param reference_id: reference_id of a dataset item.
-        :return: List[Union[BoxPrediction, PolygonPrediction]],
+        :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         """
         response = self._client.predictions_ref_id(
             self.model_run_id, reference_id
@@ -167,7 +174,14 @@ def _format_prediction_response(
         self, response: dict
     ) -> Union[
         dict,
-        List[Union[BoxPrediction, PolygonPrediction, SegmentationPrediction]],
+        List[
+            Union[
+                BoxPrediction,
+                PolygonPrediction,
+                CuboidPrediction,
+                SegmentationPrediction,
+            ]
+        ],
     ]:
         annotation_payload = response.get(ANNOTATIONS_KEY, None)
         if not annotation_payload:
@@ -179,11 +193,13 @@ def _format_prediction_response(
             Union[
                 Type[BoxPrediction],
                 Type[PolygonPrediction],
+                Type[CuboidPrediction],
                 Type[SegmentationPrediction],
             ],
         ] = {
             BOX_TYPE: BoxPrediction,
             POLYGON_TYPE: PolygonPrediction,
+            CUBOID_TYPE: CuboidPrediction,
             SEGMENTATION_TYPE: SegmentationPrediction,
         }
         for type_key in annotation_payload:

nucleus/payload_constructor.py

@@ -2,11 +2,13 @@
 from .dataset_item import DatasetItem
 from .annotation import (
     BoxAnnotation,
+    CuboidAnnotation,
     PolygonAnnotation,
     SegmentationAnnotation,
 )
 from .prediction import (
     BoxPrediction,
+    CuboidPrediction,
     PolygonPrediction,
     SegmentationPrediction,
 )
@@ -39,7 +41,14 @@ def construct_append_payload(
 
 
 def construct_annotation_payload(
-    annotation_items: List[Union[BoxAnnotation, PolygonAnnotation]],
+    annotation_items: List[
+        Union[
+            BoxAnnotation,
+            PolygonAnnotation,
+            CuboidAnnotation,
+            SegmentationAnnotation,
+        ]
+    ],
     update: bool,
 ) -> dict:
     annotations = []
@@ -63,7 +72,9 @@ def construct_segmentation_payload(
 
 
 def construct_box_predictions_payload(
-    box_predictions: List[Union[BoxPrediction, PolygonPrediction]],
+    box_predictions: List[
+        Union[BoxPrediction, PolygonPrediction, CuboidPrediction]
+    ],
     update: bool,
 ) -> dict:
     predictions = []