Skip to content

Commit c11e182

Browse files
ktro2828ktro2828
authored
refactor: modify module and class names (#126)
* refactor: gather append methods and reduce duplicated lines Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp> * refactor: add inner class to represent records Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp> * perf: accelerate segmentation image generation Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp> * refactor: rename module to record Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp> * refactor: rename classes to BatchXXX Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp> * test: add unit test for render segmentation 2d Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp> --------- Signed-off-by: ktro2828 <kotaro.uetake@tier4.jp>
1 parent 0b034d9 commit c11e182

5 files changed

Lines changed: 93 additions & 80 deletions

File tree

t4_devkit/viewer/rendering_data/__init__.py renamed to t4_devkit/viewer/record/__init__.py

File renamed without changes.

t4_devkit/viewer/rendering_data/box.py renamed to t4_devkit/viewer/record/box.py

Lines changed: 35 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,16 @@
1111
from t4_devkit.dataclass import Box2D, Box3D, Future
1212
from t4_devkit.typing import QuaternionLike, RoiLike, Vector3Like
1313

14-
__all__ = ["BoxData3D", "BoxData2D"]
14+
__all__ = ["BatchBox3D", "BatchBox2D"]
1515

1616

1717
@define
18-
class BoxData3D:
18+
class BatchBox3D:
1919
"""A class to store 3D boxes data for rendering.
2020
2121
Attributes:
2222
label2id (dict[str, int]): Key-value of map of label name and its ID.
23-
records (list[_AssetBox3D]): List of 3D box records for rendering.
23+
records (list[Record]): List of 3D box records for rendering.
2424
"""
2525

2626
label2id: dict[str, int] = field(factory=dict)
@@ -78,23 +78,17 @@ def append(self, *args, **kwargs) -> None:
7878
self._append_with_elements(*args, **kwargs)
7979

8080
def _append_with_box(self, box: Box3D) -> None:
81-
rotation_xyzw = np.roll(box.rotation.q, shift=-1)
82-
83-
width, length, height = box.size
84-
8581
if box.semantic_label.name not in self.label2id:
8682
self.label2id[box.semantic_label.name] = len(self.label2id)
8783

88-
self.records.append(
89-
self.Record(
90-
center=box.position,
91-
rotation=rr.Quaternion(xyzw=rotation_xyzw),
92-
size=(length, width, height),
93-
class_id=self.label2id[box.semantic_label.name],
94-
uuid=box.uuid,
95-
velocity=box.velocity,
96-
future=box.future,
97-
)
84+
self._append_with_elements(
85+
center=box.position,
86+
rotation=box.rotation,
87+
size=box.size,
88+
class_id=self.label2id[box.semantic_label.name],
89+
uuid=box.uuid,
90+
velocity=box.velocity,
91+
future=box.future,
9892
)
9993

10094
def _append_with_elements(
@@ -189,20 +183,24 @@ def as_linestrips3d(self) -> rr.LineStrips3D:
189183

190184

191185
@define
192-
class BoxData2D:
186+
class BatchBox2D:
193187
"""A class to store 2D boxes data for rendering.
194188
195189
Attributes:
196190
label2id (dict[str, int]): Key-value of map of label name and its ID.
197-
rois (list[RoiLike]): List of ROIs in the order of (xmin, ymin, xmax, ymax).
198-
class_ids (list[int]): List of label class IDs.
199-
uuids (list[str]): List of unique identifier IDs.
191+
records (list[Record]): List of 2D box records for rendering.
200192
"""
201193

202194
label2id: dict[str, int] = field(factory=dict)
203-
rois: list[RoiLike] = field(init=False, factory=list)
204-
class_ids: list[int] = field(init=False, factory=list)
205-
uuids: list[str] = field(init=False, factory=list)
195+
records: list[Record] = field(init=False, factory=list)
196+
197+
@define
198+
class Record:
199+
"""Inner class to represent a record of 2D box instance for rendering."""
200+
201+
roi: RoiLike
202+
class_id: int
203+
uuid: str | None = field(default=None)
206204

207205
@overload
208206
def append(self, box: Box2D) -> None:
@@ -231,35 +229,33 @@ def append(self, *args, **kwargs) -> None:
231229
self._append_with_elements(*args, **kwargs)
232230

233231
def _append_with_box(self, box: Box2D) -> None:
234-
self.rois.append(box.roi.roi)
235-
236232
if box.semantic_label.name not in self.label2id:
237233
self.label2id[box.semantic_label.name] = len(self.label2id)
238234

239-
self.class_ids.append(self.label2id[box.semantic_label.name])
240-
241-
if box.uuid is not None:
242-
self.uuids.append(box.uuid)
235+
if box.roi is not None:
236+
self._append_with_elements(
237+
roi=box.roi.roi,
238+
class_id=self.label2id[box.semantic_label.name],
239+
uuid=box.uuid,
240+
)
243241

244242
def _append_with_elements(self, roi: RoiLike, class_id: int, uuid: str | None = None) -> None:
245-
self.rois.append(roi)
246-
247-
self.class_ids.append(class_id)
248-
249-
if uuid is not None:
250-
self.uuids.append(uuid)
243+
self.records.append(self.Record(roi=roi, class_id=class_id, uuid=uuid))
251244

252245
def as_boxes2d(self) -> rr.Boxes2D:
253246
"""Return 2D boxes data as a `rr.Boxes2D`.
254247
255248
Returns:
256249
`rr.Boxes2D` object.
257250
"""
258-
labels = None if len(self.uuids) == 0 else self.uuids
251+
rois, class_ids, labels = map(
252+
list, zip(*((r.roi, r.class_id, r.uuid) for r in self.records))
253+
)
254+
259255
return rr.Boxes2D(
260-
array=self.rois,
256+
array=rois,
261257
array_format=rr.Box2DFormat.XYXY,
262258
labels=labels,
263-
class_ids=self.class_ids,
259+
class_ids=class_ids,
264260
show_labels=False,
265261
)

t4_devkit/viewer/rendering_data/segmentation.py renamed to t4_devkit/viewer/record/segmentation.py

Lines changed: 19 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -9,25 +9,29 @@
99
if TYPE_CHECKING:
1010
from t4_devkit.typing import NDArrayU8
1111

12-
__all__ = ["SegmentationData2D"]
12+
__all__ = ["BatchSegmentation2D"]
1313

1414

1515
@define
16-
class SegmentationData2D:
16+
class BatchSegmentation2D:
1717
"""A class to store 2D segmentation image data for rendering.
1818
1919
Attributes:
20-
masks (list[NDArray]): List of segmentation masks in the shape of (H, W).
21-
class_ids (list[int]): List of label class IDs.
22-
uuids (list[str]): List of unique identifier IDs.
20+
records (list[Record]): List of 2D segmentation mask records for rendering.
2321
size (tuple[int, int] | None): Size of image in the order of (height, width).
2422
"""
2523

26-
masks: list[NDArrayU8] = field(init=False, factory=list)
27-
class_ids: list[int] = field(init=False, factory=list)
28-
uuids: list[str] = field(init=False, factory=list)
24+
records: list[Record] = field(init=False, factory=list)
2925
size: tuple[int, int] | None = field(init=False, default=None)
3026

27+
@define
28+
class Record:
29+
"""Inner class to represent a record of 2D segmentation mask instance for rendering."""
30+
31+
mask: NDArrayU8
32+
class_id: int
33+
uuid: str | None = field(default=None)
34+
3135
def append(self, mask: NDArrayU8, class_id: int, uuid: str | None = None) -> None:
3236
"""Append a segmentation mask and its class ID.
3337
@@ -50,11 +54,8 @@ def append(self, mask: NDArrayU8, class_id: int, uuid: str | None = None) -> Non
5054
raise ValueError(
5155
f"All masks must be the same size. Expected: {self.size}, but got {mask.shape}."
5256
)
53-
self.masks.append(mask)
54-
self.class_ids.append(class_id)
5557

56-
if uuid is not None:
57-
self.uuids.append(uuid[:6])
58+
self.records.append(self.Record(mask=mask, class_id=class_id, uuid=uuid))
5859

5960
def as_segmentation_image(self) -> rr.SegmentationImage:
6061
"""Return mask data as a `rr.SegmentationImage`.
@@ -65,9 +66,12 @@ def as_segmentation_image(self) -> rr.SegmentationImage:
6566
TODO:
6667
Add support of instance segmentation.
6768
"""
68-
image = np.zeros(self.size, dtype=np.uint8)
69+
# (N, H, W)
70+
masks = np.stack([r.mask for r in self.records]).astype(np.uint8)
71+
72+
# (N, 1, 1)
73+
class_ids = np.asarray([r.class_id for r in self.records], dtype=np.uint8)[:, None, None]
6974

70-
for mask, class_id in zip(self.masks, self.class_ids, strict=True):
71-
image[mask == 1] = class_id
75+
image = np.max(masks * class_ids, axis=0)
7276

7377
return rr.SegmentationImage(image=image)

‎t4_devkit/viewer/viewer.py‎

Lines changed: 26 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
from .color import distance_color
1616
from .geography import calculate_geodetic_point
17-
from .rendering_data import BoxData2D, BoxData3D, SegmentationData2D
17+
from .record import BatchBox2D, BatchBox3D, BatchSegmentation2D
1818

1919
if TYPE_CHECKING:
2020
from t4_devkit.dataclass import Box2D, Box3D, Future, PointCloudLike
@@ -240,27 +240,27 @@ def _render_box3ds_with_boxes(self, seconds: float, boxes: Sequence[Box3D]) -> N
240240

241241
rr.set_time_seconds(self.timeline, seconds)
242242

243-
box_data: dict[str, BoxData3D] = {}
243+
batches: dict[str, BatchBox3D] = {}
244244
for box in boxes:
245-
if box.frame_id not in box_data:
246-
box_data[box.frame_id] = BoxData3D(label2id=self.label2id)
247-
box_data[box.frame_id].append(box)
245+
if box.frame_id not in batches:
246+
batches[box.frame_id] = BatchBox3D(label2id=self.label2id)
247+
batches[box.frame_id].append(box)
248248

249-
for frame_id, data in box_data.items():
249+
for frame_id, batch in batches.items():
250250
# record boxes 3d
251251
rr.log(
252252
format_entity(self.map_entity, frame_id, "box"),
253-
data.as_boxes3d(),
253+
batch.as_boxes3d(),
254254
)
255255
# record velocities
256256
rr.log(
257257
format_entity(self.map_entity, frame_id, "velocity"),
258-
data.as_arrows3d(),
258+
batch.as_arrows3d(),
259259
)
260260
# record futures
261261
rr.log(
262262
format_entity(self.map_entity, frame_id, "future"),
263-
data.as_linestrips3d(),
263+
batch.as_linestrips3d(),
264264
)
265265

266266
def _render_box3ds_with_elements(
@@ -289,7 +289,7 @@ def _render_box3ds_with_elements(
289289
else:
290290
show_futures = True
291291

292-
box_data = BoxData3D(label2id=self.label2id)
292+
batch = BatchBox3D(label2id=self.label2id)
293293
for center, rotation, size, class_id, velocity, uuid, future in zip(
294294
centers,
295295
rotations,
@@ -300,7 +300,7 @@ def _render_box3ds_with_elements(
300300
futures,
301301
strict=True,
302302
):
303-
box_data.append(
303+
batch.append(
304304
center=center,
305305
rotation=rotation,
306306
size=size,
@@ -312,13 +312,13 @@ def _render_box3ds_with_elements(
312312

313313
rr.set_time_seconds(self.timeline, seconds)
314314

315-
rr.log(format_entity(self.ego_entity, "box"), box_data.as_boxes3d())
315+
rr.log(format_entity(self.ego_entity, "box"), batch.as_boxes3d())
316316

317317
if show_arrows:
318-
rr.log(format_entity(self.ego_entity, "velocity"), box_data.as_arrows3d())
318+
rr.log(format_entity(self.ego_entity, "velocity"), batch.as_arrows3d())
319319

320320
if show_futures:
321-
rr.log(format_entity(self.ego_entity, "future"), box_data.as_linestrips3d())
321+
rr.log(format_entity(self.ego_entity, "future"), batch.as_linestrips3d())
322322

323323
@overload
324324
def render_box2ds(self, seconds: float, boxes: Sequence[Box2D]) -> None:
@@ -365,16 +365,16 @@ def _render_box2ds_with_boxes(self, seconds: float, boxes: Sequence[Box2D]) -> N
365365

366366
rr.set_time_seconds(self.timeline, seconds)
367367

368-
box_data: dict[str, BoxData2D] = {}
368+
batches: dict[str, BatchBox2D] = {}
369369
for box in boxes:
370-
if box.frame_id not in box_data:
371-
box_data[box.frame_id] = BoxData2D(label2id=self.label2id)
372-
box_data[box.frame_id].append(box)
370+
if box.frame_id not in batches:
371+
batches[box.frame_id] = BatchBox2D(label2id=self.label2id)
372+
batches[box.frame_id].append(box)
373373

374-
for frame_id, data in box_data.items():
374+
for frame_id, batch in batches.items():
375375
rr.log(
376376
format_entity(self.ego_entity, frame_id, "box"),
377-
data.as_boxes2d(),
377+
batch.as_boxes2d(),
378378
)
379379

380380
def _render_box2ds_with_elements(
@@ -392,12 +392,12 @@ def _render_box2ds_with_elements(
392392
if uuids is None:
393393
uuids = [None] * len(rois)
394394

395-
box_data = BoxData2D(label2id=self.label2id)
395+
batch = BatchBox2D(label2id=self.label2id)
396396
for roi, class_id, uuid in zip(rois, class_ids, uuids, strict=True):
397-
box_data.append(roi=roi, class_id=class_id, uuid=uuid)
397+
batch.append(roi=roi, class_id=class_id, uuid=uuid)
398398

399399
rr.set_time_seconds(self.timeline, seconds)
400-
rr.log(format_entity(self.ego_entity, camera, "box"), box_data.as_boxes2d())
400+
rr.log(format_entity(self.ego_entity, camera, "box"), batch.as_boxes2d())
401401

402402
def render_segmentation2d(
403403
self,
@@ -423,15 +423,15 @@ def render_segmentation2d(
423423

424424
rr.set_time_seconds(self.timeline, seconds)
425425

426-
segmentation_data = SegmentationData2D()
426+
batch = BatchSegmentation2D()
427427
if uuids is None:
428428
uuids = [None] * len(masks)
429429
for mask, class_id, uuid in zip(masks, class_ids, uuids, strict=True):
430-
segmentation_data.append(mask, class_id, uuid)
430+
batch.append(mask, class_id, uuid)
431431

432432
rr.log(
433433
format_entity(self.ego_entity, camera, "segmentation"),
434-
segmentation_data.as_segmentation_image(),
434+
batch.as_segmentation_image(),
435435
)
436436

437437
def render_pointcloud(self, seconds: float, channel: str, pointcloud: PointCloudLike) -> None:

‎tests/viewer/test_viewer.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,19 @@ def test_render_box2ds(dummy_viewer, dummy_box2ds) -> None:
5454
dummy_viewer.render_box2ds(seconds, camera=camera, rois=rois, class_ids=class_ids, uuids=uuids)
5555

5656

57+
def test_render_segmentation2d(dummy_viewer) -> None:
58+
"""Test rendering 2D segmentation mask with `RerunViewer`."""
59+
seconds = 1.0 # [sec]
60+
61+
camera = "camera"
62+
masks = [np.zeros((100, 200), dtype=np.uint8) for _ in range(2)]
63+
masks[0][0:10, 0:10] = 1
64+
masks[1][20:30, 20:30] = 1
65+
class_ids = [1, 2]
66+
67+
dummy_viewer.render_segmentation2d(seconds, camera=camera, masks=masks, class_ids=class_ids)
68+
69+
5770
def test_render_pointcloud(dummy_viewer) -> None:
5871
"""Test rendering pointcloud with `RerunViewer`."""
5972
seconds = 1.0 # [sec]

0 commit comments

Comments
 (0)