Merge pull request #101 from scaleapi/add_getters_to_scene_class

sasha-scale · web-flow · commit 203e22ab16b3 · 2021-08-17T21:07:27.000-07:00
Add getters to scene class
diff --git a/nucleus/annotation.py b/nucleus/annotation.py
@@ -295,19 +295,23 @@ def from_json(cls, payload: dict):
         )
 
     def to_payload(self) -> dict:
-        return {
+        payload = {
             LABEL_KEY: self.label,
             TYPE_KEY: CUBOID_TYPE,
             GEOMETRY_KEY: {
                 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,
         }
+        if self.reference_id:
+            payload[REFERENCE_ID_KEY] = self.reference_id
+        if self.annotation_id:
+            payload[ANNOTATION_ID_KEY] = self.annotation_id
+        if self.metadata:
+            payload[METADATA_KEY] = self.metadata
+
+        return payload
 
 
 def is_local_path(path: str) -> bool:
diff --git a/nucleus/constants.py b/nucleus/constants.py
@@ -52,6 +52,7 @@
 ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
 JOB_ID_KEY = "job_id"
 KEEP_HISTORY_KEY = "keep_history"
+LENGTH_KEY = "length"
 JOB_STATUS_KEY = "job_status"
 JOB_LAST_KNOWN_STATUS_KEY = "job_last_known_status"
 JOB_TYPE_KEY = "job_type"
@@ -65,7 +66,9 @@
 NAME_KEY = "name"
 NEW_ITEMS = "new_items"
 NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
+NUM_SENSORS_KEY = "num_sensors"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
+POINTCLOUD_KEY = "pointcloud"
 POINTCLOUD_LOCATION_KEY = "pointcloud_location"
 POINTCLOUD_URL_KEY = "pointcloud_url"
 POSITION_KEY = "position"
diff --git a/nucleus/dataset_item.py b/nucleus/dataset_item.py
@@ -96,10 +96,12 @@ class DatasetItem:  # pylint: disable=R0902
     pointcloud_location: Optional[str] = None
 
     def __post_init__(self):
-        self.local = is_local_path(self.image_location)
         assert bool(self.image_location) != bool(
             self.pointcloud_location
         ), "Must specify exactly one of the image_location, pointcloud_location parameters"
+        self.local = (
+            is_local_path(self.image_location) if self.image_location else None
+        )
         self.type = (
             DatasetItemType.IMAGE
             if self.image_location
@@ -116,14 +118,14 @@ def __post_init__(self):
 
     @classmethod
     def from_json(cls, payload: dict, is_scene=False):
-        image_url = payload.get(IMAGE_URL_KEY, "") or payload.get(
-            ORIGINAL_IMAGE_URL_KEY, ""
+        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, ""),
+                pointcloud_location=payload.get(POINTCLOUD_URL_KEY, None),
                 reference_id=payload.get(REFERENCE_ID_KEY, None),
                 item_id=payload.get(DATASET_ITEM_ID_KEY, None),
                 metadata=payload.get(METADATA_KEY, {}),
@@ -143,26 +145,25 @@ def to_payload(self, is_scene=False) -> dict:
         payload: Dict[str, Any] = {
             METADATA_KEY: self.metadata or {},
         }
+        if self.reference_id:
+            payload[REFERENCE_ID_KEY] = self.reference_id
+        if self.item_id:
+            payload[DATASET_ITEM_ID_KEY] = self.item_id
 
         if is_scene:
             if self.image_location:
                 payload[URL_KEY] = self.image_location
             elif self.pointcloud_location:
                 payload[URL_KEY] = self.pointcloud_location
             payload[TYPE_KEY] = self.type.value
+            if self.camera_params:
+                payload[CAMERA_PARAMS_KEY] = self.camera_params.to_payload()
         else:
             assert (
                 self.image_location
-            ), "Must specify image_location for DatasetItems not in a Scene"
+            ), "Must specify image_location for DatasetItems not in a LidarScene"
             payload[IMAGE_URL_KEY] = self.image_location
 
-        if self.reference_id:
-            payload[REFERENCE_ID_KEY] = self.reference_id
-        if self.item_id:
-            payload[DATASET_ITEM_ID_KEY] = self.item_id
-        if self.camera_params:
-            payload[CAMERA_PARAMS_KEY] = self.camera_params.to_payload()
-
         return payload
 
     def to_json(self) -> str:
diff --git a/nucleus/scene.py b/nucleus/scene.py
@@ -4,7 +4,9 @@
 from typing import Optional, Union, Any, Dict, List
 from nucleus.constants import (
     FRAMES_KEY,
+    LENGTH_KEY,
     METADATA_KEY,
+    NUM_SENSORS_KEY,
     REFERENCE_ID_KEY,
     POINTCLOUD_LOCATION_KEY,
     IMAGE_LOCATION_KEY,
@@ -25,9 +27,28 @@ def __post_init__(self):
                 value, DatasetItem
             ), "All values must be DatasetItems"
 
+    def __repr__(self) -> str:
+        return f"Frame(index={self.index}, items={self.items})"
+
     def add_item(self, item: DatasetItem, sensor_name: str):
         self.items[sensor_name] = item
 
+    def get_item(self, sensor_name: str):
+        if sensor_name not in self.items:
+            raise ValueError(
+                f"This frame does not have a {sensor_name} sensor"
+            )
+        return self.items[sensor_name]
+
+    def get_items(self):
+        return list(self.items.values())
+
+    def get_sensors(self):
+        return list(self.items.keys())
+
+    def get_index(self):
+        return self.index
+
     @classmethod
     def from_json(cls, payload: dict):
         items = {
@@ -51,6 +72,9 @@ class Scene(ABC):
 
     def __post_init__(self):
         self.check_valid_frame_indices()
+        self.sensors = set(
+            flatten([frame.get_sensors() for frame in self.frames])
+        )
         if all((frame.index is not None for frame in self.frames)):
             self.frames_dict = {frame.index: frame for frame in self.frames}
         else:
@@ -60,6 +84,14 @@ def __post_init__(self):
             ]
             self.frames_dict = dict(enumerate(indexed_frames))
 
+    @property
+    def length(self) -> int:
+        return len(self.frames_dict)
+
+    @property
+    def num_sensors(self) -> int:
+        return len(self.get_sensors())
+
     def check_valid_frame_indices(self):
         infer_from_list_position = all(
             (frame.index is None for frame in self.frames)
@@ -72,15 +104,14 @@ def check_valid_frame_indices(self):
         ), "Must specify index explicitly for all frames or infer from list position for all frames"
 
     def validate(self):
-        assert (
-            len(self.frames_dict) > 0
-        ), "Must have at least 1 frame in a scene"
+        assert self.length > 0, "Must have at least 1 frame in a scene"
         for frame in self.frames_dict.values():
             assert isinstance(
                 frame, Frame
             ), "Each frame in a scene must be a Frame object"
 
     def add_item(self, index: int, sensor_name: str, item: DatasetItem):
+        self.sensors.add(sensor_name)
         if index not in self.frames_dict:
             new_frame = Frame(index=index, items={sensor_name: item})
             self.frames_dict[index] = new_frame
@@ -97,6 +128,54 @@ def add_frame(self, frame: Frame, update: bool = False):
             and update
         ):
             self.frames_dict[frame.index] = frame
+            self.sensors.update(frame.get_sensors())
+
+    def get_frame(self, index: int):
+        if index not in self.frames_dict:
+            raise ValueError(
+                f"This scene does not have a frame at index {index}"
+            )
+        return self.frames_dict[index]
+
+    def get_frames(self):
+        return [
+            frame
+            for _, frame in sorted(
+                self.frames_dict.items(), key=lambda x: x[0]
+            )
+        ]
+
+    def get_sensors(self):
+        return list(self.sensors)
+
+    def get_item(self, index: int, sensor_name: str):
+        frame = self.get_frame(index)
+        return frame.get_item(sensor_name)
+
+    def get_items_from_sensor(self, sensor_name: str):
+        if sensor_name not in self.sensors:
+            raise ValueError(
+                f"This scene does not have a {sensor_name} sensor"
+            )
+        items_from_sensor = []
+        for frame in self.frames_dict.values():
+            try:
+                sensor_item = frame.get_item(sensor_name)
+                items_from_sensor.append(sensor_item)
+            except ValueError:
+                # This sensor is not present at current frame
+                items_from_sensor.append(None)
+        return items_from_sensor
+
+    def get_items(self):
+        return flatten([frame.get_items() for frame in self.get_frames()])
+
+    def info(self):
+        return {
+            REFERENCE_ID_KEY: self.reference_id,
+            LENGTH_KEY: self.length,
+            NUM_SENSORS_KEY: self.num_sensors,
+        }
 
     def validate_frames_dict(self):
         is_continuous = set(list(range(len(self.frames_dict)))) == set(
@@ -118,12 +197,7 @@ def from_json(cls, payload: dict):
 
     def to_payload(self) -> dict:
         self.validate_frames_dict()
-        ordered_frames = [
-            frame
-            for _, frame in sorted(
-                self.frames_dict.items(), key=lambda x: x[0]
-            )
-        ]
+        ordered_frames = self.get_frames()
         frames_payload = [frame.to_payload() for frame in ordered_frames]
         payload: Dict[str, Any] = {
             REFERENCE_ID_KEY: self.reference_id,
@@ -139,27 +213,30 @@ def to_json(self) -> str:
 
 @dataclass
 class LidarScene(Scene):
+    def __repr__(self) -> str:
+        return f"LidarScene(reference_id='{self.reference_id}', frames={self.get_frames()}, metadata={self.metadata})"
+
     def validate(self):
         super().validate()
-        lidar_sources = flatten(
+        lidar_sensors = flatten(
             [
                 [
-                    source
-                    for source in frame.items.keys()
-                    if frame.items[source].type == DatasetItemType.POINTCLOUD
+                    sensor
+                    for sensor in frame.items.keys()
+                    if frame.items[sensor].type == DatasetItemType.POINTCLOUD
                 ]
                 for frame in self.frames_dict.values()
             ]
         )
         assert (
-            len(set(lidar_sources)) == 1
-        ), "Each lidar scene must have exactly one lidar source"
+            len(set(lidar_sensors)) == 1
+        ), "Each lidar scene must have exactly one lidar sensor"
 
         for frame in self.frames_dict.values():
             num_pointclouds = sum(
                 [
                     int(item.type == DatasetItemType.POINTCLOUD)
-                    for item in frame.items.values()
+                    for item in frame.get_items()
                 ]
             )
             assert (
@@ -173,17 +250,16 @@ def flatten(t):
 
 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():
-                pointcloud_location = getattr(item, POINTCLOUD_LOCATION_KEY)
-                if pointcloud_location and is_local_path(pointcloud_location):
-                    raise ValueError(
-                        f"All paths for DatasetItems in a Scene must be remote, but {item.pointcloud_location} is either "
-                        "local, or a remote URL type that is not supported."
-                    )
-                image_location = getattr(item, IMAGE_LOCATION_KEY)
-                if image_location and is_local_path(image_location):
-                    raise ValueError(
-                        f"All paths for DatasetItems in a Scene must be remote, but {item.image_location} is either "
-                        "local, or a remote URL type that is not supported."
-                    )
+        for item in scene.get_items():
+            pointcloud_location = getattr(item, POINTCLOUD_LOCATION_KEY)
+            if pointcloud_location and is_local_path(pointcloud_location):
+                raise ValueError(
+                    f"All paths for DatasetItems in a Scene must be remote, but {item.pointcloud_location} is either "
+                    "local, or a remote URL type that is not supported."
+                )
+            image_location = getattr(item, IMAGE_LOCATION_KEY)
+            if image_location and is_local_path(image_location):
+                raise ValueError(
+                    f"All paths for DatasetItems in a Scene must be remote, but {item.image_location} is either "
+                    "local, or a remote URL type that is not supported."
+                )
diff --git a/tests/helpers.py b/tests/helpers.py
@@ -24,6 +24,10 @@
 
 TEST_POINTCLOUD_URLS = [
     "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-0.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-1.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-2.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-3.json",
+    "https://scaleapi-cust-lidar.s3.us-west-1.amazonaws.com/test-scale/frame-4.json",
 ]
 
 TEST_LIDAR_SCENES = {
@@ -56,7 +60,14 @@
                             }
                         },
                     },
-                }
+                },
+                {
+                    "lidar": {
+                        "pointcloud_url": TEST_POINTCLOUD_URLS[0],
+                        "reference_id": "lidar_frame_2",
+                        "metadata": {},
+                    },
+                },
             ],
             "metadata": {},
         },
@@ -71,6 +82,14 @@
     DatasetItem(TEST_IMG_URLS[3], "4"),
 ]
 
+TEST_LIDAR_ITEMS = [
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[0], reference_id="1"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[1], reference_id="2"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[2], reference_id="3"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[3], reference_id="4"),
+    DatasetItem(pointcloud_location=TEST_POINTCLOUD_URLS[4], reference_id="5"),
+]
+
 LOCAL_FILENAME = "tests/test_img.jpg"
 TEST_PREDS = [
     BoxPrediction("[Pytest Box Prediction 1]", 0, 0, 100, 100, "1"),
diff --git a/tests/test_scene.py b/tests/test_scene.py
