Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add support of rendering boxes with its elements #57

Merged
merged 3 commits into from
Nov 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions docs/tutorials/render.md
Original file line number Diff line number Diff line change
Expand Up @@ -67,3 +67,12 @@ For details, please refer to the API references.
# Rendering 2D boxes
>>> viewer.render_box2ds(seconds, box2ds)
```

It allows us to render boxes by specifying elements of boxes directly.

```python
# Rendering 3D boxes
>>> viewer.render_box3ds(seconds, centers, rotations, sizes, class_ids)
# Rendering 2D boxes
>>> viewer.render_box2ds(seconds, rois, class_ids)
```
98 changes: 92 additions & 6 deletions t4_devkit/viewer/rendering_data/box.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
from __future__ import annotations

from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, overload

import numpy as np
import rerun as rr

if TYPE_CHECKING:
from t4_devkit.dataclass import Box2D, Box3D
from t4_devkit.typing import RoiType, SizeType, TranslationType, VelocityType
from t4_devkit.typing import RoiType, RotationType, SizeType, TranslationType, VelocityType

__all__ = ["BoxData3D", "BoxData2D"]

Expand All @@ -30,12 +30,44 @@ def __init__(self, label2id: dict[str, int] | None = None) -> None:

self._label2id: dict[str, int] = {} if label2id is None else label2id

@overload
def append(self, box: Box3D) -> None:
"""Append a 3D box data.
"""Append a 3D box data with a Box3D object.

Args:
box (Box3D): `Box3D` object.
"""
pass

@overload
def append(
self,
center: TranslationType,
rotation: RotationType,
size: SizeType,
class_id: int,
uuid: str | None = None,
velocity: VelocityType | None = None,
) -> None:
"""Append a 3D box data with its elements.

Args:
center (TranslationType): 3D position in the order of (x, y, z).
rotation (RotationType): Quaternion.
size (SizeType): Box size in the order of (width, height, length).
class_id (int): Class ID.
velocity (VelocityType | None, optional): Box velocity. Defaults to None.
uuid (str | None, optional): Unique identifier.
"""
pass

def append(self, *args, **kwargs) -> None:
if len(args) + len(kwargs) == 1:
self._append_with_box(*args, **kwargs)
else:
self._append_with_elements(*args, **kwargs)

def _append_with_box(self, box: Box3D) -> None:
self._centers.append(box.position)

rotation_xyzw = np.roll(box.rotation.q, shift=-1)
Expand All @@ -49,11 +81,36 @@ def append(self, box: Box3D) -> None:

self._class_ids.append(self._label2id[box.semantic_label.name])

if box.velocity is not None:
self._velocities.append(box.velocity)

if box.uuid is not None:
self._uuids.append(box.uuid[:6])

if box.velocity is not None:
self._velocities.append(box.velocity)
def _append_with_elements(
self,
center: TranslationType,
rotation: RotationType,
size: SizeType,
class_id: int,
velocity: VelocityType | None = None,
uuid: str | None = None,
) -> None:
self._centers.append(center)

rotation_xyzw = np.roll(rotation.q, shift=-1)
self._rotations.append(rr.Quaternion(xyzw=rotation_xyzw))

width, length, height = size
self._sizes.append((length, width, height))

self._class_ids.append(class_id)

if velocity is not None:
self._velocities.append(velocity)

if uuid is not None:
self._uuids.append(uuid)

def as_boxes3d(self) -> rr.Boxes3D:
"""Return 3D boxes data as a `rr.Boxes3D`.
Expand Down Expand Up @@ -98,12 +155,33 @@ def __init__(self, label2id: dict[str, int] | None = None) -> None:

self._label2id: dict[str, int] = {} if label2id is None else label2id

@overload
def append(self, box: Box2D) -> None:
"""Append a 2D box data.
"""Append a 2D box data with a `Box2D` object.

Args:
box (Box2D): `Box2D` object.
"""
pass

@overload
def append(self, roi: RoiType, class_id: int, uuid: str | None = None) -> None:
"""Append a 2D box data with its elements.

Args:
roi (RoiType): ROI in the order of (xmin, ymin, xmax, ymax).
class_id (int): Class ID.
uuid (str | None, optional): Unique identifier.
"""
pass

def append(self, *args, **kwargs) -> None:
if len(args) + len(kwargs) == 1:
self._append_with_box(*args, **kwargs)
else:
self._append_with_elements(*args, **kwargs)

def _append_with_box(self, box: Box2D) -> None:
self._rois.append(box.roi.roi)

if box.semantic_label.name not in self._label2id:
Expand All @@ -114,6 +192,14 @@ def append(self, box: Box2D) -> None:
if box.uuid is not None:
self._uuids.append(box.uuid)

def _append_with_elements(self, roi: RoiType, class_id: int, uuid: str | None = None) -> None:
self._rois.append(roi)

self._class_ids.append(class_id)

if uuid is not None:
self._uuids.append(uuid)

def as_boxes2d(self) -> rr.Boxes2D:
"""Return 2D boxes data as a `rr.Boxes2D`.

Expand Down
Loading
Loading