diff --git a/raillabel/scene_builder/scene_builder.py b/raillabel/scene_builder/scene_builder.py index 9711746..d398a0c 100644 --- a/raillabel/scene_builder/scene_builder.py +++ b/raillabel/scene_builder/scene_builder.py @@ -7,7 +7,18 @@ from dataclasses import dataclass from uuid import UUID -from raillabel.format import Metadata, Object, Scene +from raillabel.format import ( + Camera, + GpsImu, + IntrinsicsPinhole, + IntrinsicsRadar, + Lidar, + Metadata, + Object, + OtherSensor, + Radar, + Scene, +) @dataclass @@ -27,7 +38,7 @@ def add_object( object_type: str | None = None, object_name: str | None = None, ) -> SceneBuilder: - """Add an object to a scene.""" + """Add an object to the scene.""" scene = deepcopy(self.result) object_type, object_name = _resolve_empty_object_name_or_type(object_type, object_name) @@ -36,6 +47,43 @@ def add_object( scene.objects[object_id] = Object(object_name, object_type) return SceneBuilder(scene) + def add_sensor(self, sensor_id: str) -> SceneBuilder: + """Add a sensor to the scene. + + The sensor type is implicitly determined by the sensor_id. Id's, that start with 'rgb_' or + 'ir_' are added as a camera. If the id is 'lidar', it is a lidar. 'radar' creates a Radar. + 'gps_imu' creates a GpsImu. If the id does not match any of these, a OtherSensor is added. + """ + scene = deepcopy(self.result) + + truncated_sensor_id = sensor_id.split("_")[0].lower() + + if truncated_sensor_id in ["rgb", "ir"]: + scene.sensors[sensor_id] = Camera( + intrinsics=IntrinsicsPinhole( + camera_matrix=(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0), + distortion=(0, 0, 0, 0, 0), + width_px=0, + height_px=0, + ) + ) + + elif truncated_sensor_id == "radar": + scene.sensors[sensor_id] = Radar( + intrinsics=IntrinsicsRadar(resolution_px_per_m=0, width_px=0, height_px=0) + ) + + elif truncated_sensor_id == "lidar": + scene.sensors[sensor_id] = Lidar() + + elif truncated_sensor_id == "gps": + scene.sensors[sensor_id] = GpsImu() + + else: + scene.sensors[sensor_id] = OtherSensor() + + return SceneBuilder(scene) + def _resolve_empty_object_name_or_type( object_type: str | None, object_name: str | None @@ -56,7 +104,7 @@ def _resolve_empty_object_name_or_type( if object_name is not None and object_type is not None: return object_type, object_name - raise RuntimeError + raise RuntimeError # this part is unreachable but this is the only way that mypy is happy def _resolve_empty_object_uid(scene: Scene, object_id: str | UUID | None) -> UUID: @@ -64,6 +112,7 @@ def _resolve_empty_object_uid(scene: Scene, object_id: str | UUID | None) -> UUI uid_index = 0 while _generate_deterministic_uuid(uid_index, "5c59aad4") in scene.objects: uid_index += 1 + object_id = _generate_deterministic_uuid(uid_index, "5c59aad4") return UUID(str(object_id)) diff --git a/tests/test_raillabel/scene_builder/test_scene_builder.py b/tests/test_raillabel/scene_builder/test_scene_builder.py index 3923ea1..8a8162d 100644 --- a/tests/test_raillabel/scene_builder/test_scene_builder.py +++ b/tests/test_raillabel/scene_builder/test_scene_builder.py @@ -7,8 +7,19 @@ import pytest import raillabel -from raillabel.format import Scene, Metadata, Object from raillabel.scene_builder.scene_builder import SceneBuilder +from raillabel.format import ( + Scene, + Metadata, + Object, + Camera, + Lidar, + Radar, + GpsImu, + OtherSensor, + IntrinsicsPinhole, + IntrinsicsRadar, +) def test_empty(): @@ -91,5 +102,76 @@ def test_add_object__object_id_iteration(): ) +def test_add_sensor__camera_rgb(): + actual = SceneBuilder.empty().add_sensor("rgb_middle") + assert actual.result == Scene( + metadata=Metadata(schema_version="1.0.0"), + sensors={ + "rgb_middle": Camera( + intrinsics=IntrinsicsPinhole( + camera_matrix=tuple([0] * 12), + distortion=tuple([0] * 5), + width_px=0, + height_px=0, + ) + ) + }, + ) + + +def test_add_sensor__camera_ir(): + actual = SceneBuilder.empty().add_sensor("ir_left") + assert actual.result == Scene( + metadata=Metadata(schema_version="1.0.0"), + sensors={ + "ir_left": Camera( + intrinsics=IntrinsicsPinhole( + camera_matrix=tuple([0] * 12), + distortion=tuple([0] * 5), + width_px=0, + height_px=0, + ) + ) + }, + ) + + +def test_add_sensor__radar(): + actual = SceneBuilder.empty().add_sensor("radar") + assert actual.result == Scene( + metadata=Metadata(schema_version="1.0.0"), + sensors={ + "radar": Radar( + intrinsics=IntrinsicsRadar( + resolution_px_per_m=0, + width_px=0, + height_px=0, + ) + ) + }, + ) + + +def test_add_sensor__lidar(): + actual = SceneBuilder.empty().add_sensor("lidar") + assert actual.result == Scene( + metadata=Metadata(schema_version="1.0.0"), sensors={"lidar": Lidar()} + ) + + +def test_add_sensor__gps_imu(): + actual = SceneBuilder.empty().add_sensor("gps_imu") + assert actual.result == Scene( + metadata=Metadata(schema_version="1.0.0"), sensors={"gps_imu": GpsImu()} + ) + + +def test_add_sensor__other(): + actual = SceneBuilder.empty().add_sensor("SOMETHING_ELSE") + assert actual.result == Scene( + metadata=Metadata(schema_version="1.0.0"), sensors={"SOMETHING_ELSE": OtherSensor()} + ) + + if __name__ == "__main__": - pytest.main([__file__, "-v"]) + pytest.main([__file__, "-vv"])