feat: add value arguments to SceneBuilder annotation addition functions

DSD-DBS · Dec 9, 2024 · c0ac631 · c0ac631
1 parent 5252610
commit c0ac631
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 21 deletions.
diff --git 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
@@ -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
@@ -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},
                     )
                 }