implement append_scenes initial

Author: Drew Kaul

nucleus/constants.py

@@ -71,7 +71,7 @@
 REFERENCE_IDS_KEY = "reference_ids"
 REFERENCE_ID_KEY = "reference_id"
 REQUEST_ID_KEY = "requestId"
-SCENES = "scenes"
+SCENES_KEY = "scenes"
 SEGMENTATIONS_KEY = "segmentations"
 SLICE_ID_KEY = "slice_id"
 STATUS_CODE_KEY = "status_code"

nucleus/dataset.py

@@ -27,7 +27,7 @@
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
-    SCENES,
+    SCENES_KEY,
     UPDATE_KEY,
     URL,
 )
@@ -36,7 +36,11 @@
     check_all_paths_remote,
     check_for_duplicate_reference_ids,
 )
-from .payload_constructor import construct_model_run_creation_payload
+from .scene import LidarScene, check_all_scene_paths_remote
+from .payload_constructor import (
+    construct_append_scenes_payload,
+    construct_model_run_creation_payload,
+)
 
 WARN_FOR_LARGE_UPLOAD = 50000
 
@@ -199,16 +203,16 @@ def ingest_tasks(self, task_ids: dict):
 
     def append(
         self,
-        dataset_items: List[DatasetItem],
+        items: Union[List[DatasetItem], List[LidarScene]],
         update: Optional[bool] = False,
         batch_size: Optional[int] = 20,
         asynchronous=False,
     ) -> Union[dict, AsyncJob]:
         """
-        Appends images with metadata (dataset items) to the dataset. Overwrites images on collision if forced.
+        Appends scenes or images with metadata (dataset items) to the dataset. Overwrites images on collision if forced.
 
         Parameters:
-        :param dataset_items: items to upload
+        :param items: items to upload
         :param update: if True overwrites images and metadata on collision
         :param batch_size: batch parameter for long uploads
         :param aynchronous: if True, return a job object representing asynchronous ingestion job.
@@ -220,9 +224,20 @@ def append(
             'ignored_items': int,
         }
         """
-        check_for_duplicate_reference_ids(dataset_items)
+        all_dataset_items = all(
+            (isinstance(item, DatasetItem) for item in items)
+        )
+        all_scenes = all((isinstance(item, LidarScene) for item in items))
+        if not all_dataset_items and not all_scenes:
+            raise Exception(
+                "You must append either DatasetItems or Scenes to the dataset."
+            )
+        if all_scenes:
+            return self.append_scenes(items, update, asynchronous)
+
+        check_for_duplicate_reference_ids(items)
 
