Merge pull request #104 from scaleapi/nucleus_frame_args

Small tweaks to improve Frame UX

Author: sasha-scale

nucleus/scene.py

@@ -1,7 +1,7 @@
 import json
 from abc import ABC
 from dataclasses import dataclass, field
-from typing import Optional, Union, Any, Dict, List
+from typing import Optional, Any, Dict, List
 from nucleus.constants import (
     FRAMES_KEY,
     LENGTH_KEY,
@@ -15,10 +15,11 @@
 from .dataset_item import DatasetItemType, DatasetItem
 
 
-@dataclass
 class Frame:
-    items: Dict[str, DatasetItem] = field(default_factory=dict)
-    index: Union[int, None] = None
+    def __init__(self, **kwargs):
+        self.items = {}
+        for key, value in kwargs.items():
+            self.items[key] = value
 
     def __post_init__(self):
         for key, value in self.items.items():
@@ -28,7 +29,7 @@ def __post_init__(self):
             ), "All values must be DatasetItems"
 
     def __repr__(self) -> str:
-        return f"Frame(index={self.index}, items={self.items})"
+        return f"Frame(items={self.items})"
 
     def add_item(self, item: DatasetItem, sensor_name: str):
         self.items[sensor_name] = item
@@ -46,16 +47,13 @@ def get_items(self):
     def get_sensors(self):
         return list(self.items.keys())
 
-    def get_index(self):
-        return self.index
-
     @classmethod
     def from_json(cls, payload: dict):
         items = {
             sensor: DatasetItem.from_json(item, is_scene=True)
             for sensor, item in payload.items()
         }
-        return cls(items=items)
+        return cls(**items)
 
     def to_payload(self) -> dict:
         return {
@@ -71,18 +69,10 @@ class Scene(ABC):
     metadata: Optional[dict] = None
 
     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:
-            indexed_frames = [
-                Frame(index=i, items=frame.items)
-                for i, frame in enumerate(self.frames)
-            ]
-            self.frames_dict = dict(enumerate(indexed_frames))
+        self.frames_dict = dict(enumerate(self.frames))
 
     @property
     def length(self) -> int:
@@ -92,17 +82,6 @@ def length(self) -> int:
     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)
-        )
-        explicit_frame_order = all(
-            (frame.index is not None for frame in self.frames)
-        )
-        assert (
-            infer_from_list_position or explicit_frame_order
-        ), "Must specify index explicitly for all frames or infer from list position for all frames"
-
     def validate(self):
         assert self.length > 0, "Must have at least 1 frame in a scene"
         for frame in self.frames_dict.values():
@@ -113,21 +92,18 @@ def validate(self):
     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})
+            new_frame = Frame(**{sensor_name: item})
             self.frames_dict[index] = new_frame
         else:
             self.frames_dict[index].items[sensor_name] = item
 
-    def add_frame(self, frame: Frame, update: bool = False):
-        assert (
-            frame.index is not None
-        ), "Must specify index explicitly when calling add_frame"
+    def add_frame(self, frame: Frame, index: int, update: bool = False):
         if (
-            frame.index not in self.frames_dict
-            or frame.index in self.frames_dict
+            index not in self.frames_dict
+            or index in self.frames_dict
             and update
         ):
-            self.frames_dict[frame.index] = frame
+            self.frames_dict[index] = frame
             self.sensors.update(frame.get_sensors())
 
     def get_frame(self, index: int):

tests/test_scene.py

@@ -47,11 +47,10 @@ def dataset(CLIENT):
 
 
 def test_frame_add_item():
-    frame = Frame(index=0)
+    frame = Frame()
     frame.add_item(TEST_DATASET_ITEMS[0], "camera")
     frame.add_item(TEST_LIDAR_ITEMS[0], "lidar")
 
-    assert frame.get_index() == 0
     assert frame.get_sensors() == ["camera", "lidar"]
     for item in frame.get_items():
         assert item in [TEST_DATASET_ITEMS[0], TEST_LIDAR_ITEMS[0]]
@@ -104,14 +103,19 @@ def test_scene_from_json():
         "lidar": lidar_item_f2,
     }
 
-    expected_frames = [Frame(expected_items_1), Frame(expected_items_2)]
+    expected_frames = [Frame(**expected_items_1), Frame(**expected_items_2)]
     expected_scene = LidarScene(
         scene_json[REFERENCE_ID_KEY], expected_frames, metadata={}
     )
+
     assert sorted(
         scene.get_items(), key=lambda item: item.reference_id
     ) == sorted(expected_scene.get_items(), key=lambda item: item.reference_id)
-    assert scene.get_frames() == expected_scene.get_frames()
+    scene_frames = [frame.to_payload() for frame in scene.get_frames()]
+    expected_scene_frames = [
+        frame.to_payload() for frame in expected_scene.get_frames()
+    ]
+    assert scene_frames == expected_scene_frames
     assert set(scene.get_sensors()) == set(expected_scene.get_sensors())
     assert scene.to_payload() == expected_scene.to_payload()
 
@@ -195,29 +199,26 @@ def test_scene_add_frame():
     frames = [frame_1]
     scene = LidarScene(scene_ref_id, frames=frames)
 
-    frame_2 = Frame(index=1)
+    frame_2 = Frame()
     frame_2.add_item(TEST_LIDAR_ITEMS[1], "lidar")
-    scene.add_frame(frame_2)
+    scene.add_frame(frame_2, index=1)
     frames.append(frame_2)
 
     assert scene.length == len(frames)
     assert set(scene.get_sensors()) == set(["camera", "lidar"])
     expected_frame_1 = Frame(
-        index=0,
-        items={
+        **{
             "camera": TEST_DATASET_ITEMS[0],
             "lidar": TEST_LIDAR_ITEMS[0],
         },
     )
-    assert scene.get_frame(0) == expected_frame_1
+    assert scene.get_frame(0).to_payload() == expected_frame_1.to_payload()
     expected_frame_2 = Frame(
-        index=1,
-        items={
+        **{
             "lidar": TEST_LIDAR_ITEMS[1],
         },
     )
-    expected_frames = [expected_frame_1, expected_frame_2]
-    assert scene.get_frames() == expected_frames
+    assert scene.get_frame(1).to_payload() == expected_frame_2.to_payload()
     for item in scene.get_items_from_sensor("lidar"):
         assert item in [TEST_LIDAR_ITEMS[0], TEST_LIDAR_ITEMS[1]]