Skip to content

Commit 5d3c53a

Browse files
Drew KaulDrew Kaul
authored andcommitted
add tests for Scene methods
1 parent a42ccc7 commit 5d3c53a

File tree

1 file changed

+151
-12
lines changed

1 file changed

+151
-12
lines changed

tests/test_scene.py

Lines changed: 151 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,17 @@
11
import pytest
22
from nucleus.constants import (
33
ANNOTATIONS_KEY,
4-
DATASET_ITEM_ID_KEY,
54
FRAMES_KEY,
65
IMAGE_KEY,
7-
IMAGE_URL_KEY,
8-
ITEM_KEY,
6+
LENGTH_KEY,
7+
METADATA_KEY,
8+
NUM_SENSORS_KEY,
99
POINTCLOUD_KEY,
10-
POINTCLOUD_URL_KEY,
1110
REFERENCE_ID_KEY,
1211
SCENES_KEY,
12+
TYPE_KEY,
1313
UPDATE_KEY,
14+
URL_KEY,
1415
)
1516

1617
from nucleus import (
@@ -19,6 +20,10 @@
1920
Frame,
2021
)
2122

23+
from nucleus.scene import (
24+
flatten,
25+
)
26+
2227
from .helpers import (
2328
TEST_DATASET_3D_NAME,
2429
TEST_CUBOID_ANNOTATIONS,
@@ -50,20 +55,154 @@ def test_frame_add_item(dataset):
5055
assert frame.get_item("lidar") == TEST_LIDAR_ITEMS[0]
5156
assert frame.to_payload() == {
5257
"camera": {
53-
"url": TEST_DATASET_ITEMS[0].image_location,
54-
"reference_id": TEST_DATASET_ITEMS[0].reference_id,
55-
"type": IMAGE_KEY,
56-
"metadata": TEST_DATASET_ITEMS[0].metadata or {},
58+
URL_KEY: TEST_DATASET_ITEMS[0].image_location,
59+
REFERENCE_ID_KEY: TEST_DATASET_ITEMS[0].reference_id,
60+
TYPE_KEY: IMAGE_KEY,
61+
METADATA_KEY: TEST_DATASET_ITEMS[0].metadata or {},
5762
},
5863
"lidar": {
59-
"url": TEST_LIDAR_ITEMS[0].pointcloud_location,
60-
"reference_id": TEST_LIDAR_ITEMS[0].reference_id,
61-
"type": POINTCLOUD_KEY,
62-
"metadata": TEST_LIDAR_ITEMS[0].metadata or {},
64+
URL_KEY: TEST_LIDAR_ITEMS[0].pointcloud_location,
65+
REFERENCE_ID_KEY: TEST_LIDAR_ITEMS[0].reference_id,
66+
TYPE_KEY: POINTCLOUD_KEY,
67+
METADATA_KEY: TEST_LIDAR_ITEMS[0].metadata or {},
6368
},
6469
}
6570

6671

72+
def test_scene_property_methods(dataset):
73+
payload = TEST_LIDAR_SCENES
74+
scene_json = payload[SCENES_KEY][0]
75+
scene = LidarScene.from_json(scene_json)
76+
77+
expected_length = len(scene_json[FRAMES_KEY])
78+
assert scene.length == expected_length
79+
sensors = flatten([list(frame.keys()) for frame in scene_json[FRAMES_KEY]])
80+
expected_num_sensors = len(set(sensors))
81+
assert scene.num_sensors == expected_num_sensors
82+
assert scene.info() == {
83+
REFERENCE_ID_KEY: scene_json[REFERENCE_ID_KEY],
84+
LENGTH_KEY: expected_length,
85+
NUM_SENSORS_KEY: expected_num_sensors,
86+
}
87+
88+
89+
def test_scene_add_item(dataset):
90+
scene_ref_id = "scene_1"
91+
scene = LidarScene(scene_ref_id)
92+
scene.add_item(0, "camera", TEST_DATASET_ITEMS[0])
93+
scene.add_item(0, "lidar", TEST_LIDAR_ITEMS[0])
94+
scene.add_item(1, "lidar", TEST_LIDAR_ITEMS[1])
95+
96+
assert set(scene.get_sensors()) == set(["camera", "lidar"])
97+
assert scene.get_item(1, "lidar") == TEST_LIDAR_ITEMS[1]
98+
assert scene.get_items_from_sensor("lidar") == [
99+
TEST_LIDAR_ITEMS[0],
100+
TEST_LIDAR_ITEMS[1],
101+
]
102+
assert scene.get_items_from_sensor("camera") == [
103+
TEST_DATASET_ITEMS[0],
104+
None,
105+
]
106+
for item in scene.get_items():
107+
assert item in [
108+
TEST_DATASET_ITEMS[0],
109+
TEST_LIDAR_ITEMS[0],
110+
TEST_LIDAR_ITEMS[1],
111+
]
112+
113+
assert scene.to_payload() == {
114+
REFERENCE_ID_KEY: scene_ref_id,
115+
FRAMES_KEY: [
116+
{
117+
"camera": {
118+
URL_KEY: TEST_DATASET_ITEMS[0].image_location,
119+
REFERENCE_ID_KEY: TEST_DATASET_ITEMS[0].reference_id,
120+
TYPE_KEY: IMAGE_KEY,
121+
METADATA_KEY: TEST_DATASET_ITEMS[0].metadata or {},
122+
},
123+
"lidar": {
124+
URL_KEY: TEST_LIDAR_ITEMS[0].pointcloud_location,
125+
REFERENCE_ID_KEY: TEST_LIDAR_ITEMS[0].reference_id,
126+
TYPE_KEY: POINTCLOUD_KEY,
127+
METADATA_KEY: TEST_LIDAR_ITEMS[0].metadata or {},
128+
},
129+
},
130+
{
131+
"lidar": {
132+
URL_KEY: TEST_LIDAR_ITEMS[1].pointcloud_location,
133+
REFERENCE_ID_KEY: TEST_LIDAR_ITEMS[1].reference_id,
134+
TYPE_KEY: POINTCLOUD_KEY,
135+
METADATA_KEY: TEST_LIDAR_ITEMS[1].metadata or {},
136+
}
137+
},
138+
],
139+
}
140+
141+
142+
def test_scene_add_frame(dataset):
143+
frame_1 = Frame()
144+
frame_1.add_item(TEST_DATASET_ITEMS[0], "camera")
145+
frame_1.add_item(TEST_LIDAR_ITEMS[0], "lidar")
146+
147+
scene_ref_id = "scene_1"
148+
frames = [frame_1]
149+
scene = LidarScene(scene_ref_id, frames=frames)
150+
151+
frame_2 = Frame(index=1)
152+
frame_2.add_item(TEST_LIDAR_ITEMS[1], "lidar")
153+
scene.add_frame(frame_2)
154+
frames.append(frame_2)
155+
156+
assert scene.length == len(frames)
157+
assert set(scene.get_sensors()) == set(["camera", "lidar"])
158+
expected_frame_1 = Frame(
159+
index=0,
160+
items={
161+
"camera": TEST_DATASET_ITEMS[0],
162+
"lidar": TEST_LIDAR_ITEMS[0],
163+
},
164+
)
165+
assert scene.get_frame(0) == expected_frame_1
166+
expected_frame_2 = Frame(
167+
index=1,
168+
items={
169+
"lidar": TEST_LIDAR_ITEMS[1],
170+
},
171+
)
172+
expected_frames = [expected_frame_1, expected_frame_2]
173+
assert scene.get_frames() == expected_frames
174+
for item in scene.get_items_from_sensor("lidar"):
175+
assert item in [TEST_LIDAR_ITEMS[0], TEST_LIDAR_ITEMS[1]]
176+
177+
assert scene.to_payload() == {
178+
REFERENCE_ID_KEY: scene_ref_id,
179+
FRAMES_KEY: [
180+
{
181+
"camera": {
182+
URL_KEY: TEST_DATASET_ITEMS[0].image_location,
183+
REFERENCE_ID_KEY: TEST_DATASET_ITEMS[0].reference_id,
184+
TYPE_KEY: IMAGE_KEY,
185+
METADATA_KEY: TEST_DATASET_ITEMS[0].metadata or {},
186+
},
187+
"lidar": {
188+
URL_KEY: TEST_LIDAR_ITEMS[0].pointcloud_location,
189+
REFERENCE_ID_KEY: TEST_LIDAR_ITEMS[0].reference_id,
190+
TYPE_KEY: POINTCLOUD_KEY,
191+
METADATA_KEY: TEST_LIDAR_ITEMS[0].metadata or {},
192+
},
193+
},
194+
{
195+
"lidar": {
196+
URL_KEY: TEST_LIDAR_ITEMS[1].pointcloud_location,
197+
REFERENCE_ID_KEY: TEST_LIDAR_ITEMS[1].reference_id,
198+
TYPE_KEY: POINTCLOUD_KEY,
199+
METADATA_KEY: TEST_LIDAR_ITEMS[1].metadata or {},
200+
}
201+
},
202+
],
203+
}
204+
205+
67206
def test_scene_upload_sync(dataset):
68207
payload = TEST_LIDAR_SCENES
69208
scenes = [

0 commit comments

Comments
 (0)