Added test for new endpoint, and some eq methods to make that test easier

Author: Ubuntu

nucleus/dataset.py

@@ -32,7 +32,7 @@
     check_all_paths_remote,
     check_for_duplicate_reference_ids,
 )
-from .scene import LidarScene, check_all_scene_paths_remote
+from .scene import LidarScene, Scene, check_all_scene_paths_remote
 from .payload_constructor import (
     construct_append_scenes_payload,
     construct_model_run_creation_payload,
@@ -511,3 +511,17 @@ def delete_annotations(
             self.id, reference_ids, keep_history
         )
         return AsyncJob.from_json(response, self._client)
+
+    def get_scene(self, reference_id) -> Scene:
+        """Returns a scene by reference id
+
+        Returns:
+            a Scene object representing all dataset items organized into frames
+        """
+        return Scene.from_json(
+            self._client.make_request(
+                payload=None,
+                route=f"dataset/{self.id}/scene/{reference_id}",
+                requests_command=requests.get,
+            )
+        )

nucleus/dataset_item.py

@@ -122,24 +122,16 @@ def __post_init__(self):
         )
 
     @classmethod
-    def from_json(cls, payload: dict, is_scene=False):
+    def from_json(cls, payload: dict):
         image_url = payload.get(IMAGE_URL_KEY, None) or payload.get(
             ORIGINAL_IMAGE_URL_KEY, None
         )
-
-        if is_scene:
-            return cls(
-                image_location=image_url,
-                pointcloud_location=payload.get(POINTCLOUD_URL_KEY, None),
-                reference_id=payload.get(REFERENCE_ID_KEY, None),
-                metadata=payload.get(METADATA_KEY, {}),
-            )
-
         return cls(
             image_location=image_url,
+            pointcloud_location=payload.get(POINTCLOUD_URL_KEY, None),
             reference_id=payload.get(REFERENCE_ID_KEY, None),
             metadata=payload.get(METADATA_KEY, {}),
-            upload_to_scale=payload.get(UPLOAD_TO_SCALE_KEY, None),
+            upload_to_scale=payload.get(UPLOAD_TO_SCALE_KEY, True),
         )
 
     def local_file_exists(self):

nucleus/scene.py

@@ -31,6 +31,14 @@ def __post_init__(self):
     def __repr__(self) -> str:
         return f"Frame(items={self.items})"
 
+    def __eq__(self, other):
+        for key, value in self.items.items():
+            if key not in other.items:
+                return False
+            if value != other.items[key]:
+                return False
+        return True
+
     def add_item(self, item: DatasetItem, sensor_name: str):
         self.items[sensor_name] = item
 
@@ -50,7 +58,7 @@ def get_sensors(self):
     @classmethod
     def from_json(cls, payload: dict):
         items = {
-            sensor: DatasetItem.from_json(item, is_scene=True)
+            sensor: DatasetItem.from_json(item)
             for sensor, item in payload.items()
         }
         return cls(**items)
@@ -66,13 +74,24 @@ def to_payload(self) -> dict:
 class Scene(ABC):
     reference_id: str
     frames: List[Frame] = field(default_factory=list)
-    metadata: Optional[dict] = None
+    metadata: Optional[dict] = field(default_factory=dict)
 
     def __post_init__(self):
         self.sensors = set(
             flatten([frame.get_sensors() for frame in self.frames])
         )
         self.frames_dict = dict(enumerate(self.frames))
+        if self.metadata is None:
+            self.metadata = {}
+
+    def __eq__(self, other):
+        return all(
+            [
+                self.reference_id == other.reference_id,
+                self.frames == other.frames,
+                self.metadata == other.metadata,
+            ]
+        )
 
     @property
     def length(self) -> int:
@@ -168,7 +187,7 @@ def from_json(cls, payload: dict):
         return cls(
             reference_id=payload[REFERENCE_ID_KEY],
             frames=frames,
-            metadata=payload.get(METADATA_KEY, None),
+            metadata=payload.get(METADATA_KEY, {}),
         )
 
     def to_payload(self) -> dict:

tests/test_scene.py

@@ -260,6 +260,10 @@ def test_scene_upload_sync(dataset):
 
     response = dataset.append(scenes, update=update)
 
+    first_scene = dataset.get_scene(scenes[0].reference_id)
+
+    assert first_scene == scenes[0]
+
     assert response["dataset_id"] == dataset.id
     assert response["new_scenes"] == len(scenes)