|
1 | 1 | import pytest
|
2 | 2 | from nucleus.constants import (
|
3 | 3 | ANNOTATIONS_KEY,
|
4 |
| - DATASET_ITEM_ID_KEY, |
5 | 4 | FRAMES_KEY,
|
6 | 5 | IMAGE_KEY,
|
7 |
| - IMAGE_URL_KEY, |
8 |
| - ITEM_KEY, |
| 6 | + LENGTH_KEY, |
| 7 | + METADATA_KEY, |
| 8 | + NUM_SENSORS_KEY, |
9 | 9 | POINTCLOUD_KEY,
|
10 |
| - POINTCLOUD_URL_KEY, |
11 | 10 | REFERENCE_ID_KEY,
|
12 | 11 | SCENES_KEY,
|
| 12 | + TYPE_KEY, |
13 | 13 | UPDATE_KEY,
|
| 14 | + URL_KEY, |
14 | 15 | )
|
15 | 16 |
|
16 | 17 | from nucleus import (
|
|
19 | 20 | Frame,
|
20 | 21 | )
|
21 | 22 |
|
| 23 | +from nucleus.scene import ( |
| 24 | + flatten, |
| 25 | +) |
| 26 | + |
22 | 27 | from .helpers import (
|
23 | 28 | TEST_DATASET_3D_NAME,
|
24 | 29 | TEST_CUBOID_ANNOTATIONS,
|
@@ -50,20 +55,154 @@ def test_frame_add_item(dataset):
|
50 | 55 | assert frame.get_item("lidar") == TEST_LIDAR_ITEMS[0]
|
51 | 56 | assert frame.to_payload() == {
|
52 | 57 | "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 {}, |
57 | 62 | },
|
58 | 63 | "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 {}, |
63 | 68 | },
|
64 | 69 | }
|
65 | 70 |
|
66 | 71 |
|
| 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 | + |
67 | 206 | def test_scene_upload_sync(dataset):
|
68 | 207 | payload = TEST_LIDAR_SCENES
|
69 | 208 | scenes = [
|
|
0 commit comments