diff --git a/.flake8 b/.flake8 index f2b974ae..c77f18b6 100644 --- a/.flake8 +++ b/.flake8 @@ -16,4 +16,7 @@ ignore = E731 # Local variable name is assigned to but never used F841 +per-file-ignores = + # imported but unused + __init__.py: F401 max-line-length = 120 diff --git a/.gitignore b/.gitignore index 22cfdee9..bdc0a21f 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,10 @@ *.csv *.log *.mp4 + + +# Sphinx documentation +docs/_build/ +docs/build/ +_build/ +.doctrees/ diff --git a/d123/common/datatypes/camera/camera_parameters.py b/d123/common/datatypes/camera/camera_parameters.py deleted file mode 100644 index db06db93..00000000 --- a/d123/common/datatypes/camera/camera_parameters.py +++ /dev/null @@ -1,31 +0,0 @@ -from dataclasses import dataclass - -import numpy as np -import numpy.typing as npt - - -@dataclass -class CameraParameters: - - intrinsics: npt.NDArray[np.float64] # 3x3 matrix - distortion: npt.NDArray[np.float64] # 5x1 vector - translation: npt.NDArray[np.float64] # 3x1 vector - rotation: npt.NDArray[np.float64] # 3x3 matrix - - -def get_nuplan_camera_parameters() -> CameraParameters: - # return CameraParameters(focal_length=focal_length, image_size=image_size) - return CameraParameters( - intrinsics=np.array( - [[1.545e03, 0.000e00, 9.600e02], [0.000e00, 1.545e03, 5.600e02], [0.000e00, 0.000e00, 1.000e00]] - ), - distortion=np.array([-0.356123, 0.172545, -0.00213, 0.000464, -0.05231]), - translation=np.array([1.66433035e00, -1.32379618e-03, 1.57190200e00]), - rotation=np.array( - [ - [-0.00395669, -0.03969443, 0.99920403], - [-0.99971496, -0.02336898, -0.00488707], - [0.02354437, -0.99893856, -0.03959065], - ] - ), - ) diff --git a/d123/common/datatypes/detection/detection.py b/d123/common/datatypes/detection/detection.py index 2f6f0a71..a836fd4f 100644 --- a/d123/common/datatypes/detection/detection.py +++ b/d123/common/datatypes/detection/detection.py @@ -57,6 +57,10 @@ def shapely_polygon(self) -> shapely.geometry.Polygon: def center(self) -> StateSE3: return self.bounding_box_se3.center + @property + def center_se3(self) -> StateSE3: + return self.bounding_box_se3.center_se3 + @property def bounding_box(self) -> BoundingBoxSE3: return self.bounding_box_se3 diff --git a/d123/common/datatypes/detection/detection_types.py b/d123/common/datatypes/detection/detection_types.py index 8df27d2f..a22a614c 100644 --- a/d123/common/datatypes/detection/detection_types.py +++ b/d123/common/datatypes/detection/detection_types.py @@ -24,6 +24,7 @@ class DetectionType(SerialIntEnum): GENERIC_OBJECT = 6 # Animals, debris, pushable/pullable objects, permanent poles. EGO = 7 + SIGN = 8 # TODO: Remove or extent DYNAMIC_DETECTION_TYPES: set[DetectionType] = { diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py new file mode 100644 index 00000000..56fe6f07 --- /dev/null +++ b/d123/common/datatypes/sensor/camera.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Dict + +import numpy as np +import numpy.typing as npt + +from d123.common.utils.enums import SerialIntEnum + + +class CameraType(SerialIntEnum): + """ + Enum for cameras in d123. + """ + + CAM_F0 = 0 + CAM_B0 = 1 + CAM_L0 = 2 + CAM_L1 = 3 + CAM_L2 = 4 + CAM_R0 = 5 + CAM_R1 = 6 + CAM_R2 = 7 + CAM_STEREO_L = 8 + CAM_STEREO_R = 9 + + +@dataclass +class CameraMetadata: + + camera_type: CameraType + width: int + height: int + intrinsic: npt.NDArray[np.float64] # 3x3 matrix # TODO: don't store matrix but values. + distortion: npt.NDArray[np.float64] # 5x1 vector # TODO: don't store matrix but values. + + def to_dict(self) -> Dict[str, Any]: + # TODO: remove None types. Only a placeholder for now. + return { + "camera_type": int(self.camera_type), + "width": self.width, + "height": self.height, + "intrinsic": self.intrinsic.tolist() if self.intrinsic is not None else None, + "distortion": self.distortion.tolist() if self.distortion is not None else None, + } + + @classmethod + def from_dict(cls, json_dict: Dict[str, Any]) -> CameraMetadata: + # TODO: remove None types. Only a placeholder for now. + return cls( + camera_type=CameraType(json_dict["camera_type"]), + width=json_dict["width"], + height=json_dict["height"], + intrinsic=np.array(json_dict["intrinsic"]) if json_dict["intrinsic"] is not None else None, + distortion=np.array(json_dict["distortion"]) if json_dict["distortion"] is not None else None, + ) + + @property + def aspect_ratio(self) -> float: + return self.width / self.height + + @property + def fov_x(self) -> float: + """ + Calculates the horizontal field of view (FOV) in radian. + """ + fx = self.intrinsic[0, 0] + fov_x_rad = 2 * np.arctan(self.width / (2 * fx)) + return fov_x_rad + + @property + def fov_y(self) -> float: + """ + Calculates the vertical field of view (FOV) in radian. + """ + fy = self.intrinsic[1, 1] + fov_y_rad = 2 * np.arctan(self.height / (2 * fy)) + return fov_y_rad + + +def camera_metadata_dict_to_json(camera_metadata: Dict[CameraType, CameraMetadata]) -> Dict[str, Dict[str, Any]]: + """ + Converts a dictionary of CameraMetadata to a JSON-serializable format. + :param camera_metadata: Dictionary of CameraMetadata. + :return: JSON-serializable dictionary. + """ + camera_metadata_dict = { + camera_type.serialize(): metadata.to_dict() for camera_type, metadata in camera_metadata.items() + } + return json.dumps(camera_metadata_dict) + + +def camera_metadata_dict_from_json(json_dict: Dict[str, Dict[str, Any]]) -> Dict[CameraType, CameraMetadata]: + """ + Converts a JSON-serializable dictionary back to a dictionary of CameraMetadata. + :param json_dict: JSON-serializable dictionary. + :return: Dictionary of CameraMetadata. + """ + camera_metadata_dict = json.loads(json_dict) + return { + CameraType.deserialize(camera_type): CameraMetadata.from_dict(metadata) + for camera_type, metadata in camera_metadata_dict.items() + } + + +@dataclass +class Camera: + + metadata: CameraMetadata + image: npt.NDArray[np.uint8] + extrinsic: npt.NDArray[np.float64] # 4x4 matrix + + def get_view_matrix(self) -> np.ndarray: + # Compute the view matrix based on the camera's position and orientation + pass diff --git a/d123/common/datatypes/sensor/lidar.py b/d123/common/datatypes/sensor/lidar.py new file mode 100644 index 00000000..fe178760 --- /dev/null +++ b/d123/common/datatypes/sensor/lidar.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Dict, Optional, Type + +import numpy as np +import numpy.typing as npt + +from d123.common.datatypes.sensor.lidar_index import LIDAR_INDEX_REGISTRY, LiDARIndex +from d123.common.utils.enums import SerialIntEnum + + +class LiDARType(SerialIntEnum): + + LIDAR_UNKNOWN = 0 + LIDAR_MERGED = 1 + LIDAR_TOP = 2 + LIDAR_FRONT = 3 + LIDAR_SIDE_LEFT = 4 + LIDAR_SIDE_RIGHT = 5 + LIDAR_BACK = 6 + + +@dataclass +class LiDARMetadata: + + lidar_type: LiDARType + lidar_index: Type[LiDARIndex] + extrinsic: Optional[npt.NDArray[np.float64]] = None # 4x4 matrix + + # TODO: add identifier if point cloud is returned in lidar or ego frame. + + def to_dict(self) -> dict: + return { + "lidar_type": self.lidar_type.name, + "lidar_index": self.lidar_index.__name__, + "extrinsic": self.extrinsic.tolist() if self.extrinsic is not None else None, + } + + @classmethod + def from_dict(cls, json_dict: dict) -> LiDARMetadata: + lidar_type = LiDARType[json_dict["lidar_type"]] + if json_dict["lidar_index"] not in LIDAR_INDEX_REGISTRY: + raise ValueError(f"Unknown lidar index: {json_dict['lidar_index']}") + lidar_index_class = LIDAR_INDEX_REGISTRY[json_dict["lidar_index"]] + extrinsic = np.array(json_dict["extrinsic"]) if json_dict["extrinsic"] is not None else None + return cls(lidar_type=lidar_type, lidar_index=lidar_index_class, extrinsic=extrinsic) + + +def lidar_metadata_dict_to_json(lidar_metadata: Dict[LiDARType, LiDARMetadata]) -> str: + """ + Converts a dictionary of LiDARMetadata to a JSON-serializable format. + :param lidar_metadata: Dictionary of LiDARMetadata. + :return: JSON string. + """ + lidar_metadata_dict = { + lidar_type.serialize(): metadata.to_dict() for lidar_type, metadata in lidar_metadata.items() + } + return json.dumps(lidar_metadata_dict) + + +def lidar_metadata_dict_from_json(json_str: str) -> Dict[LiDARType, LiDARMetadata]: + """ + Converts a JSON string back to a dictionary of LiDARMetadata. + :param json_str: JSON string. + :return: Dictionary of LiDARMetadata. + """ + lidar_metadata_dict = json.loads(json_str) + return { + LiDARType.deserialize(lidar_type): LiDARMetadata.from_dict(metadata) + for lidar_type, metadata in lidar_metadata_dict.items() + } + + +@dataclass +class LiDAR: + + metadata: LiDARMetadata + point_cloud: npt.NDArray[np.float32] + + @property + def xyz(self) -> npt.NDArray[np.float32]: + """ + Returns the point cloud as an Nx3 array of x, y, z coordinates. + """ + return self.point_cloud[self.metadata.lidar_index.XYZ].T + + @property + def xy(self) -> npt.NDArray[np.float32]: + """ + Returns the point cloud as an Nx2 array of x, y coordinates. + """ + return self.point_cloud[self.metadata.lidar_index.XY].T + + @property + def intensity(self) -> Optional[npt.NDArray[np.float32]]: + """ + Returns the intensity values of the LiDAR point cloud if available. + Returns None if intensity is not part of the point cloud. + """ + if hasattr(self.metadata.lidar_index, "INTENSITY"): + return self.point_cloud[self.metadata.lidar_index.INTENSITY] + return None + + @property + def range(self) -> Optional[npt.NDArray[np.float32]]: + """ + Returns the range values of the LiDAR point cloud if available. + Returns None if range is not part of the point cloud. + """ + if hasattr(self.metadata.lidar_index, "RANGE"): + return self.point_cloud[self.metadata.lidar_index.RANGE] + return None + + @property + def elongation(self) -> Optional[npt.NDArray[np.float32]]: + """ + Returns the elongation values of the LiDAR point cloud if available. + Returns None if elongation is not part of the point cloud. + """ + if hasattr(self.metadata.lidar_index, "ELONGATION"): + return self.point_cloud[self.metadata.lidar_index.ELONGATION] + return None diff --git a/d123/common/datatypes/sensor/lidar_index.py b/d123/common/datatypes/sensor/lidar_index.py new file mode 100644 index 00000000..0df92cff --- /dev/null +++ b/d123/common/datatypes/sensor/lidar_index.py @@ -0,0 +1,62 @@ +from enum import IntEnum + +from d123.common.utils.enums import classproperty + +LIDAR_INDEX_REGISTRY = {} + + +def register_lidar_index(enum_class): + LIDAR_INDEX_REGISTRY[enum_class.__name__] = enum_class + return enum_class + + +class LiDARIndex(IntEnum): + + @classproperty + def XY(self) -> slice: + """ + Returns a slice for the XY coordinates of the LiDAR point cloud. + """ + return slice(self.X, self.Y + 1) + + @classproperty + def XYZ(self) -> slice: + """ + Returns a slice for the XYZ coordinates of the LiDAR point cloud. + """ + return slice(self.X, self.Z + 1) + + +@register_lidar_index +class DefaultLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + + +@register_lidar_index +class NuplanLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + RING = 4 + ID = 5 + + +@register_lidar_index +class CarlaLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + + +@register_lidar_index +class WopdLidarIndex(LiDARIndex): + RANGE = 0 + INTENSITY = 1 + ELONGATION = 2 + X = 3 + Y = 4 + Z = 5 diff --git a/d123/common/datatypes/time/time_point.py b/d123/common/datatypes/time/time_point.py index 9460f78b..437bf967 100644 --- a/d123/common/datatypes/time/time_point.py +++ b/d123/common/datatypes/time/time_point.py @@ -206,6 +206,16 @@ def __post_init__(self) -> None: """ assert self.time_us >= 0, "Time point has to be positive!" + @classmethod + def from_ns(cls, t_ns: int) -> TimePoint: + """ + Constructs a TimePoint from a value in nanoseconds. + :param t_ns: Time in nanoseconds. + :return: TimePoint. + """ + assert isinstance(t_ns, int), "Nanoseconds must be an integer!" + return TimePoint(time_us=t_ns // 1000) + @classmethod def from_us(cls, t_us: int) -> TimePoint: """ diff --git a/d123/common/datatypes/vehicle_state/vehicle_parameters.py b/d123/common/datatypes/vehicle_state/vehicle_parameters.py index c09e9b6d..3b4f04b9 100644 --- a/d123/common/datatypes/vehicle_state/vehicle_parameters.py +++ b/d123/common/datatypes/vehicle_state/vehicle_parameters.py @@ -20,10 +20,10 @@ class VehicleParameters: rear_axle_to_center_longitudinal: float -def get_nuplan_pacifica_parameters() -> VehicleParameters: +def get_nuplan_chrysler_pacifica_parameters() -> VehicleParameters: # NOTE: use parameters from nuPlan dataset return VehicleParameters( - vehicle_name="nuplan_pacifica", + vehicle_name="nuplan_chrysler_pacifica", width=2.297, length=5.176, height=1.777, @@ -46,6 +46,35 @@ def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: ) +def get_wopd_chrysler_pacifica_parameters() -> VehicleParameters: + # NOTE: use parameters from nuPlan dataset + # Find better parameters for WOPD ego vehicle + return VehicleParameters( + vehicle_name="wopd_chrysler_pacifica", + width=2.297, + length=5.176, + height=1.777, + wheel_base=3.089, + rear_axle_to_center_vertical=1.777 / 2, + rear_axle_to_center_longitudinal=1.461, + ) + + +def get_av2_ford_fusion_hybrid_parameters() -> VehicleParameters: + # NOTE: Parameters are estimated from the vehicle model. + # https://en.wikipedia.org/wiki/Ford_Fusion_Hybrid#Second_generation + # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/tests/unit/map/test_map_api.py#L375 + return VehicleParameters( + vehicle_name="av2_ford_fusion_hybrid", + width=1.852 + 0.275, # 0.275 is the estimated width of the side mirrors + length=4.869, + height=1.476, + wheel_base=2.850, + rear_axle_to_center_vertical=0.438, + rear_axle_to_center_longitudinal=1.339, + ) + + def center_se3_to_rear_axle_se3(center_se3: StateSE3, vehicle_parameters: VehicleParameters) -> StateSE3: """ Converts a center state to a rear axle state. diff --git a/d123/common/geometry/base.py b/d123/common/geometry/base.py index 460cff7d..7fdec5cc 100644 --- a/d123/common/geometry/base.py +++ b/d123/common/geometry/base.py @@ -8,6 +8,7 @@ import numpy.typing as npt import shapely.geometry as geom +# from d123.common.geometry.transform.se3 import get_rotation_matrix from d123.common.utils.enums import classproperty # TODO: Reconsider if 2D/3D or SE2/SE3 structure would be better hierarchical, e.g. inheritance or composition. @@ -144,8 +145,10 @@ class Point3D: @classmethod def from_array(cls, array: npt.NDArray[np.float64]) -> "Point3D": - assert array.ndim == 1 - assert array.shape[0] == len(Point3DIndex) + assert array.ndim == 1, f"Array must be 1-dimensional, got shape {array.shape}" + assert array.shape[0] == len( + Point3DIndex + ), f"Array must have the same length as Point3DIndex, got shape {array.shape}" return cls(array[Point3DIndex.X], array[Point3DIndex.Y], array[Point3DIndex.Z]) @property @@ -180,7 +183,6 @@ def __hash__(self) -> int: class StateSE3Index(IntEnum): - # TODO: implement X = 0 Y = 1 @@ -227,6 +229,21 @@ def from_array(cls, array: npt.NDArray[np.float64]) -> StateSE3: array[StateSE3Index.YAW], ) + @classmethod + def from_matrix(cls, array: npt.NDArray[np.float64]) -> StateSE3: + assert array.ndim == 2 + assert array.shape == (4, 4) + translation = array[:3, 3] + rotation = array[:3, :3] + return StateSE3( + x=translation[0], + y=translation[1], + z=translation[2], + roll=np.arctan2(rotation[2, 1], rotation[2, 2]), + pitch=np.arctan2(-rotation[2, 0], np.sqrt(rotation[2, 1] ** 2 + rotation[2, 2] ** 2)), + yaw=np.arctan2(rotation[1, 0], rotation[0, 0]), + ) + @property def array(self) -> npt.NDArray[np.float64]: array = np.zeros(len(StateSE3Index), dtype=np.float64) @@ -238,6 +255,16 @@ def array(self) -> npt.NDArray[np.float64]: array[StateSE3Index.YAW] = self.yaw return array + # @property + # def matrix(self) -> npt.NDArray[np.float64]: + # """Convert SE3 state to 4x4 transformation matrix.""" + # R = get_rotation_matrix(self) + # translation = np.array([self.x, self.y, self.z], dtype=np.float64) + # matrix = np.eye(4, dtype=np.float64) + # matrix[:3, :3] = R + # matrix[:3, 3] = translation + # return matrix + @property def state_se2(self) -> StateSE2: return StateSE2(self.x, self.y, self.yaw) diff --git a/d123/common/geometry/base_index.py b/d123/common/geometry/base_index.py new file mode 100644 index 00000000..cb258618 --- /dev/null +++ b/d123/common/geometry/base_index.py @@ -0,0 +1 @@ +# TODO: Move base index here to avoid circular imports. diff --git a/d123/common/geometry/bounding_box/bounding_box.py b/d123/common/geometry/bounding_box/bounding_box.py index de7daf93..625d24d9 100644 --- a/d123/common/geometry/bounding_box/bounding_box.py +++ b/d123/common/geometry/bounding_box/bounding_box.py @@ -69,6 +69,10 @@ def bounding_box_se2(self) -> BoundingBoxSE2: width=self.width, ) + @property + def center_se3(self) -> StateSE3: + return self.center + @property def array(self) -> npt.NDArray[np.float64]: array = np.zeros(len(BoundingBoxSE3Index), dtype=np.float64) diff --git a/d123/common/geometry/bounding_box/bounding_box_index.py b/d123/common/geometry/bounding_box/bounding_box_index.py index d5a1ee26..c282b07e 100644 --- a/d123/common/geometry/bounding_box/bounding_box_index.py +++ b/d123/common/geometry/bounding_box/bounding_box_index.py @@ -48,3 +48,22 @@ def STATE_SE3(cls) -> slice: @classproperty def ROTATION_XYZ(cls) -> slice: return slice(cls.ROLL, cls.YAW + 1) + + +class Corners3DIndex(IntEnum): + FRONT_LEFT_BOTTOM = 0 + FRONT_RIGHT_BOTTOM = 1 + BACK_RIGHT_BOTTOM = 2 + BACK_LEFT_BOTTOM = 3 + FRONT_LEFT_TOP = 4 + FRONT_RIGHT_TOP = 5 + BACK_RIGHT_TOP = 6 + BACK_LEFT_TOP = 7 + + @classproperty + def BOTTOM(cls) -> slice: + return slice(cls.FRONT_LEFT_BOTTOM, cls.BACK_LEFT_BOTTOM + 1) + + @classproperty + def TOP(cls) -> slice: + return slice(cls.FRONT_LEFT_TOP, cls.BACK_LEFT_TOP + 1) diff --git a/d123/common/geometry/line/polylines.py b/d123/common/geometry/line/polylines.py index 881eef87..37b799b5 100644 --- a/d123/common/geometry/line/polylines.py +++ b/d123/common/geometry/line/polylines.py @@ -9,7 +9,7 @@ import shapely.geometry as geom from scipy.interpolate import interp1d -from d123.common.geometry.base import Point2D, Point2DIndex, Point3D, StateSE2, StateSE2Index +from d123.common.geometry.base import Point2D, Point2DIndex, Point3D, Point3DIndex, StateSE2, StateSE2Index from d123.common.geometry.constants import DEFAULT_Z from d123.common.geometry.line.helper import get_linestring_yaws, get_path_progress from d123.common.geometry.utils import normalize_angle @@ -33,7 +33,15 @@ def from_linestring(cls, linestring: geom.LineString) -> Polyline2D: @classmethod def from_array(cls, polyline_array: npt.NDArray[np.float32]) -> Polyline2D: - raise NotImplementedError + assert polyline_array.ndim == 2 + linestring: Optional[geom.LineString] = None + if polyline_array.shape[-1] == len(Point2DIndex): + linestring = geom_creation.linestrings(polyline_array) + elif polyline_array.shape[-1] == len(Point3DIndex): + linestring = geom_creation.linestrings(polyline_array[:, Point3DIndex.XY]) + else: + raise ValueError("Array must have shape (N, 2) or (N, 3) for Point2D or Point3D respectively.") + return Polyline2D(linestring) @property def array(self) -> npt.NDArray[np.float64]: @@ -79,6 +87,7 @@ def __post_init__(self): if self.linestring is None: self.linestring = geom_creation.linestrings(self.se2_array[..., StateSE2Index.XY]) + self.se2_array[:, StateSE2Index.YAW] = np.unwrap(self.se2_array[:, StateSE2Index.YAW], axis=0) self._progress = get_path_progress(self.se2_array) self._interpolator = interp1d(self._progress, self.se2_array, axis=0, bounds_error=False, fill_value=0.0) @@ -149,10 +158,15 @@ def from_linestring(cls, linestring: geom.LineString) -> Polyline3D: else Polyline3D(geom_creation.linestrings(*linestring.xy, z=DEFAULT_Z)) ) + @classmethod + def from_array(cls, array: npt.NDArray[np.float64]) -> Polyline3D: + assert array.ndim == 2 and array.shape[1] == 3, "Array must be 2D with shape (N, 3)" + linestring = geom_creation.linestrings(*array.T) + return Polyline3D(linestring) + @property def polyline_2d(self) -> Polyline2D: - Polyline2D(geom_creation.linestrings(*self.linestring.xy)) - return Polyline2D.from_linestring(self.linestring) + return Polyline2D(geom_creation.linestrings(*self.linestring.xy)) @property def polyline_se2(self) -> PolylineSE2: diff --git a/d123/common/geometry/occupancy_map.py b/d123/common/geometry/occupancy_map.py index ed80ae23..2a8085d6 100644 --- a/d123/common/geometry/occupancy_map.py +++ b/d123/common/geometry/occupancy_map.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Dict, List, Literal, Optional, Sequence, Union import numpy as np @@ -34,6 +36,12 @@ def __init__( self._node_capacity = node_capacity self._str_tree = STRtree(self._geometries, node_capacity) + @classmethod + def from_dict(cls, geometry_dict: Dict[Union[str, int], BaseGeometry], node_capacity: int = 10) -> OccupancyMap2D: + ids = list(geometry_dict.keys()) + geometries = list(geometry_dict.values()) + return cls(geometries=geometries, ids=ids, node_capacity=node_capacity) + def __getitem__(self, id: Union[str, int]) -> BaseGeometry: """ Retrieves geometry of token. diff --git a/d123/common/geometry/transform/rotation.py b/d123/common/geometry/transform/rotation.py new file mode 100644 index 00000000..03072609 --- /dev/null +++ b/d123/common/geometry/transform/rotation.py @@ -0,0 +1,27 @@ +# import numpy as np +# import numpy.typing as npt + +# from d123.common.geometry.base import Point3DIndex, StateSE3, StateSE3Index +# from d123.common.geometry.vector import Vector3D + + +# def get_roll_pitch_yaw_from_rotation_matrix( +# rotation_matrix: npt.NDArray[np.float64], +# ) -> Vector3D: +# """Extract roll, pitch, and yaw angles from a rotation matrix.""" +# assert rotation_matrix.shape == (3, 3), "Rotation matrix must be 3x3." + +# sy = np.sqrt(rotation_matrix[0, 0] ** 2 + rotation_matrix[1, 0] ** 2) + +# singular = sy < 1e-6 + +# if not singular: +# x = np.arctan2(rotation_matrix[2, 1], rotation_matrix[2, 2]) +# y = np.arctan2(-rotation_matrix[2, 0], sy) +# z = np.arctan2(rotation_matrix[1, 0], rotation_matrix[0, 0]) +# else: +# x = np.arctan2(-rotation_matrix[1, 2], rotation_matrix[1, 1]) +# y = np.arctan2(-rotation_matrix[2, 0], sy) +# z = 0.0 + +# return Vector3D(x=x, y=y, z=z) diff --git a/d123/common/geometry/transform/se3.py b/d123/common/geometry/transform/se3.py index 5bcf75b5..e61451b4 100644 --- a/d123/common/geometry/transform/se3.py +++ b/d123/common/geometry/transform/se3.py @@ -1,11 +1,39 @@ import numpy as np import numpy.typing as npt -from d123.common.geometry.base import StateSE3 +from d123.common.geometry.base import Point3DIndex, StateSE3, StateSE3Index from d123.common.geometry.vector import Vector3D +# def get_rotation_matrix(state_se3: StateSE3) -> npt.NDArray[np.float64]: +# R_x = np.array( +# [ +# [1, 0, 0], +# [0, np.cos(state_se3.roll), -np.sin(state_se3.roll)], +# [0, np.sin(state_se3.roll), np.cos(state_se3.roll)], +# ], +# dtype=np.float64, +# ) +# R_y = np.array( +# [ +# [np.cos(state_se3.pitch), 0, np.sin(state_se3.pitch)], +# [0, 1, 0], +# [-np.sin(state_se3.pitch), 0, np.cos(state_se3.pitch)], +# ], +# dtype=np.float64, +# ) +# R_z = np.array( +# [ +# [np.cos(state_se3.yaw), -np.sin(state_se3.yaw), 0], +# [np.sin(state_se3.yaw), np.cos(state_se3.yaw), 0], +# [0, 0, 1], +# ], +# dtype=np.float64, +# ) +# return R_z @ R_y @ R_x + def get_rotation_matrix(state_se3: StateSE3) -> npt.NDArray[np.float64]: + # Intrinsic Z-Y'-X'' rotation: R = R_x(roll) @ R_y(pitch) @ R_z(yaw) R_x = np.array( [ [1, 0, 0], @@ -30,7 +58,7 @@ def get_rotation_matrix(state_se3: StateSE3) -> npt.NDArray[np.float64]: ], dtype=np.float64, ) - return R_z @ R_y @ R_x + return R_x @ R_y @ R_z def translate_se3_along_z(state_se3: StateSE3, distance: float) -> StateSE3: @@ -85,3 +113,95 @@ def translate_body_frame(state_se3: StateSE3, vector_3d: Vector3D) -> StateSE3: state_se3.pitch, state_se3.yaw, ) + + +def convert_relative_to_absolute_points_3d_array( + origin: StateSE3, points_3d_array: npt.NDArray[np.float64] +) -> npt.NDArray[np.float64]: + + # TODO: implement function for origin as np.ndarray + + R = get_rotation_matrix(origin) + absolute_points = points_3d_array @ R.T + origin.point_3d.array + return absolute_points + + +def convert_absolute_to_relative_se3_array( + origin: StateSE3, se3_array: npt.NDArray[np.float64] +) -> npt.NDArray[np.float64]: + assert se3_array.shape[-1] == len(StateSE3Index) + # TODO: remove transform for-loop, use vectorized operations + + # Extract rotation and translation of origin + R_origin = get_rotation_matrix(origin) + t_origin = origin.point_3d.array + + # Prepare output array + rel_se3_array = np.empty_like(se3_array) + + # For each SE3 in the array + for i in range(se3_array.shape[0]): + abs_se3 = se3_array[i] + abs_pos = abs_se3[StateSE3Index.XYZ] + abs_rpy = abs_se3[StateSE3Index.ROLL : StateSE3Index.YAW + 1] + + # Relative position: rotate and translate + rel_pos = R_origin.T @ (abs_pos - t_origin) + + # Relative orientation: subtract origin's rpy + rel_rpy = abs_rpy - np.array([origin.roll, origin.pitch, origin.yaw], dtype=np.float64) + + rel_se3_array[i, StateSE3Index.X : StateSE3Index.Z + 1] = rel_pos + rel_se3_array[i, StateSE3Index.ROLL : StateSE3Index.YAW + 1] = rel_rpy + + return rel_se3_array + + +def convert_relative_to_absolute_se3_array( + origin: StateSE3, se3_array: npt.NDArray[np.float64] +) -> npt.NDArray[np.float64]: + assert se3_array.shape[-1] == len(StateSE3Index) + # TODO: remove transform for-loop, use vectorized operations + + # Extract rotation and translation of origin + R_origin = get_rotation_matrix(origin) + t_origin = origin.point_3d.array + + # Prepare output array + abs_se3_array = np.empty_like(se3_array) + + # For each SE3 in the array + for i in range(se3_array.shape[0]): + rel_se3 = se3_array[i] + rel_pos = rel_se3[StateSE3Index.XYZ] + rel_rpy = rel_se3[StateSE3Index.ROLL : StateSE3Index.YAW + 1] + + # Absolute position: rotate and translate + abs_pos = R_origin @ rel_pos + t_origin + + # Absolute orientation: add origin's rpy + abs_rpy = rel_rpy + np.array([origin.roll, origin.pitch, origin.yaw], dtype=np.float64) + + abs_se3_array[i, StateSE3Index.X : StateSE3Index.Z + 1] = abs_pos + abs_se3_array[i, StateSE3Index.ROLL : StateSE3Index.YAW + 1] = abs_rpy + + return abs_se3_array + + +def translate_points_3d_along_z( + state_se3: StateSE3, + points_3d: npt.NDArray[np.float64], + distance: float, +) -> npt.NDArray[np.float64]: + assert points_3d.shape[-1] == len(Point3DIndex) + + R = get_rotation_matrix(state_se3) + z_axis = R[:, 2] + + translated_points = np.zeros_like(points_3d) + + translated_points[..., Point3DIndex.X] = points_3d[..., Point3DIndex.X] + distance * z_axis[0] + translated_points[..., Point3DIndex.Y] = points_3d[..., Point3DIndex.Y] + distance * z_axis[1] + translated_points[..., Point3DIndex.Z] = points_3d[..., Point3DIndex.Z] + distance * z_axis[2] + + return translated_points diff --git a/d123/common/multithreading/ray_execution.py b/d123/common/multithreading/ray_execution.py new file mode 100644 index 00000000..4e44c56f --- /dev/null +++ b/d123/common/multithreading/ray_execution.py @@ -0,0 +1,122 @@ +import logging +import traceback +from functools import partial +from pathlib import Path +from typing import Any, Callable, Iterable, Iterator, List, Optional, Tuple, cast +from uuid import uuid1 + +import ray +from ray.exceptions import RayTaskError +from ray.remote_function import RemoteFunction +from tqdm import tqdm + +from d123.common.multithreading.worker_pool import Task + + +def _ray_object_iterator(initial_ids: List[ray.ObjectRef]) -> Iterator[Tuple[ray.ObjectRef, Any]]: + """ + Iterator that waits for each ray object in the input object list to be completed and fetches the result. + :param initial_ids: list of ray object ids + :yield: result of worker + """ + next_ids = initial_ids + + while next_ids: + ready_ids, not_ready_ids = ray.wait(next_ids) + next_id = ready_ids[0] + + yield next_id, ray.get(next_id) + + next_ids = not_ready_ids + + +def wrap_function(fn: Callable[..., Any], log_dir: Optional[Path] = None) -> Callable[..., Any]: + """ + Wraps a function to save its logs to a unique file inside the log directory. + :param fn: function to be wrapped. + :param log_dir: directory to store logs (wrapper function does nothing if it's not set). + :return: wrapped function which changes logging settings while it runs. + """ + + def wrapped_fn(*args: Any, **kwargs: Any) -> Any: + if log_dir is None: + return fn(*args, **kwargs) + + log_dir.mkdir(parents=True, exist_ok=True) + log_path = log_dir / f"{uuid1().hex}__{fn.__name__}.log" + + logging.basicConfig() + logger = logging.getLogger() + fh = logging.FileHandler(log_path, delay=True) + fh.setLevel(logging.INFO) + logger.addHandler(fh) + logger.setLevel(logging.INFO) + + # Silent botocore which is polluting the terminal because of serialization and deserialization + # with following message: INFO:botocore.credentials:Credentials found in config file: ~/.aws/config + logging.getLogger("botocore").setLevel(logging.WARNING) + + result = fn(*args, **kwargs) + + fh.flush() + fh.close() + logger.removeHandler(fh) + + return result + + return wrapped_fn + + +def _ray_map_items(task: Task, *item_lists: Iterable[List[Any]], log_dir: Optional[Path] = None) -> List[Any]: + """ + Map each item of a list of arguments to a callable and executes in parallel. + :param fn: callable to be run + :param item_list: items to be parallelized + :param log_dir: directory to store worker logs + :return: list of outputs + """ + assert len(item_lists) > 0, "No map arguments received for mapping" + assert all(isinstance(items, list) for items in item_lists), "All map arguments must be lists" + assert all( + len(cast(List, items)) == len(item_lists[0]) for items in item_lists # type: ignore + ), "All lists must have equal size" + fn = task.fn + # Wrap function in remote decorator and create ray objects + if isinstance(fn, partial): + _, _, pack = fn.__reduce__() # type: ignore + fn, _, args, _ = pack + fn = wrap_function(fn, log_dir=log_dir) + remote_fn: RemoteFunction = ray.remote(fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus) + object_ids = [remote_fn.remote(*items, **args) for items in zip(*item_lists)] + else: + fn = wrap_function(fn, log_dir=log_dir) + remote_fn = ray.remote(fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus) + object_ids = [remote_fn.remote(*items) for items in zip(*item_lists)] + + # Create ordered map to track order of objects inserted in the queue + object_result_map = dict.fromkeys(object_ids, None) + + # Asynchronously iterate through the object and track progress + for object_id, output in tqdm(_ray_object_iterator(object_ids), total=len(object_ids), desc="Ray objects"): + object_result_map[object_id] = output + + results = list(object_result_map.values()) + + return results + + +def ray_map(task: Task, *item_lists: Iterable[List[Any]], log_dir: Optional[Path] = None) -> List[Any]: + """ + Initialize ray, align item lists and map each item of a list of arguments to a callable and executes in parallel. + :param task: callable to be run + :param item_lists: items to be parallelized + :param log_dir: directory to store worker logs + :return: list of outputs + """ + try: + results = _ray_map_items(task, *item_lists, log_dir=log_dir) + return results + except (RayTaskError, Exception) as exc: + ray.shutdown() + traceback.print_exc() + raise RuntimeError(exc) diff --git a/d123/common/multithreading/worker_parallel.py b/d123/common/multithreading/worker_parallel.py new file mode 100644 index 00000000..737bd486 --- /dev/null +++ b/d123/common/multithreading/worker_parallel.py @@ -0,0 +1,60 @@ +import concurrent +import concurrent.futures +import logging +from concurrent.futures import Future +from typing import Any, Iterable, List, Optional + +from tqdm import tqdm + +from d123.common.multithreading.worker_pool import ( + Task, + WorkerPool, + WorkerResources, + get_max_size_of_arguments, +) + +logger = logging.getLogger(__name__) + + +class SingleMachineParallelExecutor(WorkerPool): + """ + This worker distributes all tasks across multiple threads on this machine. + """ + + def __init__(self, use_process_pool: bool = False, max_workers: Optional[int] = None): + """ + Create worker with limited threads. + :param use_process_pool: if true, ProcessPoolExecutor will be used as executor, otherwise ThreadPoolExecutor. + :param max_workers: if available, use this number as used number of threads. + """ + # Set the number of available threads + number_of_cpus_per_node = max_workers if max_workers else WorkerResources.current_node_cpu_count() + + super().__init__( + WorkerResources( + number_of_nodes=1, number_of_cpus_per_node=number_of_cpus_per_node, number_of_gpus_per_node=0 + ) + ) + + # Create executor + self._executor = ( + concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) + if use_process_pool + else concurrent.futures.ThreadPoolExecutor(max_workers=number_of_cpus_per_node) + ) + + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """Inherited, see superclass.""" + return list( + tqdm( + self._executor.map(task.fn, *item_lists), + leave=False, + total=get_max_size_of_arguments(*item_lists), + desc="SingleMachineParallelExecutor", + disable=not verbose, + ) + ) + + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """Inherited, see superclass.""" + return self._executor.submit(task.fn, *args, **kwargs) diff --git a/d123/common/multithreading/worker_pool.py b/d123/common/multithreading/worker_pool.py new file mode 100644 index 00000000..8ffea8ed --- /dev/null +++ b/d123/common/multithreading/worker_pool.py @@ -0,0 +1,156 @@ +import abc +import logging +from concurrent.futures import Future +from dataclasses import dataclass +from typing import Any, Callable, Iterable, List, Optional, Tuple, Union + +from psutil import cpu_count + +logger = logging.getLogger(__name__) + + +def get_max_size_of_arguments(*item_lists: Iterable[List[Any]]) -> int: + """ + Find the argument with most elements. + e.g. [db, [arg1, arg2] -> 2. + :param item_lists: arguments where some of the arguments is a list. + :return: size of largest list. + """ + lengths = [len(items) for items in item_lists if isinstance(items, list)] + if len(list(set(lengths))) > 1: + raise RuntimeError(f"There exists lists with different element size = {lengths}!") + return max(lengths) if len(lengths) != 0 else 1 + + +def align_size_of_arguments(*item_lists: Iterable[List[Any]]) -> Tuple[int, Iterable[List[Any]]]: + """ + Align item lists by repeating elements in order to achieve the same size. + eg. [db, [arg1, arg2] -> [[db, db], [arg1, arg2]]. + :param item_lists: multiple arguments which will be used to call a function. + :return: arguments with same dimension, e.g., [[db, db], [arg1, arg1]]. + """ + max_size = get_max_size_of_arguments(*item_lists) + aligned_item_lists = [items if isinstance(items, list) else [items] * max_size for items in item_lists] + return max_size, aligned_item_lists + + +@dataclass(frozen=True) +class Task: + """This class represents a task that can be submitted to a worker with specific resource requirements.""" + + fn: Callable[..., Any] # Function that should be called with arguments + + # Number of CPUs required for this task + # if None, the default number of CPUs will be used that were specified in initialization + num_cpus: Optional[int] = None + + # Fraction of GPUs required for this task, this number can be also smaller than 1 + # It num_gpus is smaller than 1, only a part of a GPU is allocated to this task. + # NOTE: it is the responsibility of a user to make sure that a model fits into num_gpus + # if None, no GPU will be used + num_gpus: Optional[Union[int, float]] = None + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """ + Call function with args. + :return: output from fn. + """ + return self.fn(*args, **kwargs) + + +@dataclass(frozen=True) +class WorkerResources: + """Data class to indicate resources used by workers.""" + + number_of_nodes: int # Number of available independent nodes + number_of_gpus_per_node: int # Number of GPUs per node + # Number of CPU logical cores per node, this can be smaller than + # current_node_cpu_count if user will decide to limit it + number_of_cpus_per_node: int + + @property + def number_of_threads(self) -> int: + """ + :return: the number of available threads across all nodes. + """ + return self.number_of_nodes * self.number_of_cpus_per_node + + @staticmethod + def current_node_cpu_count() -> int: + """ + :return: the number of logical cores on the current machine. + """ + return cpu_count(logical=True) # type: ignore + + +class WorkerPool(abc.ABC): + """ + This class executed function on list of arguments. This can either be distributed/parallel or sequential worker. + """ + + def __init__(self, config: WorkerResources): + """ + Initialize worker with resource description. + :param config: setup of this worker. + """ + self.config = config + + if self.config.number_of_threads < 1: + raise RuntimeError(f"Number of threads can not be 0, and it is {self.config.number_of_threads}!") + + logger.info(f"Worker: {self.__class__.__name__}") + logger.info(f"{self}") + + def map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """ + Run function with arguments from item_lists, this function will make sure all arguments have the same + number of elements. + :param task: function to be run. + :param item_lists: arguments to the function. + :param verbose: Whether to increase logger verbosity. + :return: type from the fn. + """ + max_size, aligned_item_lists = align_size_of_arguments(*item_lists) + + if verbose: + logger.info(f"Submitting {max_size} tasks!") + return self._map(task, *aligned_item_lists, verbose=verbose) + + @abc.abstractmethod + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """ + Run function with arguments from item_lists. This function can assume that all the args in item_lists have + the same number of elements. + :param fn: function to be run. + :param item_lists: arguments to the function. + :param number_of_elements: number of calls to the function. + :return: type from the fn. + """ + + @abc.abstractmethod + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """ + Submit a task to the worker. + :param task: to be submitted. + :param args: arguments for the task. + :param kwargs: keyword arguments for the task. + :return: future. + """ + + @property + def number_of_threads(self) -> int: + """ + :return: the number of available threads across all nodes. + """ + return self.config.number_of_threads + + def __str__(self) -> str: + """ + :return: string with information about this worker. + """ + return ( + f"Number of nodes: {self.config.number_of_nodes}\n" + f"Number of CPUs per node: {self.config.number_of_cpus_per_node}\n" + f"Number of GPUs per node: {self.config.number_of_gpus_per_node}\n" + f"Number of threads across all nodes: {self.config.number_of_threads}" + ) diff --git a/d123/common/multithreading/worker_ray.py b/d123/common/multithreading/worker_ray.py new file mode 100644 index 00000000..8bb12b04 --- /dev/null +++ b/d123/common/multithreading/worker_ray.py @@ -0,0 +1,166 @@ +import logging +import os +from concurrent.futures import Future +from pathlib import Path +from typing import Any, Iterable, List, Optional, Union + +import ray +from psutil import cpu_count + +from d123.common.multithreading.ray_execution import ray_map +from d123.common.multithreading.worker_pool import Task, WorkerPool, WorkerResources + +logger = logging.getLogger(__name__) + +# Silent botocore which is polluting the terminal because of serialization and deserialization +# with following message: INFO:botocore.credentials:Credentials found in config file: ~/.aws/config +logging.getLogger("botocore").setLevel(logging.WARNING) + + +def initialize_ray( + master_node_ip: Optional[str] = None, + threads_per_node: Optional[int] = None, + local_mode: bool = False, + log_to_driver: bool = True, + use_distributed: bool = False, +) -> WorkerResources: + """ + Initialize ray worker. + ENV_VAR_MASTER_NODE_IP="master node IP". + ENV_VAR_MASTER_NODE_PASSWORD="password to the master node". + ENV_VAR_NUM_NODES="number of nodes available". + :param master_node_ip: if available, ray will connect to remote cluster. + :param threads_per_node: Number of threads to use per node. + :param log_to_driver: If true, the output from all of the worker + processes on all nodes will be directed to the driver. + :param local_mode: If true, the code will be executed serially. This + is useful for debugging. + :param use_distributed: If true, and the env vars are available, + ray will launch in distributed mode + :return: created WorkerResources. + """ + # Env variables which are set through SLURM script + env_var_master_node_ip = "ip_head" + env_var_master_node_password = "redis_password" + env_var_num_nodes = "num_nodes" + + # Read number of CPU cores on current machine + number_of_cpus_per_node = threads_per_node if threads_per_node else cpu_count(logical=True) + + try: + import torch + + number_of_gpus_per_node = torch.cuda.device_count() if torch.cuda.is_available() else 0 + except ImportError: + number_of_gpus_per_node = 0 + if not number_of_gpus_per_node: + logger.info("Not using GPU in ray") + + # Find a way in how the ray should be initialized + if master_node_ip and use_distributed: + # Connect to ray remotely to node ip + logger.info(f"Connecting to cluster at: {master_node_ip}!") + ray.init(address=f"ray://{master_node_ip}:10001", local_mode=local_mode, log_to_driver=log_to_driver) + number_of_nodes = 1 + elif env_var_master_node_ip in os.environ and use_distributed: + # In this way, we started ray on the current machine which generated password and master node ip: + # It was started with "ray start --head" + number_of_nodes = int(os.environ[env_var_num_nodes]) + master_node_ip = os.environ[env_var_master_node_ip].split(":")[0] + redis_password = os.environ[env_var_master_node_password].split(":")[0] + logger.info(f"Connecting as part of a cluster at: {master_node_ip} with password: {redis_password}!") + # Connect to cluster, follow to https://docs.ray.io/en/latest/package-ref.html for more info + ray.init( + address="auto", + _node_ip_address=master_node_ip, + _redis_password=redis_password, + log_to_driver=log_to_driver, + local_mode=local_mode, + ) + else: + # In this case, we will just start ray directly from this script + number_of_nodes = 1 + logger.info("Starting ray local!") + ray.init( + num_cpus=number_of_cpus_per_node, + dashboard_host="0.0.0.0", + local_mode=local_mode, + log_to_driver=log_to_driver, + ) + + return WorkerResources( + number_of_nodes=number_of_nodes, + number_of_cpus_per_node=number_of_cpus_per_node, + number_of_gpus_per_node=number_of_gpus_per_node, + ) + + +class RayDistributed(WorkerPool): + """ + This worker uses ray to distribute work across all available threads. + """ + + def __init__( + self, + master_node_ip: Optional[str] = None, + threads_per_node: Optional[int] = None, + debug_mode: bool = False, + log_to_driver: bool = True, + output_dir: Optional[Union[str, Path]] = None, + logs_subdir: Optional[str] = "logs", + use_distributed: bool = False, + ): + """ + Initialize ray worker. + :param master_node_ip: if available, ray will connect to remote cluster. + :param threads_per_node: Number of threads to use per node. + :param debug_mode: If true, the code will be executed serially. This + is useful for debugging. + :param log_to_driver: If true, the output from all of the worker + processes on all nodes will be directed to the driver. + :param output_dir: Experiment output directory. + :param logs_subdir: Subdirectory inside experiment dir to store worker logs. + :param use_distributed: Boolean flag to explicitly enable/disable distributed computation + """ + self._master_node_ip = master_node_ip + self._threads_per_node = threads_per_node + self._local_mode = debug_mode + self._log_to_driver = log_to_driver + self._log_dir: Optional[Path] = Path(output_dir) / (logs_subdir or "") if output_dir is not None else None + self._use_distributed = use_distributed + super().__init__(self.initialize()) + + def initialize(self) -> WorkerResources: + """ + Initialize ray. + :return: created WorkerResources. + """ + # In case ray was already running, shut it down. This occurs mainly in tests + if ray.is_initialized(): + logger.warning("Ray is running, we will shut it down before starting again!") + ray.shutdown() + + return initialize_ray( + master_node_ip=self._master_node_ip, + threads_per_node=self._threads_per_node, + local_mode=self._local_mode, + log_to_driver=self._log_to_driver, + use_distributed=self._use_distributed, + ) + + def shutdown(self) -> None: + """ + Shutdown the worker and clear memory. + """ + ray.shutdown() + + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """Inherited, see superclass.""" + del verbose + return ray_map(task, *item_lists, log_dir=self._log_dir) # type: ignore + + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """Inherited, see superclass.""" + remote_fn = ray.remote(task.fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus) + object_ids: ray._raylet.ObjectRef = remote_fn.remote(*args, **kwargs) + return object_ids.future() # type: ignore diff --git a/d123/common/multithreading/worker_sequential.py b/d123/common/multithreading/worker_sequential.py new file mode 100644 index 00000000..e3d436cb --- /dev/null +++ b/d123/common/multithreading/worker_sequential.py @@ -0,0 +1,46 @@ +import logging +from concurrent.futures import Future +from typing import Any, Iterable, List + +from tqdm import tqdm + +from d123.common.multithreading.worker_pool import ( + Task, + WorkerPool, + WorkerResources, + get_max_size_of_arguments, +) + +logger = logging.getLogger(__name__) + + +class Sequential(WorkerPool): + """ + This function does execute all functions sequentially. + """ + + def __init__(self) -> None: + """ + Initialize simple sequential worker. + """ + super().__init__(WorkerResources(number_of_nodes=1, number_of_cpus_per_node=1, number_of_gpus_per_node=0)) + + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """Inherited, see superclass.""" + if task.num_cpus not in [None, 1]: + raise ValueError(f"Expected num_cpus to be 1 or unset for Sequential worker, got {task.num_cpus}") + output = [ + task.fn(*args) + for args in tqdm( + zip(*item_lists), + leave=False, + total=get_max_size_of_arguments(*item_lists), + desc="Sequential", + disable=not verbose, + ) + ] + return output + + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """Inherited, see superclass.""" + raise NotImplementedError diff --git a/d123/common/multithreading/worker_utils.py b/d123/common/multithreading/worker_utils.py new file mode 100644 index 00000000..fbe0a753 --- /dev/null +++ b/d123/common/multithreading/worker_utils.py @@ -0,0 +1,36 @@ +from typing import Any, Callable, List, Optional + +import numpy as np +from psutil import cpu_count + +from d123.common.multithreading.worker_pool import Task, WorkerPool + + +def chunk_list(input_list: List[Any], num_chunks: Optional[int] = None) -> List[List[Any]]: + """ + Chunks a list to equal sized lists. The size of the last list might be truncated. + :param input_list: List to be chunked. + :param num_chunks: Number of chunks, equals to the number of cores if set to None. + :return: List of equal sized lists. + """ + num_chunks = num_chunks if num_chunks else cpu_count(logical=True) + chunks = np.array_split(input_list, num_chunks) # type: ignore + return [chunk.tolist() for chunk in chunks if len(chunk) != 0] + + +def worker_map(worker: WorkerPool, fn: Callable[..., List[Any]], input_objects: List[Any]) -> List[Any]: + """ + Map a list of objects through a worker. + :param worker: Worker pool to use for parallelization. + :param fn: Function to use when mapping. + :param input_objects: List of objects to map. + :return: List of mapped objects. + """ + if worker.number_of_threads == 0: + return fn(input_objects) + + object_chunks = chunk_list(input_objects, worker.number_of_threads) + scattered_objects = worker.map(Task(fn=fn), object_chunks) + output_objects = [result for results in scattered_objects for result in results] + + return output_objects diff --git a/d123/common/utils/enums.py b/d123/common/utils/enums.py index e09103e1..33300c00 100644 --- a/d123/common/utils/enums.py +++ b/d123/common/utils/enums.py @@ -18,7 +18,7 @@ def serialize(self, lower: bool = True) -> str: return self.name.lower() if lower else self.name @classmethod - def deserialize(cls, key: str) -> type[SerialIntEnum]: + def deserialize(cls, key: str) -> SerialIntEnum: """Deserialize the type when loading from a string.""" # Allow for lower/upper case letters during deserialize return cls.__members__[key.upper()] if key.islower() else cls.__members__[key] diff --git a/d123/common/visualization/color/color.py b/d123/common/visualization/color/color.py index b151037a..5ba3c866 100644 --- a/d123/common/visualization/color/color.py +++ b/d123/common/visualization/color/color.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from dataclasses import dataclass from typing import Dict, Tuple @@ -9,6 +11,11 @@ class Color: hex: str + @classmethod + def from_rgb(cls, rgb: Tuple[int, int, int]) -> Color: + r, g, b = rgb + return cls(f"#{r:02x}{g:02x}{b:02x}") + @property def rgb(self) -> Tuple[int, int, int]: return ImageColor.getcolor(self.hex, "RGB") @@ -25,6 +32,16 @@ def rgb_norm(self) -> Tuple[float, float, float]: def rgba_norm(self) -> Tuple[float, float, float]: return tuple([c / 255 for c in self.rgba]) + def set_brightness(self, factor: float) -> Color: + r, g, b = self.rgb + return Color.from_rgb( + ( + max(min(int(r * factor), 255), 0), + max(min(int(g * factor), 255), 0), + max(min(int(b * factor), 255), 0), + ) + ) + def __str__(self) -> str: return self.hex diff --git a/d123/common/visualization/color/config.py b/d123/common/visualization/color/config.py index 69eb5f3b..6d610844 100644 --- a/d123/common/visualization/color/config.py +++ b/d123/common/visualization/color/config.py @@ -9,7 +9,7 @@ class PlotConfig: fill_color: Color = field(default_factory=lambda: BLACK) fill_color_alpha: float = 1.0 line_color: Color = BLACK - line_color_alpha: float = 0.0 + line_color_alpha: float = 1.0 line_width: float = 1.0 line_style: str = "-" marker_style: str = "o" @@ -18,5 +18,3 @@ class PlotConfig: zorder: int = 0 smoothing_radius: Optional[float] = None - - shadow: bool = False diff --git a/d123/common/visualization/color/default.py b/d123/common/visualization/color/default.py index 83921339..076a1c81 100644 --- a/d123/common/visualization/color/default.py +++ b/d123/common/visualization/color/default.py @@ -13,12 +13,12 @@ Color, ) from d123.common.visualization.color.config import PlotConfig -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer HEADING_MARKER_STYLE: str = "^" # "^": triangle, "-": line -MAP_SURFACE_CONFIG: Dict[MapSurfaceType, PlotConfig] = { - MapSurfaceType.LANE: PlotConfig( +MAP_SURFACE_CONFIG: Dict[MapLayer, PlotConfig] = { + MapLayer.LANE: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -27,7 +27,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.LANE_GROUP: PlotConfig( + MapLayer.LANE_GROUP: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -36,7 +36,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.INTERSECTION: PlotConfig( + MapLayer.INTERSECTION: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -45,7 +45,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.CROSSWALK: PlotConfig( + MapLayer.CROSSWALK: PlotConfig( fill_color=Color("#c69fbb"), fill_color_alpha=1.0, line_color=Color("#c69fbb"), @@ -54,7 +54,7 @@ line_style="-", zorder=2, ), - MapSurfaceType.WALKWAY: PlotConfig( + MapLayer.WALKWAY: PlotConfig( fill_color=Color("#d4d19e"), fill_color_alpha=1.0, line_color=Color("#d4d19e"), @@ -63,7 +63,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.CARPARK: PlotConfig( + MapLayer.CARPARK: PlotConfig( fill_color=Color("#b9d3b4"), fill_color_alpha=1.0, line_color=Color("#b9d3b4"), @@ -72,7 +72,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.GENERIC_DRIVABLE: PlotConfig( + MapLayer.GENERIC_DRIVABLE: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -157,6 +157,16 @@ marker_style=None, zorder=2, ), + DetectionType.SIGN: PlotConfig( + fill_color=NEW_TAB_10[8], + fill_color_alpha=1.0, + line_color=BLACK, + line_color_alpha=1.0, + line_width=1.0, + line_style="-", + marker_style=None, + zorder=2, + ), } EGO_VEHICLE_CONFIG: PlotConfig = PlotConfig( diff --git a/d123/common/visualization/matplotlib/camera copy.py b/d123/common/visualization/matplotlib/camera copy.py new file mode 100644 index 00000000..ea758366 --- /dev/null +++ b/d123/common/visualization/matplotlib/camera copy.py @@ -0,0 +1,330 @@ +# from typing import List, Optional, Tuple + +from typing import List, Optional, Tuple + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import numpy.typing as npt + +# from PIL import ImageColor +from pyquaternion import Quaternion + +from d123.common.datatypes.detection.detection import BoxDetectionSE3, BoxDetectionWrapper +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import Camera +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 +from d123.common.geometry.base import StateSE3 +from d123.common.geometry.bounding_box.bounding_box_index import BoundingBoxSE3Index, Corners3DIndex +from d123.common.geometry.transform.se3 import convert_absolute_to_relative_se3_array, get_rotation_matrix +from d123.common.visualization.color.default import BOX_DETECTION_CONFIG + +# from navsim.common.dataclasses import Annotations, Camera, Lidar +# from navsim.common.enums import BoundingBoxIndex, LidarIndex +# from navsim.planning.scenario_builder.navsim_scenario_utils import tracked_object_types +# from navsim.visualization.config import AGENT_CONFIG +# from navsim.visualization.lidar import filter_lidar_pc, get_lidar_pc_color + + +def add_camera_ax(ax: plt.Axes, camera: Camera) -> plt.Axes: + """ + Adds camera image to matplotlib ax object + :param ax: matplotlib ax object + :param camera: navsim camera dataclass + :return: ax object with image + """ + ax.imshow(camera.image) + return ax + + +# FIXME: +# def add_lidar_to_camera_ax(ax: plt.Axes, camera: Camera, lidar: Lidar) -> plt.Axes: +# """ +# Adds camera image with lidar point cloud on matplotlib ax object +# :param ax: matplotlib ax object +# :param camera: navsim camera dataclass +# :param lidar: navsim lidar dataclass +# :return: ax object with image +# """ + +# image, lidar_pc = camera.image.copy(), lidar.lidar_pc.copy() +# image_height, image_width = image.shape[:2] + +# lidar_pc = filter_lidar_pc(lidar_pc) +# lidar_pc_colors = np.array(get_lidar_pc_color(lidar_pc)) + +# pc_in_cam, pc_in_fov_mask = _transform_pcs_to_images( +# lidar_pc, +# camera.sensor2lidar_rotation, +# camera.sensor2lidar_translation, +# camera.intrinsics, +# img_shape=(image_height, image_width), +# ) + +# for (x, y), color in zip(pc_in_cam[pc_in_fov_mask], lidar_pc_colors[pc_in_fov_mask]): +# color = (int(color[0]), int(color[1]), int(color[2])) +# cv2.circle(image, (int(x), int(y)), 5, color, -1) + +# ax.imshow(image) +# return ax + + +def add_box_detections_to_camera_ax( + ax: plt.Axes, + camera: Camera, + box_detections: BoxDetectionWrapper, + ego_state_se3: EgoStateSE3, +) -> plt.Axes: + + box_detection_array = np.zeros((len(box_detections.box_detections), len(BoundingBoxSE3Index)), dtype=np.float64) + detection_types = np.array( + [detection.metadata.detection_type for detection in box_detections.box_detections], dtype=object + ) + for idx, box_detection in enumerate(box_detections.box_detections): + assert isinstance( + box_detection, BoxDetectionSE3 + ), f"Box detection must be of type BoxDetectionSE3, got {type(box_detection)}" + box_detection_array[idx] = box_detection.bounding_box_se3.array + + box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] = convert_absolute_to_relative_se3_array( + ego_state_se3.rear_axle_se3, box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] + ) + detection_positions, detection_extents, detection_yaws = _transform_annotations_to_camera( + box_detection_array, camera.extrinsic + ) + + corners_norm = np.stack(np.unravel_index(np.arange(len(Corners3DIndex)), [2] * 3), axis=1) + corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] + corners_norm = corners_norm - np.array([0.5, 0.5, 0.5]) + corners = detection_extents.reshape([-1, 1, 3]) * corners_norm.reshape([1, 8, 3]) + + corners = _rotation_3d_in_axis(corners, detection_yaws, axis=1) + corners += detection_positions.reshape(-1, 1, 3) + + # Then draw project corners to image. + box_corners, corners_pc_in_fov = _transform_points_to_image(corners.reshape(-1, 3), camera.metadata.intrinsic) + box_corners = box_corners.reshape(-1, 8, 2) + corners_pc_in_fov = corners_pc_in_fov.reshape(-1, 8) + valid_corners = corners_pc_in_fov.any(-1) + + box_corners, detection_types = box_corners[valid_corners], detection_types[valid_corners] + image = _plot_rect_3d_on_img(camera.image.copy(), box_corners, detection_types) + + ax.imshow(image) + return ax + + +def _transform_annotations_to_camera( + boxes: npt.NDArray[np.float32], extrinsic: npt.NDArray[np.float64] +) -> npt.NDArray[np.float32]: + """ + Helper function to transform bounding boxes into camera frame + TODO: Refactor + :param boxes: array representation of bounding boxes + :param sensor2lidar_rotation: camera rotation + :param sensor2lidar_translation: camera translation + :return: bounding boxes in camera coordinates + """ + sensor2ego_rotation = extrinsic[:3, :3] + sensor2ego_translation = extrinsic[:3, 3] + + locs, rots = ( + boxes[:, BoundingBoxSE3Index.XYZ], + boxes[:, BoundingBoxSE3Index.YAW], + ) + dims_cam = boxes[ + :, [BoundingBoxSE3Index.LENGTH, BoundingBoxSE3Index.HEIGHT, BoundingBoxSE3Index.WIDTH] + ] # l, w, h -> l, h, w + + rots_cam = np.zeros_like(rots) + for idx, state_se3_array in enumerate(boxes[:, BoundingBoxSE3Index.STATE_SE3]): + rot = Quaternion(matrix=get_rotation_matrix(StateSE3.from_array(state_se3_array))) + rot = Quaternion(matrix=sensor2ego_rotation).inverse * rot + rots_cam[idx] = -rot.yaw_pitch_roll[0] + + lidar2cam_r = np.linalg.inv(sensor2ego_rotation) + lidar2cam_t = sensor2ego_translation @ lidar2cam_r.T + lidar2cam_rt = np.eye(4) + lidar2cam_rt[:3, :3] = lidar2cam_r.T + lidar2cam_rt[3, :3] = -lidar2cam_t + + locs_cam = np.concatenate([locs, np.ones_like(locs)[:, :1]], -1) # -1, 4 + locs_cam = lidar2cam_rt.T @ locs_cam.T + locs_cam = locs_cam.T + locs_cam = locs_cam[:, :-1] + return locs_cam, dims_cam, rots_cam + + +def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np.float32], axis: int = 0): + """ + Rotate 3D points by angles according to axis. + TODO: Refactor + :param points: array of points + :param angles: array of angles + :param axis: axis to perform rotation, defaults to 0 + :raises value: _description_ + :raises ValueError: if axis invalid + :return: rotated points + """ + rot_sin = np.sin(angles) + rot_cos = np.cos(angles) + ones = np.ones_like(rot_cos) + zeros = np.zeros_like(rot_cos) + if axis == 1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, zeros, -rot_sin]), + np.stack([zeros, ones, zeros]), + np.stack([rot_sin, zeros, rot_cos]), + ] + ) + elif axis == 2 or axis == -1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, -rot_sin, zeros]), + np.stack([rot_sin, rot_cos, zeros]), + np.stack([zeros, zeros, ones]), + ] + ) + elif axis == 0: + rot_mat_T = np.stack( + [ + np.stack([zeros, rot_cos, -rot_sin]), + np.stack([zeros, rot_sin, rot_cos]), + np.stack([ones, zeros, zeros]), + ] + ) + else: + raise ValueError(f"axis should in range [0, 1, 2], got {axis}") + return np.einsum("aij,jka->aik", points, rot_mat_T) + + +def _plot_rect_3d_on_img( + image: npt.NDArray[np.float32], + box_corners: npt.NDArray[np.float32], + detection_types: List[DetectionType], + thickness: int = 1, +) -> npt.NDArray[np.uint8]: + """ + Plot the boundary lines of 3D rectangular on 2D images. + TODO: refactor + :param image: The numpy array of image. + :param box_corners: Coordinates of the corners of 3D, shape of [N, 8, 2]. + :param box_labels: labels of boxes for coloring + :param thickness: pixel width of liens, defaults to 3 + :return: image with 3D bounding boxes + """ + line_indices = ( + (0, 1), + (0, 3), + (0, 4), + (1, 2), + (1, 5), + (3, 2), + (3, 7), + (4, 5), + (4, 7), + (2, 6), + (5, 6), + (6, 7), + ) + for i in range(len(box_corners)): + color = BOX_DETECTION_CONFIG[detection_types[i]].fill_color.rgb + corners = box_corners[i].astype(np.int64) + for start, end in line_indices: + cv2.line( + image, + (corners[start, 0], corners[start, 1]), + (corners[end, 0], corners[end, 1]), + color, + thickness, + cv2.LINE_AA, + ) + return image.astype(np.uint8) + + +def _transform_points_to_image( + points: npt.NDArray[np.float32], + intrinsic: npt.NDArray[np.float32], + image_shape: Optional[Tuple[int, int]] = None, + eps: float = 1e-3, +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: + """ + Transforms points in camera frame to image pixel coordinates + TODO: refactor + :param points: points in camera frame + :param intrinsic: camera intrinsics + :param image_shape: shape of image in pixel + :param eps: lower threshold of points, defaults to 1e-3 + :return: points in pixel coordinates, mask of values in frame + """ + points = points[:, :3] + + viewpad = np.eye(4) + viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic + + pc_img = np.concatenate([points, np.ones_like(points)[:, :1]], -1) + pc_img = viewpad @ pc_img.T + pc_img = pc_img.T + + cur_pc_in_fov = pc_img[:, 2] > eps + pc_img = pc_img[..., 0:2] / np.maximum(pc_img[..., 2:3], np.ones_like(pc_img[..., 2:3]) * eps) + if image_shape is not None: + img_h, img_w = image_shape + cur_pc_in_fov = ( + cur_pc_in_fov + & (pc_img[:, 0] < (img_w - 1)) + & (pc_img[:, 0] > 0) + & (pc_img[:, 1] < (img_h - 1)) + & (pc_img[:, 1] > 0) + ) + return pc_img, cur_pc_in_fov + + +# def _transform_pcs_to_images( +# lidar_pc: npt.NDArray[np.float32], +# sensor2lidar_rotation: npt.NDArray[np.float32], +# sensor2lidar_translation: npt.NDArray[np.float32], +# intrinsic: npt.NDArray[np.float32], +# img_shape: Optional[Tuple[int, int]] = None, +# eps: float = 1e-3, +# ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: +# """ +# Transforms points in camera frame to image pixel coordinates +# TODO: refactor +# :param lidar_pc: lidar point cloud +# :param sensor2lidar_rotation: camera rotation +# :param sensor2lidar_translation: camera translation +# :param intrinsic: camera intrinsics +# :param img_shape: image shape in pixels, defaults to None +# :param eps: threshold for lidar pc height, defaults to 1e-3 +# :return: lidar pc in pixel coordinates, mask of values in frame +# """ +# pc_xyz = lidar_pc[LidarIndex.POSITION, :].T + +# lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) +# lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T +# lidar2cam_rt = np.eye(4) +# lidar2cam_rt[:3, :3] = lidar2cam_r.T +# lidar2cam_rt[3, :3] = -lidar2cam_t + +# viewpad = np.eye(4) +# viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic +# lidar2img_rt = viewpad @ lidar2cam_rt.T + +# cur_pc_xyz = np.concatenate([pc_xyz, np.ones_like(pc_xyz)[:, :1]], -1) +# cur_pc_cam = lidar2img_rt @ cur_pc_xyz.T +# cur_pc_cam = cur_pc_cam.T +# cur_pc_in_fov = cur_pc_cam[:, 2] > eps +# cur_pc_cam = cur_pc_cam[..., 0:2] / np.maximum(cur_pc_cam[..., 2:3], np.ones_like(cur_pc_cam[..., 2:3]) * eps) + +# if img_shape is not None: +# img_h, img_w = img_shape +# cur_pc_in_fov = ( +# cur_pc_in_fov +# & (cur_pc_cam[:, 0] < (img_w - 1)) +# & (cur_pc_cam[:, 0] > 0) +# & (cur_pc_cam[:, 1] < (img_h - 1)) +# & (cur_pc_cam[:, 1] > 0) +# ) +# return cur_pc_cam, cur_pc_in_fov diff --git a/d123/common/visualization/matplotlib/camera.py b/d123/common/visualization/matplotlib/camera.py new file mode 100644 index 00000000..52d873ae --- /dev/null +++ b/d123/common/visualization/matplotlib/camera.py @@ -0,0 +1,346 @@ +# from typing import List, Optional, Tuple + +from typing import List, Optional, Tuple + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import numpy.typing as npt + +# from PIL import ImageColor +from pyquaternion import Quaternion + +from d123.common.datatypes.detection.detection import BoxDetectionSE3, BoxDetectionWrapper +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import Camera +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 +from d123.common.geometry.bounding_box.bounding_box_index import BoundingBoxSE3Index, Corners3DIndex +from d123.common.geometry.transform.se3 import convert_absolute_to_relative_se3_array +from d123.common.visualization.color.default import BOX_DETECTION_CONFIG + +# from navsim.common.dataclasses import Annotations, Camera, Lidar +# from navsim.common.enums import BoundingBoxIndex, LidarIndex +# from navsim.planning.scenario_builder.navsim_scenario_utils import tracked_object_types +# from navsim.visualization.config import AGENT_CONFIG +# from navsim.visualization.lidar import filter_lidar_pc, get_lidar_pc_color + + +def add_camera_ax(ax: plt.Axes, camera: Camera) -> plt.Axes: + """ + Adds camera image to matplotlib ax object + :param ax: matplotlib ax object + :param camera: navsim camera dataclass + :return: ax object with image + """ + ax.imshow(camera.image) + return ax + + +# FIXME: +# def add_lidar_to_camera_ax(ax: plt.Axes, camera: Camera, lidar: Lidar) -> plt.Axes: +# """ +# Adds camera image with lidar point cloud on matplotlib ax object +# :param ax: matplotlib ax object +# :param camera: navsim camera dataclass +# :param lidar: navsim lidar dataclass +# :return: ax object with image +# """ + +# image, lidar_pc = camera.image.copy(), lidar.lidar_pc.copy() +# image_height, image_width = image.shape[:2] + +# lidar_pc = filter_lidar_pc(lidar_pc) +# lidar_pc_colors = np.array(get_lidar_pc_color(lidar_pc)) + +# pc_in_cam, pc_in_fov_mask = _transform_pcs_to_images( +# lidar_pc, +# camera.sensor2lidar_rotation, +# camera.sensor2lidar_translation, +# camera.intrinsics, +# img_shape=(image_height, image_width), +# ) + +# for (x, y), color in zip(pc_in_cam[pc_in_fov_mask], lidar_pc_colors[pc_in_fov_mask]): +# color = (int(color[0]), int(color[1]), int(color[2])) +# cv2.circle(image, (int(x), int(y)), 5, color, -1) + +# ax.imshow(image) +# return ax + + +def add_box_detections_to_camera_ax( + ax: plt.Axes, + camera: Camera, + box_detections: BoxDetectionWrapper, + ego_state_se3: EgoStateSE3, +) -> plt.Axes: + + # box_labels = annotations.names + # boxes = _transform_annotations_to_camera( + # annotations.boxes, + # camera.sensor2lidar_rotation, + # camera.sensor2lidar_translation, + # ) + # box_positions, box_dimensions, box_heading = ( + # boxes[:, BoundingBoxIndex.POSITION], + # boxes[:, BoundingBoxIndex.DIMENSION], + # boxes[:, BoundingBoxIndex.HEADING], + # ) + + box_detection_array = np.zeros((len(box_detections.box_detections), len(BoundingBoxSE3Index)), dtype=np.float64) + detection_types = np.array( + [detection.metadata.detection_type for detection in box_detections.box_detections], dtype=object + ) + for idx, box_detection in enumerate(box_detections.box_detections): + assert isinstance( + box_detection, BoxDetectionSE3 + ), f"Box detection must be of type BoxDetectionSE3, got {type(box_detection)}" + box_detection_array[idx] = box_detection.bounding_box_se3.array + + # FIXME + box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] = convert_absolute_to_relative_se3_array( + ego_state_se3.rear_axle_se3, box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] + ) + # box_detection_array[..., BoundingBoxSE3Index.XYZ] -= ego_state_se3.rear_axle_se3.point_3d.array + detection_positions, detection_extents, detection_yaws = _transform_annotations_to_camera( + box_detection_array, camera.extrinsic + ) + + corners_norm = np.stack(np.unravel_index(np.arange(len(Corners3DIndex)), [2] * 3), axis=1) + corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] + corners_norm = corners_norm - np.array([0.5, 0.5, 0.5]) + corners = detection_extents.reshape([-1, 1, 3]) * corners_norm.reshape([1, 8, 3]) + + corners = _rotation_3d_in_axis(corners, detection_yaws, axis=1) + corners += detection_positions.reshape(-1, 1, 3) + + # Then draw project corners to image. + box_corners, corners_pc_in_fov = _transform_points_to_image(corners.reshape(-1, 3), camera.metadata.intrinsic) + box_corners = box_corners.reshape(-1, 8, 2) + corners_pc_in_fov = corners_pc_in_fov.reshape(-1, 8) + valid_corners = corners_pc_in_fov.any(-1) + + box_corners, detection_types = box_corners[valid_corners], detection_types[valid_corners] + image = _plot_rect_3d_on_img(camera.image.copy(), box_corners, detection_types) + + ax.imshow(image) + return ax + + +def _transform_annotations_to_camera( + boxes: npt.NDArray[np.float32], + # sensor2lidar_rotation: npt.NDArray[np.float32], + # sensor2lidar_translation: npt.NDArray[np.float32], + extrinsic: npt.NDArray[np.float64], +) -> npt.NDArray[np.float32]: + """ + Helper function to transform bounding boxes into camera frame + TODO: Refactor + :param boxes: array representation of bounding boxes + :param sensor2lidar_rotation: camera rotation + :param sensor2lidar_translation: camera translation + :return: bounding boxes in camera coordinates + """ + sensor2lidar_rotation = extrinsic[:3, :3] + sensor2lidar_translation = extrinsic[:3, 3] + + locs, rots = ( + boxes[:, BoundingBoxSE3Index.XYZ], + boxes[:, BoundingBoxSE3Index.YAW], + ) + dims_cam = boxes[ + :, [BoundingBoxSE3Index.LENGTH, BoundingBoxSE3Index.HEIGHT, BoundingBoxSE3Index.WIDTH] + ] # l, w, h -> l, h, w + + rots_cam = np.zeros_like(rots) + for idx, rot in enumerate(rots): + rot = Quaternion(axis=[0, 0, 1], radians=rot) + rot = Quaternion(matrix=sensor2lidar_rotation).inverse * rot + rots_cam[idx] = -rot.yaw_pitch_roll[0] + + lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) + lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T + lidar2cam_rt = np.eye(4) + lidar2cam_rt[:3, :3] = lidar2cam_r.T + lidar2cam_rt[3, :3] = -lidar2cam_t + + locs_cam = np.concatenate([locs, np.ones_like(locs)[:, :1]], -1) # -1, 4 + locs_cam = lidar2cam_rt.T @ locs_cam.T + locs_cam = locs_cam.T + locs_cam = locs_cam[:, :-1] + return locs_cam, dims_cam, rots_cam + + +def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np.float32], axis: int = 0): + """ + Rotate 3D points by angles according to axis. + TODO: Refactor + :param points: array of points + :param angles: array of angles + :param axis: axis to perform rotation, defaults to 0 + :raises value: _description_ + :raises ValueError: if axis invalid + :return: rotated points + """ + rot_sin = np.sin(angles) + rot_cos = np.cos(angles) + ones = np.ones_like(rot_cos) + zeros = np.zeros_like(rot_cos) + if axis == 1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, zeros, -rot_sin]), + np.stack([zeros, ones, zeros]), + np.stack([rot_sin, zeros, rot_cos]), + ] + ) + elif axis == 2 or axis == -1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, -rot_sin, zeros]), + np.stack([rot_sin, rot_cos, zeros]), + np.stack([zeros, zeros, ones]), + ] + ) + elif axis == 0: + rot_mat_T = np.stack( + [ + np.stack([zeros, rot_cos, -rot_sin]), + np.stack([zeros, rot_sin, rot_cos]), + np.stack([ones, zeros, zeros]), + ] + ) + else: + raise ValueError(f"axis should in range [0, 1, 2], got {axis}") + return np.einsum("aij,jka->aik", points, rot_mat_T) + + +def _plot_rect_3d_on_img( + image: npt.NDArray[np.float32], + box_corners: npt.NDArray[np.float32], + detection_types: List[DetectionType], + thickness: int = 3, +) -> npt.NDArray[np.uint8]: + """ + Plot the boundary lines of 3D rectangular on 2D images. + TODO: refactor + :param image: The numpy array of image. + :param box_corners: Coordinates of the corners of 3D, shape of [N, 8, 2]. + :param box_labels: labels of boxes for coloring + :param thickness: pixel width of liens, defaults to 3 + :return: image with 3D bounding boxes + """ + line_indices = ( + (0, 1), + (0, 3), + (0, 4), + (1, 2), + (1, 5), + (3, 2), + (3, 7), + (4, 5), + (4, 7), + (2, 6), + (5, 6), + (6, 7), + ) + for i in range(len(box_corners)): + color = BOX_DETECTION_CONFIG[detection_types[i]].fill_color.rgb + corners = box_corners[i].astype(np.int64) + for start, end in line_indices: + cv2.line( + image, + (corners[start, 0], corners[start, 1]), + (corners[end, 0], corners[end, 1]), + color, + thickness, + cv2.LINE_AA, + ) + return image.astype(np.uint8) + + +def _transform_points_to_image( + points: npt.NDArray[np.float32], + intrinsic: npt.NDArray[np.float32], + image_shape: Optional[Tuple[int, int]] = None, + eps: float = 1e-3, +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: + """ + Transforms points in camera frame to image pixel coordinates + TODO: refactor + :param points: points in camera frame + :param intrinsic: camera intrinsics + :param image_shape: shape of image in pixel + :param eps: lower threshold of points, defaults to 1e-3 + :return: points in pixel coordinates, mask of values in frame + """ + points = points[:, :3] + + viewpad = np.eye(4) + viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic + + pc_img = np.concatenate([points, np.ones_like(points)[:, :1]], -1) + pc_img = viewpad @ pc_img.T + pc_img = pc_img.T + + cur_pc_in_fov = pc_img[:, 2] > eps + pc_img = pc_img[..., 0:2] / np.maximum(pc_img[..., 2:3], np.ones_like(pc_img[..., 2:3]) * eps) + if image_shape is not None: + img_h, img_w = image_shape + cur_pc_in_fov = ( + cur_pc_in_fov + & (pc_img[:, 0] < (img_w - 1)) + & (pc_img[:, 0] > 0) + & (pc_img[:, 1] < (img_h - 1)) + & (pc_img[:, 1] > 0) + ) + return pc_img, cur_pc_in_fov + + +# def _transform_pcs_to_images( +# lidar_pc: npt.NDArray[np.float32], +# sensor2lidar_rotation: npt.NDArray[np.float32], +# sensor2lidar_translation: npt.NDArray[np.float32], +# intrinsic: npt.NDArray[np.float32], +# img_shape: Optional[Tuple[int, int]] = None, +# eps: float = 1e-3, +# ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: +# """ +# Transforms points in camera frame to image pixel coordinates +# TODO: refactor +# :param lidar_pc: lidar point cloud +# :param sensor2lidar_rotation: camera rotation +# :param sensor2lidar_translation: camera translation +# :param intrinsic: camera intrinsics +# :param img_shape: image shape in pixels, defaults to None +# :param eps: threshold for lidar pc height, defaults to 1e-3 +# :return: lidar pc in pixel coordinates, mask of values in frame +# """ +# pc_xyz = lidar_pc[LidarIndex.POSITION, :].T + +# lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) +# lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T +# lidar2cam_rt = np.eye(4) +# lidar2cam_rt[:3, :3] = lidar2cam_r.T +# lidar2cam_rt[3, :3] = -lidar2cam_t + +# viewpad = np.eye(4) +# viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic +# lidar2img_rt = viewpad @ lidar2cam_rt.T + +# cur_pc_xyz = np.concatenate([pc_xyz, np.ones_like(pc_xyz)[:, :1]], -1) +# cur_pc_cam = lidar2img_rt @ cur_pc_xyz.T +# cur_pc_cam = cur_pc_cam.T +# cur_pc_in_fov = cur_pc_cam[:, 2] > eps +# cur_pc_cam = cur_pc_cam[..., 0:2] / np.maximum(cur_pc_cam[..., 2:3], np.ones_like(cur_pc_cam[..., 2:3]) * eps) + +# if img_shape is not None: +# img_h, img_w = img_shape +# cur_pc_in_fov = ( +# cur_pc_in_fov +# & (cur_pc_cam[:, 0] < (img_w - 1)) +# & (cur_pc_cam[:, 0] > 0) +# & (cur_pc_cam[:, 1] < (img_h - 1)) +# & (cur_pc_cam[:, 1] > 0) +# ) +# return cur_pc_cam, cur_pc_in_fov diff --git a/d123/common/visualization/matplotlib/lidar.py b/d123/common/visualization/matplotlib/lidar.py new file mode 100644 index 00000000..29ffeea8 --- /dev/null +++ b/d123/common/visualization/matplotlib/lidar.py @@ -0,0 +1,70 @@ +# TODO: implement +# from typing import Any, List + +# import matplotlib +# import matplotlib.pyplot as plt +# import numpy as np +# import numpy.typing as npt + +# from navsim.common.enums import LidarIndex +# from navsim.visualization.config import LIDAR_CONFIG + + +# def filter_lidar_pc(lidar_pc: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: +# """ +# Filter lidar point cloud according to global configuration +# :param lidar_pc: numpy array of shape (6,n) +# :return: filtered point cloud +# """ + +# pc = lidar_pc.T +# mask = ( +# np.ones((len(pc)), dtype=bool) +# & (pc[:, LidarIndex.X] > LIDAR_CONFIG["x_lim"][0]) +# & (pc[:, LidarIndex.X] < LIDAR_CONFIG["x_lim"][1]) +# & (pc[:, LidarIndex.Y] > LIDAR_CONFIG["y_lim"][0]) +# & (pc[:, LidarIndex.Y] < LIDAR_CONFIG["y_lim"][1]) +# & (pc[:, LidarIndex.Z] > LIDAR_CONFIG["z_lim"][0]) +# & (pc[:, LidarIndex.Z] < LIDAR_CONFIG["z_lim"][1]) +# ) +# pc = pc[mask] +# return pc.T + + +# def get_lidar_pc_color(lidar_pc: npt.NDArray[np.float32], as_hex: bool = False) -> List[Any]: +# """ +# Compute color map of lidar point cloud according to global configuration +# :param lidar_pc: numpy array of shape (6,n) +# :param as_hex: whether to return hex values, defaults to False +# :return: list of RGB or hex values +# """ + +# pc = lidar_pc.T +# if LIDAR_CONFIG["color_element"] == "none": +# colors_rgb = np.zeros((len(pc), 3), dtype=np.uin8) + +# else: +# if LIDAR_CONFIG["color_element"] == "distance": +# color_intensities = np.linalg.norm(pc[:, LidarIndex.POSITION], axis=-1) +# else: +# color_element_map = { +# "x": LidarIndex.X, +# "y": LidarIndex.Y, +# "z": LidarIndex.Z, +# "intensity": LidarIndex.INTENSITY, +# "ring": LidarIndex.RING, +# "id": LidarIndex.ID, +# } +# color_intensities = pc[:, color_element_map[LIDAR_CONFIG["color_element"]]] + +# min, max = color_intensities.min(), color_intensities.max() +# norm_intensities = [(value - min) / (max - min) for value in color_intensities] +# colormap = plt.get_cmap("viridis") +# colors_rgb = np.array([colormap(value) for value in norm_intensities]) +# colors_rgb = (colors_rgb[:, :3] * 255).astype(np.uint8) + +# assert len(colors_rgb) == len(pc) +# if as_hex: +# return [matplotlib.colors.to_hex(tuple(c / 255.0 for c in rgb)) for rgb in colors_rgb] + +# return [tuple(value) for value in colors_rgb] diff --git a/d123/common/visualization/matplotlib/observation.py b/d123/common/visualization/matplotlib/observation.py index 26cb65b6..7e225cc7 100644 --- a/d123/common/visualization/matplotlib/observation.py +++ b/d123/common/visualization/matplotlib/observation.py @@ -5,6 +5,7 @@ import shapely.geometry as geom from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper +from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2, EgoStateSE3 from d123.common.geometry.base import Point2D from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE2, BoundingBoxSE3 @@ -26,7 +27,8 @@ ) from d123.dataset.maps.abstract_map import AbstractMap from d123.dataset.maps.abstract_map_objects import AbstractLane -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer +from d123.dataset.scene.abstract_scene import AbstractScene def add_default_map_on_ax( @@ -36,14 +38,14 @@ def add_default_map_on_ax( radius: float, route_lane_group_ids: Optional[List[int]] = None, ) -> None: - layers: List[MapSurfaceType] = [ - MapSurfaceType.LANE, - MapSurfaceType.LANE_GROUP, - MapSurfaceType.GENERIC_DRIVABLE, - MapSurfaceType.CARPARK, - MapSurfaceType.CROSSWALK, - MapSurfaceType.INTERSECTION, - MapSurfaceType.WALKWAY, + layers: List[MapLayer] = [ + MapLayer.LANE, + MapLayer.LANE_GROUP, + MapLayer.GENERIC_DRIVABLE, + MapLayer.CARPARK, + MapLayer.CROSSWALK, + MapLayer.INTERSECTION, + MapLayer.WALKWAY, ] x_min, x_max = point_2d.x - radius, point_2d.x + radius y_min, y_max = point_2d.y - radius, point_2d.y + radius @@ -53,20 +55,20 @@ def add_default_map_on_ax( for layer, map_objects in map_objects_dict.items(): for map_object in map_objects: try: - if layer in [MapSurfaceType.LANE_GROUP]: + if layer in [MapLayer.LANE_GROUP]: if route_lane_group_ids is not None and int(map_object.id) in route_lane_group_ids: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, ROUTE_CONFIG) else: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer]) if layer in [ - MapSurfaceType.GENERIC_DRIVABLE, - MapSurfaceType.CARPARK, - MapSurfaceType.CROSSWALK, - MapSurfaceType.INTERSECTION, - MapSurfaceType.WALKWAY, + MapLayer.GENERIC_DRIVABLE, + MapLayer.CARPARK, + MapLayer.CROSSWALK, + MapLayer.INTERSECTION, + MapLayer.WALKWAY, ]: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer]) - if layer in [MapSurfaceType.LANE]: + if layer in [MapLayer.LANE]: map_object: AbstractLane add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG) except Exception: @@ -87,6 +89,39 @@ def add_box_detections_to_ax(ax: plt.Axes, box_detections: BoxDetectionWrapper) add_bounding_box_to_ax(ax, box_detection.bounding_box, plot_config) +def add_box_future_detections_to_ax(ax: plt.Axes, scene: AbstractScene, iteration: int) -> None: + + # TODO: Refactor this function + initial_agents = scene.get_box_detections_at_iteration(iteration) + agents_poses = { + agent.metadata.track_token: [agent.center_se3] + for agent in initial_agents + if agent.metadata.detection_type == DetectionType.VEHICLE + } + frequency = 1 + for iteration in range(iteration + frequency, scene.get_number_of_iterations(), frequency): + agents = scene.get_box_detections_at_iteration(iteration) + for agent in agents: + if agent.metadata.track_token in agents_poses: + agents_poses[agent.metadata.track_token].append(agent.center_se3) + + for track_token, poses in agents_poses.items(): + if len(poses) < 2: + continue + poses = np.array([pose.point_2d.array for pose in poses]) + num_poses = poses.shape[0] + alphas = 1 - np.linspace(0.2, 1.0, num_poses) # Start low, end high + for i in range(num_poses - 1): + ax.plot( + poses[i : i + 2, 0], + poses[i : i + 2, 1], + color=BOX_DETECTION_CONFIG[DetectionType.VEHICLE].fill_color.hex, + alpha=alphas[i + 1], + linewidth=BOX_DETECTION_CONFIG[DetectionType.VEHICLE].line_width * 5, + zorder=BOX_DETECTION_CONFIG[DetectionType.VEHICLE].zorder, + ) + + def add_ego_vehicle_to_ax(ax: plt.Axes, ego_vehicle_state: Union[EgoStateSE3, EgoStateSE2]) -> None: add_bounding_box_to_ax(ax, ego_vehicle_state.bounding_box, EGO_VEHICLE_CONFIG) @@ -95,7 +130,7 @@ def add_traffic_lights_to_ax( ax: plt.Axes, traffic_light_detections: TrafficLightDetectionWrapper, map_api: AbstractMap ) -> None: for traffic_light_detection in traffic_light_detections: - lane: AbstractLane = map_api.get_map_object(str(traffic_light_detection.lane_id), MapSurfaceType.LANE) + lane: AbstractLane = map_api.get_map_object(str(traffic_light_detection.lane_id), MapLayer.LANE) if lane is not None: add_shapely_linestring_to_ax( ax, diff --git a/d123/common/visualization/matplotlib/utils.py b/d123/common/visualization/matplotlib/utils.py index 1af6f9f5..144530ab 100644 --- a/d123/common/visualization/matplotlib/utils.py +++ b/d123/common/visualization/matplotlib/utils.py @@ -1,13 +1,14 @@ from functools import lru_cache from typing import Union +import matplotlib.patches as patches import matplotlib.pyplot as plt import numpy as np import shapely.affinity as affinity import shapely.geometry as geom +from matplotlib.path import Path from d123.common.geometry.base import StateSE2, StateSE3 -from d123.common.visualization.color.color import WHITE from d123.common.visualization.color.config import PlotConfig @@ -18,54 +19,48 @@ def add_shapely_polygon_to_ax( disable_smoothing: bool = False, ) -> plt.Axes: """ - Adds shapely polygon to birds-eye-view visualization + Adds shapely polygon to birds-eye-view visualization with proper hole handling :param ax: matplotlib ax object :param polygon: shapely Polygon - :param config: dictionary containing plot parameters + :param plot_config: dictionary containing plot parameters + :param disable_smoothing: whether to overwrite smoothing of the polygon :return: ax with plot """ def _add_element_helper(element: geom.Polygon): - """Helper to add single polygon to ax""" + """Helper to add single polygon to ax with proper holes""" if plot_config.smoothing_radius is not None and not disable_smoothing: element = element.buffer(-plot_config.smoothing_radius).buffer(plot_config.smoothing_radius) - exterior_x, exterior_y = element.exterior.xy - - if plot_config.shadow: - shadow_offset = 0.5 - shadow_x = [x + shadow_offset for x in exterior_x] - shadow_y = [y - shadow_offset for y in exterior_y] - ax.fill( - shadow_x, - shadow_y, - color="gray", - alpha=1.0, - edgecolor=None, - linewidth=0, - zorder=plot_config.zorder, - ) - - ax.fill( - exterior_x, - exterior_y, - color=plot_config.fill_color.hex, + + # Create path with exterior and interior rings + def create_polygon_path(polygon): + # Get exterior coordinates + exterior_coords = list(polygon.exterior.coords) + + # Start with exterior ring + vertices = exterior_coords + codes = [Path.MOVETO] + [Path.LINETO] * (len(exterior_coords) - 2) + [Path.CLOSEPOLY] + + # Add interior rings (holes) + for interior in polygon.interiors: + interior_coords = list(interior.coords) + vertices.extend(interior_coords) + codes.extend([Path.MOVETO] + [Path.LINETO] * (len(interior_coords) - 2) + [Path.CLOSEPOLY]) + + return Path(vertices, codes) + + path = create_polygon_path(element) + + # Add main polygon with holes + patch = patches.PathPatch( + path, + facecolor=plot_config.fill_color.hex, alpha=plot_config.fill_color_alpha, edgecolor=plot_config.line_color.hex, linewidth=plot_config.line_width, zorder=plot_config.zorder, ) - - for interior in element.interiors: - x_interior, y_interior = interior.xy - ax.fill( - x_interior, - y_interior, - color=WHITE.hex, - alpha=plot_config.fill_color_alpha, - edgecolor=plot_config.line_color.hex, - linewidth=plot_config.line_width, - zorder=plot_config.zorder, - ) + ax.add_patch(patch) if isinstance(polygon, geom.Polygon): _add_element_helper(polygon) diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 56a0ecfb..087a7e8b 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -1,34 +1,114 @@ import time +from typing import Dict, List, Literal -import numpy as np import trimesh import viser +from d123.common.datatypes.sensor.camera import CameraType +from d123.common.datatypes.sensor.lidar import LiDARType from d123.common.visualization.viser.utils import ( get_bounding_box_meshes, + get_camera_values, + get_lidar_points, get_map_meshes, ) +from d123.common.visualization.viser.utils_v2 import get_bounding_box_outlines from d123.dataset.scene.abstract_scene import AbstractScene # TODO: Try to fix performance issues. # TODO: Refactor this file. +all_camera_types: List[CameraType] = [ + CameraType.CAM_F0, + CameraType.CAM_B0, + CameraType.CAM_L0, + CameraType.CAM_L1, + CameraType.CAM_L2, + CameraType.CAM_R0, + CameraType.CAM_R1, + CameraType.CAM_R2, +] + +# MISC config: +LINE_WIDTH: float = 4.0 + +# Bounding box config: +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" + +# Map config: +MAP_AVAILABLE: bool = True + + +# Cameras config: + +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_STEREO_L, CameraType.CAM_STEREO_R] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] +VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] +CAMERA_SCALE: float = 1.0 + +# Lidar config: +LIDAR_AVAILABLE: bool = False + +LIDAR_TYPES: List[LiDARType] = [ + LiDARType.LIDAR_MERGED, + LiDARType.LIDAR_TOP, + LiDARType.LIDAR_FRONT, + LiDARType.LIDAR_SIDE_LEFT, + LiDARType.LIDAR_SIDE_RIGHT, + LiDARType.LIDAR_BACK, +] +# LIDAR_TYPES: List[LiDARType] = [ +# LiDARType.LIDAR_TOP, +# ] +LIDAR_POINT_SIZE: float = 0.05 + class ViserVisualizationServer: def __init__( self, + scenes: List[AbstractScene], + scene_index: int = 0, host: str = "localhost", port: int = 8080, label: str = "D123 Viser Server", ): + assert len(scenes) > 0, "At least one scene must be provided." + self.scenes = scenes + self.scene_index = scene_index + + self.host = host + self.port = port + self.label = label + + self.server = viser.ViserServer(host=self.host, port=self.port, label=self.label) + self.set_scene(self.scenes[self.scene_index % len(self.scenes)]) - self.server = viser.ViserServer(host=host, port=port, label=label) + def next(self) -> None: + self.server.flush() + self.server.gui.reset() + self.server.scene.reset() + self.scene_index = (self.scene_index + 1) % len(self.scenes) + print(f"Viser server started at {self.host}:{self.port}") + self.set_scene(self.scenes[self.scene_index]) def set_scene(self, scene: AbstractScene) -> None: num_frames = scene.get_number_of_iterations() + # print(scene.available_camera_types) + + self.server.gui.configure_theme(dark_mode=False, control_width="large") + + # TODO: Fix lighting. Environment map can help, but cannot be freely configured. + # self.server.scene.configure_environment_map( + # hdri="warehouse", + # background=False, + # background_intensity=0.25, + # environment_intensity=0.5, + # ) - self.server.gui.configure_theme(control_width="large") with self.server.gui.add_folder("Playback"): + server_playing = True gui_timestep = self.server.gui.add_slider( "Timestep", @@ -40,6 +120,7 @@ def set_scene(self, scene: AbstractScene) -> None: ) gui_next_frame = self.server.gui.add_button("Next Frame", disabled=True) gui_prev_frame = self.server.gui.add_button("Prev Frame", disabled=True) + gui_next_scene = self.server.gui.add_button("Next Scene", disabled=False) gui_playing = self.server.gui.add_checkbox("Playing", True) gui_framerate = self.server.gui.add_slider("FPS", min=1, max=60, step=0.1, initial_value=10) gui_framerate_options = self.server.gui.add_button_group("FPS options", ("10", "20", "30", "60")) @@ -53,6 +134,11 @@ def _(_) -> None: def _(_) -> None: gui_timestep.value = (gui_timestep.value - 1) % num_frames + @gui_next_scene.on_click + def _(_) -> None: + nonlocal server_playing + server_playing = False + # Disable frame controls when we're playing. @gui_playing.on_update def _(_) -> None: @@ -76,77 +162,136 @@ def _(_) -> None: start = time.time() # with self.server.atomic(): mew_frame_handle = self.server.scene.add_frame(f"/frame{gui_timestep.value}", show_axes=False) - meshes = [] - for name, mesh in get_bounding_box_meshes(scene, gui_timestep.value).items(): - meshes.append(mesh) - self.server.scene.add_mesh_trimesh( - f"/frame{gui_timestep.value}/detections", - trimesh.util.concatenate(meshes), - visible=True, - ) - gui_image_handle.image = np.array(scene.get_front_cam_demo(gui_timestep.value)) + if BOUNDING_BOX_TYPE == "mesh": + meshes = [] + for _, mesh in get_bounding_box_meshes(scene, gui_timestep.value).items(): + meshes.append(mesh) + self.server.scene.add_mesh_trimesh( + f"/frame{gui_timestep.value}/detections", + trimesh.util.concatenate(meshes), + visible=True, + ) + elif BOUNDING_BOX_TYPE == "lines": + lines, colors = get_bounding_box_outlines(scene, gui_timestep.value) + self.server.scene.add_line_segments( + f"/frame{gui_timestep.value}/detections", + points=lines, + colors=colors, + line_width=LINE_WIDTH, + ) + else: + raise ValueError(f"Unknown bounding box type: {BOUNDING_BOX_TYPE}") + + current_frame_handle.remove() + current_frame_handle = mew_frame_handle + + for camera_type in VISUALIZE_CAMERA_GUI: + if camera_type in scene.available_camera_types: + camera_gui_handles[camera_type].image = scene.get_camera_at_iteration( + gui_timestep.value, camera_type + ).image + + for camera_type in VISUALIZE_CAMERA_FRUSTUM: + if camera_type in scene.available_camera_types: + camera_position, camera_quaternion, camera = get_camera_values( + scene, camera_type, gui_timestep.value + ) - # camera_pose = _get_camera_pose_demo(scene, gui_timestep.value) - # frustum_handle.position = camera_pose.point_3d.array - # frustum_handle.wxyz = euler_to_quaternion_scipy(camera_pose.roll, camera_pose.pitch, camera_pose.yaw) - # frustum_handle.image = np.array(scene.get_front_cam_demo(gui_timestep.value)) + camera_frustum_handles[camera_type].position = camera_position.array + camera_frustum_handles[camera_type].wxyz = camera_quaternion.q + camera_frustum_handles[camera_type].image = camera.image - # ego_frame_pose = _get_ego_frame_pose(scene, gui_timestep.value) - # ego_frame_handle.position = ego_frame_pose.point_3d.array - # ego_frame_handle.wxyz = euler_to_quaternion_scipy(ego_frame_pose.roll, ego_frame_pose.pitch, ego_frame_pose.yaw) + if LIDAR_AVAILABLE: + points, colors = get_lidar_points(scene, gui_timestep.value, LIDAR_TYPES) + gui_lidar.points = points + gui_lidar.colors = colors prev_timestep = current_timestep rendering_time = time.time() - start sleep_time = 1.0 / gui_framerate.value - rendering_time time.sleep(max(sleep_time, 0.0)) - current_frame_handle.remove() - current_frame_handle = mew_frame_handle self.server.flush() # Optional! # Load in frames. current_frame_handle = self.server.scene.add_frame(f"/frame{gui_timestep.value}", show_axes=False) self.server.scene.add_frame("/map", show_axes=False) - with self.server.gui.add_folder("Camera"): - gui_image_handle = self.server.gui.add_image( - image=np.array(scene.get_front_cam_demo(gui_timestep.value)), - label="front_cam_demo", - format="jpeg", + camera_gui_handles: Dict[CameraType, viser.GuiImageHandle] = {} + camera_frustum_handles: Dict[CameraType, viser.CameraFrustumHandle] = {} + + for camera_type in VISUALIZE_CAMERA_GUI: + if camera_type in scene.available_camera_types: + with self.server.gui.add_folder(f"Camera {camera_type.serialize()}"): + camera_gui_handles[camera_type] = self.server.gui.add_image( + image=scene.get_camera_at_iteration(gui_timestep.value, camera_type).image, + label=camera_type.serialize(), + format="jpeg", + ) + + for camera_type in VISUALIZE_CAMERA_FRUSTUM: + if camera_type in scene.available_camera_types: + camera_position, camera_quaternion, camera = get_camera_values( + scene, camera_type, gui_timestep.value + ) + camera_frustum_handles[camera_type] = self.server.scene.add_camera_frustum( + f"camera_frustum_{camera_type.serialize()}", + fov=camera.metadata.fov_y, + aspect=camera.metadata.aspect_ratio, + scale=CAMERA_SCALE, + image=camera.image, + position=camera_position.array, + wxyz=camera_quaternion.q, + ) + + if LIDAR_AVAILABLE: + points, colors = get_lidar_points(scene, gui_timestep.value, LIDAR_TYPES) + gui_lidar = self.server.scene.add_point_cloud( + name="LiDAR", + points=points, + colors=colors, + point_size=LIDAR_POINT_SIZE, + point_shape="circle", ) - for name, mesh in get_map_meshes(scene).items(): - self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) - - # camera_pose = _get_camera_pose_demo(scene, gui_timestep.value) - # frustum_handle = self.server.scene.add_camera_frustum( - # "camera_frustum", - # fov=0.6724845869242845, - # aspect=16 / 9, - # scale=0.30, - # image=np.array(scene.get_front_cam_demo(gui_timestep.value)), - # position=camera_pose.point_3d.array, - # wxyz=euler_to_quaternion_scipy(camera_pose.roll, camera_pose.pitch, camera_pose.yaw), - # ) - - # ego_frame_pose = _get_ego_frame_pose(scene, gui_timestep.value) - # ego_frame_handle = self.server.scene.add_frame( - # "ego_frame_handle", - # position=ego_frame_pose.point_3d.array, - # wxyz=euler_to_quaternion_scipy(ego_frame_pose.roll, ego_frame_pose.pitch, ego_frame_pose.yaw) - # ) + if MAP_AVAILABLE: + for name, mesh in get_map_meshes(scene).items(): + self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) + + # centerlines, __, __, road_edges = get_map_lines(scene) + # for i, centerline in enumerate(centerlines): + # self.server.scene.add_line_segments( + # "/map/centerlines", + # centerlines, + # colors=[[BLACK.rgb]], + # line_width=LINE_WIDTH, + # ) + # self.server.scene.add_line_segments( + # "/map/left_boundary", + # left_boundaries, + # colors=[[TAB_10[2].rgb]], + # line_width=LINE_WIDTH, + # ) + # self.server.scene.add_line_segments( + # "/map/right_boundary",clear + # right_boundaries, + # colors=[[TAB_10[3].rgb]], + # line_width=LINE_WIDTH, + # ) + # print(centerlines.shape, road_edges.shape) + # self.server.scene.add_line_segments( + # "/map/road_edges", + # road_edges, + # colors=[[BLACK.rgb]], + # line_width=LINE_WIDTH, + # ) # Playback update loop. prev_timestep = gui_timestep.value - while True: + while server_playing: # Update the timestep if we're playing. if gui_playing.value: gui_timestep.value = (gui_timestep.value + 1) % num_frames - # Update point size of both this timestep and the next one! There's - # redundancy here, but this will be optimized out internally by viser. - # - # We update the point size for the next timestep so that it will be - # immediately available when we toggle the visibility. - # dynamic_meshes[gui_timestep.value].point_size = gui_point_size.value - # dynamic_meshes[(gui_timestep.value + 1) % num_frames].point_size = gui_point_size.value + self.server.flush() + self.next() diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 1ab23c95..561c3676 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -1,20 +1,49 @@ +from typing import List, Tuple + import numpy as np import numpy.typing as npt import trimesh +from pyquaternion import Quaternion +from typing_extensions import Final -from d123.common.datatypes.camera.camera_parameters import get_nuplan_camera_parameters +# from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters +from d123.common.datatypes.sensor.camera import Camera, CameraType +from d123.common.datatypes.sensor.lidar import LiDARType from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_points_3d_array +from d123.common.visualization.color.color import TAB_10, Color from d123.common.visualization.color.config import PlotConfig from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG, MAP_SURFACE_CONFIG -from d123.dataset.maps.abstract_map import MapSurfaceType -from d123.dataset.maps.abstract_map_objects import AbstractSurfaceMapObject +from d123.dataset.maps.abstract_map import MapLayer +from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractSurfaceMapObject from d123.dataset.scene.abstract_scene import AbstractScene # TODO: Refactor this file. # TODO: Add general utilities for 3D primitives and mesh support. +MAP_RADIUS: Final[float] = 500 +BRIGHTNESS_FACTOR: Final[float] = 1.0 + + +def configure_trimesh(mesh: trimesh.Trimesh, color: Color): + # base_color = [r / 255.0 for r in color.rgba] + mesh.visual.face_colors = color.rgba + + # pbr_material = trimesh.visual.material.PBRMaterial( + # baseColorFactor=base_color, # Your desired color (RGBA, 0-1 range) + # metallicFactor=0.0, # 0.0 = non-metallic (more matte) + # roughnessFactor=1.0, # 0.8 = quite rough (less shiny, 0=mirror, 1=completely rough) + # emissiveFactor=[0.0, 0.0, 0.0], # No emission + # alphaCutoff=0.9, # Alpha threshold for transparency + # doubleSided=True, # Single-sided material + # ) + # mesh.visual.material = pbr_material + # mesh.visual = mesh.visual.to_texture() + + return mesh + def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> trimesh.Trimesh: @@ -22,26 +51,14 @@ def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> tr box_mesh = trimesh.creation.box(extents=[bbox.length, bbox.width, bbox.height]) # Apply rotations in order: roll, pitch, yaw - box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.roll, [1, 0, 0])) - box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.pitch, [0, 1, 0])) box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.yaw, [0, 0, 1])) + box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.pitch, [0, 1, 0])) + box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.roll, [1, 0, 0])) # Apply translation box_mesh = box_mesh.apply_translation([bbox.center.x, bbox.center.y, bbox.center.z]) - base_color = [r / 255.0 for r in plot_config.fill_color.rgba] - box_mesh.visual.face_colors = plot_config.fill_color.rgba - - pbr_material = trimesh.visual.material.PBRMaterial( - baseColorFactor=base_color, # Your desired color (RGBA, 0-1 range) - metallicFactor=1.0, # 0.0 = non-metallic (more matte) - roughnessFactor=0.9, # 0.8 = quite rough (less shiny, 0=mirror, 1=completely rough) - emissiveFactor=[0.0, 0.0, 0.0], # No emission - alphaCutoff=0.75, # Alpha threshold for transparency - doubleSided=False, # Single-sided material - ) - box_mesh.visual.material = pbr_material - return box_mesh + return configure_trimesh(box_mesh, plot_config.fill_color) def translate_bounding_box_se3(bounding_box_se3: BoundingBoxSE3, point_3d: Point3D) -> BoundingBoxSE3: @@ -78,27 +95,33 @@ def get_bounding_box_meshes(scene: AbstractScene, iteration: int): def get_map_meshes(scene: AbstractScene): initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 - map_surface_types = [ - MapSurfaceType.LANE_GROUP, - MapSurfaceType.WALKWAY, - MapSurfaceType.CROSSWALK, - MapSurfaceType.CARPARK, + map_layers = [ + MapLayer.LANE_GROUP, + MapLayer.LANE, + MapLayer.WALKWAY, + MapLayer.CROSSWALK, + MapLayer.CARPARK, + MapLayer.GENERIC_DRIVABLE, ] - radius = 500 - map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=radius, layers=map_surface_types) + map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) + print(map_objects_dict.keys()) output = {} - for map_surface_type in map_objects_dict.keys(): + for map_layer in map_objects_dict.keys(): surface_meshes = [] - for map_surface in map_objects_dict[map_surface_type]: + for map_surface in map_objects_dict[map_layer]: map_surface: AbstractSurfaceMapObject trimesh_mesh = map_surface.trimesh_mesh - if map_surface_type in [MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK]: - # trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=1.777 / 2 - 0.05).array - trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.05).array + if map_layer in [ + MapLayer.WALKWAY, + MapLayer.CROSSWALK, + MapLayer.GENERIC_DRIVABLE, + MapLayer.CARPARK, + ]: + # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. + trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.1).array else: - # trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=1.777 / 2).array trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z).array if not scene.log_metadata.map_has_z: @@ -106,13 +129,50 @@ def get_map_meshes(scene: AbstractScene): x=0, y=0, z=center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 ).array - trimesh_mesh.visual.face_colors = MAP_SURFACE_CONFIG[map_surface_type].fill_color.rgba + trimesh_mesh = configure_trimesh(trimesh_mesh, MAP_SURFACE_CONFIG[map_layer].fill_color) surface_meshes.append(trimesh_mesh) - output[f"{map_surface_type.serialize()}"] = trimesh.util.concatenate(surface_meshes) - + output[f"{map_layer.serialize()}"] = trimesh.util.concatenate(surface_meshes) return output +def get_map_lines(scene: AbstractScene): + map_layers = [MapLayer.LANE, MapLayer.ROAD_EDGE] + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) + center = initial_ego_vehicle_state.center_se3 + map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) + + def polyline_to_segments(polyline: Polyline3D) -> npt.NDArray[np.float64]: + polyline_array = polyline.array - center.point_3d.array + polyline_array = polyline_array.reshape(-1, 1, 3) + polyline_array = np.concatenate([polyline_array[:-1], polyline_array[1:]], axis=1) + return polyline_array + + centerlines, right_boundaries, left_boundaries, road_edges = [], [], [], [] + for lane in map_objects_dict[MapLayer.LANE]: + lane: AbstractLane + + centerlines.append(polyline_to_segments(lane.centerline)) + right_boundaries.append(polyline_to_segments(lane.right_boundary)) + left_boundaries.append(polyline_to_segments(lane.left_boundary)) + + for road_edge in map_objects_dict[MapLayer.ROAD_EDGE]: + road_edges.append(polyline_to_segments(road_edge.polyline_3d)) + + centerlines = np.concatenate(centerlines, axis=0) + left_boundaries = np.concatenate(left_boundaries, axis=0) + right_boundaries = np.concatenate(right_boundaries, axis=0) + road_edges = np.concatenate(road_edges, axis=0) + + if not scene.log_metadata.map_has_z: + # If the map does not have a z-coordinate, we set it to the height of the ego vehicle. + centerlines[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + left_boundaries[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + right_boundaries[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + road_edges[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + + return centerlines, left_boundaries, right_boundaries, road_edges + + def get_trimesh_from_boundaries( left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 1.0 ) -> trimesh.Trimesh: @@ -126,6 +186,8 @@ def get_trimesh_from_boundaries( def _interpolate_polyline(polyline_3d: Polyline3D, num_samples: int = 20) -> npt.NDArray[np.float64]: + if num_samples < 2: + num_samples = 2 distances = np.linspace(0, polyline_3d.length, num=num_samples, endpoint=True, dtype=np.float64) return polyline_3d.interpolate(distances) @@ -152,52 +214,38 @@ def _create_lane_mesh_from_boundary_arrays( faces = np.array(faces) mesh = trimesh.Trimesh(vertices=vertices, faces=faces) - mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapSurfaceType.LANE].fill_color.rgba + mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapLayer.LANE].fill_color.rgba return mesh -def _get_camera_pose_demo(scene: AbstractScene, iteration: int) -> StateSE3: - # NOTE: This function does not work. - +def get_camera_values( + scene: AbstractScene, camera_type: CameraType, iteration: int +) -> Tuple[Point3D, Quaternion, Camera]: initial_point_3d = scene.get_ego_state_at_iteration(0).center_se3.point_3d rear_axle = scene.get_ego_state_at_iteration(iteration).rear_axle_se3 + camera = scene.get_camera_at_iteration(iteration, camera_type) + rear_axle_array = rear_axle.array rear_axle_array[:3] -= initial_point_3d.array - rear_axle = StateSE3.from_array(rear_axle_array) - camera_parameters = get_nuplan_camera_parameters() - - # Get camera parameters - camera_translation = camera_parameters.translation # 3x1 vector as array - camera_rotation = camera_parameters.rotation # 3x3 matrix as array + camera_to_ego = camera.extrinsic # 4x4 transformation from camera to ego frame # Get the rotation matrix of the rear axle pose from d123.common.geometry.transform.se3 import get_rotation_matrix - rear_axle_rotation = get_rotation_matrix(rear_axle) - - # Apply camera translation in the rear axle's coordinate frame - # Transform camera translation from rear axle frame to world frame - world_translation = rear_axle_rotation @ camera_translation - - # Apply camera rotation in the rear axle's coordinate frame - world_rotation = rear_axle_rotation @ camera_rotation - - # Extract Euler angles from the composed rotation matrix using scipy - # Use 'ZYX' convention to match the rotation matrix composition order - from scipy.spatial.transform import Rotation + ego_transform = np.eye(4, dtype=np.float64) + ego_transform[:3, :3] = get_rotation_matrix(rear_axle) + ego_transform[:3, 3] = rear_axle.point_3d.array - r = Rotation.from_matrix(world_rotation) - roll, pitch, yaw = r.as_euler("XYZ", degrees=False) + camera_transform = ego_transform @ camera_to_ego - # Calculate camera position in world coordinates - camera_x = rear_axle.x + world_translation[0] - camera_y = rear_axle.y + world_translation[1] - camera_z = rear_axle.z + world_translation[2] + # Camera transformation in ego frame + camera_position = Point3D(*camera_transform[:3, 3]) + camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) - return StateSE3(camera_x, camera_y, camera_z, roll, pitch, yaw) + return camera_position, camera_rotation, camera def _get_ego_frame_pose(scene: AbstractScene, iteration: int) -> StateSE3: @@ -218,3 +266,52 @@ def euler_to_quaternion_scipy(roll: float, pitch: float, yaw: float) -> npt.NDAr r = Rotation.from_euler("xyz", [roll, pitch, yaw], degrees=False) quat = r.as_quat(scalar_first=True) return quat + + +def get_lidar_points( + scene: AbstractScene, iteration: int, lidar_types: List[LiDARType] +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]: + + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) + current_ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) + + def float_to_rgb(values: npt.NDArray[np.float32], cmap_name: str = "viridis") -> npt.NDArray[np.float32]: + """ + Converts an array of float values to RGB colors using a matplotlib colormap. + Normalizes values to [0, 1] using min-max scaling. + Returns an array of shape (N, 3) with RGB values in [0, 1]. + """ + import matplotlib.pyplot as plt + + vmin = np.min(values) + vmax = np.max(values) + if vmax > vmin: + normed = (values - vmin) / (vmax - vmin) + else: + normed = np.zeros_like(values) + cmap = plt.get_cmap(cmap_name) + rgb = cmap(normed)[:, :3] # Ignore alpha channel + return rgb.astype(np.float32) + + points_ = [] + colors_ = [] + for lidar_idx, lidar_type in enumerate(lidar_types): + if lidar_type not in scene.available_lidar_types: + continue + lidar = scene.get_lidar_at_iteration(iteration, lidar_type) + + # 1. convert points to ego frame + points = lidar.xyz + + # 2. convert points to world frame + origin = current_ego_vehicle_state.rear_axle_se3 + points = convert_relative_to_absolute_points_3d_array(origin, points) + points = points - initial_ego_vehicle_state.center_se3.point_3d.array + points_.append(points) + colors_.append([TAB_10[lidar_idx % len(TAB_10)].rgb] * points.shape[0]) + # colors_.append(float_to_rgb(lidar.intensity, cmap_name="viridis")) + + points_ = np.concatenate(points_, axis=0) if points_ else np.empty((0, 3), dtype=np.float32) + colors_ = np.concatenate(colors_, axis=0) if colors_ else np.empty((0, 3), dtype=np.float32) + + return points_, colors_ diff --git a/d123/common/visualization/viser/utils_v2.py b/d123/common/visualization/viser/utils_v2.py new file mode 100644 index 00000000..5db06ab8 --- /dev/null +++ b/d123/common/visualization/viser/utils_v2.py @@ -0,0 +1,99 @@ +import numpy as np +import numpy.typing as npt + +# from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters +from d123.common.geometry.base import Point3D, Point3DIndex +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 +from d123.common.geometry.bounding_box.bounding_box_index import Corners3DIndex +from d123.common.geometry.transform.se3 import translate_body_frame +from d123.common.geometry.vector import Vector3D +from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG +from d123.common.visualization.viser.utils import BRIGHTNESS_FACTOR +from d123.dataset.scene.abstract_scene import AbstractScene + +# TODO: Refactor this file. +# TODO: Add general utilities for 3D primitives and mesh support. + + +def _get_bounding_box_corners(bounding_box: BoundingBoxSE3) -> npt.NDArray[np.float64]: + """ + Get the vertices of a bounding box in 3D space. + """ + corner_extent_factors = { + Corners3DIndex.FRONT_LEFT_BOTTOM: Vector3D(+0.5, -0.5, -0.5), + Corners3DIndex.FRONT_RIGHT_BOTTOM: Vector3D(+0.5, +0.5, -0.5), + Corners3DIndex.BACK_RIGHT_BOTTOM: Vector3D(-0.5, +0.5, -0.5), + Corners3DIndex.BACK_LEFT_BOTTOM: Vector3D(-0.5, -0.5, -0.5), + Corners3DIndex.FRONT_LEFT_TOP: Vector3D(+0.5, -0.5, +0.5), + Corners3DIndex.FRONT_RIGHT_TOP: Vector3D(+0.5, +0.5, +0.5), + Corners3DIndex.BACK_RIGHT_TOP: Vector3D(-0.5, +0.5, +0.5), + Corners3DIndex.BACK_LEFT_TOP: Vector3D(-0.5, -0.5, +0.5), + } + corners = np.zeros((len(Corners3DIndex), len(Point3DIndex)), dtype=np.float64) + bounding_box_extent = np.array([bounding_box.length, bounding_box.width, bounding_box.height], dtype=np.float64) + for idx, vec in corner_extent_factors.items(): + vector_3d = Vector3D.from_array(bounding_box_extent * vec.array) + corners[idx] = translate_body_frame(bounding_box.center, vector_3d).point_3d.array + return corners + + +def _get_bounding_box_lines(bounding_box: BoundingBoxSE3) -> npt.NDArray[np.float64]: + """ + Get the edges of a bounding box in 3D space as a Polyline3D. + """ + corners = _get_bounding_box_corners(bounding_box) + index_pairs = [ + (Corners3DIndex.FRONT_LEFT_BOTTOM, Corners3DIndex.FRONT_RIGHT_BOTTOM), + (Corners3DIndex.FRONT_RIGHT_BOTTOM, Corners3DIndex.BACK_RIGHT_BOTTOM), + (Corners3DIndex.BACK_RIGHT_BOTTOM, Corners3DIndex.BACK_LEFT_BOTTOM), + (Corners3DIndex.BACK_LEFT_BOTTOM, Corners3DIndex.FRONT_LEFT_BOTTOM), + (Corners3DIndex.FRONT_LEFT_TOP, Corners3DIndex.FRONT_RIGHT_TOP), + (Corners3DIndex.FRONT_RIGHT_TOP, Corners3DIndex.BACK_RIGHT_TOP), + (Corners3DIndex.BACK_RIGHT_TOP, Corners3DIndex.BACK_LEFT_TOP), + (Corners3DIndex.BACK_LEFT_TOP, Corners3DIndex.FRONT_LEFT_TOP), + (Corners3DIndex.FRONT_LEFT_BOTTOM, Corners3DIndex.FRONT_LEFT_TOP), + (Corners3DIndex.FRONT_RIGHT_BOTTOM, Corners3DIndex.FRONT_RIGHT_TOP), + (Corners3DIndex.BACK_RIGHT_BOTTOM, Corners3DIndex.BACK_RIGHT_TOP), + (Corners3DIndex.BACK_LEFT_BOTTOM, Corners3DIndex.BACK_LEFT_TOP), + ] + lines = np.zeros((len(index_pairs), 2, len(Point3DIndex)), dtype=np.float64) + for i, (start_idx, end_idx) in enumerate(index_pairs): + lines[i, 0] = corners[start_idx] + lines[i, 1] = corners[end_idx] + return lines + + +def translate_points_3d(points_3d: npt.NDArray[np.float64], point_3d: Point3D) -> npt.NDArray[np.float64]: + return points_3d - point_3d.array + + +def get_bounding_box_outlines(scene: AbstractScene, iteration: int): + + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) + ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) + box_detections = scene.get_box_detections_at_iteration(iteration) + + lines = [] + colors = [] + for box_detection in box_detections: + bbox: BoundingBoxSE3 = box_detection.bounding_box_se3 + bbox_lines = _get_bounding_box_lines(bbox) + bbox_lines = translate_points_3d(bbox_lines, initial_ego_vehicle_state.center_se3.point_3d) + bbox_color = np.zeros(bbox_lines.shape, dtype=np.float32) + bbox_color[..., :] = ( + BOX_DETECTION_CONFIG[box_detection.metadata.detection_type] + .fill_color.set_brightness(BRIGHTNESS_FACTOR) + .rgb_norm + ) + + lines.append(bbox_lines) + colors.append(bbox_color) + + ego_bbox_lines = _get_bounding_box_lines(ego_vehicle_state.bounding_box_se3) + ego_bbox_lines = translate_points_3d(ego_bbox_lines, initial_ego_vehicle_state.center_se3.point_3d) + ego_bbox_color = np.zeros(ego_bbox_lines.shape, dtype=np.float32) + ego_bbox_color[..., :] = EGO_VEHICLE_CONFIG.fill_color.set_brightness(BRIGHTNESS_FACTOR).rgb_norm + + lines.append(ego_bbox_lines) + colors.append(ego_bbox_color) + return np.concatenate(lines, axis=0), np.concatenate(colors, axis=0) diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 10c5489c..68251f82 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -1,8 +1,14 @@ # TODO: rename this file and potentially move somewhere more appropriate. +import io +import os +from pathlib import Path +from typing import Dict, Optional import numpy as np +import numpy.typing as npt import pyarrow as pa +from PIL import Image from d123.common.datatypes.detection.detection import ( BoxDetection, @@ -14,13 +20,22 @@ TrafficLightStatus, ) from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import Camera, CameraMetadata +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.vector import Vector3D +from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import List +DATASET_SENSOR_ROOT: Dict[str, Path] = { + "nuplan": Path(os.environ["NUPLAN_DATA_ROOT"]) / "nuplan-v1.1" / "sensor_blobs", + "carla": Path(os.environ["CARLA_DATA_ROOT"]) / "sensor_blobs", + # "av2-sensor": Path(os.environ["AV2_SENSOR_DATA_ROOT"]) / "sensor", +} + def get_timepoint_from_arrow_table(arrow_table: pa.Table, index: int) -> TimePoint: return TimePoint.from_us(arrow_table["timestamp"][index].as_py()) @@ -76,3 +91,77 @@ def get_traffic_light_detections_from_arrow_table(arrow_table: pa.Table, index: traffic_light_detections.append(traffic_light_detection) return TrafficLightDetectionWrapper(traffic_light_detections=traffic_light_detections) + + +def get_camera_from_arrow_table( + arrow_table: pa.Table, + index: int, + camera_metadata: CameraMetadata, + log_metadata: LogMetadata, +) -> Camera: + + table_data = arrow_table[camera_metadata.camera_type.serialize()][index].as_py() + extrinsic = arrow_table[f"{camera_metadata.camera_type.serialize()}_extrinsic"][index].as_py() + extrinsic = np.array(extrinsic).reshape((4, 4)) if extrinsic else None + + if table_data is None or extrinsic is None: + return None + + image: Optional[npt.NDArray[np.uint8]] = None + + if isinstance(table_data, str): + sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] + full_image_path = sensor_root / table_data + assert full_image_path.exists(), f"Camera file not found: {full_image_path}" + img = Image.open(full_image_path) + img.load() + image = np.asarray(img, dtype=np.uint8) + elif isinstance(table_data, bytes): + image = np.array(Image.open(io.BytesIO(table_data))) + else: + raise NotImplementedError("Only string file paths for camera data are supported.") + + return Camera( + metadata=camera_metadata, + image=image, + extrinsic=extrinsic, + ) + + +def get_lidar_from_arrow_table( + arrow_table: pa.Table, + index: int, + lidar_metadata: LiDARMetadata, + log_metadata: LogMetadata, +) -> LiDAR: + assert ( + lidar_metadata.lidar_type.serialize() in arrow_table.schema.names + ), f'"{lidar_metadata.lidar_type.serialize()}" field not found in Arrow table schema.' + lidar_data = arrow_table[lidar_metadata.lidar_type.serialize()][index].as_py() + if isinstance(lidar_data, str): + sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] + full_lidar_path = sensor_root / lidar_data + assert full_lidar_path.exists(), f"LiDAR file not found: {full_lidar_path}" + + # NOTE: We move data specific import into if-else block, to avoid data specific import errors + if log_metadata.dataset == "nuplan": + from d123.dataset.dataset_specific.nuplan.load_sensor import load_nuplan_lidar_from_path + + lidar = load_nuplan_lidar_from_path(full_lidar_path, lidar_metadata) + elif log_metadata.dataset == "carla": + from d123.dataset.dataset_specific.carla.load_sensor import load_carla_lidar_from_path + + lidar = load_carla_lidar_from_path(full_lidar_path, lidar_metadata) + elif log_metadata.dataset == "wopd": + raise NotImplementedError + else: + raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") + + else: + if log_metadata.dataset == "wopd": + lidar_data = np.array(lidar_data, dtype=np.float64) + lidar = LiDAR(metadata=lidar_metadata, point_cloud=lidar_data.T) + else: + raise NotImplementedError("Only string file paths for lidar data are supported.") + + return lidar diff --git a/d123/script/config/dataset_caching/__init__.py b/d123/dataset/conversion/__init__.py similarity index 100% rename from d123/script/config/dataset_caching/__init__.py rename to d123/dataset/conversion/__init__.py diff --git a/d123/simulation/gym/environment/trajectory_builder/__init__.py b/d123/dataset/conversion/map/__init__.py similarity index 100% rename from d123/simulation/gym/environment/trajectory_builder/__init__.py rename to d123/dataset/conversion/map/__init__.py diff --git a/d123/dataset/dataset_specific/waymo/.gitkeep b/d123/dataset/conversion/map/opendrive/__init__.py similarity index 100% rename from d123/dataset/dataset_specific/waymo/.gitkeep rename to d123/dataset/conversion/map/opendrive/__init__.py diff --git a/d123/dataset/conversion/map/opendrive/opendrive_converter.py b/d123/dataset/conversion/map/opendrive/opendrive_converter.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py b/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py new file mode 100644 index 00000000..56dfa65c --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py @@ -0,0 +1,422 @@ +import logging +import os +import warnings +from pathlib import Path +from typing import Dict, Final, List + +import geopandas as gpd +import numpy as np +import pandas as pd +import shapely +from shapely.ops import polygonize, unary_union + +from d123.dataset.conversion.map.opendrive.parser.opendrive import Junction, OpenDrive +from d123.dataset.conversion.map.opendrive.utils.collection import collect_element_helpers +from d123.dataset.conversion.map.opendrive.utils.id_mapping import IntIDMapping +from d123.dataset.conversion.map.opendrive.utils.lane_helper import ( + OpenDriveLaneGroupHelper, + OpenDriveLaneHelper, +) +from d123.dataset.conversion.map.opendrive.utils.objects_helper import ( + OpenDriveObjectHelper, +) +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length +from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import get_road_edges_3d_from_gdf +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType + +logger = logging.getLogger(__name__) +D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) + +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # [m] + + +def convert_from_xodr( + xordr_file: Path, + map_name: str, + interpolation_step_size: float, + connection_distance_threshold: float, +) -> None: + + opendrive = OpenDrive.parse_from_file(xordr_file) + + _, junction_dict, lane_helper_dict, lane_group_helper_dict, object_helper_dict = collect_element_helpers( + opendrive, interpolation_step_size, connection_distance_threshold + ) + + # Collect data frames and store + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + dataframes[MapLayer.LANE] = _extract_lane_dataframe(lane_group_helper_dict) + dataframes[MapLayer.LANE_GROUP] = _extract_lane_group_dataframe(lane_group_helper_dict) + dataframes[MapLayer.WALKWAY] = _extract_walkways_dataframe(lane_helper_dict) + dataframes[MapLayer.CARPARK] = _extract_carpark_dataframe(lane_helper_dict) + dataframes[MapLayer.GENERIC_DRIVABLE] = _extract_generic_drivable_dataframe(lane_helper_dict) + dataframes[MapLayer.INTERSECTION] = _extract_intersections_dataframe(junction_dict, lane_group_helper_dict) + dataframes[MapLayer.CROSSWALK] = _extract_crosswalk_dataframe(object_helper_dict) + + _convert_ids_to_int( + dataframes[MapLayer.LANE], + dataframes[MapLayer.WALKWAY], + dataframes[MapLayer.CARPARK], + dataframes[MapLayer.GENERIC_DRIVABLE], + dataframes[MapLayer.LANE_GROUP], + dataframes[MapLayer.INTERSECTION], + dataframes[MapLayer.CROSSWALK], + ) + dataframes[MapLayer.ROAD_EDGE] = _extract_road_edge_df( + dataframes[MapLayer.LANE], + dataframes[MapLayer.CARPARK], + dataframes[MapLayer.GENERIC_DRIVABLE], + dataframes[MapLayer.LANE_GROUP], + ) + dataframes[MapLayer.ROAD_LINE] = _extract_road_line_df( + dataframes[MapLayer.LANE], + dataframes[MapLayer.LANE_GROUP], + ) + map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_name, layer=layer.serialize(), driver="GPKG", mode="a") + + +def _extract_lane_dataframe(lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper]) -> gpd.GeoDataFrame: + + ids = [] + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] + baseline_paths = [] + geometries = [] + + for lane_group_helper in lane_group_helper_dict.values(): + lane_group_id = lane_group_helper.lane_group_id + lane_helpers = lane_group_helper.lane_helpers + num_lanes = len(lane_helpers) + # NOTE: Lanes are going left to right, ie. inner to outer + for lane_idx, lane_helper in enumerate(lane_helpers): + ids.append(lane_helper.lane_id) + lane_group_ids.append(lane_group_id) + speed_limits_mps.append(lane_helper.speed_limit_mps) + predecessor_ids.append(lane_helper.predecessor_lane_ids) + successor_ids.append(lane_helper.successor_lane_ids) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + baseline_paths.append(shapely.LineString(lane_helper.center_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + left_lane_id = lane_helpers[lane_idx - 1].lane_id if lane_idx > 0 else None + right_lane_id = lane_helpers[lane_idx + 1].lane_id if lane_idx < num_lanes - 1 else None + left_lane_ids.append(left_lane_id) + right_lane_ids.append(right_lane_id) + + data = pd.DataFrame( + { + "id": ids, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, + "baseline_path": baseline_paths, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_lane_group_dataframe(lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper]) -> gpd.GeoDataFrame: + + ids = [] + lane_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + intersection_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + for lane_group_helper in lane_group_helper_dict.values(): + lane_group_helper: OpenDriveLaneGroupHelper + ids.append(lane_group_helper.lane_group_id) + lane_ids.append([lane_helper.lane_id for lane_helper in lane_group_helper.lane_helpers]) + predecessor_lane_group_ids.append(lane_group_helper.predecessor_lane_group_ids) + successor_lane_group_ids.append(lane_group_helper.successor_lane_group_ids) + intersection_ids.append(lane_group_helper.junction_id) + left_boundaries.append(shapely.LineString(lane_group_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_group_helper.outer_polyline_3d)) + geometries.append(lane_group_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "predecessor_ids": predecessor_lane_group_ids, + "successor_ids": successor_lane_group_ids, + "intersection_id": intersection_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + + return gdf + + +def _extract_walkways_dataframe(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> gpd.GeoDataFrame: + + ids = [] + left_boundaries = [] + right_boundaries = [] + outlines = [] + geometries = [] + + for lane_helper in lane_helper_dict.values(): + if lane_helper.type == "sidewalk": + ids.append(lane_helper.lane_id) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "outline": outlines, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_carpark_dataframe(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> gpd.GeoDataFrame: + + ids = [] + left_boundaries = [] + right_boundaries = [] + outlines = [] + geometries = [] + + for lane_helper in lane_helper_dict.values(): + if lane_helper.type == "parking": + ids.append(lane_helper.lane_id) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "outline": outlines, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_generic_drivable_dataframe(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> gpd.GeoDataFrame: + + ids = [] + left_boundaries = [] + right_boundaries = [] + outlines = [] + geometries = [] + + for lane_helper in lane_helper_dict.values(): + if lane_helper.type in ["none", "border", "bidirectional"]: + ids.append(lane_helper.lane_id) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "left_boundary": left_boundaries, + "right_boundary": left_boundaries, + "outline": outlines, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_intersections_dataframe( + junction_dict: Dict[str, Junction], + lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper], +) -> gpd.GeoDataFrame: + def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriveLaneGroupHelper]: + return [ + lane_group_helper + for lane_group_helper in lane_group_helper_dict.values() + if lane_group_helper.junction_id == junction_id + ] + + ids = [] + lane_group_ids = [] + geometries = [] + for junction in junction_dict.values(): + lane_group_helpers = _find_lane_group_helpers_with_junction_id(junction.id) + lane_group_ids_ = [lane_group_helper.lane_group_id for lane_group_helper in lane_group_helpers] + if len(lane_group_ids_) == 0: + logger.debug(f"Skipping Junction {junction.id} without lane groups!") + continue + + polygon = extract_exteriors_polygon(lane_group_helpers) + ids.append(junction.id) + lane_group_ids.append(lane_group_ids_) + geometries.append(polygon) + + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_crosswalk_dataframe(object_helper_dict: Dict[int, OpenDriveObjectHelper]) -> gpd.GeoDataFrame: + ids = [] + outlines = [] + geometries = [] + for object_helper in object_helper_dict.values(): + ids.append(object_helper.object_id) + outlines.append(shapely.LineString(object_helper.outline_3d)) + geometries.append(object_helper.shapely_polygon) + + data = pd.DataFrame({"id": ids, "outline": outlines}) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _convert_ids_to_int( + lane_df: gpd.GeoDataFrame, + walkways_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, + intersections_df: gpd.GeoDataFrame, + crosswalk_df: gpd.GeoDataFrame, +) -> None: + + # NOTE: intersection and crosswalk ids are already integers + + # initialize id mappings + lane_id_mapping = IntIDMapping.from_series(lane_df["id"]) + walkway_id_mapping = IntIDMapping.from_series(walkways_df["id"]) + carpark_id_mapping = IntIDMapping.from_series(carpark_df["id"]) + generic_drivable_id_mapping = IntIDMapping.from_series(generic_drivable_area_df["id"]) + lane_group_id_mapping = IntIDMapping.from_series(lane_group_df["id"]) + + # Adjust cross reference in lane_df and lane_group_df + lane_df["lane_group_id"] = lane_df["lane_group_id"].map(lane_group_id_mapping.str_to_int) + lane_group_df["lane_ids"] = lane_group_df["lane_ids"].apply(lambda x: lane_id_mapping.map_list(x)) + + # Adjust predecessor/successor in lane_df and lane_group_df + for column in ["predecessor_ids", "successor_ids"]: + lane_df[column] = lane_df[column].apply(lambda x: lane_id_mapping.map_list(x)) + lane_group_df[column] = lane_group_df[column].apply(lambda x: lane_group_id_mapping.map_list(x)) + + for column in ["left_lane_id", "right_lane_id"]: + lane_df[column] = lane_df[column].apply( + lambda x: str(lane_id_mapping.str_to_int[x]) if pd.notna(x) and x is not None else x + ) + + lane_df["id"] = lane_df["id"].map(lane_id_mapping.str_to_int) + walkways_df["id"] = walkways_df["id"].map(walkway_id_mapping.str_to_int) + carpark_df["id"] = carpark_df["id"].map(carpark_id_mapping.str_to_int) + generic_drivable_area_df["id"] = generic_drivable_area_df["id"].map(generic_drivable_id_mapping.str_to_int) + lane_group_df["id"] = lane_group_df["id"].map(lane_group_id_mapping.str_to_int) + + intersections_df["lane_group_ids"] = intersections_df["lane_group_ids"].apply( + lambda x: lane_group_id_mapping.map_list(x) + ) + + +def _extract_road_line_df( + lane_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> None: + + lane_group_on_intersection = { + lane_group_id: str(intersection_id) != "nan" + for lane_group_id, intersection_id in zip(lane_group_df.id.tolist(), lane_group_df.intersection_id.tolist()) + } + ids = [] + road_line_types = [] + geometries = [] + + running_id = 0 + for lane_row in lane_df.itertuples(): + on_intersection = lane_group_on_intersection.get(lane_row.lane_group_id, False) + if on_intersection: + # Skip road lines on intersections + continue + if str(lane_row.right_lane_id) in ["nan", "None"]: + # This is a boundary lane, e.g. a border or sidewalk + ids.append(running_id) + road_line_types.append(int(RoadLineType.SOLID_WHITE)) + geometries.append(lane_row.right_boundary) + running_id += 1 + else: + # This is a regular lane + ids.append(running_id) + road_line_types.append(int(RoadLineType.DASHED_WHITE)) + geometries.append(lane_row.right_boundary) + running_id += 1 + if str(lane_row.left_lane_id) in ["nan", "None"]: + # This is a boundary lane, e.g. a border or sidewalk + ids.append(running_id) + road_line_types.append(int(RoadLineType.DASHED_WHITE)) + geometries.append(lane_row.left_boundary) + running_id += 1 + + data = pd.DataFrame({"id": ids, "road_line_type": road_line_types}) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_road_edge_df( + lane_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> None: + road_edges = get_road_edges_3d_from_gdf(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) + road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) + + ids = np.arange(len(road_edges), dtype=np.int64).tolist() + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types = [int(RoadEdgeType.ROAD_EDGE_BOUNDARY)] * len(road_edges) + geometries = road_edges + return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) + + +# TODO: move this somewhere else and improve +def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper]) -> shapely.Polygon: + + # Step 1: Extract all boundary line segments + all_polygons = [] + for lane_group_helper in lane_group_helpers: + all_polygons.append(lane_group_helper.shapely_polygon) + + # Step 2: Merge all boundaries and extract the enclosed polygons + # try: + merged_boundaries = unary_union(all_polygons) + # except Exception as e: + # warnings.warn(f"Topological error during polygon union: {e}") + # print([(helper.lane_group_id, poly.is_valid) for poly, helper in zip(all_polygons, lane_group_helpers)]) + # merged_boundaries = unary_union([poly for poly in all_polygons if poly.is_valid]) + + # Step 3: Generate polygons from the merged lines + polygons = list(polygonize(merged_boundaries)) + + # Step 4: Select the polygon that represents the intersection + # Usually it's the largest polygon + if len(polygons) == 1: + return polygons[0] + else: + # Take the largest polygon if there are multiple + return max(polygons, key=lambda p: p.area) diff --git a/d123/dataset/conversion/map/opendrive/parser/__init__.py b/d123/dataset/conversion/map/opendrive/parser/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/elevation.py b/d123/dataset/conversion/map/opendrive/parser/elevation.py similarity index 56% rename from d123/dataset/dataset_specific/carla/opendrive/elements/elevation.py rename to d123/dataset/conversion/map/opendrive/parser/elevation.py index 782c4311..7ee7aa67 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/elevation.py +++ b/d123/dataset/conversion/map/opendrive/parser/elevation.py @@ -4,12 +4,17 @@ from typing import List, Optional from xml.etree.ElementTree import Element -import numpy as np -import numpy.typing as npt +from d123.dataset.conversion.map.opendrive.parser.polynomial import Polynomial @dataclass class ElevationProfile: + """ + Models elevation along s-axis of reference line. + + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-1d876c00-d69e-46d9-bbcd-709ab48f14b1 + """ + elevations: List[Elevation] def __post_init__(self): @@ -27,27 +32,21 @@ def parse(cls, elevation_profile_element: Optional[Element]) -> ElevationProfile @dataclass -class Elevation: - """TODO: Refactor and merge with other elements, e.g. LaneOffset""" - - s: float - a: float - b: float - c: float - d: float - - @classmethod - def parse(cls, elevation_element: Element) -> Elevation: - args = {key: float(elevation_element.get(key)) for key in ["s", "a", "b", "c", "d"]} - return Elevation(**args) +class Elevation(Polynomial): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + Represents an elevation profile in OpenDRIVE. + """ @dataclass class LateralProfile: + """ + Models elevation along t-axis of reference line. + + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 + """ super_elevations: List[SuperElevation] shapes: List[Shape] @@ -76,28 +75,21 @@ def parse(cls, lateral_profile_element: Optional[Element]) -> LateralProfile: @dataclass -class SuperElevation: - """TODO: Refactor and merge with other elements, e.g. Elevation, LaneOffset""" - - s: float - a: float - b: float - c: float - d: float +class SuperElevation(Polynomial): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-4abf7baf-fb2f-4263-8133-ad0f64f0feac - @classmethod - def parse(cls, super_elevation_element: Element) -> SuperElevation: - args = {key: float(super_elevation_element.get(key)) for key in ["s", "a", "b", "c", "d"]} - return SuperElevation(**args) - - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + superelevation (ds) = a + b*ds + c*ds² + d*ds³ + """ @dataclass class Shape: - """TODO: Refactor and merge with other elements, e.g. Elevation, LaneOffset""" + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/1.8.0/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 + + hShape (dt)= a + b*dt + c*dt² + d*dt³ + """ s: float t: float @@ -111,6 +103,8 @@ def parse(cls, shape_element: Element) -> Shape: args = {key: float(shape_element.get(key)) for key in ["s", "t", "a", "b", "c", "d"]} return Shape(**args) - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + def get_value(self, dt: float) -> float: + """ + Evaluate the polynomial at a given t value. + """ + return self.a + self.b * dt + self.c * dt**2 + self.d * dt**3 diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/geometry.py b/d123/dataset/conversion/map/opendrive/parser/geometry.py similarity index 76% rename from d123/dataset/dataset_specific/carla/opendrive/elements/geometry.py rename to d123/dataset/conversion/map/opendrive/parser/geometry.py index 2dd33c4e..c1f337c0 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/geometry.py +++ b/d123/dataset/conversion/map/opendrive/parser/geometry.py @@ -13,6 +13,10 @@ @dataclass class Geometry: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_02_road_reference_line.html + """ + s: float x: float y: float @@ -33,6 +37,10 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: @dataclass class Line(Geometry): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_03_straight_line.html + """ + @classmethod def parse(cls, geometry_element: Element) -> Geometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} @@ -53,9 +61,16 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: @dataclass class Arc(Geometry): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_05_arc.html + """ curvature: float + def __post_init__(self): + if self.curvature == 0.0: + raise ValueError("Curvature cannot be zero for Arc geometry.") + @classmethod def parse(cls, geometry_element: Element) -> Geometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} @@ -64,19 +79,20 @@ def parse(cls, geometry_element: Element) -> Geometry: def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - c = self.curvature - hdg = self.hdg - np.pi / 2 + kappa = self.curvature + radius = 1.0 / kappa if kappa != 0 else float("inf") - a = 2 / c * np.sin(s * c / 2) - alpha = (np.pi - s * c) / 2 - hdg + initial_heading = self.hdg + final_heading = initial_heading + s * kappa - dx = -1 * a * np.cos(alpha) - dy = a * np.sin(alpha) + # Parametric circle equations + dx = radius * (np.sin(final_heading) - np.sin(initial_heading)) + dy = -radius * (np.cos(final_heading) - np.cos(initial_heading)) interpolated_se2 = self.start_se2.copy() interpolated_se2[StateSE2Index.X] += dx interpolated_se2[StateSE2Index.Y] += dy - interpolated_se2[StateSE2Index.YAW] += s * self.curvature + interpolated_se2[StateSE2Index.YAW] = final_heading if t != 0.0: interpolated_se2[StateSE2Index.X] += t * np.cos(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) @@ -85,32 +101,21 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: return interpolated_se2 -@dataclass -class Geometry: - s: float - x: float - y: float - hdg: float - length: float - - @property - def start_se2(self) -> npt.NDArray[np.float64]: - start_se2 = np.zeros(len(StateSE2Index), dtype=np.float64) - start_se2[StateSE2Index.X] = self.x - start_se2[StateSE2Index.Y] = self.y - start_se2[StateSE2Index.YAW] = self.hdg - return start_se2 - - def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - raise NotImplementedError - - @dataclass class Spiral(Geometry): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_04_spiral.html + https://en.wikipedia.org/wiki/Euler_spiral + """ curvature_start: float curvature_end: float + def __post_init__(self): + gamma = (self.curvature_end - self.curvature_start) / self.length + if abs(gamma) < 1e-10: + raise ValueError("Curvature change is too small, cannot define a valid spiral.") + @classmethod def parse(cls, geometry_element: Element) -> Geometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} @@ -120,18 +125,9 @@ def parse(cls, geometry_element: Element) -> Geometry: return cls(**args) def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - """ - https://en.wikipedia.org/wiki/Euler_spiral - :param s: _description_ - :param t: _description_, defaults to 0.0 - :return: _description_ - """ interpolated_se2 = self.start_se2.copy() gamma = (self.curvature_end - self.curvature_start) / self.length - if abs(gamma) < 1e-10: - print(gamma) - # NOTE: doesn't consider case where gamma == 0 dx, dy = self._compute_spiral_position(s, gamma) diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/lane.py b/d123/dataset/conversion/map/opendrive/parser/lane.py similarity index 77% rename from d123/dataset/dataset_specific/carla/opendrive/elements/lane.py rename to d123/dataset/conversion/map/opendrive/parser/lane.py index 9afb0f25..ccbfc9ff 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/lane.py +++ b/d123/dataset/conversion/map/opendrive/parser/lane.py @@ -4,12 +4,15 @@ from typing import List, Optional from xml.etree.ElementTree import Element -import numpy as np -import numpy.typing as npt +from d123.dataset.conversion.map.opendrive.parser.polynomial import Polynomial @dataclass class Lanes: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_01_introduction.html + """ + lane_offsets: List[LaneOffset] lane_sections: List[LaneSection] @@ -32,40 +35,29 @@ def parse(cls, lanes_element: Optional[Element]) -> Lanes: return Lanes(**args) @property - def num_lane_sections(self): + def num_lane_sections(self) -> int: return len(self.lane_sections) @property - def last_lane_section_idx(self): + def last_lane_section_idx(self) -> int: return self.num_lane_sections - 1 @dataclass -class LaneOffset: - """Section 11.4""" - - s: float - a: float - b: float - c: float - d: float - - def __post_init__(self): - # NOTE: added assertion/filtering to check for element type or consistency - pass +class LaneOffset(Polynomial): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_04_lane_offset.html - @classmethod - def parse(cls, lane_offset_element: Element) -> LaneOffset: - args = {key: float(lane_offset_element.get(key)) for key in ["s", "a", "b", "c", "d"]} - return LaneOffset(**args) - - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + offset (ds) = a + b*ds + c*ds² + d*ds³ + """ @dataclass class LaneSection: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_03_lane_sections.html + """ + s: float left_lanes: List[Lane] center_lanes: List[Lane] @@ -104,6 +96,9 @@ def parse(cls, lane_section_element: Optional[Element]) -> LaneSection: @dataclass class Lane: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_05_lane_link.html + """ id: int type: str @@ -147,6 +142,10 @@ def parse(cls, lane_element: Optional[Element]) -> Lane: @dataclass class Width: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_06_lane_geometry.html#sec-8d8ac2e0-b3d6-4048-a9ed-d5191af5c74b + """ + s_offset: Optional[float] = None a: Optional[float] = None b: Optional[float] = None @@ -163,13 +162,25 @@ def parse(cls, width_element: Optional[Element]) -> Width: args["d"] = float(width_element.get("d")) return Width(**args) - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + def get_polynomial(self, t_sign: float = 1.0) -> Polynomial: + """ + Returns the polynomial representation of the width. + """ + return Polynomial( + s=self.s_offset, + a=self.a * t_sign, + b=self.b * t_sign, + c=self.c * t_sign, + d=self.d * t_sign, + ) @dataclass class RoadMark: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_08_road_markings.html + """ + s_offset: float = None type: str = None material: Optional[str] = None diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/objects.py b/d123/dataset/conversion/map/opendrive/parser/objects.py similarity index 85% rename from d123/dataset/dataset_specific/carla/opendrive/elements/objects.py rename to d123/dataset/conversion/map/opendrive/parser/objects.py index 4ca2f730..4c409c54 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/objects.py +++ b/d123/dataset/conversion/map/opendrive/parser/objects.py @@ -7,6 +7,9 @@ @dataclass class Object: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/13_objects/13_01_introduction.html + """ id: int name: str @@ -55,6 +58,10 @@ def parse(cls, object_element: Optional[Element]) -> Object: @dataclass class CornerLocal: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/13_objects/13_03_object_outline.html#sec-cc00c8a6-eea6-49e6-b90c-37b21524c748 + """ + u: float v: float z: float diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/opendrive.py b/d123/dataset/conversion/map/opendrive/parser/opendrive.py similarity index 79% rename from d123/dataset/dataset_specific/carla/opendrive/elements/opendrive.py rename to d123/dataset/conversion/map/opendrive/parser/opendrive.py index f06876ef..2a07bb5f 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/opendrive.py +++ b/d123/dataset/conversion/map/opendrive/parser/opendrive.py @@ -3,14 +3,17 @@ import traceback from dataclasses import dataclass from pathlib import Path -from typing import List, Optional +from typing import List, Literal, Optional from xml.etree.ElementTree import Element, parse -from d123.dataset.dataset_specific.carla.opendrive.elements.road import Road +from d123.dataset.conversion.map.opendrive.parser.road import Road @dataclass class OpenDrive: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/06_general_architecture/06_03_root_element.html + """ header: Header @@ -55,7 +58,9 @@ def parse_from_file(cls, file_path: Path) -> OpenDrive: @dataclass class Header: - """Section 4.4.2""" + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/06_general_architecture/06_04_header.html + """ rev_major: Optional[int] = None rev_minor: Optional[int] = None @@ -95,6 +100,10 @@ def parse(cls, header_element: Optional[Element]) -> Header: @dataclass class Controller: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/14_signals/14_06_controllers.html + """ + name: str id: int sequence: int @@ -118,6 +127,9 @@ def parse(cls, controller_element: Optional[Element]) -> Junction: @dataclass class Control: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/14_signals/14_06_controllers.html + """ signal_id: str type: str @@ -132,6 +144,10 @@ def parse(cls, control_element: Optional[Element]) -> Control: @dataclass class Junction: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/12_junctions/12_02_common_junctions.html + """ + id: int name: str connections: List[Connection] @@ -153,12 +169,19 @@ def parse(cls, junction_element: Optional[Element]) -> Junction: @dataclass class Connection: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/12_junctions/12_02_common_junctions.html + """ + id: int incoming_road: int connecting_road: int - contact_point: str + contact_point: Literal["start", "end"] lane_links: List[LaneLink] + def __post_init__(self): + assert self.contact_point in ["start", "end"] + @classmethod def parse(cls, connection_element: Optional[Element]) -> Connection: args = {} @@ -178,6 +201,10 @@ def parse(cls, connection_element: Optional[Element]) -> Connection: @dataclass class LaneLink: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/1.8.0/specification/12_junctions/12_04_connecting_roads.html#top-3e9bb97e-f2ab-4751-906a-c25e9fb7ac4e + """ + start: int # NOTE: named "from" in xml end: int # NOTE: named "to" in xml diff --git a/d123/dataset/conversion/map/opendrive/parser/polynomial.py b/d123/dataset/conversion/map/opendrive/parser/polynomial.py new file mode 100644 index 00000000..497ffed5 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/parser/polynomial.py @@ -0,0 +1,32 @@ +from dataclasses import dataclass +from xml.etree.ElementTree import Element + + +@dataclass +class Polynomial: + """ + Multiple OpenDRIVE elements use polynomial coefficients, e.g. Elevation, LaneOffset, etc. + This class provides a common interface to parse and access polynomial coefficients. + + Stores a polynomial function of the third order: + value(ds) = a + b*ds + c*ds² + d*ds³ + """ + + s: float + a: float + b: float + c: float + d: float + + @classmethod + def parse(cls: type["Polynomial"], element: Element) -> "Polynomial": + args = {key: float(element.get(key)) for key in ["s", "a", "b", "c", "d"]} + return cls(**args) + + def get_value(self, ds: float) -> float: + """ + Evaluate the polynomial at a given ds value. + :param ds: The distance along the road. + :return: The evaluated polynomial value. + """ + return self.a + self.b * ds + self.c * ds**2 + self.d * ds**3 diff --git a/d123/dataset/conversion/map/opendrive/parser/reference.py b/d123/dataset/conversion/map/opendrive/parser/reference.py new file mode 100644 index 00000000..b4c88dc1 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/parser/reference.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from functools import cached_property +from typing import Final, List, Optional, Union +from xml.etree.ElementTree import Element + +import numpy as np +import numpy.typing as npt + +from d123.common.geometry.base import Point3DIndex, StateSE2Index +from d123.dataset.conversion.map.opendrive.parser.elevation import Elevation +from d123.dataset.conversion.map.opendrive.parser.geometry import Arc, Geometry, Line, Spiral +from d123.dataset.conversion.map.opendrive.parser.lane import LaneOffset, Width +from d123.dataset.conversion.map.opendrive.parser.polynomial import Polynomial + +TOLERANCE: Final[float] = 1e-3 + + +@dataclass +class PlanView: + + geometries: List[Geometry] + + def __post_init__(self): + # Ensure geometries are sorted by their starting position 's' + self.geometries.sort(key=lambda x: x.s) + + @classmethod + def parse(cls, plan_view_element: Optional[Element]) -> PlanView: + geometries: List[Geometry] = [] + for geometry_element in plan_view_element.findall("geometry"): + if geometry_element.find("line") is not None: + geometry = Line.parse(geometry_element) + elif geometry_element.find("arc") is not None: + geometry = Arc.parse(geometry_element) + elif geometry_element.find("spiral") is not None: + geometry = Spiral.parse(geometry_element) + else: + geometry_str = ET.tostring(geometry_element, encoding="unicode") + raise NotImplementedError(f"Geometry not implemented: {geometry_str}") + geometries.append(geometry) + return PlanView(geometries=geometries) + + @cached_property + def geometry_lengths(self) -> npt.NDArray[np.float64]: + return np.cumsum([0.0] + [geo.length for geo in self.geometries], dtype=np.float64) + + @property + def length(self) -> float: + return float(self.geometry_lengths[-1]) + + def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: + """ + Interpolates the SE2 state at a given longitudinal position s along the plan view. + """ + if s > self.length: + if np.isclose(s, self.length, atol=TOLERANCE): + s = self.length + else: + raise ValueError( + f"s={s} is beyond the end of the plan view (length={self.length}) with tolerance={TOLERANCE}." + ) + + # Find the geometry segment containing s + geo_idx = np.searchsorted(self.geometry_lengths, s, side="right") - 1 + geo_idx = int(np.clip(geo_idx, 0, len(self.geometries) - 1)) + + return self.geometries[geo_idx].interpolate_se2(s - self.geometry_lengths[geo_idx], t) + + +@dataclass +class ReferenceLine: + + reference_line: Union[ReferenceLine, PlanView] + width_polynomials: List[Polynomial] + elevations: List[Elevation] + s_offset: float + + @property + def length(self) -> float: + return float(self.reference_line.length) + + @classmethod + def from_plan_view( + cls, + plan_view: PlanView, + lane_offsets: List[LaneOffset], + elevations: List[Elevation], + ) -> ReferenceLine: + args = {} + args["reference_line"] = plan_view + args["width_polynomials"] = lane_offsets + args["elevations"] = elevations + args["s_offset"] = 0.0 + return ReferenceLine(**args) + + @classmethod + def from_reference_line( + cls, + reference_line: ReferenceLine, + widths: List[Width], + s_offset: float = 0.0, + t_sign: float = 1.0, + ) -> ReferenceLine: + assert t_sign in [1.0, -1.0], "t_sign must be either 1.0 or -1.0" + + args = {} + args["reference_line"] = reference_line + width_polynomials: List[Polynomial] = [] + for width in widths: + width_polynomials.append(width.get_polynomial(t_sign=t_sign)) + args["width_polynomials"] = width_polynomials + args["s_offset"] = s_offset + args["elevations"] = reference_line.elevations + + return ReferenceLine(**args) + + @staticmethod + def _find_polynomial(s: float, polynomials: List[Polynomial], lane_section_end: bool = False) -> Polynomial: + + out_polynomial = polynomials[-1] + for polynomial in polynomials[::-1]: + if lane_section_end: + if polynomial.s < s: + out_polynomial = polynomial + break + else: + if polynomial.s <= s: + out_polynomial = polynomial + break + + return out_polynomial + + def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: + + width_polynomial = self._find_polynomial(s, self.width_polynomials, lane_section_end=lane_section_end) + t_offset = width_polynomial.get_value(s - width_polynomial.s) + se2 = self.reference_line.interpolate_se2(self.s_offset + s, t=t_offset + t, lane_section_end=lane_section_end) + + return se2 + + def interpolate_3d(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: + + se2 = self.interpolate_se2(s, t, lane_section_end=lane_section_end) + + elevation_polynomial = self._find_polynomial(s, self.elevations, lane_section_end=lane_section_end) + point_3d = np.zeros(len(Point3DIndex), dtype=np.float64) + point_3d[Point3DIndex.XY] = se2[StateSE2Index.XY] + point_3d[Point3DIndex.Z] = elevation_polynomial.get_value(s - elevation_polynomial.s) + + return point_3d diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/road.py b/d123/dataset/conversion/map/opendrive/parser/road.py similarity index 77% rename from d123/dataset/dataset_specific/carla/opendrive/elements/road.py rename to d123/dataset/conversion/map/opendrive/parser/road.py index b22c77ed..e0171429 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/road.py +++ b/d123/dataset/conversion/map/opendrive/parser/road.py @@ -4,17 +4,21 @@ from typing import List, Optional from xml.etree.ElementTree import Element -from d123.dataset.dataset_specific.carla.opendrive.elements.elevation import ElevationProfile, LateralProfile -from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lanes -from d123.dataset.dataset_specific.carla.opendrive.elements.objects import Object -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import PlanView +from d123.dataset.conversion.map.opendrive.parser.elevation import ElevationProfile, LateralProfile +from d123.dataset.conversion.map.opendrive.parser.lane import Lanes +from d123.dataset.conversion.map.opendrive.parser.objects import Object +from d123.dataset.conversion.map.opendrive.parser.reference import PlanView @dataclass class Road: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_01_introduction.html + """ + id: int junction: Optional[str] - length: float # [m] + length: float name: Optional[str] link: Link @@ -66,7 +70,9 @@ def parse(cls, road_element: Element) -> Road: @dataclass class Link: - """Section 8.2""" + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_03_road_linkage.html + """ predecessor: Optional[PredecessorSuccessor] = None successor: Optional[PredecessorSuccessor] = None @@ -84,6 +90,10 @@ def parse(cls, link_element: Optional[Element]) -> PlanView: @dataclass class PredecessorSuccessor: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_03_road_linkage.html + """ + element_type: Optional[str] = None element_id: Optional[int] = None contact_point: Optional[str] = None @@ -103,6 +113,9 @@ def parse(cls, element: Element) -> PredecessorSuccessor: @dataclass class RoadType: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_04_road_type.html + """ s: Optional[float] = None type: Optional[str] = None @@ -120,6 +133,10 @@ def parse(cls, road_type_element: Optional[Element]) -> RoadType: @dataclass class Speed: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_04_road_type.html#sec-33dc6899-854e-4533-a3d9-76e9e1518ee7 + """ + max: Optional[float] = None unit: Optional[str] = None diff --git a/d123/dataset/conversion/map/opendrive/utils/collection.py b/d123/dataset/conversion/map/opendrive/utils/collection.py new file mode 100644 index 00000000..be631445 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/utils/collection.py @@ -0,0 +1,325 @@ +import logging +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from d123.dataset.conversion.map.opendrive.parser.opendrive import Junction, OpenDrive +from d123.dataset.conversion.map.opendrive.parser.reference import ReferenceLine +from d123.dataset.conversion.map.opendrive.parser.road import Road +from d123.dataset.conversion.map.opendrive.utils.id_system import ( + build_lane_id, + derive_lane_section_id, + lane_group_id_from_lane_id, + road_id_from_lane_group_id, +) +from d123.dataset.conversion.map.opendrive.utils.lane_helper import ( + OpenDriveLaneGroupHelper, + OpenDriveLaneHelper, + lane_section_to_lane_helpers, +) +from d123.dataset.conversion.map.opendrive.utils.objects_helper import OpenDriveObjectHelper, get_object_helper + +logger = logging.getLogger(__name__) + + +def collect_element_helpers( + opendrive: OpenDrive, + interpolation_step_size: float, + connection_distance_threshold: float, +) -> Tuple[ + Dict[int, Road], + Dict[int, Junction], + Dict[str, OpenDriveLaneHelper], + Dict[str, OpenDriveLaneGroupHelper], + Dict[int, OpenDriveObjectHelper], +]: + + # 1. Fill the road and junction dictionaries + road_dict: Dict[int, Road] = {road.id: road for road in opendrive.roads} + junction_dict: Dict[int, Junction] = {junction.id: junction for junction in opendrive.junctions} + + # 2. Create lane helpers from the roads + lane_helper_dict: Dict[str, OpenDriveLaneHelper] = {} + for road in opendrive.roads: + reference_line = ReferenceLine.from_plan_view( + road.plan_view, + road.lanes.lane_offsets, + road.elevation_profile.elevations, + ) + lane_section_lengths: List[float] = [ls.s for ls in road.lanes.lane_sections] + [road.length] + for idx, lane_section in enumerate(road.lanes.lane_sections): + lane_section_id = derive_lane_section_id(road.id, idx) + lane_helpers_ = lane_section_to_lane_helpers( + lane_section_id, + lane_section, + reference_line, + lane_section_lengths[idx], + lane_section_lengths[idx + 1], + road.road_types, + interpolation_step_size, + ) + lane_helper_dict.update(lane_helpers_) + + # 3. Update the connections and fill the lane helpers: + # 3.1. From links of the roads + _update_connection_from_links(lane_helper_dict, road_dict) + # 3.2. From junctions + _update_connection_from_junctions(lane_helper_dict, junction_dict, road_dict) + # 3.3. Flip the connections to align to the lane direction + _flip_and_set_connections(lane_helper_dict) + # 3.4. Remove invalid connections based on centerline distances + _post_process_connections(lane_helper_dict, connection_distance_threshold) + + # 4. Collect lane groups from lane helpers + lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = _collect_lane_groups( + lane_helper_dict, junction_dict, road_dict + ) + + # 5. Collect objects, i.e. crosswalks + crosswalk_dict = _collect_crosswalks(opendrive) + + return (road_dict, junction_dict, lane_helper_dict, lane_group_helper_dict, crosswalk_dict) + + +def _update_connection_from_links(lane_helper_dict: Dict[str, OpenDriveLaneHelper], road_dict: Dict[int, Road]) -> None: + """ + Uses the links of the roads to update the connections between lane helpers. + :param lane_helper_dict: Dictionary of lane helpers indexed by lane id. + :param road_dict: Dictionary of roads indexed by road id. + """ + + for lane_id in lane_helper_dict.keys(): + road_idx, lane_section_idx, _, lane_idx = lane_id.split("_") + road_idx, lane_section_idx, lane_idx = int(road_idx), int(lane_section_idx), int(lane_idx) + + road = road_dict[road_idx] + + successor_lane_idx = lane_helper_dict[lane_id].open_drive_lane.successor + if successor_lane_idx is not None: + successor_lane_id: Optional[str] = None + + # Not the last lane section -> Next lane section in same road + if lane_section_idx < road.lanes.last_lane_section_idx: + successor_lane_id = build_lane_id( + road_idx, + lane_section_idx + 1, + successor_lane_idx, + ) + + # Last lane section -> Next road in first lane section + # Try to get next road + elif road.link.successor is not None and road.link.successor.element_type != "junction": + successor_road = road_dict[road.link.successor.element_id] + successor_lane_section_idx = ( + 0 if road.link.successor.contact_point == "start" else successor_road.lanes.last_lane_section_idx + ) + + successor_lane_id = build_lane_id( + successor_road.id, + successor_lane_section_idx, + successor_lane_idx, + ) + + # assert successor_lane_id in lane_helper_dict.keys() + if successor_lane_id is None or successor_lane_id not in lane_helper_dict.keys(): + continue + lane_helper_dict[lane_id].successor_lane_ids.append(successor_lane_id) + lane_helper_dict[successor_lane_id].predecessor_lane_ids.append(lane_id) + + predecessor_lane_idx = lane_helper_dict[lane_id].open_drive_lane.predecessor + if predecessor_lane_idx is not None: + predecessor_lane_id: Optional[str] = None + + # Not the first lane section -> Previous lane section in same road + if lane_section_idx > 0: + predecessor_lane_id = build_lane_id( + road_idx, + lane_section_idx - 1, + predecessor_lane_idx, + ) + + # First lane section -> Previous road + # Try to get previous road + elif road.link.predecessor is not None and road.link.predecessor.element_type != "junction": + predecessor_road = road_dict[road.link.predecessor.element_id] + predecessor_lane_section_idx = ( + 0 + if road.link.predecessor.contact_point == "start" + else predecessor_road.lanes.last_lane_section_idx + ) + + predecessor_lane_id = build_lane_id( + predecessor_road.id, + predecessor_lane_section_idx, + predecessor_lane_idx, + ) + + # assert predecessor_lane_id in lane_helper_dict.keys() + if predecessor_lane_id is None or predecessor_lane_id not in lane_helper_dict.keys(): + continue + lane_helper_dict[lane_id].predecessor_lane_ids.append(predecessor_lane_id) + lane_helper_dict[predecessor_lane_id].successor_lane_ids.append(lane_id) + + +def _update_connection_from_junctions( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], + junction_dict: Dict[int, Junction], + road_dict: Dict[int, Road], +) -> None: + """ + Helper function to update the lane connections based on junctions. + :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. + :param junction_dict: Dictionary mapping junction ids to their objects. + :param road_dict: Dictionary mapping road ids to their objects. + :raises ValueError: If a connection is invalid. + """ + + for junction_idx, junction in junction_dict.items(): + for connection in junction.connections: + incoming_road = road_dict[connection.incoming_road] + connecting_road = road_dict[connection.connecting_road] + + for lane_link in connection.lane_links: + + incoming_lane_id: Optional[str] = None + connecting_lane_id: Optional[str] = None + + if connection.contact_point == "start": + incoming_lane_section_idx = incoming_road.lanes.last_lane_section_idx if lane_link.start < 0 else 0 + incoming_lane_id = build_lane_id(incoming_road.id, incoming_lane_section_idx, lane_link.start) + connecting_lane_id = build_lane_id(connecting_road.id, 0, lane_link.end) + elif connection.contact_point == "end": + incoming_lane_id = build_lane_id(incoming_road.id, 0, lane_link.start) + connecting_lane_id = build_lane_id( + connecting_road.id, + connecting_road.lanes.last_lane_section_idx, + lane_link.end, + ) + else: + raise ValueError(f"Unknown contact point: {connection.contact_point} in junction {junction_idx}") + + if incoming_lane_id is None or connecting_lane_id is None: + logger.debug(f"OpenDRIVE: Lane connection {incoming_lane_id} -> {connecting_lane_id} not valid") + continue + if incoming_lane_id not in lane_helper_dict.keys() or connecting_lane_id not in lane_helper_dict.keys(): + logger.debug( + f"OpenDRIVE: Lane connection {incoming_lane_id} -> {connecting_lane_id} not found in lane_helper_dict" + ) + continue + lane_helper_dict[incoming_lane_id].successor_lane_ids.append(connecting_lane_id) + lane_helper_dict[connecting_lane_id].predecessor_lane_ids.append(incoming_lane_id) + + +def _flip_and_set_connections(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> None: + """ + Helper function to flip the connections of the lane helpers, to align them with the lane direction + :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. + """ + + for lane_id in lane_helper_dict.keys(): + if lane_helper_dict[lane_id].id > 0: + successors_temp = lane_helper_dict[lane_id].successor_lane_ids + lane_helper_dict[lane_id].successor_lane_ids = lane_helper_dict[lane_id].predecessor_lane_ids + lane_helper_dict[lane_id].predecessor_lane_ids = successors_temp + lane_helper_dict[lane_id].successor_lane_ids = list(set(lane_helper_dict[lane_id].successor_lane_ids)) + lane_helper_dict[lane_id].predecessor_lane_ids = list(set(lane_helper_dict[lane_id].predecessor_lane_ids)) + + +def _post_process_connections( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], + connection_distance_threshold: float, +) -> None: + """ + Helper function to post-process the connections of the lane helpers, removing invalid connections based on centerline distances. + :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. + :param connection_distance_threshold: Threshold distance for valid connections. + """ + + for lane_id in lane_helper_dict.keys(): + lane_helper_dict[lane_id] + + centerline = lane_helper_dict[lane_id].center_polyline_se2 + + valid_successor_lane_ids: List[str] = [] + for successor_lane_id in lane_helper_dict[lane_id].successor_lane_ids: + successor_centerline = lane_helper_dict[successor_lane_id].center_polyline_se2 + distance = np.linalg.norm(centerline[-1, :2] - successor_centerline[0, :2]) + if distance > connection_distance_threshold: + + logger.debug( + f"OpenDRIVE: Removing connection {lane_id} -> {successor_lane_id} with distance {distance}" + ) + else: + valid_successor_lane_ids.append(successor_lane_id) + lane_helper_dict[lane_id].successor_lane_ids = valid_successor_lane_ids + + valid_predecessor_lane_ids: List[str] = [] + for predecessor_lane_id in lane_helper_dict[lane_id].predecessor_lane_ids: + predecessor_centerline = lane_helper_dict[predecessor_lane_id].center_polyline_se2 + distance = np.linalg.norm(centerline[0, :2] - predecessor_centerline[-1, :2]) + if distance > connection_distance_threshold: + logger.debug( + f"OpenDRIVE: Removing connection {predecessor_lane_id} -> {successor_lane_id} with distance {distance}" + ) + else: + valid_predecessor_lane_ids.append(predecessor_lane_id) + lane_helper_dict[lane_id].predecessor_lane_ids = valid_predecessor_lane_ids + + +def _collect_lane_groups( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], + junction_dict: Dict[int, Junction], + road_dict: Dict[int, Road], +) -> None: + + lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = {} + + def _collect_lane_helper_of_id(lane_group_id: str) -> List[OpenDriveLaneHelper]: + lane_helpers: List[OpenDriveLaneHelper] = [] + for lane_id, lane_helper in lane_helper_dict.items(): + if (lane_helper.type in ["driving"]) and (lane_group_id_from_lane_id(lane_id) == lane_group_id): + lane_helpers.append(lane_helper) + return lane_helpers + + def _collect_lane_group_ids_of_road(road_id: int) -> List[str]: + lane_group_ids: List[str] = [] + for lane_group_id in lane_group_helper_dict.keys(): + if int(road_id_from_lane_group_id(lane_group_id)) == road_id: + lane_group_ids.append(lane_group_id) + return lane_group_ids + + all_lane_group_ids = list(set([lane_group_id_from_lane_id(lane_id) for lane_id in lane_helper_dict.keys()])) + + for lane_group_id in all_lane_group_ids: + lane_group_lane_helper = _collect_lane_helper_of_id(lane_group_id) + if len(lane_group_lane_helper) >= 1: + lane_group_helper_dict[lane_group_id] = OpenDriveLaneGroupHelper(lane_group_id, lane_group_lane_helper) + + for junction in junction_dict.values(): + for connection in junction.connections: + connecting_road = road_dict[connection.connecting_road] + connecting_lane_group_ids = _collect_lane_group_ids_of_road(connecting_road.id) + for connecting_lane_group_id in connecting_lane_group_ids: + if connecting_lane_group_id in lane_group_helper_dict.keys(): + lane_group_helper_dict[connecting_lane_group_id].junction_id = junction.id + + return lane_group_helper_dict + + +def _collect_crosswalks(opendrive: OpenDrive) -> Dict[int, OpenDriveObjectHelper]: + + object_helper_dict: Dict[int, OpenDriveObjectHelper] = {} + for road in opendrive.roads: + if len(road.objects) == 0: + continue + reference_line = ReferenceLine.from_plan_view( + road.plan_view, + road.lanes.lane_offsets, + road.elevation_profile.elevations, + ) + for object in road.objects: + if object.type in ["crosswalk"]: + object_helper = get_object_helper(object, reference_line) + object_helper_dict[object_helper.object_id] = object_helper + + return object_helper_dict diff --git a/d123/dataset/dataset_specific/carla/opendrive/id_mapping.py b/d123/dataset/conversion/map/opendrive/utils/id_mapping.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/id_mapping.py rename to d123/dataset/conversion/map/opendrive/utils/id_mapping.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/id_system.py b/d123/dataset/conversion/map/opendrive/utils/id_system.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/conversion/id_system.py rename to d123/dataset/conversion/map/opendrive/utils/id_system.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py b/d123/dataset/conversion/map/opendrive/utils/lane_helper.py similarity index 70% rename from d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py rename to d123/dataset/conversion/map/opendrive/utils/lane_helper.py index 845f1cfa..79e0f0c1 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py +++ b/d123/dataset/conversion/map/opendrive/utils/lane_helper.py @@ -8,19 +8,14 @@ from d123.common.geometry.base import StateSE2Index from d123.common.geometry.units import kmph_to_mps, mph_to_mps -from d123.dataset.dataset_specific.carla.opendrive.conversion.id_system import ( +from d123.dataset.conversion.map.opendrive.parser.lane import Lane, LaneSection +from d123.dataset.conversion.map.opendrive.parser.reference import ReferenceLine +from d123.dataset.conversion.map.opendrive.parser.road import RoadType +from d123.dataset.conversion.map.opendrive.utils.id_system import ( derive_lane_group_id, derive_lane_id, lane_group_id_from_lane_id, ) -from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lane, LaneSection -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border -from d123.dataset.dataset_specific.carla.opendrive.elements.road import RoadType - -# TODO: Add to config -STEP_SIZE = 1.0 - -# TODO: make naming consistent with objects_collections.py @dataclass @@ -30,10 +25,10 @@ class OpenDriveLaneHelper: open_drive_lane: Lane s_inner_offset: float s_range: Tuple[float, float] - inner_border: Border - outer_border: Border + inner_boundary: ReferenceLine + outer_boundary: ReferenceLine speed_limit_mps: Optional[float] - reverse: bool = False + interpolation_step_size: float # lazy loaded predecessor_lane_ids: Optional[List[str]] = None @@ -55,23 +50,27 @@ def type(self) -> str: def _s_positions(self) -> npt.NDArray[np.float64]: length = self.s_range[1] - self.s_range[0] _s_positions = np.linspace( - self.s_range[0], self.s_range[1], int(np.ceil(length / STEP_SIZE)) + 1, endpoint=True, dtype=np.float64 + self.s_range[0], + self.s_range[1], + int(np.ceil(length / self.interpolation_step_size)) + 1, + endpoint=True, + dtype=np.float64, ) _s_positions[..., -1] = np.clip(_s_positions[..., -1], 0.0, self.s_range[-1]) return _s_positions @cached_property - def _is_last_mask(self) -> npt.NDArray[np.float64]: - is_last_mask = np.zeros(len(self._s_positions), dtype=bool) - is_last_mask[-1] = True - return is_last_mask + def _lane_section_end_mask(self) -> npt.NDArray[np.float64]: + lane_section_end_mask = np.zeros(len(self._s_positions), dtype=bool) + lane_section_end_mask[-1] = True + return lane_section_end_mask @cached_property def inner_polyline_se2(self) -> npt.NDArray[np.float64]: inner_polyline = np.array( [ - self.inner_border.interpolate_se2(self.s_inner_offset + s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.inner_boundary.interpolate_se2(self.s_inner_offset + s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -81,8 +80,8 @@ def inner_polyline_se2(self) -> npt.NDArray[np.float64]: def inner_polyline_3d(self) -> npt.NDArray[np.float64]: inner_polyline = np.array( [ - self.inner_border.interpolate_3d(self.s_inner_offset + s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.inner_boundary.interpolate_3d(self.s_inner_offset + s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -92,8 +91,8 @@ def inner_polyline_3d(self) -> npt.NDArray[np.float64]: def outer_polyline_se2(self) -> npt.NDArray[np.float64]: outer_polyline = np.array( [ - self.outer_border.interpolate_se2(s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.outer_boundary.interpolate_se2(s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -103,8 +102,8 @@ def outer_polyline_se2(self) -> npt.NDArray[np.float64]: def outer_polyline_3d(self) -> npt.NDArray[np.float64]: outer_polyline = np.array( [ - self.outer_border.interpolate_3d(s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.outer_boundary.interpolate_3d(s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -126,7 +125,7 @@ def center_polyline_3d(self) -> npt.NDArray[np.float64]: def outline_polyline_3d(self) -> npt.NDArray[np.float64]: inner_polyline = self.inner_polyline_3d[::-1] outer_polyline = self.outer_polyline_3d - return np.concatenate([inner_polyline, outer_polyline], axis=0, dtype=np.float64) + return np.concatenate([inner_polyline, outer_polyline, inner_polyline[None, 0]], axis=0, dtype=np.float64) @property def shapely_polygon(self) -> shapely.Polygon: @@ -199,75 +198,53 @@ def shapely_polygon(self) -> shapely.Polygon: def lane_section_to_lane_helpers( lane_section_id: str, lane_section: LaneSection, - reference_border: Border, + reference_line: ReferenceLine, s_min: float, s_max: float, road_types: List[RoadType], + interpolation_step_size: float, ) -> Dict[str, OpenDriveLaneHelper]: lane_helpers: Dict[str, OpenDriveLaneHelper] = {} - for side in ["right", "left"]: + for lanes, t_sign, side in zip([lane_section.left_lanes, lane_section.right_lanes], [1.0, -1.0], ["left", "right"]): lane_group_id = derive_lane_group_id(lane_section_id, side) - lanes = lane_section.right_lanes if side == "right" else lane_section.left_lanes - coeff_factor = -1.0 if side == "right" else 1.0 - - lane_borders = [reference_border] - + lane_boundaries = [reference_line] for lane in lanes: lane_id = derive_lane_id(lane_group_id, lane.id) - s_inner_offset = lane_section.s if len(lane_borders) == 1 else 0.0 - lane_borders.append(_create_outer_lane_border(lane_borders, lane_section, lane, coeff_factor)) + s_inner_offset = lane_section.s if len(lane_boundaries) == 1 else 0.0 + lane_boundaries.append( + ReferenceLine.from_reference_line( + reference_line=lane_boundaries[-1], + widths=lane.widths, + s_offset=s_inner_offset, + t_sign=t_sign, + ) + ) lane_helper = OpenDriveLaneHelper( lane_id=lane_id, open_drive_lane=lane, s_inner_offset=s_inner_offset, s_range=(s_min, s_max), - inner_border=lane_borders[-2], - outer_border=lane_borders[-1], - speed_limit_mps=_get_speed_limit_mps(s_min, s_max, road_types), + inner_boundary=lane_boundaries[-2], + outer_boundary=lane_boundaries[-1], + speed_limit_mps=_get_speed_limit_mps(s_min, road_types), + interpolation_step_size=interpolation_step_size, ) lane_helpers[lane_id] = lane_helper return lane_helpers -def _create_outer_lane_border( - lane_borders: List[Border], - lane_section: LaneSection, - lane: Lane, - coeff_factor: float, -) -> Border: - - args = {} - if len(lane_borders) == 1: - args["s_offset"] = lane_section.s - - args["reference"] = lane_borders[-1] - args["elevation_profile"] = lane_borders[-1].elevation_profile - - width_coefficient_offsets = [] - width_coefficients = [] - - for width in lane.widths: - width_coefficient_offsets.append(width.s_offset) - width_coefficients.append([x * coeff_factor for x in width.polynomial_coefficients]) - - args["width_coefficient_offsets"] = width_coefficient_offsets - args["width_coefficients"] = width_coefficients - return Border(**args) - - -def _get_speed_limit_mps(s_min: float, s_max: float, road_types: List[RoadType]) -> Optional[float]: +def _get_speed_limit_mps(s: float, road_types: List[RoadType]) -> Optional[float]: # NOTE: Likely not correct way to extract speed limit from CARLA maps, but serves as a placeholder speed_limit_mps: Optional[float] = None s_road_types = [road_type.s for road_type in road_types] + [float("inf")] if len(road_types) > 0: - # 1. Find current road type for road_type_idx, road_type in enumerate(road_types): - if s_min >= road_type.s and s_min < s_road_types[road_type_idx + 1]: + if s >= road_type.s and s < s_road_types[road_type_idx + 1]: if road_type.speed is not None: if road_type.speed.unit == "mps": speed_limit_mps = road_type.speed.max diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/objects_collections.py b/d123/dataset/conversion/map/opendrive/utils/objects_helper.py similarity index 85% rename from d123/dataset/dataset_specific/carla/opendrive/conversion/objects_collections.py rename to d123/dataset/conversion/map/opendrive/utils/objects_helper.py index d94d3eaf..3cbb569e 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/conversion/objects_collections.py +++ b/d123/dataset/conversion/map/opendrive/utils/objects_helper.py @@ -8,8 +8,8 @@ from d123.common.geometry.base import Point2D, Point3D, Point3DIndex, StateSE2 from d123.common.geometry.transform.tranform_2d import translate_along_yaw from d123.common.geometry.utils import normalize_angle -from d123.dataset.dataset_specific.carla.opendrive.elements.objects import Object -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border +from d123.dataset.conversion.map.opendrive.parser.objects import Object +from d123.dataset.conversion.map.opendrive.parser.reference import ReferenceLine # TODO: make naming consistent with group_collections.py @@ -29,14 +29,14 @@ def shapely_polygon(self) -> shapely.Polygon: return shapely.geometry.Polygon(self.outline_3d[:, Point3DIndex.XY]) -def get_object_helper(object: Object, reference_border: Border) -> OpenDriveObjectHelper: +def get_object_helper(object: Object, reference_line: ReferenceLine) -> OpenDriveObjectHelper: object_helper: Optional[OpenDriveObjectHelper] = None # 1. Extract object position in frenet frame of the reference line - object_se2: StateSE2 = StateSE2.from_array(reference_border.interpolate_se2(s=object.s, t=object.t)) - object_3d: Point3D = Point3D.from_array(reference_border.interpolate_3d(s=object.s, t=object.t)) + object_se2: StateSE2 = StateSE2.from_array(reference_line.interpolate_se2(s=object.s, t=object.t)) + object_3d: Point3D = Point3D.from_array(reference_line.interpolate_3d(s=object.s, t=object.t)) # Adjust yaw angle from object data object_se2.yaw = normalize_angle(object_se2.yaw + object.hdg) diff --git a/d123/dataset/conversion/map/road_edge/__init__.py b/d123/dataset/conversion/map/road_edge/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py b/d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py new file mode 100644 index 00000000..f0eb580d --- /dev/null +++ b/d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py @@ -0,0 +1,60 @@ +from typing import Final, List, Union + +import numpy as np +import shapely +from shapely import LinearRing, LineString, Polygon, union_all + +ROAD_EDGE_BUFFER: Final[float] = 0.05 + + +def get_road_edge_linear_rings(drivable_polygons: List[Polygon]) -> List[LinearRing]: + + def _polygon_to_linear_rings(polygon: Polygon) -> List[LinearRing]: + assert polygon.geom_type == "Polygon" + linear_ring_list = [] + linear_ring_list.append(polygon.exterior) + for interior in polygon.interiors: + linear_ring_list.append(interior) + return linear_ring_list + + union_polygon = union_all([polygon.buffer(ROAD_EDGE_BUFFER, join_style=2) for polygon in drivable_polygons]).buffer( + -ROAD_EDGE_BUFFER, join_style=2 + ) + + linear_ring_list = [] + if union_polygon.geom_type == "Polygon": + for polyline in _polygon_to_linear_rings(union_polygon): + linear_ring_list.append(LinearRing(polyline)) + elif union_polygon.geom_type == "MultiPolygon": + for polygon in union_polygon.geoms: + for polyline in _polygon_to_linear_rings(polygon): + linear_ring_list.append(LinearRing(polyline)) + + return linear_ring_list + + +def split_line_geometry_by_max_length( + geometries: List[Union[LineString, LinearRing]], + max_length_meters: float, +) -> List[LineString]: + # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. + + if not isinstance(geometries, list): + geometries = [geometries] + + all_segments = [] + for geom in geometries: + if geom.length <= max_length_meters: + all_segments.append(LineString(geom.coords)) + continue + + num_segments = int(np.ceil(geom.length / max_length_meters)) + segment_length = geom.length / num_segments + + for i in range(num_segments): + start_dist = i * segment_length + end_dist = min((i + 1) * segment_length, geom.length) + segment = shapely.ops.substring(geom, start_dist, end_dist) + all_segments.append(segment) + + return all_segments diff --git a/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py b/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py new file mode 100644 index 00000000..93a34526 --- /dev/null +++ b/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py @@ -0,0 +1,298 @@ +import logging +from collections import defaultdict +from typing import Dict, List, Set + +import geopandas as gpd +import networkx as nx +import numpy as np +import numpy.typing as npt +import shapely +from shapely.geometry import LineString + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings + +logger = logging.getLogger(__name__) + + +def get_road_edges_3d_from_generic_drivable_area_df(generic_drivable_area_df: gpd.GeoDataFrame) -> List[LineString]: + """ + Extracts 3D road edges from the generic drivable area GeoDataFrame. + """ + # NOTE: this is a simplified version that assumes the generic drivable area covers all areas. + # This is the case for argoverse 2. + road_edges_2d = get_road_edge_linear_rings(generic_drivable_area_df.geometry.tolist()) + outlines = generic_drivable_area_df.outline.tolist() + non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, outlines) + return non_conflicting_road_edges + + +def get_road_edges_3d_from_gdf( + lane_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> List[LineString]: + + # 1. Find conflicting lane groups, e.g. groups of lanes that overlap in 2D but have different Z-values (bridges) + conflicting_lane_groups = _get_conflicting_lane_groups(lane_group_df, lane_df) + + # 2. Extract road edges in 2D (including conflicting lane groups) + drivable_polygons = ( + lane_group_df.geometry.tolist() + carpark_df.geometry.tolist() + generic_drivable_area_df.geometry.tolist() + ) + road_edges_2d = get_road_edge_linear_rings(drivable_polygons) + + # 3. Collect 3D boundaries of non-conflicting lane groups and other drivable areas + non_conflicting_boundaries: List[LineString] = [] + for lane_group_id, lane_group_helper in lane_group_df.iterrows(): + if lane_group_id not in conflicting_lane_groups.keys(): + non_conflicting_boundaries.append(lane_group_helper["left_boundary"]) + non_conflicting_boundaries.append(lane_group_helper["right_boundary"]) + for outline in carpark_df.outline.tolist() + generic_drivable_area_df.outline.tolist(): + non_conflicting_boundaries.append(outline) + + # 4. Lift road edges to 3D using the boundaries of non-conflicting elements + non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, non_conflicting_boundaries) + + # 5. Add road edges from conflicting lane groups + resolved_road_edges = _resolve_conflicting_lane_groups(conflicting_lane_groups, lane_group_df) + + all_road_edges = non_conflicting_road_edges + resolved_road_edges + + return all_road_edges + + +def _get_conflicting_lane_groups(lane_group_df: gpd.GeoDataFrame, lane_df: gpd.GeoDataFrame) -> Dict[int, List[int]]: + """ + Even more optimized version using vectorized operations where possible. + """ + Z_THRESHOLD = 5.0 # [m] Z-value threshold for conflict detection + + # Convert to regular dictionaries for faster access + lane_group_dict = lane_group_df.set_index("id").to_dict("index") + lane_baseline_dict = dict(zip(lane_df.id.values, lane_df.baseline_path.values)) + + # Pre-compute all centerlines + centerlines_cache = {} + polygons = [] + ids = [] + + for lane_group_id, data in lane_group_dict.items(): + geometry = data["geometry"] + lane_ids = data["lane_ids"] + + # Vectorized centerline computation + centerlines = [np.array(lane_baseline_dict[lane_id].coords, dtype=np.float64) for lane_id in lane_ids] + centerlines_3d = np.concatenate(centerlines, axis=0) + + centerlines_cache[lane_group_id] = centerlines_3d + polygons.append(geometry) + ids.append(lane_group_id) + + occupancy_map = OccupancyMap2D(polygons, ids) + conflicting_lane_groups: Dict[int, List[int]] = defaultdict(list) + processed_pairs = set() + + for i, lane_group_id in enumerate(ids): + lane_group_polygon = polygons[i] + lane_group_centerlines = centerlines_cache[lane_group_id] + + # Get all intersecting geometries at once + intersecting_ids = occupancy_map.intersects(lane_group_polygon) + intersecting_ids.remove(lane_group_id) + + for intersecting_id in intersecting_ids: + pair_key = tuple(sorted([lane_group_id, intersecting_id])) + if pair_key in processed_pairs: + continue + processed_pairs.add(pair_key) + + intersecting_geometry = occupancy_map[intersecting_id] + if intersecting_geometry.geom_type != "Polygon": + continue + try: + # Compute actual intersection for better centroid + intersection = lane_group_polygon.intersection(intersecting_geometry) + except shapely.errors.GEOSException as e: + logger.debug(f"Error computing intersection for {pair_key}: {e}") + continue + + if intersection.is_empty: + continue + + intersection_centroid = np.array(intersection.centroid.coords[0], dtype=np.float64) + intersecting_centerlines = centerlines_cache[intersecting_id] + + z_at_intersecting = _get_nearest_z_from_points_3d(intersecting_centerlines, intersection_centroid) + z_at_lane_group = _get_nearest_z_from_points_3d(lane_group_centerlines, intersection_centroid) + + if np.abs(z_at_lane_group - z_at_intersecting) >= Z_THRESHOLD: + conflicting_lane_groups[lane_group_id].append(intersecting_id) + conflicting_lane_groups[intersecting_id].append(lane_group_id) + + return conflicting_lane_groups + + +def lift_road_edges_to_3d( + road_edges_2d: List[shapely.LinearRing], + boundaries: List[LineString], + max_distance: float = 0.01, +) -> List[LineString]: + """ + Even faster version using batch processing and optimized data structures. + """ + if not road_edges_2d or not boundaries: + return [] + + # 1. Build comprehensive spatial index with all boundary segments + boundary_segments = [] + + for boundary_idx, boundary in enumerate(boundaries): + coords = np.array(boundary.coords, dtype=np.float64).reshape(-1, 1, 3) + segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1) + boundary_segments.append(segment_coords_boundary) + + boundary_segments = np.concatenate(boundary_segments, axis=0) + boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments) + + occupancy_map = OccupancyMap2D(boundary_segment_linestrings) + + road_edges_3d = [] + for linear_ring in road_edges_2d: + points_2d = np.array(linear_ring.coords, dtype=np.float64) + points_3d = np.zeros((len(points_2d), 3), dtype=np.float64) + points_3d[:, :2] = points_2d + + # 3. Batch query for all points + query_points = shapely.creation.points(points_2d) + results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True) + + for query_idx, geometry_idx in zip(*results): + query_point = query_points[query_idx] + segment_coords = boundary_segments[geometry_idx] + best_z = _interpolate_z_on_segment(query_point, segment_coords) + points_3d[query_idx, 2] = best_z + + continuous_segments = _find_continuous_segments(np.array(results[0])) + + for segment_indices in continuous_segments: + if len(segment_indices) >= 2: + segment_points = points_3d[segment_indices] + road_edges_3d.append(LineString(segment_points)) + + return road_edges_3d + + +def _get_nearest_z_from_points_3d(points_3d: npt.NDArray[np.float64], query_point: npt.NDArray[np.float64]) -> float: + assert points_3d.ndim == 2 and points_3d.shape[1] == len( + Point3DIndex + ), "points_3d must be a 2D array with shape (N, 3)" + distances = np.linalg.norm(points_3d[..., Point3DIndex.XY] - query_point[..., Point3DIndex.XY], axis=1) + closest_point = points_3d[np.argmin(distances)] + return closest_point[2] + + +def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float: + """Interpolate Z coordinate along a 3D line segment.""" + p1, p2 = segment_coords[0], segment_coords[1] + + # Project point onto segment + segment_vec = p2[:2] - p1[:2] + point_vec = np.array([point.x, point.y]) - p1[:2] + + # Handle degenerate case + segment_length_sq = np.dot(segment_vec, segment_vec) + if segment_length_sq == 0: + return p1[2] + + # Calculate projection parameter + t = np.dot(point_vec, segment_vec) / segment_length_sq + t = np.clip(t, 0, 1) # Clamp to segment bounds + + # Interpolate Z + return p1[2] + t * (p2[2] - p1[2]) + + +def _find_continuous_segments(indices: np.ndarray) -> List[np.ndarray]: + """Vectorized version of finding continuous segments.""" + if len(indices) == 0: + return [] + + # Find breaks in continuity + breaks = np.where(np.diff(indices) != 1)[0] + 1 + segments = np.split(indices, breaks) + + # Filter segments with at least 2 points + return [seg for seg in segments if len(seg) >= 2] + + +def _resolve_conflicting_lane_groups( + conflicting_lane_groups: Dict[int, List[int]], lane_group_df: gpd.GeoDataFrame +) -> List[LineString]: + + # Split conflicting lane groups into non-conflicting sets for further merging + non_conflicting_sets = _create_non_conflicting_sets(conflicting_lane_groups) + + road_edges_3d: List[LineString] = [] + for non_conflicting_set in non_conflicting_sets: + + # Collect 2D polygons of non-conflicting lane group set + set_lane_group_rows = lane_group_df[lane_group_df.id.isin(non_conflicting_set)] + connected_lane_group = [] + for row in set_lane_group_rows.itertuples(): + connected_lane_group.extend(row.predecessor_ids) + connected_lane_group.extend(row.successor_ids) + connected_lane_group_rows = lane_group_df[lane_group_df.id.isin(connected_lane_group)] + + set_polygons = set_lane_group_rows.geometry.tolist() + connected_lane_group_rows.geometry.tolist() + + # Get 2D road edge linestrings for the non-conflicting set + set_road_edges_2d = get_road_edge_linear_rings(set_polygons) + + # Collect 3D boundaries of non-conflicting lane groups + set_boundaries_3d: List[LineString] = [] + for lane_group_id in non_conflicting_set: + lane_group_helper = lane_group_df[lane_group_df.id == lane_group_id] + set_boundaries_3d.append(lane_group_helper.left_boundary.values[0]) + set_boundaries_3d.append(lane_group_helper.right_boundary.values[0]) + + # Lift road edges to 3D using the boundaries of non-conflicting lane groups + lifted_road_edges_3d = lift_road_edges_to_3d(set_road_edges_2d, set_boundaries_3d) + road_edges_3d.extend(lifted_road_edges_3d) + + return road_edges_3d + + +def _create_non_conflicting_sets(conflicts: Dict[int, List[int]]) -> List[Set[int]]: + """ + Creates sets of non-conflicting indices using NetworkX. + """ + # Create graph from conflicts + G = nx.Graph() + for idx, conflict_list in conflicts.items(): + for conflict_idx in conflict_list: + G.add_edge(idx, conflict_idx) + + result = [] + + # Process each connected component + for component in nx.connected_components(G): + subgraph = G.subgraph(component) + + # Try bipartite coloring first (most common case) + if nx.is_bipartite(subgraph): + sets = nx.bipartite.sets(subgraph) + result.extend([set(s) for s in sets]) + else: + # Fall back to greedy coloring for non-bipartite graphs + coloring = nx.greedy_color(subgraph, strategy="largest_first") + color_groups = {} + for node, color in coloring.items(): + if color not in color_groups: + color_groups[color] = set() + color_groups[color].add(node) + result.extend(color_groups.values()) + + return result diff --git a/d123/dataset/dataset_specific/av2/av2_constants.py b/d123/dataset/dataset_specific/av2/av2_constants.py new file mode 100644 index 00000000..836bb12c --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_constants.py @@ -0,0 +1,107 @@ +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import CameraType +from d123.common.utils.enums import SerialIntEnum +from d123.dataset.maps.map_datatypes import RoadLineType + + +class AV2SensorBoxDetectionType(SerialIntEnum): + """Sensor dataset annotation categories.""" + + ANIMAL = 1 + ARTICULATED_BUS = 2 + BICYCLE = 3 + BICYCLIST = 4 + BOLLARD = 5 + BOX_TRUCK = 6 + BUS = 7 + CONSTRUCTION_BARREL = 8 + CONSTRUCTION_CONE = 9 + DOG = 10 + LARGE_VEHICLE = 11 + MESSAGE_BOARD_TRAILER = 12 + MOBILE_PEDESTRIAN_CROSSING_SIGN = 13 + MOTORCYCLE = 14 + MOTORCYCLIST = 15 + OFFICIAL_SIGNALER = 16 + PEDESTRIAN = 17 + RAILED_VEHICLE = 18 + REGULAR_VEHICLE = 19 + SCHOOL_BUS = 20 + SIGN = 21 + STOP_SIGN = 22 + STROLLER = 23 + TRAFFIC_LIGHT_TRAILER = 24 + TRUCK = 25 + TRUCK_CAB = 26 + VEHICULAR_TRAILER = 27 + WHEELCHAIR = 28 + WHEELED_DEVICE = 29 + WHEELED_RIDER = 30 + + +# Mapping from AV2SensorBoxDetectionType to general DetectionType +# TODO: Change the detection types. Multiple mistakes, e.g. animals/dogs are not generic objects. +AV2_TO_DETECTION_TYPE = { + AV2SensorBoxDetectionType.ANIMAL: DetectionType.GENERIC_OBJECT, + AV2SensorBoxDetectionType.ARTICULATED_BUS: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.BICYCLE: DetectionType.BICYCLE, + AV2SensorBoxDetectionType.BICYCLIST: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.BOLLARD: DetectionType.BARRIER, + AV2SensorBoxDetectionType.BOX_TRUCK: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.BUS: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.CONSTRUCTION_BARREL: DetectionType.BARRIER, + AV2SensorBoxDetectionType.CONSTRUCTION_CONE: DetectionType.TRAFFIC_CONE, + AV2SensorBoxDetectionType.DOG: DetectionType.GENERIC_OBJECT, + AV2SensorBoxDetectionType.LARGE_VEHICLE: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.MESSAGE_BOARD_TRAILER: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.MOBILE_PEDESTRIAN_CROSSING_SIGN: DetectionType.CZONE_SIGN, + AV2SensorBoxDetectionType.MOTORCYCLE: DetectionType.BICYCLE, + AV2SensorBoxDetectionType.MOTORCYCLIST: DetectionType.BICYCLE, + AV2SensorBoxDetectionType.OFFICIAL_SIGNALER: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.PEDESTRIAN: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.RAILED_VEHICLE: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.REGULAR_VEHICLE: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.SCHOOL_BUS: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.SIGN: DetectionType.SIGN, + AV2SensorBoxDetectionType.STOP_SIGN: DetectionType.SIGN, + AV2SensorBoxDetectionType.STROLLER: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.TRAFFIC_LIGHT_TRAILER: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.TRUCK: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.TRUCK_CAB: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.VEHICULAR_TRAILER: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.WHEELCHAIR: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.WHEELED_DEVICE: DetectionType.GENERIC_OBJECT, + AV2SensorBoxDetectionType.WHEELED_RIDER: DetectionType.BICYCLE, +} + + +AV2_CAMERA_TYPE_MAPPING = { + "ring_front_center": CameraType.CAM_F0, + "ring_front_left": CameraType.CAM_L0, + "ring_front_right": CameraType.CAM_R0, + "ring_side_left": CameraType.CAM_L1, + "ring_side_right": CameraType.CAM_R1, + "ring_rear_left": CameraType.CAM_L2, + "ring_rear_right": CameraType.CAM_R2, + "stereo_front_left": CameraType.CAM_STEREO_L, + "stereo_front_right": CameraType.CAM_STEREO_R, +} + + +AV2_ROAD_LINE_TYPE_MAPPING = { + "NONE": RoadLineType.NONE, + "UNKNOWN": RoadLineType.UNKNOWN, + "DASH_SOLID_YELLOW": RoadLineType.DASH_SOLID_YELLOW, + "DASH_SOLID_WHITE": RoadLineType.DASH_SOLID_WHITE, + "DASHED_WHITE": RoadLineType.DASHED_WHITE, + "DASHED_YELLOW": RoadLineType.DASHED_YELLOW, + "DOUBLE_SOLID_YELLOW": RoadLineType.DOUBLE_SOLID_YELLOW, + "DOUBLE_SOLID_WHITE": RoadLineType.DOUBLE_SOLID_WHITE, + "DOUBLE_DASH_YELLOW": RoadLineType.DOUBLE_DASH_YELLOW, + "DOUBLE_DASH_WHITE": RoadLineType.DOUBLE_DASH_WHITE, + "SOLID_YELLOW": RoadLineType.SOLID_YELLOW, + "SOLID_WHITE": RoadLineType.SOLID_WHITE, + "SOLID_DASH_WHITE": RoadLineType.SOLID_DASH_WHITE, + "SOLID_DASH_YELLOW": RoadLineType.SOLID_DASH_YELLOW, + "SOLID_BLUE": RoadLineType.SOLID_BLUE, +} diff --git a/d123/dataset/dataset_specific/av2/av2_data_converter.py b/d123/dataset/dataset_specific/av2/av2_data_converter.py new file mode 100644 index 00000000..0ff84a7e --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_data_converter.py @@ -0,0 +1,551 @@ +import gc +import hashlib +import json +from dataclasses import asdict +from functools import partial +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +import pyarrow as pa +from pyquaternion import Quaternion + +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.time.time_point import TimePoint +from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index +from d123.common.datatypes.vehicle_state.vehicle_parameters import ( + get_av2_ford_fusion_hybrid_parameters, + rear_axle_se3_to_center_se3, +) +from d123.common.geometry.base import StateSE3 +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index +from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array, get_rotation_matrix +from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map +from d123.dataset.dataset_specific.av2.av2_constants import ( + AV2_CAMERA_TYPE_MAPPING, + AV2_TO_DETECTION_TYPE, + AV2SensorBoxDetectionType, +) +from d123.dataset.dataset_specific.av2.av2_helper import ( + build_sensor_dataframe, + build_synchronization_dataframe, + find_closest_target_fpath, + get_slice_with_timestamp_ns, +) +from d123.dataset.dataset_specific.av2.av2_map_conversion import convert_av2_map +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.logs.log_metadata import LogMetadata + + +def create_token(input_data: str) -> str: + # TODO: Refactor this function. + # TODO: Add a general function to create tokens from arbitrary data. + if isinstance(input_data, str): + input_data = input_data.encode("utf-8") + + hash_obj = hashlib.sha256(input_data) + return hash_obj.hexdigest()[:16] + + +class AV2SensorDataConverter(RawDataConverter): + def __init__( + self, + splits: List[str], + log_path: Union[Path, str], + data_converter_config: DataConverterConfig, + ) -> None: + super().__init__(data_converter_config) + for split in splits: + assert ( + split in self.get_available_splits() + ), f"Split {split} is not available. Available splits: {self.available_splits}" + + self._splits: List[str] = splits + self._data_root: Path = Path(log_path) + self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() + self._target_dt: float = 0.1 + + def _collect_log_paths(self) -> Dict[str, List[Path]]: + log_paths_per_split: Dict[str, List[Path]] = {} + + for split in self._splits: + subsplit = split.split("_")[-1] + assert subsplit in ["train", "val", "test"] + + if "av2_sensor" in split: + log_folder = self._data_root / "sensor" / subsplit + elif "av2_lidar" in split: + log_folder = self._data_root / "lidar" / subsplit + elif "av2_motion" in split: + log_folder = self._data_root / "motion-forecasting" / subsplit + elif "av2-sensor-mini" in split: + log_folder = self._data_root / "sensor_mini" / subsplit + + log_paths_per_split[split] = list(log_folder.iterdir()) + + return log_paths_per_split + + def get_available_splits(self) -> List[str]: + return [ + "av2-sensor_train", + "av2-sensor_val", + "av2-sensor_test", + "av2-sensor-mini_train", + "av2-sensor-mini_val", + "av2-sensor-mini_test", + ] + + def convert_maps(self, worker: WorkerPool) -> None: + log_args = [ + { + "log_path": log_path, + "split": split, + } + for split, log_paths in self._log_paths_per_split.items() + for log_path in log_paths + ] + worker_map( + worker, + partial(convert_av2_map_to_gpkg, data_converter_config=self.data_converter_config), + log_args, + ) + + def convert_logs(self, worker: WorkerPool) -> None: + log_args = [ + { + "log_path": log_path, + "split": split, + } + for split, log_paths in self._log_paths_per_split.items() + for log_path in log_paths + ] + + worker_map( + worker, + partial( + convert_av2_log_to_arrow, + data_converter_config=self.data_converter_config, + ), + log_args, + ) + + +def convert_av2_map_to_gpkg( + args: List[Dict[str, Union[List[str], List[Path]]]], + data_converter_config: DataConverterConfig, +) -> List[Any]: + for log_info in args: + source_log_path: Path = log_info["log_path"] + split: str = log_info["split"] + + source_log_name = source_log_path.name + + map_path = data_converter_config.output_path / "maps" / split / f"{source_log_name}.gpkg" + if data_converter_config.force_map_conversion or not map_path.exists(): + map_path.unlink(missing_ok=True) + convert_av2_map(source_log_path, map_path) + return [] + + +def convert_av2_log_to_arrow( + args: List[Dict[str, Union[List[str], List[Path]]]], + data_converter_config: DataConverterConfig, +) -> List[Any]: + for log_info in args: + log_path: Path = log_info["log_path"] + split: str = log_info["split"] + + if not log_path.exists(): + raise FileNotFoundError(f"Log path {log_path} does not exist.") + + log_file_path = data_converter_config.output_path / split / f"{log_path.stem}.arrow" + + if data_converter_config.force_log_conversion or not log_file_path.exists(): + log_file_path.unlink(missing_ok=True) + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + sensor_df = build_sensor_dataframe(log_path) + synchronization_df = build_synchronization_dataframe(sensor_df) + + metadata = LogMetadata( + dataset="av2-sensor", + log_name=log_path.name, + location=None, + timestep_seconds=0.1, # TODO: verify this + map_has_z=True, + ) + vehicle_parameters = get_av2_ford_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters + camera_metadata = get_av2_camera_metadata(log_path) + lidar_metadata = get_av2_lidar_metadata(log_path) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + schema_column_list.append((lidar_type.serialize(), pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + if data_converter_config.camera_store_option is not None: + for camera_type in camera_metadata.keys(): + if data_converter_config.camera_store_option == "path": + schema_column_list.append((camera_type.serialize(), pa.string())) + + elif data_converter_config.camera_store_option == "binary": + schema_column_list.append((camera_type.serialize(), pa.binary())) + + schema_column_list.append((f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4))) + + recording_schema = pa.schema(schema_column_list) + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), + } + ) + + _write_recording_table( + sensor_df, + synchronization_df, + recording_schema, + log_file_path, + log_path, + data_converter_config, + ) + del recording_schema, vehicle_parameters + gc.collect() + return [] + + +def get_av2_camera_metadata(log_path: Path) -> Dict[CameraType, CameraMetadata]: + + intrinsics_file = log_path / "calibration" / "intrinsics.feather" + intrinsics_df = pd.read_feather(intrinsics_file) + + camera_metadata: Dict[CameraType, CameraMetadata] = {} + for _, row in intrinsics_df.iterrows(): + row = row.to_dict() + + camera_type = AV2_CAMERA_TYPE_MAPPING[row["sensor_name"]] + camera_metadata[camera_type] = CameraMetadata( + camera_type=camera_type, + width=row["width_px"], + height=row["height_px"], + intrinsic=np.array( + [[row["fx_px"], 0, row["cx_px"]], [0, row["fy_px"], row["cy_px"]], [0, 0, 1]], dtype=np.float64 + ), + distortion=np.array([row["k1"], row["k2"], row["k3"], 0, 0], dtype=np.float64), + ) + + return camera_metadata + + +def get_av2_lidar_metadata(log_path: Path) -> Dict[LiDARType, LiDARMetadata]: + # metadata: Dict[LiDARType, LiDARMetadata] = {} + # metadata[LiDARType.LIDAR_MERGED] = LiDARMetadata( + # lidar_type=LiDARType.LIDAR_MERGED, + # lidar_index=NuplanLidarIndex, + # extrinsic=None, # NOTE: LiDAR extrinsic are unknown + # ) + # return metadata + return {} + + +def _write_recording_table( + sensor_df: pd.DataFrame, + synchronization_df: pd.DataFrame, + recording_schema: pa.schema, + log_file_path: Path, + source_log_path: Path, + data_converter_config: DataConverterConfig, +) -> None: + + # NOTE: Similar to other datasets, we use the lidar timestamps as reference timestamps. + lidar_sensor = sensor_df.xs(key="lidar", level=2) + lidar_timestamps_ns = np.sort([int(idx_tuple[2]) for idx_tuple in lidar_sensor.index]) + + # NOTE: The annotation dataframe is not available for the test split. + annotations_df = ( + pd.read_feather(source_log_path / "annotations.feather") + if (source_log_path / "annotations.feather").exists() + else None + ) + + city_se3_egovehicle_df = pd.read_feather(source_log_path / "city_SE3_egovehicle.feather") + + egovehicle_se3_sensor_df = ( + pd.read_feather(source_log_path / "calibration" / "egovehicle_SE3_sensor.feather") + if data_converter_config.camera_store_option is not None + else None + ) + + # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: + with pa.OSFile(str(log_file_path), "wb") as sink: + with pa.ipc.new_file(sink, recording_schema) as writer: + + for lidar_timestamp_ns in lidar_timestamps_ns: + + ego_state_se3 = _extract_ego_state(city_se3_egovehicle_df, lidar_timestamp_ns) + ( + detections_state, + detections_velocity, + detections_token, + detections_types, + ) = _extract_box_detections(annotations_df, lidar_timestamp_ns, ego_state_se3) + traffic_light_ids, traffic_light_types = _extract_traffic_lights() + route_lane_group_ids = None # TODO: Add route lane group ids extraction ? + row_data = { + "token": [create_token(str(lidar_timestamp_ns))], + "timestamp": [TimePoint.from_ns(int(lidar_timestamp_ns)).time_us], + "detections_state": [detections_state], + "detections_velocity": [detections_velocity], + "detections_token": [detections_token], + "detections_type": [detections_types], + "ego_states": [ego_state_se3.array.tolist()], + "traffic_light_ids": [traffic_light_ids], + "traffic_light_types": [traffic_light_types], + "scenario_tag": [_extract_scenario_tag()], + "route_lane_group_ids": [route_lane_group_ids], + } + + # TODO: add lidar data + + # if data_converter_config.lidar_store_option is not None: + # lidar_data_dict = _extract_lidar(lidar_pc, data_converter_config) + # for lidar_type, lidar_data in lidar_data_dict.items(): + # if lidar_data is not None: + # row_data[lidar_type.serialize()] = [lidar_data] + # else: + # row_data[lidar_type.serialize()] = [None] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_camera( + lidar_timestamp_ns, + city_se3_egovehicle_df, + egovehicle_se3_sensor_df, + ego_state_se3, + synchronization_df, + source_log_path, + data_converter_config, + ) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] + else: + row_data[camera_type.serialize()] = [None] + row_data[f"{camera_type.serialize()}_extrinsic"] = [None] + + batch = pa.record_batch(row_data, schema=recording_schema) + writer.write_batch(batch) + del batch, row_data, detections_state, detections_velocity, detections_token, detections_types + + +def _extract_box_detections( + annotations_df: Optional[pd.DataFrame], + lidar_timestamp_ns: int, + ego_state_se3: EgoStateSE3, +) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: + + # TODO: Extract velocity from annotations_df if available. + + if annotations_df is None: + return [], [], [], [] + + annotations_slice = get_slice_with_timestamp_ns(annotations_df, lidar_timestamp_ns) + num_detections = len(annotations_slice) + + detections_state = np.zeros((num_detections, len(BoundingBoxSE3Index)), dtype=np.float64) + detections_velocity = np.zeros((num_detections, len(Vector3DIndex)), dtype=np.float64) + detections_token: List[str] = annotations_slice["track_uuid"].tolist() + detections_types: List[int] = [] + + for detection_idx, (_, row) in enumerate(annotations_slice.iterrows()): + row = row.to_dict() + yaw, pitch, roll = Quaternion( + w=row["qw"], + x=row["qx"], + y=row["qy"], + z=row["qz"], + ).yaw_pitch_roll + + detections_state[detection_idx, BoundingBoxSE3Index.X] = row["tx_m"] + detections_state[detection_idx, BoundingBoxSE3Index.Y] = row["ty_m"] + detections_state[detection_idx, BoundingBoxSE3Index.Z] = row["tz_m"] + detections_state[detection_idx, BoundingBoxSE3Index.ROLL] = roll + detections_state[detection_idx, BoundingBoxSE3Index.PITCH] = pitch + detections_state[detection_idx, BoundingBoxSE3Index.YAW] = yaw + detections_state[detection_idx, BoundingBoxSE3Index.LENGTH] = row["length_m"] + detections_state[detection_idx, BoundingBoxSE3Index.WIDTH] = row["width_m"] + detections_state[detection_idx, BoundingBoxSE3Index.HEIGHT] = row["height_m"] + + av2_detection_type = AV2SensorBoxDetectionType.deserialize(row["category"]) + detections_types.append(int(AV2_TO_DETECTION_TYPE[av2_detection_type])) + + detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( + origin=ego_state_se3.rear_axle_se3, se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3] + ) + + ZERO_BOX_ROLL_PITCH = False # TODO: Add config option or remove + if ZERO_BOX_ROLL_PITCH: + detections_state[:, BoundingBoxSE3Index.ROLL] = DEFAULT_ROLL + detections_state[:, BoundingBoxSE3Index.PITCH] = DEFAULT_PITCH + + return detections_state.tolist(), detections_velocity.tolist(), detections_token, detections_types + + +def _extract_ego_state(city_se3_egovehicle_df: pd.DataFrame, lidar_timestamp_ns: int) -> EgoStateSE3: + ego_state_slice = get_slice_with_timestamp_ns(city_se3_egovehicle_df, lidar_timestamp_ns) + assert ( + len(ego_state_slice) == 1 + ), f"Expected exactly one ego state for timestamp {lidar_timestamp_ns}, got {len(ego_state_slice)}." + + ego_pose_dict = ego_state_slice.iloc[0].to_dict() + + ego_pose_quat = Quaternion( + w=ego_pose_dict["qw"], + x=ego_pose_dict["qx"], + y=ego_pose_dict["qy"], + z=ego_pose_dict["qz"], + ) + + yaw, pitch, roll = ego_pose_quat.yaw_pitch_roll + + rear_axle_pose = StateSE3( + x=ego_pose_dict["tx_m"], + y=ego_pose_dict["ty_m"], + z=ego_pose_dict["tz_m"], + roll=roll, + pitch=pitch, + yaw=yaw, + ) + vehicle_parameters = get_av2_ford_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters + center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) + # TODO: Add script to calculate the dynamic state from log sequence. + dynamic_state = DynamicStateSE3( + velocity=Vector3D( + x=0.0, + y=0.0, + z=0.0, + ), + acceleration=Vector3D( + x=0.0, + y=0.0, + z=0.0, + ), + angular_velocity=Vector3D( + x=0.0, + y=0.0, + z=0.0, + ), + ) + + return EgoStateSE3( + center_se3=center, + dynamic_state_se3=dynamic_state, + vehicle_parameters=vehicle_parameters, + timepoint=None, + ) + + +def _extract_traffic_lights() -> Tuple[List[int], List[int]]: + return [], [] + + +def _extract_scenario_tag() -> List[str]: + return ["unknown"] + + +def _extract_camera( + lidar_timestamp_ns: int, + city_se3_egovehicle_df: pd.DataFrame, + egovehicle_se3_sensor_df: pd.DataFrame, + ego_state_se3: EgoStateSE3, + synchronization_df: pd.DataFrame, + source_log_path: Path, + data_converter_config: DataConverterConfig, +) -> Dict[CameraType, Union[str, bytes]]: + + camera_dict: Dict[CameraType, Union[str, bytes]] = { + camera_type: None for camera_type in AV2_CAMERA_TYPE_MAPPING.values() + } + split = source_log_path.parent.name + log_id = source_log_path.name + + source_dataset_dir = source_log_path.parent.parent + + rear_axle_se3 = ego_state_se3.rear_axle_se3 + ego_transform = np.zeros((4, 4), dtype=np.float64) + ego_transform[:3, :3] = get_rotation_matrix(ego_state_se3.rear_axle_se3) + ego_transform[:3, 3] = rear_axle_se3.point_3d.array + + for _, row in egovehicle_se3_sensor_df.iterrows(): + row = row.to_dict() + if row["sensor_name"] not in AV2_CAMERA_TYPE_MAPPING: + continue + + camera_name = row["sensor_name"] + camera_type = AV2_CAMERA_TYPE_MAPPING[camera_name] + + relative_image_path = find_closest_target_fpath( + split=split, + log_id=log_id, + src_sensor_name="lidar", + src_timestamp_ns=lidar_timestamp_ns, + target_sensor_name=camera_name, + synchronization_df=synchronization_df, + ) + if relative_image_path is None: + camera_dict[camera_type] = None + else: + absolute_image_path = source_dataset_dir / relative_image_path + assert absolute_image_path.exists() + # TODO: Adjust for finer IMU timestamps to correct the camera extrinsic. + camera_extrinsic = np.eye(4, dtype=np.float64) + camera_extrinsic[:3, :3] = Quaternion( + w=row["qw"], + x=row["qx"], + y=row["qy"], + z=row["qz"], + ).rotation_matrix + camera_extrinsic[:3, 3] = np.array([row["tx_m"], row["ty_m"], row["tz_m"]], dtype=np.float64) + # camera_extrinsic = camera_extrinsic @ ego_transform + camera_extrinsic = camera_extrinsic.flatten().tolist() + + if data_converter_config.camera_store_option == "path": + camera_dict[camera_type] = (str(relative_image_path), camera_extrinsic) + elif data_converter_config.camera_store_option == "binary": + with open(absolute_image_path, "rb") as f: + camera_dict[camera_type] = (f.read(), camera_extrinsic) + + return camera_dict + + +def _extract_lidar(lidar_pc, data_converter_config: DataConverterConfig) -> Dict[LiDARType, Optional[str]]: + + # lidar: Optional[str] = None + # lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / lidar_pc.filename + # if lidar_full_path.exists(): + # lidar = lidar_pc.filename + + # return {LiDARType.LIDAR_MERGED: lidar} + return {} diff --git a/d123/dataset/dataset_specific/av2/av2_helper.py b/d123/dataset/dataset_specific/av2/av2_helper.py new file mode 100644 index 00000000..8ad881cc --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_helper.py @@ -0,0 +1,179 @@ +from pathlib import Path +from typing import Final, List, Literal, Optional + +import pandas as pd + +from d123.dataset.dataset_specific.av2.av2_constants import AV2_CAMERA_TYPE_MAPPING + +AV2_SENSOR_CAM_SHUTTER_INTERVAL_MS: Final[float] = 50.0 +AV2_SENSOR_LIDAR_SWEEP_INTERVAL_W_BUFFER_NS: Final[float] = 102000000.0 + + +AV2_SENSOR_CAMERA_NAMES = [ + "ring_front_center", + "ring_front_left", + "ring_front_right", + "ring_rear_left", + "ring_rear_right", + "ring_side_left", + "ring_side_right", + "stereo_front_left", + "stereo_front_right", +] + + +def get_dataframe_from_file(file_path: Path) -> pd.DataFrame: + if file_path.suffix == ".parquet": + import pyarrow.parquet as pq + + return pq.read_table(file_path) + elif file_path.suffix == ".feather": + import pyarrow.feather as feather + + return feather.read_feather(file_path) + else: + raise ValueError(f"Unsupported file type: {file_path.suffix}") + + +def get_slice_with_timestamp_ns(dataframe: pd.DataFrame, timestamp_ns: int): + """Get the index of the closest timestamp to the target timestamp.""" + return dataframe[dataframe["timestamp_ns"] == timestamp_ns] + + +def build_sensor_dataframe(raw_log_path: Path) -> pd.DataFrame: + + # https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/sensor_dataloader.py#L209 + + split = raw_log_path.parent.name + log_id = raw_log_path.name + + lidar_path = raw_log_path / "sensors" / "lidar" + cameras_path = raw_log_path / "sensors" / "cameras" + + # Find all the lidar records and timestamps from file names. + lidar_records = populate_sensor_records(lidar_path, split, log_id) + + # Find all the camera records and timestamps from file names. + camera_records = [] + for camera_folder in cameras_path.iterdir(): + assert camera_folder.name in AV2_SENSOR_CAMERA_NAMES + camera_record = populate_sensor_records(camera_folder, split, log_id) + camera_records.append(camera_record) + + # Concatenate all the camera records into a single DataFrame. + sensor_records = [lidar_records] + camera_records + sensor_dataframe = pd.concat(sensor_records) + + # Set index as tuples of the form: (split, log_id, sensor_name, timestamp_ns) and sort the index. + # sorts by split, log_id, and then by sensor name, and then by timestamp. + sensor_dataframe.set_index(["split", "log_id", "sensor_name", "timestamp_ns"], inplace=True) + sensor_dataframe.sort_index(inplace=True) + + return sensor_dataframe + + +def build_synchronization_dataframe( + sensor_dataframe: pd.DataFrame, + matching_criterion: Literal["nearest", "forward"] = "nearest", +) -> pd.DataFrame: + + # https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/sensor_dataloader.py#L382 + + # Create list to store synchronized data frames. + sync_list: List[pd.DataFrame] = [] + unique_sensor_names: List[str] = sensor_dataframe.index.unique(level=2).tolist() + + # Associate a 'source' sensor to a 'target' sensor for all available sensors. + # For example, we associate the lidar sensor with each ring camera which + # produces a mapping from lidar -> all-other-sensors. + for src_sensor_name in unique_sensor_names: + src_records = sensor_dataframe.xs(src_sensor_name, level=2, drop_level=False).reset_index() + src_records = src_records.rename({"timestamp_ns": src_sensor_name}, axis=1).sort_values(src_sensor_name) + + # _Very_ important to convert to timedelta. Tolerance below causes precision loss otherwise. + src_records[src_sensor_name] = pd.to_timedelta(src_records[src_sensor_name]) + for target_sensor_name in unique_sensor_names: + if src_sensor_name == target_sensor_name: + continue + target_records = sensor_dataframe.xs(target_sensor_name, level=2).reset_index() + target_records = target_records.rename({"timestamp_ns": target_sensor_name}, axis=1).sort_values( + target_sensor_name + ) + + # Merge based on matching criterion. + # _Very_ important to convert to timedelta. Tolerance below causes precision loss otherwise. + target_records[target_sensor_name] = pd.to_timedelta(target_records[target_sensor_name]) + tolerance = pd.to_timedelta(AV2_SENSOR_CAM_SHUTTER_INTERVAL_MS / 2 * 1e6) + if "ring" in src_sensor_name: + tolerance = pd.to_timedelta(AV2_SENSOR_LIDAR_SWEEP_INTERVAL_W_BUFFER_NS / 2) + src_records = pd.merge_asof( + src_records, + target_records, + left_on=src_sensor_name, + right_on=target_sensor_name, + by=["split", "log_id"], + direction=matching_criterion, + tolerance=tolerance, + ) + sync_list.append(src_records) + + sync_records = pd.concat(sync_list).reset_index(drop=True) + sync_records.set_index(keys=["split", "log_id", "sensor_name"], inplace=True) + sync_records.sort_index(inplace=True) + + return sync_records + + +def populate_sensor_records(sensor_path: Path, split: str, log_id: str) -> pd.DataFrame: + + sensor_name = sensor_path.name + sensor_files = list(sensor_path.iterdir()) + sensor_records = [] + + for sensor_file in sensor_files: + assert sensor_file.suffix in [ + ".feather", + ".jpg", + ], f"Unsupported file type: {sensor_file.suffix} for {str(sensor_file)}" + row = {} + row["split"] = split + row["log_id"] = log_id + row["sensor_name"] = sensor_name + row["timestamp_ns"] = int(sensor_file.stem) + sensor_records.append(row) + + return pd.DataFrame(sensor_records) + + +def find_closest_target_fpath( + split: str, + log_id: str, + src_sensor_name: str, + src_timestamp_ns: int, + target_sensor_name: str, + synchronization_df: pd.DataFrame, +) -> Optional[Path]: + """Find the file path to the target sensor from a source sensor.""" + # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/datasets/sensor/sensor_dataloader.py#L448 + + src_timedelta_ns = pd.Timedelta(src_timestamp_ns) + src_to_target_records = synchronization_df.loc[(split, log_id, src_sensor_name)].set_index(src_sensor_name) + index = src_to_target_records.index + if src_timedelta_ns not in index: + # This timestamp does not correspond to any lidar sweep. + return None + + # Grab the synchronization record. + target_timestamp_ns = src_to_target_records.loc[src_timedelta_ns, target_sensor_name] + if pd.isna(target_timestamp_ns): + # No match was found within tolerance. + return None + + sensor_dir = Path(split) / log_id / "sensors" + valid_cameras = list(AV2_CAMERA_TYPE_MAPPING.keys()) + timestamp_ns_str = str(target_timestamp_ns.asm8.item()) + if target_sensor_name in valid_cameras: + target_path = sensor_dir / "cameras" / target_sensor_name / f"{timestamp_ns_str}.jpg" + else: + target_path = sensor_dir / target_sensor_name / f"{timestamp_ns_str}.feather" + return target_path diff --git a/d123/dataset/dataset_specific/av2/av2_map_conversion.py b/d123/dataset/dataset_specific/av2/av2_map_conversion.py new file mode 100644 index 00000000..390cbdb2 --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_map_conversion.py @@ -0,0 +1,527 @@ +import warnings +from pathlib import Path +from typing import Any, Dict, Final, List + +import geopandas as gpd +import numpy as np +import numpy.typing as npt +import pandas as pd +import shapely +import shapely.geometry as geom +from flask import json + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.line.polylines import Polyline2D, Polyline3D +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length +from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import ( + get_road_edges_3d_from_generic_drivable_area_df, +) +from d123.dataset.dataset_specific.av2.av2_constants import AV2_ROAD_LINE_TYPE_MAPPING +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType + +LANE_GROUP_MARK_TYPES: List[str] = [ + "DASHED_WHITE", + "DOUBLE_DASH_WHITE", + "DASH_SOLID_WHITE", + "SOLID_DASH_WHITE", + "SOLID_WHITE", +] +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # TODO: Add to config + + +def convert_av2_map(source_log_path: Path, map_file_path: Path) -> None: + + def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Polyline3D: + polyline = np.array([[p["x"], p["y"], p["z"]] for p in data], dtype=np.float64) + if close: + polyline = np.vstack([polyline, polyline[0]]) + + return Polyline3D.from_array(polyline) + + map_folder = source_log_path / "map" + log_map_archive_path = next(map_folder.glob("log_map_archive_*.json")) + + with open(log_map_archive_path, "r") as f: + log_map_archive = json.load(f) + + drivable_areas: Dict[int, Polyline3D] = {} + + for drivable_area_id, drivable_area_dict in log_map_archive["drivable_areas"].items(): + # keys: ["area_boundary", "id"] + drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict["area_boundary"], close=True) + + for lane_segment_id, lane_segment_dict in log_map_archive["lane_segments"].items(): + # keys = [ + # "id", + # "is_intersection", + # "lane_type", + # "left_lane_boundary", + # "left_lane_mark_type", + # "right_lane_boundary", + # "right_lane_mark_type", + # "successors", + # "predecessors", + # "right_neighbor_id", + # "left_neighbor_id", + # ] + lane_segment_dict["left_lane_boundary"] = _extract_polyline(lane_segment_dict["left_lane_boundary"]) + lane_segment_dict["right_lane_boundary"] = _extract_polyline(lane_segment_dict["right_lane_boundary"]) + + for crosswalk_id, crosswalk_dict in log_map_archive["pedestrian_crossings"].items(): + # keys = ["id", "outline"] + # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/map/pedestrian_crossing.py + + p1, p2 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge1"]], dtype=np.float64) + p3, p4 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge2"]], dtype=np.float64) + crosswalk_dict["outline"] = Polyline3D.from_array(np.array([p1, p2, p4, p3, p1], dtype=np.float64)) + + lane_group_dict = _extract_lane_group_dict(log_map_archive["lane_segments"]) + intersection_dict = _extract_intersection_dict(log_map_archive["lane_segments"], lane_group_dict) + + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + + dataframes[MapLayer.LANE] = get_lane_df(log_map_archive["lane_segments"]) + dataframes[MapLayer.LANE_GROUP] = get_lane_group_df(lane_group_dict) + dataframes[MapLayer.INTERSECTION] = get_intersections_df(intersection_dict) + dataframes[MapLayer.CROSSWALK] = get_crosswalk_df(log_map_archive["pedestrian_crossings"]) + dataframes[MapLayer.GENERIC_DRIVABLE] = get_generic_drivable_df(drivable_areas) + dataframes[MapLayer.ROAD_EDGE] = get_road_edge_df(dataframes[MapLayer.GENERIC_DRIVABLE]) + dataframes[MapLayer.ROAD_LINE] = get_road_line_df(log_map_archive["lane_segments"]) + # NOTE: AV2 does not provide walkways or carparks, so we create an empty DataFrame. + dataframes[MapLayer.WALKWAY] = get_empty_gdf() + dataframes[MapLayer.CARPARK] = get_empty_gdf() + + map_file_path.unlink(missing_ok=True) + if not map_file_path.parent.exists(): + map_file_path.parent.mkdir(parents=True, exist_ok=True) + + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_path, layer=layer.serialize(), driver="GPKG", mode="a") + + +def get_empty_gdf() -> gpd.GeoDataFrame: + ids = [] + outlines = [] + geometries = [] + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_lane_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: + + ids = [int(lane_id) for lane_id in lanes.keys()] + lane_types = [0] * len(ids) # TODO: Add lane types + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] + baseline_paths = [] + geometries = [] + + def _get_centerline_from_boundaries( + left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 0.1 + ) -> Polyline3D: + + points_per_meter = 1 / resolution + num_points = int(np.ceil(max([right_boundary.length, left_boundary.length]) * points_per_meter)) + right_array = right_boundary.interpolate(np.linspace(0, right_boundary.length, num_points, endpoint=True)) + left_array = left_boundary.interpolate(np.linspace(0, left_boundary.length, num_points, endpoint=True)) + + return Polyline3D.from_array(np.mean([right_array, left_array], axis=0)) + + for lane_id, lane_dict in lanes.items(): + # keys = [ + # "id", + # "is_intersection", + # "lane_type", + # "left_lane_boundary", + # "left_lane_mark_type", + # "right_lane_boundary", + # "right_lane_mark_type", + # "successors", + # "predecessors", + # "right_neighbor_id", + # "left_neighbor_id", + # ] + lane_centerline = _get_centerline_from_boundaries( + left_boundary=lane_dict["left_lane_boundary"], + right_boundary=lane_dict["right_lane_boundary"], + ) + lane_speed_limit_mps = None # TODO: Consider using geo reference to retrieve speed limits. + lane_group_ids.append(lane_id) + speed_limits_mps.append(lane_speed_limit_mps) + predecessor_ids.append(lane_dict["predecessors"]) + successor_ids.append(lane_dict["successors"]) + left_boundaries.append(lane_dict["left_lane_boundary"].linestring) + right_boundaries.append(lane_dict["right_lane_boundary"].linestring) + left_lane_ids.append(lane_dict["left_neighbor_id"]) + right_lane_ids.append(lane_dict["right_neighbor_id"]) + baseline_paths.append(lane_centerline.linestring) + + geometry = geom.Polygon( + np.vstack( + [ + lane_dict["left_lane_boundary"].array[:, :2], + lane_dict["right_lane_boundary"].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_type": lane_types, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, + "baseline_path": baseline_paths, + } + ) + + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_lane_group_df(lane_group_dict: Dict[int, Any]) -> gpd.GeoDataFrame: + + ids = list(lane_group_dict.keys()) + lane_ids = [] + intersection_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + for lane_group_id, lane_group_values in lane_group_dict.items(): + + lane_ids.append(lane_group_values["lane_ids"]) + intersection_ids.append(lane_group_values["intersection_id"]) + + predecessor_lane_group_ids.append(lane_group_values["predecessor_ids"]) + successor_lane_group_ids.append(lane_group_values["successor_ids"]) + left_boundaries.append(lane_group_values["left_boundary"].linestring) + right_boundaries.append(lane_group_values["right_boundary"].linestring) + geometry = geom.Polygon( + np.vstack( + [ + lane_group_values["left_boundary"].array[:, :2], + lane_group_values["right_boundary"].array[:, :2][::-1], + lane_group_values["left_boundary"].array[0, :2][None, ...], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "intersection_id": intersection_ids, + "predecessor_ids": predecessor_lane_group_ids, + "successor_ids": successor_lane_group_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_intersections_df(intersection_dict: Dict[int, Any]) -> gpd.GeoDataFrame: + ids = [] + lane_group_ids = [] + outlines = [] + geometries = [] + + for intersection_id, intersection_values in intersection_dict.items(): + ids.append(intersection_id) + lane_group_ids.append(intersection_values["lane_group_ids"]) + outlines.append(intersection_values["outline_3d"].linestring) + geometries.append(geom.Polygon(intersection_values["outline_3d"].array[:, Point3DIndex.XY])) + + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_carpark_df(carparks) -> gpd.GeoDataFrame: + ids = list(carparks.keys()) + outlines = [geom.LineString(outline) for outline in carparks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in carparks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_walkway_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + # NOTE: WOPD does not provide walkways, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(crosswalks.keys()) + outlines = [] + geometries = [] + for crosswalk_dict in crosswalks.values(): + outline = crosswalk_dict["outline"] + outlines.append(outline.linestring) + geometries.append(geom.Polygon(outline.array[:, Point3DIndex.XY])) + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_generic_drivable_df(drivable_areas: Dict[int, Polyline3D]) -> gpd.GeoDataFrame: + ids = list(drivable_areas.keys()) + outlines = [drivable_area.linestring for drivable_area in drivable_areas.values()] + geometries = [geom.Polygon(drivable_area.array[:, Point3DIndex.XY]) for drivable_area in drivable_areas.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_road_edge_df(generic_drivable_area_df: gpd.GeoDataFrame) -> gpd.GeoDataFrame: + road_edges = get_road_edges_3d_from_generic_drivable_area_df(generic_drivable_area_df) + road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) + + ids = np.arange(len(road_edges), dtype=np.int64).tolist() + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types = [int(RoadEdgeType.ROAD_EDGE_BOUNDARY)] * len(road_edges) + geometries = road_edges + return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) + + +def get_road_line_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: + + # TODO @DanielDauner: Allow lanes to reference road line dataframe. + + ids = [] + road_lines_type = [] + geometries = [] + + running_id = 0 + for lane in lanes.values(): + for side in ["left", "right"]: + # NOTE: We currently ignore lane markings that are NONE in the AV2 dataset. + # TODO: Review if the road line system should be changed in the future. + if lane[f"{side}_lane_mark_type"] == "NONE": + continue + + ids.append(running_id) + road_lines_type.append(AV2_ROAD_LINE_TYPE_MAPPING[lane[f"{side}_lane_mark_type"]]) + geometries.append(lane[f"{side}_lane_boundary"].linestring) + running_id += 1 + + data = pd.DataFrame({"id": ids, "road_line_type": road_lines_type}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def _extract_lane_group_dict(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: + + lane_group_sets = _extract_lane_group(lanes) + lane_group_set_dict = {i: lane_group for i, lane_group in enumerate(lane_group_sets)} + + lane_group_dict: Dict[int, Dict[str, Any]] = {} + + def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: + """Helper to find lane group ids that contain any of the given lane ids.""" + lane_group_ids_ = [] + for lane_group_id_, lane_group_set_ in lane_group_set_dict.items(): + if any(str(lane_id) in lane_group_set_ for lane_id in lane_ids): + lane_group_ids_.append(lane_group_id_) + return list(set(lane_group_ids_)) + + for lane_group_id, lane_group_set in lane_group_set_dict.items(): + + lane_group_dict[lane_group_id] = {} + lane_group_dict[lane_group_id]["id"] = lane_group_id + lane_group_dict[lane_group_id]["lane_ids"] = [int(lane_id) for lane_id in lane_group_set] + + successor_lanes = [] + predecessor_lanes = [] + for lane_id in lane_group_set: + lane_dict = lanes[str(lane_id)] + successor_lanes.extend(lane_dict["successors"]) + predecessor_lanes.extend(lane_dict["predecessors"]) + + left_boundary = lanes[lane_group_set[0]]["left_lane_boundary"] + right_boundary = lanes[lane_group_set[-1]]["right_lane_boundary"] + + lane_group_dict[lane_group_id]["intersection_id"] = None + lane_group_dict[lane_group_id]["predecessor_ids"] = _get_lane_group_ids_of_lanes_ids(predecessor_lanes) + lane_group_dict[lane_group_id]["successor_ids"] = _get_lane_group_ids_of_lanes_ids(successor_lanes) + lane_group_dict[lane_group_id]["left_boundary"] = left_boundary + lane_group_dict[lane_group_id]["right_boundary"] = right_boundary + outline_array = np.vstack( + [ + left_boundary.array[:, :3], + right_boundary.array[:, :3][::-1], + left_boundary.array[0, :3][None, ...], + ] + ) + + lane_group_dict[lane_group_id]["outline"] = Polyline3D.from_array(outline_array) + + return lane_group_dict + + +def _extract_lane_group(lanes) -> List[List[str]]: + + visited = set() + lane_groups = [] + + def _get_valid_neighbor_id(lane_data, direction): + """Helper function to safely get neighbor ID""" + neighbor_key = f"{direction}_neighbor_id" + neighbor_id = str(lane_data.get(neighbor_key)) + mark_type = lane_data.get(f"{direction}_lane_mark_type", None) + + if (neighbor_id is not None) and (neighbor_id in lanes) and (mark_type in LANE_GROUP_MARK_TYPES): + return neighbor_id + return None + + def _traverse_group(start_lane_id): + """ + Traverse left and right from a starting lane to find all connected parallel lanes + """ + group = [start_lane_id] + queue = [start_lane_id] + + while queue: + current_id = queue.pop(0) + if current_id in visited: + continue + + visited.add(current_id) + + # Check left neighbor + left_neighbor = _get_valid_neighbor_id(lanes[current_id], "left") + if left_neighbor is not None and left_neighbor not in visited: + queue.append(left_neighbor) + group = [left_neighbor] + group + + # Check right neighbor + right_neighbor = _get_valid_neighbor_id(lanes[current_id], "right") + if right_neighbor is not None and right_neighbor not in visited: + queue.append(right_neighbor) + group = group + [right_neighbor] + + return group + + # Find all lane groups + for lane_id in lanes: + if lane_id not in visited: + group = _traverse_group(lane_id) + lane_groups.append(group) + + return lane_groups + + +def _extract_intersection_dict( + lanes: Dict[int, Any], lane_group_dict: Dict[int, Any], max_distance: float = 0.01 +) -> Dict[str, Any]: + + def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float: + """Interpolate Z coordinate along a 3D line segment.""" + p1, p2 = segment_coords[0], segment_coords[1] + + # Project point onto segment + segment_vec = p2[:2] - p1[:2] + point_vec = np.array([point.x, point.y]) - p1[:2] + + # Handle degenerate case + segment_length_sq = np.dot(segment_vec, segment_vec) + if segment_length_sq == 0: + return p1[2] + + # Calculate projection parameter + t = np.dot(point_vec, segment_vec) / segment_length_sq + t = np.clip(t, 0, 1) # Clamp to segment bounds + + # Interpolate Z + return p1[2] + t * (p2[2] - p1[2]) + + # 1. Collect all lane groups where at least one lane is marked as an intersection. + lane_group_intersection_dict = {} + for lane_group_id, lane_group in lane_group_dict.items(): + is_intersection_lanes = [lanes[str(lane_id)]["is_intersection"] for lane_id in lane_group["lane_ids"]] + if any(is_intersection_lanes): + lane_group_intersection_dict[lane_group_id] = lane_group + + # 2. Merge polygons of lane groups that are marked as intersections. + lane_group_intersection_geometry = { + lane_group_id: shapely.Polygon(lane_group["outline"].array[:, Point3DIndex.XY]) + for lane_group_id, lane_group in lane_group_intersection_dict.items() + } + intersection_polygons = gpd.GeoSeries(lane_group_intersection_geometry).union_all() + + # 3. Collect all intersection polygons and their lane group IDs. + intersection_dict = {} + for intersection_idx, intersection_polygon in enumerate(intersection_polygons.geoms): + if intersection_polygon.is_empty: + continue + lane_group_ids = [ + lane_group_id + for lane_group_id, lane_group_polygon in lane_group_intersection_geometry.items() + if intersection_polygon.intersects(lane_group_polygon) + ] + for lane_group_id in lane_group_ids: + lane_group_dict[lane_group_id]["intersection_id"] = intersection_idx + + intersection_dict[intersection_idx] = { + "id": intersection_idx, + "outline_2d": Polyline2D.from_array(np.array(list(intersection_polygon.exterior.coords), dtype=np.float64)), + "lane_group_ids": lane_group_ids, + } + + # 4. Lift intersection outlines to 3D. + boundary_segments = [] + for lane_group in lane_group_intersection_dict.values(): + coords = np.array(lane_group["outline"].linestring.coords, dtype=np.float64).reshape(-1, 1, 3) + segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1) + boundary_segments.append(segment_coords_boundary) + + boundary_segments = np.concatenate(boundary_segments, axis=0) + boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments) + occupancy_map = OccupancyMap2D(boundary_segment_linestrings) + + for intersection_id, intersection_data in intersection_dict.items(): + points_2d = intersection_data["outline_2d"].array + points_3d = np.zeros((len(points_2d), 3), dtype=np.float64) + points_3d[:, :2] = points_2d + + query_points = shapely.creation.points(points_2d) + results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True) + for query_idx, geometry_idx in zip(*results): + query_point = query_points[query_idx] + segment_coords = boundary_segments[geometry_idx] + best_z = _interpolate_z_on_segment(query_point, segment_coords) + points_3d[query_idx, 2] = best_z + + intersection_dict[intersection_id]["outline_3d"] = Polyline3D.from_array(points_3d) + + return intersection_dict diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py new file mode 100644 index 00000000..f598bb1e --- /dev/null +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -0,0 +1,455 @@ +import gc +import gzip +import hashlib +import json +import os +from dataclasses import asdict +from functools import partial +from pathlib import Path +from typing import Any, Dict, Final, List, Optional, Tuple, Union + +import numpy as np +import pyarrow as pa + +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar_index import CarlaLidarIndex +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3Index +from d123.common.datatypes.vehicle_state.vehicle_parameters import get_carla_lincoln_mkz_2020_parameters +from d123.common.geometry.base import Point2D, Point3D +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index +from d123.common.geometry.vector import Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map +from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table +from d123.dataset.conversion.map.opendrive.opendrive_map_conversion import convert_from_xodr +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.logs.log_metadata import LogMetadata +from d123.dataset.maps.abstract_map import AbstractMap, MapLayer +from d123.dataset.maps.abstract_map_objects import AbstractLane +from d123.dataset.scene.arrow_scene import get_map_api_from_names + +AVAILABLE_CARLA_MAP_LOCATIONS: Final[List[str]] = [ + "Town01", # A small, simple town with a river and several bridges. + "Town02", # A small simple town with a mixture of residential and commercial buildings. + "Town03", # A larger, urban map with a roundabout and large junctions. + "Town04", # A small town embedded in the mountains with a special "figure of 8" infinite highway. + "Town05", # Squared-grid town with cross junctions and a bridge. It has multiple lanes per direction. Useful to perform lane changes. + "Town06", # Long many lane highways with many highway entrances and exits. It also has a Michigan left. + "Town07", # A rural environment with narrow roads, corn, barns and hardly any traffic lights. + "Town10HD", # A downtown urban environment with skyscrapers, residential buildings and an ocean promenade. + "Town11", # A Large Map that is undecorated. Serves as a proof of concept for the Large Maps feature. + "Town12", # A Large Map with numerous different regions, including high-rise, residential and rural environments. + "Town13", # ??? + "Town15", # ??? +] + +CARLA_DT: Final[float] = 0.1 # [s] +TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE: Final[float] = 1.0 # [m] +SORT_BY_TIMESTAMP: Final[bool] = True + +CARLA_CAMERA_TYPES = {CameraType.CAM_F0} + +CARLA_DATA_ROOT: Final[Path] = Path(os.environ["CARLA_DATA_ROOT"]) + + +# TODO: Refactor this files and convert coordinate systems more elegantly. +# NOTE: Currently some coordinate transforms from Unreal to ISO 8855 are done in the data agent of PDM-Lite. +# Ideally a general function to transform poses and points between coordinate systems would be nice + + +def _load_json_gz(path: Path) -> Dict: + """Helper function to load a gzipped JSON file.""" + with gzip.open(path, "rt") as f: + data = json.load(f) + return data + + +def create_token(input_data: str) -> str: + # TODO: Refactor this function. + # TODO: Add a general function to create tokens from arbitrary data. + if isinstance(input_data, str): + input_data = input_data.encode("utf-8") + + hash_obj = hashlib.sha256(input_data) + return hash_obj.hexdigest()[:16] + + +class CarlaDataConverter(RawDataConverter): + + def __init__( + self, + splits: List[str], + log_path: Union[Path, str], + data_converter_config: DataConverterConfig, + ) -> None: + super().__init__(data_converter_config) + for split in splits: + assert ( + split in self.get_available_splits() + ), f"Split {split} is not available. Available splits: {self.available_splits}" + + self._splits: str = splits + self._log_path: Path = Path(log_path) + self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() + + def _collect_log_paths(self) -> Dict[str, List[Path]]: + # TODO: fix "carla" split placeholder and add support for other splits + log_paths_per_split: Dict[str, List[Path]] = {} + log_paths = [ + log_path for log_path in self._log_path.iterdir() if log_path.is_dir() and log_path.stem != "sensor_blobs" + ] + log_paths_per_split["carla"] = log_paths + return log_paths_per_split + + def get_available_splits(self) -> List[str]: + """Returns a list of available raw data types.""" + return ["carla"] # TODO: fix the placeholder + + def convert_maps(self, worker: WorkerPool) -> None: + worker_map( + worker, + partial( + convert_carla_map_to_gpkg, + data_converter_config=self.data_converter_config, + ), + list(AVAILABLE_CARLA_MAP_LOCATIONS), + ) + + def convert_logs(self, worker: WorkerPool) -> None: + + log_args = [ + {"log_path": log_path, "split": split} + for split, log_paths in self._log_paths_per_split.items() + for log_path in log_paths + ] + + worker_map( + worker, partial(convert_carla_log_to_arrow, data_converter_config=self.data_converter_config), log_args + ) + + +def convert_carla_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: + + # TODO: add to config + _interpolation_step_size = 0.5 # [m] + _connection_distance_threshold = 0.1 # [m] + for map_name in map_names: + map_path = data_converter_config.output_path / "maps" / f"carla_{map_name.lower()}.gpkg" + if data_converter_config.force_map_conversion or not map_path.exists(): + map_path.unlink(missing_ok=True) + assert os.environ["CARLA_ROOT"] is not None + CARLA_ROOT = Path(os.environ["CARLA_ROOT"]) + + if map_name not in ["Town11", "Town12", "Town13", "Town15"]: + carla_maps_root = CARLA_ROOT / "CarlaUE4" / "Content" / "Carla" / "Maps" / "OpenDrive" + carla_map_path = carla_maps_root / f"{map_name}.xodr" + else: + carla_map_path = ( + CARLA_ROOT / "CarlaUE4" / "Content" / "Carla" / "Maps" / map_name / "OpenDrive" / f"{map_name}.xodr" + ) + + convert_from_xodr( + carla_map_path, + f"carla_{map_name.lower()}", + _interpolation_step_size, + _connection_distance_threshold, + ) + return [] + + +def convert_carla_log_to_arrow( + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig +) -> List[Any]: + def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> None: + for log_info in args: + log_path: Path = log_info["log_path"] + split: str = log_info["split"] + output_path: Path = data_converter_config.output_path + + log_file_path = output_path / split / f"{log_path.stem}.arrow" + + if data_converter_config.force_log_conversion or not log_file_path.exists(): + log_file_path.unlink(missing_ok=True) + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + bounding_box_paths = sorted([bb_path for bb_path in (log_path / "boxes").iterdir()]) + first_log_dict = _load_json_gz(bounding_box_paths[0]) + map_name = first_log_dict["location"] + map_api = get_map_api_from_names("carla", map_name) + + metadata = _get_metadata(map_name, str(log_path.stem)) + vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() + camera_metadata = get_carla_camera_metadata(first_log_dict) + lidar_metadata = get_carla_lidar_metadata(first_log_dict) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + schema_column_list.append((lidar_type.serialize(), pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + # TODO: Adjust how cameras are added + if data_converter_config.camera_store_option is not None: + for camera_type in camera_metadata.keys(): + if data_converter_config.camera_store_option == "path": + schema_column_list.append((camera_type.serialize(), pa.string())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 16)) + ) + elif data_converter_config.camera_store_option == "binary": + raise NotImplementedError("Binary camera storage is not implemented.") + + recording_schema = pa.schema(schema_column_list) + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), + } + ) + + _write_recording_table( + bounding_box_paths, + map_api, + recording_schema, + log_file_path, + data_converter_config, + ) + + gc.collect() + + convert_log_internal(args) + gc.collect() + return [] + + +def _get_metadata(location: str, log_name: str) -> LogMetadata: + return LogMetadata( + dataset="carla", + log_name=log_name, + location=location, + timestep_seconds=CARLA_DT, + map_has_z=True, + ) + + +def get_carla_camera_metadata(first_log_dict: Dict[str, Any]) -> Dict[CameraType, CameraMetadata]: + + # FIXME: This is a placeholder function to return camera metadata. + + intrinsic = np.array( + first_log_dict[f"{CameraType.CAM_F0.serialize()}_intrinsics"], + dtype=np.float64, + ) + camera_metadata = { + CameraType.CAM_F0: CameraMetadata( + camera_type=CameraType.CAM_F0, + width=1024, + height=512, + intrinsic=intrinsic, + distortion=np.zeros((5,), dtype=np.float64), + ) + } + return camera_metadata + + +def get_carla_lidar_metadata(first_log_dict: Dict[str, Any]) -> Dict[LiDARType, LiDARMetadata]: + + # TODO: add lidar extrinsic + lidar_metadata = { + LiDARType.LIDAR_TOP: LiDARMetadata( + lidar_type=LiDARType.LIDAR_TOP, + lidar_index=CarlaLidarIndex, + extrinsic=None, + ) + } + return lidar_metadata + + +def _write_recording_table( + bounding_box_paths: List[Path], + map_api: AbstractMap, + recording_schema: pa.Schema, + log_file_path: Path, + data_converter_config: DataConverterConfig, +) -> pa.Table: + # TODO: Refactor this function to be more readable + log_name = str(bounding_box_paths[0].parent.parent.stem) + + with pa.OSFile(str(log_file_path), "wb") as sink: + with pa.ipc.new_file(sink, recording_schema) as writer: + for box_path in bounding_box_paths: + sample_name = box_path.stem.split(".")[0] + + data = _load_json_gz(box_path) + traffic_light_ids, traffic_light_types = _extract_traffic_light_data( + data["traffic_light_states"], data["traffic_light_positions"], map_api + ) + route_lane_group_ids = _extract_route_lane_group_ids(data["route"], map_api) if "route" in data else [] + + row_data = { + "token": [create_token(f"{log_name}_{box_path.stem}")], + "timestamp": [data["timestamp"]], + "detections_state": [_extract_detection_states(data["detections_state"])], + "detections_velocity": [ + ( + data["detections_velocity"] + if "detections_velocity" in data + else np.zeros((len(data["detections_types"]), 3)).tolist() + ) + ], + "detections_token": [data["detections_token"]], + "detections_type": [data["detections_types"]], + "ego_states": [_extract_ego_vehicle_state(data["ego_state"])], + "traffic_light_ids": [traffic_light_ids], + "traffic_light_types": [traffic_light_types], + "scenario_tag": [data["scenario_tag"]], + "route_lane_group_ids": [route_lane_group_ids], + } + if data_converter_config.lidar_store_option is not None: + lidar_data_dict = _extract_lidar(log_name, sample_name, data_converter_config) + for lidar_type, lidar_data in lidar_data_dict.items(): + row_data[lidar_type.serialize()] = [lidar_data] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_cameras(data, log_name, sample_name, data_converter_config) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] + else: + row_data[camera_type.serialize()] = [None] + + batch = pa.record_batch(row_data, schema=recording_schema) + writer.write_batch(batch) + del batch, row_data + + if SORT_BY_TIMESTAMP: + recording_table = open_arrow_table(log_file_path) + recording_table = recording_table.sort_by([("timestamp", "ascending")]) + write_arrow_table(recording_table, log_file_path) + + +def _extract_ego_vehicle_state(ego_state_list: List[float]) -> List[float]: + # NOTE: This function used to convert coordinate systems, but it is not needed anymore. + assert len(ego_state_list) == len(EgoStateSE3Index), "Ego state list has incorrect length." + return ego_state_list + + +def _extract_detection_states(detection_states: List[List[float]]) -> List[List[float]]: + # NOTE: This function used to convert coordinate systems, but it is not needed anymore. + detection_state_converted = [] + for detection_state in detection_states: + assert len(detection_state) == len(BoundingBoxSE3Index), "Detection state has incorrect length." + detection_state_converted.append(detection_state) + return detection_state_converted + + +def _extract_traffic_light_data( + traffic_light_states: List[int], traffic_light_positions: List[List[float]], map_api: AbstractMap +) -> Tuple[List[int], List[int]]: + traffic_light_types: List[int] = [] + traffic_light_ids: List[int] = [] + for traffic_light_state, traffic_light_waypoints in zip(traffic_light_states, traffic_light_positions): + for traffic_light_waypoint in traffic_light_waypoints: + point_3d = Point3D(*traffic_light_waypoint) + nearby_lanes = map_api.get_proximal_map_objects( + point_3d, TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE, [MapLayer.LANE] + )[MapLayer.LANE] + + for lane in nearby_lanes: + lane: AbstractLane + lane_start_point = lane.centerline.array[0] + distance_to_lane_start = np.linalg.norm(lane_start_point - point_3d.array) + if distance_to_lane_start < TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE: + traffic_light_ids.append(int(lane.id)) + traffic_light_types.append(traffic_light_state) + return traffic_light_ids, traffic_light_types + + +def _extract_route_lane_group_ids(route: List[List[float]], map_api: AbstractMap) -> List[int]: + + # FIXME: Carla route is very buggy. No check if lanes are connected. + route = np.array(route, dtype=np.float64) + route[..., 1] = -route[..., 1] # Unreal coordinate system to ISO 8855 + route = route[::2] + + route_lane_group_ids: List[int] = [] + + for point in route[:200]: + point_2d = Point2D(point[0], point[1]) + nearby_lane_groups = map_api.query(point_2d.shapely_point, [MapLayer.LANE_GROUP], predicate="intersects")[ + MapLayer.LANE_GROUP + ] + if len(nearby_lane_groups) == 0: + continue + elif len(nearby_lane_groups) > 1: + possible_lane_group_ids = [lane_group.id for lane_group in nearby_lane_groups] + if len(route_lane_group_ids) > 0: + prev_lane_group_id = route_lane_group_ids[-1] + if prev_lane_group_id in possible_lane_group_ids: + continue + else: + # TODO: Choose with least heading difference? + route_lane_group_ids.append(int(nearby_lane_groups[0].id)) + else: + # TODO: Choose with least heading difference? + route_lane_group_ids.append(int(nearby_lane_groups[0].id)) + elif len(nearby_lane_groups) == 1: + route_lane_group_ids.append(int(nearby_lane_groups[0].id)) + + return list(dict.fromkeys(route_lane_group_ids)) # Remove duplicates while preserving order + + +def _extract_cameras( + data: Dict[str, Any], log_name: str, sample_name: str, data_converter_config: DataConverterConfig +) -> Dict[CameraType, Optional[str]]: + camera_dict: Dict[str, Union[str, bytes]] = {} + for camera_type in CARLA_CAMERA_TYPES: + camera_full_path = CARLA_DATA_ROOT / "sensor_blobs" / log_name / camera_type.name / f"{sample_name}.jpg" + if camera_full_path.exists(): + if data_converter_config.camera_store_option == "path": + path = f"{log_name}/{camera_type.name}/{sample_name}.jpg" + extrinsics = data.get(f"{camera_type.serialize()}_transform", None) + camera_dict[camera_type] = path, ( + np.array(extrinsics, dtype=np.float64).flatten() if extrinsics is not None else None + ) + + elif data_converter_config.camera_store_option == "binary": + raise NotImplementedError("Binary camera storage is not implemented.") + else: + camera_dict[camera_type] = None + return camera_dict + + +def _extract_lidar( + log_name: str, sample_name: str, data_converter_config: DataConverterConfig +) -> Dict[LiDARType, Optional[str]]: + + lidar: Optional[str] = None + lidar_full_path = CARLA_DATA_ROOT / "sensor_blobs" / log_name / "lidar" / f"{sample_name}.npy" + if lidar_full_path.exists(): + if data_converter_config.lidar_store_option == "path": + lidar = f"{log_name}/lidar/{sample_name}.npy" + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + else: + raise FileNotFoundError(f"LiDAR file not found: {lidar_full_path}") + + return {LiDARType.LIDAR_TOP: lidar} if lidar else None diff --git a/d123/dataset/dataset_specific/carla/carla_data_processor.py b/d123/dataset/dataset_specific/carla/carla_data_processor.py deleted file mode 100644 index 0bec510b..00000000 --- a/d123/dataset/dataset_specific/carla/carla_data_processor.py +++ /dev/null @@ -1,331 +0,0 @@ -import gc -import gzip -import hashlib -import json -from dataclasses import asdict -from functools import partial -from pathlib import Path -from typing import Dict, Final, List, Tuple, Union - -import numpy as np -import pyarrow as pa -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map - -from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3Index -from d123.common.datatypes.vehicle_state.vehicle_parameters import get_carla_lincoln_mkz_2020_parameters -from d123.common.geometry.base import Point2D, Point3D, StateSE3 -from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index -from d123.common.geometry.transform.se3 import translate_se3_along_z -from d123.common.geometry.vector import Vector3DIndex -from d123.dataset.arrow.conversion import VehicleParameters -from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor -from d123.dataset.logs.log_metadata import LogMetadata -from d123.dataset.maps.abstract_map import AbstractMap, MapSurfaceType -from d123.dataset.maps.abstract_map_objects import AbstractLane -from d123.dataset.scene.arrow_scene import get_map_api_from_names - -CARLA_DT: Final[float] = 0.1 # [s] -TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE: Final[float] = 1.0 # [m] -SORT_BY_TIMESTAMP: Final[bool] = True - - -# TODO: Refactor this files and convert coordinate systems more elegantly. -# NOTE: Currently some coordinate transforms from Unreal to ISO 8855 are done in the data agent of PDM-Lite. -# Ideally a general function to transform poses and points between coordinate systems would be nice - - -def _load_json_gz(path: Path) -> Dict: - """Helper function to load a gzipped JSON file.""" - with gzip.open(path, "rt") as f: - data = json.load(f) - return data - - -def create_token(input_data: str) -> str: - # TODO: Refactor this function. - # TODO: Add a general function to create tokens from arbitrary data. - if isinstance(input_data, str): - input_data = input_data.encode("utf-8") - - hash_obj = hashlib.sha256(input_data) - return hash_obj.hexdigest()[:16] - - -class CarlaDataProcessor(RawDataProcessor): - def __init__( - self, - splits: List[str], - log_path: Union[Path, str], - output_path: Union[Path, str], - force_data_conversion: bool, - ) -> None: - super().__init__(force_data_conversion) - for split in splits: - assert ( - split in self.get_available_splits() - ), f"Split {split} is not available. Available splits: {self.available_splits}" - - self._splits: str = splits - self._log_path: Path = Path(log_path) - self._output_path: Path = Path(output_path) - self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() - - def _collect_log_paths(self) -> Dict[str, List[Path]]: - # TODO: fix "carla" split placeholder and add support for other splits - log_paths_per_split: Dict[str, List[Path]] = {} - log_paths = list(self._log_path.iterdir()) - log_paths_per_split["carla"] = log_paths - return log_paths_per_split - - def get_available_splits(self) -> List[str]: - """Returns a list of available raw data types.""" - return ["carla"] # TODO: fix the placeholder - - def convert(self, worker: WorkerPool) -> None: - - log_args = [ - { - "log_path": log_path, - "output_path": self._output_path, - "split": split, - } - for split, log_paths in self._log_paths_per_split.items() - for log_path in log_paths - ] - - worker_map( - worker, partial(convert_carla_log_to_arrow, force_data_conversion=self.force_data_conversion), log_args - ) - - -def convert_carla_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], force_data_conversion: bool -) -> None: - def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> None: - for log_info in args: - log_path: Path = log_info["log_path"] - output_path: Path = log_info["output_path"] - split: str = log_info["split"] - - log_file_path = output_path / split / f"{log_path.stem}.arrow" - - if force_data_conversion or not log_file_path.exists(): - log_file_path.unlink(missing_ok=True) - if not log_file_path.parent.exists(): - log_file_path.parent.mkdir(parents=True, exist_ok=True) - - bounding_box_paths = sorted([bb_path for bb_path in (log_path / "boxes").iterdir()]) - map_name = _load_json_gz(bounding_box_paths[0])["location"] - map_api = get_map_api_from_names("carla", map_name) - - recording_schema = pa.schema( - [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ("front_cam_demo", pa.binary()), - ("front_cam_transform", pa.list_(pa.float64())), - ] - ) - metadata = _get_metadata(map_name, str(log_path.stem)) - vehicle_parameters = _get_vehicle_parameters_() - recording_schema = recording_schema.with_metadata( - { - "log_metadata": json.dumps(asdict(metadata)), - "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), - } - ) - - _write_recording_table(bounding_box_paths, map_api, recording_schema, log_file_path) - - gc.collect() - - convert_log_internal(args) - gc.collect() - return [] - - -def _get_metadata(location: str, log_name: str) -> LogMetadata: - return LogMetadata( - dataset="carla", - log_name=log_name, - location=location, - timestep_seconds=CARLA_DT, - map_has_z=True, - ) - - -def _get_vehicle_parameters(vehicle_parameter_dict: Dict[str, float]) -> VehicleParameters: - # NOTE: @DanielDauner extracting the vehicle parameters from CARLA is somewhat tricky. - # Need to extract the coordinates (for wheels, wheelbase, etc.) from CARLA which is somewhat noise. - # Thus, we hardcode the parameters for the Lincoln MKZ 2020 (default). - assert ( - vehicle_parameter_dict["vehicle_name"] == "vehicle.lincoln.mkz_2020" - ), "Currently only supports MKZ 2020 in CARLA." - return get_carla_lincoln_mkz_2020_parameters() - - -def _get_vehicle_parameters_() -> VehicleParameters: - return get_carla_lincoln_mkz_2020_parameters() - - -def _write_recording_table( - bounding_box_paths: List[Path], - map_api: AbstractMap, - recording_schema: pa.Schema, - log_file_path: Path, -) -> pa.Table: - log_path = bounding_box_paths[0].parent.parent.stem - - with pa.OSFile(str(log_file_path), "wb") as sink: - with pa.ipc.new_file(sink, recording_schema) as writer: - for box_path in bounding_box_paths: - data = _load_json_gz(box_path) - traffic_light_ids, traffic_light_types = _extract_traffic_light_data( - data["traffic_light_states"], data["traffic_light_positions"], map_api - ) - route_lane_group_ids = _extract_route_lane_group_ids(data["route"], map_api) if "route" in data else [] - front_cam_demo, front_cam_transform = _extract_front_cam_demo(box_path) - - row_data = { - "token": [create_token(f"{str(log_path)}_{box_path.stem}")], - "timestamp": [data["timestamp"]], - "detections_state": [_extract_detection_states(data["detections_state"])], - "detections_velocity": [ - ( - data["detections_velocity"] - if "detections_velocity" in data - else np.zeros((len(data["detections_types"]), 3)).tolist() - ) - ], - "detections_token": [data["detections_token"]], - "detections_type": [data["detections_types"]], - "ego_states": [_extract_ego_vehicle_state(data["ego_state"])], - "traffic_light_ids": [traffic_light_ids], - "traffic_light_types": [traffic_light_types], - "scenario_tag": [data["scenario_tag"]], - "route_lane_group_ids": [route_lane_group_ids], - "front_cam_demo": [front_cam_demo], - "front_cam_transform": [front_cam_transform], - } - batch = pa.record_batch(row_data, schema=recording_schema) - writer.write_batch(batch) - del batch, row_data - - if SORT_BY_TIMESTAMP: - recording_table = open_arrow_table(log_file_path) - recording_table = recording_table.sort_by([("timestamp", "ascending")]) - write_arrow_table(recording_table, log_file_path) - - -def _extract_ego_vehicle_state(ego_state_list: List[float]) -> List[float]: - - # NOTE: @DanielDauner This function is a temporary workaround. - # CARLAs bounding box location starts at bottom vertically. - # Need to translate half the height along the z-axis. - ego_state_array = np.array(ego_state_list, dtype=np.float64) - vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() - center = StateSE3.from_array(ego_state_array[EgoStateSE3Index.SE3]) - center = translate_se3_along_z(center, vehicle_parameters.height / 2) - ego_state_array[EgoStateSE3Index.SE3] = center.array - - return ego_state_array.tolist() - - -def _extract_detection_states(detection_states: List[List[float]]) -> List[float]: - - # NOTE: @DanielDauner This function is a temporary workaround. - # CARLAs bounding box location starts at bottom vertically. - # Need to translate half the height along the z-axis. - - detection_state_converted = [] - - for detection_state in detection_states: - detection_state_array = np.array(detection_state, dtype=np.float64) - center = StateSE3.from_array(detection_state_array[BoundingBoxSE3Index.STATE_SE3]) - center = translate_se3_along_z(center, detection_state_array[BoundingBoxSE3Index.HEIGHT] / 2) - detection_state_array[EgoStateSE3Index.SE3] = center.array - detection_state_converted.append(detection_state_array.tolist()) - - return detection_state_converted - - -def _extract_traffic_light_data( - traffic_light_states: List[int], traffic_light_positions: List[List[float]], map_api: AbstractMap -) -> Tuple[List[int], List[int]]: - traffic_light_types: List[int] = [] - traffic_light_ids: List[int] = [] - for traffic_light_state, traffic_light_waypoints in zip(traffic_light_states, traffic_light_positions): - for traffic_light_waypoint in traffic_light_waypoints: - point_3d = Point3D(*traffic_light_waypoint) - nearby_lanes = map_api.get_proximal_map_objects( - point_3d, TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE, [MapSurfaceType.LANE] - )[MapSurfaceType.LANE] - - for lane in nearby_lanes: - lane: AbstractLane - lane_start_point = lane.centerline.array[0] - distance_to_lane_start = np.linalg.norm(lane_start_point - point_3d.array) - if distance_to_lane_start < TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE: - traffic_light_ids.append(int(lane.id)) - traffic_light_types.append(traffic_light_state) - return traffic_light_ids, traffic_light_types - - -def _extract_route_lane_group_ids(route: List[List[float]], map_api: AbstractMap) -> List[int]: - - route = np.array(route, dtype=np.float64) - route[..., 1] = -route[..., 1] # Unreal coordinate system to ISO 8855 - route = route[::2] - - route_lane_group_ids: List[int] = [] - - for point in route[:200]: - point_2d = Point2D(point[0], point[1]) - nearby_lane_groups = map_api.query(point_2d.shapely_point, [MapSurfaceType.LANE_GROUP], predicate="intersects")[ - MapSurfaceType.LANE_GROUP - ] - if len(nearby_lane_groups) == 0: - continue - elif len(nearby_lane_groups) > 1: - possible_lane_group_ids = [lane_group.id for lane_group in nearby_lane_groups] - if len(route_lane_group_ids) > 0: - prev_lane_group_id = route_lane_group_ids[-1] - if prev_lane_group_id in possible_lane_group_ids: - continue - else: - # TODO: Choose with least heading difference? - route_lane_group_ids.append(int(nearby_lane_groups[0].id)) - else: - # TODO: Choose with least heading difference? - route_lane_group_ids.append(int(nearby_lane_groups[0].id)) - elif len(nearby_lane_groups) == 1: - route_lane_group_ids.append(int(nearby_lane_groups[0].id)) - - return list(dict.fromkeys(route_lane_group_ids)) # Remove duplicates while preserving order - - -def _extract_front_cam_demo(box_path: Path) -> Tuple[bytes, List[float]]: - - sample_name = str(box_path.stem).split(".")[0] - sensor_root = Path(box_path.parent.parent) / "rgb" - - front_cam_demo: bytes = None - front_cam_transform: List[float] = [] - - jpg_path = sensor_root / f"{sample_name}.jpg" - - if jpg_path.exists(): - with open(jpg_path, "rb") as f: - front_cam_demo = f.read() - - return front_cam_demo, front_cam_transform diff --git a/d123/dataset/dataset_specific/carla/load_sensor.py b/d123/dataset/dataset_specific/carla/load_sensor.py new file mode 100644 index 00000000..56ff68fa --- /dev/null +++ b/d123/dataset/dataset_specific/carla/load_sensor.py @@ -0,0 +1,10 @@ +from pathlib import Path + +import numpy as np + +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata + + +def load_carla_lidar_from_path(filepath: Path, lidar_metadata: LiDARMetadata) -> LiDAR: + assert filepath.exists(), f"LiDAR file not found: {filepath}" + return LiDAR(metadata=lidar_metadata, point_cloud=np.load(filepath)) diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/enums.py b/d123/dataset/dataset_specific/carla/opendrive/elements/enums.py deleted file mode 100644 index cd8ebe60..00000000 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/enums.py +++ /dev/null @@ -1,45 +0,0 @@ -from enum import Enum, IntEnum - - -class NamedEnum(Enum): - def __init__(self, index, names): - self._index = index - self._names = names - - @property - def index(self): - return self._index - - @property - def names(self): - return self._names - - @classmethod - def from_name(cls, name): - for member in cls: - if name in member.names: - return member - raise ValueError(f"No enum member with name '{name}'") - - -class OpenDriveLaneType(IntEnum): - - BIKING = (0, ("biking",)) - BORDER = (1, ("border",)) - CONNECTING_RAMP = (2, ("connectingRamp",)) - CURB = (3, ("curb",)) - DRIVING = (4, ("driving",)) - ENTRY = (5, ("entry",)) - EXIT = (6, ("exit",)) - MEDIAN = (7, ("median",)) - NONE = (8, ("none",)) - OFF_RAMP = (9, ("offRamp",)) - ON_RAMP = (10, ("onRamp",)) - PARKING = (11, ("parking",)) - RAIL = (12, ("rail",)) - RESTRICTED = (13, ("restricted",)) - SHOULDER = (14, ("shoulder",)) - SIDEWALK = (15, ("sidewalk", "walking")) - SLIP_LANE = (16, ("slipLane",)) - STOP = (17, ("stop",)) - TRAM = (18, ("tram",)) diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/reference.py b/d123/dataset/dataset_specific/carla/opendrive/elements/reference.py deleted file mode 100644 index ebca6ecb..00000000 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/reference.py +++ /dev/null @@ -1,209 +0,0 @@ -from __future__ import annotations - -import xml.etree.ElementTree as ET -from dataclasses import dataclass -from functools import cached_property -from typing import List, Optional, Union -from xml.etree.ElementTree import Element - -import numpy as np -import numpy.typing as npt - -from d123.common.geometry.base import Point3DIndex, StateSE2Index -from d123.common.geometry.utils import normalize_angle -from d123.dataset.dataset_specific.carla.opendrive.elements.elevation import ElevationProfile -from d123.dataset.dataset_specific.carla.opendrive.elements.geometry import Arc, Geometry, Line, Spiral -from d123.dataset.dataset_specific.carla.opendrive.elements.lane import LaneOffset - - -@dataclass -class PlanView: - - geometries: List[Geometry] - - def __post_init__(self): - # NOTE: added assertion/filtering to check for element type or consistency - self.geometries.sort(key=lambda x: x.s, reverse=False) - - @classmethod - def parse(cls, plan_view_element: Optional[Element]) -> PlanView: - - args = {} - geometries: List[Geometry] = [] - for geometry_element in plan_view_element.findall("geometry"): - if geometry_element.find("line") is not None: - geometry = Line.parse(geometry_element) - elif geometry_element.find("arc") is not None: - geometry = Arc.parse(geometry_element) - elif geometry_element.find("spiral") is not None: - geometry = Spiral.parse(geometry_element) - else: - geometry_str = ET.tostring(geometry_element, encoding="unicode") - raise NotImplementedError(f"Geometry not implemented: {geometry_str}") - geometries.append(geometry) - args["geometries"] = geometries - return PlanView(**args) - - @cached_property - def geometry_lengths(self) -> npt.NDArray[np.float64]: - return np.cumsum([0.0] + [geo.length for geo in self.geometries], dtype=np.float64) - - @property - def length(self) -> float: - return float(self.geometry_lengths[-1]) - - def interpolate_se2(self, s: float, t: float = 0.0, is_last_pos: bool = False) -> npt.NDArray[np.float64]: - - try: - # get index of geometry which is at s_pos - mask = self.geometry_lengths > s - sub_idx = np.argmin(self.geometry_lengths[mask] - s) - geo_idx = np.arange(self.geometry_lengths.shape[0])[mask][sub_idx] - 1 - except ValueError: - # s_pos is after last geometry because of rounding error - if np.isclose(s, self.geometry_lengths[-1], 0.01, 0.01): # todo parameter - geo_idx = self.geometry_lengths.size - 2 - else: - raise Exception( - f"Tried to calculate a position outside of the borders of the reference path at s={s}" - f", but path has only length of l={self.length}" - ) - - return self.geometries[geo_idx].interpolate_se2(s - self.geometry_lengths[geo_idx], t) - - -@dataclass -class Border: - - reference: Union[Border, PlanView] - elevation_profile: ElevationProfile - - width_coefficient_offsets: List[float] - width_coefficients: List[List[float]] - - s_offset: float = 0.0 - - # NOTE: loaded in __post_init__ - length: Optional[float] = None - - def __post_init__(self): - # NOTE: added assertion/filtering to check for element type or consistency - self.length = float(self.reference.length) - - @classmethod - def from_plan_view( - cls, - plan_view: PlanView, - lane_offsets: List[LaneOffset], - elevation_profile: ElevationProfile, - ) -> Border: - args = {} - args["reference"] = plan_view - args["elevation_profile"] = elevation_profile - - width_coefficient_offsets = [] - width_coefficients = [] - - # Lane offsets will be coeffs - # this has to be done if the reference path has the laneoffset attribute - # and thus is different to the geometry described in the plan_view - # openDRIVE lets multiple laneOffsets start at the same position - # but only the last one counts -> delete all previous ones - if any(lane_offsets): - for lane_offset in lane_offsets: - if lane_offset.s in width_coefficient_offsets: - # offset is already there, delete previous entries - idx = width_coefficient_offsets.index(lane_offset.s) - del width_coefficient_offsets[idx] - del width_coefficients[idx] - width_coefficient_offsets.append(lane_offset.s) - width_coefficients.append(lane_offset.polynomial_coefficients) - else: - width_coefficient_offsets.append(0.0) - width_coefficients.append([0.0]) - - args["width_coefficient_offsets"] = width_coefficient_offsets - args["width_coefficients"] = width_coefficients - - return Border(**args) - - def _get_width_index(self, s: float, is_last_pos: bool) -> float: - """Get the index of the width which applies at position s_pos. - - :param s_pos: Position on border in curve_parameter ds - :param is_last_pos: Whether s_pos is the last position - :return: Index of width that applies at position s_pos - """ - - return next( - ( - self.width_coefficient_offsets.index(n) - for n in self.width_coefficient_offsets[::-1] - if ((n <= s and (not is_last_pos or s == 0)) or (n < s and is_last_pos)) - ), - len(self.width_coefficient_offsets) - 1, - ) - - def _get_height_index(self, s: float, is_last_pos: bool) -> float: - height_offsets = [elevation.s for elevation in self.elevation_profile.elevations] - - return next( - ( - height_offsets.index(n) - for n in height_offsets[::-1] - if ((n <= s and (not is_last_pos or s == 0)) or (n < s and is_last_pos)) - ), - len(height_offsets) - 1, - ) - - def interpolate_se2(self, s: float, t: float = 0.0, is_last_pos: bool = False) -> npt.NDArray[np.float64]: - # Last reference has to be a reference geometry (PlanView) - # Offset of all inner lanes (Border) - # calculate position of reference border - if np.isclose(s, 0): - s = 0 - - try: - se2 = self.reference.interpolate_se2(self.s_offset + s, is_last_pos=is_last_pos) - except TypeError: - se2 = self.reference.interpolate_se2(np.round(self.s_offset + s, 3), is_last_pos=is_last_pos) - - if len(self.width_coefficients) == 0 or len(self.width_coefficient_offsets) == 0: - raise Exception("No entries for width definitions.") - - # Find correct coefficients - # find which width segment is at s_pos - width_idx = self._get_width_index(s, is_last_pos=is_last_pos) - # width_idx = min(width_idx, len(self.width_coefficient_offsets)-1) - # Calculate width at s_pos - distance = ( - np.polynomial.polynomial.polyval( - s - self.width_coefficient_offsets[width_idx], - self.width_coefficients[width_idx], - ) - + t - ) - ortho = normalize_angle(se2[StateSE2Index.YAW] + np.pi / 2) - se2[StateSE2Index.X] += distance * np.cos(ortho) - se2[StateSE2Index.Y] += distance * np.sin(ortho) - se2[StateSE2Index.YAW] = normalize_angle(se2[StateSE2Index.YAW]) - - return se2 - - def interpolate_3d(self, s: float, t: float = 0.0, is_last_pos: bool = False) -> npt.NDArray[np.float64]: - - point_3d = np.zeros(len(Point3DIndex), dtype=np.float64) - - se2 = self.interpolate_se2(s, t, is_last_pos=is_last_pos) - point_3d[Point3DIndex.X] = se2[StateSE2Index.X] - point_3d[Point3DIndex.Y] = se2[StateSE2Index.Y] - - if len(self.elevation_profile.elevations) > 0: - height_index = self._get_height_index(s, is_last_pos=is_last_pos) - elevation = self.elevation_profile.elevations[height_index] - - point_3d[Point3DIndex.Z] = np.polynomial.polynomial.polyval( - s - elevation.s, elevation.polynomial_coefficients - ) - - return point_3d diff --git a/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py b/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py deleted file mode 100644 index cd22aebe..00000000 --- a/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py +++ /dev/null @@ -1,583 +0,0 @@ -import os -import warnings -from pathlib import Path -from typing import Dict, List, Optional - -import geopandas as gpd -import numpy as np -import pandas as pd -import shapely -from shapely.ops import polygonize, unary_union - -from d123.dataset.dataset_specific.carla.opendrive.conversion.group_collections import ( - OpenDriveLaneGroupHelper, - OpenDriveLaneHelper, - lane_section_to_lane_helpers, -) -from d123.dataset.dataset_specific.carla.opendrive.conversion.id_system import ( - build_lane_id, - derive_lane_section_id, - lane_group_id_from_lane_id, - road_id_from_lane_group_id, -) -from d123.dataset.dataset_specific.carla.opendrive.conversion.objects_collections import ( - OpenDriveObjectHelper, - get_object_helper, -) -from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import Junction, OpenDrive, Road -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border -from d123.dataset.dataset_specific.carla.opendrive.id_mapping import IntIDMapping -from d123.dataset.maps.map_datatypes import MapSurfaceType - -ENABLE_WARNING: bool = False -CONNECTION_DISTANCE_THRESHOLD: float = 0.1 # [m] -D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) - - -class OpenDriveConverter: - def __init__(self, opendrive: OpenDrive): - - self.opendrive: OpenDrive = opendrive - - self.road_dict: Dict[int, Road] = {road.id: road for road in opendrive.roads} - self.junction_dict: Dict[int, Junction] = {junction.id: junction for junction in opendrive.junctions} - - # loaded during conversion - self.lane_helper_dict: Dict[str, OpenDriveLaneHelper] = {} - self.lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = {} - self.object_helper_dict: Dict[str, OpenDriveObjectHelper] = {} - - def run(self, map_name: str) -> None: - - # Run processing for map elements - self._collect_lane_helpers() - self._update_connection_from_links() - self._update_connection_from_junctions() - self._flip_and_set_connections() - self._post_process_connections() - self._collect_lane_groups() - self._collect_crosswalks() - - # Collect data frames and store - lane_df = self._extract_lane_dataframe() - walkways_df = self._extract_walkways_dataframe() - carpark_df = self._extract_carpark_dataframe() - generic_drivable_area_df = self._extract_generic_drivable_dataframe() - intersections_df = self._extract_intersections_dataframe() - lane_group_df = self._extract_lane_group_dataframe() - crosswalk_df = self._extract_crosswalk_dataframe() - - self._convert_ids_to_int( - lane_df, - walkways_df, - carpark_df, - generic_drivable_area_df, - lane_group_df, - intersections_df, - crosswalk_df, - ) - - # Store dataframes - map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" - lane_df.to_file(map_file_name, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") - walkways_df.to_file(map_file_name, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_name, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_area_df.to_file( - map_file_name, - layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), - driver="GPKG", - mode="a", - ) - intersections_df.to_file(map_file_name, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") - lane_group_df.to_file(map_file_name, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_name, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") - - def _collect_lane_helpers(self) -> None: - for road in self.opendrive.roads: - reference_border = Border.from_plan_view(road.plan_view, road.lanes.lane_offsets, road.elevation_profile) - lane_section_lengths: List[float] = [ls.s for ls in road.lanes.lane_sections] + [road.length] - for idx, lane_section in enumerate(road.lanes.lane_sections): - lane_section_id = derive_lane_section_id(road.id, idx) - lane_helpers_ = lane_section_to_lane_helpers( - lane_section_id, - lane_section, - reference_border, - lane_section_lengths[idx], - lane_section_lengths[idx + 1], - road.road_types, - ) - self.lane_helper_dict.update(lane_helpers_) - - def _update_connection_from_links(self) -> None: - - for lane_id in self.lane_helper_dict.keys(): - road_idx, lane_section_idx, side, lane_idx = lane_id.split("_") - road_idx, lane_section_idx, lane_idx = int(road_idx), int(lane_section_idx), int(lane_idx) - - road = self.road_dict[road_idx] - - successor_lane_idx = self.lane_helper_dict[lane_id].open_drive_lane.successor - if successor_lane_idx is not None: - successor_lane_id: Optional[str] = None - - # Not the last lane section -> Next lane section in same road - if lane_section_idx < road.lanes.last_lane_section_idx: - successor_lane_id = build_lane_id( - road_idx, - lane_section_idx + 1, - successor_lane_idx, - ) - - # Last lane section -> Next road in first lane section - # Try to get next road - elif road.link.successor is not None and road.link.successor.element_type != "junction": - successor_road = self.road_dict[road.link.successor.element_id] - - if road.link.successor.contact_point == "start": - successor_lane_section_idx = 0 - else: - successor_lane_section_idx = successor_road.lanes.last_lane_section_idx - - successor_lane_id = build_lane_id( - successor_road.id, - successor_lane_section_idx, - successor_lane_idx, - ) - - # assert successor_lane_id in self.lane_helper_dict.keys() - if successor_lane_id is None or successor_lane_id not in self.lane_helper_dict.keys(): - continue - self.lane_helper_dict[lane_id].successor_lane_ids.append(successor_lane_id) - self.lane_helper_dict[successor_lane_id].predecessor_lane_ids.append(lane_id) - - predecessor_lane_idx = self.lane_helper_dict[lane_id].open_drive_lane.predecessor - if predecessor_lane_idx is not None: - predecessor_lane_id: Optional[str] = None - - # Not the first lane section -> Previous lane section in same road - if lane_section_idx > 0: - predecessor_lane_id = build_lane_id( - road_idx, - lane_section_idx - 1, - predecessor_lane_idx, - ) - - # First lane section -> Previous road - # Try to get previous road - elif road.link.predecessor is not None and road.link.predecessor.element_type != "junction": - predecessor_road = self.road_dict[road.link.predecessor.element_id] - - if road.link.predecessor.contact_point == "start": - predecessor_lane_section_idx = 0 - else: - predecessor_lane_section_idx = predecessor_road.lanes.last_lane_section_idx - - predecessor_lane_id = build_lane_id( - predecessor_road.id, - predecessor_lane_section_idx, - predecessor_lane_idx, - ) - - # assert predecessor_lane_id in self.lane_helper_dict.keys() - if predecessor_lane_id is None or predecessor_lane_id not in self.lane_helper_dict.keys(): - continue - self.lane_helper_dict[lane_id].predecessor_lane_ids.append(predecessor_lane_id) - self.lane_helper_dict[predecessor_lane_id].successor_lane_ids.append(lane_id) - - def _update_connection_from_junctions(self) -> None: - - # add junctions to link_index - # if contact_point is start, and laneId from connecting_road is negative - # the connecting_road is the successor - # if contact_point is start, and laneId from connecting_road is positive - # the connecting_road is the predecessor - # for contact_point == end it's exactly the other way - for junction_idx, junction in self.junction_dict.items(): - for connection in junction.connections: - incoming_road = self.road_dict[connection.incoming_road] - connecting_road = self.road_dict[connection.connecting_road] - contact_point = connection.contact_point - for lane_link in connection.lane_links: - - incoming_lane_id: Optional[str] = None - connecting_lane_id: Optional[str] = None - - if contact_point == "start": - if lane_link.start < 0: - incoming_lane_section_idx = incoming_road.lanes.last_lane_section_idx - else: - incoming_lane_section_idx = 0 - incoming_lane_id = build_lane_id(incoming_road.id, incoming_lane_section_idx, lane_link.start) - connecting_lane_id = build_lane_id(connecting_road.id, 0, lane_link.end) - else: - incoming_lane_id = build_lane_id(incoming_road.id, 0, lane_link.start) - connecting_lane_id = build_lane_id( - connecting_road.id, connecting_road.lanes.last_lane_section_idx, lane_link.end - ) - - if incoming_lane_id is None or connecting_lane_id is None: - continue - if ( - incoming_lane_id not in self.lane_helper_dict.keys() - or connecting_lane_id not in self.lane_helper_dict.keys() - ): - if ENABLE_WARNING: - warnings.warn( - f"Warning..... Lane connection {incoming_lane_id} -> {connecting_lane_id} not found in lane_helper_dict" - ) - continue - # assert incoming_lane_id in self.lane_helper_dict.keys() - # assert connecting_lane_id in self.lane_helper_dict.keys() - self.lane_helper_dict[incoming_lane_id].successor_lane_ids.append(connecting_lane_id) - self.lane_helper_dict[connecting_lane_id].predecessor_lane_ids.append(incoming_lane_id) - - def _flip_and_set_connections(self) -> None: - - for lane_id in self.lane_helper_dict.keys(): - if self.lane_helper_dict[lane_id].id > 0: - successors_temp = self.lane_helper_dict[lane_id].successor_lane_ids - self.lane_helper_dict[lane_id].successor_lane_ids = self.lane_helper_dict[lane_id].predecessor_lane_ids - self.lane_helper_dict[lane_id].predecessor_lane_ids = successors_temp - self.lane_helper_dict[lane_id].successor_lane_ids = list( - set(self.lane_helper_dict[lane_id].successor_lane_ids) - ) - self.lane_helper_dict[lane_id].predecessor_lane_ids = list( - set(self.lane_helper_dict[lane_id].predecessor_lane_ids) - ) - - def _post_process_connections(self) -> None: - - for lane_id in self.lane_helper_dict.keys(): - self.lane_helper_dict[lane_id] - - centerline = self.lane_helper_dict[lane_id].center_polyline_se2 - - valid_successor_lane_ids: List[str] = [] - for successor_lane_id in self.lane_helper_dict[lane_id].successor_lane_ids: - successor_centerline = self.lane_helper_dict[successor_lane_id].center_polyline_se2 - distance = np.linalg.norm(centerline[-1, :2] - successor_centerline[0, :2]) - if distance > CONNECTION_DISTANCE_THRESHOLD: - if ENABLE_WARNING: - warnings.warn( - f"Warning..... Removing connection {lane_id} -> {successor_lane_id} with distance {distance}" - ) - else: - valid_successor_lane_ids.append(successor_lane_id) - self.lane_helper_dict[lane_id].successor_lane_ids = valid_successor_lane_ids - - valid_predecessor_lane_ids: List[str] = [] - for predecessor_lane_id in self.lane_helper_dict[lane_id].predecessor_lane_ids: - predecessor_centerline = self.lane_helper_dict[predecessor_lane_id].center_polyline_se2 - distance = np.linalg.norm(centerline[0, :2] - predecessor_centerline[-1, :2]) - if distance > CONNECTION_DISTANCE_THRESHOLD: - if ENABLE_WARNING: - warnings.warn( - f"Warning..... Removing connection {predecessor_lane_id} -> {successor_lane_id} with distance {distance}" - ) - else: - valid_predecessor_lane_ids.append(predecessor_lane_id) - self.lane_helper_dict[lane_id].predecessor_lane_ids = valid_predecessor_lane_ids - - def _collect_lane_groups(self) -> None: - def _collect_lane_helper_of_id(lane_group_id: str) -> List[OpenDriveLaneHelper]: - lane_helpers: List[OpenDriveLaneHelper] = [] - for lane_id, lane_helper in self.lane_helper_dict.items(): - if (lane_helper.type in ["driving"]) and (lane_group_id_from_lane_id(lane_id) == lane_group_id): - lane_helpers.append(lane_helper) - return lane_helpers - - all_lane_group_ids = list( - set([lane_group_id_from_lane_id(lane_id) for lane_id in self.lane_helper_dict.keys()]) - ) - - for lane_group_id in all_lane_group_ids: - lane_group_lane_helper = _collect_lane_helper_of_id(lane_group_id) - if len(lane_group_lane_helper) >= 1: - self.lane_group_helper_dict[lane_group_id] = OpenDriveLaneGroupHelper( - lane_group_id, lane_group_lane_helper - ) - - def _collect_lane_group_ids_of_road(road_id: int) -> List[str]: - lane_group_ids: List[str] = [] - for lane_group_id in self.lane_group_helper_dict.keys(): - if int(road_id_from_lane_group_id(lane_group_id)) == road_id: - lane_group_ids.append(lane_group_id) - return lane_group_ids - - for junction in self.junction_dict.values(): - for connection in junction.connections: - connecting_road = self.road_dict[connection.connecting_road] - connecting_lane_group_ids = _collect_lane_group_ids_of_road(connecting_road.id) - for connecting_lane_group_id in connecting_lane_group_ids: - if connecting_lane_group_id in self.lane_group_helper_dict.keys(): - self.lane_group_helper_dict[connecting_lane_group_id].junction_id = junction.id - - def _collect_crosswalks(self) -> None: - for road in self.opendrive.roads: - if len(road.objects) == 0: - continue - reference_border = Border.from_plan_view(road.plan_view, road.lanes.lane_offsets, road.elevation_profile) - for object in road.objects: - if object.type in ["crosswalk"]: - object_helper = get_object_helper(object, reference_border) - self.object_helper_dict[object_helper.object_id] = object_helper - - def _extract_lane_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - lane_group_ids = [] - speed_limits_mps = [] - predecessor_ids = [] - successor_ids = [] - left_boundaries = [] - right_boundaries = [] - baseline_paths = [] - geometries = [] - - # TODO: Extract speed limit and convert to mps - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type == "driving": - ids.append(lane_helper.lane_id) - lane_group_ids.append(lane_group_id_from_lane_id(lane_helper.lane_id)) - speed_limits_mps.append(lane_helper.speed_limit_mps) - predecessor_ids.append(lane_helper.predecessor_lane_ids) - successor_ids.append(lane_helper.successor_lane_ids) - left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - baseline_paths.append(shapely.LineString(lane_helper.center_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - "lane_group_id": lane_group_ids, - "speed_limit_mps": speed_limits_mps, - "predecessor_ids": predecessor_ids, - "successor_ids": successor_ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - "baseline_path": baseline_paths, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_walkways_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - # left_boundaries = [] - # right_boundaries = [] - outlines = [] - geometries = [] - - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type == "sidewalk": - ids.append(lane_helper.lane_id) - # left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - # right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - # "left_boundary": left_boundaries, - # "right_boundary": left_boundaries, - "outline": outlines, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_carpark_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - # left_boundaries = [] - # right_boundaries = [] - outlines = [] - geometries = [] - - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type == "parking": - ids.append(lane_helper.lane_id) - # left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - # right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - # "left_boundary": left_boundaries, - # "right_boundary": right_boundaries, - "outline": outlines, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_generic_drivable_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - # left_boundaries = [] - # right_boundaries = [] - outlines = [] - geometries = [] - - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type in ["none", "border", "bidirectional"]: - ids.append(lane_helper.lane_id) - # left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - # right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - # "left_boundary": left_boundaries, - # "right_boundary": left_boundaries, - "outline": outlines, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_intersections_dataframe(self) -> gpd.GeoDataFrame: - def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriveLaneGroupHelper]: - return [ - lane_group_helper - for lane_group_helper in self.lane_group_helper_dict.values() - if lane_group_helper.junction_id == junction_id - ] - - ids = [] - lane_group_ids = [] - geometries = [] - for junction in self.junction_dict.values(): - lane_group_helpers = _find_lane_group_helpers_with_junction_id(junction.id) - lane_group_ids_ = [lane_group_helper.lane_group_id for lane_group_helper in lane_group_helpers] - if len(lane_group_ids_) == 0: - warnings.warn(f"Skipped Junction {junction.id} without drivable lanes!") - continue - - polygon = extract_exteriors_polygon(lane_group_helpers) - ids.append(junction.id) - lane_group_ids.append(lane_group_ids_) - geometries.append(polygon) - - data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_lane_group_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - lane_ids = [] - predecessor_lane_group_ids = [] - successor_lane_group_ids = [] - intersection_ids = [] - left_boundaries = [] - right_boundaries = [] - geometries = [] - - for lane_group_helper in self.lane_group_helper_dict.values(): - lane_group_helper: OpenDriveLaneGroupHelper - ids.append(lane_group_helper.lane_group_id) - lane_ids.append([lane_helper.lane_id for lane_helper in lane_group_helper.lane_helpers]) - predecessor_lane_group_ids.append(lane_group_helper.predecessor_lane_group_ids) - successor_lane_group_ids.append(lane_group_helper.successor_lane_group_ids) - intersection_ids.append(lane_group_helper.junction_id) - left_boundaries.append(shapely.LineString(lane_group_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_group_helper.outer_polyline_3d)) - geometries.append(lane_group_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - "lane_ids": lane_ids, - "predecessor_ids": predecessor_lane_group_ids, - "successor_ids": successor_lane_group_ids, - "intersection_id": intersection_ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - } - ) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - - return gdf - - def _extract_crosswalk_dataframe(self) -> gpd.GeoDataFrame: - ids = [] - outlines = [] - geometries = [] - for object_helper in self.object_helper_dict.values(): - ids.append(object_helper.object_id) - outlines.append(shapely.LineString(object_helper.outline_3d)) - geometries.append(object_helper.shapely_polygon) - - data = pd.DataFrame({"id": ids, "outline": outlines}) - return gpd.GeoDataFrame(data, geometry=geometries) - - @staticmethod - def _convert_ids_to_int( - lane_df: gpd.GeoDataFrame, - walkways_df: gpd.GeoDataFrame, - carpark_df: gpd.GeoDataFrame, - generic_drivable_area_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, - intersections_df: gpd.GeoDataFrame, - crosswalk_df: gpd.GeoDataFrame, - ) -> None: - - # NOTE: intersection and crosswalk ids are already integers - - # initialize id mappings - lane_id_mapping = IntIDMapping.from_series(lane_df["id"]) - walkway_id_mapping = IntIDMapping.from_series(walkways_df["id"]) - carpark_id_mapping = IntIDMapping.from_series(carpark_df["id"]) - # TODO: add id mapping for intersections - generic_drivable_id_mapping = IntIDMapping.from_series(generic_drivable_area_df["id"]) - lane_group_id_mapping = IntIDMapping.from_series(lane_group_df["id"]) - - # Adjust cross reference in lane_df and lane_group_df - lane_df["lane_group_id"] = lane_df["lane_group_id"].map(lane_group_id_mapping.str_to_int) - lane_group_df["lane_ids"] = lane_group_df["lane_ids"].apply(lambda x: lane_id_mapping.map_list(x)) - - # Adjust predecessor/successor in lane_df and lane_group_df - for column in ["predecessor_ids", "successor_ids"]: - lane_df[column] = lane_df[column].apply(lambda x: lane_id_mapping.map_list(x)) - lane_group_df[column] = lane_group_df[column].apply(lambda x: lane_group_id_mapping.map_list(x)) - - lane_df["id"] = lane_df["id"].map(lane_id_mapping.str_to_int) - walkways_df["id"] = walkways_df["id"].map(walkway_id_mapping.str_to_int) - carpark_df["id"] = carpark_df["id"].map(carpark_id_mapping.str_to_int) - generic_drivable_area_df["id"] = generic_drivable_area_df["id"].map(generic_drivable_id_mapping.str_to_int) - lane_group_df["id"] = lane_group_df["id"].map(lane_group_id_mapping.str_to_int) - - intersections_df["lane_group_ids"] = intersections_df["lane_group_ids"].apply( - lambda x: lane_group_id_mapping.map_list(x) - ) - - -# TODO: move this somewhere else and improve -def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper]) -> shapely.Polygon: - - # Step 1: Extract all boundary line segments - all_polygons = [] - for lane_group_helper in lane_group_helpers: - all_polygons.append(lane_group_helper.shapely_polygon) - - # Step 2: Merge all boundaries and extract the enclosed polygons - # try: - merged_boundaries = unary_union(all_polygons) - # except Exception as e: - # warnings.warn(f"Topological error during polygon union: {e}") - # print([(helper.lane_group_id, poly.is_valid) for poly, helper in zip(all_polygons, lane_group_helpers)]) - # merged_boundaries = unary_union([poly for poly in all_polygons if poly.is_valid]) - - # Step 3: Generate polygons from the merged lines - polygons = list(polygonize(merged_boundaries)) - - # Step 4: Select the polygon that represents the intersection - # Usually it's the largest polygon - if len(polygons) == 1: - return polygons[0] - else: - # Take the largest polygon if there are multiple - return max(polygons, key=lambda p: p.area) diff --git a/d123/dataset/dataset_specific/kitti_360/.gitkeep b/d123/dataset/dataset_specific/kitti_360/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/dataset_specific/nuplan/load_sensor.py b/d123/dataset/dataset_specific/nuplan/load_sensor.py new file mode 100644 index 00000000..0bbdc406 --- /dev/null +++ b/d123/dataset/dataset_specific/nuplan/load_sensor.py @@ -0,0 +1,21 @@ +import io +from pathlib import Path + +from nuplan.database.utils.pointclouds.lidar import LidarPointCloud + +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata + + +def load_nuplan_lidar_from_path(filepath: Path, lidar_metadata: LiDARMetadata) -> LiDAR: + assert filepath.exists(), f"LiDAR file not found: {filepath}" + with open(filepath, "rb") as fp: + buffer = io.BytesIO(fp.read()) + return LiDAR(metadata=lidar_metadata, point_cloud=LidarPointCloud.from_buffer(buffer, "pcd").points) + + +# def load_camera_from_path(filename: str, metadata: CameraMetadata) -> Camera: +# camera_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / filename +# assert camera_full_path.exists(), f"Camera file not found: {camera_full_path}" +# img = Image.open(camera_full_path) +# img.load() +# return Camera(metadata=metadata, image=np.asarray(img, dtype=np.uint8)) diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py similarity index 53% rename from d123/dataset/dataset_specific/nuplan/nuplan_data_processor.py rename to d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index e345e3fc..23a4d702 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -1,50 +1,50 @@ import gc import json import os +import pickle from dataclasses import asdict from functools import partial from pathlib import Path -from typing import Dict, Final, List, Tuple, Union +from typing import Any, Dict, Final, List, Optional, Tuple, Union -import psutil +import numpy as np import pyarrow as pa import yaml -from nuplan.database.nuplan_db.nuplan_scenario_queries import get_images_from_lidar_tokens +from nuplan.database.nuplan_db.nuplan_scenario_queries import get_cameras, get_images_from_lidar_tokens +from nuplan.database.nuplan_db_orm.ego_pose import EgoPose from nuplan.database.nuplan_db_orm.lidar_box import LidarBox from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB -from nuplan.planning.simulation.observation.observation_type import ( - CameraChannel, -) -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map +from nuplan.planning.simulation.observation.observation_type import CameraChannel +from pyquaternion import Quaternion +from sqlalchemy import func import d123.dataset.dataset_specific.nuplan.utils as nuplan_utils from d123.common.datatypes.detection.detection import TrafficLightStatus from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar_index import NuplanLidarIndex +from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( - get_nuplan_pacifica_parameters, + get_nuplan_chrysler_pacifica_parameters, rear_axle_se3_to_center_se3, ) from d123.common.geometry.base import StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3, BoundingBoxSE3Index from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor +from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import MAP_LOCATIONS, NuPlanMapConverter +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata TARGET_DT: Final[float] = 0.1 NUPLAN_DT: Final[float] = 0.05 SORT_BY_TIMESTAMP: Final[bool] = True -NUPLAN_FULL_MAP_NAME_DICT: Final[Dict[str, str]] = { - "boston": "us-ma-boston", - "singapore": "sg-one-north", - "las_vegas": "us-nv-las-vegas-strip", - "pittsburgh": "us-pa-pittsburgh-hazelwood", -} - NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { "green": TrafficLightStatus.GREEN, "red": TrafficLightStatus.RED, @@ -60,7 +60,19 @@ "generic_object": DetectionType.GENERIC_OBJECT, } +NUPLAN_CAMERA_TYPES = { + CameraType.CAM_F0: CameraChannel.CAM_F0, + CameraType.CAM_B0: CameraChannel.CAM_B0, + CameraType.CAM_L0: CameraChannel.CAM_L0, + CameraType.CAM_L1: CameraChannel.CAM_L1, + CameraType.CAM_L2: CameraChannel.CAM_L2, + CameraType.CAM_R0: CameraChannel.CAM_R0, + CameraType.CAM_R1: CameraChannel.CAM_R1, + CameraType.CAM_R2: CameraChannel.CAM_R2, +} + NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) +NUPLAN_ROLLING_SHUTTER_S: Final[TimePoint] = TimePoint.from_s(1 / 60) def create_splits_logs() -> Dict[str, List[str]]: @@ -71,16 +83,14 @@ def create_splits_logs() -> Dict[str, List[str]]: return splits["log_splits"] -class NuplanDataProcessor(RawDataProcessor): +class NuplanDataConverter(RawDataConverter): def __init__( self, splits: List[str], log_path: Union[Path, str], - output_path: Union[Path, str], - sensor_path: Union[Path, str], - force_data_conversion: bool, + data_converter_config: DataConverterConfig, ) -> None: - super().__init__(force_data_conversion) + super().__init__(data_converter_config) for split in splits: assert ( split in self.get_available_splits() @@ -88,8 +98,6 @@ def __init__( self._splits: List[str] = splits self._log_path: Path = Path(log_path) - self._output_path: Path = Path(output_path) - self._sensor_path: Path = Path(sensor_path) self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() self._target_dt: float = 0.1 @@ -131,7 +139,14 @@ def get_available_splits(self) -> List[str]: "nuplan_private_test", ] - def convert(self, worker: WorkerPool) -> None: + def convert_maps(self, worker: WorkerPool) -> None: + worker_map( + worker, + partial(convert_nuplan_map_to_gpkg, data_converter_config=self.data_converter_config), + list(MAP_LOCATIONS), + ) + + def convert_logs(self, worker: WorkerPool) -> None: log_args = [ { "log_path": log_path, @@ -145,19 +160,24 @@ def convert(self, worker: WorkerPool) -> None: worker, partial( convert_nuplan_log_to_arrow, - output_path=self._output_path, - force_data_conversion=self.force_data_conversion, + data_converter_config=self.data_converter_config, ), log_args, ) +def convert_nuplan_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: + for map_name in map_names: + map_path = data_converter_config.output_path / "maps" / f"nuplan_{map_name}.gpkg" + if data_converter_config.force_map_conversion or not map_path.exists(): + map_path.unlink(missing_ok=True) + NuPlanMapConverter(data_converter_config.output_path / "maps").convert(map_name=map_name) + return [] + + def convert_nuplan_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], - output_path: Path, - force_data_conversion: bool, -) -> None: - process = psutil.Process(os.getpid()) + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig +) -> List[Any]: for log_info in args: log_path: Path = log_info["log_path"] split: str = log_info["split"] @@ -166,30 +186,13 @@ def convert_nuplan_log_to_arrow( raise FileNotFoundError(f"Log path {log_path} does not exist.") log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None) - log_file_path = output_path / split / f"{log_path.stem}.arrow" + log_file_path = data_converter_config.output_path / split / f"{log_path.stem}.arrow" - if force_data_conversion or not log_file_path.exists(): + if data_converter_config.force_log_conversion or not log_file_path.exists(): log_file_path.unlink(missing_ok=True) if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) - recording_schema = pa.schema( - [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ("front_cam_demo", pa.binary()), - ("front_cam_transform", pa.list_(pa.float64())), - ] - ) metadata = LogMetadata( dataset="nuplan", log_name=log_db.log_name, @@ -197,29 +200,99 @@ def convert_nuplan_log_to_arrow( timestep_seconds=TARGET_DT, map_has_z=False, ) - vehicle_parameters = get_nuplan_pacifica_parameters() + vehicle_parameters = get_nuplan_chrysler_pacifica_parameters() + camera_metadata = get_nuplan_camera_metadata(log_path) + lidar_metadata = get_nuplan_lidar_metadata(log_db) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + schema_column_list.append((lidar_type.serialize(), pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + if data_converter_config.camera_store_option is not None: + for camera_type in camera_metadata.keys(): + if data_converter_config.camera_store_option == "path": + schema_column_list.append((camera_type.serialize(), pa.string())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) + ) + + elif data_converter_config.camera_store_option == "binary": + raise NotImplementedError("Binary camera storage is not implemented.") + + recording_schema = pa.schema(schema_column_list) recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), } ) - _write_recording_table(log_db, recording_schema, log_file_path, log_path) + _write_recording_table(log_db, recording_schema, log_file_path, log_path, data_converter_config) - log_db.detach_tables() - log_db.remove_ref() - del recording_schema, vehicle_parameters, log_db + log_db.detach_tables() + log_db.remove_ref() + del recording_schema, vehicle_parameters, log_db gc.collect() - print(f"{os.getpid()} Memory usage: {process.memory_info().rss / 1024 / 1024:.1f} MB") return [] +def get_nuplan_camera_metadata(log_path: Path) -> Dict[CameraType, CameraMetadata]: + + def _get_camera_metadata(camera_type: CameraType) -> CameraMetadata: + cam = list(get_cameras(log_path, [str(NUPLAN_CAMERA_TYPES[camera_type].value)]))[0] + intrinsic = np.array(pickle.loads(cam.intrinsic)) + rotation = np.array(pickle.loads(cam.rotation)) + rotation = Quaternion(rotation).rotation_matrix + distortion = np.array(pickle.loads(cam.distortion)) + return CameraMetadata( + camera_type=camera_type, + width=cam.width, + height=cam.height, + intrinsic=intrinsic, + distortion=distortion, + ) + + log_cam_infos: Dict[str, CameraMetadata] = {} + for camera_type in NUPLAN_CAMERA_TYPES.keys(): + log_cam_infos[camera_type] = _get_camera_metadata(camera_type) + + return log_cam_infos + + +def get_nuplan_lidar_metadata(log_db: NuPlanDB) -> Dict[LiDARType, LiDARMetadata]: + metadata: Dict[LiDARType, LiDARMetadata] = {} + metadata[LiDARType.LIDAR_MERGED] = LiDARMetadata( + lidar_type=LiDARType.LIDAR_MERGED, + lidar_index=NuplanLidarIndex, + extrinsic=None, # NOTE: LiDAR extrinsic are unknown + ) + return metadata + + def _write_recording_table( log_db: NuPlanDB, recording_schema: pa.schema, log_file_path: Path, source_log_path: Path, + data_converter_config: DataConverterConfig, ) -> None: # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: @@ -228,16 +301,18 @@ def _write_recording_table( step_interval: float = int(TARGET_DT / NUPLAN_DT) for lidar_pc in log_db.lidar_pc[::step_interval]: lidar_pc_token: str = lidar_pc.token - detections_state, detections_velocity, detections_token, detections_types = _extract_detections( - lidar_pc - ) + ( + detections_state, + detections_velocity, + detections_token, + detections_types, + ) = _extract_detections(lidar_pc) traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) route_lane_group_ids = [ int(roadblock_id) for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") if len(roadblock_id) > 0 ] - front_cam_demo, front_cam_transform = _extract_front_cam_demo(lidar_pc, source_log_path) row_data = { "token": [lidar_pc_token], @@ -251,9 +326,26 @@ def _write_recording_table( "traffic_light_types": [traffic_light_types], "scenario_tag": [_extract_scenario_tag(log_db, lidar_pc_token)], "route_lane_group_ids": [route_lane_group_ids], - "front_cam_demo": [front_cam_demo], - "front_cam_transform": [front_cam_transform], } + + if data_converter_config.lidar_store_option is not None: + lidar_data_dict = _extract_lidar(lidar_pc, data_converter_config) + for lidar_type, lidar_data in lidar_data_dict.items(): + if lidar_data is not None: + row_data[lidar_type.serialize()] = [lidar_data] + else: + row_data[lidar_type.serialize()] = [None] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_camera(log_db, lidar_pc, source_log_path, data_converter_config) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] + else: + row_data[camera_type.serialize()] = [None] + row_data[f"{camera_type.serialize()}_extrinsic"] = [None] + batch = pa.record_batch(row_data, schema=recording_schema) writer.write_batch(batch) del batch, row_data, detections_state, detections_velocity, detections_token, detections_types @@ -293,7 +385,7 @@ def _extract_detections(lidar_pc: LidarPc) -> Tuple[List[List[float]], List[List def _extract_ego_state(lidar_pc: LidarPc) -> List[float]: yaw, pitch, roll = lidar_pc.ego_pose.quaternion.yaw_pitch_roll - vehicle_parameters = get_nuplan_pacifica_parameters() + vehicle_parameters = get_nuplan_chrysler_pacifica_parameters() # vehicle_parameters = get_pacifica_parameters() rear_axle_pose = StateSE3( @@ -352,22 +444,57 @@ def _extract_scenario_tag(log_db: NuPlanDB, lidar_pc_token: str) -> List[str]: return scenario_tags -def _extract_front_cam_demo(lidar_pc: LidarPc, source_log_path: Path) -> Tuple[bytes, List[float]]: +def _extract_camera( + log_db: NuPlanDB, + lidar_pc: LidarPc, + source_log_path: Path, + data_converter_config: DataConverterConfig, +) -> Dict[CameraType, Union[str, bytes]]: + + camera_dict: Dict[str, Union[str, bytes]] = {} + sensor_root = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" - sensor_root = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor" - front_image_class = list( - get_images_from_lidar_tokens(source_log_path, [lidar_pc.token], [str(CameraChannel.CAM_F0.value)]) - ) + log_cam_infos = {camera.token: camera for camera in log_db.log.cameras} + + for camera_type, camera_channel in NUPLAN_CAMERA_TYPES.items(): + camera_data: Optional[Union[str, bytes]] = None + c2e: Optional[List[float]] = None + image_class = list(get_images_from_lidar_tokens(source_log_path, [lidar_pc.token], [str(camera_channel.value)])) + if len(image_class) != 0: + image = image_class[0] + filename_jpg = sensor_root / image.filename_jpg + if filename_jpg.exists(): + + # Code taken from MTGS + # https://github.com/OpenDriveLab/MTGS/blob/main/nuplan_scripts/utils/nuplan_utils_custom.py#L117 + + timestamp = image.timestamp + NUPLAN_ROLLING_SHUTTER_S.time_us + img_ego_pose: EgoPose = ( + log_db.log._session.query(EgoPose).order_by(func.abs(EgoPose.timestamp - timestamp)).first() + ) + img_e2g = img_ego_pose.trans_matrix + g2e = lidar_pc.ego_pose.trans_matrix_inv + img_e2e = g2e @ img_e2g + cam_info = log_cam_infos[image.camera_token] + c2img_e = cam_info.trans_matrix + c2e = img_e2e @ c2img_e + + if data_converter_config.camera_store_option == "path": + camera_data = str(filename_jpg), c2e.flatten().tolist() + elif data_converter_config.camera_store_option == "binary": + with open(filename_jpg, "rb") as f: + camera_data = f.read(), c2e + + camera_dict[camera_type] = camera_data - front_cam_demo: bytes = None - front_cam_transform: List[float] = [] + return camera_dict - if len(front_image_class) != 0: - filename_jpg = sensor_root / front_image_class[0].filename_jpg +def _extract_lidar(lidar_pc: LidarPc, data_converter_config: DataConverterConfig) -> Dict[LiDARType, Optional[str]]: - if filename_jpg.exists(): - with open(filename_jpg, "rb") as f: - front_cam_demo = f.read() + lidar: Optional[str] = None + lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / lidar_pc.filename + if lidar_full_path.exists(): + lidar = lidar_pc.filename - return front_cam_demo, front_cam_transform + return {LiDARType.LIDAR_MERGED: lidar} diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py deleted file mode 100644 index ea690921..00000000 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py +++ /dev/null @@ -1,332 +0,0 @@ -import gc -import json -import os -from dataclasses import asdict -from functools import partial -from pathlib import Path -from typing import Dict, Final, List, Tuple, Union - -import psutil -import pyarrow as pa -import yaml -from nuplan.database.nuplan_db_orm.lidar_box import LidarBox -from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc -from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map - -import d123.dataset.dataset_specific.nuplan.utils as nuplan_utils -from d123.common.datatypes.detection.detection import TrafficLightStatus -from d123.common.datatypes.detection.detection_types import DetectionType -from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index -from d123.common.datatypes.vehicle_state.vehicle_parameters import ( - get_nuplan_pacifica_parameters, - rear_axle_se3_to_center_se3, -) -from d123.common.geometry.base import StateSE3 -from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3, BoundingBoxSE3Index -from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL -from d123.common.geometry.vector import Vector3D, Vector3DIndex -from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor -from d123.dataset.logs.log_metadata import LogMetadata - -TARGET_DT: Final[float] = 0.1 -NUPLAN_DT: Final[float] = 0.05 -SORT_BY_TIMESTAMP: Final[bool] = True - -NUPLAN_FULL_MAP_NAME_DICT: Final[Dict[str, str]] = { - "boston": "us-ma-boston", - "singapore": "sg-one-north", - "las_vegas": "us-nv-las-vegas-strip", - "pittsburgh": "us-pa-pittsburgh-hazelwood", -} - -NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { - "green": TrafficLightStatus.GREEN, - "red": TrafficLightStatus.RED, - "unknown": TrafficLightStatus.UNKNOWN, -} -NUPLAN_DETECTION_NAME_DICT = { - "vehicle": DetectionType.VEHICLE, - "bicycle": DetectionType.BICYCLE, - "pedestrian": DetectionType.PEDESTRIAN, - "traffic_cone": DetectionType.TRAFFIC_CONE, - "barrier": DetectionType.BARRIER, - "czone_sign": DetectionType.CZONE_SIGN, - "generic_object": DetectionType.GENERIC_OBJECT, -} - -NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) - - -def create_splits_logs() -> Dict[str, List[str]]: - yaml_filepath = Path(nuplan_utils.__path__[0]) / "log_splits.yaml" - with open(yaml_filepath, "r") as stream: - splits = yaml.safe_load(stream) - - return splits["log_splits"] - - -class NuplanDataProcessor(RawDataProcessor): - def __init__( - self, - splits: List[str], - log_path: Union[Path, str], - output_path: Union[Path, str], - force_data_conversion: bool, - ) -> None: - super().__init__(force_data_conversion) - for split in splits: - assert ( - split in self.get_available_splits() - ), f"Split {split} is not available. Available splits: {self.available_splits}" - - self._splits: List[str] = splits - self._log_path: Path = Path(log_path) - self._output_path: Path = Path(output_path) - self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() - self._target_dt: float = 0.1 - - def _collect_log_paths(self) -> Dict[str, List[Path]]: - # NOTE: the nuplan mini folder has an internal train, val, test structure, all stored in "mini". - # The complete dataset is saved in the "trainval" folder (train and val), or in the "test" folder (for test). - subsplit_log_names: Dict[str, List[str]] = create_splits_logs() - log_paths_per_split: Dict[str, List[Path]] = {} - - for split in self._splits: - subsplit = split.split("_")[-1] - assert subsplit in ["train", "val", "test"] - if split in ["nuplan_train", "nuplan_val"]: - log_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "splits" / "trainval" - elif split in ["nuplan_test"]: - log_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "splits" / "test" - elif split in ["nuplan_mini_train", "nuplan_mini_val", "nuplan_mini_test"]: - log_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "splits" / "mini" - - all_log_files_in_path = [log_file for log_file in log_path.glob("*.db")] - all_log_names = set([str(log_file.stem) for log_file in all_log_files_in_path]) - split_log_names = set(subsplit_log_names[subsplit]) - log_paths = [log_path / f"{log_name}.db" for log_name in list(all_log_names & split_log_names)] - log_paths_per_split[split] = log_paths - - return log_paths_per_split - - def get_available_splits(self) -> List[str]: - return [ - "nuplan_train", - "nuplan_val", - "nuplan_test", - "nuplan_mini_train", - "nuplan_mini_val", - "nuplan_mini_test", - ] - - def convert(self, worker: WorkerPool) -> None: - log_args = [ - { - "log_path": log_path, - "split": split, - } - for split, log_paths in self._log_paths_per_split.items() - for log_path in log_paths - ] - - worker_map( - worker, - partial( - convert_nuplan_log_to_arrow, - output_path=self._output_path, - force_data_conversion=self.force_data_conversion, - ), - log_args, - ) - - -def convert_nuplan_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], - output_path: Path, - force_data_conversion: bool, -) -> None: - process = psutil.Process(os.getpid()) - for log_info in args: - log_path: Path = log_info["log_path"] - split: str = log_info["split"] - - if not log_path.exists(): - raise FileNotFoundError(f"Log path {log_path} does not exist.") - - log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None) - log_file_path = output_path / split / f"{log_path.stem}.arrow" - - if force_data_conversion or not log_file_path.exists(): - log_file_path.unlink(missing_ok=True) - if not log_file_path.parent.exists(): - log_file_path.parent.mkdir(parents=True, exist_ok=True) - - recording_schema = pa.schema( - [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ] - ) - metadata = LogMetadata( - dataset="nuplan", - log_name=log_db.log_name, - location=log_db.log.map_version, - timestep_seconds=TARGET_DT, - map_has_z=False, - ) - vehicle_parameters = get_nuplan_pacifica_parameters() - recording_schema = recording_schema.with_metadata( - { - "log_metadata": json.dumps(asdict(metadata)), - "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), - } - ) - - _write_recording_table(log_db, recording_schema, log_file_path) - - log_db.detach_tables() - log_db.remove_ref() - del recording_schema, vehicle_parameters, log_db - gc.collect() - print(f"{os.getpid()} Memory usage: {process.memory_info().rss / 1024 / 1024:.1f} MB") - return [] - - -def _write_recording_table(log_db: NuPlanDB, recording_schema: pa.schema, log_file_path: Path) -> None: - - # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: - with pa.OSFile(str(log_file_path), "wb") as sink: - with pa.ipc.new_file(sink, recording_schema) as writer: - step_interval: float = int(TARGET_DT / NUPLAN_DT) - for lidar_pc in log_db.lidar_pc[::step_interval]: - lidar_pc_token: str = lidar_pc.token - detections_state, detections_velocity, detections_token, detections_types = _extract_detections( - lidar_pc - ) - traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) - route_lane_group_ids = [ - int(roadblock_id) - for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") - if len(roadblock_id) > 0 - ] - - row_data = { - "token": [lidar_pc_token], - "timestamp": [lidar_pc.timestamp], - "detections_state": [detections_state], - "detections_velocity": [detections_velocity], - "detections_token": [detections_token], - "detections_type": [detections_types], - "ego_states": [_extract_ego_state(lidar_pc)], - "traffic_light_ids": [traffic_light_ids], - "traffic_light_types": [traffic_light_types], - "scenario_tag": [_extract_scenario_tag(log_db, lidar_pc_token)], - "route_lane_group_ids": [route_lane_group_ids], - } - batch = pa.record_batch(row_data, schema=recording_schema) - writer.write_batch(batch) - del batch, row_data, detections_state, detections_velocity, detections_token, detections_types - - if SORT_BY_TIMESTAMP: - recording_table = open_arrow_table(log_file_path) - recording_table = recording_table.sort_by([("timestamp", "ascending")]) - write_arrow_table(recording_table, log_file_path) - - -def _extract_detections(lidar_pc: LidarPc) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: - detections_state: List[List[float]] = [] - detections_velocity: List[List[float]] = [] - detections_token: List[str] = [] - detections_types: List[int] = [] - - for lidar_box in lidar_pc.lidar_boxes: - lidar_box: LidarBox - center = StateSE3( - x=lidar_box.x, - y=lidar_box.y, - z=lidar_box.z, - roll=DEFAULT_ROLL, - pitch=DEFAULT_PITCH, - yaw=lidar_box.yaw, - ) - bounding_box_se3 = BoundingBoxSE3(center, lidar_box.length, lidar_box.width, lidar_box.height) - - detections_state.append(bounding_box_se3.array) - detections_velocity.append(lidar_box.velocity) - detections_token.append(lidar_box.track_token) - detections_types.append(int(NUPLAN_DETECTION_NAME_DICT[lidar_box.category.name])) - - return detections_state, detections_velocity, detections_token, detections_types - - -def _extract_ego_state(lidar_pc: LidarPc) -> List[float]: - - yaw, pitch, roll = lidar_pc.ego_pose.quaternion.yaw_pitch_roll - vehicle_parameters = get_nuplan_pacifica_parameters() - # vehicle_parameters = get_pacifica_parameters() - - rear_axle_pose = StateSE3( - x=lidar_pc.ego_pose.x, - y=lidar_pc.ego_pose.y, - z=lidar_pc.ego_pose.z, - roll=roll, - pitch=pitch, - yaw=yaw, - ) - # NOTE: The height to rear axle is not provided the dataset and is merely approximated. - center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) - dynamic_state = DynamicStateSE3( - velocity=Vector3D( - x=lidar_pc.ego_pose.vx, - y=lidar_pc.ego_pose.vy, - z=lidar_pc.ego_pose.vz, - ), - acceleration=Vector3D( - x=lidar_pc.ego_pose.acceleration_x, - y=lidar_pc.ego_pose.acceleration_y, - z=lidar_pc.ego_pose.acceleration_z, - ), - angular_velocity=Vector3D( - x=lidar_pc.ego_pose.angular_rate_x, - y=lidar_pc.ego_pose.angular_rate_y, - z=lidar_pc.ego_pose.angular_rate_z, - ), - ) - - return EgoStateSE3( - center_se3=center, - dynamic_state_se3=dynamic_state, - vehicle_parameters=vehicle_parameters, - timepoint=None, - ).array.tolist() - - -def _extract_traffic_lights(log_db: NuPlanDB, lidar_pc_token: str) -> Tuple[List[int], List[int]]: - traffic_light_ids: List[int] = [] - traffic_light_types: List[int] = [] - traffic_lights = log_db.traffic_light_status.select_many(lidar_pc_token=lidar_pc_token) - for traffic_light in traffic_lights: - traffic_light_ids.append(int(traffic_light.lane_connector_id)) - traffic_light_types.append(int(NUPLAN_TRAFFIC_STATUS_DICT[traffic_light.status].value)) - return traffic_light_ids, traffic_light_types - - -def _extract_scenario_tag(log_db: NuPlanDB, lidar_pc_token: str) -> List[str]: - - scenario_tags = [ - scenario_tag.type for scenario_tag in log_db.scenario_tag.select_many(lidar_pc_token=lidar_pc_token) - ] - if len(scenario_tags) == 0: - scenario_tags = ["unknown"] - return scenario_tags diff --git a/d123/dataset/dataset_specific/nuplan/map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py similarity index 81% rename from d123/dataset/dataset_specific/nuplan/map_conversion.py rename to d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py index f3861aa9..78f5d3f6 100644 --- a/d123/dataset/dataset_specific/nuplan/map_conversion.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py @@ -11,8 +11,12 @@ import pyogrio from shapely.geometry import LineString +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import ( + get_road_edge_linear_rings, + split_line_geometry_by_max_length, +) from d123.dataset.maps.gpkg.utils import get_all_rows_with_value, get_row_with_value -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType MAP_FILES = { "sg-one-north": "sg-one-north/9.17.1964/map.gpkg", @@ -42,6 +46,18 @@ "gen_lane_connectors_scaled_width_polygons", ] +# 0: generic lane I guess. +# 1: ending? +# 3: bike lanes. + +MAX_ROAD_EDGE_LENGTH = 100.0 # meters, used to filter out very long road edges + +NUPLAN_ROAD_LINE_CONVERSION = { + 0: RoadLineType.DASHED_WHITE, + 2: RoadLineType.SOLID_WHITE, + 3: RoadLineType.UNKNOWN, +} + class NuPlanMapConverter: def __init__(self, map_path: Path) -> None: @@ -55,27 +71,25 @@ def convert(self, map_name: str = "us-pa-pittsburgh-hazelwood") -> None: map_file_path = Path(NUPLAN_MAPS_ROOT) / MAP_FILES[map_name] self._load_dataframes(map_file_path) - lane_df = self._extract_lane_dataframe() - lane_group_df = self._extract_lane_group_dataframe() - intersection_df = self._extract_intersection_dataframe() - crosswalk_df = self._extract_crosswalk_dataframe() - walkway_df = self._extract_walkway_dataframe() - carpark_df = self._extract_carpark_dataframe() - generic_drivable_df = self._extract_generic_drivable_dataframe() + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + dataframes[MapLayer.LANE] = self._extract_lane_dataframe() + dataframes[MapLayer.LANE_GROUP] = self._extract_lane_group_dataframe() + dataframes[MapLayer.INTERSECTION] = self._extract_intersection_dataframe() + dataframes[MapLayer.CROSSWALK] = self._extract_crosswalk_dataframe() + dataframes[MapLayer.WALKWAY] = self._extract_walkway_dataframe() + dataframes[MapLayer.CARPARK] = self._extract_carpark_dataframe() + dataframes[MapLayer.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() + dataframes[MapLayer.ROAD_EDGE] = self._extract_road_edge_dataframe() + dataframes[MapLayer.ROAD_LINE] = self._extract_road_line_dataframe() if not self._map_path.exists(): self._map_path.mkdir(parents=True, exist_ok=True) map_file_name = self._map_path / f"nuplan_{map_name}.gpkg" - lane_df.to_file(map_file_name, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") - lane_group_df.to_file(map_file_name, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") - intersection_df.to_file(map_file_name, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_name, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") - walkway_df.to_file(map_file_name, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_name, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_df.to_file( - map_file_name, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" - ) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_name, layer=layer.serialize(), driver="GPKG", mode="a") def _load_dataframes(self, map_file_path: Path) -> None: @@ -118,6 +132,8 @@ def _extract_nuplan_lane_dataframe(self) -> gpd.GeoDataFrame: successor_ids = [] left_boundaries = [] right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] baseline_paths = [] geometries = self._gdf["lanes_polygons"].geometry.to_list() @@ -145,6 +161,16 @@ def _extract_nuplan_lane_dataframe(self) -> gpd.GeoDataFrame: right_boundary_fid = lane_series["right_boundary_fid"] right_boundary = get_row_with_value(self._gdf["boundaries"], "fid", str(right_boundary_fid))["geometry"] + # 3. left_lane_ids, right_lane_ids + lane_index = lane_series["lane_index"] + all_group_lanes = get_all_rows_with_value( + self._gdf["lanes_polygons"], "lane_group_fid", lane_series["lane_group_fid"] + ) + left_lane_id = all_group_lanes[all_group_lanes["lane_index"] == int(lane_index) - 1]["fid"] + right_lane_id = all_group_lanes[all_group_lanes["lane_index"] == int(lane_index) + 1]["fid"] + left_lane_ids.append(left_lane_id.item() if not left_lane_id.empty else None) + right_lane_ids.append(right_lane_id.item() if not right_lane_id.empty else None) + # 3. baseline_paths baseline_path = get_row_with_value(self._gdf["baseline_paths"], "lane_fid", float(lane_id))["geometry"] @@ -164,6 +190,8 @@ def _extract_nuplan_lane_dataframe(self) -> gpd.GeoDataFrame: "successor_ids": successor_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, "baseline_path": baseline_paths, } ) @@ -225,6 +253,8 @@ def _extract_nuplan_lane_connector_dataframe(self) -> None: "successor_ids": successor_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, + "left_lane_id": [None] * len(ids), + "right_lane_id": [None] * len(ids), "baseline_path": baseline_paths, } ) @@ -384,6 +414,48 @@ def _extract_generic_drivable_dataframe(self) -> gpd.GeoDataFrame: data = pd.DataFrame({"id": self._gdf["generic_drivable_areas"].fid.to_list()}) return gpd.GeoDataFrame(data, geometry=self._gdf["generic_drivable_areas"].geometry.to_list()) + def _extract_road_edge_dataframe(self) -> gpd.GeoDataFrame: + drivable_polygons = ( + self._gdf["intersections"].geometry.to_list() + + self._gdf["lane_groups_polygons"].geometry.to_list() + + self._gdf["carpark_areas"].geometry.to_list() + + self._gdf["generic_drivable_areas"].geometry.to_list() + ) + road_edge_linear_rings = get_road_edge_linear_rings(drivable_polygons) + road_edges = split_line_geometry_by_max_length(road_edge_linear_rings, MAX_ROAD_EDGE_LENGTH) + + ids = [] + road_edge_types = [] + for idx in range(len(road_edges)): + ids.append(idx) + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types.append(int(RoadEdgeType.ROAD_EDGE_BOUNDARY)) + + data = pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}) + return gpd.GeoDataFrame(data, geometry=road_edges) + + def _extract_road_line_dataframe(self) -> gpd.GeoDataFrame: + boundaries = self._gdf["boundaries"].geometry.to_list() + fids = self._gdf["boundaries"].fid.to_list() + boundary_types = self._gdf["boundaries"].boundary_type_fid.to_list() + + ids = [] + road_line_types = [] + geometries = [] + + for idx in range(len(boundary_types)): + ids.append(fids[idx]) + road_line_types.append(int(NUPLAN_ROAD_LINE_CONVERSION[boundary_types[idx]])) + geometries.append(boundaries[idx]) + + data = pd.DataFrame( + { + "id": ids, + "road_line_type": road_line_types, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + def flip_linestring(linestring: LineString) -> LineString: # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. diff --git a/d123/dataset/dataset_specific/raw_data_converter.py b/d123/dataset/dataset_specific/raw_data_converter.py new file mode 100644 index 00000000..179f9ef5 --- /dev/null +++ b/d123/dataset/dataset_specific/raw_data_converter.py @@ -0,0 +1,44 @@ +import abc +from dataclasses import dataclass +from pathlib import Path +from typing import List, Literal, Optional, Union + +from d123.common.multithreading.worker_utils import WorkerPool + + +@dataclass +class DataConverterConfig: + + output_path: Union[str, Path] + force_log_conversion: bool = False + force_map_conversion: bool = False + camera_store_option: Optional[Literal["path", "binary"]] = None + lidar_store_option: Optional[Literal["path", "binary"]] = None + + def __post_init__(self): + if isinstance(self.output_path, str): + self.output_path = Path(self.output_path) + + +class RawDataConverter(abc.ABC): + + def __init__(self, data_converter_config: DataConverterConfig) -> None: + self.data_converter_config = data_converter_config + + @abc.abstractmethod + def get_available_splits(self) -> List[str]: + """Returns a list of available raw data types.""" + + @abc.abstractmethod + def convert_maps(self, worker: WorkerPool) -> None: + """ + Convert maps in raw data format to the uniform 123D format. + :param worker: The worker pool to use for parallel processing. + """ + + @abc.abstractmethod + def convert_logs(self, worker: WorkerPool) -> None: + """ + Convert logs in raw data format to the uniform 123D format. + :param worker: The worker pool to use for parallel processing. + """ diff --git a/d123/dataset/dataset_specific/raw_data_processor.py b/d123/dataset/dataset_specific/raw_data_processor.py deleted file mode 100644 index 74292c1d..00000000 --- a/d123/dataset/dataset_specific/raw_data_processor.py +++ /dev/null @@ -1,18 +0,0 @@ -import abc -from typing import List - -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool - - -class RawDataProcessor(abc.ABC): - - def __init__(self, force_data_conversion: bool) -> None: - self.force_data_conversion = force_data_conversion - - @abc.abstractmethod - def get_available_splits(self) -> List[str]: - """Returns a list of available raw data types.""" - - @abc.abstractmethod - def convert(self, worker: WorkerPool) -> None: - pass diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py new file mode 100644 index 00000000..a2ba1ee2 --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py @@ -0,0 +1,296 @@ +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import numpy as np +import numpy.typing as npt +import shapely.geometry as geom + +from d123.common.geometry.base import Point3D, StateSE2 +from d123.common.geometry.line.polylines import Polyline3D, PolylineSE2 +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.common.geometry.transform.tranform_2d import translate_along_yaw +from d123.common.geometry.utils import normalize_angle +from d123.common.geometry.vector import Vector2D + +MAX_LANE_WIDTH = 25.0 # meters +MIN_LANE_WIDTH = 2.0 +DEFAULT_LANE_WIDTH = 3.7 +BOUNDARY_STEP_SIZE = 0.1 # meters +MAX_Z_DISTANCE = 1.0 # meters + +PERP_START_OFFSET = 0.1 # meters + +MIN_HIT_DISTANCE = 0.1 # meters +MIN_AVERAGE_DISTANCE = 1.5 +MAX_AVERAGE_DISTANCE = 7.0 + + +def get_type_and_id_from_token(token: str) -> Tuple[str, int]: + """Extract type and id from token.""" + line_type, line_id = token.split("_") + return line_type, int(line_id) + + +def get_polyline_from_token(polyline_dict: Dict[str, Dict[int, Polyline3D]], token: str) -> Polyline3D: + """Extract polyline from token.""" + line_type, line_id = get_type_and_id_from_token(token) + return polyline_dict[line_type][line_id] + + +@dataclass +class PerpendicularHit: + distance_along_perp_2d: float + hit_point_3d: Point3D + hit_polyline_token: str + centerline_hit_crossing: bool + heading_error: Optional[float] = None + + @property + def hit_polyline_id(self) -> int: + """Extract polyline id from token.""" + return get_type_and_id_from_token(self.hit_polyline_token)[1] + + @property + def hit_polyline_type(self) -> int: + """Extract polyline id from token.""" + return get_type_and_id_from_token(self.hit_polyline_token)[0] + + +def _collect_perpendicular_hits( + lane_query_se2: StateSE2, + lane_token: str, + polyline_dict: Dict[str, Dict[int, Polyline3D]], + lane_polyline_se2_dict: Dict[int, PolylineSE2], + occupancy_2d: OccupancyMap2D, + sign: float, +) -> List[PerpendicularHit]: + assert sign in [1.0, -1.0], "Sign must be either 1.0 (left) or -1.0 (right)" + # perp_start_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * PERP_START_OFFSET)) + perp_start_point = lane_query_se2 + perp_end_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)) + perp_linestring = geom.LineString([[perp_start_point.x, perp_start_point.y], [perp_end_point.x, perp_end_point.y]]) + + lane_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[lane_token]] + + # 1. find intersecting lines, compute 3D distance + intersecting_tokens = occupancy_2d.intersects(perp_linestring) + + perpendicular_hits: List[PerpendicularHit] = [] + for intersecting_token in intersecting_tokens: + intersecting_polyline_3d = get_polyline_from_token(polyline_dict, intersecting_token) + intersecting_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[intersecting_token]] + centerline_hit_crossing: bool = ( + lane_linestring.intersects(intersecting_linestring) if intersecting_token.startswith("lane_") else False + ) + + intersecting_geom_points: List[geom.Point] = [] + intersecting_geometries = perp_linestring.intersection(intersecting_linestring) + if isinstance(intersecting_geometries, geom.Point): + intersecting_geom_points.append(intersecting_geometries) + elif isinstance(intersecting_geometries, geom.MultiPoint): + intersecting_geom_points.extend([geom for geom in intersecting_geometries.geoms]) + + for intersecting_geom_point in intersecting_geom_points: + distance_along_perp_2d = perp_linestring.project(intersecting_geom_point) + + distance_along_intersecting_norm = intersecting_linestring.project(intersecting_geom_point, normalized=True) + intersecting_point_3d = intersecting_polyline_3d.interpolate( + distance_along_intersecting_norm * intersecting_polyline_3d.length + ) + + heading_error = None + if intersecting_token.startswith("lane_"): + # Compute heading error if the intersecting token is a lane + intersecting_polyline_se2 = lane_polyline_se2_dict[intersecting_token] + lane_heading = intersecting_polyline_se2.interpolate( + distance_along_intersecting_norm * intersecting_polyline_se2.length + ) + heading_error = normalize_angle(lane_query_se2.yaw - lane_heading.yaw) + + perpendicular_hits.append( + PerpendicularHit( + distance_along_perp_2d=distance_along_perp_2d, + hit_point_3d=intersecting_point_3d, + hit_polyline_token=intersecting_token, + centerline_hit_crossing=centerline_hit_crossing, + heading_error=heading_error, + ) + ) + + return perpendicular_hits + + +def _filter_perpendicular_hits( + perpendicular_hits: List[PerpendicularHit], + lane_point_3d: Point3D, +) -> List[PerpendicularHit]: + + filtered_hits = [] + for hit in perpendicular_hits: + + # 1. filter hits too far in the vertical direction + z_distance = np.abs(hit.hit_point_3d.z - lane_point_3d.z) + if z_distance > MAX_Z_DISTANCE: + continue + + # 2. filter hits that are too close and not with the road edge (e.g. close lane lines) + if hit.distance_along_perp_2d < MIN_HIT_DISTANCE and hit.hit_polyline_type != "roadedge": + continue + + filtered_hits.append(hit) + + # Sort hits by distance_along_perp_2d + filtered_hits.sort(key=lambda hit: hit.distance_along_perp_2d) + + return filtered_hits + + +def extract_lane_boundaries( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + road_lines: Dict[int, npt.NDArray[np.float64]], + road_edges: Dict[int, npt.NDArray[np.float64]], +) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: + polyline_dict: Dict[str, Dict[int, Polyline3D]] = {"lane": {}, "roadline": {}, "roadedge": {}} + lane_polyline_se2_dict: Dict[int, PolylineSE2] = {} + + for lane_id, lane_polyline in lanes.items(): + if lane_polyline.ndim == 2 and lane_polyline.shape[1] == 3 and len(lane_polyline) > 0: + polyline_dict["lane"][lane_id] = Polyline3D.from_array(lane_polyline) + lane_polyline_se2_dict[f"lane_{lane_id}"] = polyline_dict["lane"][lane_id].polyline_se2 + + # for road_line_id, road_line_polyline in road_lines.items(): + # if road_line_polyline.ndim == 2 and road_line_polyline.shape[1] == 3 and len(road_line_polyline) > 0: + # polyline_dict["roadline"][road_line_id] = Polyline3D.from_array(road_line_polyline) + + for road_edge_id, road_edge_polyline in road_edges.items(): + if road_edge_polyline.ndim == 2 and road_edge_polyline.shape[1] == 3 and len(road_edge_polyline) > 0: + polyline_dict["roadedge"][road_edge_id] = Polyline3D.from_array(road_edge_polyline) + + geometries = [] + tokens = [] + for line_type, polylines in polyline_dict.items(): + for polyline_id, polyline in polylines.items(): + geometries.append(polyline.polyline_2d.linestring) + tokens.append(f"{line_type}_{polyline_id}") + + occupancy_2d = OccupancyMap2D(geometries, tokens) + + left_boundaries = {} + right_boundaries = {} + + for lane_id, lane_polyline in polyline_dict["lane"].items(): + current_lane_token = f"lane_{lane_id}" + lane_polyline_se2 = lane_polyline_se2_dict[current_lane_token] + + # 1. sample poses along centerline + distances_se2 = np.linspace( + 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_queries_se2 = [ + StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) + ] + distances_3d = np.linspace( + 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_queries_3d = [ + Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) + ] + assert len(lane_queries_se2) == len(lane_queries_3d) + + for sign in [1.0, -1.0]: + boundary_points_3d: List[Optional[Point3D]] = [] + for lane_query_se2, lane_query_3d in zip(lane_queries_se2, lane_queries_3d): + + perpendicular_hits = _collect_perpendicular_hits( + lane_query_se2=lane_query_se2, + lane_token=current_lane_token, + polyline_dict=polyline_dict, + lane_polyline_se2_dict=lane_polyline_se2_dict, + occupancy_2d=occupancy_2d, + sign=sign, + ) + perpendicular_hits = _filter_perpendicular_hits( + perpendicular_hits=perpendicular_hits, lane_point_3d=lane_query_3d + ) + + boundary_point_3d: Optional[Point3D] = None + # 1. First, try to find the boundary point from the perpendicular hits + if len(perpendicular_hits) > 0: + first_hit = perpendicular_hits[0] + + # 1.1. If the first hit is a road edge, use it as the boundary point + if first_hit.hit_polyline_type == "roadedge": + boundary_point_3d = first_hit.hit_point_3d + elif first_hit.hit_polyline_type == "roadline": + boundary_point_3d = first_hit.hit_point_3d + elif first_hit.hit_polyline_type == "lane": + + for hit in perpendicular_hits: + if hit.hit_polyline_type == "roadedge": + continue + if hit.hit_polyline_type == "lane": + + has_same_predecessor = ( + len(set(lanes_predecessors[hit.hit_polyline_id]) & set(lanes_predecessors[lane_id])) + > 0 + ) + has_same_successor = ( + len(set(lanes_successors[hit.hit_polyline_id]) & set(lanes_successors[lane_id])) > 0 + ) + heading_min = np.pi / 8.0 + invalid_heading_error = heading_min < abs(hit.heading_error) < (np.pi - heading_min) + if ( + not has_same_predecessor + and not has_same_successor + and not hit.centerline_hit_crossing + and MAX_AVERAGE_DISTANCE > hit.distance_along_perp_2d + and MIN_AVERAGE_DISTANCE < hit.distance_along_perp_2d + and not invalid_heading_error + ): + # 2. if first hit is lane line, use it as boundary point + boundary_point_3d = Point3D.from_array( + (hit.hit_point_3d.array + lane_query_3d.array) / 2.0 + ) + break + + boundary_points_3d.append(boundary_point_3d) + + no_boundary_ratio = boundary_points_3d.count(None) / len(boundary_points_3d) + final_boundary_points_3d = [] + + def _get_default_boundary_point_3d( + lane_query_se2: StateSE2, lane_query_3d: Point3D, sign: float + ) -> Point3D: + perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 + boundary_point_se2 = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance)) + return Point3D(boundary_point_se2.x, boundary_point_se2.y, lane_query_3d.z) + + if no_boundary_ratio > 0.8: + for lane_query_se2, lane_query_3d in zip(lane_queries_se2, lane_queries_3d): + boundary_point_3d = _get_default_boundary_point_3d(lane_query_se2, lane_query_3d, sign) + final_boundary_points_3d.append(boundary_point_3d.array) + + else: + for boundary_idx, (lane_query_se2, lane_query_3d) in enumerate(zip(lane_queries_se2, lane_queries_3d)): + if boundary_points_3d[boundary_idx] is None: + boundary_point_3d = _get_default_boundary_point_3d(lane_query_se2, lane_query_3d, sign) + else: + boundary_point_3d = boundary_points_3d[boundary_idx] + final_boundary_points_3d.append(boundary_point_3d.array) + + # # 2. If no boundary point was found, use the lane query point as the boundary point + # if boundary_point_3d is None: + + if len(final_boundary_points_3d) > 1: + if sign == 1.0: + left_boundaries[lane_id] = Polyline3D.from_array( + np.array(final_boundary_points_3d, dtype=np.float64) + ) + else: + right_boundaries[lane_id] = Polyline3D.from_array( + np.array(final_boundary_points_3d, dtype=np.float64) + ) + + return left_boundaries, right_boundaries diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py new file mode 100644 index 00000000..6d3ef7ab --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py @@ -0,0 +1,388 @@ +from collections import defaultdict +from pathlib import Path +from typing import Dict, List, Optional + +import geopandas as gpd +import numpy as np +import numpy.typing as npt +import pandas as pd +import shapely.geometry as geom +from waymo_open_dataset import dataset_pb2 + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.units import mph_to_mps +from d123.dataset.dataset_specific.wopd.waymo_map_utils.womp_boundary_utils import extract_lane_boundaries +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType + +# TODO: +# - Implement stop signs +# - Implement speed bumps +# - Implement driveways with a different semantic type if needed +# - Implement intersections and lane group logic + +WAYMO_ROAD_LINE_CONVERSION = { + 0: RoadLineType.UNKNOWN, # aka. UNKNOWN + 1: RoadLineType.DASHED_WHITE, # aka. BROKEN_SINGLE_WHITE + 2: RoadLineType.SOLID_WHITE, # aka. SOLID_SINGLE_WHITE + 3: RoadLineType.DOUBLE_SOLID_WHITE, # aka. SOLID_DOUBLE_WHITE + 4: RoadLineType.DASHED_YELLOW, # aka. BROKEN_SINGLE_YELLOW + 5: RoadLineType.DOUBLE_DASH_YELLOW, # aka. BROKEN_DOUBLE_YELLOW + 6: RoadLineType.SOLID_YELLOW, # aka. SOLID_SINGLE_YELLOW + 7: RoadLineType.DOUBLE_SOLID_YELLOW, # aka. SOLID_DOUBLE_YELLOW + 8: RoadLineType.DOUBLE_DASH_YELLOW, # aka. PASSING_DOUBLE_YELLOW +} + +WAYMO_ROAD_EDGE_CONVERSION = { + 0: RoadEdgeType.UNKNOWN, + 1: RoadEdgeType.ROAD_EDGE_BOUNDARY, + 2: RoadEdgeType.ROAD_EDGE_MEDIAN, +} + + +def convert_wopd_map(frame: dataset_pb2.Frame, map_file_path: Path) -> None: + + def _extract_polyline(data) -> npt.NDArray[np.float64]: + polyline = np.array([[p.x, p.y, p.z] for p in data.polyline], dtype=np.float64) + return polyline + + def _extract_polygon(data) -> npt.NDArray[np.float64]: + polygon = np.array([[p.x, p.y, p.z] for p in data.polygon], dtype=np.float64) + assert polygon.shape[0] >= 3, "Polygon must have at least 3 points" + assert polygon.shape[1] == 3, "Polygon must have 3 coordinates (x, y, z)" + return polygon + + def _extract_neighbors(data) -> List[Dict[str, int]]: + neighbors = [] + for neighbor in data: + neighbors.append( + { + "lane_id": neighbor.feature_id, + "self_start_index": neighbor.self_start_index, + "self_end_index": neighbor.self_end_index, + "neighbor_start_index": neighbor.neighbor_start_index, + "neighbor_end_index": neighbor.neighbor_end_index, + } + ) + return neighbors + + lanes: Dict[int, npt.NDArray[np.float64]] = {} + lanes_successors = defaultdict(list) + lanes_predecessors = defaultdict(list) + lanes_speed_limit_mps: Dict[int, float] = {} + lanes_type: Dict[int, int] = {} + lanes_left_neighbors: Dict[int, List[Dict[str, int]]] = {} + lanes_right_neighbors: Dict[int, List[Dict[str, int]]] = {} + + road_lines: Dict[int, npt.NDArray[np.float64]] = {} + road_lines_type: Dict[int, RoadLineType] = {} + + road_edges: Dict[int, npt.NDArray[np.float64]] = {} + road_edges_type: Dict[int, int] = {} + + crosswalks: Dict[int, npt.NDArray[np.float64]] = {} + carparks: Dict[int, npt.NDArray[np.float64]] = {} + + for map_feature in frame.map_features: + if map_feature.HasField("lane"): + polyline = _extract_polyline(map_feature.lane) + # Ignore lanes with less than 2 points or not 2D + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + lanes[map_feature.id] = polyline + for lane_id_ in map_feature.lane.exit_lanes: + lanes_successors[map_feature.id].append(lane_id_) + for lane_id_ in map_feature.lane.exit_lanes: + lanes_predecessors[map_feature.id].append(lane_id_) + lanes_speed_limit_mps[map_feature.id] = mph_to_mps(map_feature.lane.speed_limit_mph) + lanes_type[map_feature.id] = map_feature.lane.type + lanes_left_neighbors[map_feature.id] = _extract_neighbors(map_feature.lane.left_neighbors) + lanes_right_neighbors[map_feature.id] = _extract_neighbors(map_feature.lane.right_neighbors) + elif map_feature.HasField("road_line"): + polyline = _extract_polyline(map_feature.road_line) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + road_lines[map_feature.id] = polyline + road_lines_type[map_feature.id] = WAYMO_ROAD_LINE_CONVERSION.get( + map_feature.road_line.type, RoadLineType.UNKNOWN + ) + elif map_feature.HasField("road_edge"): + polyline = _extract_polyline(map_feature.road_edge) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + road_edges[map_feature.id] = polyline + road_edges_type[map_feature.id] = WAYMO_ROAD_EDGE_CONVERSION.get( + map_feature.road_edge.type, RoadEdgeType.UNKNOWN + ) + elif map_feature.HasField("stop_sign"): + # TODO: implement stop signs + pass + elif map_feature.HasField("crosswalk"): + crosswalks[map_feature.id] = _extract_polygon(map_feature.crosswalk) + elif map_feature.HasField("speed_bump"): + # TODO: implement speed bumps + pass + elif map_feature.HasField("driveway"): + # NOTE: Determine whether to use a different semantic type for driveways. + carparks[map_feature.id] = _extract_polygon(map_feature.driveway) + + lane_left_boundaries_3d, lane_right_boundaries_3d = extract_lane_boundaries( + lanes, lanes_successors, lanes_predecessors, road_lines, road_edges + ) + + lane_df = get_lane_df( + lanes, + lanes_successors, + lanes_predecessors, + lanes_speed_limit_mps, + lane_left_boundaries_3d, + lane_right_boundaries_3d, + lanes_type, + lanes_left_neighbors, + lanes_right_neighbors, + ) + lane_group_df = get_lane_group_df( + lanes, + lanes_successors, + lanes_predecessors, + lane_left_boundaries_3d, + lane_right_boundaries_3d, + ) + intersection_df = get_intersections_df() + crosswalk_df = get_crosswalk_df(crosswalks) + walkway_df = get_walkway_df() + carpark_df = get_carpark_df(carparks) + generic_drivable_df = get_generic_drivable_df() + road_edge_df = get_road_edge_df(road_edges, road_edges_type) + road_line_df = get_road_line_df(road_lines, road_lines_type) + + map_file_path.unlink(missing_ok=True) + if not map_file_path.parent.exists(): + map_file_path.parent.mkdir(parents=True, exist_ok=True) + + lane_df.to_file(map_file_path, layer=MapLayer.LANE.serialize(), driver="GPKG") + lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") + intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") + crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") + walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") + carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") + generic_drivable_df.to_file(map_file_path, layer=MapLayer.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a") + road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + road_line_df.to_file(map_file_path, layer=MapLayer.ROAD_LINE.serialize(), driver="GPKG", mode="a") + + +def get_lane_df( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + lanes_speed_limit_mps: Dict[int, float], + lanes_left_boundaries_3d: Dict[int, Polyline3D], + lanes_right_boundaries_3d: Dict[int, Polyline3D], + lanes_type: Dict[int, int], + lanes_left_neighbors: Dict[int, List[Dict[str, int]]], + lanes_right_neighbors: Dict[int, List[Dict[str, int]]], +) -> gpd.GeoDataFrame: + + ids = [] + lane_types = [] + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] + baseline_paths = [] + geometries = [] + + def _get_majority_neighbor(neighbors: List[Dict[str, int]]) -> Optional[int]: + if len(neighbors) == 0: + return None + length = { + neighbor["lane_id"]: neighbor["self_end_index"] - neighbor["self_start_index"] for neighbor in neighbors + } + return str(max(length, key=length.get)) + + for lane_id, lane_centerline_array in lanes.items(): + if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: + continue + lane_centerline = Polyline3D.from_array(lane_centerline_array) + lane_speed_limit_mps = lanes_speed_limit_mps[lane_id] if lanes_speed_limit_mps[lane_id] > 0.0 else None + + ids.append(lane_id) + lane_types.append(lanes_type[lane_id]) + lane_group_ids.append([lane_id]) + speed_limits_mps.append(lane_speed_limit_mps) + predecessor_ids.append(lanes_predecessors[lane_id]) + successor_ids.append(lanes_successors[lane_id]) + left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) + right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + left_lane_ids.append(_get_majority_neighbor(lanes_left_neighbors[lane_id])) + right_lane_ids.append(_get_majority_neighbor(lanes_right_neighbors[lane_id])) + baseline_paths.append(lane_centerline.linestring) + + geometry = geom.Polygon( + np.vstack( + [ + lanes_left_boundaries_3d[lane_id].array[:, :2], + lanes_right_boundaries_3d[lane_id].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_type": lane_types, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, + "baseline_path": baseline_paths, + } + ) + + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_lane_group_df( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + lanes_left_boundaries_3d: Dict[int, Polyline3D], + lanes_right_boundaries_3d: Dict[int, Polyline3D], +) -> gpd.GeoDataFrame: + + ids = [] + lane_ids = [] + intersection_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + # NOTE: WOPD does not provide lane groups, so we create a lane group for each lane. + for lane_id in lanes.keys(): + if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: + continue + ids.append(lane_id) + lane_ids.append([lane_id]) + intersection_ids.append(None) # WOPD does not provide intersections + predecessor_lane_group_ids.append(lanes_predecessors[lane_id]) + successor_lane_group_ids.append(lanes_successors[lane_id]) + left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) + right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + geometry = geom.Polygon( + np.vstack( + [ + lanes_left_boundaries_3d[lane_id].array[:, :2], + lanes_right_boundaries_3d[lane_id].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "intersection_id": intersection_ids, + "predecessor_lane_group_ids": predecessor_lane_group_ids, + "successor_lane_group_ids": successor_lane_group_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_intersections_df() -> gpd.GeoDataFrame: + ids = [] + lane_group_ids = [] + geometries = [] + + # NOTE: WOPD does not provide intersections, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_carpark_df(carparks) -> gpd.GeoDataFrame: + ids = list(carparks.keys()) + outlines = [geom.LineString(outline) for outline in carparks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in carparks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_walkway_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + # NOTE: WOPD does not provide walkways, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(crosswalks.keys()) + outlines = [geom.LineString(outline) for outline in crosswalks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in crosswalks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_generic_drivable_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + # NOTE: WOPD does not provide generic drivable areas, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_road_edge_df( + road_edges: Dict[int, npt.NDArray[np.float64]], road_edges_type: Dict[int, RoadEdgeType] +) -> gpd.GeoDataFrame: + ids = list(road_edges.keys()) + geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_edges.values()] + + data = pd.DataFrame( + { + "id": ids, + "road_edge_type": [int(road_edge_type) for road_edge_type in road_edges_type.values()], + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_road_line_df( + road_lines: Dict[int, npt.NDArray[np.float64]], road_lines_type: Dict[int, RoadLineType] +) -> gpd.GeoDataFrame: + ids = list(road_lines.keys()) + geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_lines.values()] + + data = pd.DataFrame( + { + "id": ids, + "road_line_type": [int(road_line_type) for road_line_type in road_lines_type.values()], + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py new file mode 100644 index 00000000..f1145c8e --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -0,0 +1,546 @@ +import gc +import hashlib +import json +import os +from dataclasses import asdict +from functools import partial +from pathlib import Path +from typing import Any, Dict, Final, List, Literal, Tuple, Union + +import numpy as np +import numpy.typing as npt +import pyarrow as pa +import tensorflow as tf +from pyquaternion import Quaternion +from waymo_open_dataset import dataset_pb2 + +from d123.common.datatypes.detection.detection import TrafficLightStatus +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar_index import WopdLidarIndex +from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index +from d123.common.datatypes.vehicle_state.vehicle_parameters import get_wopd_chrysler_pacifica_parameters +from d123.common.geometry.base import Point3D, StateSE3 +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index +from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array, get_rotation_matrix +from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map +from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.dataset_specific.wopd.waymo_map_utils.wopd_map_utils import convert_wopd_map +from d123.dataset.dataset_specific.wopd.wopd_utils import parse_range_image_and_camera_projection +from d123.dataset.logs.log_metadata import LogMetadata + +os.environ["CUDA_VISIBLE_DEVICES"] = "-1" +D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) + +TARGET_DT: Final[float] = 0.1 +SORT_BY_TIMESTAMP: Final[bool] = False + +NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { + "green": TrafficLightStatus.GREEN, + "red": TrafficLightStatus.RED, + "unknown": TrafficLightStatus.UNKNOWN, +} + +# https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/label.proto#L63 +WOPD_DETECTION_NAME_DICT: Dict[int, DetectionType] = { + 0: DetectionType.GENERIC_OBJECT, # TYPE_UNKNOWN + 1: DetectionType.VEHICLE, # TYPE_VEHICLE + 2: DetectionType.PEDESTRIAN, # TYPE_PEDESTRIAN + 3: DetectionType.SIGN, # TYPE_SIGN + 4: DetectionType.BICYCLE, # TYPE_CYCLIST +} + +# https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L50 +WOPD_CAMERA_TYPES: Dict[int, CameraType] = { + 1: CameraType.CAM_F0, # front_camera + 2: CameraType.CAM_L0, # front_left_camera + 3: CameraType.CAM_R0, # front_right_camera + 4: CameraType.CAM_L1, # left_camera + 5: CameraType.CAM_R1, # right_camera +} + +# https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L66 +WOPD_LIDAR_TYPES: Dict[int, LiDARType] = { + 0: LiDARType.LIDAR_UNKNOWN, # UNKNOWN + 1: LiDARType.LIDAR_TOP, # TOP + 2: LiDARType.LIDAR_FRONT, # FRONT + 3: LiDARType.LIDAR_SIDE_LEFT, # SIDE_LEFT + 4: LiDARType.LIDAR_SIDE_RIGHT, # SIDE_RIGHT + 5: LiDARType.LIDAR_BACK, # REAR +} + +WOPD_DATA_ROOT = Path("/media/nvme1/waymo_perception") # TODO: set as environment variable + +# Whether to use ego or zero roll and pitch values for bounding box detections (after global conversion) +DETECTION_ROLL_PITCH: Final[Literal["ego", "zero"]] = "zero" + + +def create_token(input_data: str) -> str: + # TODO: Refactor this function. + # TODO: Add a general function to create tokens from arbitrary data. + if isinstance(input_data, str): + input_data = input_data.encode("utf-8") + + hash_obj = hashlib.sha256(input_data) + return hash_obj.hexdigest()[:16] + + +class WOPDDataConverter(RawDataConverter): + def __init__( + self, + splits: List[str], + log_path: Union[Path, str], + data_converter_config: DataConverterConfig, + ) -> None: + super().__init__(data_converter_config) + for split in splits: + assert ( + split in self.get_available_splits() + ), f"Split {split} is not available. Available splits: {self.available_splits}" + + self._splits: List[str] = splits + self._tf_records_per_split: Dict[str, List[Path]] = self._collect_tf_records() + self._target_dt: float = 0.1 + + def _collect_tf_records(self) -> Dict[str, List[Path]]: + tf_records_per_split: Dict[str, List[Path]] = {} + + for split in self._splits: + if split in ["wopd_train"]: + log_path = WOPD_DATA_ROOT / "training" + else: + raise ValueError(f"Split {split} is not supported.") + + log_paths = [log_file for log_file in log_path.glob("*.tfrecord")] + tf_records_per_split[split] = log_paths + + return tf_records_per_split + + def get_available_splits(self) -> List[str]: + # TODO: Add more splits if available + return [ + "wopd_train", + ] + + def convert_maps(self, worker: WorkerPool) -> None: + log_args = [ + { + "tf_record": tf_record, + "split": split, + } + for split, tf_record_paths in self._tf_records_per_split.items() + for tf_record in tf_record_paths + ] + + worker_map( + worker, + partial(convert_wopd_tfrecord_map_to_gpkg, data_converter_config=self.data_converter_config), + log_args, + ) + + def convert_logs(self, worker: WorkerPool) -> None: + log_args = [ + { + "tf_record": tf_record, + "split": split, + } + for split, tf_record_paths in self._tf_records_per_split.items() + for tf_record in tf_record_paths + ] + + worker_map( + worker, + partial(convert_wopd_tfrecord_log_to_arrow, data_converter_config=self.data_converter_config), + log_args, + ) + + +def convert_wopd_tfrecord_map_to_gpkg( + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig +) -> List[Any]: + + for log_info in args: + tf_record_path: Path = log_info["tf_record"] + split: str = log_info["split"] + + if not tf_record_path.exists(): + raise FileNotFoundError(f"TFRecord path {tf_record_path} does not exist.") + + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") + for data in dataset: + initial_frame = dataset_pb2.Frame() + initial_frame.ParseFromString(data.numpy()) + break + log_name = str(initial_frame.context.name) + map_file_path = D123_MAPS_ROOT / split / f"{log_name}.gpkg" + + if data_converter_config.force_map_conversion or not map_file_path.exists(): + map_file_path.unlink(missing_ok=True) + convert_wopd_map(initial_frame, map_file_path) + return [] + + +def convert_wopd_tfrecord_log_to_arrow( + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig +) -> List[Any]: + for log_info in args: + try: + + tf_record_path: Path = log_info["tf_record"] + split: str = log_info["split"] + + if not tf_record_path.exists(): + raise FileNotFoundError(f"TFRecord path {tf_record_path} does not exist.") + + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") + for data in dataset: + initial_frame = dataset_pb2.Frame() + initial_frame.ParseFromString(data.numpy()) + break + + log_name = str(initial_frame.context.name) + log_file_path = data_converter_config.output_path / split / f"{log_name}.arrow" + + if data_converter_config.force_log_conversion or not log_file_path.exists(): + log_file_path.unlink(missing_ok=True) + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + metadata = LogMetadata( + dataset="wopd", + log_name=log_name, + location=None, # TODO: implement map name + timestep_seconds=TARGET_DT, # TODO: Check if correct. Maybe not hardcode + map_has_z=True, + ) + vehicle_parameters = get_wopd_chrysler_pacifica_parameters() + camera_metadata = get_wopd_camera_metadata(initial_frame, data_converter_config) + lidar_metadata = get_wopd_lidar_metadata(initial_frame, data_converter_config) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + # TODO: Adjust how cameras are added + if data_converter_config.camera_store_option is not None: + for camera_type in camera_metadata.keys(): + if data_converter_config.camera_store_option == "path": + raise NotImplementedError("Path camera storage is not implemented.") + elif data_converter_config.camera_store_option == "binary": + schema_column_list.append((camera_type.serialize(), pa.binary())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) + ) + + if data_converter_config.lidar_store_option is not None: + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + raise NotImplementedError("Filepath lidar storage is not implemented.") + elif data_converter_config.lidar_store_option == "binary": + schema_column_list.append( + (lidar_type.serialize(), pa.list_(pa.list_(pa.float32(), len(WopdLidarIndex)))) + ) + + recording_schema = pa.schema(schema_column_list) + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), + } + ) + + _write_recording_table(dataset, recording_schema, log_file_path, tf_record_path, data_converter_config) + + del recording_schema, vehicle_parameters, dataset + except Exception as e: + import traceback + + print(f"Error processing log {str(tf_record_path)}: {e}") + traceback.print_exc() + gc.collect() + return [] + + +def get_wopd_camera_metadata( + initial_frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Dict[CameraType, CameraMetadata]: + + cam_metadatas: Dict[CameraType, CameraMetadata] = {} + if data_converter_config.camera_store_option is not None: + for calibration in initial_frame.context.camera_calibrations: + camera_type = WOPD_CAMERA_TYPES[calibration.name] + + # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L96 + # https://github.com/waymo-research/waymo-open-dataset/issues/834#issuecomment-2134995440 + fx, fy, cx, cy, k1, k2, p1, p2, k3 = calibration.intrinsic + _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) + _distortions = np.array([k1, k2, p1, p2, k3]) + + if camera_type in WOPD_CAMERA_TYPES.values(): + cam_metadatas[camera_type] = CameraMetadata( + camera_type=camera_type, + width=calibration.width, + height=calibration.height, + intrinsic=_intrinsics, + distortion=_distortions, + ) + + return cam_metadatas + + +def get_wopd_lidar_metadata( + initial_frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Dict[LiDARType, LiDARMetadata]: + + laser_metadatas: Dict[LiDARType, LiDARMetadata] = {} + if data_converter_config.lidar_store_option is not None: + for laser_calibration in initial_frame.context.laser_calibrations: + lidar_type = WOPD_LIDAR_TYPES[laser_calibration.name] + extrinsic = ( + np.array(laser_calibration.extrinsic.transform, dtype=np.float64).reshape(4, 4) + if laser_calibration.extrinsic + else None + ) + laser_metadatas[lidar_type] = LiDARMetadata( + lidar_type=lidar_type, + lidar_index=WopdLidarIndex, + extrinsic=extrinsic, + ) + + return laser_metadatas + + +def _write_recording_table( + dataset: tf.data.TFRecordDataset, + recording_schema: pa.schema, + log_file_path: Path, + tf_record_path: Path, + data_converter_config: DataConverterConfig, +) -> None: + + # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: + with pa.OSFile(str(log_file_path), "wb") as sink: + with pa.ipc.new_file(sink, recording_schema) as writer: + + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") + for frame_idx, data in enumerate(dataset): + frame = dataset_pb2.Frame() + frame.ParseFromString(data.numpy()) + + (detections_state, detections_velocity, detections_token, detections_types) = _extract_detections(frame) + + # TODO: Implement traffic light extraction + traffic_light_ids = [] + traffic_light_types = [] + + # TODO: Implement detections + row_data = { + "token": [create_token(f"{frame.context.name}_{int(frame.timestamp_micros)}")], + "timestamp": [int(frame.timestamp_micros)], + "detections_state": [detections_state], + "detections_velocity": [detections_velocity], + "detections_token": [detections_token], + "detections_type": [detections_types], + "ego_states": [_extract_ego_state(frame)], + "traffic_light_ids": [traffic_light_ids], + "traffic_light_types": [traffic_light_types], + "scenario_tag": ["unknown"], + "route_lane_group_ids": [None], + } + + # TODO: Implement lidar extraction + if data_converter_config.lidar_store_option is not None: + lidar_data_dict = _extract_lidar(frame, data_converter_config) + for lidar_type, lidar_data in lidar_data_dict.items(): + if lidar_data is not None: + row_data[lidar_type.serialize()] = [lidar_data.tolist()] + else: + row_data[lidar_type.serialize()] = [None] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_camera(frame, data_converter_config) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] + else: + row_data[camera_type.serialize()] = [None] + row_data[f"{camera_type.serialize()}_extrinsic"] = [None] + + batch = pa.record_batch(row_data, schema=recording_schema) + writer.write_batch(batch) + del batch, row_data, detections_state, detections_velocity, detections_token, detections_types + + if SORT_BY_TIMESTAMP: + recording_table = open_arrow_table(log_file_path) + recording_table = recording_table.sort_by([("timestamp", "ascending")]) + write_arrow_table(recording_table, log_file_path) + + +def _get_ego_pose_se3(frame: dataset_pb2.Frame) -> StateSE3: + ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) + yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll + ego_point_3d = Point3D.from_array(ego_pose_matrix[:3, 3]) + + return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=roll, pitch=pitch, yaw=yaw) + + +def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: + # TODO: implement + + ego_rear_axle = _get_ego_pose_se3(frame) + + num_detections = len(frame.laser_labels) + detections_state = np.zeros((num_detections, len(BoundingBoxSE3Index)), dtype=np.float64) + detections_velocity = np.zeros((num_detections, len(Vector3DIndex)), dtype=np.float64) + detections_token: List[str] = [] + detections_types: List[int] = [] + + for detection_idx, detection in enumerate(frame.laser_labels): + if detection.type not in WOPD_DETECTION_NAME_DICT: + continue + + # 1. SS3 Bounding Box + detections_state[detection_idx, BoundingBoxSE3Index.X] = detection.box.center_x + detections_state[detection_idx, BoundingBoxSE3Index.Y] = detection.box.center_y + detections_state[detection_idx, BoundingBoxSE3Index.Z] = detection.box.center_z + detections_state[detection_idx, BoundingBoxSE3Index.ROLL] = DEFAULT_ROLL # not provided in WOPD + detections_state[detection_idx, BoundingBoxSE3Index.PITCH] = DEFAULT_PITCH # not provided in WOPD + detections_state[detection_idx, BoundingBoxSE3Index.YAW] = detection.box.heading + detections_state[detection_idx, BoundingBoxSE3Index.LENGTH] = detection.box.length + detections_state[detection_idx, BoundingBoxSE3Index.WIDTH] = detection.box.width + detections_state[detection_idx, BoundingBoxSE3Index.HEIGHT] = detection.box.height + + # 2. Velocity TODO: check if velocity needs to be rotated + detections_velocity[detection_idx] = Vector3D( + x=detection.metadata.speed_x, + y=detection.metadata.speed_y, + z=detection.metadata.speed_z, + ).array + + # 3. Type and track token + detections_token.append(str(detection.id)) + detections_types.append(int(WOPD_DETECTION_NAME_DICT[detection.type])) + + detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( + origin=ego_rear_axle, se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3] + ) + if DETECTION_ROLL_PITCH == "ego": + pass + if DETECTION_ROLL_PITCH == "zero": + detections_state[:, BoundingBoxSE3Index.ROLL] = DEFAULT_ROLL + detections_state[:, BoundingBoxSE3Index.PITCH] = DEFAULT_PITCH + else: + raise ValueError(f"Invalid DETECTION_ROLL_PITCH value: {DETECTION_ROLL_PITCH}. Must be 'ego' or 'zero'.") + + return detections_state.tolist(), detections_velocity.tolist(), detections_token, detections_types + + +def _extract_ego_state(frame: dataset_pb2.Frame) -> List[float]: + rear_axle_pose = _get_ego_pose_se3(frame) + + vehicle_parameters = get_wopd_chrysler_pacifica_parameters() + # FIXME: Find dynamic state in waymo open perception dataset + # https://github.com/waymo-research/waymo-open-dataset/issues/55#issuecomment-546152290 + dynamic_state = DynamicStateSE3( + velocity=Vector3D(*np.zeros(3)), + acceleration=Vector3D(*np.zeros(3)), + angular_velocity=Vector3D(*np.zeros(3)), + ) + + return EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_pose, + dynamic_state_se3=dynamic_state, + vehicle_parameters=vehicle_parameters, + time_point=None, + ).array.tolist() + + +def _extract_traffic_lights() -> Tuple[List[int], List[int]]: + pass + + +def _extract_camera( + frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Dict[CameraType, Union[str, bytes]]: + + camera_dict: Dict[str, Union[str, bytes]] = {} # TODO: Fix wrong type hint + np.array(frame.pose.transform).reshape(4, 4) + + # NOTE: The extrinsic matrix in frame.context.camera_calibration is fixed to model the ego to camera transformation. + # The poses in frame.images[idx] are the motion compensated ego poses when the camera triggers. + # + + context_extrinsic: Dict[str, npt.NDArray] = {} + for calibration in frame.context.camera_calibrations: + camera_type = WOPD_CAMERA_TYPES[calibration.name] + + transform = np.array(calibration.extrinsic.transform).reshape(4, 4) + + # FIXME: This is an ugly hack to convert to uniform camera convention. + flip_camera = get_rotation_matrix( + StateSE3( + x=0.0, + y=0.0, + z=0.0, + roll=np.deg2rad(0.0), + pitch=np.deg2rad(90.0), + yaw=np.deg2rad(-90.0), + ) + ) + transform[:3, :3] = transform[:3, :3] @ flip_camera + context_extrinsic[camera_type] = transform + + for image_proto in frame.images: + camera_type = WOPD_CAMERA_TYPES[image_proto.name] + + np.array(image_proto.pose.transform).reshape(4, 4) + camera_bytes = image_proto.image + + # # Compute the transform from ego_global_transform to ego_at_camera_transform + # # ego_global_transform * T = ego_at_camera_transform => T = ego_global_transform^-1 * ego_at_camera_transform + # np.linalg.inv(ego_global_transform) @ ego_at_trigger_transform + + # TODO: figure out the correct transform + camera_dict[camera_type] = camera_bytes, context_extrinsic[camera_type].flatten().tolist() + + return camera_dict + + +def _extract_lidar( + frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Dict[LiDARType, npt.NDArray[np.float32]]: + from waymo_open_dataset.utils import frame_utils + + assert data_converter_config.lidar_store_option == "binary", "Lidar store option must be 'binary' for WOPD." + (range_images, camera_projections, _, range_image_top_pose) = parse_range_image_and_camera_projection(frame) + + points, cp_points = frame_utils.convert_range_image_to_point_cloud( + frame=frame, + range_images=range_images, + camera_projections=camera_projections, + range_image_top_pose=range_image_top_pose, + keep_polar_features=True, + ) + + lidar_data: Dict[LiDARType, npt.NDArray[np.float32]] = {} + for lidar_idx, frame_lidar in enumerate(frame.lasers): + lidar_type = WOPD_LIDAR_TYPES[frame_lidar.name] + lidar_data[lidar_type] = np.array(points[lidar_idx], dtype=np.float32) + + return lidar_data diff --git a/d123/dataset/dataset_specific/wopd/wopd_utils.py b/d123/dataset/dataset_specific/wopd/wopd_utils.py new file mode 100644 index 00000000..c9c302ef --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/wopd_utils.py @@ -0,0 +1,76 @@ +from typing import Dict, List, Tuple + +import tensorflow as tf +from waymo_open_dataset import dataset_pb2 + +RangeImages = Dict["dataset_pb2.LaserName.Name", List[dataset_pb2.MatrixFloat]] +CameraProjections = Dict["dataset_pb2.LaserName.Name", List[dataset_pb2.MatrixInt32]] +SegmentationLabels = Dict["dataset_pb2.LaserName.Name", List[dataset_pb2.MatrixInt32]] +ParsedFrame = Tuple[RangeImages, CameraProjections, SegmentationLabels, dataset_pb2.MatrixFloat] + + +def parse_range_image_and_camera_projection(frame: dataset_pb2.Frame) -> ParsedFrame: + """Parse range images and camera projections given a frame. + + Args: + frame: open dataset frame proto + + Returns: + range_images: A dict of {laser_name, + [range_image_first_return, range_image_second_return]}. + camera_projections: A dict of {laser_name, + [camera_projection_from_first_return, + camera_projection_from_second_return]}. + seg_labels: segmentation labels, a dict of {laser_name, + [seg_label_first_return, seg_label_second_return]} + range_image_top_pose: range image pixel pose for top lidar. + """ + range_images = {} + camera_projections = {} + seg_labels = {} + range_image_top_pose: dataset_pb2.MatrixFloat = dataset_pb2.MatrixFloat() + for laser in frame.lasers: + if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed(laser.ri_return1.range_image_compressed, "ZLIB") + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(range_image_str_tensor.numpy()) + range_images[laser.name] = [ri] + + if laser.name == dataset_pb2.LaserName.TOP: + range_image_top_pose_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_pose_compressed, "ZLIB" + ) + range_image_top_pose = dataset_pb2.MatrixFloat() + range_image_top_pose.ParseFromString(range_image_top_pose_str_tensor.numpy()) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return1.camera_projection_compressed, "ZLIB" + ) + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(camera_projection_str_tensor.numpy()) + camera_projections[laser.name] = [cp] + + if len(laser.ri_return1.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test + seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return1.segmentation_label_compressed, "ZLIB") + seg_label = dataset_pb2.MatrixInt32() + seg_label.ParseFromString(seg_label_str_tensor.numpy()) + seg_labels[laser.name] = [seg_label] + if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed(laser.ri_return2.range_image_compressed, "ZLIB") + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(range_image_str_tensor.numpy()) + range_images[laser.name].append(ri) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return2.camera_projection_compressed, "ZLIB" + ) + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(camera_projection_str_tensor.numpy()) + camera_projections[laser.name].append(cp) + + if len(laser.ri_return2.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test + seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return2.segmentation_label_compressed, "ZLIB") + seg_label = dataset_pb2.MatrixInt32() + seg_label.ParseFromString(seg_label_str_tensor.numpy()) + seg_labels[laser.name].append(seg_label) + return range_images, camera_projections, seg_labels, range_image_top_pose diff --git a/d123/dataset/logs/log_metadata.py b/d123/dataset/logs/log_metadata.py index 018156e4..818bc495 100644 --- a/d123/dataset/logs/log_metadata.py +++ b/d123/dataset/logs/log_metadata.py @@ -8,6 +8,10 @@ @dataclass class LogMetadata: + # TODO: add + # - split + # - global/local map + dataset: str log_name: str location: str diff --git a/d123/dataset/maps/abstract_map.py b/d123/dataset/maps/abstract_map.py index 5db687be..681c90cf 100644 --- a/d123/dataset/maps/abstract_map.py +++ b/d123/dataset/maps/abstract_map.py @@ -7,7 +7,7 @@ from d123.common.geometry.base import Point2D from d123.dataset.maps.abstract_map_objects import AbstractMapObject -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer # TODO: # - add docstrings @@ -17,6 +17,7 @@ class AbstractMap(abc.ABC): + @property @abc.abstractmethod def map_name(self) -> str: @@ -27,57 +28,57 @@ def initialize(self) -> None: pass @abc.abstractmethod - def get_available_map_objects(self) -> List[MapSurfaceType]: + def get_available_map_objects(self) -> List[MapLayer]: pass @abc.abstractmethod - def get_map_object(self, object_id: str, layer: MapSurfaceType) -> Optional[AbstractMapObject]: + def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[AbstractMapObject]: pass @abc.abstractmethod - def get_all_map_objects(self, point_2d: Point2D, layer: MapSurfaceType) -> List[AbstractMapObject]: + def get_all_map_objects(self, point_2d: Point2D, layer: MapLayer) -> List[AbstractMapObject]: pass @abc.abstractmethod - def is_in_layer(self, point: Point2D, layer: MapSurfaceType) -> bool: + def is_in_layer(self, point: Point2D, layer: MapLayer) -> bool: pass @abc.abstractmethod def get_proximal_map_objects( - self, point: Point2D, radius: float, layers: List[MapSurfaceType] - ) -> Dict[MapSurfaceType, List[AbstractMapObject]]: + self, point: Point2D, radius: float, layers: List[MapLayer] + ) -> Dict[MapLayer, List[AbstractMapObject]]: pass @abc.abstractmethod def query( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: + ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: pass @abc.abstractmethod def query_object_ids( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]]: + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: pass @abc.abstractmethod def query_nearest( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], return_all: bool = True, max_distance: Optional[float] = None, return_distance: bool = False, exclusive: bool = False, - ) -> Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: + ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: pass diff --git a/d123/dataset/maps/abstract_map_objects.py b/d123/dataset/maps/abstract_map_objects.py index e8e33f8a..96ce79de 100644 --- a/d123/dataset/maps/abstract_map_objects.py +++ b/d123/dataset/maps/abstract_map_objects.py @@ -1,13 +1,13 @@ from __future__ import annotations import abc -from typing import List, Optional +from typing import List, Optional, Tuple import shapely.geometry as geom import trimesh -from d123.common.geometry.line.polylines import Polyline2D, Polyline3D -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.common.geometry.line.polylines import Polyline2D, Polyline3D, PolylineSE2 +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType class AbstractMapObject(abc.ABC): @@ -22,6 +22,14 @@ def __init__(self, object_id: str): """ self.id = str(object_id) + @property + @abc.abstractmethod + def layer(self) -> MapLayer: + """ + Returns map layer type, e.g. LANE, ROAD_EDGE. + :return: map layer type + """ + class AbstractSurfaceMapObject(AbstractMapObject): """ @@ -36,14 +44,6 @@ def shapely_polygon(self) -> geom.Polygon: :return: shapely polygon """ - @property - @abc.abstractmethod - def surface_type(self) -> MapSurfaceType: - """ - Returns map surface type, e.g. LANE. - :return: map surface type - """ - @property @abc.abstractmethod def outline_3d(self) -> Polyline3D: @@ -64,12 +64,39 @@ def outline_2d(self) -> Polyline2D: return self.outline_3d.polyline_2d +class AbstractLineMapObject(AbstractMapObject): + + @property + @abc.abstractmethod + def polyline_3d(self) -> Polyline3D: + """ + Returns the 3D polyline of the road edge. + :return: 3D polyline + """ + + @property + def polyline_2d(self) -> Polyline2D: + """ + Returns the 2D polyline of the road line. + :return: 2D polyline + """ + return self.polyline_3d.polyline_2d + + @property + def polyline_se2(self) -> PolylineSE2: + """ + Returns the 2D polyline of the road line in SE(2) coordinates. + :return: 2D polyline in SE(2) + """ + return self.polyline_3d.polyline_se2 + + class AbstractLane(AbstractSurfaceMapObject): """Abstract interface for lane objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.LANE + def layer(self) -> MapLayer: + return MapLayer.LANE @property @abc.abstractmethod @@ -111,6 +138,22 @@ def right_boundary(self) -> Polyline3D: :return: returns 3D polyline """ + @property + @abc.abstractmethod + def left_lane(self) -> Optional[AbstractLane]: + """ + Property of left lane of lane. + :return: returns left lane or none, if no left lane + """ + + @property + @abc.abstractmethod + def right_lane(self) -> Optional[AbstractLane]: + """ + Property of right lane of lane. + :return: returns right lane or none, if no right lane + """ + @property @abc.abstractmethod def centerline(self) -> Polyline3D: @@ -127,13 +170,21 @@ def lane_group(self) -> AbstractLaneGroup: :return: returns lane group """ + @property + def boundaries(self) -> Tuple[Polyline3D, Polyline3D]: + """ + Property of left and right boundary. + :return: returns tuple of left and right boundary polylines + """ + return self.left_boundary, self.right_boundary + class AbstractLaneGroup(AbstractSurfaceMapObject): """Abstract interface lane groups (nearby lanes going in the same direction).""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.LANE_GROUP + def layer(self) -> MapLayer: + return MapLayer.LANE_GROUP @property @abc.abstractmethod @@ -188,8 +239,8 @@ class AbstractIntersection(AbstractSurfaceMapObject): """Abstract interface for intersection objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.INTERSECTION + def layer(self) -> MapLayer: + return MapLayer.INTERSECTION @property @abc.abstractmethod @@ -204,38 +255,69 @@ class AbstractCrosswalk(AbstractSurfaceMapObject): """Abstract interface for crosswalk objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.CROSSWALK + def layer(self) -> MapLayer: + return MapLayer.CROSSWALK class AbstractWalkway(AbstractSurfaceMapObject): """Abstract interface for walkway objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.WALKWAY + def layer(self) -> MapLayer: + return MapLayer.WALKWAY class AbstractCarpark(AbstractSurfaceMapObject): """Abstract interface for carpark objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.CARPARK + def layer(self) -> MapLayer: + return MapLayer.CARPARK class AbstractGenericDrivable(AbstractSurfaceMapObject): """Abstract interface for generic drivable objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.GENERIC_DRIVABLE + def layer(self) -> MapLayer: + return MapLayer.GENERIC_DRIVABLE class AbstractStopLine(AbstractSurfaceMapObject): """Abstract interface for stop line objects.""" @property - def surface_type(self) -> MapSurfaceType: - # return MapSurfaceType.STOP_LINE - raise NotImplementedError + def layer(self) -> MapLayer: + return MapLayer.STOP_LINE + + +class AbstractRoadEdge(AbstractLineMapObject): + """Abstract interface for road edge objects.""" + + @property + def layer(self) -> MapLayer: + return MapLayer.ROAD_EDGE + + @property + @abc.abstractmethod + def road_edge_type(self) -> RoadEdgeType: + """ + Returns the road edge type. + :return: RoadEdgeType + """ + + +class AbstractRoadLine(AbstractLineMapObject): + """Abstract interface for road line objects.""" + + @property + def layer(self) -> MapLayer: + return MapLayer.ROAD_LINE + + @property + @abc.abstractmethod + def road_line_type(self) -> RoadLineType: + """ + Returns the road line type. + :return: RoadLineType + """ diff --git a/d123/dataset/maps/gpkg/gpkg_map.py b/d123/dataset/maps/gpkg/gpkg_map.py index d202f2ce..32b02d41 100644 --- a/d123/dataset/maps/gpkg/gpkg_map.py +++ b/d123/dataset/maps/gpkg/gpkg_map.py @@ -21,10 +21,12 @@ GPKGIntersection, GPKGLane, GPKGLaneGroup, + GPKGRoadEdge, + GPKGRoadLine, GPKGWalkway, ) from d123.dataset.maps.gpkg.utils import load_gdf_with_geometry_columns -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer USE_ARROW: bool = True @@ -33,18 +35,20 @@ class GPKGMap(AbstractMap): def __init__(self, file_path: Path) -> None: self._file_path = file_path - self._map_object_getter: Dict[MapSurfaceType, Callable[[str], Optional[AbstractMapObject]]] = { - MapSurfaceType.LANE: self._get_lane, - MapSurfaceType.LANE_GROUP: self._get_lane_group, - MapSurfaceType.INTERSECTION: self._get_intersection, - MapSurfaceType.CROSSWALK: self._get_crosswalk, - MapSurfaceType.WALKWAY: self._get_walkway, - MapSurfaceType.CARPARK: self._get_carpark, - MapSurfaceType.GENERIC_DRIVABLE: self._get_generic_drivable, + self._map_object_getter: Dict[MapLayer, Callable[[str], Optional[AbstractMapObject]]] = { + MapLayer.LANE: self._get_lane, + MapLayer.LANE_GROUP: self._get_lane_group, + MapLayer.INTERSECTION: self._get_intersection, + MapLayer.CROSSWALK: self._get_crosswalk, + MapLayer.WALKWAY: self._get_walkway, + MapLayer.CARPARK: self._get_carpark, + MapLayer.GENERIC_DRIVABLE: self._get_generic_drivable, + MapLayer.ROAD_EDGE: self._get_road_edge, + MapLayer.ROAD_LINE: self._get_road_line, } # loaded during `.initialize()` - self._gpd_dataframes: Dict[MapSurfaceType, gpd.GeoDataFrame] = {} + self._gpd_dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} @property def map_name(self) -> str: @@ -55,7 +59,7 @@ def initialize(self) -> None: """Inherited, see superclass.""" if len(self._gpd_dataframes) == 0: available_layers = list(gpd.list_layers(self._file_path).name) - for map_layer in list(MapSurfaceType): + for map_layer in list(MapLayer): map_layer_name = map_layer.serialize() if map_layer_name in available_layers: self._gpd_dataframes[map_layer] = gpd.read_file( @@ -70,21 +74,22 @@ def initialize(self) -> None: self._gpd_dataframes[map_layer]["id"] = self._gpd_dataframes[map_layer]["id"].astype(str) else: warnings.warn(f"GPKGMap: {map_layer_name} not available in {str(self._file_path)}") + self._gpd_dataframes[map_layer] = None def _assert_initialize(self) -> None: "Checks if `.initialize()` was called, before retrieving data." assert len(self._gpd_dataframes) > 0, "GPKGMap: Call `.initialize()` before retrieving data!" - def _assert_layer_available(self, layer: MapSurfaceType) -> None: + def _assert_layer_available(self, layer: MapLayer) -> None: "Checks if layer is available." - assert layer in self.get_available_map_objects(), f"GPKGMap: MapSurfaceType {layer.name} is unavailable." + assert layer in self.get_available_map_objects(), f"GPKGMap: MapLayer {layer.name} is unavailable." - def get_available_map_objects(self) -> List[MapSurfaceType]: + def get_available_map_objects(self) -> List[MapLayer]: """Inherited, see superclass.""" self._assert_initialize() return list(self._gpd_dataframes.keys()) - def get_map_object(self, object_id: str, layer: MapSurfaceType) -> Optional[AbstractMapObject]: + def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[AbstractMapObject]: """Inherited, see superclass.""" self._assert_initialize() @@ -94,17 +99,17 @@ def get_map_object(self, object_id: str, layer: MapSurfaceType) -> Optional[Abst except KeyError: raise ValueError(f"Object representation for layer: {layer.name} object: {object_id} is unavailable") - def get_all_map_objects(self, point_2d: Point2D, layer: MapSurfaceType) -> List[AbstractMapObject]: + def get_all_map_objects(self, point_2d: Point2D, layer: MapLayer) -> List[AbstractMapObject]: """Inherited, see superclass.""" raise NotImplementedError - def is_in_layer(self, point: Point2D, layer: MapSurfaceType) -> bool: + def is_in_layer(self, point: Point2D, layer: MapLayer) -> bool: """Inherited, see superclass.""" raise NotImplementedError def get_proximal_map_objects( - self, point_2d: Point2D, radius: float, layers: List[MapSurfaceType] - ) -> Dict[MapSurfaceType, List[AbstractMapObject]]: + self, point_2d: Point2D, radius: float, layers: List[MapLayer] + ) -> Dict[MapLayer, List[AbstractMapObject]]: """Inherited, see superclass.""" center_point = geom.Point(point_2d.x, point_2d.y) patch = center_point.buffer(radius) @@ -113,16 +118,16 @@ def get_proximal_map_objects( def query( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: + ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: supported_layers = self.get_available_map_objects() unsupported_layers = [layer for layer in layers if layer not in supported_layers] assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]] = ( - defaultdict(list) + object_map: Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]] = defaultdict( + list ) for layer in layers: object_map[layer] = self._query_layer(geometry, layer, predicate, sort, distance) @@ -131,15 +136,15 @@ def query( def query_object_ids( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]]: + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: supported_layers = self.get_available_map_objects() unsupported_layers = [layer for layer in layers if layer not in supported_layers] assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) + object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) for layer in layers: object_map[layer] = self._query_layer_objects_ids(geometry, layer, predicate, sort, distance) return object_map @@ -147,16 +152,16 @@ def query_object_ids( def query_nearest( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], return_all: bool = True, max_distance: Optional[float] = None, return_distance: bool = False, exclusive: bool = False, - ) -> Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]]: + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]]: supported_layers = self.get_available_map_objects() unsupported_layers = [layer for layer in layers if layer not in supported_layers] assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) + object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) for layer in layers: object_map[layer] = self._query_layer_nearest( geometry, layer, return_all, max_distance, return_distance, exclusive @@ -166,7 +171,7 @@ def query_nearest( def _query_layer( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapSurfaceType, + layer: MapLayer, predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, @@ -188,7 +193,7 @@ def _query_layer( def _query_layer_objects_ids( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapSurfaceType, + layer: MapLayer, predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, @@ -211,7 +216,7 @@ def _query_layer_objects_ids( def _query_layer_nearest( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapSurfaceType, + layer: MapLayer, return_all: bool = True, max_distance: Optional[float] = None, return_distance: bool = False, @@ -247,11 +252,11 @@ def _get_lane(self, id: str) -> Optional[GPKGLane]: return ( GPKGLane( id, - self._gpd_dataframes[MapSurfaceType.LANE], - self._gpd_dataframes[MapSurfaceType.LANE_GROUP], - self._gpd_dataframes[MapSurfaceType.INTERSECTION], + self._gpd_dataframes[MapLayer.LANE], + self._gpd_dataframes[MapLayer.LANE_GROUP], + self._gpd_dataframes[MapLayer.INTERSECTION], ) - if id in self._gpd_dataframes[MapSurfaceType.LANE]["id"].tolist() + if id in self._gpd_dataframes[MapLayer.LANE]["id"].tolist() else None ) @@ -259,11 +264,11 @@ def _get_lane_group(self, id: str) -> Optional[GPKGLaneGroup]: return ( GPKGLaneGroup( id, - self._gpd_dataframes[MapSurfaceType.LANE_GROUP], - self._gpd_dataframes[MapSurfaceType.LANE], - self._gpd_dataframes[MapSurfaceType.INTERSECTION], + self._gpd_dataframes[MapLayer.LANE_GROUP], + self._gpd_dataframes[MapLayer.LANE], + self._gpd_dataframes[MapLayer.INTERSECTION], ) - if id in self._gpd_dataframes[MapSurfaceType.LANE_GROUP]["id"].tolist() + if id in self._gpd_dataframes[MapLayer.LANE_GROUP]["id"].tolist() else None ) @@ -271,46 +276,60 @@ def _get_intersection(self, id: str) -> Optional[GPKGIntersection]: return ( GPKGIntersection( id, - self._gpd_dataframes[MapSurfaceType.INTERSECTION], - self._gpd_dataframes[MapSurfaceType.LANE], - self._gpd_dataframes[MapSurfaceType.LANE_GROUP], + self._gpd_dataframes[MapLayer.INTERSECTION], + self._gpd_dataframes[MapLayer.LANE], + self._gpd_dataframes[MapLayer.LANE_GROUP], ) - if id in self._gpd_dataframes[MapSurfaceType.INTERSECTION]["id"].tolist() + if id in self._gpd_dataframes[MapLayer.INTERSECTION]["id"].tolist() else None ) def _get_crosswalk(self, id: str) -> Optional[GPKGCrosswalk]: return ( - GPKGCrosswalk(id, self._gpd_dataframes[MapSurfaceType.CROSSWALK]) - if id in self._gpd_dataframes[MapSurfaceType.CROSSWALK]["id"].tolist() + GPKGCrosswalk(id, self._gpd_dataframes[MapLayer.CROSSWALK]) + if id in self._gpd_dataframes[MapLayer.CROSSWALK]["id"].tolist() else None ) def _get_walkway(self, id: str) -> Optional[GPKGWalkway]: return ( - GPKGWalkway(id, self._gpd_dataframes[MapSurfaceType.WALKWAY]) - if id in self._gpd_dataframes[MapSurfaceType.WALKWAY]["id"].tolist() + GPKGWalkway(id, self._gpd_dataframes[MapLayer.WALKWAY]) + if id in self._gpd_dataframes[MapLayer.WALKWAY]["id"].tolist() else None ) def _get_carpark(self, id: str) -> Optional[GPKGCarpark]: return ( - GPKGCarpark(id, self._gpd_dataframes[MapSurfaceType.CARPARK]) - if id in self._gpd_dataframes[MapSurfaceType.CARPARK]["id"].tolist() + GPKGCarpark(id, self._gpd_dataframes[MapLayer.CARPARK]) + if id in self._gpd_dataframes[MapLayer.CARPARK]["id"].tolist() else None ) def _get_generic_drivable(self, id: str) -> Optional[GPKGGenericDrivable]: return ( - GPKGGenericDrivable(id, self._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE]) - if id in self._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE]["id"].tolist() + GPKGGenericDrivable(id, self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE]) + if id in self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE]["id"].tolist() + else None + ) + + def _get_road_edge(self, id: str) -> Optional[GPKGRoadEdge]: + return ( + GPKGRoadEdge(id, self._gpd_dataframes[MapLayer.ROAD_EDGE]) + if id in self._gpd_dataframes[MapLayer.ROAD_EDGE]["id"].tolist() + else None + ) + + def _get_road_line(self, id: str) -> Optional[GPKGRoadLine]: + return ( + GPKGRoadLine(id, self._gpd_dataframes[MapLayer.ROAD_LINE]) + if id in self._gpd_dataframes[MapLayer.ROAD_LINE]["id"].tolist() else None ) # def _query_layer( # self, # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - # layer: MapSurfaceType, + # layer: MapLayer, # predicate: Optional[str] = None, # sort: bool = False, # distance: Optional[float] = None, @@ -332,7 +351,7 @@ def _get_generic_drivable(self, id: str) -> Optional[GPKGGenericDrivable]: # def _query_layer_objects_ids( # self, # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - # layer: MapSurfaceType, + # layer: MapLayer, # predicate: Optional[str] = None, # sort: bool = False, # distance: Optional[float] = None, @@ -359,3 +378,13 @@ def get_map_api_from_names(dataset: str, location: str) -> GPKGMap: map_api = GPKGMap(gpkg_path) map_api.initialize() return map_api + + +def get_local_map_api(split_name: str, log_name: str) -> GPKGMap: + print(split_name, log_name) + D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) + gpkg_path = D123_MAPS_ROOT / split_name / f"{log_name}.gpkg" + assert gpkg_path.is_file(), f"{log_name}.gpkg not found in {str(D123_MAPS_ROOT)}." + map_api = GPKGMap(gpkg_path) + map_api.initialize() + return map_api diff --git a/d123/dataset/maps/gpkg/gpkg_map_objects.py b/d123/dataset/maps/gpkg/gpkg_map_objects.py index 6965b885..1fa6de6a 100644 --- a/d123/dataset/maps/gpkg/gpkg_map_objects.py +++ b/d123/dataset/maps/gpkg/gpkg_map_objects.py @@ -6,6 +6,7 @@ import geopandas as gpd import numpy as np +import pandas as pd import shapely.geometry as geom import trimesh @@ -19,10 +20,14 @@ AbstractIntersection, AbstractLane, AbstractLaneGroup, + AbstractLineMapObject, + AbstractRoadEdge, + AbstractRoadLine, AbstractSurfaceMapObject, AbstractWalkway, ) from d123.dataset.maps.gpkg.utils import get_row_with_value +from d123.dataset.maps.map_datatypes import RoadEdgeType, RoadLineType class GPKGSurfaceObject(AbstractSurfaceMapObject): @@ -30,7 +35,6 @@ class GPKGSurfaceObject(AbstractSurfaceMapObject): Base interface representation of all map objects. """ - # TODO: Extend for 3D outline def __init__(self, object_id: str, surface_df: gpd.GeoDataFrame) -> None: """ Constructor of the base surface map object type. @@ -62,10 +66,49 @@ def outline_3d(self) -> Polyline3D: @property def trimesh_mesh(self) -> trimesh.Trimesh: """Inherited, see superclass.""" - outline_3d_array = self.outline_3d.array - _, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) - # NOTE: Optional add color information to the mesh - return trimesh.Trimesh(vertices=outline_3d_array, faces=faces) + + trimesh_mesh: Optional[trimesh.Trimesh] = None + if "right_boundary" in self._object_df.columns and "left_boundary" in self._object_df.columns: + left_boundary = Polyline3D.from_linestring(self._object_row.left_boundary) + right_boundary = Polyline3D.from_linestring(self._object_row.right_boundary) + trimesh_mesh = get_trimesh_from_boundaries(left_boundary, right_boundary) + else: + # Fallback to geometry if no boundaries are available + outline_3d_array = self.outline_3d.array + vertices_2d, faces = trimesh.creation.triangulate_polygon( + geom.Polygon(outline_3d_array[:, Point3DIndex.XY]) + ) + if len(vertices_2d) == len(outline_3d_array): + # Regular case, where vertices match outline_3d_array + vertices_3d = outline_3d_array + elif len(vertices_2d) == len(outline_3d_array) + 1: + # outline array was not closed, so we need to add the first vertex again + vertices_3d = np.vstack((outline_3d_array, outline_3d_array[0])) + else: + raise ValueError("No vertices found for triangulation.") + trimesh_mesh = trimesh.Trimesh(vertices=vertices_3d, faces=faces) + return trimesh_mesh + + +class GPKGLineObject(AbstractLineMapObject): + + def __init__(self, object_id: str, line_df: gpd.GeoDataFrame) -> None: + """ + Constructor of the base line map object type. + :param object_id: unique identifier of a line map object. + """ + super().__init__(object_id) + # TODO: add assertion if columns are available + self._object_df = line_df + + @cached_property + def _object_row(self) -> gpd.GeoSeries: + return get_row_with_value(self._object_df, "id", self.id) + + @property + def polyline_3d(self) -> Polyline3D: + """Inherited, see superclass.""" + return Polyline3D.from_linestring(self._object_row.geometry) class GPKGLane(GPKGSurfaceObject, AbstractLane): @@ -107,6 +150,26 @@ def right_boundary(self) -> Polyline3D: """Inherited, see superclass.""" return Polyline3D.from_linestring(self._object_row.right_boundary) + @property + def left_lane(self) -> Optional[GPKGLane]: + """Inherited, see superclass.""" + left_lane_id = self._object_row.left_lane_id + return ( + GPKGLane(left_lane_id, self._object_df, self._lane_group_df, self._intersection_df) + if left_lane_id is not None and not pd.isna(left_lane_id) + else None + ) + + @property + def right_lane(self) -> Optional[GPKGLane]: + """Inherited, see superclass.""" + right_lane_id = self._object_row.right_lane_id + return ( + GPKGLane(right_lane_id, self._object_df, self._lane_group_df, self._intersection_df) + if right_lane_id is not None and not pd.isna(right_lane_id) + else None + ) + @property def centerline(self) -> Polyline3D: """Inherited, see superclass.""" @@ -119,11 +182,6 @@ def outline_3d(self) -> Polyline3D: outline_array = np.vstack((outline_array, outline_array[0])) return Polyline3D.from_linestring(geom.LineString(outline_array)) - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - return get_trimesh_from_boundaries(self.left_boundary, self.right_boundary) - @property def lane_group(self) -> GPKGLaneGroup: """Inherited, see superclass.""" @@ -152,13 +210,19 @@ def __init__( def successors(self) -> List[GPKGLaneGroup]: """Inherited, see superclass.""" successor_ids = ast.literal_eval(self._object_row.successor_ids) - return [GPKGLaneGroup(lane_group_id, self._object_df) for lane_group_id in successor_ids] + return [ + GPKGLaneGroup(lane_group_id, self._object_df, self._lane_df, self._intersection_df) + for lane_group_id in successor_ids + ] @property def predecessors(self) -> List[GPKGLaneGroup]: """Inherited, see superclass.""" predecessor_ids = ast.literal_eval(self._object_row.predecessor_ids) - return [GPKGLaneGroup(lane_group_id, self._object_df) for lane_group_id in predecessor_ids] + return [ + GPKGLaneGroup(lane_group_id, self._object_df, self._lane_df, self._intersection_df) + for lane_group_id in predecessor_ids + ] @property def left_boundary(self) -> Polyline3D: @@ -176,11 +240,6 @@ def outline_3d(self) -> Polyline3D: outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) return Polyline3D.from_linestring(geom.LineString(outline_array)) - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - return get_trimesh_from_boundaries(self.left_boundary, self.right_boundary) - @property def lanes(self) -> List[GPKGLane]: """Inherited, see superclass.""" @@ -206,7 +265,7 @@ def intersection(self) -> Optional[GPKGIntersection]: self._lane_df, self._object_df, ) - if intersection_id is not None + if intersection_id is not None and not pd.isna(intersection_id) else None ) @@ -256,3 +315,31 @@ def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): class GPKGGenericDrivable(GPKGSurfaceObject, AbstractGenericDrivable): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): super().__init__(object_id, object_df) + + +class GPKGRoadEdge(GPKGLineObject, AbstractRoadEdge): + def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): + super().__init__(object_id, object_df) + + @cached_property + def _object_row(self) -> gpd.GeoSeries: + return get_row_with_value(self._object_df, "id", self.id) + + @property + def road_edge_type(self) -> RoadEdgeType: + """Inherited, see superclass.""" + return RoadEdgeType(int(self._object_row.road_edge_type)) + + +class GPKGRoadLine(GPKGLineObject, AbstractRoadLine): + def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): + super().__init__(object_id, object_df) + + @cached_property + def _object_row(self) -> gpd.GeoSeries: + return get_row_with_value(self._object_df, "id", self.id) + + @property + def road_line_type(self) -> RoadLineType: + """Inherited, see superclass.""" + return RoadLineType(int(self._object_row.road_line_type)) diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index ca48497e..f948422f 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -5,7 +5,7 @@ # TODO: Add stop pads or stop lines. -class MapSurfaceType(SerialIntEnum): +class MapLayer(SerialIntEnum): """ Enum for AbstractMapSurface. """ @@ -18,3 +18,54 @@ class MapSurfaceType(SerialIntEnum): CARPARK = 5 GENERIC_DRIVABLE = 6 STOP_LINE = 7 + ROAD_EDGE = 8 + ROAD_LINE = 9 + + +class LaneType(SerialIntEnum): + """ + Enum for LaneType. + https://github.com/waymo-research/waymo-open-dataset/blob/99a4cb3ff07e2fe06c2ce73da001f850f628e45a/src/waymo_open_dataset/protos/map.proto#L147 + """ + + UNDEFINED = 0 + FREEWAY = 1 + SURFACE_STREET = 2 + BIKE_LANE = 3 + + +class RoadEdgeType(SerialIntEnum): + """ + Enum for RoadEdgeType. + NOTE: We use the road line types from Waymo. + https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L188 + """ + + UNKNOWN = 0 + ROAD_EDGE_BOUNDARY = 1 + ROAD_EDGE_MEDIAN = 2 + + +class RoadLineType(SerialIntEnum): + """ + Enum for RoadLineType. + NOTE: We use the road line types from Waymo. + TODO: Use the Argoverse 2 road line types instead. + https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L208 + """ + + NONE = 0 + UNKNOWN = 1 + DASH_SOLID_YELLOW = 2 + DASH_SOLID_WHITE = 3 + DASHED_WHITE = 4 + DASHED_YELLOW = 5 + DOUBLE_SOLID_YELLOW = 6 + DOUBLE_SOLID_WHITE = 7 + DOUBLE_DASH_YELLOW = 8 + DOUBLE_DASH_WHITE = 9 + SOLID_YELLOW = 10 + SOLID_WHITE = 11 + SOLID_DASH_WHITE = 12 + SOLID_DASH_YELLOW = 13 + SOLID_BLUE = 14 diff --git a/d123/dataset/scene/abstract_scene.py b/d123/dataset/scene/abstract_scene.py index 95bf7ddf..4877fc83 100644 --- a/d123/dataset/scene/abstract_scene.py +++ b/d123/dataset/scene/abstract_scene.py @@ -2,12 +2,12 @@ import abc from dataclasses import dataclass -from typing import List - -from PIL import Image +from typing import List, Optional from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording +from d123.common.datatypes.sensor.camera import Camera, CameraType +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARType from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.dataset.logs.log_metadata import LogMetadata @@ -17,10 +17,6 @@ class AbstractScene(abc.ABC): - @property - @abc.abstractmethod - def map_api(self) -> AbstractMap: - raise NotImplementedError @property @abc.abstractmethod @@ -37,6 +33,21 @@ def token(self) -> str: def log_metadata(self) -> LogMetadata: raise NotImplementedError + @property + @abc.abstractmethod + def available_camera_types(self) -> List[CameraType]: + raise NotImplementedError + + @property + @abc.abstractmethod + def available_lidar_types(self) -> List[LiDARType]: + raise NotImplementedError + + @property + @abc.abstractmethod + def map_api(self) -> Optional[AbstractMap]: + raise NotImplementedError + @abc.abstractmethod def get_number_of_iterations(self) -> int: raise NotImplementedError @@ -70,7 +81,11 @@ def get_route_lane_group_ids(self, iteration: int) -> List[int]: raise NotImplementedError @abc.abstractmethod - def get_front_cam_demo(self, iteration: int) -> Image: + def get_camera_at_iteration(self, iteration: int, camera_type: CameraType) -> Camera: + raise NotImplementedError + + @abc.abstractmethod + def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> LiDAR: raise NotImplementedError def open(self) -> None: diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index b0c7a28d..ecd68111 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -1,32 +1,36 @@ -import io import json from pathlib import Path -from typing import List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import pyarrow as pa -# TODO: Remove or improve open/close dynamic of Scene object. -from PIL import Image - from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording +from d123.common.datatypes.sensor.camera import Camera, CameraMetadata, CameraType, camera_metadata_dict_from_json +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata, LiDARType, lidar_metadata_dict_from_json from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.dataset.arrow.conversion import ( get_box_detections_from_arrow_table, + get_camera_from_arrow_table, get_ego_vehicle_state_from_arrow_table, + get_lidar_from_arrow_table, get_timepoint_from_arrow_table, get_traffic_light_detections_from_arrow_table, ) from d123.dataset.arrow.helper import open_arrow_table from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import AbstractMap -from d123.dataset.maps.gpkg.gpkg_map import get_map_api_from_names +from d123.dataset.maps.gpkg.gpkg_map import get_local_map_api, get_map_api_from_names from d123.dataset.scene.abstract_scene import AbstractScene, SceneExtractionInfo +# TODO: Remove or improve open/close dynamic of Scene object. + -def _get_scene_data(arrow_file_path: Union[Path, str]) -> Tuple[LogMetadata, VehicleParameters]: +def _get_scene_data( + arrow_file_path: Union[Path, str], +) -> Tuple[LogMetadata, VehicleParameters, Dict[CameraType, CameraMetadata]]: """ Extracts the metadata and vehicle parameters from the arrow file. """ @@ -34,8 +38,19 @@ def _get_scene_data(arrow_file_path: Union[Path, str]) -> Tuple[LogMetadata, Veh table = open_arrow_table(arrow_file_path) metadata = LogMetadata(**json.loads(table.schema.metadata[b"log_metadata"].decode())) vehicle_parameters = VehicleParameters(**json.loads(table.schema.metadata[b"vehicle_parameters"].decode())) + + if b"camera_metadata" in table.schema.metadata: + camera_metadata = camera_metadata_dict_from_json(table.schema.metadata[b"camera_metadata"].decode()) + else: + camera_metadata = {} + + if b"lidar_metadata" in table.schema.metadata: + lidar_metadata = lidar_metadata_dict_from_json(table.schema.metadata[b"lidar_metadata"].decode()) + else: + lidar_metadata = {} + del table - return metadata, vehicle_parameters + return metadata, vehicle_parameters, camera_metadata, lidar_metadata class ArrowScene(AbstractScene): @@ -47,9 +62,16 @@ def __init__( self._recording_table: pa.Table = None - _metadata, _vehicle_parameters = _get_scene_data(arrow_file_path) + ( + _metadata, + _vehicle_parameters, + _camera_metadata, + _lidar_metadata, + ) = _get_scene_data(arrow_file_path) self._metadata: LogMetadata = _metadata self._vehicle_parameters: VehicleParameters = _vehicle_parameters + self._camera_metadata: Dict[CameraType, CameraMetadata] = _camera_metadata + self._lidar_metadata: Dict[LiDARType, LiDARMetadata] = _lidar_metadata self._map_api: Optional[AbstractMap] = None @@ -83,6 +105,14 @@ def token(self) -> str: def log_metadata(self) -> LogMetadata: return self._metadata + @property + def available_camera_types(self) -> List[CameraType]: + return list(self._camera_metadata.keys()) + + @property + def available_lidar_types(self) -> List[LiDARType]: + return list(self._lidar_metadata.keys()) + def _get_table_index(self, iteration: int) -> int: self._lazy_initialize() assert ( @@ -128,19 +158,43 @@ def get_route_lane_group_ids(self, iteration: int) -> List[int]: table_index = self._get_table_index(iteration) return self._recording_table["route_lane_group_ids"][table_index].as_py() - def get_front_cam_demo(self, iteration: int) -> Image: + def get_camera_at_iteration(self, iteration: int, camera_type: CameraType) -> Camera: + self._lazy_initialize() + assert camera_type in self._camera_metadata, f"Camera type {camera_type} not found in metadata." + table_index = self._get_table_index(iteration) + return get_camera_from_arrow_table( + self._recording_table, + table_index, + self._camera_metadata[camera_type], + self.log_metadata, + ) + + def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> LiDAR: self._lazy_initialize() + assert lidar_type in self._lidar_metadata, f"LiDAR type {lidar_type} not found in metadata." table_index = self._get_table_index(iteration) - jpg_data = self._recording_table["front_cam_demo"][table_index].as_py() - return Image.open(io.BytesIO(jpg_data)) + return get_lidar_from_arrow_table( + self._recording_table, + table_index, + self._lidar_metadata[lidar_type], + self.log_metadata, + ) def _lazy_initialize(self) -> None: self.open() def open(self) -> None: if self._map_api is None: - self._map_api = get_map_api_from_names(self._metadata.dataset, self._metadata.location) - self._map_api.initialize() + try: + if self._metadata.dataset in ["wopd", "av2-sensor"]: + # FIXME: + split = str(self._arrow_log_path.parent.name) + self._map_api = get_local_map_api(split, self._metadata.log_name) + else: + self._map_api = get_map_api_from_names(self._metadata.dataset, self._metadata.location) + self._map_api.initialize() + except Exception as e: + print(f"Error initializing map API: {e}") if self._recording_table is None: self._recording_table = open_arrow_table(self._arrow_log_path) if self._scene_extraction_info is None: diff --git a/d123/dataset/scene/scene_builder.py b/d123/dataset/scene/scene_builder.py index 0ea3c2c3..228fcb13 100644 --- a/d123/dataset/scene/scene_builder.py +++ b/d123/dataset/scene/scene_builder.py @@ -5,8 +5,7 @@ from pathlib import Path from typing import Iterator, List, Optional, Set, Union -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map - +from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.scene.abstract_scene import AbstractScene @@ -85,7 +84,11 @@ def _discover_log_paths(dataset_path: Path, split_names: Set[str], log_names: Op def _extract_scenes_from_logs(log_paths: List[Path], filter: SceneFilter) -> List[AbstractScene]: scenes: List[AbstractScene] = [] for log_path in log_paths: - scene_extraction_infos = _get_scene_extraction_info(log_path, filter) + try: + scene_extraction_infos = _get_scene_extraction_info(log_path, filter) + except Exception as e: + print(f"Error extracting scenes from {log_path}: {e}") + continue for scene_extraction_info in scene_extraction_infos: scenes.append( ArrowScene( @@ -106,7 +109,7 @@ def _get_scene_extraction_info(log_path: Union[str, Path], filter: SceneFilter) if filter.map_names is not None and log_metadata.map_name not in filter.map_names: return scene_extraction_infos - start_idx = int(filter.history_s / log_metadata.timestep_seconds) # if filter.history_s is not None else 0 + start_idx = int(filter.history_s / log_metadata.timestep_seconds) if filter.history_s is not None else 0 end_idx = ( len(recording_table) - int(filter.duration_s / log_metadata.timestep_seconds) if filter.duration_s is not None @@ -143,12 +146,22 @@ def _get_scene_extraction_info(log_path: Union[str, Path], filter: SceneFilter) ) if scene_extraction_info is not None: - # TODO: add more options + # Check of timestamp threshold exceeded between previous scene, if specified in filter if filter.timestamp_threshold_s is not None and len(scene_extraction_infos) > 0: iteration_delta = idx - scene_extraction_infos[-1].initial_idx if (iteration_delta * log_metadata.timestep_seconds) < filter.timestamp_threshold_s: continue + # Check if camera data is available for the scene, if specified in filter + # NOTE: We only check camera availability at the initial index of the scene. + if filter.camera_types is not None: + cameras_available = [ + recording_table[camera_type.serialize()][start_idx].as_py() is not None + for camera_type in filter.camera_types + ] + if not all(cameras_available): + continue + scene_extraction_infos.append(scene_extraction_info) del recording_table, log_metadata diff --git a/d123/dataset/scene/scene_filter.py b/d123/dataset/scene/scene_filter.py index 69223c1b..0a58eba6 100644 --- a/d123/dataset/scene/scene_filter.py +++ b/d123/dataset/scene/scene_filter.py @@ -1,6 +1,8 @@ from dataclasses import dataclass from typing import List, Optional +from d123.common.datatypes.sensor.camera import CameraType + # TODO: Add more filter options (e.g. scene tags, ego movement, or whatever appropriate) @@ -21,5 +23,22 @@ class SceneFilter: duration_s: Optional[float] = 10.0 history_s: Optional[float] = 3.0 + camera_types: Optional[List[CameraType]] = None + max_num_scenes: Optional[int] = None shuffle: bool = False + + def __post_init__(self): + if self.camera_types is not None: + assert isinstance(self.camera_types, list), "camera_types must be a list of CameraType" + camera_types = [] + for camera_type in self.camera_types: + if isinstance(camera_type, str): + camera_type = CameraType.deserialize[camera_type] + camera_types.append(camera_type) + elif isinstance(camera_type, int): + camera_type = CameraType(camera_type) + camera_types.append(camera_type) + else: + raise ValueError(f"Invalid camera type: {camera_type}") + self.camera_types = camera_types diff --git a/d123/script/builders/data_converter_builder.py b/d123/script/builders/data_converter_builder.py new file mode 100644 index 00000000..d9c54004 --- /dev/null +++ b/d123/script/builders/data_converter_builder.py @@ -0,0 +1,22 @@ +import logging +from typing import List + +from hydra.utils import instantiate +from omegaconf import DictConfig + +from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter +from d123.script.builders.utils.utils_type import validate_type + +logger = logging.getLogger(__name__) + + +def build_data_converter(cfg: DictConfig) -> List[RawDataConverter]: + logger.info("Building RawDataProcessor...") + instantiated_datasets: List[RawDataConverter] = [] + for dataset_type in cfg.values(): + processor: RawDataConverter = instantiate(dataset_type) + validate_type(processor, RawDataConverter) + instantiated_datasets.append(processor) + + logger.info("Building RawDataProcessor...DONE!") + return instantiated_datasets diff --git a/d123/script/builders/data_processor_builder.py b/d123/script/builders/data_processor_builder.py deleted file mode 100644 index cd83cc28..00000000 --- a/d123/script/builders/data_processor_builder.py +++ /dev/null @@ -1,22 +0,0 @@ -import logging -from typing import List - -from hydra.utils import instantiate -from nuplan.planning.script.builders.utils.utils_type import validate_type -from omegaconf import DictConfig - -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor - -logger = logging.getLogger(__name__) - - -def build_data_processor(cfg: DictConfig) -> List[RawDataProcessor]: - logger.info("Building RawDataProcessor...") - instantiated_datasets: List[RawDataProcessor] = [] - for dataset_type in cfg.values(): - processor: RawDataProcessor = instantiate(dataset_type) - validate_type(processor, RawDataProcessor) - instantiated_datasets.append(processor) - - logger.info("Building RawDataProcessor...DONE!") - return instantiated_datasets diff --git a/d123/script/builders/utils/utils_type.py b/d123/script/builders/utils/utils_type.py new file mode 100644 index 00000000..0fe5e9fd --- /dev/null +++ b/d123/script/builders/utils/utils_type.py @@ -0,0 +1,73 @@ +from typing import Any, Callable, Dict, Type, Union, cast + +from hydra._internal.utils import _locate +from omegaconf import DictConfig + + +def is_TorchModuleWrapper_config(cfg: DictConfig) -> bool: + """ + Check whether the config is meant for a TorchModuleWrapper + :param cfg: config + :return: True if model_config and checkpoint_path is in the cfg, False otherwise + """ + return "model_config" in cfg and "checkpoint_path" in cfg + + +def is_target_type(cfg: DictConfig, target_type: Union[Type[Any], Callable[..., Any]]) -> bool: + """ + Check whether the config's resolved type matches the target type or callable. + :param cfg: config + :param target_type: Type or callable to check against. + :return: Whether cfg._target_ matches the target_type. + """ + return bool(_locate(cfg._target_) == target_type) + + +def validate_type(instantiated_class: Any, desired_type: Type[Any]) -> None: + """ + Validate that constructed type is indeed the desired one + :param instantiated_class: class that was created + :param desired_type: type that the created class should have + """ + assert isinstance( + instantiated_class, desired_type + ), f"Class to be of type {desired_type}, but is {type(instantiated_class)}!" + + +def are_the_same_type(lhs: Any, rhs: Any) -> None: + """ + Validate that lhs and rhs are of the same type + :param lhs: left argument + :param rhs: right argument + """ + lhs_type = type(lhs) + rhs_type = type(rhs) + assert lhs_type == rhs_type, f"Lhs and Rhs are not of the same type! {lhs_type} != {rhs_type}!" + + +def validate_dict_type(instantiated_dict: Dict[str, Any], desired_type: Type[Any]) -> None: + """ + Validate that all entries in dict is indeed the desired one + :param instantiated_dict: dictionary that was created + :param desired_type: type that the created class should have + """ + for value in instantiated_dict.values(): + if isinstance(value, dict): + validate_dict_type(value, desired_type) + else: + validate_type(value, desired_type) + + +def find_builder_in_config(cfg: DictConfig, desired_type: Type[Any]) -> DictConfig: + """ + Find the corresponding config for the desired builder + :param cfg: config structured as a dictionary + :param desired_type: desired builder type + :return: found config + @raise ValueError if the config cannot be found for the builder + """ + for cfg_builder in cfg.values(): + if is_target_type(cfg_builder, desired_type): + return cast(DictConfig, cfg_builder) + + raise ValueError(f"Config does not exist for builder type: {desired_type}!") diff --git a/d123/script/builders/worker_pool_builder.py b/d123/script/builders/worker_pool_builder.py index 4d13877d..123ec3c3 100644 --- a/d123/script/builders/worker_pool_builder.py +++ b/d123/script/builders/worker_pool_builder.py @@ -1,12 +1,11 @@ import logging from hydra.utils import instantiate -from nuplan.planning.script.builders.utils.utils_type import is_target_type, validate_type -from nuplan.planning.utils.multithreading.worker_parallel import SingleMachineParallelExecutor -from nuplan.planning.utils.multithreading.worker_pool import WorkerPool -from nuplan.planning.utils.multithreading.worker_sequential import Sequential from omegaconf import DictConfig +from d123.common.multithreading.worker_pool import WorkerPool +from d123.script.builders.utils.utils_type import validate_type + logger = logging.getLogger(__name__) @@ -17,12 +16,7 @@ def build_worker(cfg: DictConfig) -> WorkerPool: :return: Instance of WorkerPool. """ logger.info("Building WorkerPool...") - worker: WorkerPool = ( - instantiate(cfg.worker) - if (is_target_type(cfg.worker, SingleMachineParallelExecutor) or is_target_type(cfg.worker, Sequential)) - else instantiate(cfg.worker, output_dir=cfg.output_dir) - ) + worker: WorkerPool = instantiate(cfg.worker) validate_type(worker, WorkerPool) - logger.info("Building WorkerPool...DONE!") return worker diff --git a/d123/script/config/common/scene_filter/log_scenes.yaml b/d123/script/config/common/scene_filter/log_scenes.yaml new file mode 100644 index 00000000..68114df9 --- /dev/null +++ b/d123/script/config/common/scene_filter/log_scenes.yaml @@ -0,0 +1,20 @@ +_target_: d123.dataset.scene.scene_filter.SceneFilter +_convert_: 'all' + +split_types: null +split_names: ["av2-sensor-mini_train"] +log_names: null + + +map_names: null +scene_tokens: null +timestamp_threshold_s: null +ego_displacement_minimum_m: null + +duration_s: null +history_s: 1.0 + +camera_types: null + +max_num_scenes: null +shuffle: false diff --git a/d123/script/config/common/worker/ray_distributed.yaml b/d123/script/config/common/worker/ray_distributed.yaml index c51de1d3..2c101f66 100644 --- a/d123/script/config/common/worker/ray_distributed.yaml +++ b/d123/script/config/common/worker/ray_distributed.yaml @@ -1,4 +1,4 @@ -_target_: nuplan.planning.utils.multithreading.worker_ray.RayDistributed +_target_: d123.common.multithreading.worker_ray.RayDistributed _convert_: 'all' master_node_ip: null # Set to a master node IP if you desire to connect to cluster remotely threads_per_node: null # Number of CPU threads to use per node, "null" means all threads available diff --git a/d123/script/config/common/worker/sequential.yaml b/d123/script/config/common/worker/sequential.yaml index 0b7955e9..ea004415 100644 --- a/d123/script/config/common/worker/sequential.yaml +++ b/d123/script/config/common/worker/sequential.yaml @@ -1,2 +1,2 @@ -_target_: nuplan.planning.utils.multithreading.worker_sequential.Sequential +_target_: d123.common.multithreading.worker_sequential.Sequential _convert_: 'all' diff --git a/d123/script/config/common/worker/single_machine_thread_pool.yaml b/d123/script/config/common/worker/single_machine_thread_pool.yaml index ba1b9657..ace763f6 100644 --- a/d123/script/config/common/worker/single_machine_thread_pool.yaml +++ b/d123/script/config/common/worker/single_machine_thread_pool.yaml @@ -1,4 +1,4 @@ -_target_: nuplan.planning.utils.multithreading.worker_parallel.SingleMachineParallelExecutor +_target_: d123.common.multithreading.worker_parallel.SingleMachineParallelExecutor _convert_: 'all' use_process_pool: True # If true, use ProcessPoolExecutor as the backend, otherwise uses ThreadPoolExecutor max_workers: 16 # Number of CPU workers (threads/processes) to use per node, "null" means all available diff --git a/d123/script/config/dataset_conversion/__init__.py b/d123/script/config/dataset_conversion/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/script/config/dataset_caching/default_dataset_caching.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml similarity index 79% rename from d123/script/config/dataset_caching/default_dataset_caching.yaml rename to d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 13a9e6da..0a4544da 100644 --- a/d123/script/config/dataset_caching/default_dataset_caching.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -16,6 +16,9 @@ defaults: - _self_ - datasets: # - nuplan_private_dataset - - carla_dataset + # - carla_dataset + # - wopd_dataset + - av2_sensor_dataset -force_feature_computation: True +force_log_conversion: True +force_map_conversion: True diff --git a/d123/script/config/datasets/av2_sensor_dataset.yaml b/d123/script/config/datasets/av2_sensor_dataset.yaml new file mode 100644 index 00000000..4f9a95ec --- /dev/null +++ b/d123/script/config/datasets/av2_sensor_dataset.yaml @@ -0,0 +1,16 @@ +av2_sensor_dataset: + _target_: d123.dataset.dataset_specific.av2.av2_data_converter.AV2SensorDataConverter + _convert_: 'all' + + splits: ["av2-sensor-mini_train"] + log_path: "/media/nvme1/argoverse" + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "binary" + lidar_store_option: null diff --git a/d123/script/config/datasets/carla_dataset.yaml b/d123/script/config/datasets/carla_dataset.yaml index b331aeee..31184f78 100644 --- a/d123/script/config/datasets/carla_dataset.yaml +++ b/d123/script/config/datasets/carla_dataset.yaml @@ -1,9 +1,16 @@ carla_dataset: - _target_: d123.dataset.dataset_specific.carla.carla_data_processor.CarlaDataProcessor + _target_: d123.dataset.dataset_specific.carla.carla_data_converter.CarlaDataConverter _convert_: 'all' - force_data_conversion: True - splits: ["carla"] log_path: "${oc.env:HOME}/carla_workspace/data" - output_path: ${d123_data_root} + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/d123/script/config/datasets/nuplan_dataset.yaml b/d123/script/config/datasets/nuplan_dataset.yaml index 7d949028..8c104287 100644 --- a/d123/script/config/datasets/nuplan_dataset.yaml +++ b/d123/script/config/datasets/nuplan_dataset.yaml @@ -1,9 +1,16 @@ nuplan_dataset: - _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_processor.NuplanDataProcessor + _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' - force_data_conversion: True - splits: ["nuplan_train", "nuplan_val", "nuplan_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder - output_path: ${d123_data_root} + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/d123/script/config/datasets/nuplan_mini_dataset.yaml b/d123/script/config/datasets/nuplan_mini_dataset.yaml index 9e4ca222..1fdb2b54 100644 --- a/d123/script/config/datasets/nuplan_mini_dataset.yaml +++ b/d123/script/config/datasets/nuplan_mini_dataset.yaml @@ -1,12 +1,17 @@ nuplan_mini_dataset: - _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_processor.NuplanDataProcessor + _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' splits: ["nuplan_mini_train", "nuplan_mini_val", "nuplan_mini_test"] - # splits: ["nuplan_mini_val"] - log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder - output_path: ${d123_data_root} - force_data_conversion: True + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/d123/script/config/datasets/nuplan_private_dataset.yaml b/d123/script/config/datasets/nuplan_private_dataset.yaml index a618e5a0..399dcb7e 100644 --- a/d123/script/config/datasets/nuplan_private_dataset.yaml +++ b/d123/script/config/datasets/nuplan_private_dataset.yaml @@ -1,10 +1,16 @@ nuplan_private_dataset: - _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_processor.NuplanDataProcessor + _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' - force_data_conversion: True - splits: ["nuplan_private_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder - sensor_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/sensor - output_path: ${d123_data_root} + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/d123/script/config/datasets/wopd_dataset.yaml b/d123/script/config/datasets/wopd_dataset.yaml new file mode 100644 index 00000000..1abb0381 --- /dev/null +++ b/d123/script/config/datasets/wopd_dataset.yaml @@ -0,0 +1,16 @@ +wopd_dataset: + _target_: d123.dataset.dataset_specific.wopd.wopd_data_converter.WOPDDataConverter + _convert_: 'all' + + splits: ["wopd_train"] + log_path: null # TODO: implement + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "binary" + lidar_store_option: "binary" diff --git a/d123/script/config/viser/__init__.py b/d123/script/config/viser/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/script/config/viser/default_viser.yaml b/d123/script/config/viser/default_viser.yaml new file mode 100644 index 00000000..c42164dc --- /dev/null +++ b/d123/script/config/viser/default_viser.yaml @@ -0,0 +1,16 @@ +hydra: + run: + dir: . + output_subdir: null + searchpath: # Only in these paths are discoverable + - pkg://d123.script.config + - pkg://d123.script.config.common + - pkg://d123.script.config.preprocessing + job: + chdir: False +# +defaults: + - default_common + - default_dataset_paths + +port_number: 8080 diff --git a/d123/script/run_dataset_caching.py b/d123/script/run_dataset_conversion.py similarity index 63% rename from d123/script/run_dataset_caching.py rename to d123/script/run_dataset_conversion.py index dbf40de0..62ad1e30 100644 --- a/d123/script/run_dataset_caching.py +++ b/d123/script/run_dataset_conversion.py @@ -5,13 +5,13 @@ from nuplan.planning.script.builders.logging_builder import build_logger from omegaconf import DictConfig -from d123.script.builders.data_processor_builder import RawDataProcessor, build_data_processor +from d123.script.builders.data_converter_builder import RawDataConverter, build_data_converter from d123.script.builders.worker_pool_builder import build_worker logger = logging.getLogger(__name__) -CONFIG_PATH = "config/dataset_caching" -CONFIG_NAME = "default_dataset_caching" +CONFIG_PATH = "config/dataset_conversion" +CONFIG_NAME = "default_dataset_conversion" @hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None) @@ -28,10 +28,17 @@ def main(cfg: DictConfig) -> None: # Precompute and cache all features logger.info("Starting Dataset Caching...") - data_processors: List[RawDataProcessor] = build_data_processor(cfg.datasets) + data_processors: List[RawDataConverter] = build_data_converter(cfg.datasets) for data_processor in data_processors: + logger.info(f"Processing dataset: {data_processor.__class__.__name__}") - data_processor.convert(worker=worker) + + data_processor.convert_maps(worker=worker) + logger.info(f"Finished maps: {data_processor.__class__.__name__}") + + data_processor.convert_logs(worker=worker) + logger.info(f"Finished logs: {data_processor.__class__.__name__}") + logger.info(f"Finished processing dataset: {data_processor.__class__.__name__}") diff --git a/d123/script/run_preprocessing.py b/d123/script/run_preprocessing.py index 85c9974e..4ed6cf97 100644 --- a/d123/script/run_preprocessing.py +++ b/d123/script/run_preprocessing.py @@ -6,13 +6,13 @@ import hydra import lightning as L -from nuplan.planning.utils.multithreading.worker_utils import worker_map from omegaconf import DictConfig +from d123.common.multithreading.worker_utils import worker_map from d123.dataset.scene.abstract_scene import AbstractScene from d123.script.builders.scene_builder_builder import build_scene_builder from d123.script.builders.scene_filter_builder import build_scene_filter -from d123.script.run_dataset_caching import build_worker +from d123.script.run_dataset_conversion import build_worker from d123.training.feature_builder.smart_feature_builder import SMARTFeatureBuilder logger = logging.getLogger(__name__) diff --git a/d123/script/run_simulation.py b/d123/script/run_simulation.py index 84b3ec94..393e4367 100644 --- a/d123/script/run_simulation.py +++ b/d123/script/run_simulation.py @@ -6,14 +6,14 @@ import hydra import lightning as L import pandas as pd -from nuplan.planning.utils.multithreading.worker_utils import worker_map from omegaconf import DictConfig from tqdm import tqdm +from d123.common.multithreading.worker_utils import worker_map from d123.dataset.scene.abstract_scene import AbstractScene from d123.script.builders.scene_builder_builder import build_scene_builder from d123.script.builders.scene_filter_builder import build_scene_filter -from d123.script.run_dataset_caching import build_worker +from d123.script.run_dataset_conversion import build_worker from d123.simulation.gym.demo_gym_env import DemoGymEnv from d123.simulation.metrics.sim_agents.sim_agents import get_sim_agents_metrics diff --git a/d123/script/run_viser.py b/d123/script/run_viser.py new file mode 100644 index 00000000..e682a96e --- /dev/null +++ b/d123/script/run_viser.py @@ -0,0 +1,29 @@ +import logging + +import hydra +from omegaconf import DictConfig + +from d123.common.visualization.viser.server import ViserVisualizationServer +from d123.script.builders.scene_builder_builder import build_scene_builder +from d123.script.builders.scene_filter_builder import build_scene_filter +from d123.script.run_dataset_conversion import build_worker + +logger = logging.getLogger(__name__) + +CONFIG_PATH = "config/viser" +CONFIG_NAME = "default_viser" + + +@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None) +def main(cfg: DictConfig) -> None: + + worker = build_worker(cfg) + scene_filter = build_scene_filter(cfg.scene_filter) + scene_builder = build_scene_builder(cfg.scene_builder) + scenes = scene_builder.get_scenes(scene_filter, worker=worker) + + ViserVisualizationServer(scenes=scenes) + + +if __name__ == "__main__": + main() diff --git a/d123/simulation/gym/environment/environment_wrapper.py b/d123/simulation/gym/environment/environment_wrapper.py index b01e76bf..fe3cd5b9 100644 --- a/d123/simulation/gym/environment/environment_wrapper.py +++ b/d123/simulation/gym/environment/environment_wrapper.py @@ -11,15 +11,15 @@ from d123.simulation.gym.environment.gym_observation.abstract_gym_observation import ( AbstractGymObservation, ) +from d123.simulation.gym.environment.output_converter.abstract_output_converter import ( + AbstractOutputConverter, +) from d123.simulation.gym.environment.reward_builder.abstract_reward_builder import AbstractRewardBuilder from d123.simulation.gym.environment.scenario_sampler.abstract_scenario_sampler import AbstractScenarioSampler from d123.simulation.gym.environment.simulation_builder.abstract_simulation_builder import ( AbstractSimulationBuilder, ) from d123.simulation.gym.environment.simulation_wrapper import SimulationWrapper -from d123.simulation.gym.environment.trajectory_builder.abstract_trajectory_builder import ( - AbstractTrajectoryBuilder, -) logger = logging.getLogger(__name__) @@ -35,7 +35,7 @@ def __init__( self, scenario_sampler: AbstractScenarioSampler, simulation_builder: AbstractSimulationBuilder, - trajectory_builder: AbstractTrajectoryBuilder, + output_converter: AbstractOutputConverter, observation_builder: AbstractGymObservation, reward_builder: AbstractRewardBuilder, terminate_on_failure: bool = False, @@ -52,7 +52,7 @@ def __init__( self._scenario_sampler = scenario_sampler self._simulation_builder = simulation_builder - self._trajectory_builder = trajectory_builder + self._trajectory_builder = output_converter self._observation_builder = observation_builder self._reward_builder = reward_builder @@ -68,7 +68,7 @@ def __init__( # Set for super class self.observation_space = observation_builder.get_observation_space() - self.action_space = trajectory_builder.get_action_space() + self.action_space = output_converter.get_action_space() def reset(self, seed: Optional[int] = None, options: Optional[dict] = None): """Inherited, see superclass.""" diff --git a/d123/simulation/gym/environment/helper/environment_cache.py b/d123/simulation/gym/environment/helper/environment_cache.py index 8a46df52..55fbea92 100644 --- a/d123/simulation/gym/environment/helper/environment_cache.py +++ b/d123/simulation/gym/environment/helper/environment_cache.py @@ -25,7 +25,7 @@ AbstractLaneGroup, AbstractStopLine, ) -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer from d123.simulation.gym.environment.helper.environment_area import AbstractEnvironmentArea from d123.simulation.planning.abstract_planner import PlannerInitialization, PlannerInput @@ -77,12 +77,12 @@ def __init__( def _load_cache(self) -> None: - query_map_layers = [MapSurfaceType.LANE_GROUP, MapSurfaceType.CARPARK] + query_map_layers = [MapLayer.LANE_GROUP, MapLayer.CARPARK] # FIXME: Add stop lines and crosswalks to the map layers if needed # if self.load_crosswalks: - # query_map_layers.append(MapSurfaceType.CROSSWALK) + # query_map_layers.append(MapLayer.CROSSWALK) # if self.load_stop_lines: - # query_map_layers.append(MapSurfaceType.STOP_LINE) + # query_map_layers.append(MapLayer.STOP_LINE) map_object_dict = self.map_api.query( geometry=self.environment_area.get_global_polygon(self.ego_state.center), @@ -91,7 +91,7 @@ def _load_cache(self) -> None: ) # 1. load (1.1) lane groups, (1.2) lanes, (1.3) intersections - for lane_group in map_object_dict[MapSurfaceType.LANE_GROUP]: + for lane_group in map_object_dict[MapLayer.LANE_GROUP]: lane_group: AbstractLaneGroup self.lane_groups[lane_group.id] = lane_group for lane in lane_group.lanes: @@ -101,18 +101,18 @@ def _load_cache(self) -> None: self.intersections[optional_intersection.id] = optional_intersection # 2. load car parks - for car_park in map_object_dict[MapSurfaceType.CARPARK]: + for car_park in map_object_dict[MapLayer.CARPARK]: car_park: AbstractCarpark self.car_parks[car_park.id] = car_park # FIXME: Add stop lines and crosswalks to the map layers if needed # if self.load_crosswalks: - # for crosswalk in map_object_dict[MapSurfaceType.CROSSWALK]: + # for crosswalk in map_object_dict[MapLayer.CROSSWALK]: # crosswalk: AbstractCarpark # self.crosswalks[crosswalk.id] = crosswalk # if self.load_stop_lines: - # for stop_line in map_object_dict[MapSurfaceType.STOP_LINE]: + # for stop_line in map_object_dict[MapLayer.STOP_LINE]: # stop_line: AbstractStopLine # self.stop_lines[stop_line.id] = stop_line diff --git a/d123/simulation/gym/environment/output_converter/__init__.py b/d123/simulation/gym/environment/output_converter/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/simulation/gym/environment/trajectory_builder/abstract_trajectory_builder.py b/d123/simulation/gym/environment/output_converter/abstract_output_converter.py similarity index 65% rename from d123/simulation/gym/environment/trajectory_builder/abstract_trajectory_builder.py rename to d123/simulation/gym/environment/output_converter/abstract_output_converter.py index fcadf1ed..94f3ef0c 100644 --- a/d123/simulation/gym/environment/trajectory_builder/abstract_trajectory_builder.py +++ b/d123/simulation/gym/environment/output_converter/abstract_output_converter.py @@ -4,11 +4,12 @@ import numpy as np import numpy.typing as npt from gymnasium import spaces -from nuplan.common.actor_state.ego_state import EgoState -from nuplan.planning.simulation.trajectory.abstract_trajectory import AbstractTrajectory +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2 +from d123.simulation.planning.planner_output.abstract_planner_output import AbstractPlannerOutput -class AbstractTrajectoryBuilder(ABC): + +class AbstractOutputConverter(ABC): """Abstract class for building trajectories (nuPlan interface) in a Gym simulation environment.""" @abstractmethod @@ -19,13 +20,13 @@ def get_action_space(self) -> spaces.Space: """ @abstractmethod - def build_trajectory( - self, action: npt.NDArray[np.float32], ego_state: EgoState, info: Dict[str, Any] - ) -> AbstractTrajectory: + def get_planner_output( + self, action: npt.NDArray[np.float32], ego_state: EgoStateSE2, info: Dict[str, Any] + ) -> AbstractPlannerOutput: """ - Builds a trajectory based on the action and the current ego state. + Builds a planner output based on the action and the current ego state. :param action: Action taken by the agent, typically a numpy array. :param ego_state: Current state of the ego vehicle. :param info: Arbitrary information dictionary, for passing information between modules. - :return: Trajectory object of nuPlan. + :return: Planner output object. """ diff --git a/d123/simulation/gym/environment/trajectory_builder/action_trajectory_builder.py b/d123/simulation/gym/environment/output_converter/action_output_converter.py similarity index 81% rename from d123/simulation/gym/environment/trajectory_builder/action_trajectory_builder.py rename to d123/simulation/gym/environment/output_converter/action_output_converter.py index 9e92e42a..ed8aff94 100644 --- a/d123/simulation/gym/environment/trajectory_builder/action_trajectory_builder.py +++ b/d123/simulation/gym/environment/output_converter/action_output_converter.py @@ -2,17 +2,15 @@ import numpy as np import numpy.typing as npt -from carl_nuplan.planning.simulation.trajectory.action_trajectory import ActionTrajectory from gymnasium import spaces -from nuplan.common.actor_state.ego_state import EgoState -from d123.simulation.gym.environment.trajectory_builder.abstract_trajectory_builder import ( - AbstractTrajectoryBuilder, -) +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2 +from d123.simulation.gym.environment.output_converter.abstract_output_converter import AbstractOutputConverter from d123.simulation.gym.policy.ppo.ppo_config import GlobalConfig +from d123.simulation.planning.planner_output.action_planner_output import ActionPlannerOutput -class ActionTrajectoryBuilder(AbstractTrajectoryBuilder): +class ActionOutputConverter(AbstractOutputConverter): """ Default action trajectory builder for training CaRL. TODO: Refactor this class @@ -75,36 +73,22 @@ def get_action_space(self) -> spaces.Space: dtype=np.float32, ) - def build_trajectory( - self, action: npt.NDArray[np.float32], ego_state: EgoState, info: Dict[str, Any] - ) -> ActionTrajectory: + def get_planner_output( + self, action: npt.NDArray[np.float32], ego_state: EgoStateSE2, info: Dict[str, Any] + ) -> ActionPlannerOutput: """Inherited, see superclass.""" assert len(action) == self._config.action_space_dim info["last_action"] = action action_acceleration_normed, action_steering_angle_normed = action - target_steering_rate = self._scale_steering( - action_steering_angle_normed, - ego_state.tire_steering_angle, - ) + target_steering_rate = self._scale_steering(action_steering_angle_normed, ego_state.tire_steering_angle) target_acceleration = self._scale_acceleration( - action_acceleration_normed, - ego_state.dynamic_car_state.rear_axle_acceleration_2d.x, - ) - clipped_steering_rate = self._clip_steering( - target_acceleration, - target_steering_rate, - ego_state, - ) - clipped_acceleration = self._clip_acceleration( - target_acceleration, - clipped_steering_rate, - ego_state, + action_acceleration_normed, ego_state.dynamic_state_se2.acceleration.x ) - - store = info["store"] if "store" in info.keys() else None - return ActionTrajectory(clipped_acceleration, clipped_steering_rate, ego_state, list(action), store) + clipped_steering_rate = self._clip_steering(target_acceleration, target_steering_rate, ego_state) + clipped_acceleration = self._clip_acceleration(target_acceleration, clipped_steering_rate, ego_state) + return ActionPlannerOutput(clipped_acceleration, clipped_steering_rate, ego_state) def _scale_steering(self, action_steering_angle_normed: float, current_steering_angle: float) -> float: """ @@ -137,7 +121,9 @@ def _scale_acceleration(self, action_acceleration_normed: float, current_acceler target_acceleration = (target_acceleration - current_acceleration) / factor + current_acceleration return target_acceleration - def _clip_acceleration(self, target_acceleration: float, target_steering_rate: float, ego_state: EgoState) -> float: + def _clip_acceleration( + self, target_acceleration: float, target_steering_rate: float, ego_state: EgoStateSE2 + ) -> float: """ Clips the acceleration based on the target acceleration, steering rate, and current ego state. :param target_acceleration: Acceleration as targeted by the agent. @@ -146,10 +132,10 @@ def _clip_acceleration(self, target_acceleration: float, target_steering_rate: f :return: Clipped acceleration based on the target acceleration and steering rate. """ - current_acceleration = ego_state.dynamic_car_state.rear_axle_acceleration_2d.x + current_acceleration = ego_state.dynamic_state_se2.acceleration.x if self._disable_reverse_driving: - speed = ego_state.dynamic_car_state.rear_axle_velocity_2d.x + speed = ego_state.dynamic_state_se2.velocity.x updated_speed = speed + target_acceleration * self._dt_control # * self._dt_control if updated_speed < 0: @@ -168,19 +154,15 @@ def _clip_acceleration(self, target_acceleration: float, target_steering_rate: f _max_acceleration = self._scale_max_acceleration if self._clip_angular_adjustment: - rear_axle_to_center_dist = ego_state.car_footprint.rear_axle_to_center_dist + rear_axle_to_center_dist = ego_state.vehicle_parameters.rear_axle_to_center_longitudinal - next_point_velocity_x = ( - ego_state.dynamic_car_state.rear_axle_velocity_2d.x + target_acceleration * self._dt_control - ) + next_point_velocity_x = ego_state.dynamic_state_se2.velocity.x + target_acceleration * self._dt_control next_point_tire_steering_angle = ego_state.tire_steering_angle + target_steering_rate * self._dt_control next_point_angular_velocity = ( - next_point_velocity_x - * np.tan(next_point_tire_steering_angle) - / ego_state.car_footprint.vehicle_parameters.wheel_base + next_point_velocity_x * np.tan(next_point_tire_steering_angle) / ego_state.vehicle_parameters.wheel_base ) next_point_angular_acceleration = ( - next_point_angular_velocity - ego_state.dynamic_car_state.angular_velocity + next_point_angular_velocity - ego_state.dynamic_state_se2.angular_velocity ) / self._dt_control centripetal_acceleration_term = rear_axle_to_center_dist * (next_point_angular_velocity) ** 2 @@ -190,7 +172,7 @@ def _clip_acceleration(self, target_acceleration: float, target_steering_rate: f target_acceleration = np.clip(target_acceleration, -self._scale_max_deceleration, _max_acceleration) return target_acceleration - def _clip_steering(self, target_acceleration: float, target_steering_rate: float, ego_state: EgoState) -> float: + def _clip_steering(self, target_acceleration: float, target_steering_rate: float, ego_state: EgoStateSE2) -> float: """ Clips the steering rate based on the target acceleration and current ego state. :param target_acceleration: Acceleration as targeted by the agent. @@ -203,12 +185,10 @@ def _clip_steering(self, target_acceleration: float, target_steering_rate: float target_steering_angle = current_steering_angle + target_steering_rate * self._dt_control if self._clip_max_abs_yaw_accel is not None: - wheel_base = ego_state.car_footprint.vehicle_parameters.wheel_base - target_velocity = ( - ego_state.dynamic_car_state.rear_axle_velocity_2d.x + target_acceleration * self._dt_control - ) + wheel_base = ego_state.vehicle_parameters.wheel_base + target_velocity = ego_state.dynamic_state_se2.acceleration.x + target_acceleration * self._dt_control - current_angular_velocity = ego_state.dynamic_car_state.angular_velocity + current_angular_velocity = ego_state.dynamic_state_se2.angular_velocity max_abs_yaw_velocity = self._clip_max_abs_yaw_accel * self._dt_control min_angular_velocity = current_angular_velocity - max_abs_yaw_velocity diff --git a/d123/simulation/metrics/sim_agents/map_based.py b/d123/simulation/metrics/sim_agents/map_based.py index 513aca42..dbd9a16d 100644 --- a/d123/simulation/metrics/sim_agents/map_based.py +++ b/d123/simulation/metrics/sim_agents/map_based.py @@ -10,7 +10,7 @@ from d123.common.geometry.utils import normalize_angle from d123.dataset.maps.abstract_map import AbstractMap from d123.dataset.maps.abstract_map_objects import AbstractLane -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer MAX_LANE_CENTER_DISTANCE: Final[float] = 10.0 @@ -32,10 +32,10 @@ def _get_offroad_feature( output = map_api.query_object_ids( agent_shapely_corners, layers=[ - MapSurfaceType.INTERSECTION, - MapSurfaceType.LANE_GROUP, - MapSurfaceType.CARPARK, - MapSurfaceType.GENERIC_DRIVABLE, + MapLayer.INTERSECTION, + MapLayer.LANE_GROUP, + MapLayer.CARPARK, + MapLayer.GENERIC_DRIVABLE, ], predicate="within", ) @@ -61,7 +61,7 @@ def _get_road_center_distance_feature( def get_lane_by_id(lane_id: str, lane_dict: Dict[str, AbstractLane]) -> AbstractLane: if lane_id not in lane_dict.keys(): - lane_dict[lane_id] = map_api.get_map_object(lane_id, MapSurfaceType.LANE) + lane_dict[lane_id] = map_api.get_map_object(lane_id, MapLayer.LANE) return lane_dict[lane_id] assert agents_array.shape[-1] == len(BoundingBoxSE2Index) @@ -74,12 +74,12 @@ def get_lane_by_id(lane_id: str, lane_dict: Dict[str, AbstractLane]) -> Abstract nearest_query_output = map_api.query_nearest( agent_shapely_centers, - layers=[MapSurfaceType.LANE], + layers=[MapLayer.LANE], max_distance=MAX_LANE_CENTER_DISTANCE, return_all=True, return_distance=False, exclusive=False, - )[MapSurfaceType.LANE] + )[MapLayer.LANE] for object_idx in range(n_objects): for iteration in range(n_iterations): diff --git a/d123/simulation/simulation_2d_setup.py b/d123/simulation/simulation_2d_setup.py index b38df183..737d3e6c 100644 --- a/d123/simulation/simulation_2d_setup.py +++ b/d123/simulation/simulation_2d_setup.py @@ -2,17 +2,7 @@ from d123.simulation.controller.abstract_controller import AbstractEgoController from d123.simulation.observation.abstract_observation import AbstractObservation -from d123.simulation.time_controller.abstract_time_controller import ( - AbstractTimeController, -) - -# from nuplan.planning.scenario_builder.abstract_scenario import AbstractScenario -# from nuplan.planning.simulation.controller.abstract_controller import AbstractEgoController -# from nuplan.planning.simulation.observation.abstract_observation import AbstractObservation -# from nuplan.planning.simulation.planner.abstract_planner import AbstractPlanner -# from nuplan.planning.simulation.simulation_time_controller.abstract_simulation_time_controller import ( -# AbstractSimulationTimeController, -# ) +from d123.simulation.time_controller.abstract_time_controller import AbstractTimeController @dataclass diff --git a/d123/training/feature_builder/smart_feature_builder.py b/d123/training/feature_builder/smart_feature_builder.py index 8a2af0ba..829f0498 100644 --- a/d123/training/feature_builder/smart_feature_builder.py +++ b/d123/training/feature_builder/smart_feature_builder.py @@ -12,7 +12,7 @@ from d123.common.geometry.line.polylines import PolylineSE2 from d123.common.geometry.transform.se2_array import convert_absolute_to_relative_se2_array from d123.common.visualization.color.default import TrafficLightStatus -from d123.dataset.maps.abstract_map import MapSurfaceType +from d123.dataset.maps.abstract_map import MapLayer from d123.dataset.maps.abstract_map_objects import ( AbstractCarpark, AbstractCrosswalk, @@ -74,10 +74,10 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. map_objects = map_api.query( map_bounding_box.shapely_polygon, layers=[ - MapSurfaceType.LANE_GROUP, - MapSurfaceType.CROSSWALK, - MapSurfaceType.CARPARK, - MapSurfaceType.GENERIC_DRIVABLE, + MapLayer.LANE_GROUP, + MapLayer.CROSSWALK, + MapLayer.CARPARK, + MapLayer.GENERIC_DRIVABLE, ], predicate="intersects", ) @@ -91,7 +91,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types: List[int] = [] # 1. Add lane - for lane_group in map_objects[MapSurfaceType.LANE_GROUP]: + for lane_group in map_objects[MapLayer.LANE_GROUP]: lane_group: AbstractLaneGroup is_intersection = lane_group.intersection is not None @@ -129,7 +129,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types.extend([int(lane_traffic_light.status)] * len(lane_traj_se2)) # 2. Crosswalks - for crosswalk in map_objects[MapSurfaceType.CROSSWALK]: + for crosswalk in map_objects[MapLayer.CROSSWALK]: crosswalk: AbstractCrosswalk crosswalk_traj_se2 = _split_segments( crosswalk.outline_3d.polyline_se2, @@ -143,7 +143,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types.extend([int(TrafficLightStatus.OFF)] * len(crosswalk_traj_se2)) # 3. Parking - for carpark in map_objects[MapSurfaceType.CARPARK]: + for carpark in map_objects[MapLayer.CARPARK]: carpark: AbstractCarpark carpark_traj_se2 = _split_segments( carpark.outline_3d.polyline_se2, @@ -157,7 +157,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types.extend([int(TrafficLightStatus.OFF)] * len(carpark_traj_se2)) # 4. Generic drivable - for generic_drivable in map_objects[MapSurfaceType.GENERIC_DRIVABLE]: + for generic_drivable in map_objects[MapLayer.GENERIC_DRIVABLE]: generic_drivable: AbstractGenericDrivable drivable_traj_se2 = _split_segments( generic_drivable.outline_3d.polyline_se2, diff --git a/d123/training/models/sim_agent/smart/layers/attention_layer.py b/d123/training/models/sim_agent/smart/layers/attention_layer.py index e5190552..9452662e 100644 --- a/d123/training/models/sim_agent/smart/layers/attention_layer.py +++ b/d123/training/models/sim_agent/smart/layers/attention_layer.py @@ -17,7 +17,7 @@ def __init__( dropout: float, bipartite: bool, has_pos_emb: bool, - **kwargs + **kwargs, ) -> None: super(AttentionLayer, self).__init__(aggr="add", node_dim=0, **kwargs) self.num_heads = num_heads diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..d4bb2cbb --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/__init__.py b/docs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000..9d7e8c5f --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,68 @@ +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +project = "d123" +copyright = "2025, 123D Contributors" +author = "123D Contributors" +release = "v0.0.6" + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + "sphinx.ext.duration", + "sphinx.ext.doctest", + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.intersphinx", + "myst_parser", +] + +intersphinx_mapping = { + "rtd": ("https://docs.readthedocs.io/en/stable/", None), + "python": ("https://docs.python.org/3/", None), + "sphinx": ("https://www.sphinx-doc.org/en/master/", None), +} +intersphinx_disabled_domains = ["std"] + +templates_path = ["_templates"] + +# -- Options for EPUB output +epub_show_urls = "footnote" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path. +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "sphinx_rtd_theme" + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ["_static"] + +html_theme_options = { + "analytics_anonymize_ip": False, + "logo_only": False, + "display_version": True, + "prev_next_buttons_location": "bottom", + "style_external_links": False, + "vcs_pageview_mode": "", + # Toc options + "collapse_navigation": True, + "sticky_navigation": True, + "navigation_depth": 3, + "includehidden": True, + "titles_only": False, +} diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 00000000..e447d876 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,118 @@ + +# Contributing to 123D + +Thank you for your interest in contributing to 123D! This guide will help you get started with the development process. + +## Getting Started + +### 1. Clone the Repository + +```sh +git clone git@github.com:DanielDauner/d123.git +cd d123 +``` + +### 2. Install the pip-package + +```sh +conda env create -f environment.yml --name d123_dev # Optional +conda activate d123_dev +pip install -e .[dev] +pre-commit install +``` + +.. note:: + We might remove the conda environment in the future, but leave the file in the repo during development. + + +### 3. Managing dependencies + +One principal of 123D is to keep *minimal dependencies*. However, various datasets require problematic (or even outdated) dependencies in order to load or preprocess the dataset. In this case, you can add optional dependencies to the `pyproject.toml` install file. You can follow examples of Waymo or nuPlan. These optional dependencies can be install with + +```sh +pip install -e .[dev,waymo,nuplan] +``` +where you can combined the different optional dependencies. + +The optional dependencies should only be required for data pre-processing. If a dataset allows to load sensor data dynamically from the original dataset, please encapsule the import accordingly, e.g. + +```python +import numpy as np +import numpy.typing as npt + +def load_camera_from_outdated_dataset(file_path: str) -> npt.NDArray[np.uint8]: + try: + from outdated_dataset import load_camera_image + except ImportError: + raise ImportError( + "Optional dependency 'outdated_dataset' is required to load camera images from this dataset. " + "Please install it using: pip install .[outdated_dataset]" + ) + return load_camera_image(file_path) +``` + + +## Code Style and Formatting + +We maintain consistent code quality using the following tools: +- **[Black](https://black.readthedocs.io/)** - Code formatter +- **[isort](https://pycqa.github.io/isort/)** - Import statement formatter +- **[flake8](https://flake8.pycqa.github.io/)** - Style guide enforcement +- **[pytest](https://docs.pytest.org/)** - Testing framework for unit and integration tests +- **[pre-commit](https://pre-commit.com/)** - Framework for managing and running Git hooks to automate code quality checks + + +.. note:: + If you're using VSCode, it is recommended to install the [Black](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter), [isort](https://marketplace.visualstudio.com/items?itemName=ms-python.isort), and [Flake8](https://marketplace.visualstudio.com/items?itemName=ms-python.flake8) plugins. + + + +### Editor Setup + +**VS Code Users:** +If you're using VSCode, it is recommended to install the following plguins: +- [Black](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) - see above. +- [isort](https://marketplace.visualstudio.com/items?itemName=ms-python.isort) - see above. +- [Flake8](https://marketplace.visualstudio.com/items?itemName=ms-python.flake8) - see above. +- [autodocstring](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) - Creating docstrings (please set `"autoDocstring.docstringFormat": "sphinx-notypes"`). +- [Code Spell Checker](https://marketplace.visualstudio.com/items?itemName=streetsidesoftware.code-spell-checker) - A basic spell checker. + + +**Other Editors:** +Similar plugins are available for most popular editors including PyCharm, Vim, Emacs, and Sublime Text. + + +## Documentation Requirements + +### Docstrings +- **Development:** Docstrings are encouraged but not strictly required during active development +- **Release:** All public functions, classes, and modules must have comprehensive docstrings before release +- **Format:** Use [Sphinx-style docstrings](https://sphinx-rtd-tutorial.readthedocs.io/en/latest/docstrings.html) + +**VS Code Users:** The [autoDocstring extension](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) can help generate properly formatted docstrings. + +### Type Hints +- **Required:** All function parameters and return values must include type hints +- **Style:** Follow [PEP 484](https://peps.python.org/pep-0484/) conventions + +### Sphinx documentation + +All datasets should be included in the `/docs/datasets.md` documentation. Please follow the documentation format of other datasets. + +You can install relevant dependencies for editing the public documentation via: +```sh +pip install -e .[docs] +``` + +It is recommended to uses [sphinx-autobuild](https://github.com/sphinx-doc/sphinx-autobuild) (installed above) to edit and view the documentation. You can run: +```sh +sphinx-autobuild docs docs/_build/html +``` + +## Adding new datasets +TODO + + +## Questions? + +If you have any questions about contributing, please open an issue or reach out to the maintainers. diff --git a/docs/datasets/av2.rst b/docs/datasets/av2.rst new file mode 100644 index 00000000..ab01edfb --- /dev/null +++ b/docs/datasets/av2.rst @@ -0,0 +1,90 @@ +Argoverse 2 +----------- + +.. sidebar:: Dataset Name + + .. image:: https://www.argoverse.org/assets/images/reference_images/av2_vehicle.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/carla.rst b/docs/datasets/carla.rst new file mode 100644 index 00000000..3ca7ff09 --- /dev/null +++ b/docs/datasets/carla.rst @@ -0,0 +1,90 @@ +CARLA +----- + +.. sidebar:: Dataset Name + + .. image:: https://carla.org/img/services/getty_center_400_400.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/index.rst b/docs/datasets/index.rst new file mode 100644 index 00000000..ead97169 --- /dev/null +++ b/docs/datasets/index.rst @@ -0,0 +1,17 @@ +Datasets +======== + +Brief overview of the datasets section... + +This section provides comprehensive documentation for various autonomous driving and computer vision datasets. Each dataset entry includes installation instructions, available data types, known issues, and proper citation formats. + +.. toctree:: + :maxdepth: 0 + + av2 + nuplan + nuscenes + carla + kitti-360 + wopd + template diff --git a/docs/datasets/kitti-360.rst b/docs/datasets/kitti-360.rst new file mode 100644 index 00000000..76100d27 --- /dev/null +++ b/docs/datasets/kitti-360.rst @@ -0,0 +1,87 @@ +KiTTI-360 +--------- + +.. sidebar:: Dataset Name + + .. image:: https://www.cvlibs.net/datasets/kitti-360/images/example/3d/semantic/02400.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use KiTTI-360 in your research, please cite: + +.. code-block:: bibtex + + @article{Liao2022PAMI, + title = {{KITTI}-360: A Novel Dataset and Benchmarks for Urban Scene Understanding in 2D and 3D}, + author = {Yiyi Liao and Jun Xie and Andreas Geiger}, + journal = {Pattern Analysis and Machine Intelligence (PAMI)}, + year = {2022}, + } diff --git a/docs/datasets/nuplan.rst b/docs/datasets/nuplan.rst new file mode 100644 index 00000000..065c26a5 --- /dev/null +++ b/docs/datasets/nuplan.rst @@ -0,0 +1,93 @@ +nuPlan +----- + +.. sidebar:: nuPlan + + .. image:: https://www.nuplan.org/static/media/nuPlan_final.3fde7586.png + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Towards learning-based planning:The nuPlan benchmark for real-world autonomous driving `_ + | **Download:** `www.nuscenes.org/nuplan `_ + | **Code:** `www.github.com/motional/nuplan-devkit `_ + | **Documentation:** `nuPlan Documentation `_ + | **License:** `CC BY-NC-SA 4.0 `_, `nuPlan Dataset License `_ + | **Duration:** 1282 hours (120 hours of sensor data) + | **Supported Versions:** [TODO] + | **Redistribution:** [TODO] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - ✓ + - [Description of ego vehicle data] + * - Map + - ✓ + - [Description of map data] + * - Bounding Boxes + - X + - [Description of bounding boxes data] + * - Traffic Lights + - X + - [Description of traffic lights data] + * - Cameras + - X + - [Description of cameras data] + * - LiDARs + - X + - [Description of LiDARs data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{Karnchanachari2024ICRA, + title={Towards learning-based planning: The nuplan benchmark for real-world autonomous driving}, + author={Karnchanachari, Napat and Geromichalos, Dimitris and Tan, Kok Seang and Li, Nanxiang and Eriksen, Christopher and Yaghoubi, Shakiba and Mehdipour, Noushin and Bernasconi, Gianmarco and Fong, Whye Kit and Guo, Yiluan and others}, + booktitle={2024 IEEE International Conference on Robotics and Automation (ICRA)}, + year={2024}, + } + @article{Caesar2021CVPRW, + title={nuplan: A closed-loop ml-based planning benchmark for autonomous vehicles}, + author={Caesar, Holger and Kabzan, Juraj and Tan, Kok Seang and Fong, Whye Kit and Wolff, Eric and Lang, Alex and Fletcher, Luke and Beijbom, Oscar and Omari, Sammy}, + booktitle={Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR) Workshops}, + year={2021} + } diff --git a/docs/datasets/nuscenes.rst b/docs/datasets/nuscenes.rst new file mode 100644 index 00000000..1f4e1621 --- /dev/null +++ b/docs/datasets/nuscenes.rst @@ -0,0 +1,90 @@ +nuScenes +-------- + +.. sidebar:: nuScenes + + .. image:: https://ar5iv.labs.arxiv.org/html/1903.11027/assets/figures/sensors.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/template.rst b/docs/datasets/template.rst new file mode 100644 index 00000000..d38723ed --- /dev/null +++ b/docs/datasets/template.rst @@ -0,0 +1,90 @@ +Template +-------- + +.. sidebar:: Dataset Name + + .. image:: https://www.nuplan.org/static/media/nuPlan_final.3fde7586.png + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/wopd.rst b/docs/datasets/wopd.rst new file mode 100644 index 00000000..8dc276a8 --- /dev/null +++ b/docs/datasets/wopd.rst @@ -0,0 +1,88 @@ +Waymo Open Perception Dataset (WOPD) +------------------------------------ + +.. sidebar:: WOPD + + .. image:: https://images.ctfassets.net/e6t5diu0txbw/4LpraC18sHNvS87OFnEGKB/63de105d4ce623d91cfdbc23f77d6a37/Open_Dataset_Download_Hero.jpg?fm=webp&q=90 + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @inproceedings{Sun2020CVPR, + title={Scalability in perception for autonomous driving: Waymo open dataset}, + author={Sun, Pei and Kretzschmar, Henrik and Dotiwalla, Xerxes and Chouard, Aurelien and Patnaik, Vijaysai and Tsui, Paul and Guo, James and Zhou, Yin and Chai, Yuning and Caine, Benjamin and others}, + booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, + pages={2446--2454}, + year={2020} + } diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 00000000..b169f012 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,22 @@ +.. 123d documentation master file, created by + sphinx-quickstart on Wed Aug 13 16:57:48 2025. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +123D Documentation +================== + +Add your content using ``reStructuredText`` syntax. See the +`reStructuredText `_ +documentation for details. + + +.. toctree:: + :maxdepth: 1 + :caption: Contents: + + installation + datasets/index + schema + visualization + contributing diff --git a/docs/installation.md b/docs/installation.md index 6fe76e59..95d09a04 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -1,5 +1,5 @@ -# Install Code +# Installation Note, the following installation assumes the following folder structure: ``` @@ -19,13 +19,13 @@ Note, the following installation assumes the following folder structure: │ ├── ... │ └── 2021.10.06.07.26.10_veh-52_00006_00398.arrow ├── nuplan_mini_train - │ └── ... + │ └── ... └── nuplan_mini_test └── ... ``` -First you need to create a new conda environment and install `d123` as editable pip package. +First you need to create a new conda environment and install `d123` as editable pip package. ```bash conda env create -f environment.yml conda activate d123 @@ -45,4 +45,4 @@ export CARLA_SIMULATOR_ROOT="$HOME/carla_workspace/carla_garage/carla" # nuPlan export NUPLAN_DATA_ROOT="/path/to/nuplan/dataset" export NUPLAN_MAPS_ROOT="/path/to/nuplan/dataset/maps" -``` \ No newline at end of file +``` diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 00000000..32bb2452 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=. +set BUILDDIR=_build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/schema.md b/docs/schema.md new file mode 100644 index 00000000..b3d44a6e --- /dev/null +++ b/docs/schema.md @@ -0,0 +1,15 @@ + +# Schema + +TODO + diff --git a/docs/visualization.md b/docs/visualization.md new file mode 100644 index 00000000..9fe2e3cf --- /dev/null +++ b/docs/visualization.md @@ -0,0 +1,11 @@ + +# Visualization + + +## Matplotlib + + +## Viser + + +## Bokeh diff --git a/environment.yml b/environment.yml index eeb40a17..53e458bc 100644 --- a/environment.yml +++ b/environment.yml @@ -5,5 +5,3 @@ dependencies: - python=3.12 - pip=25.1.1 - nb_conda_kernels - - pip: - - -r requirements.txt diff --git a/notebooks/av2/delete_me.ipynb b/notebooks/av2/delete_me.ipynb new file mode 100644 index 00000000..e28ef5ef --- /dev/null +++ b/notebooks/av2/delete_me.ipynb @@ -0,0 +1,480 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import io\n", + "\n", + "from PIL import Image\n", + "\n", + "import pandas as pd" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"test\"\n", + "\n", + "\n", + "split = \"train\"\n", + "# split = \"val\"\n", + "\n", + "split_folder = Path(f\"/media/nvme1/argoverse/sensor_mini/{split}\")\n", + "log_folder = sorted(list(split_folder.iterdir()))[0]\n", + "\n", + "log_folder.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "def _ls(path: Path):\n", + " \"\"\"List all files in the given path.\"\"\"\n", + " return [f.name for f in path.iterdir()]\n", + "\n", + "\n", + "def get_arrow_from_file(file_path: Path):\n", + " if file_path.suffix == \".parquet\":\n", + " import pyarrow.parquet as pq\n", + " return pq.read_table(file_path)\n", + " elif file_path.suffix == \".feather\":\n", + " import pyarrow.feather as feather\n", + " return feather.read_feather(file_path)\n", + " else:\n", + " raise ValueError(f\"Unsupported file type: {file_path.suffix}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "\n", + "# 1. Calibration \n", + "calibration_folder = log_folder / \"calibration\"\n", + "\n", + "# 1.1 -> ego to sensor transformation\n", + "egovehicle_SE3_sensor_file = log_folder / \"calibration\" / \"egovehicle_SE3_sensor.feather\"\n", + "egovehicle_se3_sensor_table = get_arrow_from_file(egovehicle_SE3_sensor_file)\n", + "\n", + "egovehicle_se3_sensor_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "# 1.2 -> intrinsic parameters\n", + "intrinsics_file = log_folder / \"calibration\" / \"intrinsics.feather\"\n", + "intrinsics_table = get_arrow_from_file(intrinsics_file)\n", + "\n", + "\n", + "intrinsics_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "# 2. Ego Vehicle\n", + "city_SE3_egovehicle_file = log_folder / \"city_SE3_egovehicle.feather\"\n", + "city_se3_egovehicle_table = get_arrow_from_file(city_SE3_egovehicle_file)\n", + "\n", + "city_se3_egovehicle_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# # 3. Map\n", + "# # map_folder = log_folder / \"map\"\n", + "# # print(_ls(map_folder))\n", + "\n", + "# # # 4. sensors\n", + "# # print(_ls(log_folder))\n", + "\n", + "# # from d123.dataset.dataset_specific.av2.av2_data_converter import AV2SensorDataConverter\n", + "# from d123.dataset.dataset_specific.av2.av2_data_converter import AV2SensorDataConverter\n", + "\n", + "# # AV2SensorDataConverter([])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# 5. Annotations\n", + "annotations_file = log_folder / \"annotations.feather\"\n", + "annotations_table = get_arrow_from_file(annotations_file)\n", + "\n", + "# print(_ls(annotations_folder))\n", + "annotations_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "camera_name = \"ring_side_left\"\n", + "\n", + "camera_folder = log_folder / \"sensors\" / \"cameras\" / camera_name\n", + "camera_files = sorted(list(camera_folder.iterdir()))\n", + "\n", + "\n", + "def jpg_to_array(file_path):\n", + " image = np.array(Image.open(io.BytesIO(file_path.read_bytes())))\n", + " return image\n", + "\n", + "plt.imshow(jpg_to_array(camera_files[1]))\n", + "print(len(camera_files))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "lidar_folder = log_folder / \"sensors\" / \"lidar\" \n", + "lidar_files = sorted(list(lidar_folder.iterdir()))\n", + "\n", + "\n", + "get_arrow_from_file(lidar_files[0])\n", + "\n", + "print(lidar_files)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "# Testing sensor syn dataframes\n", + "\n", + "from typing import Optional\n", + "from d123.dataset.dataset_specific.av2.av2_constants import AV2_CAMERA_TYPE_MAPPING\n", + "from d123.dataset.dataset_specific.av2.av2_helper import build_sensor_dataframe, build_synchronization_dataframe\n", + "\n", + "\n", + "sensor_df = build_sensor_dataframe(log_folder)\n", + "synchronization_df = build_synchronization_dataframe(sensor_df)\n", + "dataset_dir = split_folder.parent\n", + "\n", + "\n", + "def find_closest_target_fpath(\n", + " split: str,\n", + " log_id: str,\n", + " src_sensor_name: str,\n", + " src_timestamp_ns: int,\n", + " target_sensor_name: str,\n", + ") -> Optional[Path]:\n", + " \"\"\"Find the file path to the target sensor from a source sensor.\"\"\"\n", + " if synchronization_df is None:\n", + " raise RuntimeError(\"Requested synchronized data, but the synchronization database has not been created.\")\n", + "\n", + " src_timedelta_ns = pd.Timedelta(src_timestamp_ns)\n", + " src_to_target_records = synchronization_df.loc[(split, log_id, src_sensor_name)].set_index(src_sensor_name)\n", + " index = src_to_target_records.index\n", + " if src_timedelta_ns not in index:\n", + " # This timestamp does not correspond to any lidar sweep.\n", + " return None\n", + "\n", + " # Grab the synchronization record.\n", + " target_timestamp_ns = src_to_target_records.loc[src_timedelta_ns, target_sensor_name]\n", + " if pd.isna(target_timestamp_ns):\n", + " # No match was found within tolerance.\n", + " return None\n", + "\n", + " sensor_dir = dataset_dir / split / log_id / \"sensors\"\n", + " valid_cameras = list(AV2_CAMERA_TYPE_MAPPING.keys())\n", + " timestamp_ns_str = str(target_timestamp_ns.asm8.item())\n", + " if target_sensor_name in valid_cameras:\n", + " target_path = sensor_dir / \"cameras\" / target_sensor_name / f\"{timestamp_ns_str}.jpg\"\n", + " else:\n", + " target_path = sensor_dir / target_sensor_name / f\"{timestamp_ns_str}.feather\"\n", + " return target_path" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "# split=\"train\"\n", + "# log_id=\"00a6ffc1-6ce9-3bc3-a060-6006e9893a1a\"\n", + "# src_sensor_name=\"ring_front_center\"\n", + "# src_timestamp_ns=315967376959702000\n", + "# target_sensor_name=\"lidar\"\n", + "\n", + "# src_to_target_records = synchronization_df.loc[(\"train\", \"\", src_sensor_name)]\n", + "# # synchronization_df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "lidar_timestamp_ns_list = [int(path.stem) for path in lidar_files]\n", + "\n", + "\n", + "\n", + "\n", + "for lidar_timestamp_ns in lidar_timestamp_ns_list:\n", + "\n", + " fpath = find_closest_target_fpath(\n", + " split=\"train\",\n", + " log_id=log_folder.name,\n", + " src_sensor_name=\"lidar\",\n", + " src_timestamp_ns=lidar_timestamp_ns,\n", + " target_sensor_name=\"ring_front_center\",\n", + " )\n", + " if fpath is None:\n", + " continue\n", + " # print(fpath)\n", + "\n", + " egovehicle_se3_sensor_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "for _, row in egovehicle_se3_sensor_table.iterrows():\n", + " row = row.to_dict()\n", + " print(row)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "import cv2\n", + "from pyquaternion import Quaternion\n", + "\n", + "lidar_timestamps = [int(f.stem) for f in lidar_files]\n", + "camera_timestamps = [int(f.stem) for f in camera_files]\n", + "\n", + "\n", + "def get_slice_with_timestamp_ns(dataframe: pd.DataFrame, timestamp_ns: int):\n", + " \"\"\"Get the index of the closest timestamp to the target timestamp.\"\"\"\n", + " return dataframe[dataframe[\"timestamp_ns\"] == timestamp_ns]\n", + "\n", + "\n", + "def find_nearest_timestamp(target_ns, timestamp_list):\n", + " timestamp_array = np.array(timestamp_list, dtype=np.int64)\n", + " idx = np.argmin(np.abs(timestamp_array - np.int64(target_ns)))\n", + " return int(timestamp_array[idx])\n", + "\n", + "# for lidar_timestamp in lidar_timestamps:\n", + "# slice = get_slice_with_timestamp_ns(annotations_table, lidar_timestamp)\n", + "# assert len(slice) >= 1\n", + "\n", + "\n", + "\n", + "# ego_pose = city_SE3_egovehicle_table[city_SE3_egovehicle_table[\"timestamp_ns\"] == lidar_timestamps[10]]\n", + "# ego_pose_dict = ego_pose.iloc[0].to_dict()\n", + "\n", + "\n", + "annotations_slice = get_slice_with_timestamp_ns(annotations_table, lidar_timestamps[10])\n", + "for _, row in annotations_slice.iterrows():\n", + "# qw = row[\"qw\"]\n", + "# qx = row[\"qx\"]\n", + "# qy = row[\"qy\"]\n", + "# qz = row[\"qz\"]\n", + "# tx_m = row[\"tx_m\"]\n", + "# ty_m = row[\"ty_m\"]\n", + "# tz_m = row[\"tz_m\"]\n", + " print(row.to_dict())\n", + "\n", + "annotations_slice\n", + "\n", + "# qw\tqx\tqy\tqz\ttx_m\tty_m\ttz_m\n", + "# # def jpg_to_array(file_path):\n", + "\n", + "# camera_frames = []\n", + "# for lidar_timestamp in lidar_timestamps:\n", + "# camera_stamp_at_lidar = find_nearest_timestamp(lidar_timestamp, camera_timestamps)\n", + "# image = jpg_to_array(camera_folder / f\"{camera_stamp_at_lidar}.jpg\")\n", + "# camera_frames.append(image)\n", + " \n", + "# print(len(camera_frames))\n", + "# height, width, _ = camera_frames[0].shape\n", + "# video_path = \"camera_frames_video.mp4\"\n", + "# fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n", + "# out = cv2.VideoWriter(video_path, fourcc, 10, (width, height))\n", + "\n", + "# for frame in camera_frames:\n", + "# out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))\n", + "\n", + "# out.release()\n", + "# print(f\"Saved video to {video_path}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "from pyquaternion import Quaternion\n", + "from d123.common.datatypes.detection.detection_types import DetectionType\n", + "from d123.common.geometry.base import StateSE2\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE2\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.color.default import BOX_DETECTION_CONFIG\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_polygon_to_ax\n", + "\n", + "\n", + "\n", + "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n", + "\n", + "for cuboid in sensor_data.annotations:\n", + " yaw, pitch, roll = Quaternion(matrix=cuboid.dst_SE3_object.rotation).yaw_pitch_roll\n", + " center = StateSE2(cuboid.dst_SE3_object.translation[0], cuboid.dst_SE3_object.translation[1], yaw)\n", + " bounding_box = BoundingBoxSE2(center, cuboid.length_m, cuboid.width_m)\n", + " add_shapely_polygon_to_ax(ax, bounding_box.shapely_polygon, BOX_DETECTION_CONFIG[DetectionType.VEHICLE])\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "\n", + "radius = 200\n", + "ax.set_xlim(-radius, radius)\n", + "ax.set_ylim(-radius, radius)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "bounding_box.shapely_polygon\n", + "\n", + "bounding_box.corners_array" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "sensor_cache = \"/home/daniel/.cache/av2/sensor_cache.feather\"\n", + "get_arrow_from_file(Path(sensor_cache))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "synchronization_cache = \"/home/daniel/.cache/av2/synchronization_cache.feather\"\n", + "synchronization_cache = get_arrow_from_file(Path(synchronization_cache))\n", + "\n", + "synchronization_cache[\"sensor_name\"].unique()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": {}, + "outputs": [], + "source": [ + "CAM_SHUTTER_INTERVAL_MS" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/av2/delete_me_map.ipynb b/notebooks/av2/delete_me_map.ipynb new file mode 100644 index 00000000..03e62838 --- /dev/null +++ b/notebooks/av2/delete_me_map.ipynb @@ -0,0 +1,313 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import io\n", + "\n", + "from PIL import Image\n", + "\n", + "import pandas as pd" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"test\"\n", + "\n", + "\n", + "split = \"train\"\n", + "# split = \"val\"\n", + "\n", + "split_folder = Path(f\"/media/nvme1/argoverse/sensor_mini/{split}\")\n", + "log_folder = sorted(list(split_folder.iterdir()))[2]\n", + "\n", + "log_folder.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "for lane_segment_id, lane_segment_dict in log_map_archive[\"lane_segments\"].items():\n", + " print(f\"Lane Segment ID: {lane_segment_id}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from typing import Dict, List\n", + "\n", + "from d123.common.geometry.line.polylines import Polyline3D\n", + "from d123.dataset.dataset_specific.av2.av2_map_conversion import _extract_lane_group_dict\n", + "\n", + "\n", + "def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Polyline3D:\n", + " polyline = np.array([[p[\"x\"], p[\"y\"], p[\"z\"]] for p in data], dtype=np.float64)\n", + " if close:\n", + " polyline = np.vstack([polyline, polyline[0]])\n", + "\n", + " return Polyline3D.from_array(polyline)\n", + "\n", + "\n", + "map_folder = log_folder / \"map\"\n", + "log_map_archive_path = next(map_folder.glob(\"log_map_archive_*.json\"))\n", + "\n", + "with open(log_map_archive_path, \"r\") as f:\n", + " log_map_archive = json.load(f)\n", + "\n", + "drivable_areas: Dict[int, Polyline3D] = {}\n", + "\n", + "for drivable_area_id, drivable_area_dict in log_map_archive[\"drivable_areas\"].items():\n", + " # keys: [\"area_boundary\", \"id\"]\n", + " drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict[\"area_boundary\"], close=True)\n", + "\n", + "for lane_segment_id, lane_segment_dict in log_map_archive[\"lane_segments\"].items():\n", + " # keys = [\n", + " # \"id\",\n", + " # \"is_intersection\",\n", + " # \"lane_type\",\n", + " # \"left_lane_boundary\",\n", + " # \"left_lane_mark_type\",\n", + " # \"right_lane_boundary\",\n", + " # \"right_lane_mark_type\",\n", + " # \"successors\",\n", + " # \"predecessors\",\n", + " # \"right_neighbor_id\",\n", + " # \"left_neighbor_id\",\n", + " # ]\n", + " lane_segment_dict[\"left_lane_boundary\"] = _extract_polyline(lane_segment_dict[\"left_lane_boundary\"])\n", + " lane_segment_dict[\"right_lane_boundary\"] = _extract_polyline(lane_segment_dict[\"right_lane_boundary\"])\n", + "\n", + "for crosswalk_id, crosswalk_dict in log_map_archive[\"pedestrian_crossings\"].items():\n", + " # keys = [\"id\", \"outline\"]\n", + " # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/map/pedestrian_crossing.py\n", + "\n", + " p1, p2 = np.array([[p[\"x\"], p[\"y\"], p[\"z\"]] for p in crosswalk_dict[\"edge1\"]], dtype=np.float64)\n", + " p3, p4 = np.array([[p[\"x\"], p[\"y\"], p[\"z\"]] for p in crosswalk_dict[\"edge2\"]], dtype=np.float64)\n", + " crosswalk_dict[\"outline\"] = Polyline3D.from_array(np.array([p1, p2, p4, p3, p1], dtype=np.float64))\n", + "\n", + "\n", + "lane_group_dict = _extract_lane_group_dict(log_map_archive[\"lane_segments\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Any\n", + "\n", + "import shapely\n", + "\n", + "from d123.common.geometry.base import Point3DIndex\n", + "import geopandas as gpd\n", + "\n", + "from d123.common.geometry.line.polylines import Polyline2D\n", + "from d123.common.geometry.occupancy_map import OccupancyMap2D\n", + "\n", + "import numpy.typing as npt\n", + "\n", + "\n", + "def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float:\n", + " \"\"\"Interpolate Z coordinate along a 3D line segment.\"\"\"\n", + " p1, p2 = segment_coords[0], segment_coords[1]\n", + "\n", + " # Project point onto segment\n", + " segment_vec = p2[:2] - p1[:2]\n", + " point_vec = np.array([point.x, point.y]) - p1[:2]\n", + "\n", + " # Handle degenerate case\n", + " segment_length_sq = np.dot(segment_vec, segment_vec)\n", + " if segment_length_sq == 0:\n", + " return p1[2]\n", + "\n", + " # Calculate projection parameter\n", + " t = np.dot(point_vec, segment_vec) / segment_length_sq\n", + " t = np.clip(t, 0, 1) # Clamp to segment bounds\n", + "\n", + " # Interpolate Z\n", + " return p1[2] + t * (p2[2] - p1[2])\n", + "\n", + "\n", + "def _extract_intersection_dict(\n", + " lanes: Dict[int, Any], lane_group_dict: Dict[int, Any], max_distance: float = 0.01\n", + ") -> Dict[str, Any]:\n", + "\n", + " # 1. Collect all lane groups where at least one lane is marked as an intersection.\n", + " lane_group_intersection_dict = {}\n", + " for lane_group_id, lane_group in lane_group_dict.items():\n", + " is_intersection_lanes = [lanes[str(lane_id)][\"is_intersection\"] for lane_id in lane_group[\"lane_ids\"]]\n", + " if any(is_intersection_lanes):\n", + " lane_group_intersection_dict[lane_group_id] = lane_group\n", + "\n", + " # 2. Merge polygons of lane groups that are marked as intersections.\n", + " lane_group_intersection_geometry = {\n", + " lane_group_id: shapely.Polygon(lane_group[\"outline\"].array[:, Point3DIndex.XY])\n", + " for lane_group_id, lane_group in lane_group_intersection_dict.items()\n", + " }\n", + " intersection_polygons = gpd.GeoSeries(lane_group_intersection_geometry).union_all()\n", + "\n", + " # 3. Collect all intersection polygons and their lane group IDs.\n", + " intersection_dict = {}\n", + " for intersection_idx, intersection_polygon in enumerate(intersection_polygons.geoms):\n", + " if intersection_polygon.is_empty:\n", + " continue\n", + " lane_group_ids = [\n", + " lane_group_id\n", + " for lane_group_id, lane_group_polygon in lane_group_intersection_geometry.items()\n", + " if intersection_polygon.intersects(lane_group_polygon)\n", + " ]\n", + " intersection_dict[f\"intersection_{intersection_idx}\"] = {\n", + " \"id\": intersection_idx,\n", + " \"outline_2d\": Polyline2D.from_array(np.array(list(intersection_polygon.exterior.coords), dtype=np.float64)),\n", + " \"lane_group_ids\": lane_group_ids,\n", + " }\n", + "\n", + " # 4. Lift intersection outlines to 3D.\n", + " boundary_segments = []\n", + " for lane_group in lane_group_intersection_dict.values():\n", + " coords = np.array(lane_group[\"outline\"].linestring.coords, dtype=np.float64).reshape(-1, 1, 3)\n", + " segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1)\n", + " boundary_segments.append(segment_coords_boundary)\n", + "\n", + " boundary_segments = np.concatenate(boundary_segments, axis=0)\n", + " boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments)\n", + " occupancy_map = OccupancyMap2D(boundary_segment_linestrings)\n", + "\n", + " for intersection_id, intersection_data in intersection_dict.items():\n", + " points_2d = intersection_data[\"outline_2d\"].array\n", + " points_3d = np.zeros((len(points_2d), 3), dtype=np.float64)\n", + " points_3d[:, :2] = points_2d\n", + "\n", + " query_points = shapely.creation.points(points_2d)\n", + " results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True)\n", + " for query_idx, geometry_idx in zip(*results):\n", + " query_point = query_points[query_idx]\n", + " segment_coords = boundary_segments[geometry_idx]\n", + " best_z = _interpolate_z_on_segment(query_point, segment_coords)\n", + " points_3d[query_idx, 2] = best_z\n", + "\n", + " intersection_dict[intersection_id][\"outline_3d\"] = Polyline3D.from_array(points_3d)\n", + "\n", + " return intersection_dict\n", + "\n", + "\n", + "intersection_dict = _extract_intersection_dict(log_map_archive[\"lane_segments\"], lane_group_dict)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import cm\n", + "from matplotlib.colors import Normalize\n", + "\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "for intersection_id, values in intersection_dict.items():\n", + "\n", + " outline = values[\"outline_3d\"].array\n", + " print(outline[:, 2].min(), outline[:, 2].max())\n", + "\n", + " # Normalize z values to [0, 1] for colormap mapping\n", + " norm = Normalize(vmin=outline[:, 2].min(), vmax=outline[:, 2].max())\n", + " colors = cm.viridis(norm(outline[:, 2]))\n", + "\n", + " # Plot each segment with its corresponding color\n", + " for i in range(len(outline) - 1):\n", + " ax.plot(outline[i : i + 2, 0], outline[i : i + 2, 1], color=colors[i], linewidth=2)\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "\n", + "for lane in log_map_archive[\"lane_segments\"].values():\n", + " left_boundary = lane[\"left_lane_boundary\"]\n", + " right_boundary = lane[\"right_lane_boundary\"]\n", + "\n", + " polygon = np.concatenate([left_boundary.array, right_boundary.array[::-1]]).reshape(-1, 3)[:, :2]\n", + " print(polygon)\n", + " ax.fill(\n", + " polygon[:, 0], polygon[:, 1], alpha=0.5, edgecolor=\"black\", color=\"red\" if lane[\"is_intersection\"] else \"blue\"\n", + " )\n", + "\n", + " # if left_boundary and right_boundary:\n", + " # ax.plot(left_boundary.array[:, 0], left_boundary.array[:, 1], color=\"blue\", linewidth=1)\n", + " # ax.plot(right_boundary.array[:, 0], right_boundary.array[:, 1], color=\"red\", linewidth=1)\n", + "\n", + "ax.set_title(\"Lane Segments\")\n", + "ax.set_xlabel(\"X\")\n", + "ax.set_ylabel(\"Y\")\n", + "ax.set_aspect(\"equal\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/demo_gym.ipynb b/notebooks/demo_gym.ipynb deleted file mode 100644 index 26622558..00000000 --- a/notebooks/demo_gym.ipynb +++ /dev/null @@ -1,310 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", - "from d123.dataset.scene.scene_filter import SceneFilter\n", - "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "import os, psutil\n", - "\n", - "\n", - "def print_memory_usage():\n", - " process = psutil.Process(os.getpid())\n", - " memory_info = process.memory_info()\n", - " print(f\"Memory usage: {memory_info.rss / 1024 ** 2:.2f} MB\")\n", - "\n", - "\n", - "split = \"nuplan_mini_val\"\n", - "\n", - "# log_names = [\"2021.06.07.12.54.00_veh-35_01843_02314\"]\n", - "scene_tokens = None\n", - "scene_tokens = [\"2283aea39bc1505e\"]\n", - "log_names = None\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=[split], log_names=log_names, scene_tokens=scene_tokens, duration_s=15.1, history_s=1.0\n", - ")\n", - "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", - "worker = Sequential()\n", - "# worker = RayDistributed()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "print(len(scenes))\n", - "\n", - "for scene in scenes[:10]:\n", - " print(scene.log_name, scene.token)\n", - "\n", - "\n", - "# 875f407079785268\n", - "\n", - "# _Rep0_longest1_route0_06_13_17_21_21 457da031d28ba67b\n", - "# _Rep0_longest1_route0_06_13_17_21_21 1f6d4a8c9a399b3b\n", - "# _Rep0_longest1_route0_06_13_17_21_21 5b7a3e90922db277\n", - "# _Rep0_longest1_route0_06_13_17_21_21 827d6ac2ff5356d0\n", - "# _Rep0_longest1_route0_06_13_17_21_21 16717aba72175425\n", - "# _Rep0_longest1_route0_06_13_17_21_21 5b7444f86f19b444\n", - "# _Rep0_longest1_route0_06_13_17_21_21 6ee6004b304f1d3c\n", - "# _Rep0_longest1_route0_06_13_17_21_21 c98966ed60a77f7e\n", - "# _Rep0_longest1_route0_06_13_17_21_21 790b09ca88553770\n", - "# _Rep0_longest1_route0_06_13_17_21_21 b6c1a8d385648623" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [ - "from pathlib import Path\n", - "from typing import Optional, Tuple\n", - "\n", - "import matplotlib.animation as animation\n", - "import matplotlib.pyplot as plt\n", - "from tqdm import tqdm\n", - "\n", - "from d123.common.geometry.base import Point2D, StateSE2\n", - "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE2\n", - "from d123.common.visualization.color.default import EGO_VEHICLE_CONFIG\n", - "from d123.common.visualization.matplotlib.observation import (\n", - " add_bounding_box_to_ax,\n", - " add_box_detections_to_ax,\n", - " add_default_map_on_ax,\n", - " add_traffic_lights_to_ax,\n", - ")\n", - "from d123.dataset.arrow.conversion import TrafficLightDetectionWrapper\n", - "from d123.dataset.maps.abstract_map import AbstractMap\n", - "from d123.common.datatypes.detection.detection import BoxDetectionWrapper\n", - "from d123.dataset.scene.abstract_scene import AbstractScene\n", - "from nuplan.common.actor_state.ego_state import EgoState\n", - "import io\n", - "from PIL import Image\n", - "\n", - "\n", - "def add_ego_vehicle_to_ax_(ax: plt.Axes, ego_state: EgoState) -> None:\n", - " bounding_box = BoundingBoxSE2(\n", - " center=StateSE2(*ego_state.center),\n", - " length=ego_state.car_footprint.length,\n", - " width=ego_state.car_footprint.width,\n", - " )\n", - " add_bounding_box_to_ax(ax, bounding_box, EGO_VEHICLE_CONFIG)\n", - "\n", - "\n", - "def _plot_scene_on_ax(\n", - " ax: plt.Axes,\n", - " map_api: AbstractMap,\n", - " ego_state: EgoState,\n", - " initial_ego_state: Optional[EgoState],\n", - " box_detections: BoxDetectionWrapper,\n", - " traffic_light_detections: TrafficLightDetectionWrapper,\n", - " radius: float = 120,\n", - ") -> plt.Axes:\n", - "\n", - " if initial_ego_state is not None:\n", - " point_2d = Point2D(initial_ego_state.center.x, initial_ego_state.center.y)\n", - " else:\n", - "\n", - " point_2d = Point2D(ego_state.center.x, ego_state.center.y)\n", - " add_default_map_on_ax(ax, map_api, point_2d, radius=radius)\n", - " add_traffic_lights_to_ax(ax, traffic_light_detections, map_api)\n", - "\n", - " add_box_detections_to_ax(ax, box_detections)\n", - " add_ego_vehicle_to_ax_(ax, ego_state)\n", - "\n", - " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", - " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", - "\n", - " ax.set_aspect(\"equal\", adjustable=\"box\")\n", - " return ax\n", - "\n", - "\n", - "def plot_scene_to_image(\n", - " map_api: AbstractMap,\n", - " ego_state: EgoState,\n", - " initial_ego_state: Optional[EgoState],\n", - " box_detections: BoxDetectionWrapper,\n", - " traffic_light_detections: TrafficLightDetectionWrapper,\n", - " radius: float = 120,\n", - " figsize: Tuple[int, int] = (8, 8),\n", - ") -> Image:\n", - "\n", - " fig, ax = plt.subplots(figsize=figsize)\n", - " _plot_scene_on_ax(ax, map_api, ego_state, initial_ego_state, box_detections, traffic_light_detections, radius)\n", - " ax.set_aspect(\"equal\", adjustable=\"box\")\n", - " plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)\n", - " # plt.tight_layout()\n", - "\n", - " buf = io.BytesIO()\n", - " fig.savefig(buf, format=\"png\", bbox_inches=\"tight\")\n", - " plt.close(fig)\n", - " buf.seek(0)\n", - " img = Image.open(buf)\n", - " return img" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.dataset.arrow.conversion import DetectionType\n", - "from d123.simulation.gym.demo_gym_env import DemoGymEnv\n", - "from d123.simulation.observation.agents_observation import _filter_agents_by_type\n", - "\n", - "import time\n", - "\n", - "images = []\n", - "agent_rollouts = []\n", - "plot: bool = True\n", - "action = [1.0, 0.1] # Placeholder action, replace with actual action logic\n", - "env = DemoGymEnv(scenes)\n", - "\n", - "start = time.time()\n", - "\n", - "map_api, ego_state, detection_observation, current_scene = env.reset(None)\n", - "initial_ego_state = ego_state\n", - "cars, _, _ = _filter_agents_by_type(detection_observation.box_detections, detection_types=[DetectionType.VEHICLE])\n", - "agent_rollouts.append(BoxDetectionWrapper(cars))\n", - "if plot:\n", - " images.append(\n", - " plot_scene_to_image(\n", - " map_api,\n", - " ego_state,\n", - " initial_ego_state,\n", - " detection_observation.box_detections,\n", - " detection_observation.traffic_light_detections,\n", - " )\n", - " )\n", - "\n", - "\n", - "for i in range(160):\n", - " ego_state, detection_observation, end = env.step(action)\n", - " cars, _, _ = _filter_agents_by_type(detection_observation.box_detections, detection_types=[DetectionType.VEHICLE])\n", - " agent_rollouts.append(BoxDetectionWrapper(cars))\n", - " if plot:\n", - " images.append(\n", - " plot_scene_to_image(\n", - " map_api,\n", - " ego_state,\n", - " initial_ego_state,\n", - " detection_observation.box_detections,\n", - " detection_observation.traffic_light_detections,\n", - " )\n", - " )\n", - " if end:\n", - " print(\"End of scene reached.\")\n", - " break\n", - "\n", - "print(time.time() - start)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "\n", - "def create_gif(images, output_path, duration=100):\n", - " \"\"\"\n", - " Create a GIF from a list of PIL images.\n", - "\n", - " Args:\n", - " images (list): List of PIL.Image objects.\n", - " output_path (str): Path to save the GIF.\n", - " duration (int): Duration between frames in milliseconds.\n", - " \"\"\"\n", - " if images:\n", - " print(len(images))\n", - " images_p = [img.convert(\"P\", palette=Image.ADAPTIVE) for img in images]\n", - " images_p[0].save(output_path, save_all=True, append_images=images_p[1:], duration=duration, loop=0)\n", - "\n", - "\n", - "if plot:\n", - " create_gif(images, f\"{split}_{current_scene.token}.gif\", duration=20)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [ - "from typing import List\n", - "\n", - "from d123.simulation.metrics.sim_agents.sim_agents import get_sim_agents_metrics\n", - "\n", - "\n", - "\n", - "\n", - "result = get_sim_agents_metrics(current_scene, agent_rollouts)\n", - "result\n", - "\n", - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "d123", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/deprecated/arrow.ipynb b/notebooks/deprecated/arrow.ipynb index 0a2b6337..8fa5d9bb 100644 --- a/notebooks/deprecated/arrow.ipynb +++ b/notebooks/deprecated/arrow.ipynb @@ -32,7 +32,7 @@ "\n", "import shapely.geometry as geom\n", "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", - "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", "\n", "\n", "temp_gpkg = \"/home/daniel/d123_workspace/data/maps/carla_town01.gpkg\"\n", @@ -44,7 +44,7 @@ "print_memory_usage()\n", "\n", "\n", - "table = map_api._gpd_dataframes[MapSurfaceType.LANE]\n", + "table = map_api._gpd_dataframes[MapLayer.LANE]\n", "\n", "# import pyarrow as pa\n", "# pa.table(table).schema\n", diff --git a/notebooks/deprecated/collect_sim_metrics_gym.ipynb b/notebooks/deprecated/collect_sim_metrics_gym.ipynb index 2e5b0f82..40100cd6 100644 --- a/notebooks/deprecated/collect_sim_metrics_gym.ipynb +++ b/notebooks/deprecated/collect_sim_metrics_gym.ipynb @@ -9,7 +9,7 @@ "source": [ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n" + "from d123.common.multithreading.worker_sequential import Sequential\n" ] }, { diff --git a/notebooks/deprecated/extraction_testing.ipynb b/notebooks/deprecated/extraction_testing.ipynb index 652601ad..d782d7ba 100644 --- a/notebooks/deprecated/extraction_testing.ipynb +++ b/notebooks/deprecated/extraction_testing.ipynb @@ -14,7 +14,7 @@ "\n", "\n", "import pyarrow as pa\n", - "from nuplan.planning.utils.multithreading.worker_pool import WorkerPool\n", + "from d123.common.multithreading.worker_pool import WorkerPool\n", "\n", "from d123.dataset.arrow.helper import open_arrow_arrow_table\n", "from d123.dataset.dataset_specific.nuplan.nuplan_data_processor import worker_map\n", diff --git a/notebooks/map_conversion/test_nuplan_conversion.ipynb b/notebooks/deprecated/map_conversion/test_nuplan_conversion.ipynb similarity index 69% rename from notebooks/map_conversion/test_nuplan_conversion.ipynb rename to notebooks/deprecated/map_conversion/test_nuplan_conversion.ipynb index 869e55f3..32805c97 100644 --- a/notebooks/map_conversion/test_nuplan_conversion.ipynb +++ b/notebooks/deprecated/map_conversion/test_nuplan_conversion.ipynb @@ -6,37 +6,10 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", - "from shapely.geometry import LineString, Polygon, Point\n", - "import numpy as np\n", - "\n", - "from typing import List\n", - "\n", - "\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lane, LaneSection\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", "from pathlib import Path\n", - "from d123.dataset.dataset_specific.nuplan.map_conversion import NuPlanMapConverter, MAP_LOCATIONS\n", + "from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import NuPlanMapConverter, MAP_LOCATIONS\n", "\n", - "import pyogrio\n", "\n", - "import geopandas as gpd\n", "\n", "\n", "# file_to_delete = Path(\"nuplan_us-pa-pittsburgh-hazelwood.gpkg\")\n", @@ -53,7 +26,7 @@ " file_to_delete = root / f\"nuplan_{map_name}.gpkg\"\n", " if file_to_delete.exists():\n", " file_to_delete.unlink()\n", - " \n", + "\n", " NuPlanMapConverter(root).convert(map_name=map_name)\n", " print(f\"Converting {map_name} map... done\")\n", "\n", diff --git a/notebooks/map_conversion/test_opendrive_conversion.ipynb b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb similarity index 85% rename from notebooks/map_conversion/test_opendrive_conversion.ipynb rename to notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb index 0de615ec..ac7aaded 100644 --- a/notebooks/map_conversion/test_opendrive_conversion.ipynb +++ b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb @@ -6,14 +6,14 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", - "from shapely.geometry import LineString, Polygon, Point\n", - "import numpy as np\n", + "# import matplotlib.pyplot as plt\n", + "# from shapely.geometry import LineString, Polygon, Point\n", + "# import numpy as np\n", "\n", - "from typing import List\n", + "# # from typing import List\n", "\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lane, LaneSection\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border" + "# # from d123.dataset.conversion.map.opendrive.elements.lane import Lane, LaneSection\n", + "# # from d123.dataset.conversion.map.opendrive.elements.reference import Border" ] }, { @@ -23,7 +23,7 @@ "outputs": [], "source": [ "from pathlib import Path\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive\n", + "from d123.dataset.conversion.map.opendrive.parser.opendrive import OpenDrive\n", "\n", "CARLA_MAP_LOCATIONS = [\n", " \"Town01\", # A small, simple town with a river and several bridges.\n", @@ -79,13 +79,12 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.dataset.dataset_specific.carla.opendrive.opendrive_converter import OpenDriveConverter\n", + "from d123.dataset.conversion.map.opendrive.opendrive_converter import OpenDriveConverter\n", "\n", "\n", "\n", "MAP_ROOT = Path(\"/home/daniel/d123_workspace/data\") / \"maps\"\n", "\n", - "\n", "for town_name in AVAILABLE_CARLA_MAP_LOCATIONS:\n", " map_name = town_name.lower()\n", " print(f\"Start {map_name} map...\")\n", @@ -110,6 +109,7 @@ " converter = OpenDriveConverter(opendrive)\n", " converter.run(f\"carla_{town_name.lower()}\")\n", " print(f\"Converting {map_name} map... done\")\n", + " # break\n", "\n", "\n" ] @@ -129,24 +129,24 @@ "metadata": {}, "outputs": [], "source": [ - "# from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", "\n", "\n", - "# str(MapSurfaceType.GENERIC_DRIVABLE).split(\".\")[-1].lower()\n", + "str(MapLayer.GENERIC_DRIVABLE).split(\".\")[-1].lower()\n", "\n", "\n", - "# MapSurfaceType.GENERIC_DRIVABLE.name\n", + "MapLayer.GENERIC_DRIVABLE.name\n", "\n", - "# MapSurfaceType.deserialize(MapSurfaceType.GENERIC_DRIVABLE.name)\n", + "MapLayer.deserialize(MapLayer.GENERIC_DRIVABLE.name)\n", "\n", "\n", - "# MapSurfaceType.GENERIC_DRIVABLE.name.lower().islower()\n", + "MapLayer.GENERIC_DRIVABLE.name.lower().islower()\n", "\n", "\n", - "# AVAILABLE_MAP_LAYERS = list(MapSurfaceType)\n", - "# AVAILABLE_MAP_LAYERS\n", + "AVAILABLE_MAP_LAYERS = list(MapLayer)\n", + "AVAILABLE_MAP_LAYERS\n", "\n", - "# pyogrio.read_dataframe()." + "pyogrio.read_dataframe()." ] }, { diff --git a/notebooks/deprecated/nuplan_map_dataframe.ipynb b/notebooks/deprecated/nuplan_map_dataframe.ipynb index 94c372ff..113d8391 100644 --- a/notebooks/deprecated/nuplan_map_dataframe.ipynb +++ b/notebooks/deprecated/nuplan_map_dataframe.ipynb @@ -11,6 +11,16 @@ "MAP_LOCATIONS = {\"sg-one-north\", \"us-ma-boston\", \"us-nv-las-vegas-strip\", \"us-pa-pittsburgh-hazelwood\"}" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "length = {0: 1, 1: 2, 2: 3, 3: 4, 4: 0}\n", + "max(length, key=length.get)" + ] + }, { "cell_type": "code", "execution_count": null, @@ -49,7 +59,6 @@ "\n", "def load_layer(layer_name: str) -> gpd.geodataframe:\n", "\n", - "\n", " map_meta = gpd.read_file(path_to_gpkg, layer=\"meta\", engine=\"pyogrio\")\n", " projection_system = map_meta[map_meta[\"key\"] == \"projectedCoordSystem\"][\"value\"].iloc[0]\n", "\n", @@ -63,6 +72,7 @@ "\n", " return gdf_in_utm_coords\n", "\n", + "\n", "pyogrio.read_dataframe(path_to_gpkg)" ] }, @@ -72,7 +82,11 @@ "metadata": {}, "outputs": [], "source": [ - "[\n", + "import pandas as pd\n", + "\n", + "\n", + "layers = [\n", + " \"baseline_paths\",\n", " \"carpark_areas\",\n", " \"generic_drivable_areas\",\n", " \"dubins_nodes\",\n", @@ -91,18 +105,25 @@ " \"meta\",\n", "]\n", "\n", - "lanes = load_layer(\"stop_polygons\")\n", + "\n", + "def non_nan_set(series: pd.Series) -> set:\n", + " return set(series.dropna().values)\n", + "\n", + "\n", + "lanes = load_layer(\"intersections\")\n", "boundaries = load_layer(\"boundaries\")\n", "\n", "\n", "def get_lane(id: int):\n", " return lanes[lanes[\"fid\"] == str(id)]\n", "\n", + "\n", "def get_right_boundary(lane):\n", " return boundaries[boundaries[\"fid\"] == str(lane[\"right_boundary_fid\"].iloc[0])]\n", "\n", + "\n", "def get_left_boundary(lane):\n", - " return boundaries[boundaries[\"fid\"] == str(lane[\"left_boundary_fid\"].iloc[0])]\n", + " return boundaries[boundaries[\"fid\"] == str(lane[\"left_boundary_fid\"].iloc[0])]\n", "\n", "\n", "# boundaries\n", @@ -114,11 +135,12 @@ "# error_lane_df = get_lane(error_lane_fid)\n", "\n", "\n", - "\n", "# 46552,46553\n", "# test_id = 46553\n", "\n", - "lanes" + "print(lanes.keys())\n", + "non_nan_set(lanes.intersection_type_fid)\n", + "\n" ] }, { @@ -127,7 +149,114 @@ "metadata": {}, "outputs": [], "source": [ - "lanes[lanes[\"right_boundary_fid\"] == test_id]" + "baseline_paths = load_layer(\"baseline_paths\")\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from networkx import center\n", + "\n", + "# Create a figure\n", + "fig, ax = plt.subplots(figsize=(15, 12))\n", + "\n", + "# Get unique path_type_fid values\n", + "path_types = baseline_paths[\"path_type_fid\"].unique()\n", + "\n", + "# Define colors for different path types\n", + "colors = [\"blue\", \"red\", \"green\", \"brown\", \"purple\", \"orange\", \"pink\"]\n", + "\n", + "# Plot each path type with a different color\n", + "target_label = 0\n", + "for i, path_type in enumerate(path_types):\n", + " # Filter baseline_paths by path_type\n", + " paths_of_type = baseline_paths[baseline_paths[\"path_type_fid\"] == path_type]\n", + "\n", + " if path_type == target_label:\n", + " alpha = 1.0\n", + " target_fids = paths_of_type.fid.tolist()\n", + " else:\n", + " alpha = 0.5\n", + "\n", + " # Plot these paths with a specific color\n", + " paths_of_type.plot(ax=ax, color=colors[i % len(colors)], label=f\"Path Type {path_type}\", alpha=alpha)\n", + "\n", + "# Add title and legend\n", + "ax.set_title(\"Baseline Paths by Path Type\", fontsize=16)\n", + "ax.legend()\n", + "ax.set_xlabel(\"X Coordinate\")\n", + "ax.set_ylabel(\"Y Coordinate\")\n", + "ax.axis(\"equal\") # Maintain aspect ratio\n", + "\n", + "center_zoom = False # Set to True to zoom into the center of the data\n", + "fid_zoom = target_fids[25]\n", + "\n", + "if center_zoom:\n", + " # Get the bounds of the data for better focusing\n", + " x_min, y_min, x_max, y_max = baseline_paths.total_bounds\n", + "\n", + " # Calculate center coordinates\n", + " center_x = (x_min + x_max) / 2\n", + " center_y = (y_min + y_max) / 2\n", + "\n", + " # Set axis limits to zoom into the center (using 30% of the total range)\n", + " range_x = x_max - x_min\n", + " range_y = y_max - y_min\n", + " zoom_factor = 2.0\n", + "\n", + " ax.set_xlim(center_x - range_x * zoom_factor / 2, center_x + range_x * zoom_factor / 2)\n", + " ax.set_ylim(center_y - range_y * zoom_factor / 2, center_y + range_y * zoom_factor / 2)\n", + "\n", + "elif fid_zoom:\n", + " # Filter to get the specific element with the given FID\n", + " specific_path = baseline_paths[baseline_paths[\"fid\"] == str(fid_zoom)]\n", + "\n", + " if not specific_path.empty:\n", + " # Get the bounds of the specific path\n", + " x_min, y_min, x_max, y_max = specific_path.total_bounds\n", + "\n", + " # Add some padding around the element\n", + " padding = 100 # meters\n", + " ax.set_xlim(x_min - padding, x_max + padding)\n", + " ax.set_ylim(y_min - padding, y_max + padding)\n", + "\n", + " # Highlight the specific element\n", + " # specific_path.plot(ax=ax, color=\"yellow\", linewidth=3, zorder=10)\n", + "\n", + " # Update title to show we're zoomed to a specific FID\n", + " ax.set_title(f\"Baseline Path - FID {fid_zoom}\", fontsize=16)\n", + " else:\n", + " print(f\"FID {fid_zoom} not found in baseline_paths\")\n", + "\n", + "\n", + "# Add a title to indicate we're looking at the center\n", + "ax.set_title(\"Baseline Paths by Path Type (Center Zoom)\", fontsize=16)\n", + "\n", + "\n", + "# Convert the specific path coordinates to WGS84 (latitude/longitude) for Google Maps\n", + "if fid_zoom and not specific_path.empty:\n", + " # Create a copy to avoid modifying the original\n", + " wgs84_path = specific_path.copy()\n", + "\n", + " # Convert from the current projection to WGS84 (EPSG:4326)\n", + " wgs84_path = wgs84_path.to_crs(\"EPSG:4326\")\n", + "\n", + " # Get the centroid of the path for easier lookup\n", + " centroid = wgs84_path.geometry.iloc[0].centroid\n", + "\n", + " # Display the coordinates\n", + " print(f\"\\nGoogle Maps coordinates for FID {fid_zoom}:\")\n", + " print(f\"Latitude: {centroid.y}, Longitude: {centroid.x}\")\n", + " print(f\"Google Maps link: https://www.google.com/maps?q={centroid.y},{centroid.x}\")\n", + "\n", + " # Add a text annotation showing coordinates on the plot\n", + " ax.annotate(\n", + " f\"Lat: {centroid.y:.6f}, Lon: {centroid.x:.6f}\",\n", + " xy=(0.05, 0.05),\n", + " xycoords=\"axes fraction\",\n", + " bbox=dict(boxstyle=\"round,pad=0.5\", fc=\"white\", alpha=0.8),\n", + " fontsize=10,\n", + " )\n", + "plt.tight_layout()\n", + "plt.show()" ] }, { @@ -136,7 +265,113 @@ "metadata": {}, "outputs": [], "source": [ - "boundaries[boundaries[\"fid\"] == str(test_id)]" + "baseline_paths = load_layer(\"baseline_paths\")\n", + "intersections = load_layer(\"intersections\") # Load the intersections layer\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from networkx import center\n", + "\n", + "# Create a figure\n", + "fig, ax = plt.subplots(figsize=(15, 12))\n", + "\n", + "# Get unique intersection types\n", + "intersection_types = intersections[\"intersection_type_fid\"].unique()\n", + "\n", + "# Define colors for different intersection types\n", + "colors = [\"blue\", \"red\", \"green\", \"brown\", \"purple\", \"orange\", \"pink\"]\n", + "\n", + "# Plot each intersection type with a different color\n", + "target_label = 2 # Target intersection type to highlight\n", + "for i, intersection_type in enumerate(intersection_types):\n", + " # Filter intersections by type\n", + " intersections_of_type = intersections[intersections[\"intersection_type_fid\"] == intersection_type]\n", + "\n", + " if intersection_type == target_label:\n", + " alpha = 1.0\n", + " target_fids = intersections_of_type.fid.tolist()\n", + " else:\n", + " alpha = 0.5\n", + "\n", + " # Plot these intersections with a specific color\n", + " intersections_of_type.plot(ax=ax, color=colors[i % len(colors)], \n", + " label=f\"Intersection Type {intersection_type}\", alpha=alpha)\n", + "\n", + "# Add title and legend\n", + "ax.set_title(\"Intersections by Type\", fontsize=16)\n", + "\n", + "ax.set_xlabel(\"X Coordinate\")\n", + "ax.set_ylabel(\"Y Coordinate\")\n", + "ax.axis(\"equal\") # Maintain aspect ratio\n", + "\n", + "center_zoom = False # Set to True to zoom into the center of the data\n", + "fid_zoom = target_fids[0] if 'target_fids' in locals() and len(target_fids) > 0 else None\n", + "\n", + "if center_zoom:\n", + " # Get the bounds of the data for better focusing\n", + " x_min, y_min, x_max, y_max = intersections.total_bounds\n", + "\n", + " # Calculate center coordinates\n", + " center_x = (x_min + x_max) / 2\n", + " center_y = (y_min + y_max) / 2\n", + "\n", + " # Set axis limits to zoom into the center (using 30% of the total range)\n", + " range_x = x_max - x_min\n", + " range_y = y_max - y_min\n", + " zoom_factor = 2.0\n", + "\n", + " ax.set_xlim(center_x - range_x * zoom_factor / 2, center_x + range_x * zoom_factor / 2)\n", + " ax.set_ylim(center_y - range_y * zoom_factor / 2, center_y + range_y * zoom_factor / 2)\n", + "\n", + "elif fid_zoom:\n", + " # Filter to get the specific intersection with the given FID\n", + " specific_intersection = intersections[intersections[\"fid\"] == fid_zoom]\n", + "\n", + " if not specific_intersection.empty:\n", + " # Get the bounds of the specific intersection\n", + " x_min, y_min, x_max, y_max = specific_intersection.total_bounds\n", + "\n", + " # Add some padding around the element\n", + " padding = 100 # meters\n", + " ax.set_xlim(x_min - padding, x_max + padding)\n", + " ax.set_ylim(y_min - padding, y_max + padding)\n", + "\n", + " # Update title to show we're zoomed to a specific FID\n", + " ax.set_title(f\"Intersection - FID {fid_zoom} (Type {specific_intersection['intersection_type_fid'].iloc[0]})\", fontsize=16)\n", + " else:\n", + " print(f\"FID {fid_zoom} not found in intersections\")\n", + "\n", + "\n", + "# Add a title if not zooming to a specific FID\n", + "if not fid_zoom:\n", + " ax.set_title(\"Intersections by Type\", fontsize=16)\n", + "\n", + "# Convert the specific intersection coordinates to WGS84 (latitude/longitude) for Google Maps\n", + "if fid_zoom and 'specific_intersection' in locals() and not specific_intersection.empty:\n", + " # Create a copy to avoid modifying the original\n", + " wgs84_intersection = specific_intersection.copy()\n", + "\n", + " # Convert from the current projection to WGS84 (EPSG:4326)\n", + " wgs84_intersection = wgs84_intersection.to_crs(\"EPSG:4326\")\n", + "\n", + " # Get the centroid of the intersection for easier lookup\n", + " centroid = wgs84_intersection.geometry.iloc[0].centroid\n", + "\n", + " # Display the coordinates\n", + " print(f\"\\nGoogle Maps coordinates for Intersection FID {fid_zoom}:\")\n", + " print(f\"Latitude: {centroid.y}, Longitude: {centroid.x}\")\n", + " print(f\"Google Maps link: https://www.google.com/maps?q={centroid.y},{centroid.x}\")\n", + "\n", + " # Add a text annotation showing coordinates on the plot\n", + " ax.annotate(\n", + " f\"Lat: {centroid.y:.6f}, Lon: {centroid.x:.6f}\",\n", + " xy=(0.05, 0.05),\n", + " xycoords=\"axes fraction\",\n", + " bbox=dict(boxstyle=\"round,pad=0.5\", fc=\"white\", alpha=0.8),\n", + " fontsize=10,\n", + " )\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "plt.show()\n" ] }, { @@ -212,9 +447,9 @@ "metadata": {}, "outputs": [], "source": [ - "\n", "import pandas as pd\n", - "lane_df = load_layer(\"generic_drivable_areas\")\n", + "\n", + "lane_df = load_layer(\"generic_drivable_areas\")\n", "lane_df" ] }, @@ -234,7 +469,6 @@ "metadata": {}, "outputs": [], "source": [ - "\n", "gen_lane_connectors_scaled_width_polygons_df = load_layer(\"boundaries\")\n", "\n", "gen_lane_connectors_scaled_width_polygons_df[gen_lane_connectors_scaled_width_polygons_df[\"fid\"] == \"17950\"]" @@ -256,7 +490,8 @@ "outputs": [], "source": [ "import pandas as pd\n", - "lane_df = load_layer(\"baseline_paths\")\n", + "\n", + "lane_df = load_layer(\"baseline_paths\")\n", "lane_df" ] }, @@ -273,6 +508,8 @@ " :return: meters per second [m/s]\n", " \"\"\"\n", " return mph / 0.44704\n", + "\n", + "\n", "mps_to_mph(6.705409029950827)" ] }, @@ -287,7 +524,7 @@ "# type = 3\n", "\n", "# for i in np.random.choice(len(geoms[type]), 10):\n", - "# ax.plot(*geoms[type][i].coords.xy, color=\"blue\")\n" + "# ax.plot(*geoms[type][i].coords.xy, color=\"blue\")" ] }, { diff --git a/notebooks/deprecated/test_scene_builder.ipynb b/notebooks/deprecated/test_scene_builder.ipynb index 2f4a608d..00de24f9 100644 --- a/notebooks/deprecated/test_scene_builder.ipynb +++ b/notebooks/deprecated/test_scene_builder.ipynb @@ -10,7 +10,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n" + "from d123.common.multithreading.worker_sequential import Sequential\n" ] }, { diff --git a/notebooks/deprecated/test_waypoints.ipynb b/notebooks/deprecated/test_waypoints.ipynb index a2ce5642..86d2ea9c 100644 --- a/notebooks/deprecated/test_waypoints.ipynb +++ b/notebooks/deprecated/test_waypoints.ipynb @@ -18,7 +18,7 @@ "metadata": {}, "outputs": [], "source": [ - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from d123.common.multithreading.worker_sequential import Sequential" ] }, { diff --git a/notebooks/test_gym.ipynb b/notebooks/gym/test_gym.ipynb similarity index 96% rename from notebooks/test_gym.ipynb rename to notebooks/gym/test_gym.ipynb index 718dd917..bb6cf7dc 100644 --- a/notebooks/test_gym.ipynb +++ b/notebooks/gym/test_gym.ipynb @@ -10,8 +10,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" + "from d123.common.multithreading.worker_sequential import Sequential\n", + "# from d123.common.multithreading.worker_ray import RayDistributed" ] }, { @@ -68,14 +68,6 @@ "id": "3", "metadata": {}, "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], "source": [ "from pathlib import Path\n", "from typing import Optional, Tuple\n", @@ -158,7 +150,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -218,7 +210,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -228,7 +220,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -257,7 +249,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -278,7 +270,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "8", "metadata": {}, "outputs": [], "source": [] diff --git a/notebooks/test_simulation_2d.ipynb b/notebooks/gym/test_simulation_2d.ipynb similarity index 78% rename from notebooks/test_simulation_2d.ipynb rename to notebooks/gym/test_simulation_2d.ipynb index 1aa0fec7..60c22ceb 100644 --- a/notebooks/test_simulation_2d.ipynb +++ b/notebooks/gym/test_simulation_2d.ipynb @@ -10,8 +10,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" + "from d123.common.multithreading.worker_sequential import Sequential\n", + "# from d123.common.multithreading.worker_ray import RayDistributed" ] }, { @@ -33,8 +33,9 @@ " split_names=[split],\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=15.1,\n", + " duration_s=30.1,\n", " history_s=1.0,\n", + " timestamp_threshold_s=10,\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", @@ -69,7 +70,6 @@ " )\n", "\n", "\n", - "\n", "simulation_2d_setup = get_simulation_2d_setup()\n", "simulation2d = Simulation2D(simulation_2d_setup)" ] @@ -82,11 +82,20 @@ "outputs": [], "source": [ "# reset\n", - "from typing import List\n", + "import torch\n", + "from d123.simulation.gym.environment.output_converter.action_output_converter import ActionOutputConverter\n", + "from d123.simulation.gym.policy.ppo.ppo_config import GlobalConfig\n", + "from d123.simulation.gym.policy.ppo.ppo_model import PPOPolicy\n", + "\n", + "\n", + "import numpy as np\n", + "import numpy.typing as npt\n", + "\n", + "\n", "from d123.simulation.gym.environment.gym_observation.raster.raster_gym_observation import RasterGymObservation\n", "from d123.simulation.gym.environment.gym_observation.raster.raster_renderer import RasterRenderer\n", "from d123.simulation.gym.environment.helper.environment_area import RectangleEnvironmentArea\n", - "from d123.simulation.planning.abstract_planner import PlannerInput\n", + "from d123.simulation.planning.abstract_planner import PlannerInitialization, PlannerInput\n", "from d123.simulation.planning.planner_output.action_planner_output import ActionPlannerOutput\n", "\n", "\n", @@ -94,20 +103,54 @@ "gym_observation = RasterGymObservation(environment_area, RasterRenderer(environment_area), inference=True)\n", "gym_observation.reset()\n", "\n", + "checkpoint_path = \"/home/daniel/carl_workspace/carl_nuplan/checkpoints/nuplan_51892_1B/model_best.pth\"\n", + "config = GlobalConfig()\n", + "output_converter = ActionOutputConverter()\n", "\n", - "last_action = [0.0, 0.0] # placeholder\n", - "info = {\"last_action\": last_action}\n", + "device = \"cpu\"\n", + "\n", + "agent = PPOPolicy(\n", + " gym_observation.get_observation_space(),\n", + " output_converter.get_action_space(),\n", + " config=config,\n", + ").to(device)\n", + "\n", + "state_dict = torch.load(checkpoint_path, map_location=device)\n", + "agent.load_state_dict(state_dict, strict=True)\n", + "agent.to(device)\n", + "agent.eval()\n", + "\n", + "\n", + "def forward_agent(\n", + " planner_initialization: PlannerInitialization,\n", + " planner_input: PlannerInput,\n", + " last_action: npt.NDArray[np.float32],\n", + ") -> npt.NDArray[np.float32]:\n", + " info = {\"last_action\": last_action}\n", + " obs = gym_observation.get_gym_observation(planner_input, planner_initialization, info)\n", + "\n", + " obs_tensor = {\n", + " \"bev_semantics\": torch.Tensor(obs[\"bev_semantics\"][None, ...]).to(device, dtype=torch.float32),\n", + " \"measurements\": torch.Tensor(obs[\"measurements\"][None, ...]).to(device, dtype=torch.float32),\n", + " \"value_measurements\": torch.Tensor(obs[\"value_measurements\"][None, ...]).to(device, dtype=torch.float32),\n", + " }\n", + "\n", + " with torch.no_grad():\n", + " (action, _, _, _, _, _, _, _, _, _, _) = agent.forward(\n", + " obs_tensor, deterministic=True, lstm_state=None, done=None\n", + " )\n", + " action = action.squeeze().detach().cpu().numpy()\n", + " print(action)\n", + " return action\n", "\n", + "\n", + "last_action = [0.0, 0.0] # placeholder\n", "rasters = []\n", "\n", - "idx = 1338\n", + "idx = 1\n", "planner_initialization, current_planner_input = simulation2d.reset(scenes[idx])\n", - "rasters.append(gym_observation.get_gym_observation(current_planner_input, planner_initialization, info))\n", "\n", "\n", - "# TODO: Implement action planner output\n", - "# TODO: Further test the simulation object.\n", - "\n", "\n", "def _get_action(planner_input: PlannerInput) -> ActionPlannerOutput:\n", " ego_state, _ = planner_input.history.current_state\n", @@ -117,11 +160,18 @@ "while simulation2d.is_simulation_running():\n", "\n", " # 1. trigger planner\n", - " planner_output = _get_action(current_planner_input)\n", + " # planner_output = _get_action(current_planner_input)\n", + "\n", + " last_action = forward_agent(planner_initialization, current_planner_input, last_action)\n", + " planner_output = output_converter.get_planner_output(\n", + " action=last_action,\n", + " ego_state=current_planner_input.history.current_state[0],\n", + " info={},\n", + " )\n", + " print(planner_output._acceleration, planner_output._steering_rate)\n", "\n", " # 2. step simulation\n", - " current_planner_input = simulation2d.step(planner_output)\n", - " rasters.append(gym_observation.get_gym_observation(current_planner_input, planner_initialization, info))" + " current_planner_input = simulation2d.step(planner_output)" ] }, { @@ -204,11 +254,11 @@ " return rgb_image\n", "\n", "\n", - "numpy_images_to_gif(\n", - " [image_to_rbg(raster[\"bev_semantics\"]) for raster in rasters],\n", - " gif_path=\"/home/daniel/d123_workspace/d123/notebooks/simulation_2d.gif\",\n", - " duration=100,\n", - ")" + "# numpy_images_to_gif(\n", + "# [image_to_rbg(raster[\"bev_semantics\"]) for raster in rasters],\n", + "# gif_path=\"/home/daniel/d123_workspace/d123/notebooks/simulation_2d.gif\",\n", + "# duration=100,\n", + "# )" ] }, { @@ -337,6 +387,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/notebooks/nuplan_sensor_loading.ipynb b/notebooks/nuplan/nuplan_sensor_loading.ipynb similarity index 78% rename from notebooks/nuplan_sensor_loading.ipynb rename to notebooks/nuplan/nuplan_sensor_loading.ipynb index 9618c8de..3fc654d8 100644 --- a/notebooks/nuplan_sensor_loading.ipynb +++ b/notebooks/nuplan/nuplan_sensor_loading.ipynb @@ -22,6 +22,28 @@ "id": "1", "metadata": {}, "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "np.zeros((3,3)).tolist()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.dataset_specific.nuplan.nuplan_data_converter import NuplanDataConverter" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], "source": [ "def process_images_to_arrow(input_folder, output_file, batch_size=100):\n", " \"\"\"\n", @@ -190,83 +212,97 @@ { "cell_type": "code", "execution_count": null, - "id": "2", + "id": "4", "metadata": {}, "outputs": [], "source": [ - "# process_images_to_arrow(\"/mnt/nvme/nuplan/dataset/nuplan-v1.1/sensor/2021.07.01.20.35.47_veh-38_00016_00281/CAM_F0\", \"test.arrow\", 10)" + "# process_images_to_arrow(\"/mnt/nvme/nuplan/dataset/nuplan-v1.1/sensor/2021.07.01.20.35.47_veh-38_00016_00281/F0\", \"test.arrow\", 10)" ] }, { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "5", "metadata": {}, "outputs": [], "source": [ - "# Read IPC file\n", - "import io\n", + "# # Read IPC file\n", + "# import io\n", "\n", - "import time\n", + "# import time\n", "\n", - "# with pa.OSFile(\"test.arrow\", 'rb') as source:\n", - "# with ipc.open_file(source) as reader:\n", - "# table = reader.read_all()\n", + "# # with pa.OSFile(\"test.arrow\", 'rb') as source:\n", + "# # with ipc.open_file(source) as reader:\n", + "# # table = reader.read_all()\n", "\n", - "with pa.ipc.open_file(\n", - " pa.memory_map(\"/home/daniel/d123_workspace/data/nuplan_private_test/2021.07.25.16.16.23_veh-26_02446_02589.arrow\")\n", - ") as reader:\n", - " # This doesn't load data into memory yet!\n", - " table = reader.read_all()\n", + "# with pa.ipc.open_file(\n", + "# pa.memory_map(\"/home/daniel/d123_workspace/data/nuplan_private_test/2021.07.25.16.16.23_veh-26_02446_02589.arrow\")\n", + "# ) as reader:\n", + "# # This doesn't load data into memory yet!\n", + "# table = reader.read_all()\n", "\n", "\n", - "print(len(table))\n", - "start = time.time()\n", - "# Extract JPG data\n", - "jpg_data = table[\"front_cam_demo\"][500].as_py()\n", - "read_image = Image.open(io.BytesIO(jpg_data))\n", + "# print(len(table))\n", + "# start = time.time()\n", + "# # Extract JPG data\n", + "# jpg_data = table[\"front_cam_demo\"][500].as_py()\n", + "# read_image = Image.open(io.BytesIO(jpg_data))\n", "\n", - "# read_image = read_image.convert(\"RGB\") # Ensure it's in RGB format\n", - "read_image = np.array(read_image)\n", - "print(read_image.dtype)\n", - "print(f\"Image loaded in {time.time() - start:.4f} seconds\")\n", + "# # read_image = read_image.convert(\"RGB\") # Ensure it's in RGB format\n", + "# read_image = np.array(read_image)\n", + "# print(read_image.dtype)\n", + "# print(f\"Image loaded in {time.time() - start:.4f} seconds\")\n", "\n", - "import matplotlib.pyplot as plt\n", + "# import matplotlib.pyplot as plt\n", "\n", - "plt.imshow(read_image)" + "# plt.imshow(read_image)" ] }, { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "6", "metadata": {}, "outputs": [], "source": [ "from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB\n", "\n", "NUPLAN_DATA_ROOT = Path(os.environ[\"NUPLAN_DATA_ROOT\"])\n", - "log_path = \"/mnt/nvme/nuplan/dataset/nuplan-v1.1/splits/private_test/2021.09.29.17.35.58_veh-44_01671_01819.db\"\n", + "NUPLAN_SENSOR_ROOT = Path(\"/media/nvme1/nuplan/dataset/nuplan-v1.1/sensor_blobs\")\n", + "\n", + "\n", + "log_path = \"/media/nvme1/nuplan/dataset/nuplan-v1.1/splits/private_test/2021.07.01.21.22.09_veh-14_00016_00656.db\"\n", "log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None)\n", "\n", "\n", - "log_db.image" + "log_db.lidar_pc.filename" ] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", "metadata": {}, "outputs": [], "source": [ + "import io\n", "import pickle\n", "from nuplan.database.nuplan_db.nuplan_scenario_queries import get_images_from_lidar_tokens, get_cameras\n", "from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario import NuPlanScenario, CameraChannel, LidarChannel\n", + "\n", + "from nuplan.database.utils.pointclouds.lidar import LidarPointCloud, PointCloud\n", "from pyquaternion import Quaternion\n", "\n", "\n", - "NUPLAN_DB_PATH = \"/mnt/nvme/nuplan/dataset/nuplan-v1.1/splits/private_test\"\n", + "NUPLAN_DB_PATH = \"/media/nvme1/nuplan/dataset/nuplan-v1.1/splits/private_test\"\n", "\n", "\n", "def get_log_cam_info(log):\n", @@ -274,11 +310,10 @@ " log_file = os.path.join(NUPLAN_DB_PATH, log_name + \".db\")\n", "\n", " log_cam_infos = {}\n", - " for cam in get_cameras(log_file, [str(CameraChannel.CAM_F0.value)]):\n", + " for cam in get_cameras(log_file, [str(CameraChannel.F0.value)]):\n", " intrinsics = np.array(pickle.loads(cam.intrinsic))\n", " translation = np.array(pickle.loads(cam.translation))\n", " rotation = np.array(pickle.loads(cam.rotation))\n", - " print(rotation)\n", " rotation = Quaternion(rotation).rotation_matrix\n", " distortion = np.array(pickle.loads(cam.distortion))\n", " c = dict(\n", @@ -295,21 +330,49 @@ "images = []\n", "for lidar_pc in log_db.lidar_pc[::2]:\n", "\n", - " front_image = get_images_from_lidar_tokens(log_path, [lidar_pc.token], [str(CameraChannel.CAM_F0.value)])\n", - " parameters = get_log_cam_info(log_db.log)\n", - " print(parameters)\n", + " # front_image = get_images_from_lidar_tokens(log_path, [lidar_pc.token], [str(CameraChannel.F0.value)])\n", + " # parameters = get_log_cam_info(log_db.log)\n", + " # print(parameters)\n", "\n", - " images.append(list(front_image))\n", + " # images.append(list(front_image))\n", + " lidar_path = NUPLAN_SENSOR_ROOT / lidar_pc.filename\n", + " with open(lidar_path, \"rb\") as fp:\n", + " buffer = io.BytesIO(fp.read())\n", "\n", + " pcd_data = buffer.read() # type: ignore\n", "\n", - "for image in images[0]:\n", - " print(image)" + " # headers_list = []\n", + " # with io.BytesIO(pcd_data) as stream:\n", + " # while True:\n", + " # line = stream.readline().decode('utf8').strip()\n", + " # print(line)\n", + " # if line.startswith('#'):\n", + " # continue\n", + " # columns = line.split()\n", + " # key = columns[0].lower()\n", + " # val = columns[1:] if len(columns) > 2 else columns[1]\n", + " # headers_list.append((key, val))\n", + "\n", + " # if key == 'data':\n", + " # break\n", + "\n", + " pc = PointCloud.parse(pcd_data)\n", + "\n", + " # print(pc.header)\n", + " # pcd_data\n", + "\n", + " # break\n", + "\n", + " # print(isinstance(pcd_data, bytes))\n", + "\n", + " # lidar_point_cloud = LidarPointCloud.from_buffer(buffer, \"pcd\")\n", + " # print(lidar_point_cloud.points.shape)" ] }, { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -347,7 +410,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -362,7 +425,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -372,7 +435,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -408,7 +471,7 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "13", "metadata": {}, "outputs": [], "source": [] @@ -430,7 +493,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.21" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/notebooks/scene_rendering.ipynb b/notebooks/scene_rendering.ipynb index 111590a4..dad53954 100644 --- a/notebooks/scene_rendering.ipynb +++ b/notebooks/scene_rendering.ipynb @@ -10,10 +10,9 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "\n", - "\n" + "# from d123.common.multithreading.worker_ray import RayDistributed" ] }, { @@ -31,12 +30,12 @@ " memory_info = process.memory_info()\n", " print(f\"Memory usage: {memory_info.rss / 1024 ** 2:.2f} MB\")\n", "\n", - "\n", "split = \"nuplan_mini_val\"\n", + "# split = \"carla\"\n", "\n", "# log_names = [\"2021.06.07.12.54.00_veh-35_01843_02314\"]\n", "scene_tokens = None\n", - "scene_tokens = [\"2283aea39bc1505e\"]\n", + "# scene_tokens = [\"2283aea39bc1505e\"]\n", "log_names = None\n", "\n", "scene_filter = SceneFilter(\n", @@ -61,12 +60,82 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.common.visualization.matplotlib.plots import render_scene_animation\n", + "from typing import Tuple\n", + "from d123.common.datatypes.detection.detection import BoxDetection\n", + "from d123.common.datatypes.detection.detection_types import DYNAMIC_DETECTION_TYPES, STATIC_DETECTION_TYPES\n", + "from d123.common.geometry.base import StateSE2\n", + "from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", + "from d123.common.geometry.vector import Vector2D\n", + "from d123.common.visualization.matplotlib.observation import (\n", + " add_box_detections_to_ax,\n", + " add_default_map_on_ax,\n", + " add_ego_vehicle_to_ax,\n", + " add_traffic_lights_to_ax,\n", + ")\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", + "\n", + " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " traffic_light_detections = scene.get_traffic_light_detections_at_iteration(iteration)\n", + " map_api = scene.map_api\n", + "\n", + " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", + " add_default_map_on_ax(ax, map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_traffic_lights_to_ax(ax, traffic_light_detections, map_api)\n", + "\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + "\n", + " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", + " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", + "\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")\n", + " return ax\n", + "\n", + "\n", + "iteration = 20\n", + "scene = scenes[20]\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "_plot_scene_on_ax(ax, scene, iteration, 80)\n", + "\n", + "box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "\n", + "\n", + "def _get_dxy(box_detection: BoxDetection) -> Tuple[float, float]:\n", + " \"\"\"Get the change in x and y coordinates from the box detection.\"\"\"\n", + " center = box_detection.center.state_se2\n", + " # endpoint_x = translate_along_yaw(center, Vector2D(abs(box_detection.velocity.x), 0.0))\n", + " # endpoint_y = translate_along_yaw(center, Vector2D(0.0, abs(box_detection.velocity.y)))\n", + " # print(box_detection.velocity.x, box_detection.velocity.y)\n", + "\n", + " endpoint_x = StateSE2(center.x + box_detection.velocity.x, center.y, center.yaw)\n", + " endpoint_y = StateSE2(center.x, center.y + box_detection.velocity.y, center.yaw)\n", "\n", - "scene = scenes[0]\n", + " return endpoint_x, endpoint_y\n", "\n", "\n", - "render_scene_animation(scene, \"\", fps=10, end_idx=None, step=1, dpi=100, format=\"gif\")" + "for box_detection in box_detections:\n", + " if box_detection.metadata.detection_type in STATIC_DETECTION_TYPES:\n", + " continue\n", + " endpoint_x, endpoint_y = _get_dxy(box_detection)\n", + " ax.annotate(\n", + " \"\",\n", + " xytext=(box_detection.center.state_se2.point_2d.x, box_detection.center.state_se2.point_2d.y),\n", + " xy=(endpoint_x.x, endpoint_x.y),\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", + " ax.annotate(\n", + " \"\",\n", + " xytext=(box_detection.center.state_se2.point_2d.x, box_detection.center.state_se2.point_2d.y),\n", + " xy=(endpoint_y.x, endpoint_y.y),\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )" ] }, { diff --git a/notebooks/scene_sensor_loading.ipynb b/notebooks/scene_sensor_loading.ipynb new file mode 100644 index 00000000..ed143de1 --- /dev/null +++ b/notebooks/scene_sensor_loading.ipynb @@ -0,0 +1,123 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", + "# from d123.common.multithreading.worker_ray impo\n", + "# rt RayDistribute\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import os, psutil\n", + "\n", + "\n", + "def print_memory_usage():\n", + " process = psutil.Process(os.getpid())\n", + " memory_info = process.memory_info()\n", + " print(f\"Memory usage: {memory_info.rss / 1024 ** 2:.2f} MB\")\n", + "\n", + "\n", + "split = \"nuplan_private_test\"\n", + "\n", + "log_names = [\"2021.07.01.21.22.09_veh-14_00016_00656\"]\n", + "scene_tokens = None\n", + "log_names = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=[split],\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=15.1,\n", + " history_s=1.0,\n", + " timestamp_threshold_s=15.0,\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(\"num scenes found\", len(scenes))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.datatypes.sensor.camera import CameraType\n", + "from d123.dataset.scene.arrow_scene import ArrowScene\n", + "\n", + "scene: ArrowScene = scenes[12]\n", + "scene.open()\n", + "\n", + "camera = scene.get_camera_at_iteration(150, CameraType.CAM_B0)\n", + "\n", + "plt.imshow(camera.image)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/smarty/smart_feature_testing.ipynb b/notebooks/smarty/smart_feature_testing.ipynb index 2f35f7de..c134baf9 100644 --- a/notebooks/smarty/smart_feature_testing.ipynb +++ b/notebooks/smarty/smart_feature_testing.ipynb @@ -36,7 +36,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "\n", "\n", "log_names = None\n", diff --git a/notebooks/smarty/smart_rollout.ipynb b/notebooks/smarty/smart_rollout.ipynb index ae38d252..3f3aac67 100644 --- a/notebooks/smarty/smart_rollout.ipynb +++ b/notebooks/smarty/smart_rollout.ipynb @@ -17,7 +17,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from d123.common.multithreading.worker_sequential import Sequential" ] }, { diff --git a/notebooks/tools/merge_videos.ipynb b/notebooks/tools/merge_videos.ipynb index 9e1c9f1b..257bb8ad 100644 --- a/notebooks/tools/merge_videos.ipynb +++ b/notebooks/tools/merge_videos.ipynb @@ -89,7 +89,7 @@ " return False\n", "\n", "# List your MP4 files in the order you want them merged\n", - "video_folder = Path(\"/home/daniel/d123_logs_videos/nuplan_mini_val\")\n", + "video_folder = Path(\"/home/daniel/d123_logs_videos/wopd_train\")\n", "video_files = [str(file) for file in video_folder.glob(\"*.mp4\") if file.is_file()]\n", "\n", "\n", @@ -129,7 +129,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.3" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/notebooks/tools/plot_map_sizes.ipynb b/notebooks/tools/plot_map_sizes.ipynb index 029a63ef..4bea65c7 100644 --- a/notebooks/tools/plot_map_sizes.ipynb +++ b/notebooks/tools/plot_map_sizes.ipynb @@ -59,7 +59,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.3" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb new file mode 100644 index 00000000..1feef159 --- /dev/null +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -0,0 +1,355 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "# splits = [\"wopd_train\"]\n", + "# splits = [\"carla\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "log_names = None\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=10,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=20,\n", + " shuffle=True,\n", + " # camera_types=[CameraType.CAM_F0],\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import List, Optional, Tuple\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from d123.common.geometry.base import Point2D\n", + "from d123.common.visualization.color.color import BLACK, DARK_GREY, DARKER_GREY, LIGHT_GREY, NEW_TAB_10, TAB_10\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", + "from d123.common.visualization.matplotlib.observation import (\n", + " add_box_detections_to_ax,\n", + " add_default_map_on_ax,\n", + " add_ego_vehicle_to_ax,\n", + " add_traffic_lights_to_ax,\n", + ")\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", + "from d123.dataset.maps.abstract_map import AbstractMap\n", + "from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractLaneGroup\n", + "from d123.dataset.maps.gpkg.gpkg_map_objects import GPKGIntersection\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "\n", + "import shapely.geometry as geom\n", + "\n", + "LEFT_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=TAB_10[2],\n", + " fill_color_alpha=1.0,\n", + " line_color=TAB_10[2],\n", + " line_color_alpha=0.5,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "RIGHT_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=TAB_10[3],\n", + " fill_color_alpha=1.0,\n", + " line_color=TAB_10[3],\n", + " line_color_alpha=0.5,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=22,\n", + ")\n", + "\n", + "\n", + "LANE_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=BLACK,\n", + " fill_color_alpha=1.0,\n", + " line_color=BLACK,\n", + " line_color_alpha=0.0,\n", + " line_width=0.0,\n", + " line_style=\"-\",\n", + " zorder=5,\n", + ")\n", + "\n", + "ROAD_EDGE_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=DARKER_GREY.set_brightness(0.0),\n", + " fill_color_alpha=1.0,\n", + " line_color=DARKER_GREY.set_brightness(0.0),\n", + " line_color_alpha=1.0,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "ROAD_LINE_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=NEW_TAB_10[5],\n", + " fill_color_alpha=1.0,\n", + " line_color=NEW_TAB_10[5],\n", + " line_color_alpha=1.0,\n", + " line_width=1.5,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "\n", + "def add_debug_map_on_ax(\n", + " ax: plt.Axes,\n", + " map_api: AbstractMap,\n", + " point_2d: Point2D,\n", + " radius: float,\n", + " route_lane_group_ids: Optional[List[int]] = None,\n", + ") -> None:\n", + " layers: List[MapLayer] = [\n", + " MapLayer.LANE,\n", + " MapLayer.LANE_GROUP,\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.CARPARK,\n", + " MapLayer.CROSSWALK,\n", + " MapLayer.INTERSECTION,\n", + " MapLayer.WALKWAY,\n", + " MapLayer.ROAD_EDGE,\n", + " # MapLayer.ROAD_LINE,\n", + " ]\n", + " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", + " y_min, y_max = point_2d.y - radius, point_2d.y + radius\n", + " patch = geom.box(x_min, y_min, x_max, y_max)\n", + " map_objects_dict = map_api.query(geometry=patch, layers=layers, predicate=\"intersects\")\n", + "\n", + " done = False\n", + " for layer, map_objects in map_objects_dict.items():\n", + " for map_object in map_objects:\n", + " try:\n", + " if layer in [\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.CARPARK,\n", + " MapLayer.CROSSWALK,\n", + " MapLayer.INTERSECTION,\n", + " MapLayer.WALKWAY,\n", + " ]:\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + "\n", + " if layer in [MapLayer.LANE_GROUP]:\n", + " map_object: AbstractLaneGroup\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + "\n", + "\n", + " if layer in [MapLayer.LANE]:\n", + " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", + "\n", + " # if layer in [MapLayer.ROAD_EDGE]:\n", + " # add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", + "\n", + " # if layer in [MapLayer.ROAD_LINE]:\n", + " # # line_type = int(map_object.road_line_type)\n", + " # # plt_config = PlotConfig(\n", + " # # fill_color=NEW_TAB_10[line_type % (len(NEW_TAB_10) - 1)],\n", + " # # fill_color_alpha=1.0,\n", + " # # line_color=NEW_TAB_10[line_type % (len(NEW_TAB_10) - 1)],\n", + " # # line_color_alpha=1.0,\n", + " # # line_width=1.5,\n", + " # # line_style=\"-\",\n", + " # # zorder=10,\n", + " # # )\n", + " # add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_LINE_CONFIG)\n", + "\n", + " except Exception:\n", + " import traceback\n", + "\n", + " print(f\"Error adding map object of type {layer.name} and id {map_object.id}\")\n", + " traceback.print_exc()\n", + "\n", + " ax.set_title(f\"Map: {map_api.map_name}\")\n", + "\n", + "\n", + "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", + "\n", + " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "\n", + " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", + " # add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", + "\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + "\n", + " zoom = 1.0\n", + " ax.set_xlim(point_2d.x - radius * zoom, point_2d.x + radius * zoom)\n", + " ax.set_ylim(point_2d.y - radius * zoom, point_2d.y + radius * zoom)\n", + "\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")\n", + " return ax\n", + "\n", + "\n", + "def plot_scene_at_iteration(\n", + " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", + ") -> Tuple[plt.Figure, plt.Axes]:\n", + "\n", + " size = 10\n", + "\n", + " fig, ax = plt.subplots(figsize=(size, size))\n", + " _plot_scene_on_ax(ax, scene, iteration, radius)\n", + " return fig, ax\n", + "\n", + "\n", + "scene_index = 9\n", + "iteration = 99\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=iteration, radius=30)\n", + "plt.show()\n", + "\n", + "# camera = scenes[scene_index].get_camera_at_iteration(\n", + "# iteration=iteration, camera_type=CameraType.CAM_F0\n", + "# )\n", + "\n", + "# plt.imshow(camera.image, cmap=\"gray\", vmin=0, vmax=255)\n", + "# # fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")\n", + "\n", + "# scenes[scene_index].log_name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "scene = scenes[scene_index]\n", + "\n", + "\n", + "scene.get_camera_at_iteration(camera_type=CameraType.CAM_F0, iteration=0)\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", + "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", + "\n", + "\n", + "map_api: GPKGMap = scenes[scene_index].map_api\n", + "\n", + "drivable_polygons = map_api._gpd_dataframes[MapLayer.LANE]\n", + "\n", + "\n", + "\n", + "linear_rings = get_road_edge_linear_rings(drivable_polygons.geometry.tolist())\n", + "rings_lengths = [ring.length for ring in linear_rings]\n", + "max_length_idx = np.argmax(rings_lengths)\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "size = 16\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "for idx, ring in enumerate(linear_rings):\n", + " if idx == max_length_idx:\n", + " ax.plot(*ring.xy, color=\"black\", linewidth=2, label=\"Longest Road Edge\")\n", + " else:\n", + " ax.plot(*ring.xy)\n", + "\n", + "\n", + "ax.set_aspect(\"equal\", adjustable=\"box\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/bev_matplotlib_prediction.ipynb b/notebooks/viz/bev_matplotlib_prediction.ipynb new file mode 100644 index 00000000..041daa72 --- /dev/null +++ b/notebooks/viz/bev_matplotlib_prediction.ipynb @@ -0,0 +1,246 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "# splits = [\"wopd_train\"]\n", + "# splits = [\"carla\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "log_names = None\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=8.0,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=4.0,\n", + " shuffle=True,\n", + " # camera_types=[CameraType.CAM_F0],\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from typing import List, Optional, Tuple\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from d123.common.geometry.base import Point2D\n", + "from d123.common.visualization.color.color import BLACK, DARK_GREY, DARKER_GREY, LIGHT_GREY, NEW_TAB_10, TAB_10\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", + "from d123.common.visualization.matplotlib.observation import (\n", + " add_box_detections_to_ax,\n", + " add_box_future_detections_to_ax,\n", + " add_default_map_on_ax,\n", + " add_ego_vehicle_to_ax,\n", + " add_traffic_lights_to_ax,\n", + ")\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", + "from d123.dataset.maps.abstract_map import AbstractMap\n", + "from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractLaneGroup\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "\n", + "\n", + "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", + "\n", + " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "\n", + " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", + " # add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radiuss, route_lane_group_ids=None)\n", + " add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", + "\n", + " add_box_future_detections_to_ax(ax, scene, iteration=iteration)\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + "\n", + " zoom = 1.0\n", + " ax.set_xlim(point_2d.x - radius * zoom, point_2d.x + radius * zoom)\n", + " ax.set_ylim(point_2d.y - radius * zoom, point_2d.y + radius * zoom)\n", + "\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")\n", + " return ax\n", + "\n", + "\n", + "def plot_scene_at_iteration(\n", + " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", + ") -> Tuple[plt.Figure, plt.Axes]:\n", + "\n", + " size = 10\n", + "\n", + " fig, ax = plt.subplots(figsize=(size, size))\n", + " _plot_scene_on_ax(ax, scene, iteration, radius)\n", + " return fig, ax\n", + "\n", + "\n", + "scene_index = 0\n", + "\n", + "for i, scene in enumerate(scenes):\n", + "\n", + " iteration = 0\n", + " fig, ax = plot_scene_at_iteration(scenes[i], iteration=iteration, radius=60)\n", + " Path(f\"/home/daniel/examples/{splits[0]}/\").mkdir(exist_ok=True, parents=True)\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + " fig.tight_layout()\n", + " fig.savefig(\n", + " f\"/home/daniel/examples/{splits[0]}/{scene.log_name}_{i}.pdf\",\n", + " dpi=300,\n", + " bbox_inches=\"tight\",\n", + " )\n", + "\n", + "# camera = scenes[scene_index].get_camera_at_iteration(\n", + "# iteration=iteration, camera_type=CameraType.CAM_F0\n", + "# )\n", + "\n", + "# plt.imshow(camera.image, cmap=\"gray\", vmin=0, vmax=255)\n", + "# # fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")\n", + "\n", + "# scenes[scene_index].log_name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "scene = scenes[scene_index]\n", + "\n", + "\n", + "scene.get_camera_at_iteration(camera_type=CameraType.CAM_F0, iteration=0)\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", + "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", + "\n", + "\n", + "map_api: GPKGMap = scenes[scene_index].map_api\n", + "\n", + "drivable_polygons = map_api._gpd_dataframes[MapLayer.LANE]\n", + "\n", + "\n", + "\n", + "linear_rings = get_road_edge_linear_rings(drivable_polygons.geometry.tolist())\n", + "rings_lengths = [ring.length for ring in linear_rings]\n", + "max_length_idx = np.argmax(rings_lengths)\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "size = 16\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "for idx, ring in enumerate(linear_rings):\n", + " if idx == max_length_idx:\n", + " ax.plot(*ring.xy, color=\"black\", linewidth=2, label=\"Longest Road Edge\")\n", + " else:\n", + " ax.plot(*ring.xy)\n", + "\n", + "\n", + "ax.set_aspect(\"equal\", adjustable=\"box\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb new file mode 100644 index 00000000..bbd6221d --- /dev/null +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -0,0 +1,397 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Tuple\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", + "\n", + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "from typing import Dict\n", + "from d123.common.datatypes.sensor.camera import CameraType\n", + "from d123.common.visualization.matplotlib.camera import add_camera_ax\n", + "from d123.common.visualization.matplotlib.camera import add_box_detections_to_camera_ax" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "\n", + "\n", + "# splits = [\"carla\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "# splits = [\"wopd_train\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "\n", + "# splits = [\"nuplan_private_test\"]\n", + "log_names = None\n", + "\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=15,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=15,\n", + " shuffle=False,\n", + " camera_types=[CameraType.CAM_F0],\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "\n", + "def plot_cameras_frame(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 2\n", + " fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L2: ax[1, 0],\n", + " CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R2: ax[1, 2],\n", + " }\n", + "\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_camera_ax(camera_ax, camera)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + " return fig, ax\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L2: ax[1, 0],\n", + " CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R2: ax[1, 2],\n", + " }\n", + "\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " if camera_type not in scene.available_camera_types:\n", + " continue\n", + " # assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + " return ax\n", + "\n", + "\n", + "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 1\n", + " fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", + " _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "plot_cameras_with_detections(scenes[9], iteration=20)\n", + "\n", + "\n", + "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", + "\n", + "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + "# for camera_type, camera_ax in camera_ax_assignment.items():\n", + "# assert camera_type in scene.available_camera_types\n", + "# camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + "# if camera is not None:\n", + "# add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + "# return ax\n", + "\n", + "\n", + "# def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "# \"\"\"\n", + "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + "# :param scene: navsim scene dataclass\n", + "# :param frame_idx: index of selected frame\n", + "# :return: figure and ax object of matplotlib\n", + "# \"\"\"\n", + "# scale = 2\n", + "# fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", + "# _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + "# fig.tight_layout()\n", + "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + "# CameraType.CAM_L1: ax[0],\n", + "# CameraType.CAM_L0: ax[1],\n", + "# CameraType.F0: ax[2],\n", + "# CameraType.CAM_R0: ax[3],\n", + "# CameraType.CAM_R1: ax[4],\n", + "# }\n", + "\n", + "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + "# for camera_type, camera_ax in camera_ax_assignment.items():\n", + "# assert camera_type in scene.available_camera_types\n", + "# camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + "# if camera is not None:\n", + "# add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + "# return ax\n", + "\n", + "\n", + "# def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "# \"\"\"\n", + "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + "# :param scene: navsim scene dataclass\n", + "# :param frame_idx: index of selected frame\n", + "# :return: figure and ax object of matplotlib\n", + "# \"\"\"\n", + "# scale = 2\n", + "# fig, ax = plt.subplots(1, 5, figsize=(scale * 11.5, scale * 4.3))\n", + "# _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + "# fig.tight_layout()\n", + "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "# plot_cameras_with_detections(scenes[0], iteration=0)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L1: ax[1, 0],\n", + " # CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R1: ax[1, 2],\n", + " CameraType.CAM_L2: ax[2, 0],\n", + " CameraType.CAM_R2: ax[2, 2],\n", + " }\n", + "\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " if camera_type not in scene.available_camera_types:\n", + " continue\n", + " # assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + " return ax\n", + "\n", + "\n", + "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 1\n", + " fig, ax = plt.subplots(3, 3, figsize=(scale * 10, scale * 10))\n", + " _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "plot_cameras_with_detections(scenes[9], iteration=20)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "scenes[3].log_name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from typing import Optional, Union\n", + "\n", + "from matplotlib import animation\n", + "from tqdm import tqdm\n", + "\n", + "\n", + "def render_scene_animation(\n", + " scene: AbstractScene,\n", + " output_path: Union[str, Path],\n", + " start_idx: int = 0,\n", + " end_idx: Optional[int] = None,\n", + " step: int = 10,\n", + " fps: float = 20.0,\n", + " dpi: int = 300,\n", + " format: str = \"mp4\",\n", + ") -> None:\n", + " assert format in [\"mp4\", \"gif\"], \"Format must be either 'mp4' or 'gif'.\"\n", + " output_path = Path(output_path)\n", + " output_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + " scene.open()\n", + "\n", + " if end_idx is None:\n", + " end_idx = scene.get_number_of_iterations()\n", + " end_idx = min(end_idx, scene.get_number_of_iterations())\n", + "\n", + " scale = 1\n", + " fig, ax = plt.subplots(1, 1, figsize=(scale * 16, scale * 9))\n", + "\n", + " def update(i):\n", + " ax.clear()\n", + " # for a in ax.flatten():\n", + " # a.clear()\n", + " _add_camera_with_detections(ax, scene, i)\n", + " plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)\n", + " pbar.update(1)\n", + "\n", + " frames = list(range(start_idx, end_idx, step))\n", + " pbar = tqdm(total=len(frames), desc=f\"Rendering {scene.log_name} as {format}\")\n", + " ani = animation.FuncAnimation(fig, update, frames=frames, repeat=False)\n", + "\n", + " ani.save(output_path / f\"{scene.log_name}_{scene.token}.{format}\", writer=\"ffmpeg\", fps=fps, dpi=dpi)\n", + " plt.close(fig)\n", + " scene.close()\n", + "\n", + "\n", + "# for i in range(10):\n", + " # render_scene_animation(scenes[i], output_path=\"output\", start_idx=0, end_idx=None, step=1, fps=20, dpi=300, format=\"mp4\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/log_rendering.ipynb b/notebooks/viz/log_rendering.ipynb index 908d1825..fe88e006 100644 --- a/notebooks/viz/log_rendering.ipynb +++ b/notebooks/viz/log_rendering.ipynb @@ -13,8 +13,8 @@ "\n", "\n", "\n", - "log_name = \"2021.07.01.20.35.47_veh-38_00016_00281\"\n", - "log_file = Path(f\"/home/daniel/d123_workspace/data/nuplan_private_test/{log_name}.arrow\")\n", + "log_name = \"1005081002024129653_5313_150_5333_150\"\n", + "log_file = Path(f\"/home/daniel/d123_workspace/data/wopd_train/{log_name}.arrow\")\n", "scene = ArrowScene(log_file)\n", "\n", "fig, ax = plot_scene_at_iteration(scene, iteration=10)\n", @@ -47,19 +47,19 @@ "import traceback\n", "from d123.common.visualization.matplotlib.plots import render_scene_animation\n", "\n", - "\n", - "output_path = Path(\"/home/daniel/d123_logs_videos\")\n", - "log_path = Path(\"/home/daniel/d123_workspace/data/carla\")\n", + "split = \"av2-sensor-mini_train\"\n", + "output_path = Path(f\"/home/daniel/d123_logs_videos/{split}\")\n", + "log_path = Path(f\"/home/daniel/d123_workspace/data/{split}\")\n", "for log_file in log_path.iterdir():\n", - " scene = ArrowScene(log_file)\n", " try:\n", - " render_scene_animation(scene, output_path, fps=30, end_idx=None, step=5)\n", + " scene = ArrowScene(log_file)\n", + " render_scene_animation(scene, output_path, fps=20, end_idx=None, step=1)\n", + " del scene\n", " except Exception as e:\n", " traceback.print_exc()\n", " erroneous_file = output_path / f\"{log_name}.mp4\"\n", " if erroneous_file.exists():\n", " erroneous_file.unlink()\n", - " del scene\n", " # break\n", " # break" ] @@ -83,7 +83,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123", + "display_name": "d123_dev", "language": "python", "name": "python3" }, diff --git a/notebooks/viz/viser_testing_v2.ipynb b/notebooks/viz/viser_testing_v2.ipynb deleted file mode 100644 index 60b88c98..00000000 --- a/notebooks/viz/viser_testing_v2.ipynb +++ /dev/null @@ -1,87 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.common.visualization.viser.server import ViserVisualizationServer\n", - "\n", - "from pathlib import Path\n", - "from d123.dataset.scene.arrow_scene import ArrowScene" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# log_name = \"Town12_Rep0_longest6_route17_06_03_21_40_12\" # Town3\n", - "log_name = \"2021.09.29.17.35.58_veh-44_01671_01819\" # Town4\n", - "log_file = Path(f\"/home/daniel/d123_workspace/data/nuplan_private_test/{log_name}.arrow\")\n", - "\n", - "# log_name = \"_Rep0_longest1_route0_07_04_10_18_47\"\n", - "# log_file = Path(f\"/home/daniel/d123_workspace/data/carla/{log_name}.arrow\")\n", - "scene = ArrowScene(log_file)\n", - "\n", - "visualization_server = ViserVisualizationServer()\n", - "\n", - "visualization_server.set_scene(scene)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "ego_vehicle_state = scene.get_ego_state_at_iteration(0)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [ - "ego_vehicle_state.vehicle_parameters.height/2 - ego_vehicle_state.vehicle_parameters.height/3\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "d123", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 1bacd558..07d91d1d 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -7,14 +7,11 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.common.visualization.viser.server import ViserVisualizationServer\n", - "\n", - "from pathlib import Path\n", - "from d123.dataset.scene.arrow_scene import ArrowScene\n", "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from d123.common.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" ] }, { @@ -24,27 +21,36 @@ "metadata": {}, "outputs": [], "source": [ - "split = \"carla\"\n", + "import requests\n", + "from PIL import Image\n", + "from io import BytesIO\n", + "\n", + "\n", + "\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"carla\"]\n", + "# splits = [\"wopd_train\"]\n", + "# splits = [\"av2-sensor-mini_train\"]\n", + "log_names = None\n", "\n", - "log_names = [\"_Rep0_routes_devtest_route1_07_10_18_45_18\"]\n", "scene_tokens = None\n", - "# scene_tokens = [\"2283aea39bc1505e\"]\n", - "# log_names = None\n", "\n", "scene_filter = SceneFilter(\n", - " split_names=[split],\n", + " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=20.1,\n", - " history_s=1.0,\n", - " timestamp_threshold_s=0.1,\n", + " duration_s=10,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=10,\n", + " shuffle=False,\n", + " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", "# worker = RayDistributed()\n", "scenes = scene_builder.get_scenes(scene_filter, worker)\n", "\n", - "print(len(scenes))\n" + "print(f\"Found {len(scenes)} scenes\")" ] }, { @@ -53,11 +59,7 @@ "id": "2", "metadata": {}, "outputs": [], - "source": [ - "\n", - "visualization_server = ViserVisualizationServer()\n", - "visualization_server.set_scene(scenes[180])" - ] + "source": [] }, { "cell_type": "code", @@ -66,19 +68,17 @@ "metadata": {}, "outputs": [], "source": [ - "# ego_vehicle_state = scene.get_ego_vehicle_state_at_iteration(0)" + "from d123.common.visualization.viser.server import ViserVisualizationServer\n", + "\n", + "\n", + "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" ] }, { - "cell_type": "code", - "execution_count": null, + "cell_type": "markdown", "id": "4", "metadata": {}, - "outputs": [], - "source": [ - "# ego_vehicle_state.vehicle_parameters.height/2 - ego_vehicle_state.vehicle_parameters.height/3\n", - "\n" - ] + "source": [] }, { "cell_type": "code", @@ -86,24 +86,7 @@ "id": "5", "metadata": {}, "outputs": [], - "source": [ - "# from pyquaternion import Quaternion\n", - "# import numpy as np\n", - "# # Example Euler angles (in radians) and rotation sequence\n", - "# # roll = 0.5 # Rotation about the X-axis (radians)\n", - "# # pitch = 1.0 # Rotation about the Y-axis (radians)\n", - "# # yaw = 1.5 # Rotation about the Z-axis (radians)\n", - "# roll = 1.0 # Rotation about the X-axis (radians)\n", - "# pitch = 2.0 # Rotation about the Y-axis (radians)\n", - "# yaw = 1.0 # Rotation about the Z-axis (radians)\n", - "# rotation_sequence = 'xyz' # Or 'yxz', 'zxy', etc.\n", - "\n", - "# # Create a quaternion from Euler angles\n", - "# quaternion = Quaternion(axis=[roll, pitch, yaw], sequence=rotation_sequence)\n", - "\n", - "\n", - "# quaternion.yaw_pitch_roll" - ] + "source": [] }, { "cell_type": "code", @@ -111,25 +94,7 @@ "id": "6", "metadata": {}, "outputs": [], - "source": [ - "# from d123.common.datatypes.camera.camera_parameters import get_nuplan_camera_parameters\n", - "# import numpy as np\n", - "\n", - "# fy = get_nuplan_camera_parameters().intrinsics[1, 1]\n", - "\n", - "# fov_v = 2 * np.arctan(1080 / (2 * fy))\n", - "# fov_v" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [ - "# get_nuplan_camera_parameters().intrinsics, fy" - ] + "source": [] } ], "metadata": { diff --git a/notebooks/waymo_perception/lidar_testing.ipynb b/notebooks/waymo_perception/lidar_testing.ipynb new file mode 100644 index 00000000..3cd2a77e --- /dev/null +++ b/notebooks/waymo_perception/lidar_testing.ipynb @@ -0,0 +1,301 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset import dataset_pb2\n", + "\n", + "import json\n", + "import os\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "from PIL import Image\n", + "from tqdm import tqdm\n", + "from waymo_open_dataset import label_pb2\n", + "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", + "from waymo_open_dataset.utils import box_utils\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "\n", + "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", + "\n", + "\n", + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from pyquaternion import Quaternion\n", + "\n", + "from d123.common.geometry.base import StateSE3\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", + "\n", + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "\n", + "# Frame attributes:\n", + "# context: \n", + "# timestamp_micros: \n", + "# pose: \n", + "# images: List with 5 images\n", + "# lasers: \n", + "# Length: 5\n", + "# laser_labels: \n", + "# Length: 0\n", + "# projected_lidar_labels: \n", + "# Length: 0\n", + "# camera_labels: \n", + "# Length: 0\n", + "# no_label_zones: \n", + "# Length: 0\n", + "# map_features: \n", + "# Length: 0\n", + "# map_pose_offset: \n", + "\n", + "file_idx = 0\n", + "pathname = tfrecords_file_list[file_idx]\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "num_frames = sum(1 for _ in dataset)\n", + "\n", + "\n", + "def read_jpg_image(data: bytes) -> np.ndarray:\n", + " \"\"\"Read a JPEG image from bytes and return it as a numpy array.\"\"\"\n", + " image = Image.open(io.BytesIO(data))\n", + " return np.array(image)\n", + "\n", + "\n", + "ego_state_se3s = []\n", + "front_images = []\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "\n", + "boxes = []\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + "\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " # print(frame.camera_labels)\n", + " print(\"Frame attributes:\")\n", + " for field in frame.DESCRIPTOR.fields:\n", + " field_name = field.name\n", + " if hasattr(frame, field_name):\n", + " value = getattr(frame, field_name)\n", + " if field_name != \"images\": # Don't print the whole image data\n", + " print(f\" {field_name}: {type(value)}\")\n", + " if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", + " print(f\" Length: {len(value)}\")\n", + " else:\n", + " print(f\" {field_name}: List with {len(value)} images\")\n", + "\n", + "\n", + " # # 1. pose\n", + " pose = np.array(frame.pose.transform).reshape(4, 4)\n", + " yaw_pitch_roll = Quaternion(matrix=pose[:3, :3]).yaw_pitch_roll\n", + " ego_state_se3s.append(\n", + " np.array(\n", + " [\n", + " pose[0, 3], # x\n", + " pose[1, 3], # y\n", + " pose[2, 3], # z\n", + " yaw_pitch_roll[2], # yaw\n", + " yaw_pitch_roll[1], # pitch\n", + " yaw_pitch_roll[0], # roll\n", + " ],\n", + " dtype=np.float64,\n", + " )\n", + " )\n", + "\n", + " # # plt.show()\n", + " if frame_idx == 0:\n", + " break\n", + "\n", + "ego_state_se3s = np.array(ego_state_se3s, dtype=np.float64)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Dict, List, Tuple\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "\n", + "from waymo_open_dataset import dataset_pb2\n", + "from waymo_open_dataset.utils import range_image_utils\n", + "from waymo_open_dataset.utils import transform_utils\n", + "\n", + "RangeImages = Dict[\"dataset_pb2.LaserName.Name\", List[dataset_pb2.MatrixFloat]]\n", + "CameraProjections = Dict[\"dataset_pb2.LaserName.Name\", List[dataset_pb2.MatrixInt32]]\n", + "SegmentationLabels = Dict[\"dataset_pb2.LaserName.Name\", List[dataset_pb2.MatrixInt32]]\n", + "ParsedFrame = Tuple[RangeImages, CameraProjections, SegmentationLabels, dataset_pb2.MatrixFloat]\n", + "\n", + "\n", + "def parse_range_image_and_camera_projection(frame: dataset_pb2.Frame) -> ParsedFrame:\n", + " \"\"\"Parse range images and camera projections given a frame.\n", + "\n", + " Args:\n", + " frame: open dataset frame proto\n", + "\n", + " Returns:\n", + " range_images: A dict of {laser_name,\n", + " [range_image_first_return, range_image_second_return]}.\n", + " camera_projections: A dict of {laser_name,\n", + " [camera_projection_from_first_return,\n", + " camera_projection_from_second_return]}.\n", + " seg_labels: segmentation labels, a dict of {laser_name,\n", + " [seg_label_first_return, seg_label_second_return]}\n", + " range_image_top_pose: range image pixel pose for top lidar.\n", + " \"\"\"\n", + " range_images = {}\n", + " camera_projections = {}\n", + " seg_labels = {}\n", + " range_image_top_pose: dataset_pb2.MatrixFloat = dataset_pb2.MatrixFloat()\n", + " for laser in frame.lasers:\n", + " if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " range_image_str_tensor = tf.io.decode_compressed(laser.ri_return1.range_image_compressed, \"ZLIB\")\n", + " ri = dataset_pb2.MatrixFloat()\n", + " ri.ParseFromString(range_image_str_tensor.numpy())\n", + " range_images[laser.name] = [ri]\n", + "\n", + " if laser.name == dataset_pb2.LaserName.TOP:\n", + " range_image_top_pose_str_tensor = tf.io.decode_compressed(\n", + " laser.ri_return1.range_image_pose_compressed, \"ZLIB\"\n", + " )\n", + " range_image_top_pose = dataset_pb2.MatrixFloat()\n", + " range_image_top_pose.ParseFromString(range_image_top_pose_str_tensor.numpy())\n", + "\n", + " camera_projection_str_tensor = tf.io.decode_compressed(\n", + " laser.ri_return1.camera_projection_compressed, \"ZLIB\"\n", + " )\n", + " cp = dataset_pb2.MatrixInt32()\n", + " cp.ParseFromString(camera_projection_str_tensor.numpy())\n", + " camera_projections[laser.name] = [cp]\n", + "\n", + " if len(laser.ri_return1.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return1.segmentation_label_compressed, \"ZLIB\")\n", + " seg_label = dataset_pb2.MatrixInt32()\n", + " seg_label.ParseFromString(seg_label_str_tensor.numpy())\n", + " seg_labels[laser.name] = [seg_label]\n", + " if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " range_image_str_tensor = tf.io.decode_compressed(laser.ri_return2.range_image_compressed, \"ZLIB\")\n", + " ri = dataset_pb2.MatrixFloat()\n", + " ri.ParseFromString(range_image_str_tensor.numpy())\n", + " range_images[laser.name].append(ri)\n", + "\n", + " camera_projection_str_tensor = tf.io.decode_compressed(\n", + " laser.ri_return2.camera_projection_compressed, \"ZLIB\"\n", + " )\n", + " cp = dataset_pb2.MatrixInt32()\n", + " cp.ParseFromString(camera_projection_str_tensor.numpy())\n", + " camera_projections[laser.name].append(cp)\n", + "\n", + " if len(laser.ri_return2.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return2.segmentation_label_compressed, \"ZLIB\")\n", + " seg_label = dataset_pb2.MatrixInt32()\n", + " seg_label.ParseFromString(seg_label_str_tensor.numpy())\n", + " seg_labels[laser.name].append(seg_label)\n", + " return range_images, camera_projections, seg_labels, range_image_top_pose\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "dataset = tf.data.TFRecordDataset(tfrecords_file_list[file_idx], compression_type=\"\")\n", + "for data in dataset:\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy()) # No need for bytearray conversion\n", + " break\n", + "\n", + "(range_images, camera_projections, _, range_image_top_pose) = parse_range_image_and_camera_projection(frame)\n", + "points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n", + " frame=frame,\n", + " range_images=range_images,\n", + " camera_projections=camera_projections,\n", + " range_image_top_pose=range_image_top_pose,\n", + " keep_polar_features=True,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "pc = points[0]\n", + "\n", + "\n", + "plt.scatter(pc[:, 3], pc[:, 4], s=0.1, c=pc[:, 5], cmap='viridis')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "cp_points[0].shape, points[0].shape" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/waymo_perception/map_testing.ipynb b/notebooks/waymo_perception/map_testing.ipynb new file mode 100644 index 00000000..c29fc212 --- /dev/null +++ b/notebooks/waymo_perception/map_testing.ipynb @@ -0,0 +1,725 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset import dataset_pb2\n", + "\n", + "import json\n", + "import os\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "from PIL import Image\n", + "from tqdm import tqdm\n", + "from waymo_open_dataset import label_pb2\n", + "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", + "from waymo_open_dataset.utils import box_utils\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "\n", + "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", + "\n", + "\n", + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from pyquaternion import Quaternion\n", + "\n", + "from d123.common.geometry.base import StateSE3\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", + "\n", + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "\n", + "# Frame attributes:\n", + "# context: \n", + "# timestamp_micros: \n", + "# pose: \n", + "# images: List with 5 images\n", + "# lasers: \n", + "# Length: 5\n", + "# laser_labels: \n", + "# Length: 0\n", + "# projected_lidar_labels: \n", + "# Length: 0\n", + "# camera_labels: \n", + "# Length: 0\n", + "# no_label_zones: \n", + "# Length: 0\n", + "# map_features: \n", + "# Length: 0\n", + "# map_pose_offset: \n", + "\n", + "file_idx = 0\n", + "pathname = tfrecords_file_list[file_idx]\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "num_frames = sum(1 for _ in dataset)\n", + "\n", + "\n", + "def read_jpg_image(data: bytes) -> np.ndarray:\n", + " \"\"\"Read a JPEG image from bytes and return it as a numpy array.\"\"\"\n", + " image = Image.open(io.BytesIO(data))\n", + " return np.array(image)\n", + "\n", + "\n", + "ego_state_se3s = []\n", + "front_images = []\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "\n", + "boxes = []\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + "\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " print(\"Frame attributes:\")\n", + " for field in frame.DESCRIPTOR.fields:\n", + " field_name = field.name\n", + " if hasattr(frame, field_name):\n", + " value = getattr(frame, field_name)\n", + " if field_name != \"images\": # Don't print the whole image data\n", + " print(f\" {field_name}: {type(value)}\")\n", + " if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", + " print(f\" Length: {len(value)}\")\n", + " else:\n", + " print(f\" {field_name}: List with {len(value)} images\")\n", + "\n", + " # # plt.show()\n", + " if frame_idx == 0:\n", + " break\n", + "\n", + "ego_state_se3s = np.array(ego_state_se3s, dtype=np.float64)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "frame.context.stats.location" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "\n", + "from d123.common.geometry.units import mph_to_mps\n", + "\n", + "\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "for frame_idx, data in enumerate(dataset):\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " break\n", + "\n", + "driveways = {}\n", + "stop_signs = {}\n", + "crosswalks = {}\n", + "road_edges = {}\n", + "road_lines = {}\n", + "\n", + "lanes = {}\n", + "lanes_type = {}\n", + "lanes_speed_limit_mps = {}\n", + "lanes_interpolate = {}\n", + "lanes_successors = defaultdict(list)\n", + "lanes_predecessors = defaultdict(list)\n", + "\n", + "lanes_map_features = {}\n", + "\n", + "for map_feature in frame.map_features:\n", + "\n", + " if map_feature.HasField(\"lane\"):\n", + " lanes[map_feature.id] = np.array([[p.x, p.y, p.z] for p in map_feature.lane.polyline])\n", + " lanes_map_features[map_feature.id] = map_feature\n", + " lanes_type[map_feature.id] = map_feature.lane.type\n", + " lanes_speed_limit_mps[map_feature.id] = mph_to_mps(map_feature.lane.speed_limit_mph)\n", + " for lane_id_ in map_feature.lane.exit_lanes:\n", + " lanes_successors[map_feature.id].append(lane_id_)\n", + " for lane_id_ in map_feature.lane.exit_lanes:\n", + " lanes_predecessors[map_feature.id].append(lane_id_)\n", + "\n", + " elif map_feature.HasField(\"road_line\"):\n", + " road_lines[map_feature.id] = np.array([[p.x, p.y, p.z] for p in map_feature.road_line.polyline])\n", + " elif map_feature.HasField(\"road_edge\"):\n", + " road_edges[map_feature.id] = np.array([[p.x, p.y, p.z] for p in map_feature.road_edge.polyline])\n", + " elif map_feature.HasField(\"stop_sign\"):\n", + " pass\n", + " elif map_feature.HasField(\"crosswalk\"):\n", + " outline = np.array([[p.x, p.y, p.z] for p in map_feature.crosswalk.polygon])\n", + " crosswalks[map_feature.id] = outline\n", + " elif map_feature.HasField(\"speed_bump\"):\n", + " pass\n", + " elif map_feature.HasField(\"driveway\"):\n", + " # print(map_feature.driveway)\n", + " outline = np.array([[p.x, p.y, p.z] for p in map_feature.driveway.polygon])\n", + " driveways[map_feature.id] = outline\n", + "\n", + " # print(f\"Roadline: {map_feature.road_line}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", + "\n", + "\n", + "fig, ax = plt.subplots(figsize=(30, 30))\n", + "\n", + "for road_edge in road_edges.values():\n", + " # print(len(driveway))\n", + " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"blue\", label=\"road_edge\")\n", + "\n", + "for lane in lanes.values():\n", + " # print(len(driveway))\n", + " ax.plot(lane[:, 0], lane[:, 1], color=\"gray\", label=\"lane\")\n", + "\n", + "\n", + "for road_line in road_lines.values():\n", + " # print(len(driveway))\n", + " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "for driveway in driveways.values():\n", + " # print(len(driveway))\n", + " ax.plot(driveway[:, 0], driveway[:, 1], color=\"green\", label=\"driveway\")\n", + "\n", + "\n", + "for crosswalk in crosswalks.values():\n", + " # print(len(driveway))\n", + " ax.plot(crosswalk[:, 0], crosswalk[:, 1], color=\"violet\", label=\"crosswalk\")\n", + "\n", + "\n", + "add_non_repeating_legend_to_ax(ax)\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.dataset_specific.wopd.wopd_map_utils import extract_lane_boundaries\n", + "\n", + "\n", + "left_boundaries, right_boundaries = extract_lane_boundaries(\n", + " lanes=lanes,\n", + " lanes_successors=lanes_successors,\n", + " lanes_predecessors=lanes_predecessors,\n", + " road_lines=road_lines,\n", + " road_edges=road_edges,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "size = 20\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "for lane_id, lane_ in lanes.items():\n", + " if lane_id not in left_boundaries or lane_id not in right_boundaries:\n", + " continue\n", + " left_boundary = left_boundaries[lane_id].array\n", + " right_boundary = right_boundaries[lane_id].array\n", + "\n", + " assert len(left_boundary) > 0 and len(right_boundary) > 0\n", + " ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", + " ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "\n", + " ax.plot(lane_[:, 0], lane_[:, 1], color=\"black\")\n", + "\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "# " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "lane_idx = 158\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.plot(lanes[lane_id][:, 0], lanes[lane_id][:, 1], color=\"black\", label=\"Lane Centerline\", alpha=0.5)\n", + "\n", + "for left_boundary in lane.left_boundaries:\n", + " print(\"left\", left_boundary)\n", + " try:\n", + " boundary = road_lines[left_boundary.boundary_feature_id]\n", + "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {left_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[left_boundary.boundary_feature_id]\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 0],\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 1],\n", + " color=\"lime\",\n", + " linestyle=\"-\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "for right_boundary in lane.right_boundaries:\n", + " print(\"right\", right_boundary)\n", + " try:\n", + " boundary = road_lines[right_boundary.boundary_feature_id]\n", + "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {right_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[right_boundary.boundary_feature_id]\n", + "\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 0],\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 1],\n", + " color=\"red\",\n", + " linestyle=\"-\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes\n", + "\n", + "zoom_out = 10\n", + "lane = lanes[lane_id]\n", + "ax.set_xlim(lane[:, 0].min() - zoom_out, lane[:, 0].max() + zoom_out)\n", + "ax.set_ylim(lane[:, 1].min() - zoom_out, lane[:, 1].max() + zoom_out)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.geometry.line.polylines import Polyline3D\n", + "import numpy as np\n", + "\n", + "\n", + "polyline_3d = Polyline3D.from_array(lanes[0])\n", + "\n", + "\n", + "# Create left and right boundaries for the lane\n", + "\n", + "\n", + "def create_lane_boundaries(polyline_3d, width=2.0):\n", + " \"\"\"\n", + " Create left and right boundaries for a lane by offsetting perpendicular to the heading.\n", + "\n", + " Args:\n", + " polyline_3d: Polyline3D object representing the centerline\n", + " width: Total width of the lane in meters\n", + "\n", + " Returns:\n", + " left_boundary: numpy array representing the left boundary\n", + " right_boundary: numpy array representing the right boundary\n", + " \"\"\"\n", + " points = polyline_3d.array\n", + " half_width = width / 2.0\n", + "\n", + " # Calculate the direction vectors between consecutive points\n", + " directions = np.diff(points, axis=0)\n", + "\n", + " # Normalize the direction vectors\n", + " directions_norm = np.linalg.norm(directions, axis=1, keepdims=True)\n", + " directions_normalized = directions / directions_norm\n", + "\n", + " # Calculate perpendicular vectors in the xy plane (z remains 0)\n", + " perpendiculars = np.zeros_like(directions)\n", + " perpendiculars[:, 0] = -directions_normalized[:, 1] # -dy\n", + " perpendiculars[:, 1] = directions_normalized[:, 0] # dx\n", + "\n", + " # Create boundaries (need to handle the last point separately)\n", + " left_boundary = points[:-1] + perpendiculars * half_width\n", + " right_boundary = points[:-1] - perpendiculars * half_width\n", + "\n", + " # Handle the last point based on the last direction\n", + " last_perp = perpendiculars[-1]\n", + " left_boundary = np.vstack([left_boundary, points[-1] + last_perp * half_width])\n", + " right_boundary = np.vstack([right_boundary, points[-1] - last_perp * half_width])\n", + "\n", + " return left_boundary, right_boundary\n", + "\n", + "\n", + "# Create the boundaries with a 4m wide lane\n", + "left_boundary, right_boundary = create_lane_boundaries(polyline_3d, width=4.0)\n", + "\n", + "# Plot the centerline and the boundaries\n", + "plt.figure(figsize=(12, 8))\n", + "plt.plot(polyline_3d.array[:, 0], polyline_3d.array[:, 1], \"g-\", label=\"Centerline\")\n", + "plt.plot(left_boundary[:, 0], left_boundary[:, 1], \"b-\", label=\"Left Boundary\")\n", + "plt.plot(right_boundary[:, 0], right_boundary[:, 1], \"r-\", label=\"Right Boundary\")\n", + "plt.axis(\"equal\")\n", + "plt.grid(True)\n", + "plt.legend()\n", + "plt.title(\"Lane with Left and Right Boundaries\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "lane_idx = 59\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "# for right_neighbor in lane.right_neighbors:\n", + "# print(right_neighbor)\n", + "\n", + "# lane\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from torch import le\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.plot(lanes[lane_id][:, 0], lanes[lane_id][:, 1], color=\"blue\", label=\"Lane Centerline\")\n", + "\n", + "for right_neighbor in lane.right_neighbors:\n", + " # print(\"right\", right_neighbor)\n", + " ax.plot(\n", + " lanes[right_neighbor.feature_id][:, 0],\n", + " lanes[right_neighbor.feature_id][:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[right_neighbor.feature_id][right_neighbor.neighbor_start_index: right_neighbor.neighbor_end_index, 0],\n", + " lanes[right_neighbor.feature_id][right_neighbor.neighbor_start_index: right_neighbor.neighbor_end_index, 1],\n", + " color=\"red\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "for left_neighbor in lane.left_neighbors:\n", + " # print(\"left\", left_neighbor)\n", + " ax.plot(\n", + " lanes[left_neighbor.feature_id][:, 0],\n", + " lanes[left_neighbor.feature_id][:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[left_neighbor.feature_id][left_neighbor.neighbor_start_index: left_neighbor.neighbor_end_index, 0],\n", + " lanes[left_neighbor.feature_id][left_neighbor.neighbor_start_index: left_neighbor.neighbor_end_index, 1],\n", + " color=\"lime\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "lane_idx = 48\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.scatter(lanes[lane_id][0, 0], lanes[lane_id][0, 1], color=\"black\", label=\"Lane Centerline\", alpha=0.5)\n", + "ax.scatter(lanes[lane_id][-1, 0], lanes[lane_id][-1, 1], color=\"blue\", label=\"Lane Centerline\", alpha=0.5)\n", + "\n", + "for entry_lane in lane.entry_lanes:\n", + " # print(\"right\", right_neighbor)\n", + " ax.plot(\n", + " lanes[entry_lane][:, 0],\n", + " lanes[entry_lane][:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "for exit_lane in lane.exit_lanes:\n", + " # print(\"left\", left_neighbor)\n", + " ax.plot(\n", + " lanes[exit_lane][:, 0],\n", + " lanes[exit_lane][:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [ + "lane_idx = 196\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.plot(lanes[lane_id][:, 0], lanes[lane_id][:, 1], color=\"black\", label=\"Lane Centerline\", alpha=0.5)\n", + "\n", + "for left_boundary in lane.left_boundaries:\n", + " print(\"left\", left_boundary)\n", + " try:\n", + " boundary = road_lines[left_boundary.boundary_feature_id]\n", + "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {left_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[left_boundary.boundary_feature_id]\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 0],\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 1],\n", + " color=\"lime\",\n", + " linestyle=\"-\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "for right_boundary in lane.right_boundaries:\n", + " print(\"right\", right_boundary)\n", + " try:\n", + " boundary = road_lines[right_boundary.boundary_feature_id]\n", + "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {right_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[right_boundary.boundary_feature_id]\n", + "\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 0],\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 1],\n", + " color=\"red\",\n", + " linestyle=\"-\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.geometry.base import StateSE2\n", + "# from d123.common.geometry.line.polylines import PolylineSE2\n", + "# from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", + "# from d123.common.geometry.vector import Vector2D\n", + "\n", + "# size = 30\n", + "# fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "# lane_id = 114\n", + "# lane_ = lanes[lane_id]\n", + "\n", + "# lane_polyline_se2 = PolylineSE2.from_array(lane_[:, :2])\n", + "# BOUNDARY_STEP_SIZE = 0.5\n", + "\n", + "\n", + "# distances_se2 = np.linspace(\n", + "# 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True\n", + "# )\n", + "# lane_query_se2 = lane_polyline_se2.interpolate(distances_se2)\n", + "\n", + "\n", + "# left_boundary = left_boundaries[lane_id].array\n", + "# right_boundary = right_boundaries[lane_id].array\n", + "\n", + "# assert len(left_boundary) > 0 and len(right_boundary) > 0\n", + "# ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", + "# ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "\n", + "# ax.scatter(lane_query_se2[:, 0], lane_query_se2[:, 1], c=lane_query_se2[:, 2] / np.pi, cmap=\"viridis\")\n", + "\n", + "\n", + "# MAX_LANE_WIDTH = 25\n", + "# for state_se2_array in lane_query_se2:\n", + "# for sign in [1.0, -1.0]:\n", + "# perp_start_point = translate_along_yaw(StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * 0.1))\n", + "# perp_end_point = translate_along_yaw(\n", + "# StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)\n", + "# )\n", + "# ax.plot(\n", + "# [perp_start_point.x, perp_end_point.x],\n", + "# [perp_start_point.y, perp_end_point.y],\n", + "# color=\"lime\" if sign > 0 else \"red\",\n", + "# linestyle=\"--\",\n", + "# alpha=0.5,\n", + "# )\n", + "\n", + "\n", + "# for road_edge in road_edges.values():\n", + "# # print(len(driveway))\n", + "# ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"black\", label=\"road_edge\", linestyle=\"dashdot\")\n", + "\n", + "\n", + "\n", + "# for road_line in road_lines.values():\n", + "# # print(len(driveway))\n", + "# ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "# ax.set_aspect(\"equal\")\n", + "\n", + "# ax.set_xlim(lane_[:, 0].min() - 10, lane_[:, 0].max() + 10)\n", + "# ax.set_ylim(lane_[:, 1].min() - 10, lane_[:, 1].max() + 10)\n", + "\n", + "# lanes_map_features[lane_id]" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/waymo_perception/testing.ipynb b/notebooks/waymo_perception/testing.ipynb new file mode 100644 index 00000000..19745f8f --- /dev/null +++ b/notebooks/waymo_perception/testing.ipynb @@ -0,0 +1,279 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset import dataset_pb2\n", + "\n", + "import json\n", + "import os\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "from PIL import Image\n", + "from tqdm import tqdm\n", + "from waymo_open_dataset import label_pb2\n", + "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", + "from waymo_open_dataset.utils import box_utils\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "\n", + "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", + "\n", + "\n", + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from pyquaternion import Quaternion\n", + "\n", + "from d123.common.geometry.base import StateSE3\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", + "\n", + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "\n", + "# Frame attributes:\n", + "# context: \n", + "# timestamp_micros: \n", + "# pose: \n", + "# images: List with 5 images\n", + "# lasers: \n", + "# Length: 5\n", + "# laser_labels: \n", + "# Length: 0\n", + "# projected_lidar_labels: \n", + "# Length: 0\n", + "# camera_labels: \n", + "# Length: 0\n", + "# no_label_zones: \n", + "# Length: 0\n", + "# map_features: \n", + "# Length: 0\n", + "# map_pose_offset: \n", + "\n", + "file_idx = 0\n", + "pathname = tfrecords_file_list[file_idx]\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "num_frames = sum(1 for _ in dataset)\n", + "\n", + "\n", + "def read_jpg_image(data: bytes) -> np.ndarray:\n", + " \"\"\"Read a JPEG image from bytes and return it as a numpy array.\"\"\"\n", + " image = Image.open(io.BytesIO(data))\n", + " return np.array(image)\n", + "\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + "\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " print(frame.context)\n", + "\n", + " # # plt.show()\n", + " if frame_idx == 0:\n", + " break\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "for laser_calibration in frame.context.laser_calibrations:\n", + " print(np.array(laser_calibration.extrinsic.transform).reshape(4, 4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "for frame_idx, data in enumerate(dataset):\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " if frame_idx == 2:\n", + " break\n", + "\n", + "print(\"Ego\")\n", + "ego_transform = np.array(frame.pose.transform).reshape(4, 4)\n", + "print(ego_transform[:3, 3])\n", + "\n", + "# 1 [ 1.5441613 -0.02302364 2.11557864]\n", + "# 2 [1.49672397 0.0954948 2.11616463]\n", + "# 3 [ 1.49442485 -0.09637497 2.11519385]\n", + "# 4 [1.43213651 0.11612398 2.11625087]\n", + "# 5 [ 1.42936162 -0.11545043 2.1150792 ]\n", + "\n", + "\n", + "# frame.map_pose_offset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(64, 20))\n", + "\n", + "\n", + "def plot_range_image_helper(data, name, layout, vmin=0, vmax=1, cmap=\"gray\"):\n", + " \"\"\"Plots range image.\n", + "\n", + " Args:\n", + " data: range image data\n", + " name: the image title\n", + " layout: plt layout\n", + " vmin: minimum value of the passed data\n", + " vmax: maximum value of the passed data\n", + " cmap: color map\n", + " \"\"\"\n", + " plt.subplot(*layout)\n", + " plt.imshow(data, cmap=cmap, vmin=vmin, vmax=vmax)\n", + " plt.title(name)\n", + " plt.grid(False)\n", + " plt.axis(\"off\")\n", + "\n", + "\n", + "def get_range_image(laser_name, return_index):\n", + " \"\"\"Returns range image given a laser name and its return index.\"\"\"\n", + " return range_images[laser_name][return_index]\n", + "\n", + "\n", + "# def show_range_image(range_image, layout_index_start=1):\n", + "# \"\"\"Shows range image.\n", + "\n", + "# Args:\n", + "# range_image: the range image data from a given lidar of type MatrixFloat.\n", + "# layout_index_start: layout offset\n", + "# \"\"\"\n", + "# range_image_tensor = tf.convert_to_tensor(range_image.data)\n", + "# range_image_tensor = tf.reshape(range_image_tensor, range_image.shape.dims)\n", + "# lidar_image_mask = tf.greater_equal(range_image_tensor, 0)\n", + "# range_image_tensor = tf.where(lidar_image_mask, range_image_tensor, tf.ones_like(range_image_tensor) * 1e10)\n", + "# range_image_range = range_image_tensor[..., 0]\n", + "# range_image_intensity = range_image_tensor[..., 1]\n", + "# range_image_elongation = range_image_tensor[..., 2]\n", + "# plot_range_image_helper(range_image_range.numpy(), \"range\", [8, 1, layout_index_start], vmax=75, cmap=\"gray\")\n", + "# plot_range_image_helper(\n", + "# range_image_intensity.numpy(), \"intensity\", [8, 1, layout_index_start + 1], vmax=1.5, cmap=\"gray\"\n", + "# )\n", + "# plot_range_image_helper(\n", + "# range_image_elongation.numpy(), \"elongation\", [8, 1, layout_index_start + 2], vmax=1.5, cmap=\"gray\"\n", + "# )\n", + "\n", + "\n", + "def show_range_image(range_image, layout_index_start=1):\n", + " \"\"\"Shows range image.\n", + "\n", + " Args:\n", + " range_image: the range image data from a given lidar of type MatrixFloat.\n", + " layout_index_start: layout offset\n", + " \"\"\"\n", + " range_image_tensor = np.array([data for data in range_image.data]).reshape(range_image.shape.dims)\n", + " lidar_image_mask = np.greater_equal(range_image_tensor, 0)\n", + " range_image_tensor = np.where(lidar_image_mask, range_image_tensor, np.ones_like(range_image_tensor) * 1e10)\n", + " range_image_range = range_image_tensor[..., 0]\n", + " range_image_intensity = range_image_tensor[..., 1]\n", + " range_image_elongation = range_image_tensor[..., 2]\n", + " plot_range_image_helper(range_image_range, \"range\", [8, 1, layout_index_start], vmax=75, cmap=\"gray\")\n", + " plot_range_image_helper(range_image_intensity, \"intensity\", [8, 1, layout_index_start + 1], vmax=1.5, cmap=\"gray\")\n", + " plot_range_image_helper(range_image_elongation, \"elongation\", [8, 1, layout_index_start + 2], vmax=1.5, cmap=\"gray\")\n", + "\n", + "\n", + "frame.lasers.sort(key=lambda laser: laser.name)\n", + "show_range_image(get_range_image(open_dataset.LaserName.TOP, 0), 1)\n", + "# show_range_image(get_range_image(open_dataset.LaserName.TOP, 1), 4)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.datatypes.time.time_point import TimePoint\n", + "\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " if frame_idx == 4:\n", + " break\n", + " break\n", + "\n", + "\n", + "# for calibration in frame.context.camera_calibrations:\n", + "\n", + "frame.timestamp_micros, frame.images[0].pose_timestamp\n", + "# frame.images[0]\n", + "\n", + "frame_timestamp = TimePoint.from_us(frame.timestamp_micros)\n", + "image_timestamp = TimePoint.from_s(frame.images[0].pose_timestamp)\n", + "frame_timestamp.time_s, frame_timestamp.time_s" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..d5419010 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,99 @@ +[tool.isort] +profile = "black" + +[build-system] +requires = ["setuptools>=58", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3.9", + "License :: OSI Approved :: Apache Software License", +] +name = "d123" +version = "v0.0.6" +authors = [{ name = "Daniel Dauner", email = "daniel.dauner@gmail.com" }] +description = "TODO" +readme = "README.md" +requires-python = ">=3.9" +license = {text = "Apache-2.0"} +dependencies = [ + "aioboto3", + "aiofiles", + "bokeh", + "casadi", + "control", + "Fiona", + "geopandas", + "guppy3", + "joblib", + "matplotlib", + "nest_asyncio", + "numpy", + "opencv-python", + "pandas", + "Pillow", + "psutil", + "pyarrow", + "pyinstrument", + "pyogrio", + "pyquaternion", + "pytest", + "rasterio", + "ray", + "retry", + "rtree", + "scipy", + "selenium", + "setuptools", + "shapely>=2.0.0", + "SQLAlchemy==1.4.27", + "sympy", + "tornado", + "tqdm", + "ujson", + "notebook", + "pre-commit", + "cachetools", + "hydra_colorlog", + "hydra-core", + "lxml", + "trimesh", + "viser", +] + +[project.optional-dependencies] +dev = [ + "black", + "isort", + "flake8", + "pre-commit", +] +docs = [ + "Sphinx", + "sphinx-rtd-theme", + "sphinx-autobuild", + "myst-parser", +] +nuplan = [ + "nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2", +] +waymo = [ + "tensorflow==2.11.0", + "waymo-open-dataset-tf-2-11-0", + "intervaltree", +] +av2 = [ + "av2==0.2.1", +] + +[tool.setuptools.packages.find] +where = ["."] +include = ["d123*"] # Only include d123 package +exclude = ["notebooks*", "tests*", "docs*"] # Explicitly exclude notebooks + +[project.urls] +"Homepage" = "https://github.com/DanielDauner/d123" +"Bug Tracker" = "https://github.com/DanielDauner/d123/issues" diff --git a/requirements.txt b/requirements.txt index f0c697e2..4c17c209 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,59 +1,58 @@ -nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2 +# nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2 -# nuplan requirements -aioboto3 -aiofiles -bokeh -casadi -control -Fiona -geopandas -guppy3 -joblib -matplotlib -nest_asyncio -numpy -opencv-python -pandas -Pillow -psutil -pyarrow -pyinstrument -pyogrio -pyquaternion -pytest -rasterio -ray -retry -rtree -scipy -selenium -setuptools -shapely==2.1.1 -SQLAlchemy==1.4.27 -sympy -tornado -tqdm -ujson -notebook -pre-commit -cachetools +# # # nuplan requirements +# aioboto3 +# aiofiles +# bokeh +# casadi +# control +# Fiona +# geopandas +# guppy3 +# joblib +# matplotlib +# nest_asyncio +# numpy +# opencv-python +# pandas +# Pillow +# psutil +# pyarrow +# pyinstrument +# pyogrio +# pyquaternion +# pytest +# rasterio +# ray +# retry +# rtree +# scipy +# selenium +# setuptools +# shapely>=2.0.0 +# SQLAlchemy==1.4.27 +# sympy +# tornado +# tqdm +# ujson +# notebook +# pre-commit +# cachetools -# hydra -hydra_colorlog -hydra-core +# # hydra +# hydra_colorlog +# hydra-core -# d123 -lxml -trimesh -viser -gym==0.17.2 -gymnasium==0.26.3 +# # d123 +# lxml +# trimesh +# viser +# gym==0.17.2 +# gymnasium==0.26.3 -# torch & lighting -torch==2.6.0 -torchvision==0.21.0 -lightning -tensorboard -protobuf==4.25.3 -pre-commit +# # torch & lighting +# torch==2.6.0 +# torchvision==0.21.0 +# lightning +# tensorboard +# protobuf==4.25.3 diff --git a/scripts/dataset/run_log_caching.sh b/scripts/dataset/run_log_caching.sh index ef62ff58..34ff923c 100644 --- a/scripts/dataset/run_log_caching.sh +++ b/scripts/dataset/run_log_caching.sh @@ -1,7 +1,7 @@ -python $D123_DEVKIT_ROOT/d123/script/run_dataset_caching.py \ +python $D123_DEVKIT_ROOT/d123/script/run_dataset_conversion.py \ experiment_name="caching" \ -worker.threads_per_node=10 +worker.threads_per_node=32 # worker=single_machine_thread_pool diff --git a/scripts/download/download_av2.sh b/scripts/download/download_av2.sh new file mode 100644 index 00000000..039de648 --- /dev/null +++ b/scripts/download/download_av2.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash + +# Dataset URIs +# s3://argoverse/datasets/av2/sensor/ +# s3://argoverse/datasets/av2/lidar/ +# s3://argoverse/datasets/av2/motion-forecasting/ +# s3://argoverse/datasets/av2/tbv/ + +DATASET_NAMES=("sensor" "lidar" "motion-forecasting" "tbv") +TARGET_DIR="/path/to/argoverse" + +for DATASET_NAME in "${DATASET_NAMES[@]}"; do + mkdir -p "$TARGET_DIR/$DATASET_NAME" + s5cmd --no-sign-request cp "s3://argoverse/datasets/av2/$DATASET_NAME/*" "$TARGET_DIR/$DATASET_NAME" +done + + +# wget -r s3://argoverse/datasets/av2/sensor/test/0f0cdd79-bc6c-35cd-9d99-7ae2fc7e165c/sensors/cameras/ring_front_center/315965893599927217.jpg +# wget http://argoverse.s3.amazonaws.com/datasets/av2/sensor/test/0f0cdd79-bc6c-35cd-9d99-7ae2fc7e165c/sensors/cameras/ring_front_center/315965893599927217.jpg diff --git a/scripts/download/download_nuplan.sh b/scripts/download/download_nuplan_logs.sh similarity index 100% rename from scripts/download/download_nuplan.sh rename to scripts/download/download_nuplan_logs.sh diff --git a/scripts/download/download_nuplan_sensor.sh b/scripts/download/download_nuplan_sensor.sh new file mode 100644 index 00000000..6cd3d246 --- /dev/null +++ b/scripts/download/download_nuplan_sensor.sh @@ -0,0 +1,14 @@ +# NOTE: Please check the LICENSE file when downloading the nuPlan dataset +# wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/LICENSE + +# train: nuplan_train + +# val: nuplan_val + +# test: nuplan_test + +# mini: nuplan_mini_train, nuplan_mini_val, nuplan_mini_test +for split in {0..8}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_camera_${split}.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_lidar_${split}.zip +done diff --git a/scripts/viz/run_viser.sh b/scripts/viz/run_viser.sh new file mode 100644 index 00000000..0dafc0a3 --- /dev/null +++ b/scripts/viz/run_viser.sh @@ -0,0 +1,6 @@ + + +python $D123_DEVKIT_ROOT/d123/script/run_viser.py \ +scene_filter=log_scenes \ +scene_filter.shuffle=True \ +worker=sequential diff --git a/setup.py b/setup.py deleted file mode 100644 index 8c516b45..00000000 --- a/setup.py +++ /dev/null @@ -1,30 +0,0 @@ -import os - -import setuptools - -# Change directory to allow installation from anywhere -script_folder = os.path.dirname(os.path.realpath(__file__)) -os.chdir(script_folder) - -with open("requirements.txt") as f: - requirements = f.read().splitlines() - -# Installs -setuptools.setup( - name="d123", - version="0.0.5", - author="Daniel Dauner", - author_email="daniel.dauner@gmail.com", - description="TODO", - url="https://github.com/autonomousvision/d123", - python_requires=">=3.9", - packages=["d123"], - package_dir={"": "."}, - classifiers=[ - "Programming Language :: Python :: 3.9", - "Operating System :: OS Independent", - "License :: Free for non-commercial use", - ], - license="apache-2.0", - install_requires=requirements, -)