Skip to content

Commit

Permalink
refacor: refactor viewer class (#50)
Browse files Browse the repository at this point in the history
* refacor: refactor viewer class

Signed-off-by: ktro2828 <[email protected]>

* fix: resolve unexpected label is assigned to boxes

Signed-off-by: ktro2828 <[email protected]>

* refactor: refactor condition

Signed-off-by: ktro2828 <[email protected]>

---------

Signed-off-by: ktro2828 <[email protected]>
  • Loading branch information
ktro2828 authored Nov 21, 2024
1 parent 68ecfcc commit 517ebc2
Show file tree
Hide file tree
Showing 3 changed files with 65 additions and 62 deletions.
2 changes: 1 addition & 1 deletion t4_devkit/tier4.py
Original file line number Diff line number Diff line change
Expand Up @@ -904,7 +904,7 @@ def _init_viewer(
else None
)

viewer = Tier4Viewer(application_id, cameras=cameras, without_3d=not render_3d, spawn=spawn)
viewer = Tier4Viewer(application_id, cameras=cameras, with_3d=render_3d, spawn=spawn)

if render_annotation:
viewer = viewer.with_labels(self._label2id)
Expand Down
18 changes: 14 additions & 4 deletions t4_devkit/viewer/rendering_data/box.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,15 +15,20 @@
class BoxData3D:
"""A class to store 3D boxes data for rendering."""

def __init__(self) -> None:
def __init__(self, label2id: dict[str, int] | None = None) -> None:
"""Construct a new object.
Args:
label2id (dict[str, int] | None, optional): Key-value mapping which maps label name to its class ID.
"""
self._centers: list[TranslationType] = []
self._rotations: list[rr.Quaternion] = []
self._sizes: list[SizeType] = []
self._class_ids: list[int] = []
self._uuids: list[int] = []
self._velocities: list[VelocityType] = []

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

def append(self, box: Box3D) -> None:
"""Append a 3D box data.
Expand Down Expand Up @@ -81,12 +86,17 @@ def as_arrows3d(self) -> rr.Arrows3D:
class BoxData2D:
"""A class to store 2D boxes data for rendering."""

def __init__(self) -> None:
def __init__(self, label2id: dict[str, int] | None = None) -> None:
"""Construct a new object.
Args:
label2id (dict[str, int] | None, optional): Key-value mapping which maps label name to its class ID.
"""
self._rois: list[RoiType] = []
self._uuids: list[str] = []
self._class_ids: list[int] = []

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

def append(self, box: Box2D) -> None:
"""Append a 2D box data.
Expand Down
107 changes: 50 additions & 57 deletions t4_devkit/viewer/viewer.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,9 +71,9 @@ class Tier4Viewer:
def __init__(
self,
app_id: str,
cameras: Sequence[str] | None = None,
*,
without_3d: bool = False,
cameras: Sequence[str] | None = None,
with_3d: bool = True,
spawn: bool = True,
) -> None:
"""Construct a new object.
Expand All @@ -82,7 +82,7 @@ def __init__(
app_id (str): Application ID.
cameras (Sequence[str] | None, optional): Sequence of camera names.
If `None`, any 2D spaces will not be visualized.
without_3d (bool, optional): Whether to render objects without the 3D space.
with_3d (bool, optional): Whether to render objects with the 3D space.
spawn (bool, optional): Whether to spawn the viewer.
Examples:
Expand All @@ -92,16 +92,19 @@ def __init__(
# Rendering 3D space only
>>> viewer = Tier4Viewer("myapp")
# Rendering 2D space only
>>> viewer = Tier4Viewer("myapp", cameras=["camera0", "camera1"], without_3d=True)
>>> viewer = Tier4Viewer("myapp", cameras=["camera0", "camera1"], with_3d=False)
"""
self.app_id = app_id
self.without_3d = without_3d
self.cameras = cameras
self.with_3d = with_3d
self.with_2d = self.cameras is not None
self.label2id: dict[str, int] | None = None

assert self.cameras is not None or not without_3d
if not (self.with_3d or self.with_2d):
raise ValueError("At least one of 3D or 2D spaces must be rendered.")

view_container = []
if not without_3d:
view_container: list[rrb.Container | rrb.SpaceView] = []
if self.with_3d:
view_container.extend(
[
rrb.Horizontal(
Expand All @@ -112,7 +115,7 @@ def __init__(
]
)

if self.cameras is not None:
if self.with_2d:
camera_space_views = [
rrb.Spatial2DView(name=name, origin=format_entity(self.ego_entity, name))
for name in self.cameras
Expand All @@ -136,7 +139,7 @@ def with_labels(self, label2id: dict[str, int]) -> Self:
"""Return myself after creating `rr.AnnotationContext` on the recording.
Args:
label2id (dict[str, int]): Key-value mapping which maps label name to its ID.
label2id (dict[str, int]): Key-value mapping which maps label name to its class ID.
Returns:
Self instance.
Expand All @@ -145,13 +148,14 @@ def with_labels(self, label2id: dict[str, int]) -> Self:
>>> label2id = {"car": 0, "pedestrian": 1}
>>> viewer = Tier4Viewer("myapp").with_labels(label2id)
"""
self.label2id = label2id

rr.log(
self.map_entity,
rr.AnnotationContext(
[
rr.AnnotationInfo(id=label_id, label=label)
for label, label_id in label2id.items()
for label, label_id in self.label2id.items()
]
),
static=True,
Expand All @@ -169,18 +173,23 @@ def save(self, save_dir: str) -> None:
rr.save(filepath, default_blueprint=self.blueprint)

def render_box3ds(self, seconds: float, boxes: Sequence[Box3D]) -> None:
"""Render 3D boxes.
"""Render 3D boxes. Note that if the viewer initialized with `with_3d=False`,
no 3D box will be rendered.
Args:
seconds (float): Timestamp in [sec].
boxes (Sequence[Box3D]): Sequence of `Box3D`s.
"""
if not self.with_3d:
warnings.warn("There is no camera space.")
return

rr.set_time_seconds(self.timeline, seconds)

box_data: dict[str, BoxData3D] = {}
for box in boxes:
if box.frame_id not in box_data:
box_data[box.frame_id] = BoxData3D()
box_data[box.frame_id] = BoxData3D(label2id=self.label2id)
else:
box_data[box.frame_id].append(box)

Expand All @@ -197,13 +206,14 @@ def render_box3ds(self, seconds: float, boxes: Sequence[Box3D]) -> None:
)

def render_box2ds(self, seconds: float, boxes: Sequence[Box2D]) -> None:
"""Render 2D boxes.
"""Render 2D boxes. Note that if the viewer initialized without `cameras=None`,
no 2D box will be rendered.
Args:
seconds (float): Timestamp in [sec].
boxes (Sequence[Box2D]): Sequence of `Box2D`s.
"""
if self.cameras is None:
if not self.with_2d:
warnings.warn("There is no camera space.")
return

Expand All @@ -212,7 +222,7 @@ def render_box2ds(self, seconds: float, boxes: Sequence[Box2D]) -> None:
box_data: dict[str, BoxData2D] = {}
for box in boxes:
if box.frame_id not in box_data:
box_data[box.frame_id] = BoxData2D()
box_data[box.frame_id] = BoxData2D(label2id=self.label2id)
else:
box_data[box.frame_id].append(box)

Expand Down Expand Up @@ -240,7 +250,7 @@ def render_segmentation2d(
class_ids (Sequence[int]): Sequence of label ids.
uuids (Sequence[str | None] | None, optional): Sequence of each instance ID.
"""
if self.cameras is None or camera not in self.cameras:
if not self.with_2d or camera not in self.cameras:
warnings.warn(f"There is no camera space: {camera}")
return

Expand All @@ -265,6 +275,7 @@ def render_pointcloud(self, seconds: float, channel: str, pointcloud: PointCloud
channel (str): Name of the pointcloud sensor channel.
pointcloud (PointCloudLike): Inherence object of `PointCloud`.
"""
# TODO(ktro2828): add support of rendering pointcloud on images
rr.set_time_seconds(self.timeline, seconds)

colors = distance_color(np.linalg.norm(pointcloud.points[:3].T, axis=1))
Expand All @@ -284,7 +295,7 @@ def render_image(self, seconds: float, camera: str, image: str | NDArrayU8) -> N
camera (str): Name of the camera channel.
image (str | NDArrayU8): Image tensor or path of the image file.
"""
if self.cameras is None or camera not in self.cameras:
if not self.with_2d or camera not in self.cameras:
warnings.warn(f"There is no camera space: {camera}")
return

Expand All @@ -306,25 +317,13 @@ def _render_ego_with_schema(self, ego_pose: EgoPose) -> None:
Args:
ego_pose (EgoPose): `EgoPose` object.
"""
rr.set_time_seconds(self.timeline, us2sec(ego_pose.timestamp))

rotation_xyzw = np.roll(ego_pose.rotation.q, shift=-1)
rr.log(
self.ego_entity,
rr.Transform3D(
translation=ego_pose.translation,
rotation=rr.Quaternion(xyzw=rotation_xyzw),
relation=rr.TransformRelation.ParentFromChild,
),
self._render_ego_without_schema(
seconds=us2sec(ego_pose.timestamp),
translation=ego_pose.translation,
rotation=ego_pose.rotation,
geocoordinate=ego_pose.geocoordinate,
)

if ego_pose.geocoordinate is not None:
latitude, longitude, _ = ego_pose.geocoordinate
rr.log(
self.geocoordinate_entity,
rr.GeoPoints(lat_lon=(latitude, longitude), radii=rr.Radius.ui_points(10.0)),
)

@render_ego.register
def _render_ego_without_schema(
self,
Expand All @@ -333,6 +332,16 @@ def _render_ego_without_schema(
rotation: RotationType,
geocoordinate: GeoCoordinateType | None = None,
) -> None:
"""Render an ego pose.
Args:
seconds (float): Timestamp in [sec].
translation (TranslationType): 3D position in the map coordinate system
, in the order of (x, y, z) in [m].
rotation (RotationType): Rotation in the map coordinate system.
geocoordinate (GeoCoordinateType | None, optional): Coordinates in the WGS 84
reference ellipsoid (latitude, longitude, altitude) in degrees and meters.
"""
rr.set_time_seconds(self.timeline, seconds)

rotation_xyzw = np.roll(rotation.q, shift=-1)
Expand Down Expand Up @@ -368,24 +377,14 @@ def _render_calibration_with_schema(
sensor (Sensor): `Sensor` object.
calibration (CalibratedSensor): `CalibratedSensor` object.
"""
rotation_xyzw = np.roll(calibration.rotation.q, shift=-1)

rr.log(
format_entity(self.ego_entity, sensor.channel),
rr.Transform3D(
translation=calibration.translation,
rotation=rr.Quaternion(xyzw=rotation_xyzw),
),
static=True,
self._render_calibration_without_schema(
channel=sensor.channel,
modality=sensor.modality,
translation=calibration.translation,
rotation=calibration.rotation,
camera_intrinsic=calibration.camera_intrinsic,
)

if sensor.modality == SensorModality.CAMERA:
rr.log(
format_entity(self.ego_entity, sensor.channel),
rr.Pinhole(image_from_camera=calibration.camera_intrinsic),
static=True,
)

@render_calibration.register
def _render_calibration_without_schema(
self,
Expand Down Expand Up @@ -419,9 +418,3 @@ def _render_calibration_without_schema(
rr.Pinhole(image_from_camera=camera_intrinsic),
static=True,
)
if modality == SensorModality.CAMERA:
rr.log(
format_entity(self.ego_entity, channel),
rr.Pinhole(image_from_camera=camera_intrinsic),
static=True,
)

0 comments on commit 517ebc2

Please sign in to comment.