-        if len(dataset_items) > WARN_FOR_LARGE_UPLOAD and not asynchronous:
+        if len(items) > WARN_FOR_LARGE_UPLOAD and not asynchronous:
             print(
                 "Tip: for large uploads, get faster performance by importing your data "
                 "into Nucleus directly from a cloud storage provider. See "
@@ -231,9 +246,9 @@ def append(
             )
 
         if asynchronous:
-            check_all_paths_remote(dataset_items)
+            check_all_paths_remote(items)
             request_id = serialize_and_write_to_presigned_url(
-                dataset_items, self.id, self._client
+                items, self.id, self._client
             )
             response = self._client.make_request(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
@@ -243,11 +258,36 @@ def append(
 
         return self._client.populate_dataset(
             self.id,
-            dataset_items,
+            items,
             update=update,
             batch_size=batch_size,
         )
 
+    def append_scenes(
+        self,
+        scenes: List[LidarScene],
+        update: Optional[bool] = False,
+        asynchronous: Optional[bool] = False,
+    ) -> Union[dict, AsyncJob]:
+        """TODO: Add updated docstring here"""
+        if asynchronous:
+            check_all_scene_paths_remote(scenes)
+            request_id = serialize_and_write_to_presigned_url(
+                scenes, 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)
+
+        payload = construct_append_scenes_payload(scenes, update)
+        response = self._client.make_request(
+            payload=payload,
+            route=f"{self.id}/upload_scenes",
+        )
+        return response
+
     def upload_scenes(
         self,
         payload: dict,
@@ -268,7 +308,7 @@ def upload_scenes(
         }
         """
         if asynchronous:
-            for scene in payload[SCENES]:
+            for scene in payload[SCENES_KEY]:
                 for frame in scene[FRAMES]:
                     check_all_frame_paths_remote(frame[URL])
             request_id = serialize_and_write_to_presigned_url(

nucleus/payload_constructor.py

@@ -1,5 +1,6 @@
 from typing import List, Optional, Dict, Union
 from .dataset_item import DatasetItem
+from .scene import LidarScene
 from .annotation import (
     BoxAnnotation,
     CuboidAnnotation,
@@ -19,6 +20,7 @@
     REFERENCE_ID_KEY,
     ANNOTATIONS_KEY,
     ITEMS_KEY,
+    SCENES_KEY,
     UPDATE_KEY,
     MODEL_ID_KEY,
     ANNOTATION_METADATA_SCHEMA_KEY,
@@ -40,6 +42,15 @@ def construct_append_payload(
     )
 
 
+def construct_append_scenes_payload(
+    scene_list: List[LidarScene], update: Optional[bool] = False
+) -> dict:
+    scenes = []
+    for scene in scene_list:
+        scenes.append(scene.to_payload())
+    return {SCENES_KEY: scenes, UPDATE_KEY: update}
+
+
 def construct_annotation_payload(
     annotation_items: List[
         Union[

nucleus/scene.py

@@ -24,7 +24,7 @@
     Z_KEY,
 )
 from .annotation import Point3D
-from .utils import flatten
+from .dataset_item import is_local_path
 
 
 class DatasetItemType(Enum):
@@ -257,3 +257,18 @@ def validate(self):
             assert (
                 num_pointclouds == 1
             ), "Each frame of a lidar scene must have exactly 1 pointcloud"
+
+
+def flatten(t):
+    return [item for sublist in t for item in sublist]
+
+
+def check_all_scene_paths_remote(scenes: List[LidarScene]):
+    for scene in scenes:
+        for frame in scene.frames_dict.values():
+            for item in frame.items.values():
+                if is_local_path(getattr(item, URL_KEY)):
+                    raise ValueError(
+                        f"All paths for SceneDatasetItems must be remote, but {item.url} is either "
+                        "local, or a remote URL type that is not supported."
+                    )

nucleus/utils.py

@@ -29,16 +29,13 @@
 )
 from .dataset_item import DatasetItem
 from .prediction import BoxPrediction, CuboidPrediction, PolygonPrediction
+from .scene import LidarScene
 
 
 def _get_all_field_values(metadata_list: List[dict], key: str):
     return {metadata[key] for metadata in metadata_list if key in metadata}
 
 
-def flatten(t):
-    return [item for sublist in t for item in sublist]
-
-
 def suggest_metadata_schema(
     data: Union[
         List[DatasetItem],
@@ -125,7 +122,7 @@ def convert_export_payload(api_payload):
 
 
 def serialize_and_write(
-    upload_units: Sequence[Union[DatasetItem, Annotation, Dict]], file_pointer
+    upload_units: Sequence[Union[DatasetItem, Annotation, LidarScene]], file_pointer
 ):
     for unit in upload_units:
         try:
@@ -159,7 +156,7 @@ def upload_to_presigned_url(presigned_url: str, file_pointer: IO):
 
 
 def serialize_and_write_to_presigned_url(
-    upload_units: Sequence[Union[DatasetItem, Annotation, Dict]],
+    upload_units: Sequence[Union[DatasetItem, Annotation, LidarScene]],
     dataset_id: str,
     client,
 ):

tests/test_scene.py

@@ -1,4 +1,4 @@
-from nucleus.constants import SCENES
+from nucleus.constants import SCENES_KEY
 import pytest
 
 from .helpers import (
@@ -27,7 +27,7 @@ def test_scene_upload_sync(dataset):
     response = dataset.upload_scenes(payload)
 
     assert response["dataset_id"] == dataset.id
-    assert response["new_scenes"] == len(TEST_LIDAR_SCENES[SCENES])
+    assert response["new_scenes"] == len(TEST_LIDAR_SCENES[SCENES_KEY])
 
 
 @pytest.mark.integration
@@ -36,7 +36,7 @@ def test_scene_and_cuboid_upload_sync(dataset):
     response = dataset.upload_scenes(payload)
 
     assert response["dataset_id"] == dataset.id
-    assert response["new_scenes"] == len(TEST_LIDAR_SCENES[SCENES])
+    assert response["new_scenes"] == len(TEST_LIDAR_SCENES[SCENES_KEY])
 
     TEST_CUBOID_ANNOTATIONS[0]["dataset_item_id"] = dataset.items[0].item_id
     annotations = [CuboidAnnotation.from_json(TEST_CUBOID_ANNOTATIONS[0])]