diff --git a/pyproject.toml b/pyproject.toml index a133dbf..5e4709b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,6 +71,8 @@ ignore = [ "N802", # does not allow constant abstractproperties + "PLR0913",# constructors sometimes need many arguments + "TCH001", # adds hard to understand compexity without providing a benefit for smaller projects "TCH002", # same as TCH001 "TCH003", # same as TCH001 diff --git a/raillabel/scene_builder/scene_builder.py b/raillabel/scene_builder/scene_builder.py index feaa56d..11c2637 100644 --- a/raillabel/scene_builder/scene_builder.py +++ b/raillabel/scene_builder/scene_builder.py @@ -153,6 +153,8 @@ def add_annotation( def add_bbox( self, uid: str | UUID | None = None, + pos: Point2d | None = None, + size: Size2d | None = None, frame_id: int = 1, object_name: str = "person_0001", sensor_id: str = "rgb_middle", @@ -162,8 +164,8 @@ def add_bbox( bbox = Bbox( object_id=UUID("ffffffff-ffff-4fff-ffff-ffffffffffff"), sensor_id=sensor_id, - pos=Point2d(0, 0), - size=Size2d(0, 0), + pos=pos if pos is not None else Point2d(0, 0), + size=size if size is not None else Size2d(0, 0), attributes=attributes if attributes is not None else {}, ) return self.add_annotation( @@ -177,6 +179,9 @@ def add_bbox( def add_cuboid( self, uid: str | UUID | None = None, + pos: Point3d | None = None, + quat: Quaternion | None = None, + size: Size3d | None = None, frame_id: int = 1, object_name: str = "person_0001", sensor_id: str = "lidar", @@ -186,9 +191,9 @@ def add_cuboid( cuboid = Cuboid( object_id=UUID("ffffffff-ffff-4fff-ffff-ffffffffffff"), sensor_id=sensor_id, - pos=Point3d(0, 0, 0), - size=Size3d(0, 0, 0), - quat=Quaternion(0, 0, 0, 0), + pos=pos if pos is not None else Point3d(0, 0, 0), + size=size if size is not None else Size3d(0, 0, 0), + quat=quat if quat is not None else Quaternion(0, 0, 0, 0), attributes=attributes if attributes is not None else {}, ) return self.add_annotation( @@ -202,6 +207,7 @@ def add_cuboid( def add_poly2d( self, uid: str | UUID | None = None, + points: list[Point2d] | None = None, frame_id: int = 1, object_name: str = "person_0001", sensor_id: str = "rgb_middle", @@ -211,7 +217,7 @@ def add_poly2d( poly2d = Poly2d( object_id=UUID("ffffffff-ffff-4fff-ffff-ffffffffffff"), sensor_id=sensor_id, - points=[], + points=points if points is not None else [], closed=False, attributes=attributes if attributes is not None else {}, ) @@ -226,6 +232,7 @@ def add_poly2d( def add_poly3d( self, uid: str | UUID | None = None, + points: list[Point3d] | None = None, frame_id: int = 1, object_name: str = "person_0001", sensor_id: str = "lidar", @@ -235,7 +242,7 @@ def add_poly3d( poly3d = Poly3d( object_id=UUID("ffffffff-ffff-4fff-ffff-ffffffffffff"), sensor_id=sensor_id, - points=[], + points=points if points is not None else [], closed=False, attributes=attributes if attributes is not None else {}, ) @@ -250,6 +257,7 @@ def add_poly3d( def add_seg3d( self, uid: str | UUID | None = None, + point_ids: list[int] | None = None, frame_id: int = 1, object_name: str = "person_0001", sensor_id: str = "lidar", @@ -259,7 +267,7 @@ def add_seg3d( seg3d = Seg3d( object_id=UUID("ffffffff-ffff-4fff-ffff-ffffffffffff"), sensor_id=sensor_id, - point_ids=[], + point_ids=point_ids if point_ids is not None else [], attributes=attributes if attributes is not None else {}, ) return self.add_annotation( diff --git a/tests/scene_builder/test_scene_builder.py b/tests/scene_builder/test_scene_builder.py index 223077e..7684db0 100644 --- a/tests/scene_builder/test_scene_builder.py +++ b/tests/scene_builder/test_scene_builder.py @@ -7,19 +7,14 @@ import pytest -import raillabel from raillabel.scene_builder.scene_builder import SceneBuilder from raillabel.format import ( Scene, Metadata, Object, - Camera, Lidar, - Radar, GpsImu, OtherSensor, - IntrinsicsPinhole, - IntrinsicsRadar, Frame, Bbox, Point2d, @@ -233,6 +228,8 @@ def test_add_bbox(camera_empty): SceneBuilder.empty() .add_bbox( uid=UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"), + pos=Point2d(1, 2), + size=Size2d(3, 4), frame_id=2, object_name="person_0001", sensor_id="ir_middle", @@ -254,8 +251,8 @@ def test_add_bbox(camera_empty): UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"): Bbox( sensor_id="ir_middle", object_id=UUID("5c59aad4-0000-4000-0000-000000000000"), - pos=Point2d(0, 0), - size=Size2d(0, 0), + pos=Point2d(1, 2), + size=Size2d(3, 4), attributes={"attr": True}, ) } @@ -295,6 +292,9 @@ def test_add_cuboid(): SceneBuilder.empty() .add_cuboid( uid=UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"), + pos=Point3d(1, 2, 3), + quat=Quaternion(4, 5, 6, 7), + size=Size3d(8, 9, 10), frame_id=2, object_name="person_0001", sensor_id="lidar_left", @@ -316,9 +316,9 @@ def test_add_cuboid(): UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"): Cuboid( sensor_id="lidar_left", object_id=UUID("5c59aad4-0000-4000-0000-000000000000"), - pos=Point3d(0, 0, 0), - size=Size3d(0, 0, 0), - quat=Quaternion(0, 0, 0, 0), + pos=Point3d(1, 2, 3), + size=Size3d(8, 9, 10), + quat=Quaternion(4, 5, 6, 7), attributes={"my_attr": 5}, ) } @@ -359,6 +359,7 @@ def test_add_poly2d(camera_empty): SceneBuilder.empty() .add_poly2d( uid=UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"), + points=[Point2d(0.0, 1.0), Point2d(2.0, 3.0)], frame_id=2, object_name="person_0001", sensor_id="ir_left", @@ -380,7 +381,7 @@ def test_add_poly2d(camera_empty): UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"): Poly2d( sensor_id="ir_left", object_id=UUID("5c59aad4-0000-4000-0000-000000000000"), - points=[], + points=[Point2d(0.0, 1.0), Point2d(2.0, 3.0)], closed=False, attributes={"my_attr": 5}, ) @@ -421,6 +422,7 @@ def test_add_poly3d(): SceneBuilder.empty() .add_poly3d( uid=UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"), + points=[Point3d(1, 2, 3), Point3d(4, 5, 6)], frame_id=2, object_name="person_0001", sensor_id="lidar_right", @@ -442,7 +444,7 @@ def test_add_poly3d(): UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"): Poly3d( sensor_id="lidar_right", object_id=UUID("5c59aad4-0000-4000-0000-000000000000"), - points=[], + points=[Point3d(1, 2, 3), Point3d(4, 5, 6)], closed=False, attributes={"my_attr": 5}, ) @@ -483,6 +485,7 @@ def test_add_seg3d(): SceneBuilder.empty() .add_seg3d( uid=UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"), + point_ids=[1, 2, 3], frame_id=2, object_name="person_0001", sensor_id="lidar_right", @@ -504,7 +507,7 @@ def test_add_seg3d(): UUID("6c95543d-4d4f-43df-a52d-36bf868e09d8"): Seg3d( sensor_id="lidar_right", object_id=UUID("5c59aad4-0000-4000-0000-000000000000"), - point_ids=[], + point_ids=[1, 2, 3], attributes={"my_attr": 5}, ) }