From d4bff4bcabaf1ac11649c695d0dc78fb7bd36499 Mon Sep 17 00:00:00 2001 From: sephyli Date: Thu, 30 Oct 2025 10:49:20 +0800 Subject: [PATCH 1/4] add nuscenes dataset --- .../conversion/datasets/nuscenes/.gitkeep | 0 .../nuscenes/nuscenes_data_converter.py | 671 +++++++++++++++++ .../nuscenes/nuscenes_map_conversion.py | 708 ++++++++++++++++++ .../datatypes/detections/box_detections.py | 4 +- .../datatypes/maps/abstract_map_objects.py | 4 +- .../datatypes/maps/gpkg/gpkg_map_objects.py | 37 +- src/py123d/datatypes/maps/gpkg/gpkg_utils.py | 23 +- .../datatypes/scene/arrow/arrow_scene.py | 9 +- .../scene/arrow/utils/arrow_getters.py | 1 + .../datatypes/sensors/lidar/lidar_index.py | 9 + .../vehicle_state/vehicle_parameters.py | 10 + .../config/common/default_dataset_paths.yaml | 4 +- .../conversion/datasets/nuscenes_dataset.yaml | 35 + 13 files changed, 1497 insertions(+), 18 deletions(-) delete mode 100644 src/py123d/conversion/datasets/nuscenes/.gitkeep create mode 100644 src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py create mode 100644 src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py create mode 100644 src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml diff --git a/src/py123d/conversion/datasets/nuscenes/.gitkeep b/src/py123d/conversion/datasets/nuscenes/.gitkeep deleted file mode 100644 index e69de29b..00000000 diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py new file mode 100644 index 00000000..eb743fb0 --- /dev/null +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py @@ -0,0 +1,671 @@ +import gc +import json +import os +import numpy as np +import pyarrow as pa + +from dataclasses import asdict +from pathlib import Path +from typing import Any, Dict, Final, List, Optional, Tuple, Union +from nuscenes import NuScenes +from nuscenes.can_bus.can_bus_api import NuScenesCanBus +from nuscenes.utils.data_classes import Box +from nuscenes.utils.splits import create_splits_scenes +from pyquaternion import Quaternion + +from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter +from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter +from py123d.script.builders.worker_pool_builder import WorkerPool +from py123d.datatypes.detections.detection import ( + BoxDetectionSE3, + BoxDetectionWrapper, + BoxDetectionMetadata, + TrafficLightDetection, + TrafficLightDetectionWrapper, +) +from py123d.datatypes.detections.detection_types import DetectionType +from py123d.datatypes.maps.map_metadata import MapMetadata +from py123d.datatypes.scene.scene_metadata import LogMetadata +from py123d.datatypes.sensors.camera.pinhole_camera import ( + PinholeCameraMetadata, + PinholeCameraType, + PinholeDistortion, + PinholeIntrinsics, +) +from py123d.datatypes.sensors.lidar.lidar import LiDARMetadata, LiDARType +from py123d.conversion.utils.sensor_utils.lidar_index_registry import NuscenesLidarIndex + +from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index +from py123d.datatypes.vehicle_state.vehicle_parameters import get_nuscenes_renauly_zoe_parameters +from py123d.geometry import StateSE3, BoundingBoxSE3, BoundingBoxSE3Index +from py123d.geometry.vector import Vector3D, Vector3DIndex +from py123d.common.utils.arrow_helper import open_arrow_table, write_arrow_table +from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import write_nuscenes_map, NUSCENES_MAPS +from py123d.conversion.abstract_dataset_converter import AbstractDatasetConverter +from py123d.conversion.dataset_converter_config import DatasetConverterConfig +from py123d.datatypes.time.time_point import TimePoint + +TARGET_DT: Final[float] = 0.1 +NUSCENES_DT: Final[float] = 0.5 +SORT_BY_TIMESTAMP: Final[bool] = True +NUSCENES_DETECTION_NAME_DICT = { + # Vehicles (4+ wheels) + "vehicle.car": DetectionType.VEHICLE, + "vehicle.truck": DetectionType.VEHICLE, + "vehicle.bus.bendy": DetectionType.VEHICLE, + "vehicle.bus.rigid": DetectionType.VEHICLE, + "vehicle.construction": DetectionType.VEHICLE, + "vehicle.emergency.ambulance": DetectionType.VEHICLE, + "vehicle.emergency.police": DetectionType.VEHICLE, + "vehicle.trailer": DetectionType.VEHICLE, + + # Bicycles / Motorcycles + "vehicle.bicycle": DetectionType.BICYCLE, + "vehicle.motorcycle": DetectionType.BICYCLE, + + # Pedestrians (all subtypes) + "human.pedestrian.adult": DetectionType.PEDESTRIAN, + "human.pedestrian.child": DetectionType.PEDESTRIAN, + "human.pedestrian.construction_worker": DetectionType.PEDESTRIAN, + "human.pedestrian.personal_mobility": DetectionType.PEDESTRIAN, + "human.pedestrian.police_officer": DetectionType.PEDESTRIAN, + "human.pedestrian.stroller": DetectionType.PEDESTRIAN, + "human.pedestrian.wheelchair": DetectionType.PEDESTRIAN, + + # Traffic cone / barrier + "movable_object.trafficcone": DetectionType.TRAFFIC_CONE, + "movable_object.barrier": DetectionType.BARRIER, + + # Generic objects + "movable_object.pushable_pullable": DetectionType.GENERIC_OBJECT, + "movable_object.debris": DetectionType.GENERIC_OBJECT, + "static_object.bicycle_rack": DetectionType.GENERIC_OBJECT, + "animal": DetectionType.GENERIC_OBJECT, +} + +NUSCENES_CAMERA_TYPES = { + PinholeCameraType.CAM_F0: "CAM_FRONT", + PinholeCameraType.CAM_B0: "CAM_BACK", + PinholeCameraType.CAM_L0: "CAM_FRONT_LEFT", + PinholeCameraType.CAM_L1: "CAM_BACK_LEFT", + PinholeCameraType.CAM_R0: "CAM_FRONT_RIGHT", + PinholeCameraType.CAM_R1: "CAM_BACK_RIGHT", +} +NUSCENES_DATA_ROOT = Path(os.environ["NUSCENES_DATA_ROOT"]) + + +class NuScenesDataConverter(AbstractDatasetConverter): + def __init__( + self, + splits: List[str], + nuscenes_data_root: Union[Path, str], + nuscenes_map_root: Union[Path, str], + nuscenes_lanelet2_root: Union[Path, str], + use_lanelet2: bool, + dataset_converter_config: DatasetConverterConfig, + version: str = "v1.0-trainval", + ) -> None: + super().__init__(dataset_converter_config) + self._splits: List[str] = splits + self._nuscenes_data_root: Path = Path(nuscenes_data_root) + self._nuscenes_map_root: Path = Path(nuscenes_map_root) + self._nuscenes_lanelet2_root: Path = Path(nuscenes_lanelet2_root) + self._use_lanelet2 = use_lanelet2 + self._version = version + self._scene_tokens_per_split: Dict[str, List[str]] = self._collect_scene_tokens() + self._target_dt: float = TARGET_DT + + def _collect_scene_tokens(self) -> Dict[str, List[str]]: + scene_tokens_per_split: Dict[str, List[str]] = {} + nusc = NuScenes(version=self._version, dataroot=str(self._nuscenes_data_root), verbose=False) + + scene_splits = create_splits_scenes() + available_scenes = [scene for scene in nusc.scene] + + for split in self._splits: + # Map the split name to the division of nuScenes + nusc_split = split.replace("nuscenes_", "") + if nusc_split == "trainval": + scene_names = scene_splits['train'] + scene_splits['val'] + else: + scene_names = scene_splits.get(nusc_split, []) + + # get token + scene_tokens = [ + scene['token'] for scene in available_scenes + if scene['name'] in scene_names + ] + scene_tokens_per_split[split] = scene_tokens + + return scene_tokens_per_split + + def get_available_splits(self) -> List[str]: + return [ + "nuscenes_train", + "nuscenes_val", + "nuscenes_test", + "nuscenes_mini_train", + "nuscenes_mini_val", + ] + + def get_number_of_maps(self) -> int: + """Inherited, see superclass.""" + return len(NUSCENES_MAPS) + + def get_number_of_logs(self) -> int: + """Inherited, see superclass.""" + return sum(len(scene_tokens) for scene_tokens in self._scene_tokens_per_split.values()) + + def convert_map(self, map_index: int, map_writer: AbstractMapWriter) -> None: + """Inherited, see superclass.""" + map_name = NUSCENES_MAPS[map_index] + + map_metadata = _get_nuscenes_map_metadata(map_name) + map_needs_writing = map_writer.reset(self.dataset_converter_config, map_metadata) + + if map_needs_writing: + write_nuscenes_map( + nuscenes_maps_root=self._nuscenes_map_root, + location=map_name, + map_writer=map_writer, + use_lanelet2=self._use_lanelet2, + lanelet2_root=Path(self._nuscenes_lanelet2_root), + ) + + map_writer.close() + + def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: + """Inherited, see superclass.""" + # Find the scene token for the given log index + all_scene_tokens = [] + for split, scene_tokens in self._scene_tokens_per_split.items(): + all_scene_tokens.extend([(split, token) for token in scene_tokens]) + + if log_index >= len(all_scene_tokens): + raise ValueError(f"Log index {log_index} is out of range. Total logs: {len(all_scene_tokens)}") + + split, scene_token = all_scene_tokens[log_index] + + nusc = NuScenes(version=self._version, dataroot=str(self._nuscenes_data_root), verbose=False) + scene = nusc.get("scene", scene_token) + log_record = nusc.get("log", scene["log_token"]) + + # 1. Initialize log metadata + log_metadata = LogMetadata( + dataset="nuscenes", + split=split, + log_name=scene["name"], + location=log_record["location"], + timestep_seconds=TARGET_DT, + vehicle_parameters=get_nuscenes_renauly_zoe_parameters(), + camera_metadata=_get_nuscenes_camera_metadata(nusc, scene, self.dataset_converter_config), + lidar_metadata=_get_nuscenes_lidar_metadata(nusc, scene, self.dataset_converter_config), + map_metadata=_get_nuscenes_map_metadata(log_record["location"]), + ) + + # 2. Prepare log writer + log_needs_writing = log_writer.reset(self.dataset_converter_config, log_metadata) + + if log_needs_writing: + can_bus = NuScenesCanBus(dataroot=str(self._nuscenes_data_root)) + + step_interval = max(1, int(TARGET_DT / NUSCENES_DT)) + sample_count = 0 + + # Traverse all samples in the scene + sample_token = scene["first_sample_token"] + while sample_token: + if sample_count % step_interval == 0: + sample = nusc.get("sample", sample_token) + + log_writer.write( + timestamp=TimePoint.from_us(sample["timestamp"]), + ego_state=_extract_nuscenes_ego_state(nusc, sample, can_bus), + box_detections=_extract_nuscenes_box_detections(nusc, sample), + traffic_lights=_extract_nuscenes_traffic_lights(), # nuScenes doesn't have traffic lights + cameras=_extract_nuscenes_cameras( + nusc=nusc, + sample=sample, + dataset_converter_config=self.dataset_converter_config, + ), + lidars=_extract_nuscenes_lidars( + nusc=nusc, + sample=sample, + dataset_converter_config=self.dataset_converter_config, + ), + scenario_tags=_extract_nuscenes_scenario_tag(), # nuScenes doesn't have scenario tags + route_lane_group_ids=_extract_nuscenes_route_lane_group_ids(), + # nuScenes doesn't have route info + ) + + sample_token = sample["next"] + sample_count += 1 + + log_writer.close() + del nusc + gc.collect() + + def convert_logs(self, worker: WorkerPool) -> None: + """ + NuScenes logs conversion is handled externally through convert_log method. + This method is kept for interface compatibility. + """ + pass + + +def _get_nuscenes_camera_metadata( + nusc: NuScenes, + scene: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, +) -> Dict[PinholeCameraType, PinholeCameraMetadata]: + camera_metadata: Dict[PinholeCameraType, PinholeCameraMetadata] = {} + + if dataset_converter_config.include_cameras: + first_sample_token = scene["first_sample_token"] + first_sample = nusc.get("sample", first_sample_token) + + for camera_type, camera_channel in NUSCENES_CAMERA_TYPES.items(): + cam_token = first_sample["data"][camera_channel] + cam_data = nusc.get("sample_data", cam_token) + calib = nusc.get("calibrated_sensor", cam_data["calibrated_sensor_token"]) + + intrinsic_matrix = np.array(calib["camera_intrinsic"]) + intrinsic = PinholeIntrinsics.from_camera_matrix(intrinsic_matrix) + distortion = PinholeDistortion.from_array(np.zeros(5), copy=False) + + camera_metadata[camera_type] = PinholeCameraMetadata( + camera_type=camera_type, + width=cam_data["width"], + height=cam_data["height"], + intrinsics=intrinsic, + distortion=distortion, + ) + + return camera_metadata + + +def _get_nuscenes_lidar_metadata( + nusc: NuScenes, + scene: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, +) -> Dict[LiDARType, LiDARMetadata]: + metadata: Dict[LiDARType, LiDARMetadata] = {} + + if dataset_converter_config.include_lidars: + first_sample_token = scene["first_sample_token"] + first_sample = nusc.get("sample", first_sample_token) + lidar_token = first_sample["data"]["LIDAR_TOP"] + lidar_data = nusc.get("sample_data", lidar_token) + calib = nusc.get("calibrated_sensor", lidar_data["calibrated_sensor_token"]) + + translation = np.array(calib["translation"]) + rotation = Quaternion(calib["rotation"]).rotation_matrix + extrinsic = np.eye(4) + extrinsic[:3, :3] = rotation + extrinsic[:3, 3] = translation + extrinsic = StateSE3.from_transformation_matrix(extrinsic) + + metadata[LiDARType.LIDAR_MERGED] = LiDARMetadata( + lidar_type=LiDARType.LIDAR_MERGED, + lidar_index=NuscenesLidarIndex, + extrinsic=extrinsic, + ) + + return metadata + + +def _get_nuscenes_map_metadata(location): + return MapMetadata( + dataset="nuscenes", + split=None, + log_name=None, + location=location, + map_has_z=False, + map_is_local=False, + ) + + +def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: + lidar_data = nusc.get("sample_data", sample["data"]["LIDAR_TOP"]) + ego_pose = nusc.get("ego_pose", lidar_data["ego_pose_token"]) + + quat = Quaternion(ego_pose["rotation"]) + + vehicle_parameters = get_nuscenes_renauly_zoe_parameters() + pose = StateSE3( + x=ego_pose["translation"][0], + y=ego_pose["translation"][1], + z=ego_pose["translation"][2], + qw=quat.w, + qx=quat.x, + qy=quat.y, + qz=quat.z, + ) + + scene_name = nusc.get("scene", sample["scene_token"])["name"] + + try: + pose_msgs = can_bus.get_messages(scene_name, "pose") + except Exception as e: + pose_msgs = [] + + if pose_msgs: + closest_msg = None + min_time_diff = float('inf') + for msg in pose_msgs: + time_diff = abs(msg["utime"] - sample["timestamp"]) + if time_diff < min_time_diff: + min_time_diff = time_diff + closest_msg = msg + + if closest_msg and min_time_diff < 500000: + velocity = [*closest_msg["vel"]] + acceleration = [*closest_msg["accel"]] + angular_velocity = [*closest_msg["rotation_rate"]] + else: + velocity = acceleration = angular_velocity = [0.0, 0.0, 0.0] + else: + velocity = acceleration = angular_velocity = [0.0, 0.0, 0.0] + + dynamic_state = DynamicStateSE3( + velocity=Vector3D(*velocity), + acceleration=Vector3D(*acceleration), + angular_velocity=Vector3D(*angular_velocity), + ) + + return EgoStateSE3( + center_se3=pose, + dynamic_state_se3=dynamic_state, + vehicle_parameters=vehicle_parameters, + timepoint=TimePoint.from_us(sample["timestamp"]), + ) + + +def _extract_nuscenes_box_detections( + nusc: NuScenes, + sample: Dict[str, Any] +) -> BoxDetectionWrapper: + box_detections: List[BoxDetectionSE3] = [] + + for ann_token in sample["anns"]: + ann = nusc.get("sample_annotation", ann_token) + box = Box(ann["translation"], ann["size"], Quaternion(ann["rotation"])) + + box_quat = box.orientation + euler_angles = box_quat.yaw_pitch_roll # (yaw, pitch, roll) + + # Create StateSE3 for box center and orientation + center_quat = box.orientation + center = StateSE3( + box.center[0], + box.center[1], + box.center[2], + center_quat.w, + center_quat.x, + center_quat.y, + center_quat.z, + ) + bounding_box = BoundingBoxSE3(center, box.wlh[1], box.wlh[0], box.wlh[2]) + # Get detection type + category = ann["category_name"] + det_type = None + for key, value in NUSCENES_DETECTION_NAME_DICT.items(): + if category.startswith(key): + det_type = value + break + + if det_type is None: + print(f"Warning: Unmapped nuScenes category: {category}, skipping") + continue + + # Get velocity if available + velocity = nusc.box_velocity(ann_token) + velocity_3d = Vector3D(x=velocity[0], y=velocity[1], z=velocity[2] if len(velocity) > 2 else 0.0) + + metadata = BoxDetectionMetadata( + detection_type=det_type, + track_token=ann["instance_token"], + timepoint=TimePoint.from_us(sample["timestamp"]), + confidence=1.0, # nuScenes annotations are ground truth + num_lidar_points=ann.get("num_lidar_pts", 0), + ) + + box_detection = BoxDetectionSE3( + metadata=metadata, + bounding_box_se3=bounding_box, + velocity=velocity_3d, + ) + box_detections.append(box_detection) + + return BoxDetectionWrapper(box_detections=box_detections) + + +def _extract_nuscenes_traffic_lights() -> TrafficLightDetectionWrapper: + """nuScenes doesn't have traffic light information.""" + return TrafficLightDetectionWrapper(traffic_light_detections=[]) + + +def _extract_nuscenes_cameras( + nusc: NuScenes, + sample: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, +) -> Dict[PinholeCameraType, Tuple[Union[str, bytes], StateSE3]]: + camera_dict: Dict[PinholeCameraType, Tuple[Union[str, bytes], StateSE3]] = {} + + if dataset_converter_config.include_cameras: + for camera_type, camera_channel in NUSCENES_CAMERA_TYPES.items(): + cam_token = sample["data"][camera_channel] + cam_data = nusc.get("sample_data", cam_token) + + # Check timestamp synchronization (within 100ms) + if abs(cam_data["timestamp"] - sample["timestamp"]) > 100000: + continue + + calib = nusc.get("calibrated_sensor", cam_data["calibrated_sensor_token"]) + + translation = np.array(calib["translation"]) + rotation = Quaternion(calib["rotation"]).rotation_matrix + extrinsic_matrix = np.eye(4) + extrinsic_matrix[:3, :3] = rotation + extrinsic_matrix[:3, 3] = translation + extrinsic = StateSE3.from_transformation_matrix(extrinsic_matrix) + + cam_path = NUSCENES_DATA_ROOT / cam_data["filename"] + + if cam_path.exists() and cam_path.is_file(): + if dataset_converter_config.camera_store_option == "path": + camera_data = str(cam_path) + elif dataset_converter_config.camera_store_option == "binary": + with open(cam_path, "rb") as f: + camera_data = f.read() + else: + continue + + camera_dict[camera_type] = (camera_data, extrinsic) + + return camera_dict + + +def _extract_nuscenes_lidars( + nusc: NuScenes, + sample: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, +) -> Dict[LiDARType, Optional[str]]: + lidar_dict: Dict[LiDARType, Optional[str]] = {} + + if dataset_converter_config.include_lidars: + lidar_token = sample["data"]["LIDAR_TOP"] + lidar_data = nusc.get("sample_data", lidar_token) + lidar_path = NUSCENES_DATA_ROOT / lidar_data["filename"] + + if lidar_path.exists() and lidar_path.is_file(): + lidar_dict[LiDARType.LIDAR_MERGED] = str(lidar_path) + else: + lidar_dict[LiDARType.LIDAR_MERGED] = None + + return lidar_dict + + +def _extract_nuscenes_scenario_tag() -> List[str]: + """nuScenes doesn't have scenario tags.""" + return ["unknown"] + + +def _extract_nuscenes_route_lane_group_ids() -> List[int]: + """nuScenes doesn't have route lane group information.""" + return [] + + +# Updated arrow conversion function using the new extraction functions +def convert_nuscenes_log_to_arrow( + args: List[Dict[str, Union[str, List[str]]]], + dataset_converter_config: DatasetConverterConfig, + version: str +) -> List[Any]: + for log_info in args: + scene_token: str = log_info["scene_token"] + split: str = log_info["split"] + + nusc = NuScenes(version=version, dataroot=str(NUSCENES_DATA_ROOT), verbose=False) + scene = nusc.get("scene", scene_token) + + log_file_path = dataset_converter_config.output_path / split / f"{scene_token}.arrow" + + if dataset_converter_config.force_log_conversion or not log_file_path.exists(): + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + # Define schema + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("ego_state", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("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())), + ("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 dataset_converter_config.lidar_store_option == "path": + schema_column_list.append(("lidar", pa.string())) + + if dataset_converter_config.camera_store_option == "path": + for camera_type in NUSCENES_CAMERA_TYPES.keys(): + schema_column_list.append((camera_type.serialize(), pa.string())) + schema_column_list.append((f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4))) + + recording_schema = pa.schema(schema_column_list) + + log_record = nusc.get("log", scene["log_token"]) + location = log_record["location"] + + # Create metadata using the same functions as the new interface + metadata = LogMetadata( + dataset="nuscenes", + split=split, + log_name=scene["name"], + location=location, + timestep_seconds=TARGET_DT, + vehicle_parameters=get_nuscenes_renauly_zoe_parameters(), + camera_metadata=_get_nuscenes_camera_metadata(nusc, scene, dataset_converter_config), + lidar_metadata=_get_nuscenes_lidar_metadata(nusc, scene, dataset_converter_config), + map_metadata=_get_nuscenes_map_metadata(location), + ) + + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + } + ) + + _write_arrow_table_with_new_interface( + nusc, scene, recording_schema, log_file_path, dataset_converter_config + ) + + del nusc + gc.collect() + + return [] + + +def _write_arrow_table_with_new_interface( + nusc: NuScenes, + scene: Dict[str, Any], + recording_schema: pa.schema, + log_file_path: Path, + dataset_converter_config: DatasetConverterConfig, +) -> None: + can_bus = NuScenesCanBus(dataroot=str(NUSCENES_DATA_ROOT)) + + with pa.OSFile(str(log_file_path), "wb") as sink: + with pa.ipc.new_file(sink, recording_schema) as writer: + step_interval = max(1, int(TARGET_DT / NUSCENES_DT)) + sample_count = 0 + + sample_token = scene["first_sample_token"] + while sample_token: + if sample_count % step_interval == 0: + sample = nusc.get("sample", sample_token) + + # Use the new extraction functions for consistency + ego_state = _extract_nuscenes_ego_state(nusc, sample, can_bus) + box_detections = _extract_nuscenes_box_detections(nusc, sample) + cameras = _extract_nuscenes_cameras(nusc, sample, dataset_converter_config) + lidars = _extract_nuscenes_lidars(nusc, sample, dataset_converter_config) + + detections_state_list = [] + for det in box_detections.box_detections: + bbox_array = det.bounding_box_se3.array + + print(f"bbox_array shape: {bbox_array.shape}, ndim: {bbox_array.ndim}") + if bbox_array.ndim > 1: + detections_state_list.append(bbox_array.flatten().tolist()) + else: + detections_state_list.append(bbox_array.tolist()) + + # Prepare row data + row_data = { + "token": [sample_token], + "timestamp": [sample["timestamp"]], + "ego_state": ego_state.array.flatten().tolist(), + "detections_state": detections_state_list, + "detections_velocity": [det.velocity.array.tolist() for det in box_detections.box_detections], + "detections_token": [det.metadata.track_token for det in box_detections.box_detections], + "detections_type": [det.metadata.detection_type.value for det in box_detections.box_detections], + "traffic_light_ids": [], + "traffic_light_types": [], + "scenario_tag": ["unknown"], + "route_lane_group_ids": [], + } + + # Add lidar data if configured + if dataset_converter_config.lidar_store_option == "path": + row_data["lidar"] = [lidars.get(LiDARType.LIDAR_MERGED, None)] + + # Add camera data if configured + if dataset_converter_config.camera_store_option == "path": + for camera_type in NUSCENES_CAMERA_TYPES.keys(): + if camera_type in cameras: + camera_path, extrinsic = cameras[camera_type] + row_data[camera_type.serialize()] = [camera_path] + row_data[f"{camera_type.serialize()}_extrinsic"] = [ + extrinsic.to_transformation_matrix().flatten().tolist()] + 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) + + sample_token = sample["next"] + sample_count += 1 + + # Sort by timestamp if required + 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) \ No newline at end of file diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py new file mode 100644 index 00000000..66f7a658 --- /dev/null +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py @@ -0,0 +1,708 @@ +import lanelet2 +import numpy as np + +from pathlib import Path +from typing import List, Optional +from lanelet2.io import load +from lanelet2.projection import MercatorProjector +from shapely.geometry import Polygon, MultiPolygon, LineString +from shapely.validation import make_valid +from nuscenes.map_expansion.map_api import NuScenesMap +from nuscenes.map_expansion.arcline_path_utils import discretize_lane + +from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter +from py123d.datatypes.maps.cache.cache_map_objects import ( + CacheCarpark, + CacheCrosswalk, + CacheGenericDrivable, + CacheIntersection, + CacheLane, + CacheLaneGroup, + CacheRoadLine, + CacheWalkway, +) +from py123d.datatypes.maps.map_datatypes import RoadLineType +from py123d.geometry import Polyline3D + +NUSCENES_MAPS: List[str] = [ + "boston-seaport", + "singapore-hollandvillage", + "singapore-onenorth", + "singapore-queenstown" +] + + +def write_nuscenes_map(nuscenes_maps_root: Path, location: str, map_writer: AbstractMapWriter, use_lanelet2: bool, lanelet2_root: Optional[str] = None) -> None: + """ + Main function to convert nuscenes map to unified format and write using map_writer. + """ + assert location in NUSCENES_MAPS, f"Map name {location} is not supported." + nusc_map = NuScenesMap(dataroot=str(nuscenes_maps_root), map_name=location) + + # Write all layers + if use_lanelet2: + _write_nuscenes_lanes_lanelet2(nusc_map, map_writer, lanelet2_root) + _write_nuscenes_lane_groups_lanelet2(nusc_map, map_writer, lanelet2_root) + else: + _write_nuscenes_lanes(nusc_map, map_writer) + _write_nuscenes_lane_groups(nusc_map, map_writer) + _write_nuscenes_intersections(nusc_map, map_writer) + _write_nuscenes_crosswalks(nusc_map, map_writer) + _write_nuscenes_walkways(nusc_map, map_writer) + _write_nuscenes_carparks(nusc_map, map_writer) + _write_nuscenes_generic_drivables(nusc_map, map_writer) + _write_nuscenes_stop_lines(nusc_map, map_writer) + _write_nuscenes_road_lines(nusc_map, map_writer) + + +def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str) -> None: + map_name = nusc_map.map_name + osm_map_file = str(Path(lanelet2_root) / f"{map_name}.osm") + + if "boston" in map_name.lower(): + origin_lat, origin_lon = 42.3365, -71.0577 + elif "singapore" in map_name.lower(): + origin_lat, origin_lon = 1.3, 103.8 + else: + origin_lat, origin_lon = 49.0, 8.4 + + origin = lanelet2.io.Origin(origin_lat, origin_lon) + + try: + lanelet_map = lanelet2.io.load(osm_map_file, origin) + except Exception: + try: + projector = lanelet2.projection.MercatorProjector(origin) + lanelet_map = lanelet2.io.load(osm_map_file, projector) + except Exception: + return + + for lanelet in lanelet_map.laneletLayer: + token = lanelet.id + + try: + left_bound = [(p.x, p.y) for p in lanelet.leftBound] + right_bound = [(p.x, p.y) for p in lanelet.rightBound] + polygon_points = left_bound + right_bound[::-1] + polygon = Polygon(polygon_points) + + predecessor_ids = [int(pred.id) for pred in lanelet.previousLanelets] + successor_ids = [int(succ.id) for succ in lanelet.followingLanelets] + + left_lane_id = None + right_lane_id = None + + left_boundary = [(p.x, p.y) for p in lanelet.leftBound] + right_boundary = [(p.x, p.y) for p in lanelet.rightBound] + centerline = [] + for left_pt, right_pt in zip(lanelet.leftBound, lanelet.rightBound): + center_x = (left_pt.x + right_pt.x) / 2 + center_y = (left_pt.y + right_pt.y) / 2 + centerline.append((center_x, center_y)) + + speed_limit_mps = 0.0 + if "speed_limit" in lanelet.attributes: + try: + speed_limit_str = lanelet.attributes["speed_limit"] + if "km/h" in speed_limit_str: + speed_kmh = float(speed_limit_str.replace("km/h", "").strip()) + speed_limit_mps = speed_kmh / 3.6 + except (ValueError, TypeError): + pass + + map_writer.write_lane( + CacheLane( + object_id=token, + lane_group_id=None, + left_boundary=left_boundary, + right_boundary=right_boundary, + centerline=centerline, + left_lane_id=left_lane_id, + right_lane_id=right_lane_id, + predecessor_ids=predecessor_ids, + successor_ids=successor_ids, + speed_limit_mps=speed_limit_mps, + outline=None, + geometry=polygon, + ) + ) + except Exception: + continue + +def _write_nuscenes_lane_groups_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str) -> None: + map_name = nusc_map.map_name + osm_map_file = str(Path(lanelet2_root) / f"{map_name}.osm") + + if "boston" in map_name.lower(): + origin_lat, origin_lon = 42.3365, -71.0577 + else: + origin_lat, origin_lon = 1.3, 103.8 + + origin = lanelet2.io.Origin(origin_lat, origin_lon) + + try: + projector = MercatorProjector(origin) + lanelet_map = load(osm_map_file, projector) + except Exception: + return + + for lanelet in lanelet_map.laneletLayer: + token = lanelet.id + lane_ids = [lanelet.id] + try: + predecessor_ids = [int(lanelet.id) for lanelet in lanelet.previous] + successor_ids = [int(lanelet.id) for lanelet in lanelet.following] + except AttributeError: + predecessor_ids = [] + successor_ids = [] + try: + if hasattr(lanelet, 'left'): + for left_lane in lanelet.left: + predecessor_ids.append(int(left_lane.id)) + if hasattr(lanelet, 'right'): + for right_lane in lanelet.right: + successor_ids.append(int(right_lane.id)) + except Exception: + pass + + try: + left_bound = [(p.x, p.y) for p in lanelet.leftBound] + right_bound = [(p.x, p.y) for p in lanelet.rightBound] + polygon_points = left_bound + right_bound[::-1] + polygon = Polygon(polygon_points) + except Exception: + continue + + try: + map_writer.write_lane_group( + CacheLaneGroup( + object_id=token, + lane_ids=lane_ids, + left_boundary=None, + right_boundary=None, + intersection_id=None, + predecessor_ids=predecessor_ids, + successor_ids=successor_ids, + outline=None, + geometry=polygon, + ) + ) + except Exception: + continue + +def _get_lanelet_connections(lanelet): + """ + Helper function to extract incoming and outgoing lanelets. + """ + incoming = lanelet.incomings + outgoing = lanelet.outgoings + return incoming, outgoing + + +def _write_nuscenes_lanes(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write lane data to map_writer, including topology and boundaries. + """ + lane_records = nusc_map.lane + for lane_record in lane_records: + token = lane_record["token"] + + # Extract geometry from lane record + try: + if "polygon_token" in lane_record: + polygon = nusc_map.extract_polygon(lane_record["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + # Get topology + incoming = nusc_map.get_incoming_lane_ids(token) + outgoing = nusc_map.get_outgoing_lane_ids(token) + + # Get lane connectors + lane_connectors = [] + for connector in nusc_map.lane_connector: + if connector.get("incoming_lane") == token or connector.get("outgoing_lane") == token: + lane_connectors.append(connector["token"]) + + # Extract boundaries + left_boundary = _get_lane_boundary(token, "left", nusc_map) + right_boundary = _get_lane_boundary(token, "right", nusc_map) + + # Skip lanes without valid boundaries + if left_boundary is None or right_boundary is None: + continue + if left_boundary.is_empty or right_boundary.is_empty: + continue + + # Extract baseline path + baseline_path = None + if token in nusc_map.arcline_path_3: + arc_path = nusc_map.arcline_path_3[token] + try: + points = discretize_lane(arc_path, resolution_meters=0.1) + xy_points = [(p[0], p[1]) for p in points] + baseline_path = LineString(xy_points) + except Exception: + baseline_path = None + + # Align boundaries with baseline path direction + if baseline_path and left_boundary: + left_boundary = align_boundary_direction(baseline_path, left_boundary) + if baseline_path and right_boundary: + right_boundary = align_boundary_direction(baseline_path, right_boundary) + + # Write lane object safely + try: + map_writer.write_lane( + CacheLane( + object_id=token, + lane_group_id=lane_record.get("road_segment_token", None), + left_boundary=left_boundary, + right_boundary=right_boundary, + centerline=baseline_path, + left_lane_id=None, # Not directly available in nuscenes + right_lane_id=None, # Not directly available in nuscenes + predecessor_ids=incoming, + successor_ids=outgoing, + speed_limit_mps=0.0, # Default value + outline=None, + geometry=polygon, + ) + ) + except Exception: + continue + + +def _write_nuscenes_lane_groups(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write lane group data to map_writer. + """ + road_segments = nusc_map.road_segment + for segment in road_segments: + token = segment["token"] + + # Extract geometry + try: + if "polygon_token" in segment: + polygon = nusc_map.extract_polygon(segment["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + # Find lanes in this segment + lane_ids = [] + for lane in nusc_map.lane: + if lane.get("road_segment_token") == token: + lane_ids.append(lane["token"]) + + # Get connected segments + incoming, outgoing = _get_connected_segments(token, nusc_map) + + # Extract boundaries + left_boundary = _get_lane_group_boundary(token, "left", nusc_map) + right_boundary = _get_lane_group_boundary(token, "right", nusc_map) + + # Skip invalid boundaries + if left_boundary is None or right_boundary is None: + continue + if left_boundary.is_empty or right_boundary.is_empty: + continue + + # Use first lane's baseline path for direction alignment + baseline_path = None + if lane_ids: + first_lane_token = lane_ids[0] + if first_lane_token in nusc_map.arcline_path_3: + arc_path = nusc_map.arcline_path_3[first_lane_token] + try: + points = discretize_lane(arc_path, resolution_meters=0.1) + xy_points = [(p[0], p[1]) for p in points] + baseline_path = LineString(xy_points) + except Exception: + baseline_path = None + + if baseline_path and left_boundary: + left_boundary = align_boundary_direction(baseline_path, left_boundary) + if baseline_path and right_boundary: + right_boundary = align_boundary_direction(baseline_path, right_boundary) + + try: + map_writer.write_lane_group( + CacheLaneGroup( + object_id=token, + lane_ids=lane_ids, + left_boundary=left_boundary, + right_boundary=right_boundary, + intersection_id=None, # Handled in intersections + predecessor_ids=incoming, + successor_ids=outgoing, + outline=None, + geometry=polygon, + ) + ) + except Exception: + continue + + +def _write_nuscenes_intersections(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write intersection data to map_writer. + """ + road_blocks = nusc_map.road_block + for block in road_blocks: + token = block["token"] + try: + if "polygon_token" in block: + polygon = nusc_map.extract_polygon(block["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + # Lane group IDs are not directly available; use empty list + lane_group_ids = [] + + map_writer.write_intersection( + CacheIntersection( + object_id=token, + lane_group_ids=lane_group_ids, + geometry=polygon, + ) + ) + + +def _write_nuscenes_crosswalks(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write crosswalk data to map_writer. + """ + ped_crossings = nusc_map.ped_crossing + for crossing in ped_crossings: + token = crossing["token"] + try: + if "polygon_token" in crossing: + polygon = nusc_map.extract_polygon(crossing["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + map_writer.write_crosswalk( + CacheCrosswalk( + object_id=token, + geometry=polygon, + ) + ) + + +def _write_nuscenes_walkways(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write walkway data to map_writer. + """ + walkways = nusc_map.walkway + for walkway in walkways: + token = walkway["token"] + try: + if "polygon_token" in walkway: + polygon = nusc_map.extract_polygon(walkway["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + map_writer.write_walkway( + CacheWalkway( + object_id=token, + geometry=polygon, + ) + ) + + +def _write_nuscenes_carparks(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write carpark data to map_writer. + """ + carpark_areas = nusc_map.carpark_area + for carpark in carpark_areas: + token = carpark["token"] + try: + if "polygon_token" in carpark: + polygon = nusc_map.extract_polygon(carpark["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + map_writer.write_carpark( + CacheCarpark( + object_id=token, + geometry=polygon, + ) + ) + + +def _write_nuscenes_generic_drivables(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write generic drivable areas to map_writer. + """ + # Combine road segments, lanes, and drivable areas + all_drivables = [] + + # Add road segments + for segment in nusc_map.road_segment: + try: + if "polygon_token" in segment: + polygon = nusc_map.extract_polygon(segment["polygon_token"]) + if polygon.is_valid: + all_drivables.append((f"road_segment_{segment['token']}", polygon)) + except Exception: + continue + + # Add lanes + for lane in nusc_map.lane: + try: + if "polygon_token" in lane: + polygon = nusc_map.extract_polygon(lane["polygon_token"]) + if polygon.is_valid: + all_drivables.append((f"lane_{lane['token']}", polygon)) + except Exception: + continue + + # Add drivable areas + for road in nusc_map.drivable_area: + try: + if "polygon_token" in road: + polygon = nusc_map.extract_polygon(road["polygon_token"]) + if polygon.is_valid: + all_drivables.append((f"road_{road['token']}", polygon)) + except Exception: + continue + + for obj_id, geometry in all_drivables: + map_writer.write_generic_drivable( + CacheGenericDrivable( + object_id=obj_id, + geometry=geometry, + ) + ) + + +def _write_nuscenes_stop_lines(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write stop line data to map_writer. + """ + stop_lines = nusc_map.stop_line + for stop_line in stop_lines: + token = stop_line["token"] + try: + if "polygon_token" in stop_line: + polygon = nusc_map.extract_polygon(stop_line["polygon_token"]) + else: + continue + if not polygon.is_valid: + continue + except Exception: + continue + + # Note: Stop lines are written as generic drivable for compatibility + map_writer.write_generic_drivable( + CacheGenericDrivable( + object_id=token, + geometry=polygon, + ) + ) + + +def _write_nuscenes_road_lines(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """ + Write road line data (dividers) to map_writer. + """ + # Process road dividers + road_dividers = nusc_map.road_divider + for divider in road_dividers: + token = divider["token"] + try: + line = nusc_map.extract_line(divider["line_token"]) + if not line.is_valid: + continue + except Exception: + continue + + # Determine line type + line_type = _get_road_line_type(divider["line_token"], nusc_map) + + map_writer.write_road_line( + CacheRoadLine( + object_id=token, + road_line_type=line_type, + polyline=Polyline3D(LineString(line.coords)), + ) + ) + + # Process lane dividers + lane_dividers = nusc_map.lane_divider + for divider in lane_dividers: + token = divider["token"] + try: + line = nusc_map.extract_line(divider["line_token"]) + if not line.is_valid: + continue + except Exception: + continue + + line_type = _get_road_line_type(divider["line_token"], nusc_map) + + map_writer.write_road_line( + CacheRoadLine( + object_id=token, + road_line_type=line_type, + polyline=Polyline3D(LineString(line.coords)), + ) + ) + + +def _get_lane_boundary(lane_token: str, side: str, nusc_map: NuScenesMap) -> Optional[LineString]: + """ + Extract lane boundary geometry for a given side. + """ + lane_record = next((lr for lr in nusc_map.lane if lr["token"] == lane_token), None) + if not lane_record: + return None + + divider_segment_nodes_key = f"{side}_lane_divider_segment_nodes" + if divider_segment_nodes_key in lane_record and lane_record[divider_segment_nodes_key]: + nodes = lane_record[divider_segment_nodes_key] + boundary = LineString([(node['x'], node['y']) for node in nodes]) + return boundary + + return None + + +def _get_lane_group_boundary(segment_token: str, side: str, nusc_map: NuScenesMap) -> Optional[LineString]: + """ + Extract lane group boundary geometry (simplified). + """ + # This is a simplified implementation; in practice, may need more robust geometry extraction + boundary_type = "road_divider" if side == "left" else "lane_divider" + + # Find the segment geometry + segment = next((rs for rs in nusc_map.road_segment if rs["token"] == segment_token), None) + if not segment: + return None + + try: + segment_geom = nusc_map.extract_polygon(segment["polygon_token"]) + except Exception: + return None + + # Find nearest boundary of the specified type within a threshold + nearest = None + min_dist = float('inf') + + if boundary_type == "road_divider": + records = nusc_map.road_divider + else: + records = nusc_map.lane_divider + + for record in records: + try: + line = nusc_map.extract_line(record["line_token"]) + dist = segment_geom.distance(line) + if dist < 10.0 and dist < min_dist: + min_dist = dist + nearest = line + except Exception: + continue + + return nearest + + +def _get_connected_segments(segment_token: str, nusc_map: NuScenesMap): + """ + Get incoming and outgoing segment connections. + """ + incoming, outgoing = [], [] + + for connector in nusc_map.lane_connector: + if connector.get("outgoing_lane") == segment_token: + incoming.append(connector.get("incoming_lane")) + elif connector.get("incoming_lane") == segment_token: + outgoing.append(connector.get("outgoing_lane")) + + incoming = [id for id in incoming if id is not None] + outgoing = [id for id in outgoing if id is not None] + + return incoming, outgoing + + +def _get_road_line_type(line_token: str, nusc_map: NuScenesMap) -> RoadLineType: + """ + Map nuscenes line type to RoadLineType. + """ + nuscenes_to_road_line_type = { + "SINGLE_SOLID_WHITE": RoadLineType.SOLID_WHITE, + "DOUBLE_DASHED_WHITE": RoadLineType.DOUBLE_DASH_WHITE, + "SINGLE_SOLID_YELLOW": RoadLineType.SOLID_YELLOW, + } + + line_token_to_type = {} + for lane_record in nusc_map.lane: + for seg in lane_record.get("left_lane_divider_segments", []): + token = seg.get("line_token") + seg_type = seg.get("segment_type") + if token and seg_type: + line_token_to_type[token] = seg_type + + for seg in lane_record.get("right_lane_divider_segments", []): + token = seg.get("line_token") + seg_type = seg.get("segment_type") + if token and seg_type: + line_token_to_type[token] = seg_type + + nuscenes_type = line_token_to_type.get(line_token, "UNKNOWN") + return nuscenes_to_road_line_type.get(nuscenes_type, RoadLineType.UNKNOWN) + + +def flip_linestring(linestring: LineString) -> LineString: + """ + Flip the direction of a LineString. + """ + return LineString(linestring.coords[::-1]) + + +def lines_same_direction(centerline: LineString, boundary: LineString) -> bool: + """ + Check if centerline and boundary have the same direction. + """ + center_start = np.array(centerline.coords[0]) + center_end = np.array(centerline.coords[-1]) + boundary_start = np.array(boundary.coords[0]) + boundary_end = np.array(boundary.coords[-1]) + + same_dir_dist = np.linalg.norm(center_start - boundary_start) + np.linalg.norm(center_end - boundary_end) + opposite_dir_dist = np.linalg.norm(center_start - boundary_end) + np.linalg.norm(center_end - boundary_start) + + return same_dir_dist <= opposite_dir_dist + + +def align_boundary_direction(centerline: LineString, boundary: LineString) -> LineString: + """ + Align boundary direction with centerline. + """ + if not lines_same_direction(centerline, boundary): + return flip_linestring(boundary) + return boundary \ No newline at end of file diff --git a/src/py123d/datatypes/detections/box_detections.py b/src/py123d/datatypes/detections/box_detections.py index 2cff36a8..dccaa73b 100644 --- a/src/py123d/datatypes/detections/box_detections.py +++ b/src/py123d/datatypes/detections/box_detections.py @@ -97,8 +97,8 @@ def get_box_detections_by_types(self, detection_types: Iterable[BoxDetectionType detection for detection in self.box_detections if detection.metadata.box_detection_type in detection_types ] - def get_detection_by_track_token(self, track_token: str) -> BoxDetection | None: - box_detection: BoxDetection | None = None + def get_detection_by_track_token(self, track_token: str) -> Optional[BoxDetection]: + box_detection: Optional[BoxDetection] = None for detection in self.box_detections: if detection.metadata.track_token == track_token: box_detection = detection diff --git a/src/py123d/datatypes/maps/abstract_map_objects.py b/src/py123d/datatypes/maps/abstract_map_objects.py index 83004a94..de43bc81 100644 --- a/src/py123d/datatypes/maps/abstract_map_objects.py +++ b/src/py123d/datatypes/maps/abstract_map_objects.py @@ -1,7 +1,9 @@ from __future__ import annotations import abc -from typing import List, Optional, Tuple, TypeAlias, Union +from typing import List, Optional, Tuple, Union +from typing_extensions import TypeAlias + import shapely.geometry as geom import trimesh diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py b/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py index 97b11d73..a7e0dc36 100644 --- a/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py +++ b/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py @@ -1,6 +1,7 @@ from __future__ import annotations import ast +import trimesh from functools import cached_property from typing import List, Optional, Union @@ -193,9 +194,23 @@ def centerline(self) -> Polyline3D: @property def outline_3d(self) -> Polyline3D: - """Inherited, see superclass.""" - outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) - outline_array = np.vstack((outline_array, outline_array[0])) + left_array = self.left_boundary.array if getattr(self, "left_boundary", None) is not None else np.zeros((0, 3)) + right_array = self.right_boundary.array[::-1] if getattr(self, "right_boundary", None) is not None else np.zeros((0, 3)) + + outline_array = np.vstack((left_array, right_array)) if left_array.size + right_array.size > 0 else np.zeros((0, 3)) + + if outline_array.shape[0] == 0: + # fallback: use shapely polygon generate Polyline3D + poly = getattr(self, "shapely_polygon", None) + if poly is not None: + outline_array = np.array(poly.exterior.coords) + else: + return Polyline3D(np.zeros((0, 3))) + + # close + if outline_array.shape[0] > 0: + outline_array = np.vstack((outline_array, outline_array[0])) + return Polyline3D.from_linestring(geom.LineString(outline_array)) @property @@ -264,8 +279,20 @@ def right_boundary(self) -> Polyline3D: @property def outline_3d(self) -> Polyline3D: - """Inherited, see superclass.""" - outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) + # get left_array and right_array + left_array = self.left_boundary.array if getattr(self, "left_boundary", None) is not None else np.zeros((0, 3)) + right_array = self.right_boundary.array[::-1] if getattr(self, "right_boundary", None) is not None else np.zeros((0, 3)) + + if left_array.size + right_array.size == 0: + # fallback: use geometry polygon generate + poly = getattr(self, "shapely_polygon", None) + if poly is not None: + outline_array = np.array(poly.exterior.coords) + else: + return Polyline3D(np.zeros((0, 3))) + else: + outline_array = np.vstack((left_array, right_array)) + return Polyline3D.from_linestring(geom.LineString(outline_array)) @property diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_utils.py b/src/py123d/datatypes/maps/gpkg/gpkg_utils.py index 54dd93e6..ba587077 100644 --- a/src/py123d/datatypes/maps/gpkg/gpkg_utils.py +++ b/src/py123d/datatypes/maps/gpkg/gpkg_utils.py @@ -21,16 +21,23 @@ def load_gdf_with_geometry_columns(gdf: gpd.GeoDataFrame, geometry_column_names: def get_all_rows_with_value( - elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str -) -> gpd.geodataframe.GeoDataFrame: + elements: gpd.GeoDataFrame, column_label: str, desired_value +) -> gpd.GeoDataFrame: """ - Extract all matching elements. Note, if no matching desired_key is found and empty list is returned. - :param elements: data frame from MapsDb. - :param column_label: key to extract from a column. - :param desired_value: key which is compared with the values of column_label entry. - :return: a subset of the original GeoDataFrame containing the matching key. + Extract all matching elements by value. + Automatically handles both integer IDs and UUID strings. """ - return elements.iloc[np.where(elements[column_label].to_numpy().astype(int) == int(desired_value))] + # If the column is of integer type, attempt to convert the desired_value to an integer before comparison. + col_dtype = elements[column_label].dtype + if np.issubdtype(col_dtype, np.integer): + try: + desired_value_int = int(desired_value) + return elements[elements[column_label] == desired_value_int] + except ValueError: + raise ValueError(f"Expected an integer value for column '{column_label}', got '{desired_value}'") + else: + # Otherwise, directly compare it as a string. + return elements[elements[column_label].astype(str) == str(desired_value)] def get_row_with_value(elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str) -> gpd.GeoSeries: diff --git a/src/py123d/datatypes/scene/arrow/arrow_scene.py b/src/py123d/datatypes/scene/arrow/arrow_scene.py index 236862d5..7d89b786 100644 --- a/src/py123d/datatypes/scene/arrow/arrow_scene.py +++ b/src/py123d/datatypes/scene/arrow/arrow_scene.py @@ -36,10 +36,17 @@ def __init__( self._arrow_file_path: Path = Path(arrow_file_path) self._log_metadata: LogMetadata = get_log_metadata_from_arrow(arrow_file_path) + with pa.memory_map(str(self._arrow_file_path), "r") as source: + reader = pa.ipc.open_file(source) + table = reader.read_all() + num_rows = table.num_rows + initial_uuid = table['uuid'][0].as_py() + if scene_extraction_metadata is None: scene_extraction_metadata = SceneExtractionMetadata( + initial_uuid=initial_uuid, initial_idx=0, - duration_s=self._log_metadata.timestep_seconds * len(self._arrow_file_path), + duration_s=self._log_metadata.timestep_seconds * num_rows, history_s=0.0, iteration_duration_s=self._log_metadata.timestep_seconds, ) diff --git a/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py b/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py index 6be9d433..031c726a 100644 --- a/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py +++ b/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py @@ -38,6 +38,7 @@ "av2-sensor": DATASET_PATHS.av2_sensor_data_root, "wopd": DATASET_PATHS.wopd_data_root, "pandaset": DATASET_PATHS.pandaset_data_root, + "nuscenes": DATASET_PATHS.nuscenes_data_root, } diff --git a/src/py123d/datatypes/sensors/lidar/lidar_index.py b/src/py123d/datatypes/sensors/lidar/lidar_index.py index 7684b685..c101ba8e 100644 --- a/src/py123d/datatypes/sensors/lidar/lidar_index.py +++ b/src/py123d/datatypes/sensors/lidar/lidar_index.py @@ -83,3 +83,12 @@ class PandasetLidarIndex(LiDARIndex): Y = 1 Z = 2 INTENSITY = 3 + + +@register_lidar_index +class NuscenesLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + RING = 4 diff --git a/src/py123d/datatypes/vehicle_state/vehicle_parameters.py b/src/py123d/datatypes/vehicle_state/vehicle_parameters.py index 92ce83b0..19b050c7 100644 --- a/src/py123d/datatypes/vehicle_state/vehicle_parameters.py +++ b/src/py123d/datatypes/vehicle_state/vehicle_parameters.py @@ -52,6 +52,16 @@ def get_nuplan_chrysler_pacifica_parameters() -> VehicleParameters: rear_axle_to_center_longitudinal=1.461, ) +def get_nuscenes_renauly_zoe_parameters() -> VehicleParameters: + return VehicleParameters( + vehicle_name="nuscenes_renauly_zoe", + width=1.730, + length=4.084, + height=1.562, + wheel_base=2.588, + rear_axle_to_center_vertical=1.562 / 2, # NOTE: missing in nuscenes, TODO: find more accurate value + rear_axle_to_center_longitudinal=1.385, + ) def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: # NOTE: values are extracted from CARLA diff --git a/src/py123d/script/config/common/default_dataset_paths.yaml b/src/py123d/script/config/common/default_dataset_paths.yaml index ded971a6..c6a57f75 100644 --- a/src/py123d/script/config/common/default_dataset_paths.yaml +++ b/src/py123d/script/config/common/default_dataset_paths.yaml @@ -18,6 +18,8 @@ dataset_paths: # WOPD defaults wopd_data_root: ${oc.env:WOPD_DATA_ROOT,null} - # Pandaset defaults pandaset_data_root: ${oc.env:PANDASET_DATA_ROOT,null} + + # nuScenes defaults + nuscenes_data_root: ${oc.env:NUSCENES_DATA_ROOT,null} diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml new file mode 100644 index 00000000..3e3a8b92 --- /dev/null +++ b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml @@ -0,0 +1,35 @@ +nuscenes_dataset: + _target_: py123d.conversion.datasets.nuscenes.nuscenes_data_converter.NuScenesDataConverter + _convert_: 'all' + + splits: ["nuscenes_train", "nuscenes_val", "nuscenes_test"] + nuscenes_data_root: ${dataset_paths.nuscenes_data_root} + nuscenes_map_root: ${dataset_paths.nuscenes_data_root} + nuscenes_lanelet2_root: ${dataset_paths.nuscenes_data_root}/lanelet2 + use_lanelet2: False + + dataset_converter_config: + _target_: py123d.conversion.dataset_converter_config.DatasetConverterConfig + _convert_: 'all' + + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + # Map + include_map: true + + # Ego + include_ego: true + + # Box Detections + include_box_detections: true + + # Traffic Lights + include_traffic_lights: false + + #cameras + include_cameras: true + camera_store_option: "path" + + #lidar + include_lidars: true + lidar_store_option: "path" \ No newline at end of file From 8e1fbf6b7aa6d141ef9fe401732b9783992d4b23 Mon Sep 17 00:00:00 2001 From: Changhui Jing <1903025786@qq.com> Date: Thu, 30 Oct 2025 15:09:27 +0000 Subject: [PATCH 2/4] add nuscenes sensor io. --- pyproject.toml | 12 ++ .../nuscenes/nuscenes_data_converter.py | 114 ++++++++++-------- .../datasets/nuscenes/nuscenes_sensor_io.py | 13 ++ .../registry/lidar_index_registry.py | 8 ++ .../sensor_io/lidar/file_lidar_io.py | 5 + .../config/common/default_dataset_paths.yaml | 2 + .../conversion/datasets/nuscenes_dataset.yaml | 1 - .../config/conversion/default_conversion.yaml | 2 +- 8 files changed, 102 insertions(+), 55 deletions(-) create mode 100644 src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py diff --git a/pyproject.toml b/pyproject.toml index 6df6dff0..5f305b9e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -95,6 +95,18 @@ waymo = [ "waymo-open-dataset-tf-2-12-0==1.6.6", ] +nuscenes = [ + "nuscenes-devkit==1.2.0", + "pycocotools==2.0.10", + "laspy==2.6.1", + "embreex==2.17.7.post6", + "lanelet2==1.2.2", + "protobuf==4.25.3", + "pycollada==0.9.2", + "vhacdx==0.0.8.post2", + "yourdfpy==0.0.58", +] + [tool.setuptools.packages.find] where = ["src"] diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py index eb743fb0..62ddc0d4 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py @@ -13,17 +13,17 @@ from nuscenes.utils.splits import create_splits_scenes from pyquaternion import Quaternion -from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter +from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter from py123d.script.builders.worker_pool_builder import WorkerPool -from py123d.datatypes.detections.detection import ( +from py123d.datatypes.detections.box_detections import ( BoxDetectionSE3, BoxDetectionWrapper, - BoxDetectionMetadata, - TrafficLightDetection, + BoxDetectionMetadata,) +from py123d.datatypes.detections.traffic_light_detections import (TrafficLightDetection, TrafficLightDetectionWrapper, ) -from py123d.datatypes.detections.detection_types import DetectionType +from py123d.datatypes.detections.box_detection_types import BoxDetectionType from py123d.datatypes.maps.map_metadata import MapMetadata from py123d.datatypes.scene.scene_metadata import LogMetadata from py123d.datatypes.sensors.camera.pinhole_camera import ( @@ -33,7 +33,7 @@ PinholeIntrinsics, ) from py123d.datatypes.sensors.lidar.lidar import LiDARMetadata, LiDARType -from py123d.conversion.utils.sensor_utils.lidar_index_registry import NuscenesLidarIndex +from py123d.datatypes.sensors.lidar.lidar_index import NuscenesLidarIndex from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from py123d.datatypes.vehicle_state.vehicle_parameters import get_nuscenes_renauly_zoe_parameters @@ -50,37 +50,37 @@ SORT_BY_TIMESTAMP: Final[bool] = True NUSCENES_DETECTION_NAME_DICT = { # Vehicles (4+ wheels) - "vehicle.car": DetectionType.VEHICLE, - "vehicle.truck": DetectionType.VEHICLE, - "vehicle.bus.bendy": DetectionType.VEHICLE, - "vehicle.bus.rigid": DetectionType.VEHICLE, - "vehicle.construction": DetectionType.VEHICLE, - "vehicle.emergency.ambulance": DetectionType.VEHICLE, - "vehicle.emergency.police": DetectionType.VEHICLE, - "vehicle.trailer": DetectionType.VEHICLE, + "vehicle.car": BoxDetectionType.VEHICLE, + "vehicle.truck": BoxDetectionType.VEHICLE, + "vehicle.bus.bendy": BoxDetectionType.VEHICLE, + "vehicle.bus.rigid": BoxDetectionType.VEHICLE, + "vehicle.construction": BoxDetectionType.VEHICLE, + "vehicle.emergency.ambulance": BoxDetectionType.VEHICLE, + "vehicle.emergency.police": BoxDetectionType.VEHICLE, + "vehicle.trailer": BoxDetectionType.VEHICLE, # Bicycles / Motorcycles - "vehicle.bicycle": DetectionType.BICYCLE, - "vehicle.motorcycle": DetectionType.BICYCLE, + "vehicle.bicycle": BoxDetectionType.BICYCLE, + "vehicle.motorcycle": BoxDetectionType.BICYCLE, # Pedestrians (all subtypes) - "human.pedestrian.adult": DetectionType.PEDESTRIAN, - "human.pedestrian.child": DetectionType.PEDESTRIAN, - "human.pedestrian.construction_worker": DetectionType.PEDESTRIAN, - "human.pedestrian.personal_mobility": DetectionType.PEDESTRIAN, - "human.pedestrian.police_officer": DetectionType.PEDESTRIAN, - "human.pedestrian.stroller": DetectionType.PEDESTRIAN, - "human.pedestrian.wheelchair": DetectionType.PEDESTRIAN, + "human.pedestrian.adult": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.child": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.construction_worker": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.personal_mobility": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.police_officer": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.stroller": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.wheelchair": BoxDetectionType.PEDESTRIAN, # Traffic cone / barrier - "movable_object.trafficcone": DetectionType.TRAFFIC_CONE, - "movable_object.barrier": DetectionType.BARRIER, + "movable_object.trafficcone": BoxDetectionType.TRAFFIC_CONE, + "movable_object.barrier": BoxDetectionType.BARRIER, # Generic objects - "movable_object.pushable_pullable": DetectionType.GENERIC_OBJECT, - "movable_object.debris": DetectionType.GENERIC_OBJECT, - "static_object.bicycle_rack": DetectionType.GENERIC_OBJECT, - "animal": DetectionType.GENERIC_OBJECT, + "movable_object.pushable_pullable": BoxDetectionType.GENERIC_OBJECT, + "movable_object.debris": BoxDetectionType.GENERIC_OBJECT, + "static_object.bicycle_rack": BoxDetectionType.GENERIC_OBJECT, + "animal": BoxDetectionType.GENERIC_OBJECT, } NUSCENES_CAMERA_TYPES = { @@ -99,7 +99,6 @@ def __init__( self, splits: List[str], nuscenes_data_root: Union[Path, str], - nuscenes_map_root: Union[Path, str], nuscenes_lanelet2_root: Union[Path, str], use_lanelet2: bool, dataset_converter_config: DatasetConverterConfig, @@ -108,7 +107,6 @@ def __init__( super().__init__(dataset_converter_config) self._splits: List[str] = splits self._nuscenes_data_root: Path = Path(nuscenes_data_root) - self._nuscenes_map_root: Path = Path(nuscenes_map_root) self._nuscenes_lanelet2_root: Path = Path(nuscenes_lanelet2_root) self._use_lanelet2 = use_lanelet2 self._version = version @@ -165,7 +163,7 @@ def convert_map(self, map_index: int, map_writer: AbstractMapWriter) -> None: if map_needs_writing: write_nuscenes_map( - nuscenes_maps_root=self._nuscenes_map_root, + nuscenes_maps_root=self._nuscenes_data_root, location=map_name, map_writer=map_writer, use_lanelet2=self._use_lanelet2, @@ -423,7 +421,7 @@ def _extract_nuscenes_box_detections( velocity_3d = Vector3D(x=velocity[0], y=velocity[1], z=velocity[2] if len(velocity) > 2 else 0.0) metadata = BoxDetectionMetadata( - detection_type=det_type, + box_detection_type=det_type, track_token=ann["instance_token"], timepoint=TimePoint.from_us(sample["timestamp"]), confidence=1.0, # nuScenes annotations are ground truth @@ -446,9 +444,9 @@ def _extract_nuscenes_traffic_lights() -> TrafficLightDetectionWrapper: def _extract_nuscenes_cameras( - nusc: NuScenes, - sample: Dict[str, Any], - dataset_converter_config: DatasetConverterConfig, + nusc: NuScenes, + sample: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, ) -> Dict[PinholeCameraType, Tuple[Union[str, bytes], StateSE3]]: camera_dict: Dict[PinholeCameraType, Tuple[Union[str, bytes], StateSE3]] = {} @@ -474,7 +472,8 @@ def _extract_nuscenes_cameras( if cam_path.exists() and cam_path.is_file(): if dataset_converter_config.camera_store_option == "path": - camera_data = str(cam_path) + # TODO: should be relative path + camera_data = cam_data["filename"] elif dataset_converter_config.camera_store_option == "binary": with open(cam_path, "rb") as f: camera_data = f.read() @@ -490,8 +489,8 @@ def _extract_nuscenes_lidars( nusc: NuScenes, sample: Dict[str, Any], dataset_converter_config: DatasetConverterConfig, -) -> Dict[LiDARType, Optional[str]]: - lidar_dict: Dict[LiDARType, Optional[str]] = {} +) -> List[LiDARData]: + lidars: List[LiDARData] = [] if dataset_converter_config.include_lidars: lidar_token = sample["data"]["LIDAR_TOP"] @@ -499,12 +498,21 @@ def _extract_nuscenes_lidars( lidar_path = NUSCENES_DATA_ROOT / lidar_data["filename"] if lidar_path.exists() and lidar_path.is_file(): - lidar_dict[LiDARType.LIDAR_MERGED] = str(lidar_path) + lidar = LiDARData( + lidar_type=LiDARType.LIDAR_TOP, + relative_path=str(lidar_path), + dataset_root=NUSCENES_DATA_ROOT, + iteration=lidar_data.get("iteration"), + ) + lidars.append(lidar) else: - lidar_dict[LiDARType.LIDAR_MERGED] = None - - return lidar_dict + lidars.append(LiDARData( + lidar_type=LiDARType.LIDAR_TOP, + relative_path=None, + dataset_root=NUSCENES_DATA_ROOT, + )) + return lidars def _extract_nuscenes_scenario_tag() -> List[str]: """nuScenes doesn't have scenario tags.""" @@ -518,9 +526,9 @@ def _extract_nuscenes_route_lane_group_ids() -> List[int]: # Updated arrow conversion function using the new extraction functions def convert_nuscenes_log_to_arrow( - args: List[Dict[str, Union[str, List[str]]]], - dataset_converter_config: DatasetConverterConfig, - version: str + args: List[Dict[str, Union[str, List[str]]]], + dataset_converter_config: DatasetConverterConfig, + version: str ) -> List[Any]: for log_info in args: scene_token: str = log_info["scene_token"] @@ -593,11 +601,11 @@ def convert_nuscenes_log_to_arrow( def _write_arrow_table_with_new_interface( - nusc: NuScenes, - scene: Dict[str, Any], - recording_schema: pa.schema, - log_file_path: Path, - dataset_converter_config: DatasetConverterConfig, + nusc: NuScenes, + scene: Dict[str, Any], + recording_schema: pa.schema, + log_file_path: Path, + dataset_converter_config: DatasetConverterConfig, ) -> None: can_bus = NuScenesCanBus(dataroot=str(NUSCENES_DATA_ROOT)) @@ -635,7 +643,7 @@ def _write_arrow_table_with_new_interface( "detections_state": detections_state_list, "detections_velocity": [det.velocity.array.tolist() for det in box_detections.box_detections], "detections_token": [det.metadata.track_token for det in box_detections.box_detections], - "detections_type": [det.metadata.detection_type.value for det in box_detections.box_detections], + "detections_type": [det.metadata.box_detection_type.value for det in box_detections.box_detections], "traffic_light_ids": [], "traffic_light_types": [], "scenario_tag": ["unknown"], @@ -668,4 +676,4 @@ def _write_arrow_table_with_new_interface( 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) \ No newline at end of file + write_arrow_table(recording_table, log_file_path) diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py new file mode 100644 index 00000000..82f1ac72 --- /dev/null +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py @@ -0,0 +1,13 @@ +import numpy as np +from pathlib import Path +from typing import Dict +from py123d.datatypes.sensors.lidar.lidar import LiDARType + +def load_nuscenes_lidar_pcs_from_file(pcd_path: Path) -> Dict[LiDARType, np.ndarray]: + points = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, 5) + + lidar_pcs_dict: Dict[LiDARType, np.ndarray] = { + LiDARType.LIDAR_TOP: points + } + + return lidar_pcs_dict diff --git a/src/py123d/conversion/registry/lidar_index_registry.py b/src/py123d/conversion/registry/lidar_index_registry.py index 50f5f8a3..c5213d2a 100644 --- a/src/py123d/conversion/registry/lidar_index_registry.py +++ b/src/py123d/conversion/registry/lidar_index_registry.py @@ -82,3 +82,11 @@ class PandasetLidarIndex(LiDARIndex): Y = 1 Z = 2 INTENSITY = 3 + +@register_lidar_index +class NuscenesLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + RING = 4 diff --git a/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py b/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py index edc5c7d5..83d4da21 100644 --- a/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py +++ b/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py @@ -15,6 +15,7 @@ "av2-sensor": DATASET_PATHS.av2_sensor_data_root, "wopd": DATASET_PATHS.wopd_data_root, "pandaset": DATASET_PATHS.pandaset_data_root, + "nuscenes": DATASET_PATHS.nuscenes_sensor_root, } @@ -56,6 +57,10 @@ def load_lidar_pcs_from_file( from py123d.conversion.datasets.pandaset.pandaset_sensor_io import load_pandaset_lidars_pcs_from_file lidar_pcs_dict = load_pandaset_lidars_pcs_from_file(full_lidar_path, index) + + elif log_metadata.dataset == "nuscenes": + from py123d.conversion.datasets.nuscenes.nuscenes_sensor_io import load_nuscenes_lidar_pcs_from_file + lidar_pcs_dict = load_nuscenes_lidar_pcs_from_file(full_lidar_path) else: raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") diff --git a/src/py123d/script/config/common/default_dataset_paths.yaml b/src/py123d/script/config/common/default_dataset_paths.yaml index c6a57f75..64e0a8c0 100644 --- a/src/py123d/script/config/common/default_dataset_paths.yaml +++ b/src/py123d/script/config/common/default_dataset_paths.yaml @@ -23,3 +23,5 @@ dataset_paths: # nuScenes defaults nuscenes_data_root: ${oc.env:NUSCENES_DATA_ROOT,null} + nuscenes_map_root: ${dataset_paths.nuscenes_data_root} + nuscenes_sensor_root: ${dataset_paths.nuscenes_data_root}/samples/LIDAR_TOP \ No newline at end of file diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml index 3e3a8b92..410b910e 100644 --- a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml @@ -4,7 +4,6 @@ nuscenes_dataset: splits: ["nuscenes_train", "nuscenes_val", "nuscenes_test"] nuscenes_data_root: ${dataset_paths.nuscenes_data_root} - nuscenes_map_root: ${dataset_paths.nuscenes_data_root} nuscenes_lanelet2_root: ${dataset_paths.nuscenes_data_root}/lanelet2 use_lanelet2: False diff --git a/src/py123d/script/config/conversion/default_conversion.yaml b/src/py123d/script/config/conversion/default_conversion.yaml index aecfa9b7..daa55f12 100644 --- a/src/py123d/script/config/conversion/default_conversion.yaml +++ b/src/py123d/script/config/conversion/default_conversion.yaml @@ -22,5 +22,5 @@ defaults: terminate_on_exception: True -force_map_conversion: False +force_map_conversion: True force_log_conversion: True From c42b7388f3754bd9c9fa28729a0eb341e0dce852 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Thu, 30 Oct 2025 20:31:16 +0100 Subject: [PATCH 3/4] Add some minor changes and refactorings. --- notebooks/bev_matplotlib.ipynb | 14 +- pyproject.toml | 15 +- .../conversion/nuscenes_mini_conversion.sh | 3 + .../datasets/nuscenes/nuscenes_constants.py | 60 +++ ...ata_converter.py => nuscenes_converter.py} | 416 ++++-------------- .../nuscenes/nuscenes_map_conversion.py | 155 ++++--- .../datasets/nuscenes/nuscenes_sensor_io.py | 27 +- .../registry/lidar_index_registry.py | 15 +- .../sensor_io/lidar/file_lidar_io.py | 7 +- .../scene/arrow/utils/arrow_getters.py | 4 +- .../datatypes/sensors/lidar/lidar_index.py | 13 +- .../vehicle_state/vehicle_parameters.py | 9 +- .../config/common/default_dataset_paths.yaml | 2 +- .../conversion/datasets/nuscenes_dataset.yaml | 10 +- .../datasets/nuscenes_mini_dataset.yaml | 36 ++ test_viser.py | 13 +- 16 files changed, 351 insertions(+), 448 deletions(-) create mode 100644 scripts/conversion/nuscenes_mini_conversion.sh create mode 100644 src/py123d/conversion/datasets/nuscenes/nuscenes_constants.py rename src/py123d/conversion/datasets/nuscenes/{nuscenes_data_converter.py => nuscenes_converter.py} (50%) create mode 100644 src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml diff --git a/notebooks/bev_matplotlib.ipynb b/notebooks/bev_matplotlib.ipynb index bddd36d6..bd2dd882 100644 --- a/notebooks/bev_matplotlib.ipynb +++ b/notebooks/bev_matplotlib.ipynb @@ -24,7 +24,7 @@ "\n", "# splits = [\"wopd_val\"]\n", "# splits = [\"carla_test\"]\n", - "splits = [\"nuplan-mini_test\"]\n", + "splits = [\"nuscenes-mini_val\", \"nuscenes-mini_train\"]\n", "# splits = [\"av2-sensor-mini_train\"]\n", "# splits = [\"pandaset_train\"]\n", "# log_names = None\n", @@ -38,7 +38,7 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_uuids=scene_uuids,\n", - " duration_s=6.0,\n", + " duration_s=None,\n", " history_s=0.0,\n", " timestamp_threshold_s=20,\n", " shuffle=True,\n", @@ -144,8 +144,8 @@ " route_lane_group_ids: Optional[List[int]] = None,\n", ") -> None:\n", " layers: List[MapLayer] = [\n", - " # MapLayer.LANE,\n", - " # MapLayer.LANE_GROUP,\n", + " MapLayer.LANE,\n", + " MapLayer.LANE_GROUP,\n", " # MapLayer.GENERIC_DRIVABLE,\n", " # MapLayer.CARPARK,\n", " # MapLayer.CROSSWALK,\n", @@ -214,11 +214,11 @@ " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", " if map_api is not None:\n", " add_debug_map_on_ax(ax, scene.get_map_api(), point_2d, radius=radius, route_lane_group_ids=None)\n", - " add_default_map_on_ax(ax, map_api, point_2d, radius=radius, route_lane_group_ids=None)\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, scene.get_map_api())\n", "\n", - " # add_box_detections_to_ax(ax, box_detections)\n", - " # add_ego_vehicle_to_ax(ax, ego_vehicle_state)\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", diff --git a/pyproject.toml b/pyproject.toml index 5f305b9e..9284be71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,13 +89,11 @@ nuplan = [ "guppy3", "retry", ] -waymo = [ - "protobuf==6.30.2", - "tensorflow==2.13.0", - "waymo-open-dataset-tf-2-12-0==1.6.6", -] - nuscenes = [ + "lanelet2", + "nuscenes-devkit==1.2.0", +] +nuscenes_expanded = [ "nuscenes-devkit==1.2.0", "pycocotools==2.0.10", "laspy==2.6.1", @@ -106,6 +104,11 @@ nuscenes = [ "vhacdx==0.0.8.post2", "yourdfpy==0.0.58", ] +waymo = [ + "protobuf==6.30.2", + "tensorflow==2.13.0", + "waymo-open-dataset-tf-2-12-0==1.6.6", +] [tool.setuptools.packages.find] where = ["src"] diff --git a/scripts/conversion/nuscenes_mini_conversion.sh b/scripts/conversion/nuscenes_mini_conversion.sh new file mode 100644 index 00000000..b6d35aa4 --- /dev/null +++ b/scripts/conversion/nuscenes_mini_conversion.sh @@ -0,0 +1,3 @@ +export NUSCENES_DATA_ROOT="/home/daniel/nuscenes_mini/" + +py123d-conversion datasets=["nuscenes_mini_dataset"] diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_constants.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_constants.py new file mode 100644 index 00000000..dd04d91a --- /dev/null +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_constants.py @@ -0,0 +1,60 @@ +import os +from pathlib import Path +from typing import Final, List + +from py123d.datatypes.detections.box_detection_types import BoxDetectionType +from py123d.datatypes.sensors.camera.pinhole_camera import PinholeCameraType + +NUSCENES_MAPS: List[str] = ["boston-seaport", "singapore-hollandvillage", "singapore-onenorth", "singapore-queenstown"] + +NUSCENES_DATA_SPLITS: Final[List[str]] = [ + "nuscenes_train", + "nuscenes_val", + "nuscenes_test", + "nuscenes-mini_train", + "nuscenes-mini_val", +] + +TARGET_DT: Final[float] = 0.1 +NUSCENES_DT: Final[float] = 0.5 +SORT_BY_TIMESTAMP: Final[bool] = True +NUSCENES_DETECTION_NAME_DICT = { + # Vehicles (4+ wheels) + "vehicle.car": BoxDetectionType.VEHICLE, + "vehicle.truck": BoxDetectionType.VEHICLE, + "vehicle.bus.bendy": BoxDetectionType.VEHICLE, + "vehicle.bus.rigid": BoxDetectionType.VEHICLE, + "vehicle.construction": BoxDetectionType.VEHICLE, + "vehicle.emergency.ambulance": BoxDetectionType.VEHICLE, + "vehicle.emergency.police": BoxDetectionType.VEHICLE, + "vehicle.trailer": BoxDetectionType.VEHICLE, + # Bicycles / Motorcycles + "vehicle.bicycle": BoxDetectionType.BICYCLE, + "vehicle.motorcycle": BoxDetectionType.BICYCLE, + # Pedestrians (all subtypes) + "human.pedestrian.adult": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.child": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.construction_worker": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.personal_mobility": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.police_officer": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.stroller": BoxDetectionType.PEDESTRIAN, + "human.pedestrian.wheelchair": BoxDetectionType.PEDESTRIAN, + # Traffic cone / barrier + "movable_object.trafficcone": BoxDetectionType.TRAFFIC_CONE, + "movable_object.barrier": BoxDetectionType.BARRIER, + # Generic objects + "movable_object.pushable_pullable": BoxDetectionType.GENERIC_OBJECT, + "movable_object.debris": BoxDetectionType.GENERIC_OBJECT, + "static_object.bicycle_rack": BoxDetectionType.GENERIC_OBJECT, + "animal": BoxDetectionType.GENERIC_OBJECT, +} + +NUSCENES_CAMERA_TYPES = { + PinholeCameraType.CAM_F0: "CAM_FRONT", + PinholeCameraType.CAM_B0: "CAM_BACK", + PinholeCameraType.CAM_L0: "CAM_FRONT_LEFT", + PinholeCameraType.CAM_L1: "CAM_BACK_LEFT", + PinholeCameraType.CAM_R0: "CAM_FRONT_RIGHT", + PinholeCameraType.CAM_R1: "CAM_BACK_RIGHT", +} +NUSCENES_DATA_ROOT = Path(os.environ["NUSCENES_DATA_ROOT"]) diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py similarity index 50% rename from src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py rename to src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py index 62ddc0d4..483a6bc4 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_data_converter.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py @@ -1,29 +1,25 @@ import gc -import json -import os -import numpy as np -import pyarrow as pa - -from dataclasses import asdict from pathlib import Path -from typing import Any, Dict, Final, List, Optional, Tuple, Union -from nuscenes import NuScenes -from nuscenes.can_bus.can_bus_api import NuScenesCanBus -from nuscenes.utils.data_classes import Box -from nuscenes.utils.splits import create_splits_scenes +from typing import Any, Dict, List, Tuple, Union + +import numpy as np from pyquaternion import Quaternion +from py123d.common.utils.dependencies import check_dependencies +from py123d.conversion.abstract_dataset_converter import AbstractDatasetConverter +from py123d.conversion.dataset_converter_config import DatasetConverterConfig +from py123d.conversion.datasets.nuplan.nuplan_converter import TARGET_DT +from py123d.conversion.datasets.nuscenes.nuscenes_constants import ( + NUSCENES_CAMERA_TYPES, + NUSCENES_DATA_ROOT, + NUSCENES_DATA_SPLITS, + NUSCENES_DETECTION_NAME_DICT, + NUSCENES_DT, +) +from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import NUSCENES_MAPS, write_nuscenes_map from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.script.builders.worker_pool_builder import WorkerPool -from py123d.datatypes.detections.box_detections import ( - BoxDetectionSE3, - BoxDetectionWrapper, - BoxDetectionMetadata,) -from py123d.datatypes.detections.traffic_light_detections import (TrafficLightDetection, - TrafficLightDetectionWrapper, -) -from py123d.datatypes.detections.box_detection_types import BoxDetectionType +from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper from py123d.datatypes.maps.map_metadata import MapMetadata from py123d.datatypes.scene.scene_metadata import LogMetadata from py123d.datatypes.sensors.camera.pinhole_camera import ( @@ -33,119 +29,75 @@ PinholeIntrinsics, ) from py123d.datatypes.sensors.lidar.lidar import LiDARMetadata, LiDARType -from py123d.datatypes.sensors.lidar.lidar_index import NuscenesLidarIndex - -from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index -from py123d.datatypes.vehicle_state.vehicle_parameters import get_nuscenes_renauly_zoe_parameters -from py123d.geometry import StateSE3, BoundingBoxSE3, BoundingBoxSE3Index -from py123d.geometry.vector import Vector3D, Vector3DIndex -from py123d.common.utils.arrow_helper import open_arrow_table, write_arrow_table -from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import write_nuscenes_map, NUSCENES_MAPS -from py123d.conversion.abstract_dataset_converter import AbstractDatasetConverter -from py123d.conversion.dataset_converter_config import DatasetConverterConfig +from py123d.datatypes.sensors.lidar.lidar_index import NuScenesLidarIndex from py123d.datatypes.time.time_point import TimePoint +from py123d.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3 +from py123d.datatypes.vehicle_state.vehicle_parameters import get_nuscenes_renault_zoe_parameters +from py123d.geometry import BoundingBoxSE3, StateSE3 +from py123d.geometry.vector import Vector3D + +check_dependencies(["nuscenes"], "nuscenes") +from nuscenes import NuScenes +from nuscenes.can_bus.can_bus_api import NuScenesCanBus +from nuscenes.utils.data_classes import Box +from nuscenes.utils.splits import create_splits_scenes -TARGET_DT: Final[float] = 0.1 -NUSCENES_DT: Final[float] = 0.5 -SORT_BY_TIMESTAMP: Final[bool] = True -NUSCENES_DETECTION_NAME_DICT = { - # Vehicles (4+ wheels) - "vehicle.car": BoxDetectionType.VEHICLE, - "vehicle.truck": BoxDetectionType.VEHICLE, - "vehicle.bus.bendy": BoxDetectionType.VEHICLE, - "vehicle.bus.rigid": BoxDetectionType.VEHICLE, - "vehicle.construction": BoxDetectionType.VEHICLE, - "vehicle.emergency.ambulance": BoxDetectionType.VEHICLE, - "vehicle.emergency.police": BoxDetectionType.VEHICLE, - "vehicle.trailer": BoxDetectionType.VEHICLE, - - # Bicycles / Motorcycles - "vehicle.bicycle": BoxDetectionType.BICYCLE, - "vehicle.motorcycle": BoxDetectionType.BICYCLE, - - # Pedestrians (all subtypes) - "human.pedestrian.adult": BoxDetectionType.PEDESTRIAN, - "human.pedestrian.child": BoxDetectionType.PEDESTRIAN, - "human.pedestrian.construction_worker": BoxDetectionType.PEDESTRIAN, - "human.pedestrian.personal_mobility": BoxDetectionType.PEDESTRIAN, - "human.pedestrian.police_officer": BoxDetectionType.PEDESTRIAN, - "human.pedestrian.stroller": BoxDetectionType.PEDESTRIAN, - "human.pedestrian.wheelchair": BoxDetectionType.PEDESTRIAN, - - # Traffic cone / barrier - "movable_object.trafficcone": BoxDetectionType.TRAFFIC_CONE, - "movable_object.barrier": BoxDetectionType.BARRIER, - - # Generic objects - "movable_object.pushable_pullable": BoxDetectionType.GENERIC_OBJECT, - "movable_object.debris": BoxDetectionType.GENERIC_OBJECT, - "static_object.bicycle_rack": BoxDetectionType.GENERIC_OBJECT, - "animal": BoxDetectionType.GENERIC_OBJECT, -} - -NUSCENES_CAMERA_TYPES = { - PinholeCameraType.CAM_F0: "CAM_FRONT", - PinholeCameraType.CAM_B0: "CAM_BACK", - PinholeCameraType.CAM_L0: "CAM_FRONT_LEFT", - PinholeCameraType.CAM_L1: "CAM_BACK_LEFT", - PinholeCameraType.CAM_R0: "CAM_FRONT_RIGHT", - PinholeCameraType.CAM_R1: "CAM_BACK_RIGHT", -} -NUSCENES_DATA_ROOT = Path(os.environ["NUSCENES_DATA_ROOT"]) - - -class NuScenesDataConverter(AbstractDatasetConverter): + +class NuScenesConverter(AbstractDatasetConverter): def __init__( - self, - splits: List[str], - nuscenes_data_root: Union[Path, str], - nuscenes_lanelet2_root: Union[Path, str], - use_lanelet2: bool, - dataset_converter_config: DatasetConverterConfig, - version: str = "v1.0-trainval", + self, + splits: List[str], + nuscenes_data_root: Union[Path, str], + nuscenes_map_root: Union[Path, str], + nuscenes_lanelet2_root: Union[Path, str], + use_lanelet2: bool, + dataset_converter_config: DatasetConverterConfig, + version: str = "v1.0-mini", ) -> None: super().__init__(dataset_converter_config) + + for split in splits: + assert ( + split in NUSCENES_DATA_SPLITS + ), f"Split {split} is not available. Available splits: {NUSCENES_DATA_SPLITS}" + self._splits: List[str] = splits + self._nuscenes_data_root: Path = Path(nuscenes_data_root) + self._nuscenes_map_root: Path = Path(nuscenes_map_root) self._nuscenes_lanelet2_root: Path = Path(nuscenes_lanelet2_root) + self._use_lanelet2 = use_lanelet2 self._version = version self._scene_tokens_per_split: Dict[str, List[str]] = self._collect_scene_tokens() self._target_dt: float = TARGET_DT def _collect_scene_tokens(self) -> Dict[str, List[str]]: + scene_tokens_per_split: Dict[str, List[str]] = {} nusc = NuScenes(version=self._version, dataroot=str(self._nuscenes_data_root), verbose=False) + nuscenes_split_name_mapping = { + "nuscenes_train": "train", + "nuscenes_val": "val", + "nuscenes_test": "test", + "nuscenes-mini_train": "mini_train", + "nuscenes-mini_val": "mini_val", + } + scene_splits = create_splits_scenes() available_scenes = [scene for scene in nusc.scene] for split in self._splits: - # Map the split name to the division of nuScenes - nusc_split = split.replace("nuscenes_", "") - if nusc_split == "trainval": - scene_names = scene_splits['train'] + scene_splits['val'] - else: - scene_names = scene_splits.get(nusc_split, []) + nuscenes_split = nuscenes_split_name_mapping[split] + scene_names = scene_splits.get(nuscenes_split, []) # get token - scene_tokens = [ - scene['token'] for scene in available_scenes - if scene['name'] in scene_names - ] + scene_tokens = [scene["token"] for scene in available_scenes if scene["name"] in scene_names] scene_tokens_per_split[split] = scene_tokens return scene_tokens_per_split - def get_available_splits(self) -> List[str]: - return [ - "nuscenes_train", - "nuscenes_val", - "nuscenes_test", - "nuscenes_mini_train", - "nuscenes_mini_val", - ] - def get_number_of_maps(self) -> int: """Inherited, see superclass.""" return len(NUSCENES_MAPS) @@ -163,7 +115,7 @@ def convert_map(self, map_index: int, map_writer: AbstractMapWriter) -> None: if map_needs_writing: write_nuscenes_map( - nuscenes_maps_root=self._nuscenes_data_root, + nuscenes_maps_root=self._nuscenes_map_root, location=map_name, map_writer=map_writer, use_lanelet2=self._use_lanelet2, @@ -195,7 +147,7 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: log_name=scene["name"], location=log_record["location"], timestep_seconds=TARGET_DT, - vehicle_parameters=get_nuscenes_renauly_zoe_parameters(), + vehicle_parameters=get_nuscenes_renault_zoe_parameters(), camera_metadata=_get_nuscenes_camera_metadata(nusc, scene, self.dataset_converter_config), lidar_metadata=_get_nuscenes_lidar_metadata(nusc, scene, self.dataset_converter_config), map_metadata=_get_nuscenes_map_metadata(log_record["location"]), @@ -220,7 +172,6 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: timestamp=TimePoint.from_us(sample["timestamp"]), ego_state=_extract_nuscenes_ego_state(nusc, sample, can_bus), box_detections=_extract_nuscenes_box_detections(nusc, sample), - traffic_lights=_extract_nuscenes_traffic_lights(), # nuScenes doesn't have traffic lights cameras=_extract_nuscenes_cameras( nusc=nusc, sample=sample, @@ -231,9 +182,6 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: sample=sample, dataset_converter_config=self.dataset_converter_config, ), - scenario_tags=_extract_nuscenes_scenario_tag(), # nuScenes doesn't have scenario tags - route_lane_group_ids=_extract_nuscenes_route_lane_group_ids(), - # nuScenes doesn't have route info ) sample_token = sample["next"] @@ -243,18 +191,11 @@ def convert_log(self, log_index: int, log_writer: AbstractLogWriter) -> None: del nusc gc.collect() - def convert_logs(self, worker: WorkerPool) -> None: - """ - NuScenes logs conversion is handled externally through convert_log method. - This method is kept for interface compatibility. - """ - pass - def _get_nuscenes_camera_metadata( - nusc: NuScenes, - scene: Dict[str, Any], - dataset_converter_config: DatasetConverterConfig, + nusc: NuScenes, + scene: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, ) -> Dict[PinholeCameraType, PinholeCameraMetadata]: camera_metadata: Dict[PinholeCameraType, PinholeCameraMetadata] = {} @@ -283,9 +224,9 @@ def _get_nuscenes_camera_metadata( def _get_nuscenes_lidar_metadata( - nusc: NuScenes, - scene: Dict[str, Any], - dataset_converter_config: DatasetConverterConfig, + nusc: NuScenes, + scene: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, ) -> Dict[LiDARType, LiDARMetadata]: metadata: Dict[LiDARType, LiDARMetadata] = {} @@ -303,9 +244,9 @@ def _get_nuscenes_lidar_metadata( extrinsic[:3, 3] = translation extrinsic = StateSE3.from_transformation_matrix(extrinsic) - metadata[LiDARType.LIDAR_MERGED] = LiDARMetadata( - lidar_type=LiDARType.LIDAR_MERGED, - lidar_index=NuscenesLidarIndex, + metadata[LiDARType.LIDAR_TOP] = LiDARMetadata( + lidar_type=LiDARType.LIDAR_TOP, + lidar_index=NuScenesLidarIndex, extrinsic=extrinsic, ) @@ -329,8 +270,8 @@ def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: quat = Quaternion(ego_pose["rotation"]) - vehicle_parameters = get_nuscenes_renauly_zoe_parameters() - pose = StateSE3( + vehicle_parameters = get_nuscenes_renault_zoe_parameters() + imu_pose = StateSE3( x=ego_pose["translation"][0], y=ego_pose["translation"][1], z=ego_pose["translation"][2], @@ -344,12 +285,12 @@ def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: try: pose_msgs = can_bus.get_messages(scene_name, "pose") - except Exception as e: + except Exception: pose_msgs = [] if pose_msgs: closest_msg = None - min_time_diff = float('inf') + min_time_diff = float("inf") for msg in pose_msgs: time_diff = abs(msg["utime"] - sample["timestamp"]) if time_diff < min_time_diff: @@ -371,18 +312,21 @@ def _extract_nuscenes_ego_state(nusc, sample, can_bus) -> EgoStateSE3: angular_velocity=Vector3D(*angular_velocity), ) - return EgoStateSE3( - center_se3=pose, + # return EgoStateSE3( + # center_se3=pose, + # dynamic_state_se3=dynamic_state, + # vehicle_parameters=vehicle_parameters, + # timepoint=TimePoint.from_us(sample["timestamp"]), + # ) + return EgoStateSE3.from_rear_axle( + rear_axle_se3=imu_pose, dynamic_state_se3=dynamic_state, vehicle_parameters=vehicle_parameters, - timepoint=TimePoint.from_us(sample["timestamp"]), + time_point=TimePoint.from_us(sample["timestamp"]), ) -def _extract_nuscenes_box_detections( - nusc: NuScenes, - sample: Dict[str, Any] -) -> BoxDetectionWrapper: +def _extract_nuscenes_box_detections(nusc: NuScenes, sample: Dict[str, Any]) -> BoxDetectionWrapper: box_detections: List[BoxDetectionSE3] = [] for ann_token in sample["anns"]: @@ -438,11 +382,6 @@ def _extract_nuscenes_box_detections( return BoxDetectionWrapper(box_detections=box_detections) -def _extract_nuscenes_traffic_lights() -> TrafficLightDetectionWrapper: - """nuScenes doesn't have traffic light information.""" - return TrafficLightDetectionWrapper(traffic_light_detections=[]) - - def _extract_nuscenes_cameras( nusc: NuScenes, sample: Dict[str, Any], @@ -472,8 +411,7 @@ def _extract_nuscenes_cameras( if cam_path.exists() and cam_path.is_file(): if dataset_converter_config.camera_store_option == "path": - # TODO: should be relative path - camera_data = cam_data["filename"] + camera_data = str(cam_path) elif dataset_converter_config.camera_store_option == "binary": with open(cam_path, "rb") as f: camera_data = f.read() @@ -486,194 +424,24 @@ def _extract_nuscenes_cameras( def _extract_nuscenes_lidars( - nusc: NuScenes, - sample: Dict[str, Any], - dataset_converter_config: DatasetConverterConfig, + nusc: NuScenes, + sample: Dict[str, Any], + dataset_converter_config: DatasetConverterConfig, ) -> List[LiDARData]: lidars: List[LiDARData] = [] if dataset_converter_config.include_lidars: lidar_token = sample["data"]["LIDAR_TOP"] lidar_data = nusc.get("sample_data", lidar_token) - lidar_path = NUSCENES_DATA_ROOT / lidar_data["filename"] + absolute_lidar_path = NUSCENES_DATA_ROOT / lidar_data["filename"] - if lidar_path.exists() and lidar_path.is_file(): + if absolute_lidar_path.exists() and absolute_lidar_path.is_file(): lidar = LiDARData( - lidar_type=LiDARType.LIDAR_TOP, - relative_path=str(lidar_path), - dataset_root=NUSCENES_DATA_ROOT, + lidar_type=LiDARType.LIDAR_MERGED, + relative_path=absolute_lidar_path.relative_to(NUSCENES_DATA_ROOT), + dataset_root=NUSCENES_DATA_ROOT, iteration=lidar_data.get("iteration"), ) lidars.append(lidar) - else: - lidars.append(LiDARData( - lidar_type=LiDARType.LIDAR_TOP, - relative_path=None, - dataset_root=NUSCENES_DATA_ROOT, - )) return lidars - -def _extract_nuscenes_scenario_tag() -> List[str]: - """nuScenes doesn't have scenario tags.""" - return ["unknown"] - - -def _extract_nuscenes_route_lane_group_ids() -> List[int]: - """nuScenes doesn't have route lane group information.""" - return [] - - -# Updated arrow conversion function using the new extraction functions -def convert_nuscenes_log_to_arrow( - args: List[Dict[str, Union[str, List[str]]]], - dataset_converter_config: DatasetConverterConfig, - version: str -) -> List[Any]: - for log_info in args: - scene_token: str = log_info["scene_token"] - split: str = log_info["split"] - - nusc = NuScenes(version=version, dataroot=str(NUSCENES_DATA_ROOT), verbose=False) - scene = nusc.get("scene", scene_token) - - log_file_path = dataset_converter_config.output_path / split / f"{scene_token}.arrow" - - if dataset_converter_config.force_log_conversion or not log_file_path.exists(): - if not log_file_path.parent.exists(): - log_file_path.parent.mkdir(parents=True, exist_ok=True) - - # Define schema - schema_column_list = [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("ego_state", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("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())), - ("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 dataset_converter_config.lidar_store_option == "path": - schema_column_list.append(("lidar", pa.string())) - - if dataset_converter_config.camera_store_option == "path": - for camera_type in NUSCENES_CAMERA_TYPES.keys(): - schema_column_list.append((camera_type.serialize(), pa.string())) - schema_column_list.append((f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4))) - - recording_schema = pa.schema(schema_column_list) - - log_record = nusc.get("log", scene["log_token"]) - location = log_record["location"] - - # Create metadata using the same functions as the new interface - metadata = LogMetadata( - dataset="nuscenes", - split=split, - log_name=scene["name"], - location=location, - timestep_seconds=TARGET_DT, - vehicle_parameters=get_nuscenes_renauly_zoe_parameters(), - camera_metadata=_get_nuscenes_camera_metadata(nusc, scene, dataset_converter_config), - lidar_metadata=_get_nuscenes_lidar_metadata(nusc, scene, dataset_converter_config), - map_metadata=_get_nuscenes_map_metadata(location), - ) - - recording_schema = recording_schema.with_metadata( - { - "log_metadata": json.dumps(asdict(metadata)), - } - ) - - _write_arrow_table_with_new_interface( - nusc, scene, recording_schema, log_file_path, dataset_converter_config - ) - - del nusc - gc.collect() - - return [] - - -def _write_arrow_table_with_new_interface( - nusc: NuScenes, - scene: Dict[str, Any], - recording_schema: pa.schema, - log_file_path: Path, - dataset_converter_config: DatasetConverterConfig, -) -> None: - can_bus = NuScenesCanBus(dataroot=str(NUSCENES_DATA_ROOT)) - - with pa.OSFile(str(log_file_path), "wb") as sink: - with pa.ipc.new_file(sink, recording_schema) as writer: - step_interval = max(1, int(TARGET_DT / NUSCENES_DT)) - sample_count = 0 - - sample_token = scene["first_sample_token"] - while sample_token: - if sample_count % step_interval == 0: - sample = nusc.get("sample", sample_token) - - # Use the new extraction functions for consistency - ego_state = _extract_nuscenes_ego_state(nusc, sample, can_bus) - box_detections = _extract_nuscenes_box_detections(nusc, sample) - cameras = _extract_nuscenes_cameras(nusc, sample, dataset_converter_config) - lidars = _extract_nuscenes_lidars(nusc, sample, dataset_converter_config) - - detections_state_list = [] - for det in box_detections.box_detections: - bbox_array = det.bounding_box_se3.array - - print(f"bbox_array shape: {bbox_array.shape}, ndim: {bbox_array.ndim}") - if bbox_array.ndim > 1: - detections_state_list.append(bbox_array.flatten().tolist()) - else: - detections_state_list.append(bbox_array.tolist()) - - # Prepare row data - row_data = { - "token": [sample_token], - "timestamp": [sample["timestamp"]], - "ego_state": ego_state.array.flatten().tolist(), - "detections_state": detections_state_list, - "detections_velocity": [det.velocity.array.tolist() for det in box_detections.box_detections], - "detections_token": [det.metadata.track_token for det in box_detections.box_detections], - "detections_type": [det.metadata.box_detection_type.value for det in box_detections.box_detections], - "traffic_light_ids": [], - "traffic_light_types": [], - "scenario_tag": ["unknown"], - "route_lane_group_ids": [], - } - - # Add lidar data if configured - if dataset_converter_config.lidar_store_option == "path": - row_data["lidar"] = [lidars.get(LiDARType.LIDAR_MERGED, None)] - - # Add camera data if configured - if dataset_converter_config.camera_store_option == "path": - for camera_type in NUSCENES_CAMERA_TYPES.keys(): - if camera_type in cameras: - camera_path, extrinsic = cameras[camera_type] - row_data[camera_type.serialize()] = [camera_path] - row_data[f"{camera_type.serialize()}_extrinsic"] = [ - extrinsic.to_transformation_matrix().flatten().tolist()] - 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) - - sample_token = sample["next"] - sample_count += 1 - - # Sort by timestamp if required - 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) diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py index 66f7a658..3fbb9b62 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py @@ -1,21 +1,16 @@ -import lanelet2 -import numpy as np - from pathlib import Path -from typing import List, Optional -from lanelet2.io import load -from lanelet2.projection import MercatorProjector -from shapely.geometry import Polygon, MultiPolygon, LineString -from shapely.validation import make_valid -from nuscenes.map_expansion.map_api import NuScenesMap -from nuscenes.map_expansion.arcline_path_utils import discretize_lane +from typing import Optional + +import numpy as np +from shapely.geometry import LineString, Polygon +from py123d.common.utils.dependencies import check_dependencies +from py123d.conversion.datasets.nuscenes.nuscenes_constants import NUSCENES_MAPS from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.datatypes.maps.cache.cache_map_objects import ( +from py123d.datatypes.maps.cache.cache_map_objects import ( # CacheIntersection, CacheCarpark, CacheCrosswalk, CacheGenericDrivable, - CacheIntersection, CacheLane, CacheLaneGroup, CacheRoadLine, @@ -24,21 +19,29 @@ from py123d.datatypes.maps.map_datatypes import RoadLineType from py123d.geometry import Polyline3D -NUSCENES_MAPS: List[str] = [ - "boston-seaport", - "singapore-hollandvillage", - "singapore-onenorth", - "singapore-queenstown" -] +check_dependencies(["lanelet2", "nuscenes"], optional_name="nuscenes") +import traceback + +import lanelet2 +from lanelet2.io import load +from lanelet2.projection import MercatorProjector +from nuscenes.map_expansion.arcline_path_utils import discretize_lane +from nuscenes.map_expansion.map_api import NuScenesMap -def write_nuscenes_map(nuscenes_maps_root: Path, location: str, map_writer: AbstractMapWriter, use_lanelet2: bool, lanelet2_root: Optional[str] = None) -> None: +def write_nuscenes_map( + nuscenes_maps_root: Path, + location: str, + map_writer: AbstractMapWriter, + use_lanelet2: bool, + lanelet2_root: Optional[str] = None, +) -> None: """ Main function to convert nuscenes map to unified format and write using map_writer. """ assert location in NUSCENES_MAPS, f"Map name {location} is not supported." nusc_map = NuScenesMap(dataroot=str(nuscenes_maps_root), map_name=location) - + # Write all layers if use_lanelet2: _write_nuscenes_lanes_lanelet2(nusc_map, map_writer, lanelet2_root) @@ -46,6 +49,7 @@ def write_nuscenes_map(nuscenes_maps_root: Path, location: str, map_writer: Abst else: _write_nuscenes_lanes(nusc_map, map_writer) _write_nuscenes_lane_groups(nusc_map, map_writer) + _write_nuscenes_intersections(nusc_map, map_writer) _write_nuscenes_crosswalks(nusc_map, map_writer) _write_nuscenes_walkways(nusc_map, map_writer) @@ -58,16 +62,16 @@ def write_nuscenes_map(nuscenes_maps_root: Path, location: str, map_writer: Abst def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str) -> None: map_name = nusc_map.map_name osm_map_file = str(Path(lanelet2_root) / f"{map_name}.osm") - + if "boston" in map_name.lower(): origin_lat, origin_lon = 42.3365, -71.0577 elif "singapore" in map_name.lower(): origin_lat, origin_lon = 1.3, 103.8 else: origin_lat, origin_lon = 49.0, 8.4 - + origin = lanelet2.io.Origin(origin_lat, origin_lon) - + try: lanelet_map = lanelet2.io.load(osm_map_file, origin) except Exception: @@ -79,19 +83,19 @@ def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMa for lanelet in lanelet_map.laneletLayer: token = lanelet.id - + try: left_bound = [(p.x, p.y) for p in lanelet.leftBound] right_bound = [(p.x, p.y) for p in lanelet.rightBound] polygon_points = left_bound + right_bound[::-1] polygon = Polygon(polygon_points) - + predecessor_ids = [int(pred.id) for pred in lanelet.previousLanelets] successor_ids = [int(succ.id) for succ in lanelet.followingLanelets] - + left_lane_id = None right_lane_id = None - + left_boundary = [(p.x, p.y) for p in lanelet.leftBound] right_boundary = [(p.x, p.y) for p in lanelet.rightBound] centerline = [] @@ -99,7 +103,7 @@ def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMa center_x = (left_pt.x + right_pt.x) / 2 center_y = (left_pt.y + right_pt.y) / 2 centerline.append((center_x, center_y)) - + speed_limit_mps = 0.0 if "speed_limit" in lanelet.attributes: try: @@ -109,7 +113,7 @@ def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMa speed_limit_mps = speed_kmh / 3.6 except (ValueError, TypeError): pass - + map_writer.write_lane( CacheLane( object_id=token, @@ -127,9 +131,13 @@ def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMa ) ) except Exception: + traceback.print_exc() continue -def _write_nuscenes_lane_groups_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str) -> None: + +def _write_nuscenes_lane_groups_lanelet2( + nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str +) -> None: map_name = nusc_map.map_name osm_map_file = str(Path(lanelet2_root) / f"{map_name}.osm") @@ -137,9 +145,9 @@ def _write_nuscenes_lane_groups_lanelet2(nusc_map: NuScenesMap, map_writer: Abst origin_lat, origin_lon = 42.3365, -71.0577 else: origin_lat, origin_lon = 1.3, 103.8 - + origin = lanelet2.io.Origin(origin_lat, origin_lon) - + try: projector = MercatorProjector(origin) lanelet_map = load(osm_map_file, projector) @@ -156,10 +164,10 @@ def _write_nuscenes_lane_groups_lanelet2(nusc_map: NuScenesMap, map_writer: Abst predecessor_ids = [] successor_ids = [] try: - if hasattr(lanelet, 'left'): + if hasattr(lanelet, "left"): for left_lane in lanelet.left: predecessor_ids.append(int(left_lane.id)) - if hasattr(lanelet, 'right'): + if hasattr(lanelet, "right"): for right_lane in lanelet.right: successor_ids.append(int(right_lane.id)) except Exception: @@ -171,6 +179,7 @@ def _write_nuscenes_lane_groups_lanelet2(nusc_map: NuScenesMap, map_writer: Abst polygon_points = left_bound + right_bound[::-1] polygon = Polygon(polygon_points) except Exception: + traceback.print_exc() continue try: @@ -188,16 +197,9 @@ def _write_nuscenes_lane_groups_lanelet2(nusc_map: NuScenesMap, map_writer: Abst ) ) except Exception: + traceback.print_exc() continue -def _get_lanelet_connections(lanelet): - """ - Helper function to extract incoming and outgoing lanelets. - """ - incoming = lanelet.incomings - outgoing = lanelet.outgoings - return incoming, outgoing - def _write_nuscenes_lanes(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: """ @@ -216,6 +218,7 @@ def _write_nuscenes_lanes(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) if not polygon.is_valid: continue except Exception: + traceback.print_exc() continue # Get topology @@ -261,19 +264,20 @@ def _write_nuscenes_lanes(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) CacheLane( object_id=token, lane_group_id=lane_record.get("road_segment_token", None), - left_boundary=left_boundary, - right_boundary=right_boundary, - centerline=baseline_path, + left_boundary=Polyline3D.from_linestring(left_boundary), + right_boundary=Polyline3D.from_linestring(right_boundary), + centerline=Polyline3D.from_linestring(baseline_path), left_lane_id=None, # Not directly available in nuscenes right_lane_id=None, # Not directly available in nuscenes predecessor_ids=incoming, successor_ids=outgoing, - speed_limit_mps=0.0, # Default value + speed_limit_mps=None, # Default value outline=None, geometry=polygon, ) ) except Exception: + traceback.print_exc() continue @@ -348,6 +352,7 @@ def _write_nuscenes_lane_groups(nusc_map: NuScenesMap, map_writer: AbstractMapWr ) ) except Exception: + traceback.print_exc() continue @@ -355,29 +360,29 @@ def _write_nuscenes_intersections(nusc_map: NuScenesMap, map_writer: AbstractMap """ Write intersection data to map_writer. """ - road_blocks = nusc_map.road_block - for block in road_blocks: - token = block["token"] - try: - if "polygon_token" in block: - polygon = nusc_map.extract_polygon(block["polygon_token"]) - else: - continue - if not polygon.is_valid: - continue - except Exception: - continue - - # Lane group IDs are not directly available; use empty list - lane_group_ids = [] - - map_writer.write_intersection( - CacheIntersection( - object_id=token, - lane_group_ids=lane_group_ids, - geometry=polygon, - ) - ) + # road_blocks = nusc_map.road_block + # for block in road_blocks: + # token = block["token"] + # try: + # if "polygon_token" in block: + # polygon = nusc_map.extract_polygon(block["polygon_token"]) + # else: + # continue + # if not polygon.is_valid: + # continue + # except Exception: + # continue + + # # Lane group IDs are not directly available; use empty list + # lane_group_ids = [] + + # map_writer.write_intersection( + # CacheIntersection( + # object_id=token, + # lane_group_ids=lane_group_ids, + # geometry=polygon, + # ) + # ) def _write_nuscenes_crosswalks(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: @@ -395,6 +400,7 @@ def _write_nuscenes_crosswalks(nusc_map: NuScenesMap, map_writer: AbstractMapWri if not polygon.is_valid: continue except Exception: + traceback.print_exc() continue map_writer.write_crosswalk( @@ -420,6 +426,7 @@ def _write_nuscenes_walkways(nusc_map: NuScenesMap, map_writer: AbstractMapWrite if not polygon.is_valid: continue except Exception: + traceback.print_exc() continue map_writer.write_walkway( @@ -470,6 +477,7 @@ def _write_nuscenes_generic_drivables(nusc_map: NuScenesMap, map_writer: Abstrac if polygon.is_valid: all_drivables.append((f"road_segment_{segment['token']}", polygon)) except Exception: + traceback.print_exc() continue # Add lanes @@ -480,6 +488,7 @@ def _write_nuscenes_generic_drivables(nusc_map: NuScenesMap, map_writer: Abstrac if polygon.is_valid: all_drivables.append((f"lane_{lane['token']}", polygon)) except Exception: + traceback.print_exc() continue # Add drivable areas @@ -490,6 +499,7 @@ def _write_nuscenes_generic_drivables(nusc_map: NuScenesMap, map_writer: Abstrac if polygon.is_valid: all_drivables.append((f"road_{road['token']}", polygon)) except Exception: + traceback.print_exc() continue for obj_id, geometry in all_drivables: @@ -516,6 +526,7 @@ def _write_nuscenes_stop_lines(nusc_map: NuScenesMap, map_writer: AbstractMapWri if not polygon.is_valid: continue except Exception: + traceback.print_exc() continue # Note: Stop lines are written as generic drivable for compatibility @@ -586,7 +597,7 @@ def _get_lane_boundary(lane_token: str, side: str, nusc_map: NuScenesMap) -> Opt divider_segment_nodes_key = f"{side}_lane_divider_segment_nodes" if divider_segment_nodes_key in lane_record and lane_record[divider_segment_nodes_key]: nodes = lane_record[divider_segment_nodes_key] - boundary = LineString([(node['x'], node['y']) for node in nodes]) + boundary = LineString([(node["x"], node["y"]) for node in nodes]) return boundary return None @@ -611,7 +622,7 @@ def _get_lane_group_boundary(segment_token: str, side: str, nusc_map: NuScenesMa # Find nearest boundary of the specified type within a threshold nearest = None - min_dist = float('inf') + min_dist = float("inf") if boundary_type == "road_divider": records = nusc_map.road_divider @@ -705,4 +716,4 @@ def align_boundary_direction(centerline: LineString, boundary: LineString) -> Li """ if not lines_same_direction(centerline, boundary): return flip_linestring(boundary) - return boundary \ No newline at end of file + return boundary diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py index 82f1ac72..7592182c 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py @@ -1,13 +1,26 @@ -import numpy as np from pathlib import Path from typing import Dict + +import numpy as np + +from py123d.datatypes.scene.scene_metadata import LogMetadata from py123d.datatypes.sensors.lidar.lidar import LiDARType +from py123d.datatypes.sensors.lidar.lidar_index import NuScenesLidarIndex +from py123d.geometry.se import StateSE3 +from py123d.geometry.transform.transform_se3 import convert_points_3d_array_between_origins -def load_nuscenes_lidar_pcs_from_file(pcd_path: Path) -> Dict[LiDARType, np.ndarray]: + +def load_nuscenes_lidar_pcs_from_file(pcd_path: Path, log_metadata: LogMetadata) -> Dict[LiDARType, np.ndarray]: points = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, 5) - - lidar_pcs_dict: Dict[LiDARType, np.ndarray] = { - LiDARType.LIDAR_TOP: points - } - + + # convert lidar to ego frame + lidar_extrinsic = log_metadata.lidar_metadata[LiDARType.LIDAR_TOP].extrinsic + + points[..., NuScenesLidarIndex.XYZ] = convert_points_3d_array_between_origins( + from_origin=lidar_extrinsic, + to_origin=StateSE3(0, 0, 0, 1.0, 0, 0, 0), + points_3d_array=points[..., NuScenesLidarIndex.XYZ], + ) + lidar_pcs_dict: Dict[LiDARType, np.ndarray] = {LiDARType.LIDAR_TOP: points} + return lidar_pcs_dict diff --git a/src/py123d/conversion/registry/lidar_index_registry.py b/src/py123d/conversion/registry/lidar_index_registry.py index c5213d2a..b58e4fce 100644 --- a/src/py123d/conversion/registry/lidar_index_registry.py +++ b/src/py123d/conversion/registry/lidar_index_registry.py @@ -82,11 +82,12 @@ class PandasetLidarIndex(LiDARIndex): Y = 1 Z = 2 INTENSITY = 3 - + + @register_lidar_index -class NuscenesLidarIndex(LiDARIndex): - X = 0 - Y = 1 - Z = 2 - INTENSITY = 3 - RING = 4 +class NuScenesLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + RING = 4 diff --git a/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py b/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py index 83d4da21..82c08efc 100644 --- a/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py +++ b/src/py123d/conversion/sensor_io/lidar/file_lidar_io.py @@ -15,7 +15,7 @@ "av2-sensor": DATASET_PATHS.av2_sensor_data_root, "wopd": DATASET_PATHS.wopd_data_root, "pandaset": DATASET_PATHS.pandaset_data_root, - "nuscenes": DATASET_PATHS.nuscenes_sensor_root, + "nuscenes": DATASET_PATHS.nuscenes_sensor_root, } @@ -57,10 +57,11 @@ def load_lidar_pcs_from_file( from py123d.conversion.datasets.pandaset.pandaset_sensor_io import load_pandaset_lidars_pcs_from_file lidar_pcs_dict = load_pandaset_lidars_pcs_from_file(full_lidar_path, index) - + elif log_metadata.dataset == "nuscenes": from py123d.conversion.datasets.nuscenes.nuscenes_sensor_io import load_nuscenes_lidar_pcs_from_file - lidar_pcs_dict = load_nuscenes_lidar_pcs_from_file(full_lidar_path) + + lidar_pcs_dict = load_nuscenes_lidar_pcs_from_file(full_lidar_path, log_metadata) else: raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") diff --git a/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py b/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py index 031c726a..6345df02 100644 --- a/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py +++ b/src/py123d/datatypes/scene/arrow/utils/arrow_getters.py @@ -34,11 +34,11 @@ DATASET_PATHS: DictConfig = get_dataset_paths() DATASET_SENSOR_ROOT: Dict[str, Path] = { - "nuplan": DATASET_PATHS.nuplan_sensor_root, "av2-sensor": DATASET_PATHS.av2_sensor_data_root, + "nuplan": DATASET_PATHS.nuplan_sensor_root, + "nuscenes": DATASET_PATHS.nuscenes_data_root, "wopd": DATASET_PATHS.wopd_data_root, "pandaset": DATASET_PATHS.pandaset_data_root, - "nuscenes": DATASET_PATHS.nuscenes_data_root, } diff --git a/src/py123d/datatypes/sensors/lidar/lidar_index.py b/src/py123d/datatypes/sensors/lidar/lidar_index.py index c101ba8e..b4864cac 100644 --- a/src/py123d/datatypes/sensors/lidar/lidar_index.py +++ b/src/py123d/datatypes/sensors/lidar/lidar_index.py @@ -86,9 +86,10 @@ class PandasetLidarIndex(LiDARIndex): @register_lidar_index -class NuscenesLidarIndex(LiDARIndex): - X = 0 - Y = 1 - Z = 2 - INTENSITY = 3 - RING = 4 +class NuScenesLidarIndex(LiDARIndex): + + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + RING = 4 diff --git a/src/py123d/datatypes/vehicle_state/vehicle_parameters.py b/src/py123d/datatypes/vehicle_state/vehicle_parameters.py index 19b050c7..5e15c3b7 100644 --- a/src/py123d/datatypes/vehicle_state/vehicle_parameters.py +++ b/src/py123d/datatypes/vehicle_state/vehicle_parameters.py @@ -52,17 +52,20 @@ def get_nuplan_chrysler_pacifica_parameters() -> VehicleParameters: rear_axle_to_center_longitudinal=1.461, ) -def get_nuscenes_renauly_zoe_parameters() -> VehicleParameters: + +def get_nuscenes_renault_zoe_parameters() -> VehicleParameters: + # https://en.wikipedia.org/wiki/Renault_Zoe return VehicleParameters( - vehicle_name="nuscenes_renauly_zoe", + vehicle_name="nuscenes_renault_zoe", width=1.730, length=4.084, height=1.562, wheel_base=2.588, - rear_axle_to_center_vertical=1.562 / 2, # NOTE: missing in nuscenes, TODO: find more accurate value + rear_axle_to_center_vertical=1.562 / 2, # NOTE: missing in nuscenes, TODO: find more accurate value rear_axle_to_center_longitudinal=1.385, ) + def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: # NOTE: values are extracted from CARLA return VehicleParameters( diff --git a/src/py123d/script/config/common/default_dataset_paths.yaml b/src/py123d/script/config/common/default_dataset_paths.yaml index 64e0a8c0..b4941707 100644 --- a/src/py123d/script/config/common/default_dataset_paths.yaml +++ b/src/py123d/script/config/common/default_dataset_paths.yaml @@ -24,4 +24,4 @@ dataset_paths: # nuScenes defaults nuscenes_data_root: ${oc.env:NUSCENES_DATA_ROOT,null} nuscenes_map_root: ${dataset_paths.nuscenes_data_root} - nuscenes_sensor_root: ${dataset_paths.nuscenes_data_root}/samples/LIDAR_TOP \ No newline at end of file + nuscenes_sensor_root: ${dataset_paths.nuscenes_data_root} diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml index 410b910e..0f3ab95e 100644 --- a/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuscenes_dataset.yaml @@ -1,9 +1,10 @@ nuscenes_dataset: - _target_: py123d.conversion.datasets.nuscenes.nuscenes_data_converter.NuScenesDataConverter + _target_: py123d.conversion.datasets.nuscenes.nuscenes_converter.NuScenesConverter _convert_: 'all' splits: ["nuscenes_train", "nuscenes_val", "nuscenes_test"] nuscenes_data_root: ${dataset_paths.nuscenes_data_root} + nuscenes_map_root: ${dataset_paths.nuscenes_data_root} nuscenes_lanelet2_root: ${dataset_paths.nuscenes_data_root}/lanelet2 use_lanelet2: False @@ -13,6 +14,7 @@ nuscenes_dataset: force_log_conversion: ${force_log_conversion} force_map_conversion: ${force_map_conversion} + # Map include_map: true @@ -25,10 +27,10 @@ nuscenes_dataset: # Traffic Lights include_traffic_lights: false - #cameras + # Cameras include_cameras: true camera_store_option: "path" - + #lidar include_lidars: true - lidar_store_option: "path" \ No newline at end of file + lidar_store_option: "path" diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml new file mode 100644 index 00000000..95a3c927 --- /dev/null +++ b/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml @@ -0,0 +1,36 @@ +nuscenes_dataset: + _target_: py123d.conversion.datasets.nuscenes.nuscenes_converter.NuScenesConverter + _convert_: 'all' + + splits: ["nuscenes-mini_train", "nuscenes-mini_val"] + nuscenes_data_root: ${dataset_paths.nuscenes_data_root} + nuscenes_map_root: ${dataset_paths.nuscenes_data_root} + nuscenes_lanelet2_root: ${dataset_paths.nuscenes_data_root}/lanelet2 + use_lanelet2: False + + dataset_converter_config: + _target_: py123d.conversion.dataset_converter_config.DatasetConverterConfig + _convert_: 'all' + + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + + # Map + include_map: true + + # Ego + include_ego: true + + # Box Detections + include_box_detections: true + + # Traffic Lights + include_traffic_lights: false + + # Cameras + include_cameras: true + camera_store_option: "binary" + + #lidar + include_lidars: true + lidar_store_option: "binary" diff --git a/test_viser.py b/test_viser.py index 73896531..d5375bd7 100644 --- a/test_viser.py +++ b/test_viser.py @@ -1,15 +1,16 @@ from py123d.common.multithreading.worker_sequential import Sequential from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder from py123d.datatypes.scene.scene_filter import SceneFilter -from py123d.datatypes.sensors.camera.pinhole_camera import PinholeCameraType from py123d.visualization.viser.viser_viewer import ViserViewer -if __name__ == "__main__": +# from py123d.datatypes.sensors.camera.pinhole_camera import PinholeCameraType - splits = ["nuplan-mini_test", "nuplan-mini_train", "nuplan-mini_val"] +if __name__ == "__main__": + splits = ["nuscenes-mini_val", "nuscenes-mini_train"] + # splits = ["nuplan-mini_test", "nuplan-mini_train", "nuplan-mini_val"] # splits = ["nuplan_private_test"] # splits = ["carla_test"] - splits = ["wopd_val"] + # splits = ["wopd_val"] # splits = ["av2-sensor_train"] # splits = ["pandaset_test", "pandaset_val", "pandaset_train"] # log_names = ["2021.08.24.13.12.55_veh-45_00386_00472"] @@ -21,11 +22,11 @@ split_names=splits, log_names=log_names, scene_uuids=scene_uuids, - duration_s=10.0, + duration_s=None, history_s=0.0, timestamp_threshold_s=10.0, shuffle=True, - camera_types=[PinholeCameraType.CAM_F0], + # camera_types=[PinholeCameraType.CAM_F0], ) scene_builder = ArrowSceneBuilder() worker = Sequential() From b733b8223221f2994077f53b65b8e800e405b080 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 31 Oct 2025 23:26:06 +0100 Subject: [PATCH 4/4] Refactor nuScenes map conversion (lanelet currently missing). --- notebooks/bev_matplotlib.ipynb | 61 +- .../conversion/nuscenes_mini_conversion.sh | 2 +- .../datasets/nuscenes/nuscenes_converter.py | 4 +- .../nuscenes/nuscenes_map_conversion.py | 963 +++++++----------- .../datasets/nuscenes/nuscenes_sensor_io.py | 11 +- .../{ => utils}/nuscenes_constants.py | 0 .../nuscenes/utils/nuscenes_map_utils.py | 197 ++++ .../map_utils/road_edge/road_edge_2d_utils.py | 48 +- .../datatypes/maps/gpkg/gpkg_map_objects.py | 37 +- src/py123d/datatypes/maps/gpkg/gpkg_utils.py | 23 +- src/py123d/geometry/utils/polyline_utils.py | 22 + .../datasets/nuscenes_mini_dataset.yaml | 2 +- .../map_writer/gpkg_map_writer.yaml | 1 + 13 files changed, 693 insertions(+), 678 deletions(-) rename src/py123d/conversion/datasets/nuscenes/{ => utils}/nuscenes_constants.py (100%) create mode 100644 src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py diff --git a/notebooks/bev_matplotlib.ipynb b/notebooks/bev_matplotlib.ipynb index bd2dd882..910bf63a 100644 --- a/notebooks/bev_matplotlib.ipynb +++ b/notebooks/bev_matplotlib.ipynb @@ -10,6 +10,7 @@ "from py123d.datatypes.scene.arrow.arrow_scene_builder import ArrowSceneBuilder\n", "from py123d.datatypes.scene.scene_filter import SceneFilter\n", "\n", + "\n", "from py123d.common.multithreading.worker_sequential import Sequential\n", "# from py123d.datatypes.sensors.camera.pinhole_camera import PinholeCameraType " ] @@ -144,22 +145,22 @@ " route_lane_group_ids: Optional[List[int]] = None,\n", ") -> None:\n", " layers: List[MapLayer] = [\n", - " MapLayer.LANE,\n", + " # MapLayer.LANE,\n", " MapLayer.LANE_GROUP,\n", - " # MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.GENERIC_DRIVABLE,\n", " # MapLayer.CARPARK,\n", " # MapLayer.CROSSWALK,\n", - " MapLayer.INTERSECTION,\n", + " # MapLayer.INTERSECTION,\n", " # MapLayer.WALKWAY,\n", " MapLayer.ROAD_EDGE,\n", - " # MapLayer.ROAD_LINE,\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", + " # print(map_objects_dict[MapLayer.ROAD_EDGE])\n", "\n", - " done = False\n", " for layer, map_objects in map_objects_dict.items():\n", " for map_object in map_objects:\n", " try:\n", @@ -167,39 +168,45 @@ " MapLayer.GENERIC_DRIVABLE,\n", " MapLayer.CARPARK,\n", " MapLayer.CROSSWALK,\n", - " MapLayer.INTERSECTION,\n", + " # MapLayer.INTERSECTION,\n", " MapLayer.WALKWAY,\n", " ]:\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", - " print(f\"Added {layer.name} with id {map_object.object_id}\")\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", - " if layer in [MapLayer.LANE]:\n", - " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", + " if map_object.intersection is not None:\n", + " add_shapely_polygon_to_ax(ax, map_object.intersection.shapely_polygon, ROUTE_CONFIG)\n", + "\n", + " for lane in map_object.lanes:\n", + " add_shapely_linestring_to_ax(ax, lane.centerline.linestring, CENTERLINE_CONFIG)\n", + "\n", + " # if layer in [MapLayer.LANE]:\n", + " # add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", + " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\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", + " 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", + " print(f\"Error adding map object of type {layer.name} and id {map_object.object_id}\")\n", " traceback.print_exc()\n", "\n", " # ax.set_title(f\"Map: {map_api.map_name}\")\n", @@ -213,8 +220,10 @@ "\n", " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", " if map_api is not None:\n", - " add_debug_map_on_ax(ax, scene.get_map_api(), point_2d, radius=radius, route_lane_group_ids=None)\n", - " # add_default_map_on_ax(ax, map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_debug_map_on_ax(ax, scene.get_map_api(), point_2d, radius=radius, route_lane_group_ids=None)\n", + "\n", + "\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, scene.get_map_api())\n", "\n", " add_box_detections_to_ax(ax, box_detections)\n", @@ -239,7 +248,7 @@ " return fig, ax\n", "\n", "\n", - "# scene_index = \n", + "# scene_index =\n", "iteration = 1\n", "\n", "scale = 10\n", @@ -249,7 +258,7 @@ "# _plot_scene_on_ax(ax[1], scene, iteration, radius=50)\n", "# _plot_scene_on_ax(ax[2], scene, iteration, radius=100)\n", "\n", - "plt.show()\n" + "plt.show()" ] }, { diff --git a/scripts/conversion/nuscenes_mini_conversion.sh b/scripts/conversion/nuscenes_mini_conversion.sh index b6d35aa4..b9a9a7d1 100644 --- a/scripts/conversion/nuscenes_mini_conversion.sh +++ b/scripts/conversion/nuscenes_mini_conversion.sh @@ -1,3 +1,3 @@ export NUSCENES_DATA_ROOT="/home/daniel/nuscenes_mini/" -py123d-conversion datasets=["nuscenes_mini_dataset"] +py123d-conversion datasets=["nuscenes_mini_dataset"] map_writer.remap_ids=true diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py index 483a6bc4..c4e1627e 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_converter.py @@ -9,14 +9,14 @@ from py123d.conversion.abstract_dataset_converter import AbstractDatasetConverter from py123d.conversion.dataset_converter_config import DatasetConverterConfig from py123d.conversion.datasets.nuplan.nuplan_converter import TARGET_DT -from py123d.conversion.datasets.nuscenes.nuscenes_constants import ( +from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import NUSCENES_MAPS, write_nuscenes_map +from py123d.conversion.datasets.nuscenes.utils.nuscenes_constants import ( NUSCENES_CAMERA_TYPES, NUSCENES_DATA_ROOT, NUSCENES_DATA_SPLITS, NUSCENES_DETECTION_NAME_DICT, NUSCENES_DT, ) -from py123d.conversion.datasets.nuscenes.nuscenes_map_conversion import NUSCENES_MAPS, write_nuscenes_map from py123d.conversion.log_writer.abstract_log_writer import AbstractLogWriter, LiDARData from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter from py123d.datatypes.detections.box_detections import BoxDetectionMetadata, BoxDetectionSE3, BoxDetectionWrapper diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py index 3fbb9b62..39ae313a 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_map_conversion.py @@ -1,33 +1,45 @@ +from collections import defaultdict from pathlib import Path -from typing import Optional +from typing import Dict, Final, List, Optional import numpy as np from shapely.geometry import LineString, Polygon from py123d.common.utils.dependencies import check_dependencies -from py123d.conversion.datasets.nuscenes.nuscenes_constants import NUSCENES_MAPS +from py123d.conversion.datasets.nuscenes.utils.nuscenes_constants import NUSCENES_MAPS +from py123d.conversion.datasets.nuscenes.utils.nuscenes_map_utils import ( + extract_lane_and_boundaries, + extract_nuscenes_centerline, + order_lanes_left_to_right, +) from py123d.conversion.map_writer.abstract_map_writer import AbstractMapWriter -from py123d.datatypes.maps.cache.cache_map_objects import ( # CacheIntersection, +from py123d.conversion.utils.map_utils.road_edge.road_edge_2d_utils import ( + get_road_edge_linear_rings, + split_line_geometry_by_max_length, + split_polygon_by_grid, +) +from py123d.datatypes.maps.cache.cache_map_objects import ( CacheCarpark, CacheCrosswalk, CacheGenericDrivable, + CacheIntersection, CacheLane, CacheLaneGroup, + CacheRoadEdge, CacheRoadLine, CacheWalkway, ) -from py123d.datatypes.maps.map_datatypes import RoadLineType -from py123d.geometry import Polyline3D - -check_dependencies(["lanelet2", "nuscenes"], optional_name="nuscenes") -import traceback +from py123d.datatypes.maps.map_datatypes import RoadEdgeType, RoadLineType +from py123d.geometry import OccupancyMap2D, Polyline2D, Polyline3D +from py123d.geometry.utils.polyline_utils import offset_points_perpendicular -import lanelet2 -from lanelet2.io import load -from lanelet2.projection import MercatorProjector -from nuscenes.map_expansion.arcline_path_utils import discretize_lane +check_dependencies(["nuscenes"], optional_name="nuscenes") from nuscenes.map_expansion.map_api import NuScenesMap +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # [m] +MAX_LANE_WIDTH: Final[float] = 4.0 # [m] +MIN_LANE_WIDTH: Final[float] = 1.0 # [m] + def write_nuscenes_map( nuscenes_maps_root: Path, @@ -36,634 +48,432 @@ def write_nuscenes_map( use_lanelet2: bool, lanelet2_root: Optional[str] = None, ) -> None: + """Converts the nuScenes map types to the 123D format, and sends elements to the map writer. + FIXME @DanielDauner: Currently, Lanelet2 format is not supported for nuScenes. + + :param nuscenes_maps_root: Path to the nuScenes maps root directory + :param location: Name of the specific map location to convert + :param map_writer: Map writer instance to write the converted elements + :param use_lanelet2: Flag indicating whether to use Lanelet2 format + :param lanelet2_root: Path to the Lanelet2 root directory, defaults to None """ - Main function to convert nuscenes map to unified format and write using map_writer. - """ + assert location in NUSCENES_MAPS, f"Map name {location} is not supported." - nusc_map = NuScenesMap(dataroot=str(nuscenes_maps_root), map_name=location) - - # Write all layers - if use_lanelet2: - _write_nuscenes_lanes_lanelet2(nusc_map, map_writer, lanelet2_root) - _write_nuscenes_lane_groups_lanelet2(nusc_map, map_writer, lanelet2_root) - else: - _write_nuscenes_lanes(nusc_map, map_writer) - _write_nuscenes_lane_groups(nusc_map, map_writer) - - _write_nuscenes_intersections(nusc_map, map_writer) - _write_nuscenes_crosswalks(nusc_map, map_writer) - _write_nuscenes_walkways(nusc_map, map_writer) - _write_nuscenes_carparks(nusc_map, map_writer) - _write_nuscenes_generic_drivables(nusc_map, map_writer) - _write_nuscenes_stop_lines(nusc_map, map_writer) - _write_nuscenes_road_lines(nusc_map, map_writer) - - -def _write_nuscenes_lanes_lanelet2(nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str) -> None: - map_name = nusc_map.map_name - osm_map_file = str(Path(lanelet2_root) / f"{map_name}.osm") - - if "boston" in map_name.lower(): - origin_lat, origin_lon = 42.3365, -71.0577 - elif "singapore" in map_name.lower(): - origin_lat, origin_lon = 1.3, 103.8 - else: - origin_lat, origin_lon = 49.0, 8.4 - - origin = lanelet2.io.Origin(origin_lat, origin_lon) - - try: - lanelet_map = lanelet2.io.load(osm_map_file, origin) - except Exception: - try: - projector = lanelet2.projection.MercatorProjector(origin) - lanelet_map = lanelet2.io.load(osm_map_file, projector) - except Exception: - return - - for lanelet in lanelet_map.laneletLayer: - token = lanelet.id - - try: - left_bound = [(p.x, p.y) for p in lanelet.leftBound] - right_bound = [(p.x, p.y) for p in lanelet.rightBound] - polygon_points = left_bound + right_bound[::-1] - polygon = Polygon(polygon_points) - - predecessor_ids = [int(pred.id) for pred in lanelet.previousLanelets] - successor_ids = [int(succ.id) for succ in lanelet.followingLanelets] - - left_lane_id = None - right_lane_id = None - - left_boundary = [(p.x, p.y) for p in lanelet.leftBound] - right_boundary = [(p.x, p.y) for p in lanelet.rightBound] - centerline = [] - for left_pt, right_pt in zip(lanelet.leftBound, lanelet.rightBound): - center_x = (left_pt.x + right_pt.x) / 2 - center_y = (left_pt.y + right_pt.y) / 2 - centerline.append((center_x, center_y)) - - speed_limit_mps = 0.0 - if "speed_limit" in lanelet.attributes: - try: - speed_limit_str = lanelet.attributes["speed_limit"] - if "km/h" in speed_limit_str: - speed_kmh = float(speed_limit_str.replace("km/h", "").strip()) - speed_limit_mps = speed_kmh / 3.6 - except (ValueError, TypeError): - pass - - map_writer.write_lane( - CacheLane( - object_id=token, - lane_group_id=None, - left_boundary=left_boundary, - right_boundary=right_boundary, - centerline=centerline, - left_lane_id=left_lane_id, - right_lane_id=right_lane_id, - predecessor_ids=predecessor_ids, - successor_ids=successor_ids, - speed_limit_mps=speed_limit_mps, - outline=None, - geometry=polygon, - ) - ) - except Exception: - traceback.print_exc() - continue + nuscenes_map = NuScenesMap(dataroot=str(nuscenes_maps_root), map_name=location) + # 1. extract road edges (used later to determine lane connector widths) + road_edges = _extract_nuscenes_road_edges(nuscenes_map) -def _write_nuscenes_lane_groups_lanelet2( - nusc_map: NuScenesMap, map_writer: AbstractMapWriter, lanelet2_root: str -) -> None: - map_name = nusc_map.map_name - osm_map_file = str(Path(lanelet2_root) / f"{map_name}.osm") - - if "boston" in map_name.lower(): - origin_lat, origin_lon = 42.3365, -71.0577 - else: - origin_lat, origin_lon = 1.3, 103.8 - - origin = lanelet2.io.Origin(origin_lat, origin_lon) - - try: - projector = MercatorProjector(origin) - lanelet_map = load(osm_map_file, projector) - except Exception: - return - - for lanelet in lanelet_map.laneletLayer: - token = lanelet.id - lane_ids = [lanelet.id] - try: - predecessor_ids = [int(lanelet.id) for lanelet in lanelet.previous] - successor_ids = [int(lanelet.id) for lanelet in lanelet.following] - except AttributeError: - predecessor_ids = [] - successor_ids = [] - try: - if hasattr(lanelet, "left"): - for left_lane in lanelet.left: - predecessor_ids.append(int(left_lane.id)) - if hasattr(lanelet, "right"): - for right_lane in lanelet.right: - successor_ids.append(int(right_lane.id)) - except Exception: - pass - - try: - left_bound = [(p.x, p.y) for p in lanelet.leftBound] - right_bound = [(p.x, p.y) for p in lanelet.rightBound] - polygon_points = left_bound + right_bound[::-1] - polygon = Polygon(polygon_points) - except Exception: - traceback.print_exc() - continue - - try: - map_writer.write_lane_group( - CacheLaneGroup( - object_id=token, - lane_ids=lane_ids, - left_boundary=None, - right_boundary=None, - intersection_id=None, - predecessor_ids=predecessor_ids, - successor_ids=successor_ids, - outline=None, - geometry=polygon, - ) - ) - except Exception: - traceback.print_exc() - continue + # 2. extract lanes + lanes = _extract_nuscenes_lanes(nuscenes_map) + # 3. extract lane connectors (i.e. lanes on intersections) + lane_connectors = _extract_nuscenes_lane_connectors(nuscenes_map, road_edges) -def _write_nuscenes_lanes(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write lane data to map_writer, including topology and boundaries. - """ - lane_records = nusc_map.lane - for lane_record in lane_records: - token = lane_record["token"] + # 4. extract intersections (and store lane-connector to intersection assignment for lane groups) + intersection_assignment = _write_nuscenes_intersections(nuscenes_map, lane_connectors, map_writer) - # Extract geometry from lane record - try: - if "polygon_token" in lane_record: - polygon = nusc_map.extract_polygon(lane_record["polygon_token"]) - else: - continue - if not polygon.is_valid: - continue - except Exception: - traceback.print_exc() - continue + # 5. extract lane groups + lane_groups = _extract_nuscenes_lane_groups(nuscenes_map, lanes, lane_connectors, intersection_assignment) - # Get topology - incoming = nusc_map.get_incoming_lane_ids(token) - outgoing = nusc_map.get_outgoing_lane_ids(token) + # Write remaining map elements + _write_nuscenes_crosswalks(nuscenes_map, map_writer) + _write_nuscenes_walkways(nuscenes_map, map_writer) + _write_nuscenes_carparks(nuscenes_map, map_writer) + _write_nuscenes_generic_drivables(nuscenes_map, map_writer) + _write_nuscenes_stop_lines(nuscenes_map, map_writer) + _write_nuscenes_road_lines(nuscenes_map, map_writer) - # Get lane connectors - lane_connectors = [] - for connector in nusc_map.lane_connector: - if connector.get("incoming_lane") == token or connector.get("outgoing_lane") == token: - lane_connectors.append(connector["token"]) + for lane in lanes + lane_connectors: + map_writer.write_lane(lane) - # Extract boundaries - left_boundary = _get_lane_boundary(token, "left", nusc_map) - right_boundary = _get_lane_boundary(token, "right", nusc_map) + for road_edge in road_edges: + map_writer.write_road_edge(road_edge) - # Skip lanes without valid boundaries - if left_boundary is None or right_boundary is None: - continue - if left_boundary.is_empty or right_boundary.is_empty: - continue - - # Extract baseline path - baseline_path = None - if token in nusc_map.arcline_path_3: - arc_path = nusc_map.arcline_path_3[token] - try: - points = discretize_lane(arc_path, resolution_meters=0.1) - xy_points = [(p[0], p[1]) for p in points] - baseline_path = LineString(xy_points) - except Exception: - baseline_path = None - - # Align boundaries with baseline path direction - if baseline_path and left_boundary: - left_boundary = align_boundary_direction(baseline_path, left_boundary) - if baseline_path and right_boundary: - right_boundary = align_boundary_direction(baseline_path, right_boundary) - - # Write lane object safely - try: - map_writer.write_lane( - CacheLane( - object_id=token, - lane_group_id=lane_record.get("road_segment_token", None), - left_boundary=Polyline3D.from_linestring(left_boundary), - right_boundary=Polyline3D.from_linestring(right_boundary), - centerline=Polyline3D.from_linestring(baseline_path), - left_lane_id=None, # Not directly available in nuscenes - right_lane_id=None, # Not directly available in nuscenes - predecessor_ids=incoming, - successor_ids=outgoing, - speed_limit_mps=None, # Default value - outline=None, - geometry=polygon, - ) - ) - except Exception: - traceback.print_exc() - continue + for lane_group in lane_groups: + map_writer.write_lane_group(lane_group) -def _write_nuscenes_lane_groups(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write lane group data to map_writer. - """ - road_segments = nusc_map.road_segment - for segment in road_segments: - token = segment["token"] - - # Extract geometry - try: - if "polygon_token" in segment: - polygon = nusc_map.extract_polygon(segment["polygon_token"]) - else: - continue - if not polygon.is_valid: - continue - except Exception: - continue +def _extract_nuscenes_lanes(nuscenes_map: NuScenesMap) -> List[CacheLane]: + """Helper function to extract lanes from a nuScenes map.""" - # Find lanes in this segment - lane_ids = [] - for lane in nusc_map.lane: - if lane.get("road_segment_token") == token: - lane_ids.append(lane["token"]) + # NOTE: nuScenes does not provide explicitly provide lane groups and does not assign lanes to roadblocks. + # Therefore, we query the roadblocks given the middle-point of the centerline to assign lanes to a road block. + # Unlike road segments, road blocks outline a lane group going in the same direction. + # In case a roadblock cannot be assigned, e.g. because the lane is not located within any roadblock, or the + # roadblock data is invalid [1], we assign a new lane group with only this lane. + # [1] https://github.com/nutonomy/nuscenes-devkit/issues/862 - # Get connected segments - incoming, outgoing = _get_connected_segments(token, nusc_map) + road_blocks_invalid = nuscenes_map.map_name in ["singapore-queenstown", "singapore-hollandvillage"] - # Extract boundaries - left_boundary = _get_lane_group_boundary(token, "left", nusc_map) - right_boundary = _get_lane_group_boundary(token, "right", nusc_map) + road_block_dict: Dict[str, Polygon] = {} + if not road_blocks_invalid: + road_block_dict: Dict[str, Polygon] = { + road_block["token"]: nuscenes_map.extract_polygon(road_block["polygon_token"]) + for road_block in nuscenes_map.road_block + } + + road_block_map = OccupancyMap2D.from_dict(road_block_dict) + lanes: List[CacheLane] = [] + for lane_record in nuscenes_map.lane: + token = lane_record["token"] + + # 1. Extract centerline and boundaries + centerline, left_boundary, right_boundary = extract_lane_and_boundaries(nuscenes_map, lane_record) - # Skip invalid boundaries if left_boundary is None or right_boundary is None: - continue - if left_boundary.is_empty or right_boundary.is_empty: - continue - - # Use first lane's baseline path for direction alignment - baseline_path = None - if lane_ids: - first_lane_token = lane_ids[0] - if first_lane_token in nusc_map.arcline_path_3: - arc_path = nusc_map.arcline_path_3[first_lane_token] - try: - points = discretize_lane(arc_path, resolution_meters=0.1) - xy_points = [(p[0], p[1]) for p in points] - baseline_path = LineString(xy_points) - except Exception: - baseline_path = None - - if baseline_path and left_boundary: - left_boundary = align_boundary_direction(baseline_path, left_boundary) - if baseline_path and right_boundary: - right_boundary = align_boundary_direction(baseline_path, right_boundary) - - try: - map_writer.write_lane_group( - CacheLaneGroup( - object_id=token, - lane_ids=lane_ids, - left_boundary=left_boundary, - right_boundary=right_boundary, - intersection_id=None, # Handled in intersections - predecessor_ids=incoming, - successor_ids=outgoing, - outline=None, - geometry=polygon, - ) + continue # skip lanes without valid boundaries + + # 2. Query road block for lane group assignment + lane_group_id: str = token # default to self, override if road block found + if not road_blocks_invalid: + query_point = centerline.interpolate(0.5, normalized=True).shapely_point + intersecting_roadblock = road_block_map.query_nearest(query_point, max_distance=0.1, all_matches=False) + + # NOTE: if a lane cannot be assigned to a road block, we assume a new lane group with only this lane. + # The lane group id is set to be the same as the lane id in this case. + if len(intersecting_roadblock) > 0: + lane_group_id = road_block_map.ids[intersecting_roadblock[0]] + + # Get topology + incoming = nuscenes_map.get_incoming_lane_ids(token) + outgoing = nuscenes_map.get_outgoing_lane_ids(token) + + lanes.append( + CacheLane( + object_id=token, + lane_group_id=lane_group_id, + left_boundary=left_boundary, + right_boundary=right_boundary, + centerline=centerline, + left_lane_id=None, + right_lane_id=None, + predecessor_ids=incoming, + successor_ids=outgoing, + speed_limit_mps=None, + outline=None, + geometry=None, ) - except Exception: - traceback.print_exc() - continue + ) + return lanes -def _write_nuscenes_intersections(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write intersection data to map_writer. - """ - # road_blocks = nusc_map.road_block - # for block in road_blocks: - # token = block["token"] - # try: - # if "polygon_token" in block: - # polygon = nusc_map.extract_polygon(block["polygon_token"]) - # else: - # continue - # if not polygon.is_valid: - # continue - # except Exception: - # continue - # # Lane group IDs are not directly available; use empty list - # lane_group_ids = [] +def _extract_nuscenes_lane_connectors(nuscenes_map: NuScenesMap, road_edges: List[CacheRoadEdge]) -> List[CacheLane]: + """Helper function to extract lane connectors from a nuScenes map.""" - # map_writer.write_intersection( - # CacheIntersection( - # object_id=token, - # lane_group_ids=lane_group_ids, - # geometry=polygon, - # ) - # ) + # TODO @DanielDauner: consider using connected lanes to estimate the lane width + road_edge_map = OccupancyMap2D(geometries=[road_edge.shapely_linestring for road_edge in road_edges]) -def _write_nuscenes_crosswalks(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write crosswalk data to map_writer. - """ - ped_crossings = nusc_map.ped_crossing - for crossing in ped_crossings: - token = crossing["token"] - try: - if "polygon_token" in crossing: - polygon = nusc_map.extract_polygon(crossing["polygon_token"]) - else: - continue - if not polygon.is_valid: + lane_connectors: List[CacheLane] = [] + for lane_record in nuscenes_map.lane_connector: + lane_connector_token: str = lane_record["token"] + + centerline = extract_nuscenes_centerline(nuscenes_map, lane_record) + + _, nearest_edge_distances = road_edge_map.query_nearest( + centerline.linestring, return_distance=True, all_matches=False + ) + road_edge_distance = nearest_edge_distances[0] if nearest_edge_distances else float("inf") + + lane_half_width = np.clip(road_edge_distance, MIN_LANE_WIDTH / 2.0, MAX_LANE_WIDTH / 2.0) + + left_pts = offset_points_perpendicular(centerline.array, offset=lane_half_width) + right_pts = offset_points_perpendicular(centerline.array, offset=-lane_half_width) + + predecessor_ids = nuscenes_map.get_incoming_lane_ids(lane_connector_token) + successor_ids = nuscenes_map.get_outgoing_lane_ids(lane_connector_token) + + lane_group_id = lane_connector_token + + lane_connectors.append( + CacheLane( + object_id=lane_connector_token, + lane_group_id=lane_group_id, + left_boundary=Polyline2D.from_array(left_pts), + right_boundary=Polyline2D.from_array(right_pts), + centerline=centerline, + left_lane_id=None, # Not directly available in nuscenes + right_lane_id=None, # Not directly available in nuscenes + predecessor_ids=predecessor_ids, + successor_ids=successor_ids, + speed_limit_mps=None, # Default value + outline=None, + geometry=None, + ) + ) + + return lane_connectors + + +def _extract_nuscenes_lane_groups( + nuscenes_map: NuScenesMap, + lanes: List[CacheLane], + lane_connectors: List[CacheLane], + intersection_assignment: Dict[str, int], +) -> List[CacheLaneGroup]: + """Helper function to extract lane groups from a nuScenes map.""" + + lane_groups = [] + lanes_dict = {lane.object_id: lane for lane in lanes + lane_connectors} + + # 1. Gather all lane group ids that were previously assigned in the lanes (either roadblocks of lane themselves) + lane_group_lane_dict: Dict[str, List[str]] = defaultdict(list) + for lane in lanes + lane_connectors: + lane_group_lane_dict[lane.lane_group_id].append(lane.object_id) + + for lane_group_id, lane_ids in lane_group_lane_dict.items(): + + if len(lane_ids) > 1: + lane_centerlines: List[Polyline2D] = [lanes_dict[lane_id].centerline for lane_id in lane_ids] + ordered_lane_indices = order_lanes_left_to_right(lane_centerlines) + left_boundary = lanes_dict[lane_ids[ordered_lane_indices[0]]].left_boundary + right_boundary = lanes_dict[lane_ids[ordered_lane_indices[-1]]].right_boundary + + else: + lane_id = lane_ids[0] + lane = lanes_dict[lane_id] + left_boundary = lane.left_boundary + right_boundary = lane.right_boundary + + # 2. For each lane group, gather predecessor and successor lane groups + predecessor_ids = set() + successor_ids = set() + for lane_id in lane_ids: + lane = lanes_dict[lane_id] + if lane is None: continue - except Exception: - traceback.print_exc() - continue + for pred_id in lane.predecessor_ids: + pred_lane = lanes_dict.get(pred_id) + if pred_lane is not None: + predecessor_ids.add(pred_lane.lane_group_id) + for succ_id in lane.successor_ids: + succ_lane = lanes_dict.get(succ_id) + if succ_lane is not None: + successor_ids.add(succ_lane.lane_group_id) + + intersection_ids = set( + [int(intersection_assignment[lane_id]) for lane_id in lane_ids if lane_id in intersection_assignment] + ) + assert len(intersection_ids) <= 1, "A lane group cannot belong to multiple intersections." + intersection_id = None if len(intersection_ids) == 0 else intersection_ids.pop() + + lane_groups.append( + CacheLaneGroup( + object_id=lane_group_id, + lane_ids=lane_ids, + left_boundary=left_boundary, + right_boundary=right_boundary, + intersection_id=intersection_id, + predecessor_ids=list(predecessor_ids), + successor_ids=list(successor_ids), + outline=None, + geometry=None, + ) + ) + + return lane_groups + +def _write_nuscenes_intersections( + nuscenes_map: NuScenesMap, lane_connectors: List[CacheLane], map_writer: AbstractMapWriter +) -> None: + """Write intersection data to map_writer and return lane-connector to intersection assignment.""" + + intersection_assignment = {} + + # 1. Extract intersections and corresponding polygons + intersection_polygons = [] + for road_segment in nuscenes_map.road_segment: + if road_segment["is_intersection"]: + if "polygon_token" in road_segment: + polygon = nuscenes_map.extract_polygon(road_segment["polygon_token"]) + intersection_polygons.append(polygon) + + # 2. Find lane connectors within each intersection polygon + lane_connector_center_point_dict = { + lane_connector.object_id: lane_connector.centerline.interpolate(0.5, normalized=True).shapely_point + for lane_connector in lane_connectors + } + centerpoint_map = OccupancyMap2D.from_dict(lane_connector_center_point_dict) + for idx, intersection_polygon in enumerate(intersection_polygons): + intersecting_lane_connector_ids = centerpoint_map.intersects(intersection_polygon) + for lane_connector_id in intersecting_lane_connector_ids: + intersection_assignment[lane_connector_id] = idx + + map_writer.write_intersection( + CacheIntersection( + object_id=idx, + lane_group_ids=intersecting_lane_connector_ids, + outline=None, + geometry=intersection_polygon, + ) + ) + + return intersection_assignment + + +def _write_nuscenes_crosswalks(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """Write crosswalk data to map_writer.""" + + crosswalk_polygons = [] + for crossing in nuscenes_map.ped_crossing: + if "polygon_token" in crossing: + polygon = nuscenes_map.extract_polygon(crossing["polygon_token"]) + crosswalk_polygons.append(polygon) + + for idx, polygon in enumerate(crosswalk_polygons): map_writer.write_crosswalk( CacheCrosswalk( - object_id=token, + object_id=idx, geometry=polygon, ) ) -def _write_nuscenes_walkways(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write walkway data to map_writer. - """ - walkways = nusc_map.walkway - for walkway in walkways: - token = walkway["token"] - try: - if "polygon_token" in walkway: - polygon = nusc_map.extract_polygon(walkway["polygon_token"]) - else: - continue - if not polygon.is_valid: - continue - except Exception: - traceback.print_exc() - continue +def _write_nuscenes_walkways(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """Write walkway data to map_writer.""" + walkway_polygons = [] + for walkway_record in nuscenes_map.walkway: + if "polygon_token" in walkway_record: + polygon = nuscenes_map.extract_polygon(walkway_record["polygon_token"]) + walkway_polygons.append(polygon) + for idx, polygon in enumerate(walkway_polygons): map_writer.write_walkway( CacheWalkway( - object_id=token, + object_id=idx, geometry=polygon, ) ) -def _write_nuscenes_carparks(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write carpark data to map_writer. - """ - carpark_areas = nusc_map.carpark_area - for carpark in carpark_areas: - token = carpark["token"] - try: - if "polygon_token" in carpark: - polygon = nusc_map.extract_polygon(carpark["polygon_token"]) - else: - continue - if not polygon.is_valid: - continue - except Exception: - continue +def _write_nuscenes_carparks(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """Write carpark data to map_writer.""" + carpark_polygons = [] + for carpark_record in nuscenes_map.carpark_area: + if "polygon_token" in carpark_record: + polygon = nuscenes_map.extract_polygon(carpark_record["polygon_token"]) + carpark_polygons.append(polygon) + for idx, polygon in enumerate(carpark_polygons): map_writer.write_carpark( CacheCarpark( - object_id=token, + object_id=idx, geometry=polygon, ) ) -def _write_nuscenes_generic_drivables(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write generic drivable areas to map_writer. - """ - # Combine road segments, lanes, and drivable areas - all_drivables = [] - - # Add road segments - for segment in nusc_map.road_segment: - try: - if "polygon_token" in segment: - polygon = nusc_map.extract_polygon(segment["polygon_token"]) - if polygon.is_valid: - all_drivables.append((f"road_segment_{segment['token']}", polygon)) - except Exception: - traceback.print_exc() - continue - - # Add lanes - for lane in nusc_map.lane: - try: - if "polygon_token" in lane: - polygon = nusc_map.extract_polygon(lane["polygon_token"]) - if polygon.is_valid: - all_drivables.append((f"lane_{lane['token']}", polygon)) - except Exception: - traceback.print_exc() - continue - - # Add drivable areas - for road in nusc_map.drivable_area: - try: - if "polygon_token" in road: - polygon = nusc_map.extract_polygon(road["polygon_token"]) - if polygon.is_valid: - all_drivables.append((f"road_{road['token']}", polygon)) - except Exception: - traceback.print_exc() - continue - - for obj_id, geometry in all_drivables: - map_writer.write_generic_drivable( - CacheGenericDrivable( - object_id=obj_id, - geometry=geometry, - ) - ) +def _write_nuscenes_generic_drivables(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """Write generic drivable area data to map_writer.""" + cell_size = 10.0 + drivable_polygons = [] + for drivable_area_record in nuscenes_map.drivable_area: + drivable_area = nuscenes_map.get("drivable_area", drivable_area_record["token"]) + for polygon_token in drivable_area["polygon_tokens"]: + polygon = nuscenes_map.extract_polygon(polygon_token) + split_polygons = split_polygon_by_grid(polygon, cell_size=cell_size) + drivable_polygons.extend(split_polygons) + # drivable_polygons.append(polygon) -def _write_nuscenes_stop_lines(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write stop line data to map_writer. - """ - stop_lines = nusc_map.stop_line - for stop_line in stop_lines: - token = stop_line["token"] - try: - if "polygon_token" in stop_line: - polygon = nusc_map.extract_polygon(stop_line["polygon_token"]) - else: - continue - if not polygon.is_valid: - continue - except Exception: - traceback.print_exc() - continue + for idx, geometry in enumerate(drivable_polygons): + map_writer.write_generic_drivable(CacheGenericDrivable(object_id=idx, geometry=geometry)) - # Note: Stop lines are written as generic drivable for compatibility - map_writer.write_generic_drivable( - CacheGenericDrivable( - object_id=token, - geometry=polygon, - ) - ) +def _write_nuscenes_stop_lines(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """Write stop line data to map_writer.""" + # FIXME: Add stop lines. + # stop_lines = nuscenes_map.stop_line + # for stop_line in stop_lines: + # token = stop_line["token"] + # try: + # if "polygon_token" in stop_line: + # polygon = nuscenes_map.extract_polygon(stop_line["polygon_token"]) + # else: + # continue + # if not polygon.is_valid: + # continue + # except Exception: + # traceback.print_exc() + # continue + + # # Note: Stop lines are written as generic drivable for compatibility + # map_writer.write_generic_drivable( + # CacheGenericDrivable( + # object_id=token, + # geometry=polygon, + # ) + # ) -def _write_nuscenes_road_lines(nusc_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: - """ - Write road line data (dividers) to map_writer. - """ + +def _write_nuscenes_road_lines(nuscenes_map: NuScenesMap, map_writer: AbstractMapWriter) -> None: + """Write road line data (dividers) to map_writer.""" # Process road dividers - road_dividers = nusc_map.road_divider + road_dividers = nuscenes_map.road_divider + running_idx = 0 for divider in road_dividers: - token = divider["token"] - try: - line = nusc_map.extract_line(divider["line_token"]) - if not line.is_valid: - continue - except Exception: - continue + line = nuscenes_map.extract_line(divider["line_token"]) # Determine line type - line_type = _get_road_line_type(divider["line_token"], nusc_map) + line_type = _get_road_line_type(divider["line_token"], nuscenes_map) map_writer.write_road_line( CacheRoadLine( - object_id=token, + object_id=running_idx, road_line_type=line_type, polyline=Polyline3D(LineString(line.coords)), ) ) + running_idx += 1 # Process lane dividers - lane_dividers = nusc_map.lane_divider + lane_dividers = nuscenes_map.lane_divider for divider in lane_dividers: - token = divider["token"] - try: - line = nusc_map.extract_line(divider["line_token"]) - if not line.is_valid: - continue - except Exception: - continue - - line_type = _get_road_line_type(divider["line_token"], nusc_map) + line = nuscenes_map.extract_line(divider["line_token"]) + line_type = _get_road_line_type(divider["line_token"], nuscenes_map) map_writer.write_road_line( CacheRoadLine( - object_id=token, + object_id=running_idx, road_line_type=line_type, polyline=Polyline3D(LineString(line.coords)), ) ) + running_idx += 1 + + +def _extract_nuscenes_road_edges(nuscenes_map: NuScenesMap) -> List[CacheRoadEdge]: + """Helper function to extract road edges from a nuScenes map.""" + drivable_polygons = [] + for drivable_area_record in nuscenes_map.drivable_area: + drivable_area = nuscenes_map.get("drivable_area", drivable_area_record["token"]) + for polygon_token in drivable_area["polygon_tokens"]: + polygon = nuscenes_map.extract_polygon(polygon_token) + drivable_polygons.append(polygon) + + road_edge_linear_rings = get_road_edge_linear_rings(drivable_polygons) + road_edges_linestrings = split_line_geometry_by_max_length(road_edge_linear_rings, MAX_ROAD_EDGE_LENGTH) + + road_edges_cache: List[CacheRoadEdge] = [] + for idx in range(len(road_edges_linestrings)): + road_edges_cache.append( + CacheRoadEdge( + object_id=idx, + road_edge_type=RoadEdgeType.ROAD_EDGE_BOUNDARY, + polyline=Polyline2D.from_linestring(road_edges_linestrings[idx]), + ) + ) + return road_edges_cache -def _get_lane_boundary(lane_token: str, side: str, nusc_map: NuScenesMap) -> Optional[LineString]: - """ - Extract lane boundary geometry for a given side. - """ - lane_record = next((lr for lr in nusc_map.lane if lr["token"] == lane_token), None) - if not lane_record: - return None - - divider_segment_nodes_key = f"{side}_lane_divider_segment_nodes" - if divider_segment_nodes_key in lane_record and lane_record[divider_segment_nodes_key]: - nodes = lane_record[divider_segment_nodes_key] - boundary = LineString([(node["x"], node["y"]) for node in nodes]) - return boundary - - return None - - -def _get_lane_group_boundary(segment_token: str, side: str, nusc_map: NuScenesMap) -> Optional[LineString]: - """ - Extract lane group boundary geometry (simplified). - """ - # This is a simplified implementation; in practice, may need more robust geometry extraction - boundary_type = "road_divider" if side == "left" else "lane_divider" - - # Find the segment geometry - segment = next((rs for rs in nusc_map.road_segment if rs["token"] == segment_token), None) - if not segment: - return None - - try: - segment_geom = nusc_map.extract_polygon(segment["polygon_token"]) - except Exception: - return None - - # Find nearest boundary of the specified type within a threshold - nearest = None - min_dist = float("inf") - - if boundary_type == "road_divider": - records = nusc_map.road_divider - else: - records = nusc_map.lane_divider - - for record in records: - try: - line = nusc_map.extract_line(record["line_token"]) - dist = segment_geom.distance(line) - if dist < 10.0 and dist < min_dist: - min_dist = dist - nearest = line - except Exception: - continue - - return nearest - - -def _get_connected_segments(segment_token: str, nusc_map: NuScenesMap): - """ - Get incoming and outgoing segment connections. - """ - incoming, outgoing = [], [] - - for connector in nusc_map.lane_connector: - if connector.get("outgoing_lane") == segment_token: - incoming.append(connector.get("incoming_lane")) - elif connector.get("incoming_lane") == segment_token: - outgoing.append(connector.get("outgoing_lane")) - - incoming = [id for id in incoming if id is not None] - outgoing = [id for id in outgoing if id is not None] - - return incoming, outgoing +def _get_road_line_type(line_token: str, nuscenes_map: NuScenesMap) -> RoadLineType: + """Map nuscenes line type to RoadLineType.""" -def _get_road_line_type(line_token: str, nusc_map: NuScenesMap) -> RoadLineType: - """ - Map nuscenes line type to RoadLineType. - """ + # FIXME @DanielDauner: Store token to type mapping. Creating mapping for every call is not ideal. nuscenes_to_road_line_type = { "SINGLE_SOLID_WHITE": RoadLineType.SOLID_WHITE, "DOUBLE_DASHED_WHITE": RoadLineType.DOUBLE_DASH_WHITE, @@ -671,7 +481,7 @@ def _get_road_line_type(line_token: str, nusc_map: NuScenesMap) -> RoadLineType: } line_token_to_type = {} - for lane_record in nusc_map.lane: + for lane_record in nuscenes_map.lane: for seg in lane_record.get("left_lane_divider_segments", []): token = seg.get("line_token") seg_type = seg.get("segment_type") @@ -686,34 +496,3 @@ def _get_road_line_type(line_token: str, nusc_map: NuScenesMap) -> RoadLineType: nuscenes_type = line_token_to_type.get(line_token, "UNKNOWN") return nuscenes_to_road_line_type.get(nuscenes_type, RoadLineType.UNKNOWN) - - -def flip_linestring(linestring: LineString) -> LineString: - """ - Flip the direction of a LineString. - """ - return LineString(linestring.coords[::-1]) - - -def lines_same_direction(centerline: LineString, boundary: LineString) -> bool: - """ - Check if centerline and boundary have the same direction. - """ - center_start = np.array(centerline.coords[0]) - center_end = np.array(centerline.coords[-1]) - boundary_start = np.array(boundary.coords[0]) - boundary_end = np.array(boundary.coords[-1]) - - same_dir_dist = np.linalg.norm(center_start - boundary_start) + np.linalg.norm(center_end - boundary_end) - opposite_dir_dist = np.linalg.norm(center_start - boundary_end) + np.linalg.norm(center_end - boundary_start) - - return same_dir_dist <= opposite_dir_dist - - -def align_boundary_direction(centerline: LineString, boundary: LineString) -> LineString: - """ - Align boundary direction with centerline. - """ - if not lines_same_direction(centerline, boundary): - return flip_linestring(boundary) - return boundary diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py index 7592182c..eccf0124 100644 --- a/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py +++ b/src/py123d/conversion/datasets/nuscenes/nuscenes_sensor_io.py @@ -11,16 +11,13 @@ def load_nuscenes_lidar_pcs_from_file(pcd_path: Path, log_metadata: LogMetadata) -> Dict[LiDARType, np.ndarray]: - points = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, 5) + lidar_pc = np.fromfile(pcd_path, dtype=np.float32).reshape(-1, len(NuScenesLidarIndex)) # convert lidar to ego frame lidar_extrinsic = log_metadata.lidar_metadata[LiDARType.LIDAR_TOP].extrinsic - - points[..., NuScenesLidarIndex.XYZ] = convert_points_3d_array_between_origins( + lidar_pc[..., NuScenesLidarIndex.XYZ] = convert_points_3d_array_between_origins( from_origin=lidar_extrinsic, to_origin=StateSE3(0, 0, 0, 1.0, 0, 0, 0), - points_3d_array=points[..., NuScenesLidarIndex.XYZ], + points_3d_array=lidar_pc[..., NuScenesLidarIndex.XYZ], ) - lidar_pcs_dict: Dict[LiDARType, np.ndarray] = {LiDARType.LIDAR_TOP: points} - - return lidar_pcs_dict + return {LiDARType.LIDAR_TOP: lidar_pc} diff --git a/src/py123d/conversion/datasets/nuscenes/nuscenes_constants.py b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_constants.py similarity index 100% rename from src/py123d/conversion/datasets/nuscenes/nuscenes_constants.py rename to src/py123d/conversion/datasets/nuscenes/utils/nuscenes_constants.py diff --git a/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py new file mode 100644 index 00000000..a512dc18 --- /dev/null +++ b/src/py123d/conversion/datasets/nuscenes/utils/nuscenes_map_utils.py @@ -0,0 +1,197 @@ +from typing import Dict, List, Optional, Tuple + +import numpy as np +from scipy.spatial.distance import cdist + +from py123d.common.utils.dependencies import check_dependencies +from py123d.geometry.polyline import Polyline2D + +check_dependencies(["nuscenes"], optional_name="nuscenes") +from nuscenes.map_expansion import arcline_path_utils +from nuscenes.map_expansion.map_api import NuScenesMap + + +def extract_lane_and_boundaries( + nuscenes_map: NuScenesMap, lane_record: Dict +) -> Tuple[Polyline2D, Optional[Polyline2D], Optional[Polyline2D]]: + """Extracts the centerline and left / right boundary from a nusScenes lane. + + :param nuscenes_map: nuScenes map API instance + :param lane_record: lane record dictionary + :return: centerline, left boundary (optional), right boundary (optional) + """ + + # NOTE @DanielDauner: Code adapted from trajdata, Apache 2.0 License. Thank you! + # https://github.com/NVlabs/trajdata/blob/main/src/trajdata/dataset_specific/nusc/nusc_utils.py#L281 + + # Getting the bounding polygon vertices. + lane_polygon_obj = nuscenes_map.get("polygon", lane_record["polygon_token"]) + polygon_nodes = [nuscenes_map.get("node", node_token) for node_token in lane_polygon_obj["exterior_node_tokens"]] + polygon_outline: np.ndarray = np.array([(node["x"], node["y"]) for node in polygon_nodes]) + + # Getting the lane center's points. + centerline = extract_nuscenes_centerline(nuscenes_map, lane_record) + centerline_array = centerline.array + + # Computing the closest lane center point to each bounding polygon vertex. + closest_midlane_pt: np.ndarray = np.argmin(cdist(polygon_outline, centerline_array), axis=1) + # Computing the local direction of the lane at each lane center point. + direction_vectors: np.ndarray = np.diff( + centerline_array, + axis=0, + prepend=centerline_array[[0]] - (centerline_array[[1]] - centerline_array[[0]]), + ) + + # Selecting the direction vectors at the closest lane center point per polygon vertex. + local_dir_vecs: np.ndarray = direction_vectors[closest_midlane_pt] + # Calculating the vectors from the the closest lane center point per polygon vertex to the polygon vertex. + origin_to_polygon_vecs: np.ndarray = polygon_outline - centerline_array[closest_midlane_pt] + + # Computing the perpendicular dot product. + # See https://www.xarg.org/book/linear-algebra/2d-perp-product/ + # If perp_dot_product < 0, then the associated polygon vertex is + # on the right edge of the lane. + perp_dot_product: np.ndarray = ( + local_dir_vecs[:, 0] * origin_to_polygon_vecs[:, 1] - local_dir_vecs[:, 1] * origin_to_polygon_vecs[:, 0] + ) + + # Determining which indices are on the right of the lane center. + on_right: np.ndarray = perp_dot_product < 0 + # Determining the boundary between the left/right polygon vertices + # (they will be together in blocks due to the ordering of the polygon vertices). + idx_changes: int = np.where(np.roll(on_right, 1) < on_right)[0].item() + + if idx_changes > 0: + # If the block of left/right points spreads across the bounds of the array, + # roll it until the boundary between left/right points is at index 0. + # This is important so that the following index selection orders points + # without jumps. + polygon_outline = np.roll(polygon_outline, shift=-idx_changes, axis=0) + on_right = np.roll(on_right, shift=-idx_changes) + + left_polyline_array: np.ndarray = polygon_outline[~on_right] + right_polyline_array: np.ndarray = polygon_outline[on_right] + + # Final ordering check, ensuring that left_pts and right_pts can be combined + # into a polygon without the endpoints intersecting. + # Reversing the one lane edge that does not match the ordering of the midline. + if endpoints_intersect(left_polyline_array, right_polyline_array): + if not order_matches(left_polyline_array, centerline_array): + left_polyline_array = left_polyline_array[::-1] + else: + right_polyline_array = right_polyline_array[::-1] + + left_boundary = Polyline2D.from_array(left_polyline_array) if len(left_polyline_array) > 1 else None + right_boundary = Polyline2D.from_array(right_polyline_array) if len(right_polyline_array) > 1 else None + return centerline, left_boundary, right_boundary + + +def extract_nuscenes_centerline(nuscenes_map: NuScenesMap, lane_record: Dict) -> Polyline2D: + """Extract the centerline of a nuScenes lane. + + :param nuscenes_map: nuScenes map API instance + :param lane_record: lane record dictionary + :return: centerline 2D polyline + """ + + # NOTE @DanielDauner: Code adapted from trajdata, Apache 2.0 License. Thank you! + # https://github.com/NVlabs/trajdata/blob/main/src/trajdata/dataset_specific/nusc/nusc_utils.py#L262 + + # Getting the lane center's points. + curr_lane = nuscenes_map.arcline_path_3.get(lane_record["token"], []) + centerline_array: np.ndarray = np.array(arcline_path_utils.discretize_lane(curr_lane, resolution_meters=0.25))[ + :, :2 + ] + + # For some reason, nuScenes duplicates a few entries + # (likely how they're building their arcline representation). + # We delete those duplicate entries here. + duplicate_check: np.ndarray = np.where( + np.linalg.norm(np.diff(centerline_array, axis=0, prepend=0), axis=1) < 1e-10 + )[0] + if duplicate_check.size > 0: + centerline_array = np.delete(centerline_array, duplicate_check, axis=0) + + return Polyline2D.from_array(centerline_array) + + +def endpoints_intersect(left_edge: np.ndarray, right_edge: np.ndarray) -> bool: + """Check if the line segment connecting the endpoints of left_edge intersects + with the line segment connecting the endpoints of right_edge. + + Forms two segments: (left_edge[-1], left_edge[0]) and (right_edge[-1], right_edge[0]), + then tests if they intersect using the counter-clockwise (CCW) orientation test. + """ + + # NOTE @DanielDauner: Code adapted from trajdata, Apache 2.0 License. Thank you! + # https://github.com/NVlabs/trajdata/blob/main/src/trajdata/utils/map_utils.py#L177 + + def ccw(A, B, C): + return (C[1] - A[1]) * (B[0] - A[0]) > (B[1] - A[1]) * (C[0] - A[0]) + + A, B = left_edge[-1], right_edge[-1] + C, D = right_edge[0], left_edge[0] + return ccw(A, C, D) != ccw(B, C, D) and ccw(A, B, C) != ccw(A, B, D) + + +def order_matches(pts: np.ndarray, ref: np.ndarray) -> bool: + """Check if two polylines have the same ordering direction, by comparing + the distance of their start and end points to the start point of the reference polyline. + """ + + # NOTE @DanielDauner: Code adapted from trajdata, Apache 2.0 License. Thank you! + # https://github.com/NVlabs/trajdata/blob/main/src/trajdata/utils/map_utils.py#L162 + + return np.linalg.norm(pts[0] - ref[0]) <= np.linalg.norm(pts[-1] - ref[0]) + + +def order_lanes_left_to_right(polylines: List[Polyline2D]) -> List[int]: + """ + Order lanes from left to right based on their position. + + :param polylines: List of polylines representing lanes + :return: List of indices representing the order of lanes from left to right + """ + if len(polylines) == 0: + return [] + + # Step 1: Compute the average direction vector across all lanes + all_directions = [] + for polyline in polylines: + + polyline_array = polyline.array + if len(polyline_array) < 2: + continue + start = np.array(polyline_array[0]) + end = np.array(polyline_array[-1]) + direction = end - start + all_directions.append(direction) + + avg_direction = np.mean(all_directions, axis=0) + avg_direction = avg_direction / np.linalg.norm(avg_direction) + + # Step 2: Compute perpendicular vector (left direction) + # Rotate 90 degrees counter-clockwise: (x, y) -> (-y, x) + left_vector = np.array([-avg_direction[1], avg_direction[0]]) + + # Step 3: For each lane, use midpoint of start and end, project onto left vector + lane_positions = [] + for i, polyline in enumerate(polylines): + if len(polyline) == 0: + lane_positions.append((i, 0)) + continue + + start = np.array(polyline[0]) + end = np.array(polyline[-1]) + # Use midpoint of start and end + midpoint = (start + end) / 2 + + # Project midpoint onto the left vector + position = np.dot(midpoint, left_vector) + lane_positions.append((i, position)) + + # Step 4: Sort by position (higher values are more to the left) + lane_positions.sort(key=lambda x: x[1], reverse=True) + + # Return ordered indices + return [idx for idx, _ in lane_positions] diff --git a/src/py123d/conversion/utils/map_utils/road_edge/road_edge_2d_utils.py b/src/py123d/conversion/utils/map_utils/road_edge/road_edge_2d_utils.py index f4cbb094..f8059b12 100644 --- a/src/py123d/conversion/utils/map_utils/road_edge/road_edge_2d_utils.py +++ b/src/py123d/conversion/utils/map_utils/road_edge/road_edge_2d_utils.py @@ -2,7 +2,8 @@ import numpy as np import shapely -from shapely import LinearRing, LineString, Polygon, union_all +from shapely import LinearRing, LineString, Polygon, box, union_all +from shapely.strtree import STRtree def get_road_edge_linear_rings( @@ -10,6 +11,10 @@ def get_road_edge_linear_rings( buffer_distance: float = 0.05, add_interiors: bool = True, ) -> List[LinearRing]: + """ + Helper function to extract road edges (i.e. linear rings) from drivable area polygons. + TODO: Move and rename for general use. + """ def _polygon_to_linear_rings(polygon: Polygon) -> List[LinearRing]: assert polygon.geom_type == "Polygon" @@ -40,7 +45,10 @@ def split_line_geometry_by_max_length( geometries: Union[LineString, LinearRing, List[Union[LineString, LinearRing]]], max_length_meters: float, ) -> List[LineString]: - # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. + """ + Splits LineString or LinearRing geometries into smaller segments based on a maximum length. + TODO: Move and rename for general use. + """ if not isinstance(geometries, list): geometries = [geometries] @@ -61,3 +69,39 @@ def split_line_geometry_by_max_length( all_segments.append(segment) return all_segments + + +def split_polygon_by_grid(polygon: Polygon, cell_size: float) -> List[Polygon]: + """ + Split a polygon by grid-like cells of given size. + TODO: Move and rename for general use. + """ + + minx, miny, maxx, maxy = polygon.bounds + + # Generate all grid cells + x_coords = np.arange(minx, maxx, cell_size) + y_coords = np.arange(miny, maxy, cell_size) + + grid_cells = [box(x, y, x + cell_size, y + cell_size) for x in x_coords for y in y_coords] + + # Build spatial index for fast queries + tree = STRtree(grid_cells) + + # Query cells that potentially intersect + candidate_indices = tree.query(polygon, predicate="intersects") + + cells = [] + for idx in candidate_indices: + cell = grid_cells[idx] + intersection = polygon.intersection(cell) + + if intersection.is_empty: + continue + + if intersection.geom_type == "Polygon": + cells.append(intersection) + elif intersection.geom_type == "MultiPolygon": + cells.extend(intersection.geoms) + + return cells diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py b/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py index a7e0dc36..97b11d73 100644 --- a/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py +++ b/src/py123d/datatypes/maps/gpkg/gpkg_map_objects.py @@ -1,7 +1,6 @@ from __future__ import annotations import ast -import trimesh from functools import cached_property from typing import List, Optional, Union @@ -194,23 +193,9 @@ def centerline(self) -> Polyline3D: @property def outline_3d(self) -> Polyline3D: - left_array = self.left_boundary.array if getattr(self, "left_boundary", None) is not None else np.zeros((0, 3)) - right_array = self.right_boundary.array[::-1] if getattr(self, "right_boundary", None) is not None else np.zeros((0, 3)) - - outline_array = np.vstack((left_array, right_array)) if left_array.size + right_array.size > 0 else np.zeros((0, 3)) - - if outline_array.shape[0] == 0: - # fallback: use shapely polygon generate Polyline3D - poly = getattr(self, "shapely_polygon", None) - if poly is not None: - outline_array = np.array(poly.exterior.coords) - else: - return Polyline3D(np.zeros((0, 3))) - - # close - if outline_array.shape[0] > 0: - outline_array = np.vstack((outline_array, outline_array[0])) - + """Inherited, see superclass.""" + outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) + outline_array = np.vstack((outline_array, outline_array[0])) return Polyline3D.from_linestring(geom.LineString(outline_array)) @property @@ -279,20 +264,8 @@ def right_boundary(self) -> Polyline3D: @property def outline_3d(self) -> Polyline3D: - # get left_array and right_array - left_array = self.left_boundary.array if getattr(self, "left_boundary", None) is not None else np.zeros((0, 3)) - right_array = self.right_boundary.array[::-1] if getattr(self, "right_boundary", None) is not None else np.zeros((0, 3)) - - if left_array.size + right_array.size == 0: - # fallback: use geometry polygon generate - poly = getattr(self, "shapely_polygon", None) - if poly is not None: - outline_array = np.array(poly.exterior.coords) - else: - return Polyline3D(np.zeros((0, 3))) - else: - outline_array = np.vstack((left_array, right_array)) - + """Inherited, see superclass.""" + outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) return Polyline3D.from_linestring(geom.LineString(outline_array)) @property diff --git a/src/py123d/datatypes/maps/gpkg/gpkg_utils.py b/src/py123d/datatypes/maps/gpkg/gpkg_utils.py index ba587077..54dd93e6 100644 --- a/src/py123d/datatypes/maps/gpkg/gpkg_utils.py +++ b/src/py123d/datatypes/maps/gpkg/gpkg_utils.py @@ -21,23 +21,16 @@ def load_gdf_with_geometry_columns(gdf: gpd.GeoDataFrame, geometry_column_names: def get_all_rows_with_value( - elements: gpd.GeoDataFrame, column_label: str, desired_value -) -> gpd.GeoDataFrame: + elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str +) -> gpd.geodataframe.GeoDataFrame: """ - Extract all matching elements by value. - Automatically handles both integer IDs and UUID strings. + Extract all matching elements. Note, if no matching desired_key is found and empty list is returned. + :param elements: data frame from MapsDb. + :param column_label: key to extract from a column. + :param desired_value: key which is compared with the values of column_label entry. + :return: a subset of the original GeoDataFrame containing the matching key. """ - # If the column is of integer type, attempt to convert the desired_value to an integer before comparison. - col_dtype = elements[column_label].dtype - if np.issubdtype(col_dtype, np.integer): - try: - desired_value_int = int(desired_value) - return elements[elements[column_label] == desired_value_int] - except ValueError: - raise ValueError(f"Expected an integer value for column '{column_label}', got '{desired_value}'") - else: - # Otherwise, directly compare it as a string. - return elements[elements[column_label].astype(str) == str(desired_value)] + return elements.iloc[np.where(elements[column_label].to_numpy().astype(int) == int(desired_value))] def get_row_with_value(elements: gpd.geodataframe.GeoDataFrame, column_label: str, desired_value: str) -> gpd.GeoSeries: diff --git a/src/py123d/geometry/utils/polyline_utils.py b/src/py123d/geometry/utils/polyline_utils.py index 20bf299f..d66628e8 100644 --- a/src/py123d/geometry/utils/polyline_utils.py +++ b/src/py123d/geometry/utils/polyline_utils.py @@ -3,6 +3,7 @@ from shapely.geometry import LineString from py123d.geometry.geometry_index import Point2DIndex, StateSE2Index +from py123d.geometry.transform.transform_se2 import translate_2d_along_body_frame def get_linestring_yaws(linestring: LineString) -> npt.NDArray[np.float64]: @@ -41,3 +42,24 @@ def get_path_progress(points_array: npt.NDArray[np.float64]) -> list[float]: points_diff: npt.NDArray[np.float64] = np.concatenate(([x_diff], [y_diff]), axis=0, dtype=np.float64) progress_diff = np.append(0.0, np.linalg.norm(points_diff, axis=0)) return np.cumsum(progress_diff, dtype=np.float64) # type: ignore + + +def offset_points_perpendicular(points_array: npt.NDArray[np.float64], offset: float) -> npt.NDArray[np.float64]: + if points_array.shape[-1] == len(Point2DIndex): + xy = points_array[..., Point2DIndex.XY] + yaws = get_points_2d_yaws(points_array[..., Point2DIndex.XY]) + elif points_array.shape[-1] == len(StateSE2Index): + xy = points_array[..., StateSE2Index.XY] + yaws = points_array[..., StateSE2Index.YAW] + else: + raise ValueError( + f"Invalid points_array shape: {points_array.shape}. Expected last dimension to be {len(Point2DIndex)} or " + f"{len(StateSE2Index)}." + ) + + return translate_2d_along_body_frame( + points_2d=xy, + yaws=yaws, + y_translate=offset, + x_translate=0.0, + ) diff --git a/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml b/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml index 95a3c927..4c9ba050 100644 --- a/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml +++ b/src/py123d/script/config/conversion/datasets/nuscenes_mini_dataset.yaml @@ -6,7 +6,7 @@ nuscenes_dataset: nuscenes_data_root: ${dataset_paths.nuscenes_data_root} nuscenes_map_root: ${dataset_paths.nuscenes_data_root} nuscenes_lanelet2_root: ${dataset_paths.nuscenes_data_root}/lanelet2 - use_lanelet2: False + use_lanelet2: false dataset_converter_config: _target_: py123d.conversion.dataset_converter_config.DatasetConverterConfig diff --git a/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml b/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml index 6bfb4877..86bf8e0b 100644 --- a/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml +++ b/src/py123d/script/config/conversion/map_writer/gpkg_map_writer.yaml @@ -2,3 +2,4 @@ _target_: py123d.conversion.map_writer.gpkg_map_writer.GPKGMapWriter _convert_: 'all' maps_root: ${dataset_paths.py123d_maps_root} +remap_ids: true