diff --git a/t4_devkit/dataclass/box.py b/t4_devkit/dataclass/box.py index 8762c8e..5fc79a3 100644 --- a/t4_devkit/dataclass/box.py +++ b/t4_devkit/dataclass/box.py @@ -3,13 +3,13 @@ from typing import TYPE_CHECKING, TypeVar import numpy as np -from attrs import define, field, validators -from attrs.converters import optional +from attrs import converters, define, field, validators from shapely.geometry import Polygon from typing_extensions import Self from t4_devkit.common.converter import to_quaternion from t4_devkit.common.validator import is_vector3 +from t4_devkit.schema import VisibilityLevel from .label import SemanticLabel from .roi import Roi @@ -89,6 +89,7 @@ class Box3D(BaseBox): shape (Shape): `Shape` object. velocity (Vector3Like | None, optional): Box velocity (vx, vy, vz). num_points (int | None, optional): The number of points inside the box. + visibility (VisibilityLevel, optional): Box visibility. future (Future | None, optional): Box trajectory in the future of each mode. Examples: @@ -116,13 +117,18 @@ class Box3D(BaseBox): shape: Shape = field(validator=validators.instance_of(Shape)) velocity: Vector3Like | None = field( default=None, - converter=optional(np.array), + converter=converters.optional(np.array), validator=validators.optional(is_vector3), ) num_points: int | None = field( default=None, validator=validators.optional((validators.instance_of(int), validators.ge(0))), ) + visibility: VisibilityLevel = field( + default=VisibilityLevel.UNAVAILABLE, + converter=converters.optional(VisibilityLevel), + validator=validators.optional(validators.instance_of(VisibilityLevel)), + ) # additional attributes: set by `with_**` future: Future | None = field( diff --git a/t4_devkit/tier4.py b/t4_devkit/tier4.py index b171a1c..971ac4a 100644 --- a/t4_devkit/tier4.py +++ b/t4_devkit/tier4.py @@ -9,13 +9,7 @@ from pyquaternion import Quaternion from t4_devkit.common.geometry import is_box_in_image -from t4_devkit.dataclass import ( - Box2D, - Box3D, - SemanticLabel, - Shape, - ShapeType, -) +from t4_devkit.dataclass import Box2D, Box3D, SemanticLabel, Shape, ShapeType from t4_devkit.helper import RenderingHelper, TimeseriesHelper from t4_devkit.schema import SchemaName, SensorModality, VisibilityLevel, build_schema @@ -416,6 +410,7 @@ def get_box3d(self, sample_annotation_token: str, *, future_seconds: float = 0.0 ann: SampleAnnotation = self.get("sample_annotation", sample_annotation_token) instance: Instance = self.get("instance", ann.instance_token) sample: Sample = self.get("sample", ann.sample_token) + visibility: Visibility = self.get("visibility", ann.visibility_token) # semantic label semantic_label = self.get_semantic_label( @@ -438,6 +433,8 @@ def get_box3d(self, sample_annotation_token: str, *, future_seconds: float = 0.0 velocity=velocity, confidence=1.0, uuid=instance.token, # TODO(ktro2828): extract uuid from `instance_name`. + num_points=ann.num_lidar_pts, + visibility=visibility.level, ) if future_seconds > 0.0: @@ -549,6 +546,7 @@ def get_box3ds(self, sample_data_token: str, *, future_seconds: float = 0.0) -> instance.category_token, curr_ann.attribute_tokens ) velocity = self.box_velocity(curr_ann.token) + visibility: Visibility = self.get("visibility", curr_ann.visibility_token) box = Box3D( unix_time=t, @@ -560,6 +558,8 @@ def get_box3ds(self, sample_data_token: str, *, future_seconds: float = 0.0) -> velocity=velocity, confidence=1.0, uuid=instance.token, # TODO(ktro2828): extract uuid from `instance_name`. + num_points=curr_ann.num_lidar_pts, + visibility=visibility.level, ) else: # If not, simply grab the current annotation. diff --git a/tests/conftest.py b/tests/conftest.py index 26b1629..e6d12e9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,6 +16,7 @@ ShapeType, TransformBuffer, ) +from t4_devkit.schema import VisibilityLevel if TYPE_CHECKING: from t4_devkit.typing import NDArrayFloat @@ -58,6 +59,8 @@ def dummy_box3d() -> Box3D: velocity=(1.0, 1.0, 1.0), confidence=1.0, uuid="car3d_0", + num_points=1, + visibility=VisibilityLevel.FULL, ).with_future( timestamps=[101, 102, 103, 104], confidences=[1.0, 0.5], @@ -100,6 +103,8 @@ def dummy_box3ds() -> list[Box3D]: velocity=(1.0, 1.0, 1.0), confidence=1.0, uuid="car3d_1", + num_points=1, + visibility=VisibilityLevel.FULL, ).with_future( timestamps=[101, 102, 103, 104], confidences=[1.0], @@ -124,6 +129,8 @@ def dummy_box3ds() -> list[Box3D]: velocity=(1.0, 1.0, 1.0), confidence=1.0, uuid="bicycle3d_1", + num_points=1, + visibility=VisibilityLevel.FULL, ).with_future( timestamps=[101, 102, 103, 104], confidences=[1.0, 0.5], @@ -156,6 +163,8 @@ def dummy_box3ds() -> list[Box3D]: velocity=(1.0, 1.0, 1.0), confidence=1.0, uuid="pedestrian3d_1", + num_points=1, + visibility="full", ).with_future( timestamps=[101, 102, 103, 104], confidences=[1.0, 0.5, 0.2],