Skip to content

Commit

Permalink
feat: implement SceneBuilder.add_sensor()
Browse files Browse the repository at this point in the history
  • Loading branch information
unexcellent committed Nov 15, 2024
1 parent 45eef29 commit ab02190
Show file tree
Hide file tree
Showing 2 changed files with 136 additions and 5 deletions.
55 changes: 52 additions & 3 deletions raillabel/scene_builder/scene_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand All @@ -56,14 +104,15 @@ 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:
if object_id is None:
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))
Expand Down
86 changes: 84 additions & 2 deletions tests/test_raillabel/scene_builder/test_scene_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -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"])

0 comments on commit ab02190

Please sign in to comment.