From 9546a1c821d23ad8b63addb12734cbcfef4ef2fa Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Wed, 23 Jul 2025 22:06:38 +0200 Subject: [PATCH 01/42] Rolling out CaRL in 123D (demo) --- .../gym/environment/environment_wrapper.py | 10 +-- .../__init__.py | 0 .../abstract_output_converter.py} | 18 ++-- .../action_output_converter.py} | 70 ++++++--------- d123/simulation/simulation_2d_setup.py | 12 +-- notebooks/test_simulation_2d.ipynb | 87 +++++++++++++++---- 6 files changed, 110 insertions(+), 87 deletions(-) rename d123/simulation/gym/environment/{trajectory_builder => output_converter}/__init__.py (100%) rename d123/simulation/gym/environment/{trajectory_builder/abstract_trajectory_builder.py => output_converter/abstract_output_converter.py} (65%) rename d123/simulation/gym/environment/{trajectory_builder/action_trajectory_builder.py => output_converter/action_output_converter.py} (81%) diff --git a/d123/simulation/gym/environment/environment_wrapper.py b/d123/simulation/gym/environment/environment_wrapper.py index b01e76bf..ef0723fe 100644 --- a/d123/simulation/gym/environment/environment_wrapper.py +++ b/d123/simulation/gym/environment/environment_wrapper.py @@ -17,8 +17,8 @@ AbstractSimulationBuilder, ) from d123.simulation.gym.environment.simulation_wrapper import SimulationWrapper -from d123.simulation.gym.environment.trajectory_builder.abstract_trajectory_builder import ( - AbstractTrajectoryBuilder, +from d123.simulation.gym.environment.output_converter.abstract_output_converter import ( + AbstractOutputConverter, ) logger = logging.getLogger(__name__) @@ -35,7 +35,7 @@ def __init__( self, scenario_sampler: AbstractScenarioSampler, simulation_builder: AbstractSimulationBuilder, - trajectory_builder: AbstractTrajectoryBuilder, + output_converter: AbstractOutputConverter, observation_builder: AbstractGymObservation, reward_builder: AbstractRewardBuilder, terminate_on_failure: bool = False, @@ -52,7 +52,7 @@ def __init__( self._scenario_sampler = scenario_sampler self._simulation_builder = simulation_builder - self._trajectory_builder = trajectory_builder + self._trajectory_builder = output_converter self._observation_builder = observation_builder self._reward_builder = reward_builder @@ -68,7 +68,7 @@ def __init__( # Set for super class self.observation_space = observation_builder.get_observation_space() - self.action_space = trajectory_builder.get_action_space() + self.action_space = output_converter.get_action_space() def reset(self, seed: Optional[int] = None, options: Optional[dict] = None): """Inherited, see superclass.""" diff --git a/d123/simulation/gym/environment/trajectory_builder/__init__.py b/d123/simulation/gym/environment/output_converter/__init__.py similarity index 100% rename from d123/simulation/gym/environment/trajectory_builder/__init__.py rename to d123/simulation/gym/environment/output_converter/__init__.py diff --git a/d123/simulation/gym/environment/trajectory_builder/abstract_trajectory_builder.py b/d123/simulation/gym/environment/output_converter/abstract_output_converter.py similarity index 65% rename from d123/simulation/gym/environment/trajectory_builder/abstract_trajectory_builder.py rename to d123/simulation/gym/environment/output_converter/abstract_output_converter.py index fcadf1ed..c045d7bf 100644 --- a/d123/simulation/gym/environment/trajectory_builder/abstract_trajectory_builder.py +++ b/d123/simulation/gym/environment/output_converter/abstract_output_converter.py @@ -4,11 +4,13 @@ import numpy as np import numpy.typing as npt from gymnasium import spaces -from nuplan.common.actor_state.ego_state import EgoState -from nuplan.planning.simulation.trajectory.abstract_trajectory import AbstractTrajectory -class AbstractTrajectoryBuilder(ABC): +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2 +from d123.simulation.planning.planner_output.abstract_planner_output import AbstractPlannerOutput + + +class AbstractOutputConverter(ABC): """Abstract class for building trajectories (nuPlan interface) in a Gym simulation environment.""" @abstractmethod @@ -19,13 +21,13 @@ def get_action_space(self) -> spaces.Space: """ @abstractmethod - def build_trajectory( - self, action: npt.NDArray[np.float32], ego_state: EgoState, info: Dict[str, Any] - ) -> AbstractTrajectory: + def get_planner_output( + self, action: npt.NDArray[np.float32], ego_state: EgoStateSE2, info: Dict[str, Any] + ) -> AbstractPlannerOutput: """ - Builds a trajectory based on the action and the current ego state. + Builds a planner output based on the action and the current ego state. :param action: Action taken by the agent, typically a numpy array. :param ego_state: Current state of the ego vehicle. :param info: Arbitrary information dictionary, for passing information between modules. - :return: Trajectory object of nuPlan. + :return: Planner output object. """ diff --git a/d123/simulation/gym/environment/trajectory_builder/action_trajectory_builder.py b/d123/simulation/gym/environment/output_converter/action_output_converter.py similarity index 81% rename from d123/simulation/gym/environment/trajectory_builder/action_trajectory_builder.py rename to d123/simulation/gym/environment/output_converter/action_output_converter.py index 9e92e42a..ed8aff94 100644 --- a/d123/simulation/gym/environment/trajectory_builder/action_trajectory_builder.py +++ b/d123/simulation/gym/environment/output_converter/action_output_converter.py @@ -2,17 +2,15 @@ import numpy as np import numpy.typing as npt -from carl_nuplan.planning.simulation.trajectory.action_trajectory import ActionTrajectory from gymnasium import spaces -from nuplan.common.actor_state.ego_state import EgoState -from d123.simulation.gym.environment.trajectory_builder.abstract_trajectory_builder import ( - AbstractTrajectoryBuilder, -) +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2 +from d123.simulation.gym.environment.output_converter.abstract_output_converter import AbstractOutputConverter from d123.simulation.gym.policy.ppo.ppo_config import GlobalConfig +from d123.simulation.planning.planner_output.action_planner_output import ActionPlannerOutput -class ActionTrajectoryBuilder(AbstractTrajectoryBuilder): +class ActionOutputConverter(AbstractOutputConverter): """ Default action trajectory builder for training CaRL. TODO: Refactor this class @@ -75,36 +73,22 @@ def get_action_space(self) -> spaces.Space: dtype=np.float32, ) - def build_trajectory( - self, action: npt.NDArray[np.float32], ego_state: EgoState, info: Dict[str, Any] - ) -> ActionTrajectory: + def get_planner_output( + self, action: npt.NDArray[np.float32], ego_state: EgoStateSE2, info: Dict[str, Any] + ) -> ActionPlannerOutput: """Inherited, see superclass.""" assert len(action) == self._config.action_space_dim info["last_action"] = action action_acceleration_normed, action_steering_angle_normed = action - target_steering_rate = self._scale_steering( - action_steering_angle_normed, - ego_state.tire_steering_angle, - ) + target_steering_rate = self._scale_steering(action_steering_angle_normed, ego_state.tire_steering_angle) target_acceleration = self._scale_acceleration( - action_acceleration_normed, - ego_state.dynamic_car_state.rear_axle_acceleration_2d.x, - ) - clipped_steering_rate = self._clip_steering( - target_acceleration, - target_steering_rate, - ego_state, - ) - clipped_acceleration = self._clip_acceleration( - target_acceleration, - clipped_steering_rate, - ego_state, + action_acceleration_normed, ego_state.dynamic_state_se2.acceleration.x ) - - store = info["store"] if "store" in info.keys() else None - return ActionTrajectory(clipped_acceleration, clipped_steering_rate, ego_state, list(action), store) + clipped_steering_rate = self._clip_steering(target_acceleration, target_steering_rate, ego_state) + clipped_acceleration = self._clip_acceleration(target_acceleration, clipped_steering_rate, ego_state) + return ActionPlannerOutput(clipped_acceleration, clipped_steering_rate, ego_state) def _scale_steering(self, action_steering_angle_normed: float, current_steering_angle: float) -> float: """ @@ -137,7 +121,9 @@ def _scale_acceleration(self, action_acceleration_normed: float, current_acceler target_acceleration = (target_acceleration - current_acceleration) / factor + current_acceleration return target_acceleration - def _clip_acceleration(self, target_acceleration: float, target_steering_rate: float, ego_state: EgoState) -> float: + def _clip_acceleration( + self, target_acceleration: float, target_steering_rate: float, ego_state: EgoStateSE2 + ) -> float: """ Clips the acceleration based on the target acceleration, steering rate, and current ego state. :param target_acceleration: Acceleration as targeted by the agent. @@ -146,10 +132,10 @@ def _clip_acceleration(self, target_acceleration: float, target_steering_rate: f :return: Clipped acceleration based on the target acceleration and steering rate. """ - current_acceleration = ego_state.dynamic_car_state.rear_axle_acceleration_2d.x + current_acceleration = ego_state.dynamic_state_se2.acceleration.x if self._disable_reverse_driving: - speed = ego_state.dynamic_car_state.rear_axle_velocity_2d.x + speed = ego_state.dynamic_state_se2.velocity.x updated_speed = speed + target_acceleration * self._dt_control # * self._dt_control if updated_speed < 0: @@ -168,19 +154,15 @@ def _clip_acceleration(self, target_acceleration: float, target_steering_rate: f _max_acceleration = self._scale_max_acceleration if self._clip_angular_adjustment: - rear_axle_to_center_dist = ego_state.car_footprint.rear_axle_to_center_dist + rear_axle_to_center_dist = ego_state.vehicle_parameters.rear_axle_to_center_longitudinal - next_point_velocity_x = ( - ego_state.dynamic_car_state.rear_axle_velocity_2d.x + target_acceleration * self._dt_control - ) + next_point_velocity_x = ego_state.dynamic_state_se2.velocity.x + target_acceleration * self._dt_control next_point_tire_steering_angle = ego_state.tire_steering_angle + target_steering_rate * self._dt_control next_point_angular_velocity = ( - next_point_velocity_x - * np.tan(next_point_tire_steering_angle) - / ego_state.car_footprint.vehicle_parameters.wheel_base + next_point_velocity_x * np.tan(next_point_tire_steering_angle) / ego_state.vehicle_parameters.wheel_base ) next_point_angular_acceleration = ( - next_point_angular_velocity - ego_state.dynamic_car_state.angular_velocity + next_point_angular_velocity - ego_state.dynamic_state_se2.angular_velocity ) / self._dt_control centripetal_acceleration_term = rear_axle_to_center_dist * (next_point_angular_velocity) ** 2 @@ -190,7 +172,7 @@ def _clip_acceleration(self, target_acceleration: float, target_steering_rate: f target_acceleration = np.clip(target_acceleration, -self._scale_max_deceleration, _max_acceleration) return target_acceleration - def _clip_steering(self, target_acceleration: float, target_steering_rate: float, ego_state: EgoState) -> float: + def _clip_steering(self, target_acceleration: float, target_steering_rate: float, ego_state: EgoStateSE2) -> float: """ Clips the steering rate based on the target acceleration and current ego state. :param target_acceleration: Acceleration as targeted by the agent. @@ -203,12 +185,10 @@ def _clip_steering(self, target_acceleration: float, target_steering_rate: float target_steering_angle = current_steering_angle + target_steering_rate * self._dt_control if self._clip_max_abs_yaw_accel is not None: - wheel_base = ego_state.car_footprint.vehicle_parameters.wheel_base - target_velocity = ( - ego_state.dynamic_car_state.rear_axle_velocity_2d.x + target_acceleration * self._dt_control - ) + wheel_base = ego_state.vehicle_parameters.wheel_base + target_velocity = ego_state.dynamic_state_se2.acceleration.x + target_acceleration * self._dt_control - current_angular_velocity = ego_state.dynamic_car_state.angular_velocity + current_angular_velocity = ego_state.dynamic_state_se2.angular_velocity max_abs_yaw_velocity = self._clip_max_abs_yaw_accel * self._dt_control min_angular_velocity = current_angular_velocity - max_abs_yaw_velocity diff --git a/d123/simulation/simulation_2d_setup.py b/d123/simulation/simulation_2d_setup.py index b38df183..737d3e6c 100644 --- a/d123/simulation/simulation_2d_setup.py +++ b/d123/simulation/simulation_2d_setup.py @@ -2,17 +2,7 @@ from d123.simulation.controller.abstract_controller import AbstractEgoController from d123.simulation.observation.abstract_observation import AbstractObservation -from d123.simulation.time_controller.abstract_time_controller import ( - AbstractTimeController, -) - -# from nuplan.planning.scenario_builder.abstract_scenario import AbstractScenario -# from nuplan.planning.simulation.controller.abstract_controller import AbstractEgoController -# from nuplan.planning.simulation.observation.abstract_observation import AbstractObservation -# from nuplan.planning.simulation.planner.abstract_planner import AbstractPlanner -# from nuplan.planning.simulation.simulation_time_controller.abstract_simulation_time_controller import ( -# AbstractSimulationTimeController, -# ) +from d123.simulation.time_controller.abstract_time_controller import AbstractTimeController @dataclass diff --git a/notebooks/test_simulation_2d.ipynb b/notebooks/test_simulation_2d.ipynb index 1aa0fec7..ba913ff3 100644 --- a/notebooks/test_simulation_2d.ipynb +++ b/notebooks/test_simulation_2d.ipynb @@ -33,8 +33,9 @@ " split_names=[split],\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=15.1,\n", + " duration_s=30.1,\n", " history_s=1.0,\n", + " timestamp_threshold_s=10,\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", @@ -82,11 +83,20 @@ "outputs": [], "source": [ "# reset\n", - "from typing import List\n", + "import torch\n", + "from d123.simulation.gym.environment.output_converter.action_output_converter import ActionOutputConverter\n", + "from d123.simulation.gym.policy.ppo.ppo_config import GlobalConfig\n", + "from d123.simulation.gym.policy.ppo.ppo_model import PPOPolicy\n", + "\n", + "\n", + "import numpy as np\n", + "import numpy.typing as npt\n", + "\n", + "\n", "from d123.simulation.gym.environment.gym_observation.raster.raster_gym_observation import RasterGymObservation\n", "from d123.simulation.gym.environment.gym_observation.raster.raster_renderer import RasterRenderer\n", "from d123.simulation.gym.environment.helper.environment_area import RectangleEnvironmentArea\n", - "from d123.simulation.planning.abstract_planner import PlannerInput\n", + "from d123.simulation.planning.abstract_planner import PlannerInitialization, PlannerInput\n", "from d123.simulation.planning.planner_output.action_planner_output import ActionPlannerOutput\n", "\n", "\n", @@ -94,20 +104,54 @@ "gym_observation = RasterGymObservation(environment_area, RasterRenderer(environment_area), inference=True)\n", "gym_observation.reset()\n", "\n", + "checkpoint_path = \"/home/daniel/carl_workspace/carl_nuplan/checkpoints/nuplan_51892_1B/model_best.pth\"\n", + "config = GlobalConfig()\n", + "output_converter = ActionOutputConverter()\n", + "\n", + "device = \"cpu\"\n", + "\n", + "agent = PPOPolicy(\n", + " gym_observation.get_observation_space(),\n", + " output_converter.get_action_space(),\n", + " config=config,\n", + ").to(device)\n", + "\n", + "state_dict = torch.load(checkpoint_path, map_location=device)\n", + "agent.load_state_dict(state_dict, strict=True)\n", + "agent.to(device)\n", + "agent.eval()\n", "\n", - "last_action = [0.0, 0.0] # placeholder\n", - "info = {\"last_action\": last_action}\n", "\n", + "def forward_agent(\n", + " planner_initialization: PlannerInitialization,\n", + " planner_input: PlannerInput,\n", + " last_action: npt.NDArray[np.float32],\n", + ") -> npt.NDArray[np.float32]:\n", + " info = {\"last_action\": last_action}\n", + " obs = gym_observation.get_gym_observation(planner_input, planner_initialization, info)\n", + "\n", + " obs_tensor = {\n", + " \"bev_semantics\": torch.Tensor(obs[\"bev_semantics\"][None, ...]).to(device, dtype=torch.float32),\n", + " \"measurements\": torch.Tensor(obs[\"measurements\"][None, ...]).to(device, dtype=torch.float32),\n", + " \"value_measurements\": torch.Tensor(obs[\"value_measurements\"][None, ...]).to(device, dtype=torch.float32),\n", + " }\n", + "\n", + " with torch.no_grad():\n", + " (action, _, _, _, _, _, _, _, _, _, _) = agent.forward(\n", + " obs_tensor, deterministic=True, lstm_state=None, done=None\n", + " )\n", + " action = action.squeeze().detach().cpu().numpy()\n", + " print(action)\n", + " return action\n", + "\n", + "\n", + "last_action = [0.0, 0.0] # placeholder\n", "rasters = []\n", "\n", - "idx = 1338\n", + "idx = 120\n", "planner_initialization, current_planner_input = simulation2d.reset(scenes[idx])\n", - "rasters.append(gym_observation.get_gym_observation(current_planner_input, planner_initialization, info))\n", "\n", "\n", - "# TODO: Implement action planner output\n", - "# TODO: Further test the simulation object.\n", - "\n", "\n", "def _get_action(planner_input: PlannerInput) -> ActionPlannerOutput:\n", " ego_state, _ = planner_input.history.current_state\n", @@ -117,11 +161,18 @@ "while simulation2d.is_simulation_running():\n", "\n", " # 1. trigger planner\n", - " planner_output = _get_action(current_planner_input)\n", + " # planner_output = _get_action(current_planner_input)\n", + "\n", + " last_action = forward_agent(planner_initialization, current_planner_input, last_action)\n", + " planner_output = output_converter.get_planner_output(\n", + " action=last_action,\n", + " ego_state=current_planner_input.history.current_state[0],\n", + " info={},\n", + " )\n", + " print(planner_output._acceleration, planner_output._steering_rate)\n", "\n", " # 2. step simulation\n", - " current_planner_input = simulation2d.step(planner_output)\n", - " rasters.append(gym_observation.get_gym_observation(current_planner_input, planner_initialization, info))" + " current_planner_input = simulation2d.step(planner_output)" ] }, { @@ -204,11 +255,11 @@ " return rgb_image\n", "\n", "\n", - "numpy_images_to_gif(\n", - " [image_to_rbg(raster[\"bev_semantics\"]) for raster in rasters],\n", - " gif_path=\"/home/daniel/d123_workspace/d123/notebooks/simulation_2d.gif\",\n", - " duration=100,\n", - ")" + "# numpy_images_to_gif(\n", + "# [image_to_rbg(raster[\"bev_semantics\"]) for raster in rasters],\n", + "# gif_path=\"/home/daniel/d123_workspace/d123/notebooks/simulation_2d.gif\",\n", + "# duration=100,\n", + "# )" ] }, { From b0ea9a08fb0feb43c9a7bc25c32084d4b2882468 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 25 Jul 2025 16:47:58 +0200 Subject: [PATCH 02/42] Refactoring and move map conversion into general converter --- ...a_processor.py => carla_data_converter.py} | 72 +++- ..._processor.py => nuplan_data_converter.py} | 66 ++-- .../nuplan/nuplan_data_processor_.py | 332 ------------------ ...conversion.py => nuplan_map_conversion.py} | 0 .../dataset_specific/raw_data_converter.py | 29 ++ .../dataset_specific/raw_data_processor.py | 18 - ...r_builder.py => data_converter_builder.py} | 10 +- .../__init__.py | 0 .../default_dataset_conversion.yaml} | 0 .../script/config/datasets/carla_dataset.yaml | 7 +- .../config/datasets/nuplan_dataset.yaml | 5 +- .../config/datasets/nuplan_mini_dataset.yaml | 5 +- .../datasets/nuplan_private_dataset.yaml | 5 +- ...t_caching.py => run_dataset_conversion.py} | 17 +- d123/script/run_preprocessing.py | 2 +- d123/script/run_simulation.py | 2 +- notebooks/test_simulation_2d.ipynb | 11 +- scripts/dataset/run_log_caching.sh | 2 +- 18 files changed, 172 insertions(+), 411 deletions(-) rename d123/dataset/dataset_specific/carla/{carla_data_processor.py => carla_data_converter.py} (80%) rename d123/dataset/dataset_specific/nuplan/{nuplan_data_processor.py => nuplan_data_converter.py} (88%) delete mode 100644 d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py rename d123/dataset/dataset_specific/nuplan/{map_conversion.py => nuplan_map_conversion.py} (100%) create mode 100644 d123/dataset/dataset_specific/raw_data_converter.py delete mode 100644 d123/dataset/dataset_specific/raw_data_processor.py rename d123/script/builders/{data_processor_builder.py => data_converter_builder.py} (58%) rename d123/script/config/{dataset_caching => dataset_conversion}/__init__.py (100%) rename d123/script/config/{dataset_caching/default_dataset_caching.yaml => dataset_conversion/default_dataset_conversion.yaml} (100%) rename d123/script/{run_dataset_caching.py => run_dataset_conversion.py} (63%) diff --git a/d123/dataset/dataset_specific/carla/carla_data_processor.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py similarity index 80% rename from d123/dataset/dataset_specific/carla/carla_data_processor.py rename to d123/dataset/dataset_specific/carla/carla_data_converter.py index 0bec510b..89ba5c66 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_processor.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -2,6 +2,7 @@ import gzip import hashlib import json +import os from dataclasses import asdict from functools import partial from pathlib import Path @@ -10,6 +11,7 @@ import numpy as np import pyarrow as pa from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map +from traitlets import Any from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import get_carla_lincoln_mkz_2020_parameters @@ -19,12 +21,29 @@ from d123.common.geometry.vector import Vector3DIndex from d123.dataset.arrow.conversion import VehicleParameters from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor +from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive +from d123.dataset.dataset_specific.carla.opendrive.opendrive_converter import OpenDriveConverter +from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import AbstractMap, MapSurfaceType from d123.dataset.maps.abstract_map_objects import AbstractLane from d123.dataset.scene.arrow_scene import get_map_api_from_names +AVAILABLE_CARLA_MAP_LOCATIONS: Final[List[str]] = [ + "Town01", # A small, simple town with a river and several bridges. + "Town02", # A small simple town with a mixture of residential and commercial buildings. + "Town03", # A larger, urban map with a roundabout and large junctions. + "Town04", # A small town embedded in the mountains with a special "figure of 8" infinite highway. + "Town05", # Squared-grid town with cross junctions and a bridge. It has multiple lanes per direction. Useful to perform lane changes. + "Town06", # Long many lane highways with many highway entrances and exits. It also has a Michigan left. + "Town07", # A rural environment with narrow roads, corn, barns and hardly any traffic lights. + "Town10HD", # A downtown urban environment with skyscrapers, residential buildings and an ocean promenade. + "Town11", # A Large Map that is undecorated. Serves as a proof of concept for the Large Maps feature. + "Town12", # A Large Map with numerous different regions, including high-rise, residential and rural environments. + "Town13", # ??? + "Town15", # ??? +] + CARLA_DT: Final[float] = 0.1 # [s] TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE: Final[float] = 1.0 # [m] SORT_BY_TIMESTAMP: Final[bool] = True @@ -52,15 +71,17 @@ def create_token(input_data: str) -> str: return hash_obj.hexdigest()[:16] -class CarlaDataProcessor(RawDataProcessor): +class CarlaDataConverter(RawDataConverter): + def __init__( self, splits: List[str], log_path: Union[Path, str], output_path: Union[Path, str], - force_data_conversion: bool, + force_log_conversion: bool, + force_map_conversion: bool, ) -> None: - super().__init__(force_data_conversion) + super().__init__(force_log_conversion, force_map_conversion) for split in splits: assert ( split in self.get_available_splits() @@ -82,7 +103,18 @@ def get_available_splits(self) -> List[str]: """Returns a list of available raw data types.""" return ["carla"] # TODO: fix the placeholder - def convert(self, worker: WorkerPool) -> None: + def convert_maps(self, worker: WorkerPool) -> None: + worker_map( + worker, + partial( + convert_carla_map_to_gpkg, + output_path=self._output_path, + force_map_conversion=self.force_map_conversion, + ), + list(AVAILABLE_CARLA_MAP_LOCATIONS), + ) + + def convert_logs(self, worker: WorkerPool) -> None: log_args = [ { @@ -95,13 +127,35 @@ def convert(self, worker: WorkerPool) -> None: ] worker_map( - worker, partial(convert_carla_log_to_arrow, force_data_conversion=self.force_data_conversion), log_args + worker, partial(convert_carla_log_to_arrow, force_log_conversion=self.force_log_conversion), log_args ) +def convert_carla_map_to_gpkg(map_names: List[str], output_path: Path, force_map_conversion: bool) -> List[Any]: + for map_name in map_names: + map_path = output_path / "maps" / f"carla_{map_name}.gpkg" + if force_map_conversion or not map_path.exists(): + map_path.unlink(missing_ok=True) + assert os.environ["CARLA_ROOT"] is not None + CARLA_ROOT = Path(os.environ["CARLA_ROOT"]) + + if map_name not in ["Town11", "Town12", "Town13", "Town15"]: + carla_maps_root = CARLA_ROOT / "CarlaUE4" / "Content" / "Carla" / "Maps" / "OpenDrive" + carla_map_path = carla_maps_root / f"{map_name}.xodr" + else: + carla_map_path = ( + CARLA_ROOT / "CarlaUE4" / "Content" / "Carla" / "Maps" / map_name / "OpenDrive" / f"{map_name}.xodr" + ) + + OpenDriveConverter(OpenDrive.parse_from_file(carla_map_path)).run(f"carla_{map_name.lower()}") + + return [] + + def convert_carla_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], force_data_conversion: bool -) -> None: + args: List[Dict[str, Union[List[str], List[Path]]]], + force_log_conversion: bool, +) -> List[Any]: def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> None: for log_info in args: log_path: Path = log_info["log_path"] @@ -110,7 +164,7 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> log_file_path = output_path / split / f"{log_path.stem}.arrow" - if force_data_conversion or not log_file_path.exists(): + if force_log_conversion or not log_file_path.exists(): log_file_path.unlink(missing_ok=True) if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py similarity index 88% rename from d123/dataset/dataset_specific/nuplan/nuplan_data_processor.py rename to d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index e345e3fc..e7654074 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -4,18 +4,15 @@ from dataclasses import asdict from functools import partial from pathlib import Path -from typing import Dict, Final, List, Tuple, Union +from typing import Any, Dict, Final, List, Tuple, Union -import psutil import pyarrow as pa import yaml from nuplan.database.nuplan_db.nuplan_scenario_queries import get_images_from_lidar_tokens from nuplan.database.nuplan_db_orm.lidar_box import LidarBox from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB -from nuplan.planning.simulation.observation.observation_type import ( - CameraChannel, -) +from nuplan.planning.simulation.observation.observation_type import CameraChannel from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map import d123.dataset.dataset_specific.nuplan.utils as nuplan_utils @@ -31,20 +28,14 @@ from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor +from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import MAP_LOCATIONS, NuPlanMapConverter +from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata TARGET_DT: Final[float] = 0.1 NUPLAN_DT: Final[float] = 0.05 SORT_BY_TIMESTAMP: Final[bool] = True -NUPLAN_FULL_MAP_NAME_DICT: Final[Dict[str, str]] = { - "boston": "us-ma-boston", - "singapore": "sg-one-north", - "las_vegas": "us-nv-las-vegas-strip", - "pittsburgh": "us-pa-pittsburgh-hazelwood", -} - NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { "green": TrafficLightStatus.GREEN, "red": TrafficLightStatus.RED, @@ -71,16 +62,17 @@ def create_splits_logs() -> Dict[str, List[str]]: return splits["log_splits"] -class NuplanDataProcessor(RawDataProcessor): +class NuplanDataConverter(RawDataConverter): def __init__( self, splits: List[str], log_path: Union[Path, str], output_path: Union[Path, str], sensor_path: Union[Path, str], - force_data_conversion: bool, + force_log_conversion: bool, + force_map_conversion: bool, ) -> None: - super().__init__(force_data_conversion) + super().__init__(force_log_conversion, force_map_conversion) for split in splits: assert ( split in self.get_available_splits() @@ -131,7 +123,18 @@ def get_available_splits(self) -> List[str]: "nuplan_private_test", ] - def convert(self, worker: WorkerPool) -> None: + def convert_maps(self, worker: WorkerPool) -> None: + worker_map( + worker, + partial( + convert_nuplan_map_to_gpkg, + output_path=self._output_path, + force_map_conversion=self.force_map_conversion, + ), + list(MAP_LOCATIONS), + ) + + def convert_logs(self, worker: WorkerPool) -> None: log_args = [ { "log_path": log_path, @@ -146,18 +149,26 @@ def convert(self, worker: WorkerPool) -> None: partial( convert_nuplan_log_to_arrow, output_path=self._output_path, - force_data_conversion=self.force_data_conversion, + force_log_conversion=self.force_log_conversion, ), log_args, ) +def convert_nuplan_map_to_gpkg(map_names: List[str], output_path: Path, force_map_conversion: bool) -> List[Any]: + for map_name in map_names: + map_path = output_path / "maps" / f"nuplan_{map_name}.gpkg" + if force_map_conversion or not map_path.exists(): + map_path.unlink(missing_ok=True) + NuPlanMapConverter(output_path / "maps").convert(map_name=map_name) + return [] + + def convert_nuplan_log_to_arrow( args: List[Dict[str, Union[List[str], List[Path]]]], output_path: Path, - force_data_conversion: bool, -) -> None: - process = psutil.Process(os.getpid()) + force_log_conversion: bool, +) -> List[Any]: for log_info in args: log_path: Path = log_info["log_path"] split: str = log_info["split"] @@ -168,7 +179,7 @@ def convert_nuplan_log_to_arrow( log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None) log_file_path = output_path / split / f"{log_path.stem}.arrow" - if force_data_conversion or not log_file_path.exists(): + if force_log_conversion or not log_file_path.exists(): log_file_path.unlink(missing_ok=True) if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) @@ -186,8 +197,8 @@ def convert_nuplan_log_to_arrow( ("traffic_light_types", pa.list_(pa.int16())), ("scenario_tag", pa.list_(pa.string())), ("route_lane_group_ids", pa.list_(pa.int64())), - ("front_cam_demo", pa.binary()), - ("front_cam_transform", pa.list_(pa.float64())), + # ("front_cam_demo", pa.binary()), + # ("front_cam_transform", pa.list_(pa.float64())), ] ) metadata = LogMetadata( @@ -211,7 +222,6 @@ def convert_nuplan_log_to_arrow( log_db.remove_ref() del recording_schema, vehicle_parameters, log_db gc.collect() - print(f"{os.getpid()} Memory usage: {process.memory_info().rss / 1024 / 1024:.1f} MB") return [] @@ -237,7 +247,7 @@ def _write_recording_table( for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") if len(roadblock_id) > 0 ] - front_cam_demo, front_cam_transform = _extract_front_cam_demo(lidar_pc, source_log_path) + # front_cam_demo, front_cam_transform = _extract_front_cam_demo(lidar_pc, source_log_path) row_data = { "token": [lidar_pc_token], @@ -251,8 +261,8 @@ def _write_recording_table( "traffic_light_types": [traffic_light_types], "scenario_tag": [_extract_scenario_tag(log_db, lidar_pc_token)], "route_lane_group_ids": [route_lane_group_ids], - "front_cam_demo": [front_cam_demo], - "front_cam_transform": [front_cam_transform], + # "front_cam_demo": [front_cam_demo], + # "front_cam_transform": [front_cam_transform], } batch = pa.record_batch(row_data, schema=recording_schema) writer.write_batch(batch) diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py deleted file mode 100644 index ea690921..00000000 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_processor_.py +++ /dev/null @@ -1,332 +0,0 @@ -import gc -import json -import os -from dataclasses import asdict -from functools import partial -from pathlib import Path -from typing import Dict, Final, List, Tuple, Union - -import psutil -import pyarrow as pa -import yaml -from nuplan.database.nuplan_db_orm.lidar_box import LidarBox -from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc -from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map - -import d123.dataset.dataset_specific.nuplan.utils as nuplan_utils -from d123.common.datatypes.detection.detection import TrafficLightStatus -from d123.common.datatypes.detection.detection_types import DetectionType -from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index -from d123.common.datatypes.vehicle_state.vehicle_parameters import ( - get_nuplan_pacifica_parameters, - rear_axle_se3_to_center_se3, -) -from d123.common.geometry.base import StateSE3 -from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3, BoundingBoxSE3Index -from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL -from d123.common.geometry.vector import Vector3D, Vector3DIndex -from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor -from d123.dataset.logs.log_metadata import LogMetadata - -TARGET_DT: Final[float] = 0.1 -NUPLAN_DT: Final[float] = 0.05 -SORT_BY_TIMESTAMP: Final[bool] = True - -NUPLAN_FULL_MAP_NAME_DICT: Final[Dict[str, str]] = { - "boston": "us-ma-boston", - "singapore": "sg-one-north", - "las_vegas": "us-nv-las-vegas-strip", - "pittsburgh": "us-pa-pittsburgh-hazelwood", -} - -NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { - "green": TrafficLightStatus.GREEN, - "red": TrafficLightStatus.RED, - "unknown": TrafficLightStatus.UNKNOWN, -} -NUPLAN_DETECTION_NAME_DICT = { - "vehicle": DetectionType.VEHICLE, - "bicycle": DetectionType.BICYCLE, - "pedestrian": DetectionType.PEDESTRIAN, - "traffic_cone": DetectionType.TRAFFIC_CONE, - "barrier": DetectionType.BARRIER, - "czone_sign": DetectionType.CZONE_SIGN, - "generic_object": DetectionType.GENERIC_OBJECT, -} - -NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) - - -def create_splits_logs() -> Dict[str, List[str]]: - yaml_filepath = Path(nuplan_utils.__path__[0]) / "log_splits.yaml" - with open(yaml_filepath, "r") as stream: - splits = yaml.safe_load(stream) - - return splits["log_splits"] - - -class NuplanDataProcessor(RawDataProcessor): - def __init__( - self, - splits: List[str], - log_path: Union[Path, str], - output_path: Union[Path, str], - force_data_conversion: bool, - ) -> None: - super().__init__(force_data_conversion) - for split in splits: - assert ( - split in self.get_available_splits() - ), f"Split {split} is not available. Available splits: {self.available_splits}" - - self._splits: List[str] = splits - self._log_path: Path = Path(log_path) - self._output_path: Path = Path(output_path) - self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() - self._target_dt: float = 0.1 - - def _collect_log_paths(self) -> Dict[str, List[Path]]: - # NOTE: the nuplan mini folder has an internal train, val, test structure, all stored in "mini". - # The complete dataset is saved in the "trainval" folder (train and val), or in the "test" folder (for test). - subsplit_log_names: Dict[str, List[str]] = create_splits_logs() - log_paths_per_split: Dict[str, List[Path]] = {} - - for split in self._splits: - subsplit = split.split("_")[-1] - assert subsplit in ["train", "val", "test"] - if split in ["nuplan_train", "nuplan_val"]: - log_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "splits" / "trainval" - elif split in ["nuplan_test"]: - log_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "splits" / "test" - elif split in ["nuplan_mini_train", "nuplan_mini_val", "nuplan_mini_test"]: - log_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "splits" / "mini" - - all_log_files_in_path = [log_file for log_file in log_path.glob("*.db")] - all_log_names = set([str(log_file.stem) for log_file in all_log_files_in_path]) - split_log_names = set(subsplit_log_names[subsplit]) - log_paths = [log_path / f"{log_name}.db" for log_name in list(all_log_names & split_log_names)] - log_paths_per_split[split] = log_paths - - return log_paths_per_split - - def get_available_splits(self) -> List[str]: - return [ - "nuplan_train", - "nuplan_val", - "nuplan_test", - "nuplan_mini_train", - "nuplan_mini_val", - "nuplan_mini_test", - ] - - def convert(self, worker: WorkerPool) -> None: - log_args = [ - { - "log_path": log_path, - "split": split, - } - for split, log_paths in self._log_paths_per_split.items() - for log_path in log_paths - ] - - worker_map( - worker, - partial( - convert_nuplan_log_to_arrow, - output_path=self._output_path, - force_data_conversion=self.force_data_conversion, - ), - log_args, - ) - - -def convert_nuplan_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], - output_path: Path, - force_data_conversion: bool, -) -> None: - process = psutil.Process(os.getpid()) - for log_info in args: - log_path: Path = log_info["log_path"] - split: str = log_info["split"] - - if not log_path.exists(): - raise FileNotFoundError(f"Log path {log_path} does not exist.") - - log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None) - log_file_path = output_path / split / f"{log_path.stem}.arrow" - - if force_data_conversion or not log_file_path.exists(): - log_file_path.unlink(missing_ok=True) - if not log_file_path.parent.exists(): - log_file_path.parent.mkdir(parents=True, exist_ok=True) - - recording_schema = pa.schema( - [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ] - ) - metadata = LogMetadata( - dataset="nuplan", - log_name=log_db.log_name, - location=log_db.log.map_version, - timestep_seconds=TARGET_DT, - map_has_z=False, - ) - vehicle_parameters = get_nuplan_pacifica_parameters() - recording_schema = recording_schema.with_metadata( - { - "log_metadata": json.dumps(asdict(metadata)), - "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), - } - ) - - _write_recording_table(log_db, recording_schema, log_file_path) - - log_db.detach_tables() - log_db.remove_ref() - del recording_schema, vehicle_parameters, log_db - gc.collect() - print(f"{os.getpid()} Memory usage: {process.memory_info().rss / 1024 / 1024:.1f} MB") - return [] - - -def _write_recording_table(log_db: NuPlanDB, recording_schema: pa.schema, log_file_path: Path) -> None: - - # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: - with pa.OSFile(str(log_file_path), "wb") as sink: - with pa.ipc.new_file(sink, recording_schema) as writer: - step_interval: float = int(TARGET_DT / NUPLAN_DT) - for lidar_pc in log_db.lidar_pc[::step_interval]: - lidar_pc_token: str = lidar_pc.token - detections_state, detections_velocity, detections_token, detections_types = _extract_detections( - lidar_pc - ) - traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) - route_lane_group_ids = [ - int(roadblock_id) - for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") - if len(roadblock_id) > 0 - ] - - row_data = { - "token": [lidar_pc_token], - "timestamp": [lidar_pc.timestamp], - "detections_state": [detections_state], - "detections_velocity": [detections_velocity], - "detections_token": [detections_token], - "detections_type": [detections_types], - "ego_states": [_extract_ego_state(lidar_pc)], - "traffic_light_ids": [traffic_light_ids], - "traffic_light_types": [traffic_light_types], - "scenario_tag": [_extract_scenario_tag(log_db, lidar_pc_token)], - "route_lane_group_ids": [route_lane_group_ids], - } - batch = pa.record_batch(row_data, schema=recording_schema) - writer.write_batch(batch) - del batch, row_data, detections_state, detections_velocity, detections_token, detections_types - - if SORT_BY_TIMESTAMP: - recording_table = open_arrow_table(log_file_path) - recording_table = recording_table.sort_by([("timestamp", "ascending")]) - write_arrow_table(recording_table, log_file_path) - - -def _extract_detections(lidar_pc: LidarPc) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: - detections_state: List[List[float]] = [] - detections_velocity: List[List[float]] = [] - detections_token: List[str] = [] - detections_types: List[int] = [] - - for lidar_box in lidar_pc.lidar_boxes: - lidar_box: LidarBox - center = StateSE3( - x=lidar_box.x, - y=lidar_box.y, - z=lidar_box.z, - roll=DEFAULT_ROLL, - pitch=DEFAULT_PITCH, - yaw=lidar_box.yaw, - ) - bounding_box_se3 = BoundingBoxSE3(center, lidar_box.length, lidar_box.width, lidar_box.height) - - detections_state.append(bounding_box_se3.array) - detections_velocity.append(lidar_box.velocity) - detections_token.append(lidar_box.track_token) - detections_types.append(int(NUPLAN_DETECTION_NAME_DICT[lidar_box.category.name])) - - return detections_state, detections_velocity, detections_token, detections_types - - -def _extract_ego_state(lidar_pc: LidarPc) -> List[float]: - - yaw, pitch, roll = lidar_pc.ego_pose.quaternion.yaw_pitch_roll - vehicle_parameters = get_nuplan_pacifica_parameters() - # vehicle_parameters = get_pacifica_parameters() - - rear_axle_pose = StateSE3( - x=lidar_pc.ego_pose.x, - y=lidar_pc.ego_pose.y, - z=lidar_pc.ego_pose.z, - roll=roll, - pitch=pitch, - yaw=yaw, - ) - # NOTE: The height to rear axle is not provided the dataset and is merely approximated. - center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) - dynamic_state = DynamicStateSE3( - velocity=Vector3D( - x=lidar_pc.ego_pose.vx, - y=lidar_pc.ego_pose.vy, - z=lidar_pc.ego_pose.vz, - ), - acceleration=Vector3D( - x=lidar_pc.ego_pose.acceleration_x, - y=lidar_pc.ego_pose.acceleration_y, - z=lidar_pc.ego_pose.acceleration_z, - ), - angular_velocity=Vector3D( - x=lidar_pc.ego_pose.angular_rate_x, - y=lidar_pc.ego_pose.angular_rate_y, - z=lidar_pc.ego_pose.angular_rate_z, - ), - ) - - return EgoStateSE3( - center_se3=center, - dynamic_state_se3=dynamic_state, - vehicle_parameters=vehicle_parameters, - timepoint=None, - ).array.tolist() - - -def _extract_traffic_lights(log_db: NuPlanDB, lidar_pc_token: str) -> Tuple[List[int], List[int]]: - traffic_light_ids: List[int] = [] - traffic_light_types: List[int] = [] - traffic_lights = log_db.traffic_light_status.select_many(lidar_pc_token=lidar_pc_token) - for traffic_light in traffic_lights: - traffic_light_ids.append(int(traffic_light.lane_connector_id)) - traffic_light_types.append(int(NUPLAN_TRAFFIC_STATUS_DICT[traffic_light.status].value)) - return traffic_light_ids, traffic_light_types - - -def _extract_scenario_tag(log_db: NuPlanDB, lidar_pc_token: str) -> List[str]: - - scenario_tags = [ - scenario_tag.type for scenario_tag in log_db.scenario_tag.select_many(lidar_pc_token=lidar_pc_token) - ] - if len(scenario_tags) == 0: - scenario_tags = ["unknown"] - return scenario_tags diff --git a/d123/dataset/dataset_specific/nuplan/map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py similarity index 100% rename from d123/dataset/dataset_specific/nuplan/map_conversion.py rename to d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py diff --git a/d123/dataset/dataset_specific/raw_data_converter.py b/d123/dataset/dataset_specific/raw_data_converter.py new file mode 100644 index 00000000..d9cf8f17 --- /dev/null +++ b/d123/dataset/dataset_specific/raw_data_converter.py @@ -0,0 +1,29 @@ +import abc +from typing import List + +from nuplan.planning.utils.multithreading.worker_utils import WorkerPool + + +class RawDataConverter(abc.ABC): + + def __init__(self, force_log_conversion: bool, force_map_conversion: bool) -> None: + self.force_log_conversion = force_log_conversion + self.force_map_conversion = force_map_conversion + + @abc.abstractmethod + def get_available_splits(self) -> List[str]: + """Returns a list of available raw data types.""" + + @abc.abstractmethod + def convert_maps(self, worker: WorkerPool) -> None: + """ + Convert maps in raw data format to the uniform 123D format. + :param worker: The worker pool to use for parallel processing. + """ + + @abc.abstractmethod + def convert_logs(self, worker: WorkerPool) -> None: + """ + Convert logs in raw data format to the uniform 123D format. + :param worker: The worker pool to use for parallel processing. + """ diff --git a/d123/dataset/dataset_specific/raw_data_processor.py b/d123/dataset/dataset_specific/raw_data_processor.py deleted file mode 100644 index 74292c1d..00000000 --- a/d123/dataset/dataset_specific/raw_data_processor.py +++ /dev/null @@ -1,18 +0,0 @@ -import abc -from typing import List - -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool - - -class RawDataProcessor(abc.ABC): - - def __init__(self, force_data_conversion: bool) -> None: - self.force_data_conversion = force_data_conversion - - @abc.abstractmethod - def get_available_splits(self) -> List[str]: - """Returns a list of available raw data types.""" - - @abc.abstractmethod - def convert(self, worker: WorkerPool) -> None: - pass diff --git a/d123/script/builders/data_processor_builder.py b/d123/script/builders/data_converter_builder.py similarity index 58% rename from d123/script/builders/data_processor_builder.py rename to d123/script/builders/data_converter_builder.py index cd83cc28..6047b008 100644 --- a/d123/script/builders/data_processor_builder.py +++ b/d123/script/builders/data_converter_builder.py @@ -5,17 +5,17 @@ from nuplan.planning.script.builders.utils.utils_type import validate_type from omegaconf import DictConfig -from d123.dataset.dataset_specific.raw_data_processor import RawDataProcessor +from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter logger = logging.getLogger(__name__) -def build_data_processor(cfg: DictConfig) -> List[RawDataProcessor]: +def build_data_converter(cfg: DictConfig) -> List[RawDataConverter]: logger.info("Building RawDataProcessor...") - instantiated_datasets: List[RawDataProcessor] = [] + instantiated_datasets: List[RawDataConverter] = [] for dataset_type in cfg.values(): - processor: RawDataProcessor = instantiate(dataset_type) - validate_type(processor, RawDataProcessor) + processor: RawDataConverter = instantiate(dataset_type) + validate_type(processor, RawDataConverter) instantiated_datasets.append(processor) logger.info("Building RawDataProcessor...DONE!") diff --git a/d123/script/config/dataset_caching/__init__.py b/d123/script/config/dataset_conversion/__init__.py similarity index 100% rename from d123/script/config/dataset_caching/__init__.py rename to d123/script/config/dataset_conversion/__init__.py diff --git a/d123/script/config/dataset_caching/default_dataset_caching.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml similarity index 100% rename from d123/script/config/dataset_caching/default_dataset_caching.yaml rename to d123/script/config/dataset_conversion/default_dataset_conversion.yaml diff --git a/d123/script/config/datasets/carla_dataset.yaml b/d123/script/config/datasets/carla_dataset.yaml index b331aeee..dcab53de 100644 --- a/d123/script/config/datasets/carla_dataset.yaml +++ b/d123/script/config/datasets/carla_dataset.yaml @@ -1,9 +1,10 @@ carla_dataset: - _target_: d123.dataset.dataset_specific.carla.carla_data_processor.CarlaDataProcessor + _target_: d123.dataset.dataset_specific.carla.carla_data_converter.CarlaDataConverter _convert_: 'all' - force_data_conversion: True - splits: ["carla"] log_path: "${oc.env:HOME}/carla_workspace/data" output_path: ${d123_data_root} + + force_log_conversion: True + force_map_conversion: True diff --git a/d123/script/config/datasets/nuplan_dataset.yaml b/d123/script/config/datasets/nuplan_dataset.yaml index 7d949028..d7839c86 100644 --- a/d123/script/config/datasets/nuplan_dataset.yaml +++ b/d123/script/config/datasets/nuplan_dataset.yaml @@ -1,8 +1,9 @@ nuplan_dataset: - _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_processor.NuplanDataProcessor + _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' - force_data_conversion: True + force_log_conversion: True + force_map_conversion: True splits: ["nuplan_train", "nuplan_val", "nuplan_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder diff --git a/d123/script/config/datasets/nuplan_mini_dataset.yaml b/d123/script/config/datasets/nuplan_mini_dataset.yaml index 9e4ca222..92337524 100644 --- a/d123/script/config/datasets/nuplan_mini_dataset.yaml +++ b/d123/script/config/datasets/nuplan_mini_dataset.yaml @@ -1,5 +1,5 @@ nuplan_mini_dataset: - _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_processor.NuplanDataProcessor + _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' @@ -9,4 +9,5 @@ nuplan_mini_dataset: log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder output_path: ${d123_data_root} - force_data_conversion: True + force_log_conversion: True + force_map_conversion: True diff --git a/d123/script/config/datasets/nuplan_private_dataset.yaml b/d123/script/config/datasets/nuplan_private_dataset.yaml index a618e5a0..56402001 100644 --- a/d123/script/config/datasets/nuplan_private_dataset.yaml +++ b/d123/script/config/datasets/nuplan_private_dataset.yaml @@ -1,8 +1,9 @@ nuplan_private_dataset: - _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_processor.NuplanDataProcessor + _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' - force_data_conversion: True + force_log_conversion: True + force_map_conversion: True splits: ["nuplan_private_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder diff --git a/d123/script/run_dataset_caching.py b/d123/script/run_dataset_conversion.py similarity index 63% rename from d123/script/run_dataset_caching.py rename to d123/script/run_dataset_conversion.py index dbf40de0..62ad1e30 100644 --- a/d123/script/run_dataset_caching.py +++ b/d123/script/run_dataset_conversion.py @@ -5,13 +5,13 @@ from nuplan.planning.script.builders.logging_builder import build_logger from omegaconf import DictConfig -from d123.script.builders.data_processor_builder import RawDataProcessor, build_data_processor +from d123.script.builders.data_converter_builder import RawDataConverter, build_data_converter from d123.script.builders.worker_pool_builder import build_worker logger = logging.getLogger(__name__) -CONFIG_PATH = "config/dataset_caching" -CONFIG_NAME = "default_dataset_caching" +CONFIG_PATH = "config/dataset_conversion" +CONFIG_NAME = "default_dataset_conversion" @hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None) @@ -28,10 +28,17 @@ def main(cfg: DictConfig) -> None: # Precompute and cache all features logger.info("Starting Dataset Caching...") - data_processors: List[RawDataProcessor] = build_data_processor(cfg.datasets) + data_processors: List[RawDataConverter] = build_data_converter(cfg.datasets) for data_processor in data_processors: + logger.info(f"Processing dataset: {data_processor.__class__.__name__}") - data_processor.convert(worker=worker) + + data_processor.convert_maps(worker=worker) + logger.info(f"Finished maps: {data_processor.__class__.__name__}") + + data_processor.convert_logs(worker=worker) + logger.info(f"Finished logs: {data_processor.__class__.__name__}") + logger.info(f"Finished processing dataset: {data_processor.__class__.__name__}") diff --git a/d123/script/run_preprocessing.py b/d123/script/run_preprocessing.py index 85c9974e..d8c426d7 100644 --- a/d123/script/run_preprocessing.py +++ b/d123/script/run_preprocessing.py @@ -12,7 +12,7 @@ from d123.dataset.scene.abstract_scene import AbstractScene from d123.script.builders.scene_builder_builder import build_scene_builder from d123.script.builders.scene_filter_builder import build_scene_filter -from d123.script.run_dataset_caching import build_worker +from d123.script.run_dataset_conversion import build_worker from d123.training.feature_builder.smart_feature_builder import SMARTFeatureBuilder logger = logging.getLogger(__name__) diff --git a/d123/script/run_simulation.py b/d123/script/run_simulation.py index 84b3ec94..a1be003c 100644 --- a/d123/script/run_simulation.py +++ b/d123/script/run_simulation.py @@ -13,7 +13,7 @@ from d123.dataset.scene.abstract_scene import AbstractScene from d123.script.builders.scene_builder_builder import build_scene_builder from d123.script.builders.scene_filter_builder import build_scene_filter -from d123.script.run_dataset_caching import build_worker +from d123.script.run_dataset_conversion import build_worker from d123.simulation.gym.demo_gym_env import DemoGymEnv from d123.simulation.metrics.sim_agents.sim_agents import get_sim_agents_metrics diff --git a/notebooks/test_simulation_2d.ipynb b/notebooks/test_simulation_2d.ipynb index ba913ff3..26dad379 100644 --- a/notebooks/test_simulation_2d.ipynb +++ b/notebooks/test_simulation_2d.ipynb @@ -70,7 +70,6 @@ " )\n", "\n", "\n", - "\n", "simulation_2d_setup = get_simulation_2d_setup()\n", "simulation2d = Simulation2D(simulation_2d_setup)" ] @@ -148,7 +147,7 @@ "last_action = [0.0, 0.0] # placeholder\n", "rasters = []\n", "\n", - "idx = 120\n", + "idx = 1\n", "planner_initialization, current_planner_input = simulation2d.reset(scenes[idx])\n", "\n", "\n", @@ -388,6 +387,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/scripts/dataset/run_log_caching.sh b/scripts/dataset/run_log_caching.sh index ef62ff58..4d278d15 100644 --- a/scripts/dataset/run_log_caching.sh +++ b/scripts/dataset/run_log_caching.sh @@ -1,7 +1,7 @@ -python $D123_DEVKIT_ROOT/d123/script/run_dataset_caching.py \ +python $D123_DEVKIT_ROOT/d123/script/run_dataset_conversion.py \ experiment_name="caching" \ worker.threads_per_node=10 # worker=single_machine_thread_pool From d94bfbe7c46034741a70470feedaca36d43f7d09 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 27 Jul 2025 15:22:05 +0200 Subject: [PATCH 03/42] Push simple lidar loading for nuplan. --- d123/common/datatypes/sensor/camera.py | 13 ++ .../{camera => sensor}/camera_parameters.py | 0 d123/common/datatypes/sensor/lidar.py | 10 ++ d123/common/visualization/viser/utils.py | 2 +- d123/dataset/arrow/conversion.py | 14 ++ .../dataset_specific/nuplan/load_sensor.py | 21 +++ .../nuplan/nuplan_data_converter.py | 15 +- d123/dataset/scene/abstract_scene.py | 5 + d123/dataset/scene/arrow_scene.py | 6 + .../default_dataset_conversion.yaml | 7 +- .../script/config/datasets/carla_dataset.yaml | 4 +- .../config/datasets/nuplan_dataset.yaml | 6 +- .../config/datasets/nuplan_mini_dataset.yaml | 5 +- .../datasets/nuplan_private_dataset.yaml | 6 +- notebooks/nuplan_sensor_loading.ipynb | 126 +++++++++++------ notebooks/scene_sensor_loading.ipynb | 132 ++++++++++++++++++ 16 files changed, 314 insertions(+), 58 deletions(-) create mode 100644 d123/common/datatypes/sensor/camera.py rename d123/common/datatypes/{camera => sensor}/camera_parameters.py (100%) create mode 100644 d123/common/datatypes/sensor/lidar.py create mode 100644 d123/dataset/dataset_specific/nuplan/load_sensor.py create mode 100644 notebooks/scene_sensor_loading.ipynb diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py new file mode 100644 index 00000000..f0c371df --- /dev/null +++ b/d123/common/datatypes/sensor/camera.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass + +import numpy as np + + +@dataclass +class Camera: + + pass + + def get_view_matrix(self) -> np.ndarray: + # Compute the view matrix based on the camera's position and orientation + pass diff --git a/d123/common/datatypes/camera/camera_parameters.py b/d123/common/datatypes/sensor/camera_parameters.py similarity index 100% rename from d123/common/datatypes/camera/camera_parameters.py rename to d123/common/datatypes/sensor/camera_parameters.py diff --git a/d123/common/datatypes/sensor/lidar.py b/d123/common/datatypes/sensor/lidar.py new file mode 100644 index 00000000..c455279c --- /dev/null +++ b/d123/common/datatypes/sensor/lidar.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt + + +@dataclass +class LiDAR: + + point_cloud: npt.NDArray[np.float32] diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 1ab23c95..ef3fe732 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -2,7 +2,7 @@ import numpy.typing as npt import trimesh -from d123.common.datatypes.camera.camera_parameters import get_nuplan_camera_parameters +from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 10c5489c..4d1b7877 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -14,11 +14,13 @@ TrafficLightStatus, ) from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.lidar import LiDAR from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.vector import Vector3D +from d123.dataset.dataset_specific.nuplan.load_sensor import load_lidar_from_path from d123.dataset.maps.abstract_map import List @@ -76,3 +78,15 @@ def get_traffic_light_detections_from_arrow_table(arrow_table: pa.Table, index: traffic_light_detections.append(traffic_light_detection) return TrafficLightDetectionWrapper(traffic_light_detections=traffic_light_detections) + + +def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int) -> LiDAR: + assert "lidar" in arrow_table.schema.names, '"lidar" field not found in Arrow table schema.' + lidar_data = arrow_table["lidar"][index] + if isinstance(lidar_data.as_py(), str): + # TODO: Handle other formats and datasets + lidar = load_lidar_from_path(lidar_data.as_py()) + else: + raise NotImplementedError("Only string file paths for lidar data are supported.") + + return lidar diff --git a/d123/dataset/dataset_specific/nuplan/load_sensor.py b/d123/dataset/dataset_specific/nuplan/load_sensor.py new file mode 100644 index 00000000..878be2ff --- /dev/null +++ b/d123/dataset/dataset_specific/nuplan/load_sensor.py @@ -0,0 +1,21 @@ +import io +import os +from pathlib import Path + +from nuplan.database.utils.pointclouds.lidar import LidarPointCloud + +from d123.common.datatypes.sensor.lidar import LiDAR + +NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) + + +def load_lidar_from_path(filename: str) -> LiDAR: + lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / filename + + assert lidar_full_path.exists(), f"LiDAR file not found: {lidar_full_path}" + + with open(lidar_full_path, "rb") as fp: + buffer = io.BytesIO(fp.read()) + + points = LidarPointCloud.from_buffer(buffer, "pcd").points + return LiDAR(point_cloud=points) # Reshape to diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index e7654074..bbb17424 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -4,7 +4,7 @@ from dataclasses import asdict from functools import partial from pathlib import Path -from typing import Any, Dict, Final, List, Tuple, Union +from typing import Any, Dict, Final, List, Optional, Tuple, Union import pyarrow as pa import yaml @@ -197,6 +197,7 @@ def convert_nuplan_log_to_arrow( ("traffic_light_types", pa.list_(pa.int16())), ("scenario_tag", pa.list_(pa.string())), ("route_lane_group_ids", pa.list_(pa.int64())), + ("lidar", pa.string()), # ("front_cam_demo", pa.binary()), # ("front_cam_transform", pa.list_(pa.float64())), ] @@ -248,7 +249,6 @@ def _write_recording_table( if len(roadblock_id) > 0 ] # front_cam_demo, front_cam_transform = _extract_front_cam_demo(lidar_pc, source_log_path) - row_data = { "token": [lidar_pc_token], "timestamp": [lidar_pc.timestamp], @@ -261,6 +261,7 @@ def _write_recording_table( "traffic_light_types": [traffic_light_types], "scenario_tag": [_extract_scenario_tag(log_db, lidar_pc_token)], "route_lane_group_ids": [route_lane_group_ids], + "lidar": [_extract_lidar(lidar_pc)], # "front_cam_demo": [front_cam_demo], # "front_cam_transform": [front_cam_transform], } @@ -381,3 +382,13 @@ def _extract_front_cam_demo(lidar_pc: LidarPc, source_log_path: Path) -> Tuple[b front_cam_demo = f.read() return front_cam_demo, front_cam_transform + + +def _extract_lidar(lidar_pc: LidarPc) -> Optional[str]: + + lidar: Optional[str] = None + lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / lidar_pc.filename + if lidar_full_path.exists(): + lidar = lidar_pc.filename + + return lidar diff --git a/d123/dataset/scene/abstract_scene.py b/d123/dataset/scene/abstract_scene.py index 95bf7ddf..2d55a95f 100644 --- a/d123/dataset/scene/abstract_scene.py +++ b/d123/dataset/scene/abstract_scene.py @@ -8,6 +8,7 @@ from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording +from d123.common.datatypes.sensor.lidar import LiDAR from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.dataset.logs.log_metadata import LogMetadata @@ -73,6 +74,10 @@ def get_route_lane_group_ids(self, iteration: int) -> List[int]: def get_front_cam_demo(self, iteration: int) -> Image: raise NotImplementedError + @abc.abstractmethod + def get_lidar_at_iteration(self, iteration: int) -> LiDAR: + raise NotImplementedError + def open(self) -> None: pass diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index b0c7a28d..2718078a 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -10,12 +10,14 @@ from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording +from d123.common.datatypes.sensor.lidar import LiDAR from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.dataset.arrow.conversion import ( get_box_detections_from_arrow_table, get_ego_vehicle_state_from_arrow_table, + get_lidar_from_arrow_table, get_timepoint_from_arrow_table, get_traffic_light_detections_from_arrow_table, ) @@ -134,6 +136,10 @@ def get_front_cam_demo(self, iteration: int) -> Image: jpg_data = self._recording_table["front_cam_demo"][table_index].as_py() return Image.open(io.BytesIO(jpg_data)) + def get_lidar_at_iteration(self, iteration: int) -> LiDAR: + self._lazy_initialize() + return get_lidar_from_arrow_table(self._recording_table, self._get_table_index(iteration)) + def _lazy_initialize(self) -> None: self.open() diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 13a9e6da..07b8be63 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,7 +15,8 @@ defaults: - default_dataset_paths - _self_ - datasets: - # - nuplan_private_dataset - - carla_dataset + - nuplan_private_dataset + # - carla_dataset -force_feature_computation: True +force_log_conversion: True +force_map_conversion: False diff --git a/d123/script/config/datasets/carla_dataset.yaml b/d123/script/config/datasets/carla_dataset.yaml index dcab53de..b98bba2b 100644 --- a/d123/script/config/datasets/carla_dataset.yaml +++ b/d123/script/config/datasets/carla_dataset.yaml @@ -6,5 +6,5 @@ carla_dataset: log_path: "${oc.env:HOME}/carla_workspace/data" output_path: ${d123_data_root} - force_log_conversion: True - force_map_conversion: True + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} diff --git a/d123/script/config/datasets/nuplan_dataset.yaml b/d123/script/config/datasets/nuplan_dataset.yaml index d7839c86..c20793af 100644 --- a/d123/script/config/datasets/nuplan_dataset.yaml +++ b/d123/script/config/datasets/nuplan_dataset.yaml @@ -2,9 +2,9 @@ nuplan_dataset: _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' - force_log_conversion: True - force_map_conversion: True - splits: ["nuplan_train", "nuplan_val", "nuplan_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder output_path: ${d123_data_root} + + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} diff --git a/d123/script/config/datasets/nuplan_mini_dataset.yaml b/d123/script/config/datasets/nuplan_mini_dataset.yaml index 92337524..68bf33a0 100644 --- a/d123/script/config/datasets/nuplan_mini_dataset.yaml +++ b/d123/script/config/datasets/nuplan_mini_dataset.yaml @@ -5,9 +5,8 @@ nuplan_mini_dataset: splits: ["nuplan_mini_train", "nuplan_mini_val", "nuplan_mini_test"] # splits: ["nuplan_mini_val"] - log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder output_path: ${d123_data_root} - force_log_conversion: True - force_map_conversion: True + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} diff --git a/d123/script/config/datasets/nuplan_private_dataset.yaml b/d123/script/config/datasets/nuplan_private_dataset.yaml index 56402001..e69147cc 100644 --- a/d123/script/config/datasets/nuplan_private_dataset.yaml +++ b/d123/script/config/datasets/nuplan_private_dataset.yaml @@ -2,10 +2,10 @@ nuplan_private_dataset: _target_: d123.dataset.dataset_specific.nuplan.nuplan_data_converter.NuplanDataConverter _convert_: 'all' - force_log_conversion: True - force_map_conversion: True - splits: ["nuplan_private_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder sensor_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/sensor output_path: ${d123_data_root} + + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} diff --git a/notebooks/nuplan_sensor_loading.ipynb b/notebooks/nuplan_sensor_loading.ipynb index 9618c8de..10182cec 100644 --- a/notebooks/nuplan_sensor_loading.ipynb +++ b/notebooks/nuplan_sensor_loading.ipynb @@ -22,6 +22,16 @@ "id": "1", "metadata": {}, "outputs": [], + "source": [ + "from d123.dataset.dataset_specific.nuplan.nuplan_data_converter import NuplanDataConverter" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], "source": [ "def process_images_to_arrow(input_folder, output_file, batch_size=100):\n", " \"\"\"\n", @@ -190,7 +200,7 @@ { "cell_type": "code", "execution_count": null, - "id": "2", + "id": "3", "metadata": {}, "outputs": [], "source": [ @@ -200,73 +210,79 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], "source": [ - "# Read IPC file\n", - "import io\n", + "# # Read IPC file\n", + "# import io\n", "\n", - "import time\n", + "# import time\n", "\n", - "# with pa.OSFile(\"test.arrow\", 'rb') as source:\n", - "# with ipc.open_file(source) as reader:\n", - "# table = reader.read_all()\n", + "# # with pa.OSFile(\"test.arrow\", 'rb') as source:\n", + "# # with ipc.open_file(source) as reader:\n", + "# # table = reader.read_all()\n", "\n", - "with pa.ipc.open_file(\n", - " pa.memory_map(\"/home/daniel/d123_workspace/data/nuplan_private_test/2021.07.25.16.16.23_veh-26_02446_02589.arrow\")\n", - ") as reader:\n", - " # This doesn't load data into memory yet!\n", - " table = reader.read_all()\n", + "# with pa.ipc.open_file(\n", + "# pa.memory_map(\"/home/daniel/d123_workspace/data/nuplan_private_test/2021.07.25.16.16.23_veh-26_02446_02589.arrow\")\n", + "# ) as reader:\n", + "# # This doesn't load data into memory yet!\n", + "# table = reader.read_all()\n", "\n", "\n", - "print(len(table))\n", - "start = time.time()\n", - "# Extract JPG data\n", - "jpg_data = table[\"front_cam_demo\"][500].as_py()\n", - "read_image = Image.open(io.BytesIO(jpg_data))\n", + "# print(len(table))\n", + "# start = time.time()\n", + "# # Extract JPG data\n", + "# jpg_data = table[\"front_cam_demo\"][500].as_py()\n", + "# read_image = Image.open(io.BytesIO(jpg_data))\n", "\n", - "# read_image = read_image.convert(\"RGB\") # Ensure it's in RGB format\n", - "read_image = np.array(read_image)\n", - "print(read_image.dtype)\n", - "print(f\"Image loaded in {time.time() - start:.4f} seconds\")\n", + "# # read_image = read_image.convert(\"RGB\") # Ensure it's in RGB format\n", + "# read_image = np.array(read_image)\n", + "# print(read_image.dtype)\n", + "# print(f\"Image loaded in {time.time() - start:.4f} seconds\")\n", "\n", - "import matplotlib.pyplot as plt\n", + "# import matplotlib.pyplot as plt\n", "\n", - "plt.imshow(read_image)" + "# plt.imshow(read_image)" ] }, { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [ "from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB\n", "\n", "NUPLAN_DATA_ROOT = Path(os.environ[\"NUPLAN_DATA_ROOT\"])\n", - "log_path = \"/mnt/nvme/nuplan/dataset/nuplan-v1.1/splits/private_test/2021.09.29.17.35.58_veh-44_01671_01819.db\"\n", + "NUPLAN_SENSOR_ROOT = Path(\"/media/nvme1/nuplan/dataset/nuplan-v1.1/sensor_blobs\")\n", + "\n", + "\n", + "log_path = \"/media/nvme1/nuplan/dataset/nuplan-v1.1/splits/private_test/2021.07.01.21.22.09_veh-14_00016_00656.db\"\n", "log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None)\n", "\n", "\n", - "log_db.image" + "log_db.lidar_pc.filename" ] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [], "source": [ + "import io\n", "import pickle\n", "from nuplan.database.nuplan_db.nuplan_scenario_queries import get_images_from_lidar_tokens, get_cameras\n", "from nuplan.planning.scenario_builder.nuplan_db.nuplan_scenario import NuPlanScenario, CameraChannel, LidarChannel\n", + "\n", + "from nuplan.database.utils.pointclouds.lidar import LidarPointCloud, PointCloud\n", "from pyquaternion import Quaternion\n", "\n", "\n", - "NUPLAN_DB_PATH = \"/mnt/nvme/nuplan/dataset/nuplan-v1.1/splits/private_test\"\n", + "NUPLAN_DB_PATH = \"/media/nvme1/nuplan/dataset/nuplan-v1.1/splits/private_test\"\n", "\n", "\n", "def get_log_cam_info(log):\n", @@ -295,21 +311,49 @@ "images = []\n", "for lidar_pc in log_db.lidar_pc[::2]:\n", "\n", - " front_image = get_images_from_lidar_tokens(log_path, [lidar_pc.token], [str(CameraChannel.CAM_F0.value)])\n", - " parameters = get_log_cam_info(log_db.log)\n", - " print(parameters)\n", + " # front_image = get_images_from_lidar_tokens(log_path, [lidar_pc.token], [str(CameraChannel.CAM_F0.value)])\n", + " # parameters = get_log_cam_info(log_db.log)\n", + " # print(parameters)\n", + "\n", + " # images.append(list(front_image))\n", + " lidar_path = NUPLAN_SENSOR_ROOT / lidar_pc.filename\n", + " with open(lidar_path, \"rb\") as fp:\n", + " buffer = io.BytesIO(fp.read())\n", + "\n", + " pcd_data = buffer.read() # type: ignore\n", + "\n", + " # headers_list = []\n", + " # with io.BytesIO(pcd_data) as stream:\n", + " # while True:\n", + " # line = stream.readline().decode('utf8').strip()\n", + " # print(line)\n", + " # if line.startswith('#'):\n", + " # continue\n", + " # columns = line.split()\n", + " # key = columns[0].lower()\n", + " # val = columns[1:] if len(columns) > 2 else columns[1]\n", + " # headers_list.append((key, val))\n", + "\n", + " # if key == 'data':\n", + " # break\n", + "\n", + " pc = PointCloud.parse(pcd_data)\n", + "\n", + " # print(pc.header)\n", + " # pcd_data\n", "\n", - " images.append(list(front_image))\n", + " # break\n", "\n", + " # print(isinstance(pcd_data, bytes))\n", "\n", - "for image in images[0]:\n", - " print(image)" + " # lidar_point_cloud = LidarPointCloud.from_buffer(buffer, \"pcd\")\n", + " # print(lidar_point_cloud.points.shape)" ] }, { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -347,7 +391,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -362,7 +406,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -372,7 +416,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -408,7 +452,7 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "11", "metadata": {}, "outputs": [], "source": [] @@ -430,7 +474,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.21" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/notebooks/scene_sensor_loading.ipynb b/notebooks/scene_sensor_loading.ipynb new file mode 100644 index 00000000..4b25abab --- /dev/null +++ b/notebooks/scene_sensor_loading.ipynb @@ -0,0 +1,132 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "# from nuplan.planning.utils.multithreading.worker_ray impo\n", + "# rt RayDistributed\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "import os, psutil\n", + "\n", + "\n", + "def print_memory_usage():\n", + " process = psutil.Process(os.getpid())\n", + " memory_info = process.memory_info()\n", + " print(f\"Memory usage: {memory_info.rss / 1024 ** 2:.2f} MB\")\n", + "\n", + "\n", + "split = \"nuplan_private_test\"\n", + "\n", + "log_names = [\"2021.07.01.21.22.09_veh-14_00016_00656\"]\n", + "scene_tokens = None\n", + "log_names = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=[split],\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=15.1,\n", + " history_s=1.0,\n", + " timestamp_threshold_s=15.0,\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(\"num scenes found\", len(scenes))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "from d123.dataset.scene.arrow_scene import ArrowScene\n", + "\n", + "\n", + "scene: ArrowScene = scenes[12]\n", + "\n", + "scene.open()\n", + "\n", + "lidar = scene.get_lidar_at_iteration(0)\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "pc = lidar.point_cloud\n", + "pc.shape" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "plt.scatter(pc[0, :], pc[1, :], s=0.1, alpha=0.5)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From f574fc7fdc5aef4e02b2cc1bc3bd0a29af743759 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 27 Jul 2025 19:06:07 +0200 Subject: [PATCH 04/42] Add camera and lidar loading from path in nuPlan. --- d123/common/datatypes/sensor/camera.py | 83 +++++++- .../datatypes/sensor/camera_parameters.py | 31 --- d123/common/visualization/viser/utils.py | 4 +- d123/dataset/arrow/conversion.py | 57 ++++- .../carla/carla_data_converter.py | 9 +- .../dataset_specific/nuplan/load_sensor.py | 22 +- .../nuplan/nuplan_data_converter.py | 193 +++++++++++------ .../dataset_specific/raw_data_converter.py | 23 +- d123/dataset/scene/abstract_scene.py | 5 +- d123/dataset/scene/arrow_scene.py | 38 ++-- .../config/datasets/nuplan_dataset.yaml | 12 +- .../config/datasets/nuplan_mini_dataset.yaml | 13 +- .../datasets/nuplan_private_dataset.yaml | 13 +- docs/camera.md | 201 ++++++++++++++++++ notebooks/nuplan_sensor_loading.ipynb | 41 +++- notebooks/scene_sensor_loading.ipynb | 19 +- 16 files changed, 587 insertions(+), 177 deletions(-) delete mode 100644 d123/common/datatypes/sensor/camera_parameters.py create mode 100644 docs/camera.md diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py index f0c371df..edc8ad7c 100644 --- a/d123/common/datatypes/sensor/camera.py +++ b/d123/common/datatypes/sensor/camera.py @@ -1,12 +1,93 @@ +from __future__ import annotations + +import json from dataclasses import dataclass +from typing import Any, Dict import numpy as np +import numpy.typing as npt + +from d123.common.utils.enums import SerialIntEnum + + +class CameraType(SerialIntEnum): + """ + Enum for cameras in d123. + """ + + CAM_F0 = 0 + CAM_B0 = 1 + CAM_L0 = 2 + CAM_L1 = 3 + CAM_L2 = 4 + CAM_R0 = 5 + CAM_R1 = 6 + CAM_R2 = 7 + + +@dataclass +class CameraMetadata: + + camera_type: CameraType + width: int + height: int + intrinsic: npt.NDArray[np.float64] # 3x3 matrix + distortion: npt.NDArray[np.float64] # 5x1 vector + translation: npt.NDArray[np.float64] # 3x1 vector + rotation: npt.NDArray[np.float64] # 3x3 matrix + + def to_dict(self) -> Dict[str, Any]: + return { + "camera_type": int(self.camera_type), + "width": self.width, + "height": self.height, + "intrinsic": self.intrinsic.tolist(), + "distortion": self.distortion.tolist(), + "translation": self.translation.tolist(), + "rotation": self.rotation.tolist(), + } + + @classmethod + def from_dict(cls, json_dict: Dict[str, Any]) -> CameraMetadata: + return cls( + camera_type=CameraType(json_dict["camera_type"]), + width=json_dict["width"], + height=json_dict["height"], + intrinsic=np.array(json_dict["intrinsic"]), + distortion=np.array(json_dict["distortion"]), + translation=np.array(json_dict["translation"]), + rotation=np.array(json_dict["rotation"]), + ) + + +def camera_metadata_dict_to_json(camera_metadata: Dict[CameraType, CameraMetadata]) -> Dict[str, Dict[str, Any]]: + """ + Converts a dictionary of CameraMetadata to a JSON-serializable format. + :param camera_metadata: Dictionary of CameraMetadata. + :return: JSON-serializable dictionary. + """ + camera_metadata_dict = {str(camera_type): metadata.to_dict() for camera_type, metadata in camera_metadata.items()} + return json.dumps(camera_metadata_dict) + + +def camera_metadata_dict_from_json(json_dict: Dict[str, Dict[str, Any]]) -> Dict[CameraType, CameraMetadata]: + """ + Converts a JSON-serializable dictionary back to a dictionary of CameraMetadata. + :param json_dict: JSON-serializable dictionary. + :return: Dictionary of CameraMetadata. + """ + camera_metadata_dict = json.loads(json_dict) + return { + CameraType.deserialize(camera_type): CameraMetadata.from_dict(metadata) + for camera_type, metadata in camera_metadata_dict.items() + } @dataclass class Camera: - pass + metadata: CameraMetadata + image: npt.NDArray[np.uint8] def get_view_matrix(self) -> np.ndarray: # Compute the view matrix based on the camera's position and orientation diff --git a/d123/common/datatypes/sensor/camera_parameters.py b/d123/common/datatypes/sensor/camera_parameters.py deleted file mode 100644 index db06db93..00000000 --- a/d123/common/datatypes/sensor/camera_parameters.py +++ /dev/null @@ -1,31 +0,0 @@ -from dataclasses import dataclass - -import numpy as np -import numpy.typing as npt - - -@dataclass -class CameraParameters: - - intrinsics: npt.NDArray[np.float64] # 3x3 matrix - distortion: npt.NDArray[np.float64] # 5x1 vector - translation: npt.NDArray[np.float64] # 3x1 vector - rotation: npt.NDArray[np.float64] # 3x3 matrix - - -def get_nuplan_camera_parameters() -> CameraParameters: - # return CameraParameters(focal_length=focal_length, image_size=image_size) - return CameraParameters( - intrinsics=np.array( - [[1.545e03, 0.000e00, 9.600e02], [0.000e00, 1.545e03, 5.600e02], [0.000e00, 0.000e00, 1.000e00]] - ), - distortion=np.array([-0.356123, 0.172545, -0.00213, 0.000464, -0.05231]), - translation=np.array([1.66433035e00, -1.32379618e-03, 1.57190200e00]), - rotation=np.array( - [ - [-0.00395669, -0.03969443, 0.99920403], - [-0.99971496, -0.02336898, -0.00488707], - [0.02354437, -0.99893856, -0.03959065], - ] - ), - ) diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index ef3fe732..469b49d5 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -2,7 +2,7 @@ import numpy.typing as npt import trimesh -from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters +# from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D @@ -167,7 +167,7 @@ def _get_camera_pose_demo(scene: AbstractScene, iteration: int) -> StateSE3: rear_axle = StateSE3.from_array(rear_axle_array) - camera_parameters = get_nuplan_camera_parameters() + camera_parameters = None # Get camera parameters camera_translation = camera_parameters.translation # 3x1 vector as array diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 4d1b7877..9d8781da 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -1,8 +1,14 @@ # TODO: rename this file and potentially move somewhere more appropriate. +import os +from pathlib import Path +from typing import Dict, Optional + import numpy as np +import numpy.typing as npt import pyarrow as pa +from PIL import Image from d123.common.datatypes.detection.detection import ( BoxDetection, @@ -14,15 +20,22 @@ TrafficLightStatus, ) from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import Camera, CameraMetadata from d123.common.datatypes.sensor.lidar import LiDAR from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.vector import Vector3D -from d123.dataset.dataset_specific.nuplan.load_sensor import load_lidar_from_path +from d123.dataset.dataset_specific.nuplan.load_sensor import load_nuplan_lidar_from_path +from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import List +DATASET_SENSOR_ROOT: Dict[str, Path] = { + "nuplan": Path(os.environ["NUPLAN_DATA_ROOT"]) / "nuplan-v1.1" / "sensor_blobs", + # "carla": Path(os.environ["CARLA_DATA_ROOT"]) / "sensor_blobs", +} + def get_timepoint_from_arrow_table(arrow_table: pa.Table, index: int) -> TimePoint: return TimePoint.from_us(arrow_table["timestamp"][index].as_py()) @@ -80,12 +93,44 @@ def get_traffic_light_detections_from_arrow_table(arrow_table: pa.Table, index: return TrafficLightDetectionWrapper(traffic_light_detections=traffic_light_detections) -def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int) -> LiDAR: +def get_camera_from_arrow_table( + arrow_table: pa.Table, + index: int, + camera_metadata: CameraMetadata, + log_metadata: LogMetadata, +) -> Camera: + + table_data = arrow_table[camera_metadata.camera_type.serialize()][index].as_py() + image: Optional[npt.NDArray[np.uint8]] = None + + if isinstance(table_data, str): + sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] + full_image_path = sensor_root / table_data + assert full_image_path.exists(), f"Camera file not found: {full_image_path}" + img = Image.open(full_image_path) + img.load() + image = np.asarray(img, dtype=np.uint8) + else: + raise NotImplementedError("Only string file paths for camera data are supported.") + + return Camera( + metadata=camera_metadata, + image=image, + ) + + +def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int, log_metadata: LogMetadata) -> LiDAR: assert "lidar" in arrow_table.schema.names, '"lidar" field not found in Arrow table schema.' - lidar_data = arrow_table["lidar"][index] - if isinstance(lidar_data.as_py(), str): - # TODO: Handle other formats and datasets - lidar = load_lidar_from_path(lidar_data.as_py()) + lidar_data = arrow_table["lidar"][index].as_py() + if isinstance(lidar_data, str): + sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] + full_lidar_path = sensor_root / lidar_data + assert full_lidar_path.exists(), f"LiDAR file not found: {full_lidar_path}" + if log_metadata.dataset == "nuplan": + lidar = load_nuplan_lidar_from_path(full_lidar_path) + else: + raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") + else: raise NotImplementedError("Only string file paths for lidar data are supported.") diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 89ba5c66..4d999110 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -23,7 +23,7 @@ from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive from d123.dataset.dataset_specific.carla.opendrive.opendrive_converter import OpenDriveConverter -from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import AbstractMap, MapSurfaceType from d123.dataset.maps.abstract_map_objects import AbstractLane @@ -77,11 +77,9 @@ def __init__( self, splits: List[str], log_path: Union[Path, str], - output_path: Union[Path, str], - force_log_conversion: bool, - force_map_conversion: bool, + data_converter_config: DataConverterConfig, ) -> None: - super().__init__(force_log_conversion, force_map_conversion) + super().__init__(data_converter_config) for split in splits: assert ( split in self.get_available_splits() @@ -89,7 +87,6 @@ def __init__( self._splits: str = splits self._log_path: Path = Path(log_path) - self._output_path: Path = Path(output_path) self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() def _collect_log_paths(self) -> Dict[str, List[Path]]: diff --git a/d123/dataset/dataset_specific/nuplan/load_sensor.py b/d123/dataset/dataset_specific/nuplan/load_sensor.py index 878be2ff..dc699b2c 100644 --- a/d123/dataset/dataset_specific/nuplan/load_sensor.py +++ b/d123/dataset/dataset_specific/nuplan/load_sensor.py @@ -1,21 +1,21 @@ import io -import os from pathlib import Path from nuplan.database.utils.pointclouds.lidar import LidarPointCloud from d123.common.datatypes.sensor.lidar import LiDAR -NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) - -def load_lidar_from_path(filename: str) -> LiDAR: - lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / filename - - assert lidar_full_path.exists(), f"LiDAR file not found: {lidar_full_path}" - - with open(lidar_full_path, "rb") as fp: +def load_nuplan_lidar_from_path(filepath: Path) -> LiDAR: + assert filepath.exists(), f"LiDAR file not found: {filepath}" + with open(filepath, "rb") as fp: buffer = io.BytesIO(fp.read()) + return LiDAR(point_cloud=LidarPointCloud.from_buffer(buffer, "pcd").points) + - points = LidarPointCloud.from_buffer(buffer, "pcd").points - return LiDAR(point_cloud=points) # Reshape to +# def load_camera_from_path(filename: str, metadata: CameraMetadata) -> Camera: +# camera_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / filename +# assert camera_full_path.exists(), f"Camera file not found: {camera_full_path}" +# img = Image.open(camera_full_path) +# img.load() +# return Camera(metadata=metadata, image=np.asarray(img, dtype=np.uint8)) diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index bbb17424..fd05de11 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -1,23 +1,27 @@ import gc import json import os +import pickle from dataclasses import asdict from functools import partial from pathlib import Path from typing import Any, Dict, Final, List, Optional, Tuple, Union +import numpy as np import pyarrow as pa import yaml -from nuplan.database.nuplan_db.nuplan_scenario_queries import get_images_from_lidar_tokens +from nuplan.database.nuplan_db.nuplan_scenario_queries import get_cameras, get_images_from_lidar_tokens from nuplan.database.nuplan_db_orm.lidar_box import LidarBox from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB from nuplan.planning.simulation.observation.observation_type import CameraChannel from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map +from pyquaternion import Quaternion import d123.dataset.dataset_specific.nuplan.utils as nuplan_utils from d123.common.datatypes.detection.detection import TrafficLightStatus from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( get_nuplan_pacifica_parameters, @@ -29,7 +33,7 @@ from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import MAP_LOCATIONS, NuPlanMapConverter -from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata TARGET_DT: Final[float] = 0.1 @@ -51,6 +55,17 @@ "generic_object": DetectionType.GENERIC_OBJECT, } +NUPLAN_CAMERA_TYPES = { + CameraType.CAM_F0: CameraChannel.CAM_F0, + CameraType.CAM_B0: CameraChannel.CAM_B0, + CameraType.CAM_L0: CameraChannel.CAM_L0, + CameraType.CAM_L1: CameraChannel.CAM_L1, + CameraType.CAM_L2: CameraChannel.CAM_L2, + CameraType.CAM_R0: CameraChannel.CAM_R0, + CameraType.CAM_R1: CameraChannel.CAM_R1, + CameraType.CAM_R2: CameraChannel.CAM_R2, +} + NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) @@ -67,12 +82,9 @@ def __init__( self, splits: List[str], log_path: Union[Path, str], - output_path: Union[Path, str], - sensor_path: Union[Path, str], - force_log_conversion: bool, - force_map_conversion: bool, + data_converter_config: DataConverterConfig, ) -> None: - super().__init__(force_log_conversion, force_map_conversion) + super().__init__(data_converter_config) for split in splits: assert ( split in self.get_available_splits() @@ -80,8 +92,6 @@ def __init__( self._splits: List[str] = splits self._log_path: Path = Path(log_path) - self._output_path: Path = Path(output_path) - self._sensor_path: Path = Path(sensor_path) self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() self._target_dt: float = 0.1 @@ -126,11 +136,7 @@ def get_available_splits(self) -> List[str]: def convert_maps(self, worker: WorkerPool) -> None: worker_map( worker, - partial( - convert_nuplan_map_to_gpkg, - output_path=self._output_path, - force_map_conversion=self.force_map_conversion, - ), + partial(convert_nuplan_map_to_gpkg, data_converter_config=self.data_converter_config), list(MAP_LOCATIONS), ) @@ -148,26 +154,23 @@ def convert_logs(self, worker: WorkerPool) -> None: worker, partial( convert_nuplan_log_to_arrow, - output_path=self._output_path, - force_log_conversion=self.force_log_conversion, + data_converter_config=self.data_converter_config, ), log_args, ) -def convert_nuplan_map_to_gpkg(map_names: List[str], output_path: Path, force_map_conversion: bool) -> List[Any]: +def convert_nuplan_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: for map_name in map_names: - map_path = output_path / "maps" / f"nuplan_{map_name}.gpkg" - if force_map_conversion or not map_path.exists(): + map_path = data_converter_config.output_path / "maps" / f"nuplan_{map_name}.gpkg" + if data_converter_config.force_map_conversion or not map_path.exists(): map_path.unlink(missing_ok=True) - NuPlanMapConverter(output_path / "maps").convert(map_name=map_name) + NuPlanMapConverter(data_converter_config.output_path / "maps").convert(map_name=map_name) return [] def convert_nuplan_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], - output_path: Path, - force_log_conversion: bool, + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig ) -> List[Any]: for log_info in args: log_path: Path = log_info["log_path"] @@ -177,31 +180,41 @@ def convert_nuplan_log_to_arrow( raise FileNotFoundError(f"Log path {log_path} does not exist.") log_db = NuPlanDB(NUPLAN_DATA_ROOT, str(log_path), None) - log_file_path = output_path / split / f"{log_path.stem}.arrow" + log_file_path = data_converter_config.output_path / split / f"{log_path.stem}.arrow" - if force_log_conversion or not log_file_path.exists(): + if data_converter_config.force_log_conversion or not log_file_path.exists(): log_file_path.unlink(missing_ok=True) if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) - recording_schema = pa.schema( - [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ("lidar", pa.string()), - # ("front_cam_demo", pa.binary()), - # ("front_cam_transform", pa.list_(pa.float64())), - ] - ) + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + if data_converter_config.lidar_store_option == "path": + schema_column_list.append(("lidar", pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + # TODO: Adjust how cameras are added + if data_converter_config.camera_store_option is not None: + for camera_type in NUPLAN_CAMERA_TYPES.keys(): + if data_converter_config.camera_store_option == "path": + schema_column_list.append((camera_type.serialize(), pa.string())) + elif data_converter_config.camera_store_option == "binary": + raise NotImplementedError("Binary camera storage is not implemented.") + + recording_schema = pa.schema(schema_column_list) metadata = LogMetadata( dataset="nuplan", log_name=log_db.log_name, @@ -210,14 +223,16 @@ def convert_nuplan_log_to_arrow( map_has_z=False, ) vehicle_parameters = get_nuplan_pacifica_parameters() + camera_metadata = get_nuplan_camera_metadata(log_path) recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), } ) - _write_recording_table(log_db, recording_schema, log_file_path, log_path) + _write_recording_table(log_db, recording_schema, log_file_path, log_path, data_converter_config) log_db.detach_tables() log_db.remove_ref() @@ -226,11 +241,38 @@ def convert_nuplan_log_to_arrow( return [] +def get_nuplan_camera_metadata(log_path: Path) -> Dict[str, CameraMetadata]: + + def _get_camera_metadata(camera_type: CameraType) -> CameraMetadata: + cam = list(get_cameras(log_path, [str(NUPLAN_CAMERA_TYPES[camera_type].value)]))[0] + intrinsic = np.array(pickle.loads(cam.intrinsic)) + translation = np.array(pickle.loads(cam.translation)) + rotation = np.array(pickle.loads(cam.rotation)) + rotation = Quaternion(rotation).rotation_matrix + distortion = np.array(pickle.loads(cam.distortion)) + return CameraMetadata( + camera_type=camera_type, + width=cam.width, + height=cam.height, + intrinsic=intrinsic, + distortion=distortion, + translation=translation, + rotation=rotation, + ) + + log_cam_infos: Dict[str, CameraMetadata] = {} + for camera_type in NUPLAN_CAMERA_TYPES.keys(): + log_cam_infos[camera_type.serialize()] = _get_camera_metadata(camera_type) + + return log_cam_infos + + def _write_recording_table( log_db: NuPlanDB, recording_schema: pa.schema, log_file_path: Path, source_log_path: Path, + data_converter_config: DataConverterConfig, ) -> None: # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: @@ -239,16 +281,19 @@ def _write_recording_table( step_interval: float = int(TARGET_DT / NUPLAN_DT) for lidar_pc in log_db.lidar_pc[::step_interval]: lidar_pc_token: str = lidar_pc.token - detections_state, detections_velocity, detections_token, detections_types = _extract_detections( - lidar_pc - ) + ( + detections_state, + detections_velocity, + detections_token, + detections_types, + ) = _extract_detections(lidar_pc) traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) route_lane_group_ids = [ int(roadblock_id) for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") if len(roadblock_id) > 0 ] - # front_cam_demo, front_cam_transform = _extract_front_cam_demo(lidar_pc, source_log_path) + row_data = { "token": [lidar_pc_token], "timestamp": [lidar_pc.timestamp], @@ -261,10 +306,19 @@ def _write_recording_table( "traffic_light_types": [traffic_light_types], "scenario_tag": [_extract_scenario_tag(log_db, lidar_pc_token)], "route_lane_group_ids": [route_lane_group_ids], - "lidar": [_extract_lidar(lidar_pc)], - # "front_cam_demo": [front_cam_demo], - # "front_cam_transform": [front_cam_transform], } + + if data_converter_config.lidar_store_option is not None: + row_data["lidar"] = [_extract_lidar(lidar_pc, data_converter_config)] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_camera(lidar_pc, source_log_path, data_converter_config) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data] + else: + row_data[camera_type.serialize()] = [None] + batch = pa.record_batch(row_data, schema=recording_schema) writer.write_batch(batch) del batch, row_data, detections_state, detections_velocity, detections_token, detections_types @@ -363,28 +417,33 @@ def _extract_scenario_tag(log_db: NuPlanDB, lidar_pc_token: str) -> List[str]: return scenario_tags -def _extract_front_cam_demo(lidar_pc: LidarPc, source_log_path: Path) -> Tuple[bytes, List[float]]: - - sensor_root = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor" - front_image_class = list( - get_images_from_lidar_tokens(source_log_path, [lidar_pc.token], [str(CameraChannel.CAM_F0.value)]) - ) - - front_cam_demo: bytes = None - front_cam_transform: List[float] = [] +def _extract_camera( + lidar_pc: LidarPc, + source_log_path: Path, + data_converter_config: DataConverterConfig, +) -> Dict[CameraType, Union[str, bytes]]: - if len(front_image_class) != 0: + camera_dict: Dict[str, Union[str, bytes]] = {} + sensor_root = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" - filename_jpg = sensor_root / front_image_class[0].filename_jpg + for camera_type, camera_channel in NUPLAN_CAMERA_TYPES.items(): + camera_data: Optional[Union[str, bytes]] = None + image_class = list(get_images_from_lidar_tokens(source_log_path, [lidar_pc.token], [str(camera_channel.value)])) + if len(image_class) != 0: + filename_jpg = sensor_root / image_class[0].filename_jpg + if filename_jpg.exists(): + if data_converter_config.camera_store_option == "path": + camera_data = str(filename_jpg) + elif data_converter_config.camera_store_option == "binary": + with open(filename_jpg, "rb") as f: + camera_data = f.read() - if filename_jpg.exists(): - with open(filename_jpg, "rb") as f: - front_cam_demo = f.read() + camera_dict[camera_type] = camera_data - return front_cam_demo, front_cam_transform + return camera_dict -def _extract_lidar(lidar_pc: LidarPc) -> Optional[str]: +def _extract_lidar(lidar_pc: LidarPc, data_converter_config: DataConverterConfig) -> Optional[str]: lidar: Optional[str] = None lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / lidar_pc.filename diff --git a/d123/dataset/dataset_specific/raw_data_converter.py b/d123/dataset/dataset_specific/raw_data_converter.py index d9cf8f17..3a7f826c 100644 --- a/d123/dataset/dataset_specific/raw_data_converter.py +++ b/d123/dataset/dataset_specific/raw_data_converter.py @@ -1,14 +1,29 @@ import abc -from typing import List +from dataclasses import dataclass +from pathlib import Path +from typing import List, Literal, Optional, Union from nuplan.planning.utils.multithreading.worker_utils import WorkerPool +@dataclass +class DataConverterConfig: + + output_path: Union[str, Path] + force_log_conversion: bool = False + force_map_conversion: bool = False + camera_store_option: Optional[Literal["path", "binary"]] = None + lidar_store_option: Optional[Literal["path", "binary"]] = None + + def __post_init__(self): + if isinstance(self.output_path, str): + self.output_path = Path(self.output_path) + + class RawDataConverter(abc.ABC): - def __init__(self, force_log_conversion: bool, force_map_conversion: bool) -> None: - self.force_log_conversion = force_log_conversion - self.force_map_conversion = force_map_conversion + def __init__(self, data_converter_config: DataConverterConfig) -> None: + self.data_converter_config = data_converter_config @abc.abstractmethod def get_available_splits(self) -> List[str]: diff --git a/d123/dataset/scene/abstract_scene.py b/d123/dataset/scene/abstract_scene.py index 2d55a95f..110f2d1c 100644 --- a/d123/dataset/scene/abstract_scene.py +++ b/d123/dataset/scene/abstract_scene.py @@ -4,10 +4,9 @@ from dataclasses import dataclass from typing import List -from PIL import Image - from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording +from d123.common.datatypes.sensor.camera import Camera, CameraType from d123.common.datatypes.sensor.lidar import LiDAR from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 @@ -71,7 +70,7 @@ def get_route_lane_group_ids(self, iteration: int) -> List[int]: raise NotImplementedError @abc.abstractmethod - def get_front_cam_demo(self, iteration: int) -> Image: + def get_camera_at_iteration(self, iteration: int, camera_type: CameraType) -> Camera: raise NotImplementedError @abc.abstractmethod diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index 2718078a..835c0150 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -1,21 +1,19 @@ -import io import json from pathlib import Path -from typing import List, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import pyarrow as pa -# TODO: Remove or improve open/close dynamic of Scene object. -from PIL import Image - from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording +from d123.common.datatypes.sensor.camera import Camera, CameraMetadata, CameraType, camera_metadata_dict_from_json from d123.common.datatypes.sensor.lidar import LiDAR from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.dataset.arrow.conversion import ( get_box_detections_from_arrow_table, + get_camera_from_arrow_table, get_ego_vehicle_state_from_arrow_table, get_lidar_from_arrow_table, get_timepoint_from_arrow_table, @@ -27,8 +25,12 @@ from d123.dataset.maps.gpkg.gpkg_map import get_map_api_from_names from d123.dataset.scene.abstract_scene import AbstractScene, SceneExtractionInfo +# TODO: Remove or improve open/close dynamic of Scene object. + -def _get_scene_data(arrow_file_path: Union[Path, str]) -> Tuple[LogMetadata, VehicleParameters]: +def _get_scene_data( + arrow_file_path: Union[Path, str], +) -> Tuple[LogMetadata, VehicleParameters, Dict[CameraType, CameraMetadata]]: """ Extracts the metadata and vehicle parameters from the arrow file. """ @@ -36,8 +38,14 @@ def _get_scene_data(arrow_file_path: Union[Path, str]) -> Tuple[LogMetadata, Veh table = open_arrow_table(arrow_file_path) metadata = LogMetadata(**json.loads(table.schema.metadata[b"log_metadata"].decode())) vehicle_parameters = VehicleParameters(**json.loads(table.schema.metadata[b"vehicle_parameters"].decode())) + + if b"camera_metadata" in table.schema.metadata: + camera_metadata = camera_metadata_dict_from_json(table.schema.metadata[b"camera_metadata"].decode()) + else: + camera_metadata = {} + del table - return metadata, vehicle_parameters + return metadata, vehicle_parameters, camera_metadata class ArrowScene(AbstractScene): @@ -49,9 +57,10 @@ def __init__( self._recording_table: pa.Table = None - _metadata, _vehicle_parameters = _get_scene_data(arrow_file_path) + _metadata, _vehicle_parameters, _camera_metadata = _get_scene_data(arrow_file_path) self._metadata: LogMetadata = _metadata self._vehicle_parameters: VehicleParameters = _vehicle_parameters + self._camera_metadata: Dict[CameraType, CameraMetadata] = _camera_metadata self._map_api: Optional[AbstractMap] = None @@ -130,15 +139,20 @@ def get_route_lane_group_ids(self, iteration: int) -> List[int]: table_index = self._get_table_index(iteration) return self._recording_table["route_lane_group_ids"][table_index].as_py() - def get_front_cam_demo(self, iteration: int) -> Image: + def get_camera_at_iteration(self, iteration: int, camera_type: CameraType) -> Camera: self._lazy_initialize() + assert camera_type in self._camera_metadata, f"Camera type {camera_type} not found in metadata." table_index = self._get_table_index(iteration) - jpg_data = self._recording_table["front_cam_demo"][table_index].as_py() - return Image.open(io.BytesIO(jpg_data)) + return get_camera_from_arrow_table( + self._recording_table, + table_index, + self._camera_metadata[camera_type], + self.log_metadata, + ) def get_lidar_at_iteration(self, iteration: int) -> LiDAR: self._lazy_initialize() - return get_lidar_from_arrow_table(self._recording_table, self._get_table_index(iteration)) + return get_lidar_from_arrow_table(self._recording_table, self._get_table_index(iteration), self.log_metadata) def _lazy_initialize(self) -> None: self.open() diff --git a/d123/script/config/datasets/nuplan_dataset.yaml b/d123/script/config/datasets/nuplan_dataset.yaml index c20793af..8c104287 100644 --- a/d123/script/config/datasets/nuplan_dataset.yaml +++ b/d123/script/config/datasets/nuplan_dataset.yaml @@ -4,7 +4,13 @@ nuplan_dataset: splits: ["nuplan_train", "nuplan_val", "nuplan_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder - output_path: ${d123_data_root} - force_log_conversion: ${force_log_conversion} - force_map_conversion: ${force_map_conversion} + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/d123/script/config/datasets/nuplan_mini_dataset.yaml b/d123/script/config/datasets/nuplan_mini_dataset.yaml index 68bf33a0..1fdb2b54 100644 --- a/d123/script/config/datasets/nuplan_mini_dataset.yaml +++ b/d123/script/config/datasets/nuplan_mini_dataset.yaml @@ -4,9 +4,14 @@ nuplan_mini_dataset: splits: ["nuplan_mini_train", "nuplan_mini_val", "nuplan_mini_test"] - # splits: ["nuplan_mini_val"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder - output_path: ${d123_data_root} - force_log_conversion: ${force_log_conversion} - force_map_conversion: ${force_map_conversion} + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/d123/script/config/datasets/nuplan_private_dataset.yaml b/d123/script/config/datasets/nuplan_private_dataset.yaml index e69147cc..399dcb7e 100644 --- a/d123/script/config/datasets/nuplan_private_dataset.yaml +++ b/d123/script/config/datasets/nuplan_private_dataset.yaml @@ -4,8 +4,13 @@ nuplan_private_dataset: splits: ["nuplan_private_test"] log_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/splits # NOTE: folder including [mini, trainval, test], sometimes not inside "splits" folder - sensor_path: ${oc.env:NUPLAN_DATA_ROOT}/nuplan-v1.1/sensor - output_path: ${d123_data_root} - force_log_conversion: ${force_log_conversion} - force_map_conversion: ${force_map_conversion} + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/docs/camera.md b/docs/camera.md new file mode 100644 index 00000000..9a016aa7 --- /dev/null +++ b/docs/camera.md @@ -0,0 +1,201 @@ + +# nuPlan + +``` +0: { + "camera_name": "CAM_F0", + "original_size": (1080, 1920), + "egocar_visible": False +}, +1: { + "camera_name": "CAM_L0", + "original_size": (1080, 1920), + "egocar_visible": True +}, +2: { + "camera_name": "CAM_R0", + "original_size": (1080, 1920), + "egocar_visible": True +}, +3: { + "camera_name": "CAM_L1", + "original_size": (1080, 1920), + "egocar_visible": False +}, +4: { + "camera_name": "CAM_R1", + "original_size": (1080, 1920), + "egocar_visible": False +}, +5: { + "camera_name": "CAM_L2", + "original_size": (1080, 1920), + "egocar_visible": True +}, +6: { + "camera_name": "CAM_R2", + "original_size": (1080, 1920), + "egocar_visible": True +}, +7: { + "camera_name": "CAM_B0", + "original_size": (1080, 1920), + "egocar_visible": False +}, +``` + + +# nuScenes +``` +0: { + "camera_name": "CAM_FRONT", + "original_size": (900, 1600), + "egocar_visible": False +}, +1: { + "camera_name": "CAM_FRONT_LEFT", + "original_size": (900, 1600), + "egocar_visible": False +}, +2: { + "camera_name": "CAM_FRONT_RIGHT", + "original_size": (900, 1600), + "egocar_visible": False +}, +3: { + "camera_name": "CAM_BACK_LEFT", + "original_size": (900, 1600), + "egocar_visible": False +}, +4: { + "camera_name": "CAM_BACK_RIGHT", + "original_size": (900, 1600), + "egocar_visible": False +}, +5: { + "camera_name": "CAM_BACK", + "original_size": (900, 1600), + "egocar_visible": True +} +``` + + +# Waymo +``` +0: { + "camera_name": "front_camera", + "original_size": (1280, 1920), + "egocar_visible": False +}, +1: { + "camera_name": "front_left_camera", + "original_size": (1280, 1920), + "egocar_visible": False +}, +2: { + "camera_name": "front_right_camera", + "original_size": (1280, 1920), + "egocar_visible": False +}, +3: { + "camera_name": "left_camera", + "original_size": (866, 1920), + "egocar_visible": False +}, +4: { + "camera_name": "right_camera", + "original_size": (866, 1920), + "egocar_visible": False +}, +``` + +# Panda set + +``` +0: { + "camera_name": "front_camera", + "original_size": (1080, 1920), + "egocar_visible": False +}, +1: { + "camera_name": "front_left_camera", + "original_size": (1080, 1920), + "egocar_visible": False +}, +2: { + "camera_name": "front_right_camera", + "original_size": (1080, 1920), + "egocar_visible": False +}, +3: { + "camera_name": "left_camera", + "original_size": (1080, 1920), + "egocar_visible": False +}, +4: { + "camera_name": "right_camera", + "original_size": (1080, 1920), + "egocar_visible": False +}, +5: { + "camera_name": "back_camera", + "original_size": (1080, 1920), + "egocar_visible": True +}, +``` + + +# Argoverse + +``` +0: { + "camera_name": "ring_front_center", + "original_size": (2048, 1550), + "egocar_visible": True +}, +1: { + "camera_name": "ring_front_left", + "original_size": (1550, 2048), + "egocar_visible": False +}, +2: { + "camera_name": "ring_front_right", + "original_size": (1550, 2048), + "egocar_visible": False +}, +3: { + "camera_name": "ring_side_left", + "original_size": (1550, 2048), + "egocar_visible": False +}, +4: { + "camera_name": "ring_side_right", + "original_size": (1550, 2048), + "egocar_visible": False +}, +5: { + "camera_name": "ring_rear_left", + "original_size": (1550, 2048), + "egocar_visible": True +}, +6: { + "camera_name": "ring_rear_right", + "original_size": (1550, 2048), + "egocar_visible": True +}, +``` + + +# KITTI +``` +0: { + "camera_name": "CAM_LEFT", + "original_size": (375, 1242), + "egocar_visible": False +}, +1: { + "camera_name": "CAM_RIGHT", + "original_size": (375, 1242), + "egocar_visible": False +} +``` diff --git a/notebooks/nuplan_sensor_loading.ipynb b/notebooks/nuplan_sensor_loading.ipynb index 10182cec..0dd69b4e 100644 --- a/notebooks/nuplan_sensor_loading.ipynb +++ b/notebooks/nuplan_sensor_loading.ipynb @@ -23,7 +23,9 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.dataset.dataset_specific.nuplan.nuplan_data_converter import NuplanDataConverter" + "import numpy as np\n", + "\n", + "np.zeros((3,3)).tolist()" ] }, { @@ -32,6 +34,16 @@ "id": "2", "metadata": {}, "outputs": [], + "source": [ + "from d123.dataset.dataset_specific.nuplan.nuplan_data_converter import NuplanDataConverter" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], "source": [ "def process_images_to_arrow(input_folder, output_file, batch_size=100):\n", " \"\"\"\n", @@ -200,7 +212,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -210,7 +222,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -249,7 +261,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -269,7 +281,15 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", "metadata": {}, "outputs": [], "source": [ @@ -294,7 +314,6 @@ " intrinsics = np.array(pickle.loads(cam.intrinsic))\n", " translation = np.array(pickle.loads(cam.translation))\n", " rotation = np.array(pickle.loads(cam.rotation))\n", - " print(rotation)\n", " rotation = Quaternion(rotation).rotation_matrix\n", " distortion = np.array(pickle.loads(cam.distortion))\n", " c = dict(\n", @@ -353,7 +372,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "9", "metadata": {}, "outputs": [], "source": [ @@ -391,7 +410,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -406,7 +425,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -416,7 +435,7 @@ { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -452,7 +471,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "13", "metadata": {}, "outputs": [], "source": [] diff --git a/notebooks/scene_sensor_loading.ipynb b/notebooks/scene_sensor_loading.ipynb index 4b25abab..e10a0842 100644 --- a/notebooks/scene_sensor_loading.ipynb +++ b/notebooks/scene_sensor_loading.ipynb @@ -7,6 +7,7 @@ "metadata": {}, "outputs": [], "source": [ + "import matplotlib.pyplot as plt\n", "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", @@ -62,17 +63,15 @@ "metadata": {}, "outputs": [], "source": [ - "\n", - "\n", + "from d123.common.datatypes.sensor.camera import CameraType\n", "from d123.dataset.scene.arrow_scene import ArrowScene\n", "\n", - "\n", "scene: ArrowScene = scenes[12]\n", - "\n", "scene.open()\n", "\n", - "lidar = scene.get_lidar_at_iteration(0)\n", - "\n" + "camera = scene.get_camera_at_iteration(150, CameraType.CAM_B0)\n", + "\n", + "plt.imshow(camera.image)" ] }, { @@ -81,11 +80,7 @@ "id": "3", "metadata": {}, "outputs": [], - "source": [ - "\n", - "pc = lidar.point_cloud\n", - "pc.shape" - ] + "source": [] }, { "cell_type": "code", @@ -94,7 +89,7 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", + "\n", "\n", "plt.scatter(pc[0, :], pc[1, :], s=0.1, alpha=0.5)" ] From c8e8fa4dec9ab971812cbd68d049ea84e2ae6f0c Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Mon, 28 Jul 2025 13:44:14 +0200 Subject: [PATCH 05/42] Add lidar and back camera to viser visualization --- d123/common/datatypes/sensor/lidar.py | 7 +++ d123/common/geometry/transform/se3.py | 11 +++++ d123/common/visualization/viser/server.py | 50 ++++++++++++++++++---- d123/common/visualization/viser/utils.py | 13 ++++++ d123/dataset/scene/abstract_scene.py | 5 +++ d123/dataset/scene/arrow_scene.py | 4 ++ notebooks/viz/viser_testing_v2_scene.ipynb | 10 ++--- 7 files changed, 85 insertions(+), 15 deletions(-) diff --git a/d123/common/datatypes/sensor/lidar.py b/d123/common/datatypes/sensor/lidar.py index c455279c..7b13f261 100644 --- a/d123/common/datatypes/sensor/lidar.py +++ b/d123/common/datatypes/sensor/lidar.py @@ -8,3 +8,10 @@ class LiDAR: point_cloud: npt.NDArray[np.float32] + + @property + def xyz(self) -> npt.NDArray[np.float32]: + """ + Returns the point cloud as an Nx3 array of x, y, z coordinates. + """ + return self.point_cloud[:3].T diff --git a/d123/common/geometry/transform/se3.py b/d123/common/geometry/transform/se3.py index 5bcf75b5..ccd861c9 100644 --- a/d123/common/geometry/transform/se3.py +++ b/d123/common/geometry/transform/se3.py @@ -85,3 +85,14 @@ def translate_body_frame(state_se3: StateSE3, vector_3d: Vector3D) -> StateSE3: state_se3.pitch, state_se3.yaw, ) + + +def convert_relative_to_absolute_points_3d_array( + origin: StateSE3, points_3d_array: npt.NDArray[np.float64] +) -> npt.NDArray[np.float64]: + + # TODO: implement function for origin as np.ndarray + + R = get_rotation_matrix(origin) + absolute_points = points_3d_array @ R.T + origin.point_3d.array + return absolute_points diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 56a0ecfb..d1d24d2c 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -1,11 +1,12 @@ import time -import numpy as np import trimesh import viser +from d123.common.datatypes.sensor.camera import CameraType from d123.common.visualization.viser.utils import ( get_bounding_box_meshes, + get_lidar_points, get_map_meshes, ) from d123.dataset.scene.abstract_scene import AbstractScene @@ -13,6 +14,12 @@ # TODO: Try to fix performance issues. # TODO: Refactor this file. +FRONT_CAMERA_TYPE: CameraType = CameraType.CAM_F0 +BACK_CAMERA_TYPE: CameraType = CameraType.CAM_B0 +LIDAR_AVAILABLE: bool = True +LIDAR_POINT_SIZE: float = 0.05 +MAP_AVAILABLE: bool = False + class ViserVisualizationServer: def __init__( @@ -84,8 +91,14 @@ def _(_) -> None: trimesh.util.concatenate(meshes), visible=True, ) - gui_image_handle.image = np.array(scene.get_front_cam_demo(gui_timestep.value)) + if FRONT_CAMERA_TYPE in scene.available_camera_types: + gui_front_cam.image = scene.get_camera_at_iteration(gui_timestep.value, FRONT_CAMERA_TYPE).image + + if BACK_CAMERA_TYPE in scene.available_camera_types: + gui_back_cam.image = scene.get_camera_at_iteration(gui_timestep.value, BACK_CAMERA_TYPE).image + if LIDAR_AVAILABLE: + gui_lidar.points = get_lidar_points(scene, gui_timestep.value) # camera_pose = _get_camera_pose_demo(scene, gui_timestep.value) # frustum_handle.position = camera_pose.point_3d.array # frustum_handle.wxyz = euler_to_quaternion_scipy(camera_pose.roll, camera_pose.pitch, camera_pose.yaw) @@ -108,15 +121,34 @@ def _(_) -> None: current_frame_handle = self.server.scene.add_frame(f"/frame{gui_timestep.value}", show_axes=False) self.server.scene.add_frame("/map", show_axes=False) - with self.server.gui.add_folder("Camera"): - gui_image_handle = self.server.gui.add_image( - image=np.array(scene.get_front_cam_demo(gui_timestep.value)), - label="front_cam_demo", - format="jpeg", + if FRONT_CAMERA_TYPE in scene.available_camera_types: + with self.server.gui.add_folder("Camera"): + gui_front_cam = self.server.gui.add_image( + image=scene.get_camera_at_iteration(gui_timestep.value, FRONT_CAMERA_TYPE).image, + label=FRONT_CAMERA_TYPE.serialize(), + format="jpeg", + ) + + if BACK_CAMERA_TYPE in scene.available_camera_types: + with self.server.gui.add_folder("Camera"): + gui_back_cam = self.server.gui.add_image( + image=scene.get_camera_at_iteration(gui_timestep.value, BACK_CAMERA_TYPE).image, + label=BACK_CAMERA_TYPE.serialize(), + format="jpeg", + ) + + if LIDAR_AVAILABLE: + gui_lidar = self.server.scene.add_point_cloud( + name="LiDAR", + points=get_lidar_points(scene, gui_timestep.value), + colors=(0.0, 0.0, 0.0), + point_size=LIDAR_POINT_SIZE, + point_shape="circle", ) - for name, mesh in get_map_meshes(scene).items(): - self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) + if MAP_AVAILABLE: + for name, mesh in get_map_meshes(scene).items(): + self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) # camera_pose = _get_camera_pose_demo(scene, gui_timestep.value) # frustum_handle = self.server.scene.add_camera_frustum( diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 469b49d5..f95c54fe 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -6,6 +6,7 @@ from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_points_3d_array from d123.common.visualization.color.config import PlotConfig from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG, MAP_SURFACE_CONFIG from d123.dataset.maps.abstract_map import MapSurfaceType @@ -218,3 +219,15 @@ def euler_to_quaternion_scipy(roll: float, pitch: float, yaw: float) -> npt.NDAr r = Rotation.from_euler("xyz", [roll, pitch, yaw], degrees=False) quat = r.as_quat(scalar_first=True) return quat + + +def get_lidar_points(scene: AbstractScene, iteration: int) -> npt.NDArray[np.float32]: + + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) + current_ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) + + lidar = scene.get_lidar_at_iteration(iteration) + points = convert_relative_to_absolute_points_3d_array(current_ego_vehicle_state.rear_axle_se3, lidar.xyz) + points = points - initial_ego_vehicle_state.center_se3.point_3d.array + + return points diff --git a/d123/dataset/scene/abstract_scene.py b/d123/dataset/scene/abstract_scene.py index 110f2d1c..84d26dd8 100644 --- a/d123/dataset/scene/abstract_scene.py +++ b/d123/dataset/scene/abstract_scene.py @@ -37,6 +37,11 @@ def token(self) -> str: def log_metadata(self) -> LogMetadata: raise NotImplementedError + @property + @abc.abstractmethod + def available_camera_types(self) -> List[CameraType]: + raise NotImplementedError + @abc.abstractmethod def get_number_of_iterations(self) -> int: raise NotImplementedError diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index 835c0150..434147be 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -94,6 +94,10 @@ def token(self) -> str: def log_metadata(self) -> LogMetadata: return self._metadata + @property + def available_camera_types(self) -> List[CameraType]: + return list(self._camera_metadata.keys()) + def _get_table_index(self, iteration: int) -> int: self._lazy_initialize() assert ( diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 1bacd558..6fdb67fc 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -24,12 +24,10 @@ "metadata": {}, "outputs": [], "source": [ - "split = \"carla\"\n", + "split = \"nuplan_private_test\"\n", "\n", - "log_names = [\"_Rep0_routes_devtest_route1_07_10_18_45_18\"]\n", + "log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "scene_tokens = None\n", - "# scene_tokens = [\"2283aea39bc1505e\"]\n", - "# log_names = None\n", "\n", "scene_filter = SceneFilter(\n", " split_names=[split],\n", @@ -37,7 +35,7 @@ " scene_tokens=scene_tokens,\n", " duration_s=20.1,\n", " history_s=1.0,\n", - " timestamp_threshold_s=0.1,\n", + " timestamp_threshold_s=20.1,\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", @@ -56,7 +54,7 @@ "source": [ "\n", "visualization_server = ViserVisualizationServer()\n", - "visualization_server.set_scene(scenes[180])" + "visualization_server.set_scene(scenes[16])" ] }, { From 5b165f22e649cc4609849e8c7292680079d0be34 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Mon, 28 Jul 2025 14:41:23 +0200 Subject: [PATCH 06/42] Add some fixes and extension to load lidar and camera from paths for CARLA --- d123/common/datatypes/sensor/camera.py | 18 +- d123/dataset/arrow/conversion.py | 5 +- .../carla/carla_data_converter.py | 187 +++++++++++------- .../dataset_specific/carla/load_sensor.py | 11 ++ .../default_dataset_conversion.yaml | 4 +- .../script/config/datasets/carla_dataset.yaml | 12 +- notebooks/viz/viser_testing_v2_scene.ipynb | 20 +- 7 files changed, 167 insertions(+), 90 deletions(-) create mode 100644 d123/dataset/dataset_specific/carla/load_sensor.py diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py index edc8ad7c..cca14dd5 100644 --- a/d123/common/datatypes/sensor/camera.py +++ b/d123/common/datatypes/sensor/camera.py @@ -37,26 +37,28 @@ class CameraMetadata: rotation: npt.NDArray[np.float64] # 3x3 matrix def to_dict(self) -> Dict[str, Any]: + # TODO: remove None types. Only a placeholder for now. return { "camera_type": int(self.camera_type), "width": self.width, "height": self.height, - "intrinsic": self.intrinsic.tolist(), - "distortion": self.distortion.tolist(), - "translation": self.translation.tolist(), - "rotation": self.rotation.tolist(), + "intrinsic": self.intrinsic.tolist() if self.intrinsic is not None else None, + "distortion": self.distortion.tolist() if self.distortion is not None else None, + "translation": self.translation.tolist() if self.translation is not None else None, + "rotation": self.rotation.tolist() if self.rotation is not None else None, } @classmethod def from_dict(cls, json_dict: Dict[str, Any]) -> CameraMetadata: + # TODO: remove None types. Only a placeholder for now. return cls( camera_type=CameraType(json_dict["camera_type"]), width=json_dict["width"], height=json_dict["height"], - intrinsic=np.array(json_dict["intrinsic"]), - distortion=np.array(json_dict["distortion"]), - translation=np.array(json_dict["translation"]), - rotation=np.array(json_dict["rotation"]), + intrinsic=np.array(json_dict["intrinsic"]) if json_dict["intrinsic"] is not None else None, + distortion=np.array(json_dict["distortion"]) if json_dict["distortion"] is not None else None, + translation=np.array(json_dict["translation"]) if json_dict["translation"] is not None else None, + rotation=np.array(json_dict["rotation"]) if json_dict["rotation"] is not None else None, ) diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 9d8781da..50a6160a 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -27,13 +27,14 @@ from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.vector import Vector3D +from d123.dataset.dataset_specific.carla.load_sensor import load_carla_lidar_from_path from d123.dataset.dataset_specific.nuplan.load_sensor import load_nuplan_lidar_from_path from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import List DATASET_SENSOR_ROOT: Dict[str, Path] = { "nuplan": Path(os.environ["NUPLAN_DATA_ROOT"]) / "nuplan-v1.1" / "sensor_blobs", - # "carla": Path(os.environ["CARLA_DATA_ROOT"]) / "sensor_blobs", + "carla": Path(os.environ["CARLA_DATA_ROOT"]) / "sensor_blobs", } @@ -128,6 +129,8 @@ def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int, log_metadata: assert full_lidar_path.exists(), f"LiDAR file not found: {full_lidar_path}" if log_metadata.dataset == "nuplan": lidar = load_nuplan_lidar_from_path(full_lidar_path) + elif log_metadata.dataset == "carla": + lidar = load_carla_lidar_from_path(full_lidar_path) else: raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 4d999110..943ceda3 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -6,20 +6,19 @@ from dataclasses import asdict from functools import partial from pathlib import Path -from typing import Dict, Final, List, Tuple, Union +from typing import Any, Dict, Final, List, Optional, Tuple, Union import numpy as np import pyarrow as pa from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map -from traitlets import Any +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import get_carla_lincoln_mkz_2020_parameters from d123.common.geometry.base import Point2D, Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index from d123.common.geometry.transform.se3 import translate_se3_along_z from d123.common.geometry.vector import Vector3DIndex -from d123.dataset.arrow.conversion import VehicleParameters from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive from d123.dataset.dataset_specific.carla.opendrive.opendrive_converter import OpenDriveConverter @@ -48,6 +47,10 @@ TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE: Final[float] = 1.0 # [m] SORT_BY_TIMESTAMP: Final[bool] = True +CARLA_CAMERA_TYPES = {CameraType.CAM_F0} + +CARLA_DATA_ROOT: Final[Path] = Path(os.environ["CARLA_DATA_ROOT"]) + # TODO: Refactor this files and convert coordinate systems more elegantly. # NOTE: Currently some coordinate transforms from Unreal to ISO 8855 are done in the data agent of PDM-Lite. @@ -92,7 +95,9 @@ def __init__( def _collect_log_paths(self) -> Dict[str, List[Path]]: # TODO: fix "carla" split placeholder and add support for other splits log_paths_per_split: Dict[str, List[Path]] = {} - log_paths = list(self._log_path.iterdir()) + log_paths = [ + log_path for log_path in self._log_path.iterdir() if log_path.is_dir() and log_path.stem != "sensor_blobs" + ] log_paths_per_split["carla"] = log_paths return log_paths_per_split @@ -105,8 +110,7 @@ def convert_maps(self, worker: WorkerPool) -> None: worker, partial( convert_carla_map_to_gpkg, - output_path=self._output_path, - force_map_conversion=self.force_map_conversion, + data_converter_config=self.data_converter_config, ), list(AVAILABLE_CARLA_MAP_LOCATIONS), ) @@ -114,24 +118,20 @@ def convert_maps(self, worker: WorkerPool) -> None: def convert_logs(self, worker: WorkerPool) -> None: log_args = [ - { - "log_path": log_path, - "output_path": self._output_path, - "split": split, - } + {"log_path": log_path, "split": split} for split, log_paths in self._log_paths_per_split.items() for log_path in log_paths ] worker_map( - worker, partial(convert_carla_log_to_arrow, force_log_conversion=self.force_log_conversion), log_args + worker, partial(convert_carla_log_to_arrow, data_converter_config=self.data_converter_config), log_args ) -def convert_carla_map_to_gpkg(map_names: List[str], output_path: Path, force_map_conversion: bool) -> List[Any]: +def convert_carla_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: for map_name in map_names: - map_path = output_path / "maps" / f"carla_{map_name}.gpkg" - if force_map_conversion or not map_path.exists(): + map_path = data_converter_config.output_path / "maps" / f"carla_{map_name.lower()}.gpkg" + if data_converter_config.force_map_conversion or not map_path.exists(): map_path.unlink(missing_ok=True) assert os.environ["CARLA_ROOT"] is not None CARLA_ROOT = Path(os.environ["CARLA_ROOT"]) @@ -150,18 +150,17 @@ def convert_carla_map_to_gpkg(map_names: List[str], output_path: Path, force_map def convert_carla_log_to_arrow( - args: List[Dict[str, Union[List[str], List[Path]]]], - force_log_conversion: bool, + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig ) -> List[Any]: def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> None: for log_info in args: log_path: Path = log_info["log_path"] - output_path: Path = log_info["output_path"] split: str = log_info["split"] + output_path: Path = data_converter_config.output_path log_file_path = output_path / split / f"{log_path.stem}.arrow" - if force_log_conversion or not log_file_path.exists(): + if data_converter_config.force_log_conversion or not log_file_path.exists(): log_file_path.unlink(missing_ok=True) if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) @@ -170,33 +169,52 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> map_name = _load_json_gz(bounding_box_paths[0])["location"] map_api = get_map_api_from_names("carla", map_name) - recording_schema = pa.schema( - [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ("front_cam_demo", pa.binary()), - ("front_cam_transform", pa.list_(pa.float64())), - ] - ) + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + if data_converter_config.lidar_store_option == "path": + schema_column_list.append(("lidar", pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + # TODO: Adjust how cameras are added + if data_converter_config.camera_store_option is not None: + for camera_type in CARLA_CAMERA_TYPES: + if data_converter_config.camera_store_option == "path": + schema_column_list.append((camera_type.serialize(), pa.string())) + elif data_converter_config.camera_store_option == "binary": + raise NotImplementedError("Binary camera storage is not implemented.") + + recording_schema = pa.schema(schema_column_list) metadata = _get_metadata(map_name, str(log_path.stem)) - vehicle_parameters = _get_vehicle_parameters_() + vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() + camera_metadata = get_carla_camera_metadata() recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), } ) - _write_recording_table(bounding_box_paths, map_api, recording_schema, log_file_path) + _write_recording_table( + bounding_box_paths, + map_api, + recording_schema, + log_file_path, + data_converter_config, + ) gc.collect() @@ -215,18 +233,21 @@ def _get_metadata(location: str, log_name: str) -> LogMetadata: ) -def _get_vehicle_parameters(vehicle_parameter_dict: Dict[str, float]) -> VehicleParameters: - # NOTE: @DanielDauner extracting the vehicle parameters from CARLA is somewhat tricky. - # Need to extract the coordinates (for wheels, wheelbase, etc.) from CARLA which is somewhat noise. - # Thus, we hardcode the parameters for the Lincoln MKZ 2020 (default). - assert ( - vehicle_parameter_dict["vehicle_name"] == "vehicle.lincoln.mkz_2020" - ), "Currently only supports MKZ 2020 in CARLA." - return get_carla_lincoln_mkz_2020_parameters() - +def get_carla_camera_metadata() -> Dict[str, CameraMetadata]: -def _get_vehicle_parameters_() -> VehicleParameters: - return get_carla_lincoln_mkz_2020_parameters() + # FIXME: This is a placeholder function to return camera metadata. + camera_metadata = { + CameraType.CAM_F0.serialize(): CameraMetadata( + camera_type=CameraType.CAM_F0, + width=1024, + height=512, + intrinsic=None, + distortion=None, + translation=None, + rotation=None, + ) + } + return camera_metadata def _write_recording_table( @@ -234,21 +255,24 @@ def _write_recording_table( map_api: AbstractMap, recording_schema: pa.Schema, log_file_path: Path, + data_converter_config: DataConverterConfig, ) -> pa.Table: - log_path = bounding_box_paths[0].parent.parent.stem + # TODO: Refactor this function to be more readable + log_name = str(bounding_box_paths[0].parent.parent.stem) with pa.OSFile(str(log_file_path), "wb") as sink: with pa.ipc.new_file(sink, recording_schema) as writer: for box_path in bounding_box_paths: + sample_name = box_path.stem.split(".")[0] + data = _load_json_gz(box_path) traffic_light_ids, traffic_light_types = _extract_traffic_light_data( data["traffic_light_states"], data["traffic_light_positions"], map_api ) route_lane_group_ids = _extract_route_lane_group_ids(data["route"], map_api) if "route" in data else [] - front_cam_demo, front_cam_transform = _extract_front_cam_demo(box_path) row_data = { - "token": [create_token(f"{str(log_path)}_{box_path.stem}")], + "token": [create_token(f"{log_name}_{box_path.stem}")], "timestamp": [data["timestamp"]], "detections_state": [_extract_detection_states(data["detections_state"])], "detections_velocity": [ @@ -265,9 +289,18 @@ def _write_recording_table( "traffic_light_types": [traffic_light_types], "scenario_tag": [data["scenario_tag"]], "route_lane_group_ids": [route_lane_group_ids], - "front_cam_demo": [front_cam_demo], - "front_cam_transform": [front_cam_transform], } + if data_converter_config.lidar_store_option is not None: + row_data["lidar"] = [_extract_lidar(log_name, sample_name, data_converter_config)] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_cameras(log_name, sample_name, data_converter_config) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data] + else: + row_data[camera_type.serialize()] = [None] + batch = pa.record_batch(row_data, schema=recording_schema) writer.write_batch(batch) del batch, row_data @@ -334,6 +367,7 @@ def _extract_traffic_light_data( def _extract_route_lane_group_ids(route: List[List[float]], map_api: AbstractMap) -> List[int]: + # FIXME: Carla route is very buggy. No check if lanes are connected. route = np.array(route, dtype=np.float64) route[..., 1] = -route[..., 1] # Unreal coordinate system to ISO 8855 route = route[::2] @@ -365,18 +399,33 @@ def _extract_route_lane_group_ids(route: List[List[float]], map_api: AbstractMap return list(dict.fromkeys(route_lane_group_ids)) # Remove duplicates while preserving order -def _extract_front_cam_demo(box_path: Path) -> Tuple[bytes, List[float]]: - - sample_name = str(box_path.stem).split(".")[0] - sensor_root = Path(box_path.parent.parent) / "rgb" - - front_cam_demo: bytes = None - front_cam_transform: List[float] = [] - - jpg_path = sensor_root / f"{sample_name}.jpg" - - if jpg_path.exists(): - with open(jpg_path, "rb") as f: - front_cam_demo = f.read() - - return front_cam_demo, front_cam_transform +def _extract_cameras( + log_name: str, sample_name: str, data_converter_config: DataConverterConfig +) -> Dict[CameraType, Optional[str]]: + camera_dict: Dict[str, Union[str, bytes]] = {} + for camera_type in CARLA_CAMERA_TYPES: + camera_full_path = CARLA_DATA_ROOT / "sensor_blobs" / log_name / camera_type.name / f"{sample_name}.jpg" + if camera_full_path.exists(): + if data_converter_config.camera_store_option == "path": + camera_dict[camera_type] = f"{log_name}/{camera_type.name}/{sample_name}.jpg" + elif data_converter_config.camera_store_option == "binary": + raise NotImplementedError("Binary camera storage is not implemented.") + else: + camera_dict[camera_type] = None + print("camera_dict", camera_dict) + return camera_dict + + +def _extract_lidar(log_name: str, sample_name: str, data_converter_config: DataConverterConfig) -> Optional[str]: + + lidar: Optional[str] = None + lidar_full_path = CARLA_DATA_ROOT / "sensor_blobs" / log_name / "lidar" / f"{sample_name}.npy" + if lidar_full_path.exists(): + if data_converter_config.lidar_store_option == "path": + lidar = f"{log_name}/lidar/{sample_name}.npy" + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + else: + raise FileNotFoundError(f"LiDAR file not found: {lidar_full_path}") + print("lidar", lidar) + return lidar diff --git a/d123/dataset/dataset_specific/carla/load_sensor.py b/d123/dataset/dataset_specific/carla/load_sensor.py new file mode 100644 index 00000000..31605ec5 --- /dev/null +++ b/d123/dataset/dataset_specific/carla/load_sensor.py @@ -0,0 +1,11 @@ +from pathlib import Path + +import numpy as np + +from d123.common.datatypes.sensor.lidar import LiDAR + + +def load_carla_lidar_from_path(filepath: Path) -> LiDAR: + assert filepath.exists(), f"LiDAR file not found: {filepath}" + points = np.load(filepath).T + return LiDAR(point_cloud=points) diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 07b8be63..7d5fcd7b 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,8 +15,8 @@ defaults: - default_dataset_paths - _self_ - datasets: - - nuplan_private_dataset - # - carla_dataset + # - nuplan_private_dataset + - carla_dataset force_log_conversion: True force_map_conversion: False diff --git a/d123/script/config/datasets/carla_dataset.yaml b/d123/script/config/datasets/carla_dataset.yaml index b98bba2b..31184f78 100644 --- a/d123/script/config/datasets/carla_dataset.yaml +++ b/d123/script/config/datasets/carla_dataset.yaml @@ -4,7 +4,13 @@ carla_dataset: splits: ["carla"] log_path: "${oc.env:HOME}/carla_workspace/data" - output_path: ${d123_data_root} - force_log_conversion: ${force_log_conversion} - force_map_conversion: ${force_map_conversion} + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "path" + lidar_store_option: "path" diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 6fdb67fc..b4284bf4 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -9,8 +9,6 @@ "source": [ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", - "from pathlib import Path\n", - "from d123.dataset.scene.arrow_scene import ArrowScene\n", "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", @@ -24,9 +22,14 @@ "metadata": {}, "outputs": [], "source": [ - "split = \"nuplan_private_test\"\n", + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "\n", + "split = \"carla\"\n", + "log_names = [\"_Rep0_routes_devtest_route0_07_27_19_39_57\"]\n", "\n", - "log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "scene_tokens = None\n", "\n", "scene_filter = SceneFilter(\n", @@ -42,7 +45,9 @@ "# worker = RayDistributed()\n", "scenes = scene_builder.get_scenes(scene_filter, worker)\n", "\n", - "print(len(scenes))\n" + "print(len(scenes))\n", + "\n", + "\n" ] }, { @@ -54,7 +59,7 @@ "source": [ "\n", "visualization_server = ViserVisualizationServer()\n", - "visualization_server.set_scene(scenes[16])" + "visualization_server.set_scene(scenes[10])" ] }, { @@ -75,7 +80,8 @@ "outputs": [], "source": [ "# ego_vehicle_state.vehicle_parameters.height/2 - ego_vehicle_state.vehicle_parameters.height/3\n", - "\n" + "\n", + "asdasd\n" ] }, { From c5cfe9f6c3f0090b3af381a4480ab7eef17ca866 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Mon, 28 Jul 2025 22:24:31 +0200 Subject: [PATCH 07/42] Move some carla ue4 coordinate transforms to collection. Improve lidar. Improve circular loading in viser. --- d123/common/geometry/base.py | 1 - d123/common/geometry/transform/se3.py | 21 ++++++- d123/common/visualization/viser/server.py | 40 +++++++++---- d123/common/visualization/viser/utils.py | 20 ++++++- .../carla/carla_data_converter.py | 36 +++-------- .../dataset_specific/carla/load_sensor.py | 3 + notebooks/viz/viser_testing_v2_scene.ipynb | 59 +++++++++++++------ 7 files changed, 119 insertions(+), 61 deletions(-) diff --git a/d123/common/geometry/base.py b/d123/common/geometry/base.py index 460cff7d..28adc163 100644 --- a/d123/common/geometry/base.py +++ b/d123/common/geometry/base.py @@ -180,7 +180,6 @@ def __hash__(self) -> int: class StateSE3Index(IntEnum): - # TODO: implement X = 0 Y = 1 diff --git a/d123/common/geometry/transform/se3.py b/d123/common/geometry/transform/se3.py index ccd861c9..b71c7415 100644 --- a/d123/common/geometry/transform/se3.py +++ b/d123/common/geometry/transform/se3.py @@ -1,7 +1,7 @@ import numpy as np import numpy.typing as npt -from d123.common.geometry.base import StateSE3 +from d123.common.geometry.base import Point3DIndex, StateSE3 from d123.common.geometry.vector import Vector3D @@ -96,3 +96,22 @@ def convert_relative_to_absolute_points_3d_array( R = get_rotation_matrix(origin) absolute_points = points_3d_array @ R.T + origin.point_3d.array return absolute_points + + +def translate_points_3d_along_z( + state_se3: StateSE3, + points_3d: npt.NDArray[np.float64], + distance: float, +) -> npt.NDArray[np.float64]: + assert points_3d.shape[-1] == len(Point3DIndex) + + R = get_rotation_matrix(state_se3) + z_axis = R[:, 2] + + translated_points = np.zeros_like(points_3d) + + translated_points[..., Point3DIndex.X] = points_3d[..., Point3DIndex.X] + distance * z_axis[0] + translated_points[..., Point3DIndex.Y] = points_3d[..., Point3DIndex.Y] + distance * z_axis[1] + translated_points[..., Point3DIndex.Z] = points_3d[..., Point3DIndex.Z] + distance * z_axis[2] + + return translated_points diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index d1d24d2c..efd9841a 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -1,4 +1,5 @@ import time +from typing import List import trimesh import viser @@ -18,24 +19,42 @@ BACK_CAMERA_TYPE: CameraType = CameraType.CAM_B0 LIDAR_AVAILABLE: bool = True LIDAR_POINT_SIZE: float = 0.05 -MAP_AVAILABLE: bool = False +MAP_AVAILABLE: bool = True class ViserVisualizationServer: def __init__( self, + scenes: List[AbstractScene], + scene_index: int = 0, host: str = "localhost", port: int = 8080, label: str = "D123 Viser Server", ): + self.scenes = scenes + self.scene_index = scene_index - self.server = viser.ViserServer(host=host, port=port, label=label) + self.host = host + self.port = port + self.label = label + + self.server = viser.ViserServer(host=self.host, port=self.port, label=self.label) + self.set_scene(self.scenes[self.scene_index % len(self.scenes)]) + + def next(self) -> None: + self.server.flush() + self.server.gui.reset() + self.server.scene.reset() + self.scene_index = (self.scene_index + 1) % len(self.scenes) + print(f"Viser server started at {self.host}:{self.port}") + self.set_scene(self.scenes[self.scene_index]) def set_scene(self, scene: AbstractScene) -> None: num_frames = scene.get_number_of_iterations() self.server.gui.configure_theme(control_width="large") with self.server.gui.add_folder("Playback"): + server_playing = True gui_timestep = self.server.gui.add_slider( "Timestep", @@ -47,6 +66,7 @@ def set_scene(self, scene: AbstractScene) -> None: ) gui_next_frame = self.server.gui.add_button("Next Frame", disabled=True) gui_prev_frame = self.server.gui.add_button("Prev Frame", disabled=True) + gui_next_scene = self.server.gui.add_button("Next Scene", disabled=False) gui_playing = self.server.gui.add_checkbox("Playing", True) gui_framerate = self.server.gui.add_slider("FPS", min=1, max=60, step=0.1, initial_value=10) gui_framerate_options = self.server.gui.add_button_group("FPS options", ("10", "20", "30", "60")) @@ -60,6 +80,11 @@ def _(_) -> None: def _(_) -> None: gui_timestep.value = (gui_timestep.value - 1) % num_frames + @gui_next_scene.on_click + def _(_) -> None: + nonlocal server_playing + server_playing = False + # Disable frame controls when we're playing. @gui_playing.on_update def _(_) -> None: @@ -170,15 +195,10 @@ def _(_) -> None: # Playback update loop. prev_timestep = gui_timestep.value - while True: + while server_playing: # Update the timestep if we're playing. if gui_playing.value: gui_timestep.value = (gui_timestep.value + 1) % num_frames - # Update point size of both this timestep and the next one! There's - # redundancy here, but this will be optimized out internally by viser. - # - # We update the point size for the next timestep so that it will be - # immediately available when we toggle the visibility. - # dynamic_meshes[gui_timestep.value].point_size = gui_point_size.value - # dynamic_meshes[(gui_timestep.value + 1) % num_frames].point_size = gui_point_size.value + self.server.flush() + self.next() diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index f95c54fe..ff4604ba 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -223,11 +223,27 @@ def euler_to_quaternion_scipy(roll: float, pitch: float, yaw: float) -> npt.NDAr def get_lidar_points(scene: AbstractScene, iteration: int) -> npt.NDArray[np.float32]: + from d123.common.geometry.transform.se3 import translate_points_3d_along_z + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) current_ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) lidar = scene.get_lidar_at_iteration(iteration) - points = convert_relative_to_absolute_points_3d_array(current_ego_vehicle_state.rear_axle_se3, lidar.xyz) - points = points - initial_ego_vehicle_state.center_se3.point_3d.array + if scene.log_metadata.dataset == "nuplan": + # NOTE: nuPlan uses the rear axle as origin. + origin = current_ego_vehicle_state.rear_axle_se3 + points = lidar.xyz + points = convert_relative_to_absolute_points_3d_array(origin, points) + points = points - initial_ego_vehicle_state.center_se3.point_3d.array + elif scene.log_metadata.dataset == "carla": + # NOTE: CARLA uses the center of the vehicle as origin. + origin = current_ego_vehicle_state.center_se3 + points = translate_points_3d_along_z( + origin, lidar.xyz, -current_ego_vehicle_state.vehicle_parameters.height / 2 + ) + points = convert_relative_to_absolute_points_3d_array(origin, points) + points = points - initial_ego_vehicle_state.center_se3.point_3d.array + else: + raise ValueError(f"Unsupported dataset: {scene.log_metadata.dataset}") return points diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 943ceda3..2cd2bfbf 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -15,9 +15,8 @@ from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import get_carla_lincoln_mkz_2020_parameters -from d123.common.geometry.base import Point2D, Point3D, StateSE3 +from d123.common.geometry.base import Point2D, Point3D from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index -from d123.common.geometry.transform.se3 import translate_se3_along_z from d123.common.geometry.vector import Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive @@ -312,34 +311,17 @@ def _write_recording_table( def _extract_ego_vehicle_state(ego_state_list: List[float]) -> List[float]: + # NOTE: This function used to convert coordinate systems, but it is not needed anymore. + assert len(ego_state_list) == len(EgoStateSE3Index), "Ego state list has incorrect length." + return ego_state_list - # NOTE: @DanielDauner This function is a temporary workaround. - # CARLAs bounding box location starts at bottom vertically. - # Need to translate half the height along the z-axis. - ego_state_array = np.array(ego_state_list, dtype=np.float64) - vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() - center = StateSE3.from_array(ego_state_array[EgoStateSE3Index.SE3]) - center = translate_se3_along_z(center, vehicle_parameters.height / 2) - ego_state_array[EgoStateSE3Index.SE3] = center.array - - return ego_state_array.tolist() - - -def _extract_detection_states(detection_states: List[List[float]]) -> List[float]: - - # NOTE: @DanielDauner This function is a temporary workaround. - # CARLAs bounding box location starts at bottom vertically. - # Need to translate half the height along the z-axis. +def _extract_detection_states(detection_states: List[List[float]]) -> List[List[float]]: + # NOTE: This function used to convert coordinate systems, but it is not needed anymore. detection_state_converted = [] - for detection_state in detection_states: - detection_state_array = np.array(detection_state, dtype=np.float64) - center = StateSE3.from_array(detection_state_array[BoundingBoxSE3Index.STATE_SE3]) - center = translate_se3_along_z(center, detection_state_array[BoundingBoxSE3Index.HEIGHT] / 2) - detection_state_array[EgoStateSE3Index.SE3] = center.array - detection_state_converted.append(detection_state_array.tolist()) - + assert len(detection_state) == len(BoundingBoxSE3Index), "Detection state has incorrect length." + detection_state_converted.append(detection_state) return detection_state_converted @@ -412,7 +394,6 @@ def _extract_cameras( raise NotImplementedError("Binary camera storage is not implemented.") else: camera_dict[camera_type] = None - print("camera_dict", camera_dict) return camera_dict @@ -427,5 +408,4 @@ def _extract_lidar(log_name: str, sample_name: str, data_converter_config: DataC raise NotImplementedError("Binary lidar storage is not implemented.") else: raise FileNotFoundError(f"LiDAR file not found: {lidar_full_path}") - print("lidar", lidar) return lidar diff --git a/d123/dataset/dataset_specific/carla/load_sensor.py b/d123/dataset/dataset_specific/carla/load_sensor.py index 31605ec5..c4feeb8b 100644 --- a/d123/dataset/dataset_specific/carla/load_sensor.py +++ b/d123/dataset/dataset_specific/carla/load_sensor.py @@ -8,4 +8,7 @@ def load_carla_lidar_from_path(filepath: Path) -> LiDAR: assert filepath.exists(), f"LiDAR file not found: {filepath}" points = np.load(filepath).T + + points[1] = -points[1] # FIXME + return LiDAR(point_cloud=points) diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index b4284bf4..949d49db 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -12,7 +12,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "\n" ] }, { @@ -21,51 +22,66 @@ "id": "1", "metadata": {}, "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], "source": [ "# split = \"nuplan_private_test\"\n", "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", + "# split = \"carla\"\n", + "# log_names = [\"_Rep0_routes_123d_route1_07_28_20_46_56\"]\n", + "\n", "\n", + "splits = [\"carla\", \"nuplan_private_test\"]\n", + "log_names = [\n", + " \"_Rep0_routes_123d_route0_07_28_20_36_14\",\n", + " \"_Rep0_routes_123d_route1_07_28_20_46_56\",\n", + " \"2021.07.01.21.22.09_veh-14_00016_00656\",\n", + " \"2021.07.20.17.06.09_veh-38_04743_05687\",\n", + " \"2021.09.08.14.01.10_veh-39_01898_02179\",\n", + " \"2021.09.29.17.35.58_veh-44_00066_00432\",\n", + " \"2021.10.01.01.03.27_veh-50_01313_01390\",\n", + "]\n", "\n", - "split = \"carla\"\n", - "log_names = [\"_Rep0_routes_devtest_route0_07_27_19_39_57\"]\n", "\n", "scene_tokens = None\n", "\n", "scene_filter = SceneFilter(\n", - " split_names=[split],\n", + " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", " duration_s=20.1,\n", " history_s=1.0,\n", " timestamp_threshold_s=20.1,\n", + " shuffle=True,\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", "# worker = RayDistributed()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "print(len(scenes))\n", - "\n", - "\n" + "scenes = scene_builder.get_scenes(scene_filter, worker)" ] }, { "cell_type": "code", "execution_count": null, - "id": "2", + "id": "3", "metadata": {}, "outputs": [], "source": [ "\n", - "visualization_server = ViserVisualizationServer()\n", - "visualization_server.set_scene(scenes[10])" + "visualization_server = ViserVisualizationServer(scenes)" ] }, { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -75,19 +91,24 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [ - "# ego_vehicle_state.vehicle_parameters.height/2 - ego_vehicle_state.vehicle_parameters.height/3\n", + "import numpy as np\n", + "\n", + "path = \"/home/daniel/carla_workspace/data/sensor_blobs/_Rep0_routes_devtest_route0_07_28_15_04_50/lidar/0000000502.npy\"\n", + "lidar = np.load(path).T\n", "\n", - "asdasd\n" + "import matplotlib.pyplot as plt\n", + "\n", + "plt.scatter(lidar[0], lidar[1], s=0.1)\n" ] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -112,7 +133,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -128,7 +149,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [], "source": [ From def4cf2c9cadcb5478e9b6ed3f7c78a23d29a621 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Tue, 29 Jul 2025 11:15:04 +0200 Subject: [PATCH 08/42] Add viser visualization with lines --- .../bounding_box/bounding_box_index.py | 19 +++ d123/common/visualization/viser/server.py | 33 +++-- d123/common/visualization/viser/utils_v2.py | 113 ++++++++++++++++++ notebooks/viz/viser_testing_v2_scene.ipynb | 78 ------------ 4 files changed, 156 insertions(+), 87 deletions(-) create mode 100644 d123/common/visualization/viser/utils_v2.py diff --git a/d123/common/geometry/bounding_box/bounding_box_index.py b/d123/common/geometry/bounding_box/bounding_box_index.py index d5a1ee26..c282b07e 100644 --- a/d123/common/geometry/bounding_box/bounding_box_index.py +++ b/d123/common/geometry/bounding_box/bounding_box_index.py @@ -48,3 +48,22 @@ def STATE_SE3(cls) -> slice: @classproperty def ROTATION_XYZ(cls) -> slice: return slice(cls.ROLL, cls.YAW + 1) + + +class Corners3DIndex(IntEnum): + FRONT_LEFT_BOTTOM = 0 + FRONT_RIGHT_BOTTOM = 1 + BACK_RIGHT_BOTTOM = 2 + BACK_LEFT_BOTTOM = 3 + FRONT_LEFT_TOP = 4 + FRONT_RIGHT_TOP = 5 + BACK_RIGHT_TOP = 6 + BACK_LEFT_TOP = 7 + + @classproperty + def BOTTOM(cls) -> slice: + return slice(cls.FRONT_LEFT_BOTTOM, cls.BACK_LEFT_BOTTOM + 1) + + @classproperty + def TOP(cls) -> slice: + return slice(cls.FRONT_LEFT_TOP, cls.BACK_LEFT_TOP + 1) diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index efd9841a..1de9885b 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -1,5 +1,5 @@ import time -from typing import List +from typing import List, Literal import trimesh import viser @@ -10,6 +10,7 @@ get_lidar_points, get_map_meshes, ) +from d123.common.visualization.viser.utils_v2 import get_bounding_box_outlines from d123.dataset.scene.abstract_scene import AbstractScene # TODO: Try to fix performance issues. @@ -20,6 +21,8 @@ LIDAR_AVAILABLE: bool = True LIDAR_POINT_SIZE: float = 0.05 MAP_AVAILABLE: bool = True +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" +LINE_WIDTH: float = 4.0 class ViserVisualizationServer: @@ -108,14 +111,26 @@ def _(_) -> None: start = time.time() # with self.server.atomic(): mew_frame_handle = self.server.scene.add_frame(f"/frame{gui_timestep.value}", show_axes=False) - meshes = [] - for name, mesh in get_bounding_box_meshes(scene, gui_timestep.value).items(): - meshes.append(mesh) - self.server.scene.add_mesh_trimesh( - f"/frame{gui_timestep.value}/detections", - trimesh.util.concatenate(meshes), - visible=True, - ) + if BOUNDING_BOX_TYPE == "mesh": + meshes = [] + for _, mesh in get_bounding_box_meshes(scene, gui_timestep.value).items(): + meshes.append(mesh) + self.server.scene.add_mesh_trimesh( + f"/frame{gui_timestep.value}/detections", + trimesh.util.concatenate(meshes), + visible=True, + ) + elif BOUNDING_BOX_TYPE == "lines": + lines, colors = get_bounding_box_outlines(scene, gui_timestep.value) + self.server.scene.add_line_segments( + f"/frame{gui_timestep.value}/detections", + points=lines, + colors=colors, + line_width=LINE_WIDTH, + ) + else: + raise ValueError(f"Unknown bounding box type: {BOUNDING_BOX_TYPE}") + if FRONT_CAMERA_TYPE in scene.available_camera_types: gui_front_cam.image = scene.get_camera_at_iteration(gui_timestep.value, FRONT_CAMERA_TYPE).image diff --git a/d123/common/visualization/viser/utils_v2.py b/d123/common/visualization/viser/utils_v2.py new file mode 100644 index 00000000..9fc3a45c --- /dev/null +++ b/d123/common/visualization/viser/utils_v2.py @@ -0,0 +1,113 @@ +import numpy as np +import numpy.typing as npt + +# from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters +from d123.common.geometry.base import Point3D, Point3DIndex +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 +from d123.common.geometry.bounding_box.bounding_box_index import Corners3DIndex +from d123.common.geometry.transform.se3 import translate_body_frame +from d123.common.geometry.vector import Vector3D +from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG +from d123.dataset.scene.abstract_scene import AbstractScene + +# TODO: Refactor this file. +# TODO: Add general utilities for 3D primitives and mesh support. + + +def _get_bounding_box_corners(bounding_box: BoundingBoxSE3) -> npt.NDArray[np.float64]: + """ + Get the vertices of a bounding box in 3D space. + """ + # TODO: apply transform over array batch instead + corners = np.zeros((len(Corners3DIndex), len(Point3DIndex)), dtype=np.float64) + corners[Corners3DIndex.FRONT_LEFT_BOTTOM] = translate_body_frame( + bounding_box.center, + Vector3D(bounding_box.length / 2, -bounding_box.width / 2, -bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.FRONT_RIGHT_BOTTOM] = translate_body_frame( + bounding_box.center, + Vector3D(bounding_box.length / 2, bounding_box.width / 2, -bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.BACK_RIGHT_BOTTOM] = translate_body_frame( + bounding_box.center, + Vector3D(-bounding_box.length / 2, bounding_box.width / 2, -bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.BACK_LEFT_BOTTOM] = translate_body_frame( + bounding_box.center, + Vector3D(-bounding_box.length / 2, -bounding_box.width / 2, -bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.FRONT_LEFT_TOP] = translate_body_frame( + bounding_box.center, + Vector3D(bounding_box.length / 2, -bounding_box.width / 2, bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.FRONT_RIGHT_TOP] = translate_body_frame( + bounding_box.center, + Vector3D(bounding_box.length / 2, bounding_box.width / 2, bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.BACK_RIGHT_TOP] = translate_body_frame( + bounding_box.center, + Vector3D(-bounding_box.length / 2, bounding_box.width / 2, bounding_box.height / 2), + ).point_3d.array + corners[Corners3DIndex.BACK_LEFT_TOP] = translate_body_frame( + bounding_box.center, + Vector3D(-bounding_box.length / 2, -bounding_box.width / 2, bounding_box.height / 2), + ).point_3d.array + return corners + + +def _get_bounding_box_lines(bounding_box: BoundingBoxSE3) -> npt.NDArray[np.float64]: + """ + Get the edges of a bounding box in 3D space as a Polyline3D. + """ + corners = _get_bounding_box_corners(bounding_box) + index_pairs = [ + (Corners3DIndex.FRONT_LEFT_BOTTOM, Corners3DIndex.FRONT_RIGHT_BOTTOM), + (Corners3DIndex.FRONT_RIGHT_BOTTOM, Corners3DIndex.BACK_RIGHT_BOTTOM), + (Corners3DIndex.BACK_RIGHT_BOTTOM, Corners3DIndex.BACK_LEFT_BOTTOM), + (Corners3DIndex.BACK_LEFT_BOTTOM, Corners3DIndex.FRONT_LEFT_BOTTOM), + (Corners3DIndex.FRONT_LEFT_TOP, Corners3DIndex.FRONT_RIGHT_TOP), + (Corners3DIndex.FRONT_RIGHT_TOP, Corners3DIndex.BACK_RIGHT_TOP), + (Corners3DIndex.BACK_RIGHT_TOP, Corners3DIndex.BACK_LEFT_TOP), + (Corners3DIndex.BACK_LEFT_TOP, Corners3DIndex.FRONT_LEFT_TOP), + (Corners3DIndex.FRONT_LEFT_BOTTOM, Corners3DIndex.FRONT_LEFT_TOP), + (Corners3DIndex.FRONT_RIGHT_BOTTOM, Corners3DIndex.FRONT_RIGHT_TOP), + (Corners3DIndex.BACK_RIGHT_BOTTOM, Corners3DIndex.BACK_RIGHT_TOP), + (Corners3DIndex.BACK_LEFT_BOTTOM, Corners3DIndex.BACK_LEFT_TOP), + ] + lines = np.zeros((len(index_pairs), 2, len(Point3DIndex)), dtype=np.float64) + for i, (start_idx, end_idx) in enumerate(index_pairs): + lines[i, 0] = corners[start_idx] + lines[i, 1] = corners[end_idx] + return lines + + +def translate_points_3d(points_3d: npt.NDArray[np.float64], point_3d: Point3D) -> npt.NDArray[np.float64]: + return points_3d - point_3d.array + + +def get_bounding_box_outlines(scene: AbstractScene, iteration: int): + + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) + ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) + box_detections = scene.get_box_detections_at_iteration(iteration) + + lines = [] + colors = [] + for box_detection in box_detections: + bbox: BoundingBoxSE3 = box_detection.bounding_box_se3 + bbox_lines = _get_bounding_box_lines(bbox) + bbox_lines = translate_points_3d(bbox_lines, initial_ego_vehicle_state.center_se3.point_3d) + bbox_color = np.zeros(bbox_lines.shape, dtype=np.float32) + bbox_color[..., :] = BOX_DETECTION_CONFIG[box_detection.metadata.detection_type].fill_color.rgb_norm + + lines.append(bbox_lines) + colors.append(bbox_color) + + ego_bbox_lines = _get_bounding_box_lines(ego_vehicle_state.bounding_box_se3) + ego_bbox_lines = translate_points_3d(ego_bbox_lines, initial_ego_vehicle_state.center_se3.point_3d) + ego_bbox_color = np.zeros(ego_bbox_lines.shape, dtype=np.float32) + ego_bbox_color[..., :] = EGO_VEHICLE_CONFIG.fill_color.rgb_norm + + lines.append(ego_bbox_lines) + colors.append(ego_bbox_color) + return np.concatenate(lines, axis=0), np.concatenate(colors, axis=0) diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 949d49db..d0511676 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -77,84 +77,6 @@ "\n", "visualization_server = ViserVisualizationServer(scenes)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "# ego_vehicle_state = scene.get_ego_vehicle_state_at_iteration(0)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "path = \"/home/daniel/carla_workspace/data/sensor_blobs/_Rep0_routes_devtest_route0_07_28_15_04_50/lidar/0000000502.npy\"\n", - "lidar = np.load(path).T\n", - "\n", - "import matplotlib.pyplot as plt\n", - "\n", - "plt.scatter(lidar[0], lidar[1], s=0.1)\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [ - "# from pyquaternion import Quaternion\n", - "# import numpy as np\n", - "# # Example Euler angles (in radians) and rotation sequence\n", - "# # roll = 0.5 # Rotation about the X-axis (radians)\n", - "# # pitch = 1.0 # Rotation about the Y-axis (radians)\n", - "# # yaw = 1.5 # Rotation about the Z-axis (radians)\n", - "# roll = 1.0 # Rotation about the X-axis (radians)\n", - "# pitch = 2.0 # Rotation about the Y-axis (radians)\n", - "# yaw = 1.0 # Rotation about the Z-axis (radians)\n", - "# rotation_sequence = 'xyz' # Or 'yxz', 'zxy', etc.\n", - "\n", - "# # Create a quaternion from Euler angles\n", - "# quaternion = Quaternion(axis=[roll, pitch, yaw], sequence=rotation_sequence)\n", - "\n", - "\n", - "# quaternion.yaw_pitch_roll" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [ - "# from d123.common.datatypes.camera.camera_parameters import get_nuplan_camera_parameters\n", - "# import numpy as np\n", - "\n", - "# fy = get_nuplan_camera_parameters().intrinsics[1, 1]\n", - "\n", - "# fov_v = 2 * np.arctan(1080 / (2 * fy))\n", - "# fov_v" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8", - "metadata": {}, - "outputs": [], - "source": [ - "# get_nuplan_camera_parameters().intrinsics, fy" - ] } ], "metadata": { From 003ec8506bff56a10c51e3f362369d80486bf06e Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 1 Aug 2025 15:42:52 +0200 Subject: [PATCH 09/42] Add extrinsic as column in arrow table for nuPlan --- d123/common/datatypes/sensor/camera.py | 29 ++++- d123/common/visualization/viser/server.py | 107 ++++++++++-------- d123/common/visualization/viser/utils.py | 47 ++++---- d123/dataset/arrow/conversion.py | 7 ++ .../nuplan/nuplan_data_converter.py | 43 +++++-- .../default_dataset_conversion.yaml | 4 +- docs/schema.md | 14 +++ notebooks/scene_rendering.ipynb | 83 ++++++++++++-- notebooks/viz/viser_testing_v2_scene.ipynb | 50 ++++---- 9 files changed, 262 insertions(+), 122 deletions(-) create mode 100644 docs/schema.md diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py index cca14dd5..f03917f2 100644 --- a/d123/common/datatypes/sensor/camera.py +++ b/d123/common/datatypes/sensor/camera.py @@ -33,8 +33,6 @@ class CameraMetadata: height: int intrinsic: npt.NDArray[np.float64] # 3x3 matrix distortion: npt.NDArray[np.float64] # 5x1 vector - translation: npt.NDArray[np.float64] # 3x1 vector - rotation: npt.NDArray[np.float64] # 3x3 matrix def to_dict(self) -> Dict[str, Any]: # TODO: remove None types. Only a placeholder for now. @@ -44,8 +42,6 @@ def to_dict(self) -> Dict[str, Any]: "height": self.height, "intrinsic": self.intrinsic.tolist() if self.intrinsic is not None else None, "distortion": self.distortion.tolist() if self.distortion is not None else None, - "translation": self.translation.tolist() if self.translation is not None else None, - "rotation": self.rotation.tolist() if self.rotation is not None else None, } @classmethod @@ -57,10 +53,30 @@ def from_dict(cls, json_dict: Dict[str, Any]) -> CameraMetadata: height=json_dict["height"], intrinsic=np.array(json_dict["intrinsic"]) if json_dict["intrinsic"] is not None else None, distortion=np.array(json_dict["distortion"]) if json_dict["distortion"] is not None else None, - translation=np.array(json_dict["translation"]) if json_dict["translation"] is not None else None, - rotation=np.array(json_dict["rotation"]) if json_dict["rotation"] is not None else None, ) + @property + def aspect_ratio(self) -> float: + return self.width / self.height + + @property + def fov_x(self) -> float: + """ + Calculates the horizontal field of view (FOV) in radian. + """ + fx = self.intrinsic[0, 0] + fov_x_rad = 2 * np.arctan(self.width / (2 * fx)) + return fov_x_rad + + @property + def fov_y(self) -> float: + """ + Calculates the vertical field of view (FOV) in radian. + """ + fy = self.intrinsic[1, 1] + fov_y_rad = 2 * np.arctan(self.height / (2 * fy)) + return fov_y_rad + def camera_metadata_dict_to_json(camera_metadata: Dict[CameraType, CameraMetadata]) -> Dict[str, Dict[str, Any]]: """ @@ -90,6 +106,7 @@ class Camera: metadata: CameraMetadata image: npt.NDArray[np.uint8] + extrinsic: npt.NDArray[np.float64] # 4x4 matrix def get_view_matrix(self) -> np.ndarray: # Compute the view matrix based on the camera's position and orientation diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 1de9885b..e14214f3 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -1,5 +1,5 @@ import time -from typing import List, Literal +from typing import Dict, List, Literal import trimesh import viser @@ -7,6 +7,7 @@ from d123.common.datatypes.sensor.camera import CameraType from d123.common.visualization.viser.utils import ( get_bounding_box_meshes, + get_camera_values, get_lidar_points, get_map_meshes, ) @@ -16,14 +17,27 @@ # TODO: Try to fix performance issues. # TODO: Refactor this file. -FRONT_CAMERA_TYPE: CameraType = CameraType.CAM_F0 -BACK_CAMERA_TYPE: CameraType = CameraType.CAM_B0 -LIDAR_AVAILABLE: bool = True +# all_camera_types: List[CameraType] = [ +# CameraType.CAM_F0, +# CameraType.CAM_B0, +# CameraType.CAM_L0, +# CameraType.CAM_L1, +# CameraType.CAM_L2, +# CameraType.CAM_R0, +# CameraType.CAM_R1, +# CameraType.CAM_R2, +# ] + LIDAR_POINT_SIZE: float = 0.05 MAP_AVAILABLE: bool = True BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" LINE_WIDTH: float = 4.0 +CAMERA_SCALE: float = 1.0 +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] +VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] +LIDAR_AVAILABLE: bool = True + class ViserVisualizationServer: def __init__( @@ -54,6 +68,7 @@ def next(self) -> None: def set_scene(self, scene: AbstractScene) -> None: num_frames = scene.get_number_of_iterations() + print(scene.available_camera_types) self.server.gui.configure_theme(control_width="large") with self.server.gui.add_folder("Playback"): @@ -131,22 +146,27 @@ def _(_) -> None: else: raise ValueError(f"Unknown bounding box type: {BOUNDING_BOX_TYPE}") - if FRONT_CAMERA_TYPE in scene.available_camera_types: - gui_front_cam.image = scene.get_camera_at_iteration(gui_timestep.value, FRONT_CAMERA_TYPE).image + for camera_type in VISUALIZE_CAMERA_GUI: + if camera_type in scene.available_camera_types: + camera_gui_handles[camera_type].image = scene.get_camera_at_iteration( + gui_timestep.value, camera_type + ).image + + for camera_type in VISUALIZE_CAMERA_FRUSTUM: + if camera_type in scene.available_camera_types: + camera_position, camera_quaternion, camera = get_camera_values( + scene, camera_type, gui_timestep.value + ) - if BACK_CAMERA_TYPE in scene.available_camera_types: - gui_back_cam.image = scene.get_camera_at_iteration(gui_timestep.value, BACK_CAMERA_TYPE).image + camera_frustum_handles[camera_type].position = camera_position.array + camera_frustum_handles[camera_type].wxyz = camera_quaternion.q + camera_frustum_handles[camera_type].image = camera.image + + # if BACK_CAMERA_TYPE in scene.available_camera_types: + # gui_back_cam.image = scene.get_camera_at_iteration(gui_timestep.value, BACK_CAMERA_TYPE).image if LIDAR_AVAILABLE: gui_lidar.points = get_lidar_points(scene, gui_timestep.value) - # camera_pose = _get_camera_pose_demo(scene, gui_timestep.value) - # frustum_handle.position = camera_pose.point_3d.array - # frustum_handle.wxyz = euler_to_quaternion_scipy(camera_pose.roll, camera_pose.pitch, camera_pose.yaw) - # frustum_handle.image = np.array(scene.get_front_cam_demo(gui_timestep.value)) - - # ego_frame_pose = _get_ego_frame_pose(scene, gui_timestep.value) - # ego_frame_handle.position = ego_frame_pose.point_3d.array - # ego_frame_handle.wxyz = euler_to_quaternion_scipy(ego_frame_pose.roll, ego_frame_pose.pitch, ego_frame_pose.yaw) prev_timestep = current_timestep @@ -161,20 +181,31 @@ def _(_) -> None: current_frame_handle = self.server.scene.add_frame(f"/frame{gui_timestep.value}", show_axes=False) self.server.scene.add_frame("/map", show_axes=False) - if FRONT_CAMERA_TYPE in scene.available_camera_types: - with self.server.gui.add_folder("Camera"): - gui_front_cam = self.server.gui.add_image( - image=scene.get_camera_at_iteration(gui_timestep.value, FRONT_CAMERA_TYPE).image, - label=FRONT_CAMERA_TYPE.serialize(), - format="jpeg", + camera_gui_handles: Dict[CameraType, viser.GuiImageHandle] = {} + camera_frustum_handles: Dict[CameraType, viser.CameraFrustumHandle] = {} + + for camera_type in VISUALIZE_CAMERA_GUI: + if camera_type in scene.available_camera_types: + with self.server.gui.add_folder(f"Camera {camera_type.serialize()}"): + camera_gui_handles[camera_type] = self.server.gui.add_image( + image=scene.get_camera_at_iteration(gui_timestep.value, camera_type).image, + label=camera_type.serialize(), + format="jpeg", + ) + + for camera_type in VISUALIZE_CAMERA_FRUSTUM: + if camera_type in scene.available_camera_types: + camera_position, camera_quaternion, camera = get_camera_values( + scene, camera_type, gui_timestep.value ) - - if BACK_CAMERA_TYPE in scene.available_camera_types: - with self.server.gui.add_folder("Camera"): - gui_back_cam = self.server.gui.add_image( - image=scene.get_camera_at_iteration(gui_timestep.value, BACK_CAMERA_TYPE).image, - label=BACK_CAMERA_TYPE.serialize(), - format="jpeg", + camera_frustum_handles[camera_type] = self.server.scene.add_camera_frustum( + f"camera_frustum_{camera_type.serialize()}", + fov=camera.metadata.fov_y, + aspect=camera.metadata.aspect_ratio, + scale=CAMERA_SCALE, + image=camera.image, + position=camera_position.array, + wxyz=camera_quaternion.q, ) if LIDAR_AVAILABLE: @@ -190,24 +221,6 @@ def _(_) -> None: for name, mesh in get_map_meshes(scene).items(): self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) - # camera_pose = _get_camera_pose_demo(scene, gui_timestep.value) - # frustum_handle = self.server.scene.add_camera_frustum( - # "camera_frustum", - # fov=0.6724845869242845, - # aspect=16 / 9, - # scale=0.30, - # image=np.array(scene.get_front_cam_demo(gui_timestep.value)), - # position=camera_pose.point_3d.array, - # wxyz=euler_to_quaternion_scipy(camera_pose.roll, camera_pose.pitch, camera_pose.yaw), - # ) - - # ego_frame_pose = _get_ego_frame_pose(scene, gui_timestep.value) - # ego_frame_handle = self.server.scene.add_frame( - # "ego_frame_handle", - # position=ego_frame_pose.point_3d.array, - # wxyz=euler_to_quaternion_scipy(ego_frame_pose.roll, ego_frame_pose.pitch, ego_frame_pose.yaw) - # ) - # Playback update loop. prev_timestep = gui_timestep.value while server_playing: diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index ff4604ba..5056ae5f 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -1,8 +1,12 @@ +from typing import Tuple + import numpy as np import numpy.typing as npt import trimesh +from pyquaternion import Quaternion # from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters +from d123.common.datatypes.sensor.camera import Camera, CameraType from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D @@ -157,48 +161,35 @@ def _create_lane_mesh_from_boundary_arrays( return mesh -def _get_camera_pose_demo(scene: AbstractScene, iteration: int) -> StateSE3: - # NOTE: This function does not work. - +def get_camera_values( + scene: AbstractScene, camera_type: CameraType, iteration: int +) -> Tuple[Point3D, Quaternion, Camera]: initial_point_3d = scene.get_ego_state_at_iteration(0).center_se3.point_3d rear_axle = scene.get_ego_state_at_iteration(iteration).rear_axle_se3 + camera = scene.get_camera_at_iteration(iteration, camera_type) + rear_axle_array = rear_axle.array rear_axle_array[:3] -= initial_point_3d.array - rear_axle = StateSE3.from_array(rear_axle_array) - camera_parameters = None - - # Get camera parameters - camera_translation = camera_parameters.translation # 3x1 vector as array - camera_rotation = camera_parameters.rotation # 3x3 matrix as array + camera_to_ego = camera.extrinsic # 4x4 transformation from camera to ego frame + camera.image # Get the rotation matrix of the rear axle pose from d123.common.geometry.transform.se3 import get_rotation_matrix - rear_axle_rotation = get_rotation_matrix(rear_axle) - - # Apply camera translation in the rear axle's coordinate frame - # Transform camera translation from rear axle frame to world frame - world_translation = rear_axle_rotation @ camera_translation - - # Apply camera rotation in the rear axle's coordinate frame - world_rotation = rear_axle_rotation @ camera_rotation - - # Extract Euler angles from the composed rotation matrix using scipy - # Use 'ZYX' convention to match the rotation matrix composition order - from scipy.spatial.transform import Rotation + ego_transform = np.eye(4, dtype=np.float64) + ego_transform[:3, :3] = get_rotation_matrix(rear_axle) + ego_transform[:3, 3] = rear_axle.point_3d.array - r = Rotation.from_matrix(world_rotation) - roll, pitch, yaw = r.as_euler("XYZ", degrees=False) + # Camera transformation in ego frame + camera_transform = ego_transform @ camera_to_ego - # Calculate camera position in world coordinates - camera_x = rear_axle.x + world_translation[0] - camera_y = rear_axle.y + world_translation[1] - camera_z = rear_axle.z + world_translation[2] + camera_position = Point3D(*camera_transform[:3, 3]) + camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) - return StateSE3(camera_x, camera_y, camera_z, roll, pitch, yaw) + return camera_position, camera_rotation, camera def _get_ego_frame_pose(scene: AbstractScene, iteration: int) -> StateSE3: diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 50a6160a..ada284c4 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -102,6 +102,12 @@ def get_camera_from_arrow_table( ) -> Camera: table_data = arrow_table[camera_metadata.camera_type.serialize()][index].as_py() + extrinsic = arrow_table[f"{camera_metadata.camera_type.serialize()}_extrinsic"][index].as_py() + extrinsic = np.array(extrinsic).reshape((4, 4)) if extrinsic else np.eye(4) + + if table_data is None: + return None + image: Optional[npt.NDArray[np.uint8]] = None if isinstance(table_data, str): @@ -117,6 +123,7 @@ def get_camera_from_arrow_table( return Camera( metadata=camera_metadata, image=image, + extrinsic=extrinsic, ) diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index fd05de11..59aea551 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -11,17 +11,20 @@ import pyarrow as pa import yaml from nuplan.database.nuplan_db.nuplan_scenario_queries import get_cameras, get_images_from_lidar_tokens +from nuplan.database.nuplan_db_orm.ego_pose import EgoPose from nuplan.database.nuplan_db_orm.lidar_box import LidarBox from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB from nuplan.planning.simulation.observation.observation_type import CameraChannel from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map from pyquaternion import Quaternion +from sqlalchemy import func import d123.dataset.dataset_specific.nuplan.utils as nuplan_utils from d123.common.datatypes.detection.detection import TrafficLightStatus from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( get_nuplan_pacifica_parameters, @@ -67,6 +70,7 @@ } NUPLAN_DATA_ROOT = Path(os.environ["NUPLAN_DATA_ROOT"]) +NUPLAN_ROLLING_SHUTTER_S: Final[TimePoint] = TimePoint.from_s(1 / 60) def create_splits_logs() -> Dict[str, List[str]]: @@ -211,6 +215,10 @@ def convert_nuplan_log_to_arrow( for camera_type in NUPLAN_CAMERA_TYPES.keys(): if data_converter_config.camera_store_option == "path": schema_column_list.append((camera_type.serialize(), pa.string())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) + ) + elif data_converter_config.camera_store_option == "binary": raise NotImplementedError("Binary camera storage is not implemented.") @@ -246,7 +254,6 @@ def get_nuplan_camera_metadata(log_path: Path) -> Dict[str, CameraMetadata]: def _get_camera_metadata(camera_type: CameraType) -> CameraMetadata: cam = list(get_cameras(log_path, [str(NUPLAN_CAMERA_TYPES[camera_type].value)]))[0] intrinsic = np.array(pickle.loads(cam.intrinsic)) - translation = np.array(pickle.loads(cam.translation)) rotation = np.array(pickle.loads(cam.rotation)) rotation = Quaternion(rotation).rotation_matrix distortion = np.array(pickle.loads(cam.distortion)) @@ -256,8 +263,6 @@ def _get_camera_metadata(camera_type: CameraType) -> CameraMetadata: height=cam.height, intrinsic=intrinsic, distortion=distortion, - translation=translation, - rotation=rotation, ) log_cam_infos: Dict[str, CameraMetadata] = {} @@ -312,12 +317,14 @@ def _write_recording_table( row_data["lidar"] = [_extract_lidar(lidar_pc, data_converter_config)] if data_converter_config.camera_store_option is not None: - camera_data_dict = _extract_camera(lidar_pc, source_log_path, data_converter_config) + camera_data_dict = _extract_camera(log_db, lidar_pc, source_log_path, data_converter_config) for camera_type, camera_data in camera_data_dict.items(): if camera_data is not None: - row_data[camera_type.serialize()] = [camera_data] + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] else: row_data[camera_type.serialize()] = [None] + row_data[f"{camera_type.serialize()}_extrinsic"] = [None] batch = pa.record_batch(row_data, schema=recording_schema) writer.write_batch(batch) @@ -418,6 +425,7 @@ def _extract_scenario_tag(log_db: NuPlanDB, lidar_pc_token: str) -> List[str]: def _extract_camera( + log_db: NuPlanDB, lidar_pc: LidarPc, source_log_path: Path, data_converter_config: DataConverterConfig, @@ -426,17 +434,36 @@ def _extract_camera( camera_dict: Dict[str, Union[str, bytes]] = {} sensor_root = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" + log_cam_infos = {camera.token: camera for camera in log_db.log.cameras} + for camera_type, camera_channel in NUPLAN_CAMERA_TYPES.items(): camera_data: Optional[Union[str, bytes]] = None + c2e: Optional[List[float]] = None image_class = list(get_images_from_lidar_tokens(source_log_path, [lidar_pc.token], [str(camera_channel.value)])) if len(image_class) != 0: - filename_jpg = sensor_root / image_class[0].filename_jpg + image = image_class[0] + filename_jpg = sensor_root / image.filename_jpg if filename_jpg.exists(): + + # Code taken from MTGS + # https://github.com/OpenDriveLab/MTGS/blob/main/nuplan_scripts/utils/nuplan_utils_custom.py#L117 + + timestamp = image.timestamp + NUPLAN_ROLLING_SHUTTER_S.time_us + img_ego_pose: EgoPose = ( + log_db.log._session.query(EgoPose).order_by(func.abs(EgoPose.timestamp - timestamp)).first() + ) + img_e2g = img_ego_pose.trans_matrix + g2e = lidar_pc.ego_pose.trans_matrix_inv + img_e2e = g2e @ img_e2g + cam_info = log_cam_infos[image.camera_token] + c2img_e = cam_info.trans_matrix + c2e = img_e2e @ c2img_e + if data_converter_config.camera_store_option == "path": - camera_data = str(filename_jpg) + camera_data = str(filename_jpg), c2e.flatten().tolist() elif data_converter_config.camera_store_option == "binary": with open(filename_jpg, "rb") as f: - camera_data = f.read() + camera_data = f.read(), c2e camera_dict[camera_type] = camera_data diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 7d5fcd7b..07b8be63 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,8 +15,8 @@ defaults: - default_dataset_paths - _self_ - datasets: - # - nuplan_private_dataset - - carla_dataset + - nuplan_private_dataset + # - carla_dataset force_log_conversion: True force_map_conversion: False diff --git a/docs/schema.md b/docs/schema.md new file mode 100644 index 00000000..2669585a --- /dev/null +++ b/docs/schema.md @@ -0,0 +1,14 @@ + + + +token +timestamp +detections_state +detections_velocity +detections_token +detections_type +ego_states +traffic_light_ids +traffic_light_types +scenario_tag +route_lane_group_ids diff --git a/notebooks/scene_rendering.ipynb b/notebooks/scene_rendering.ipynb index 111590a4..e575365e 100644 --- a/notebooks/scene_rendering.ipynb +++ b/notebooks/scene_rendering.ipynb @@ -11,9 +11,8 @@ "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed\n", "\n", - "\n" + "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" ] }, { @@ -31,12 +30,12 @@ " memory_info = process.memory_info()\n", " print(f\"Memory usage: {memory_info.rss / 1024 ** 2:.2f} MB\")\n", "\n", - "\n", "split = \"nuplan_mini_val\"\n", + "# split = \"carla\"\n", "\n", "# log_names = [\"2021.06.07.12.54.00_veh-35_01843_02314\"]\n", "scene_tokens = None\n", - "scene_tokens = [\"2283aea39bc1505e\"]\n", + "# scene_tokens = [\"2283aea39bc1505e\"]\n", "log_names = None\n", "\n", "scene_filter = SceneFilter(\n", @@ -61,12 +60,82 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.common.visualization.matplotlib.plots import render_scene_animation\n", + "from typing import Tuple\n", + "from d123.common.datatypes.detection.detection import BoxDetection\n", + "from d123.common.datatypes.detection.detection_types import DYNAMIC_DETECTION_TYPES, STATIC_DETECTION_TYPES\n", + "from d123.common.geometry.base import StateSE2\n", + "from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", + "from d123.common.geometry.vector import Vector2D\n", + "from d123.common.visualization.matplotlib.observation import (\n", + " add_box_detections_to_ax,\n", + " add_default_map_on_ax,\n", + " add_ego_vehicle_to_ax,\n", + " add_traffic_lights_to_ax,\n", + ")\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", + "\n", + " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " traffic_light_detections = scene.get_traffic_light_detections_at_iteration(iteration)\n", + " map_api = scene.map_api\n", + "\n", + " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", + " add_default_map_on_ax(ax, map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_traffic_lights_to_ax(ax, traffic_light_detections, map_api)\n", + "\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + "\n", + " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", + " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", + "\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")\n", + " return ax\n", + "\n", + "\n", + "iteration = 20\n", + "scene = scenes[20]\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "_plot_scene_on_ax(ax, scene, iteration, 80)\n", + "\n", + "box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "\n", + "\n", + "def _get_dxy(box_detection: BoxDetection) -> Tuple[float, float]:\n", + " \"\"\"Get the change in x and y coordinates from the box detection.\"\"\"\n", + " center = box_detection.center.state_se2\n", + " # endpoint_x = translate_along_yaw(center, Vector2D(abs(box_detection.velocity.x), 0.0))\n", + " # endpoint_y = translate_along_yaw(center, Vector2D(0.0, abs(box_detection.velocity.y)))\n", + " # print(box_detection.velocity.x, box_detection.velocity.y)\n", + "\n", + " endpoint_x = StateSE2(center.x + box_detection.velocity.x, center.y, center.yaw)\n", + " endpoint_y = StateSE2(center.x, center.y + box_detection.velocity.y, center.yaw)\n", "\n", - "scene = scenes[0]\n", + " return endpoint_x, endpoint_y\n", "\n", "\n", - "render_scene_animation(scene, \"\", fps=10, end_idx=None, step=1, dpi=100, format=\"gif\")" + "for box_detection in box_detections:\n", + " if box_detection.metadata.detection_type in STATIC_DETECTION_TYPES:\n", + " continue\n", + " endpoint_x, endpoint_y = _get_dxy(box_detection)\n", + " ax.annotate(\n", + " \"\",\n", + " xytext=(box_detection.center.state_se2.point_2d.x, box_detection.center.state_se2.point_2d.y),\n", + " xy=(endpoint_x.x, endpoint_x.y),\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )\n", + " ax.annotate(\n", + " \"\",\n", + " xytext=(box_detection.center.state_se2.point_2d.x, box_detection.center.state_se2.point_2d.y),\n", + " xy=(endpoint_y.x, endpoint_y.y),\n", + " arrowprops=dict(arrowstyle=\"->\"),\n", + " )" ] }, { diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index d0511676..38338e68 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -7,13 +7,10 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.common.visualization.viser.server import ViserVisualizationServer\n", - "\n", "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "\n" + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" ] }, { @@ -22,14 +19,6 @@ "id": "1", "metadata": {}, "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], "source": [ "# split = \"nuplan_private_test\"\n", "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", @@ -38,16 +27,8 @@ "# log_names = [\"_Rep0_routes_123d_route1_07_28_20_46_56\"]\n", "\n", "\n", - "splits = [\"carla\", \"nuplan_private_test\"]\n", - "log_names = [\n", - " \"_Rep0_routes_123d_route0_07_28_20_36_14\",\n", - " \"_Rep0_routes_123d_route1_07_28_20_46_56\",\n", - " \"2021.07.01.21.22.09_veh-14_00016_00656\",\n", - " \"2021.07.20.17.06.09_veh-38_04743_05687\",\n", - " \"2021.09.08.14.01.10_veh-39_01898_02179\",\n", - " \"2021.09.29.17.35.58_veh-44_00066_00432\",\n", - " \"2021.10.01.01.03.27_veh-50_01313_01390\",\n", - "]\n", + "splits = [\"nuplan_private_test\"]\n", + "log_names = [\"2021.07.01.21.22.09_veh-14_00016_00656\"]\n", "\n", "\n", "scene_tokens = None\n", @@ -64,19 +45,40 @@ "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", "# worker = RayDistributed()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)" + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" ] }, { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "2", "metadata": {}, "outputs": [], "source": [ + "\n", + "from d123.common.visualization.viser.server import ViserVisualizationServer\n", + "\n", "\n", "visualization_server = ViserVisualizationServer(scenes)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { From 9c9f08f5be160a9618dd8040b13e476b05f204b5 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 1 Aug 2025 20:45:26 +0200 Subject: [PATCH 10/42] Fix several carla specific errors (coordinate system, camera paramters, visualization) --- d123/common/visualization/viser/server.py | 3 -- d123/common/visualization/viser/utils.py | 42 +++++++++++-------- .../carla/carla_data_converter.py | 35 +++++++++++----- .../dataset_specific/carla/load_sensor.py | 6 +-- .../default_dataset_conversion.yaml | 4 +- notebooks/test_gym.ipynb | 18 +++----- notebooks/viz/camera_matplotlib.ipynb | 25 +++++++++++ notebooks/viz/viser_testing_v2_scene.ipynb | 10 ++--- 8 files changed, 86 insertions(+), 57 deletions(-) create mode 100644 notebooks/viz/camera_matplotlib.ipynb diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index e14214f3..6a401593 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -162,9 +162,6 @@ def _(_) -> None: camera_frustum_handles[camera_type].wxyz = camera_quaternion.q camera_frustum_handles[camera_type].image = camera.image - # if BACK_CAMERA_TYPE in scene.available_camera_types: - # gui_back_cam.image = scene.get_camera_at_iteration(gui_timestep.value, BACK_CAMERA_TYPE).image - if LIDAR_AVAILABLE: gui_lidar.points = get_lidar_points(scene, gui_timestep.value) diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 5056ae5f..f294ba97 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -131,6 +131,8 @@ def get_trimesh_from_boundaries( def _interpolate_polyline(polyline_3d: Polyline3D, num_samples: int = 20) -> npt.NDArray[np.float64]: + if num_samples < 2: + num_samples = 2 distances = np.linspace(0, polyline_3d.length, num=num_samples, endpoint=True, dtype=np.float64) return polyline_3d.interpolate(distances) @@ -183,8 +185,22 @@ def get_camera_values( ego_transform[:3, :3] = get_rotation_matrix(rear_axle) ego_transform[:3, 3] = rear_axle.point_3d.array + CARLA_DEBUG = True + if CARLA_DEBUG: + rotation = StateSE3( + 0.0, + 0.0, + 0.0, + -np.deg2rad(90.0), + np.deg2rad(0.0), + -np.deg2rad(90.0), + ) + camera_to_ego[:3, :3] = get_rotation_matrix(rotation) + camera_transform = ego_transform @ camera_to_ego + else: + camera_transform = ego_transform @ camera_to_ego + # Camera transformation in ego frame - camera_transform = ego_transform @ camera_to_ego camera_position = Point3D(*camera_transform[:3, 3]) camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) @@ -214,27 +230,17 @@ def euler_to_quaternion_scipy(roll: float, pitch: float, yaw: float) -> npt.NDAr def get_lidar_points(scene: AbstractScene, iteration: int) -> npt.NDArray[np.float32]: - from d123.common.geometry.transform.se3 import translate_points_3d_along_z + pass initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) current_ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) lidar = scene.get_lidar_at_iteration(iteration) - if scene.log_metadata.dataset == "nuplan": - # NOTE: nuPlan uses the rear axle as origin. - origin = current_ego_vehicle_state.rear_axle_se3 - points = lidar.xyz - points = convert_relative_to_absolute_points_3d_array(origin, points) - points = points - initial_ego_vehicle_state.center_se3.point_3d.array - elif scene.log_metadata.dataset == "carla": - # NOTE: CARLA uses the center of the vehicle as origin. - origin = current_ego_vehicle_state.center_se3 - points = translate_points_3d_along_z( - origin, lidar.xyz, -current_ego_vehicle_state.vehicle_parameters.height / 2 - ) - points = convert_relative_to_absolute_points_3d_array(origin, points) - points = points - initial_ego_vehicle_state.center_se3.point_3d.array - else: - raise ValueError(f"Unsupported dataset: {scene.log_metadata.dataset}") + # if scene.log_metadata.dataset == "nuplan": + # NOTE: nuPlan uses the rear axle as origin. + origin = current_ego_vehicle_state.rear_axle_se3 + points = lidar.xyz + points = convert_relative_to_absolute_points_3d_array(origin, points) + points = points - initial_ego_vehicle_state.center_se3.point_3d.array return points diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 2cd2bfbf..fc0fddbe 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -165,7 +165,8 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> log_file_path.parent.mkdir(parents=True, exist_ok=True) bounding_box_paths = sorted([bb_path for bb_path in (log_path / "boxes").iterdir()]) - map_name = _load_json_gz(bounding_box_paths[0])["location"] + first_log_dict = _load_json_gz(bounding_box_paths[0]) + map_name = first_log_dict["location"] map_api = get_map_api_from_names("carla", map_name) schema_column_list = [ @@ -192,13 +193,16 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> for camera_type in CARLA_CAMERA_TYPES: if data_converter_config.camera_store_option == "path": schema_column_list.append((camera_type.serialize(), pa.string())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 16)) + ) elif data_converter_config.camera_store_option == "binary": raise NotImplementedError("Binary camera storage is not implemented.") recording_schema = pa.schema(schema_column_list) metadata = _get_metadata(map_name, str(log_path.stem)) vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() - camera_metadata = get_carla_camera_metadata() + camera_metadata = get_carla_camera_metadata(first_log_dict) recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), @@ -232,18 +236,21 @@ def _get_metadata(location: str, log_name: str) -> LogMetadata: ) -def get_carla_camera_metadata() -> Dict[str, CameraMetadata]: +def get_carla_camera_metadata(first_log_dict: Dict[str, Any]) -> Dict[str, CameraMetadata]: # FIXME: This is a placeholder function to return camera metadata. + + intrinsic = np.array( + first_log_dict[f"{CameraType.CAM_F0.serialize()}_intrinsics"], + dtype=np.float64, + ) camera_metadata = { CameraType.CAM_F0.serialize(): CameraMetadata( camera_type=CameraType.CAM_F0, width=1024, height=512, - intrinsic=None, - distortion=None, - translation=None, - rotation=None, + intrinsic=intrinsic, + distortion=np.zeros((5,), dtype=np.float64), ) } return camera_metadata @@ -293,10 +300,11 @@ def _write_recording_table( row_data["lidar"] = [_extract_lidar(log_name, sample_name, data_converter_config)] if data_converter_config.camera_store_option is not None: - camera_data_dict = _extract_cameras(log_name, sample_name, data_converter_config) + camera_data_dict = _extract_cameras(data, log_name, sample_name, data_converter_config) for camera_type, camera_data in camera_data_dict.items(): if camera_data is not None: - row_data[camera_type.serialize()] = [camera_data] + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] else: row_data[camera_type.serialize()] = [None] @@ -382,14 +390,19 @@ def _extract_route_lane_group_ids(route: List[List[float]], map_api: AbstractMap def _extract_cameras( - log_name: str, sample_name: str, data_converter_config: DataConverterConfig + data: Dict[str, Any], log_name: str, sample_name: str, data_converter_config: DataConverterConfig ) -> Dict[CameraType, Optional[str]]: camera_dict: Dict[str, Union[str, bytes]] = {} for camera_type in CARLA_CAMERA_TYPES: camera_full_path = CARLA_DATA_ROOT / "sensor_blobs" / log_name / camera_type.name / f"{sample_name}.jpg" if camera_full_path.exists(): if data_converter_config.camera_store_option == "path": - camera_dict[camera_type] = f"{log_name}/{camera_type.name}/{sample_name}.jpg" + path = f"{log_name}/{camera_type.name}/{sample_name}.jpg" + extrinsics = data.get(f"{camera_type.serialize()}_transform", None) + camera_dict[camera_type] = path, ( + np.array(extrinsics, dtype=np.float64).flatten() if extrinsics is not None else None + ) + elif data_converter_config.camera_store_option == "binary": raise NotImplementedError("Binary camera storage is not implemented.") else: diff --git a/d123/dataset/dataset_specific/carla/load_sensor.py b/d123/dataset/dataset_specific/carla/load_sensor.py index c4feeb8b..63f5ee35 100644 --- a/d123/dataset/dataset_specific/carla/load_sensor.py +++ b/d123/dataset/dataset_specific/carla/load_sensor.py @@ -7,8 +7,4 @@ def load_carla_lidar_from_path(filepath: Path) -> LiDAR: assert filepath.exists(), f"LiDAR file not found: {filepath}" - points = np.load(filepath).T - - points[1] = -points[1] # FIXME - - return LiDAR(point_cloud=points) + return LiDAR(point_cloud=np.load(filepath)) diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 07b8be63..7d5fcd7b 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,8 +15,8 @@ defaults: - default_dataset_paths - _self_ - datasets: - - nuplan_private_dataset - # - carla_dataset + # - nuplan_private_dataset + - carla_dataset force_log_conversion: True force_map_conversion: False diff --git a/notebooks/test_gym.ipynb b/notebooks/test_gym.ipynb index 718dd917..5a5f40f2 100644 --- a/notebooks/test_gym.ipynb +++ b/notebooks/test_gym.ipynb @@ -68,14 +68,6 @@ "id": "3", "metadata": {}, "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], "source": [ "from pathlib import Path\n", "from typing import Optional, Tuple\n", @@ -158,7 +150,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -218,7 +210,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -228,7 +220,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -257,7 +249,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -278,7 +270,7 @@ { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "8", "metadata": {}, "outputs": [], "source": [] diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb new file mode 100644 index 00000000..7c13e7f8 --- /dev/null +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -0,0 +1,25 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 38338e68..d04d3665 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -23,12 +23,12 @@ "# split = \"nuplan_private_test\"\n", "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", - "# split = \"carla\"\n", - "# log_names = [\"_Rep0_routes_123d_route1_07_28_20_46_56\"]\n", + "splits = [\"carla\"]\n", + "# log_names = [\"_Rep0_routes_123d_route0_08_01_17_32_01\", \"_Rep0_routes_123d_route1_08_01_17_33_57\"]\n", + "log_names = None\n", "\n", - "\n", - "splits = [\"nuplan_private_test\"]\n", - "log_names = [\"2021.07.01.21.22.09_veh-14_00016_00656\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "# log_names = [\"2021.07.01.21.22.09_veh-14_00016_00656\"]\n", "\n", "\n", "scene_tokens = None\n", From 689ae2c7059402fbd5a7e7bc775f3234078e8223 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 1 Aug 2025 22:38:34 +0200 Subject: [PATCH 11/42] Add simple visualization scripts to render bounding boxes on images (mostly from navsim) --- d123/common/datatypes/detection/detection.py | 4 + .../geometry/bounding_box/bounding_box.py | 4 + d123/common/geometry/transform/se3.py | 32 +- .../common/visualization/matplotlib/camera.py | 346 ++++++++++++++++++ d123/common/visualization/matplotlib/lidar.py | 70 ++++ notebooks/viz/camera_matplotlib.ipynb | 255 +++++++++++++ 6 files changed, 710 insertions(+), 1 deletion(-) create mode 100644 d123/common/visualization/matplotlib/camera.py create mode 100644 d123/common/visualization/matplotlib/lidar.py diff --git a/d123/common/datatypes/detection/detection.py b/d123/common/datatypes/detection/detection.py index 2f6f0a71..a836fd4f 100644 --- a/d123/common/datatypes/detection/detection.py +++ b/d123/common/datatypes/detection/detection.py @@ -57,6 +57,10 @@ def shapely_polygon(self) -> shapely.geometry.Polygon: def center(self) -> StateSE3: return self.bounding_box_se3.center + @property + def center_se3(self) -> StateSE3: + return self.bounding_box_se3.center_se3 + @property def bounding_box(self) -> BoundingBoxSE3: return self.bounding_box_se3 diff --git a/d123/common/geometry/bounding_box/bounding_box.py b/d123/common/geometry/bounding_box/bounding_box.py index de7daf93..625d24d9 100644 --- a/d123/common/geometry/bounding_box/bounding_box.py +++ b/d123/common/geometry/bounding_box/bounding_box.py @@ -69,6 +69,10 @@ def bounding_box_se2(self) -> BoundingBoxSE2: width=self.width, ) + @property + def center_se3(self) -> StateSE3: + return self.center + @property def array(self) -> npt.NDArray[np.float64]: array = np.zeros(len(BoundingBoxSE3Index), dtype=np.float64) diff --git a/d123/common/geometry/transform/se3.py b/d123/common/geometry/transform/se3.py index b71c7415..dc09ed83 100644 --- a/d123/common/geometry/transform/se3.py +++ b/d123/common/geometry/transform/se3.py @@ -1,7 +1,7 @@ import numpy as np import numpy.typing as npt -from d123.common.geometry.base import Point3DIndex, StateSE3 +from d123.common.geometry.base import Point3DIndex, StateSE3, StateSE3Index from d123.common.geometry.vector import Vector3D @@ -98,6 +98,36 @@ def convert_relative_to_absolute_points_3d_array( return absolute_points +def convert_absolute_to_relative_se3_array( + origin: StateSE3, se3_array: npt.NDArray[np.float64] +) -> npt.NDArray[np.float64]: + assert se3_array.shape[-1] == len(StateSE3Index) + + # Extract rotation and translation of origin + R_origin = get_rotation_matrix(origin) + t_origin = origin.point_3d.array + + # Prepare output array + rel_se3_array = np.empty_like(se3_array) + + # For each SE3 in the array + for i in range(se3_array.shape[0]): + abs_se3 = se3_array[i] + abs_pos = abs_se3[StateSE3Index.XYZ] + abs_rpy = abs_se3[StateSE3Index.ROLL : StateSE3Index.YAW + 1] + + # Relative position: rotate and translate + rel_pos = R_origin.T @ (abs_pos - t_origin) + + # Relative orientation: subtract origin's rpy + rel_rpy = abs_rpy - np.array([origin.roll, origin.pitch, origin.yaw], dtype=np.float64) + + rel_se3_array[i, StateSE3Index.X : StateSE3Index.Z + 1] = rel_pos + rel_se3_array[i, StateSE3Index.ROLL : StateSE3Index.YAW + 1] = rel_rpy + + return rel_se3_array + + def translate_points_3d_along_z( state_se3: StateSE3, points_3d: npt.NDArray[np.float64], diff --git a/d123/common/visualization/matplotlib/camera.py b/d123/common/visualization/matplotlib/camera.py new file mode 100644 index 00000000..1f3192f0 --- /dev/null +++ b/d123/common/visualization/matplotlib/camera.py @@ -0,0 +1,346 @@ +# from typing import List, Optional, Tuple + +from typing import List, Optional, Tuple + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import numpy.typing as npt + +# from PIL import ImageColor +from pyquaternion import Quaternion + +from d123.common.datatypes.detection.detection import BoxDetectionSE3, BoxDetectionWrapper +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import Camera +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 +from d123.common.geometry.bounding_box.bounding_box_index import BoundingBoxSE3Index, Corners3DIndex +from d123.common.geometry.transform.se3 import convert_absolute_to_relative_se3_array +from d123.common.visualization.color.default import BOX_DETECTION_CONFIG + +# from navsim.common.dataclasses import Annotations, Camera, Lidar +# from navsim.common.enums import BoundingBoxIndex, LidarIndex +# from navsim.planning.scenario_builder.navsim_scenario_utils import tracked_object_types +# from navsim.visualization.config import AGENT_CONFIG +# from navsim.visualization.lidar import filter_lidar_pc, get_lidar_pc_color + + +def add_camera_ax(ax: plt.Axes, camera: Camera) -> plt.Axes: + """ + Adds camera image to matplotlib ax object + :param ax: matplotlib ax object + :param camera: navsim camera dataclass + :return: ax object with image + """ + ax.imshow(camera.image) + return ax + + +# FIXME: +# def add_lidar_to_camera_ax(ax: plt.Axes, camera: Camera, lidar: Lidar) -> plt.Axes: +# """ +# Adds camera image with lidar point cloud on matplotlib ax object +# :param ax: matplotlib ax object +# :param camera: navsim camera dataclass +# :param lidar: navsim lidar dataclass +# :return: ax object with image +# """ + +# image, lidar_pc = camera.image.copy(), lidar.lidar_pc.copy() +# image_height, image_width = image.shape[:2] + +# lidar_pc = filter_lidar_pc(lidar_pc) +# lidar_pc_colors = np.array(get_lidar_pc_color(lidar_pc)) + +# pc_in_cam, pc_in_fov_mask = _transform_pcs_to_images( +# lidar_pc, +# camera.sensor2lidar_rotation, +# camera.sensor2lidar_translation, +# camera.intrinsics, +# img_shape=(image_height, image_width), +# ) + +# for (x, y), color in zip(pc_in_cam[pc_in_fov_mask], lidar_pc_colors[pc_in_fov_mask]): +# color = (int(color[0]), int(color[1]), int(color[2])) +# cv2.circle(image, (int(x), int(y)), 5, color, -1) + +# ax.imshow(image) +# return ax + + +def add_box_detections_to_camera_ax( + ax: plt.Axes, + camera: Camera, + box_detections: BoxDetectionWrapper, + ego_state_se3: EgoStateSE3, +) -> plt.Axes: + + # box_labels = annotations.names + # boxes = _transform_annotations_to_camera( + # annotations.boxes, + # camera.sensor2lidar_rotation, + # camera.sensor2lidar_translation, + # ) + # box_positions, box_dimensions, box_heading = ( + # boxes[:, BoundingBoxIndex.POSITION], + # boxes[:, BoundingBoxIndex.DIMENSION], + # boxes[:, BoundingBoxIndex.HEADING], + # ) + + box_detection_array = np.zeros((len(box_detections.box_detections), len(BoundingBoxSE3Index)), dtype=np.float64) + detection_types = np.array( + [detection.metadata.detection_type for detection in box_detections.box_detections], dtype=object + ) + for idx, box_detection in enumerate(box_detections.box_detections): + assert isinstance( + box_detection, BoxDetectionSE3 + ), f"Box detection must be of type BoxDetectionSE3, got {type(box_detection)}" + box_detection_array[idx] = box_detection.bounding_box_se3.array + + # FIXME + box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] = convert_absolute_to_relative_se3_array( + ego_state_se3.rear_axle_se3, box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] + ) + # box_detection_array[..., BoundingBoxSE3Index.XYZ] -= ego_state_se3.rear_axle_se3.point_3d.array + detection_positions, detection_extents, detection_yaws = _transform_annotations_to_camera( + box_detection_array, camera.extrinsic + ) + + corners_norm = np.stack(np.unravel_index(np.arange(len(Corners3DIndex)), [2] * 3), axis=1) + corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] + corners_norm = corners_norm - np.array([0.5, 0.5, 0.5]) + corners = detection_extents.reshape([-1, 1, 3]) * corners_norm.reshape([1, 8, 3]) + + corners = _rotation_3d_in_axis(corners, detection_yaws, axis=1) + corners += detection_positions.reshape(-1, 1, 3) + + # Then draw project corners to image. + box_corners, corners_pc_in_fov = _transform_points_to_image(corners.reshape(-1, 3), camera.metadata.intrinsic) + box_corners = box_corners.reshape(-1, 8, 2) + corners_pc_in_fov = corners_pc_in_fov.reshape(-1, 8) + valid_corners = corners_pc_in_fov.any(-1) + + box_corners, detection_types = box_corners[valid_corners], detection_types[valid_corners] + image = _plot_rect_3d_on_img(camera.image.copy(), box_corners, detection_types) + + ax.imshow(image) + return ax + + +def _transform_annotations_to_camera( + boxes: npt.NDArray[np.float32], + # sensor2lidar_rotation: npt.NDArray[np.float32], + # sensor2lidar_translation: npt.NDArray[np.float32], + extrinsic: npt.NDArray[np.float64], +) -> npt.NDArray[np.float32]: + """ + Helper function to transform bounding boxes into camera frame + TODO: Refactor + :param boxes: array representation of bounding boxes + :param sensor2lidar_rotation: camera rotation + :param sensor2lidar_translation: camera translation + :return: bounding boxes in camera coordinates + """ + sensor2lidar_rotation = extrinsic[:3, :3] + sensor2lidar_translation = extrinsic[:3, 3] + + locs, rots = ( + boxes[:, BoundingBoxSE3Index.XYZ], + boxes[:, BoundingBoxSE3Index.YAW], + ) + dims_cam = boxes[ + :, [BoundingBoxSE3Index.LENGTH, BoundingBoxSE3Index.HEIGHT, BoundingBoxSE3Index.WIDTH] + ] # l, w, h -> l, h, w + + rots_cam = np.zeros_like(rots) + for idx, rot in enumerate(rots): + rot = Quaternion(axis=[0, 0, 1], radians=rot) + rot = Quaternion(matrix=sensor2lidar_rotation).inverse * rot + rots_cam[idx] = -rot.yaw_pitch_roll[0] + + lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) + lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T + lidar2cam_rt = np.eye(4) + lidar2cam_rt[:3, :3] = lidar2cam_r.T + lidar2cam_rt[3, :3] = -lidar2cam_t + + locs_cam = np.concatenate([locs, np.ones_like(locs)[:, :1]], -1) # -1, 4 + locs_cam = lidar2cam_rt.T @ locs_cam.T + locs_cam = locs_cam.T + locs_cam = locs_cam[:, :-1] + return locs_cam, dims_cam, rots_cam + + +def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np.float32], axis: int = 0): + """ + Rotate 3D points by angles according to axis. + TODO: Refactor + :param points: array of points + :param angles: array of angles + :param axis: axis to perform rotation, defaults to 0 + :raises value: _description_ + :raises ValueError: if axis invalid + :return: rotated points + """ + rot_sin = np.sin(angles) + rot_cos = np.cos(angles) + ones = np.ones_like(rot_cos) + zeros = np.zeros_like(rot_cos) + if axis == 1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, zeros, -rot_sin]), + np.stack([zeros, ones, zeros]), + np.stack([rot_sin, zeros, rot_cos]), + ] + ) + elif axis == 2 or axis == -1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, -rot_sin, zeros]), + np.stack([rot_sin, rot_cos, zeros]), + np.stack([zeros, zeros, ones]), + ] + ) + elif axis == 0: + rot_mat_T = np.stack( + [ + np.stack([zeros, rot_cos, -rot_sin]), + np.stack([zeros, rot_sin, rot_cos]), + np.stack([ones, zeros, zeros]), + ] + ) + else: + raise ValueError(f"axis should in range [0, 1, 2], got {axis}") + return np.einsum("aij,jka->aik", points, rot_mat_T) + + +def _plot_rect_3d_on_img( + image: npt.NDArray[np.float32], + box_corners: npt.NDArray[np.float32], + detection_types: List[DetectionType], + thickness: int = 1, +) -> npt.NDArray[np.uint8]: + """ + Plot the boundary lines of 3D rectangular on 2D images. + TODO: refactor + :param image: The numpy array of image. + :param box_corners: Coordinates of the corners of 3D, shape of [N, 8, 2]. + :param box_labels: labels of boxes for coloring + :param thickness: pixel width of liens, defaults to 3 + :return: image with 3D bounding boxes + """ + line_indices = ( + (0, 1), + (0, 3), + (0, 4), + (1, 2), + (1, 5), + (3, 2), + (3, 7), + (4, 5), + (4, 7), + (2, 6), + (5, 6), + (6, 7), + ) + for i in range(len(box_corners)): + color = BOX_DETECTION_CONFIG[detection_types[i]].fill_color.rgb + corners = box_corners[i].astype(np.int64) + for start, end in line_indices: + cv2.line( + image, + (corners[start, 0], corners[start, 1]), + (corners[end, 0], corners[end, 1]), + color, + thickness, + cv2.LINE_AA, + ) + return image.astype(np.uint8) + + +def _transform_points_to_image( + points: npt.NDArray[np.float32], + intrinsic: npt.NDArray[np.float32], + image_shape: Optional[Tuple[int, int]] = None, + eps: float = 1e-3, +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: + """ + Transforms points in camera frame to image pixel coordinates + TODO: refactor + :param points: points in camera frame + :param intrinsic: camera intrinsics + :param image_shape: shape of image in pixel + :param eps: lower threshold of points, defaults to 1e-3 + :return: points in pixel coordinates, mask of values in frame + """ + points = points[:, :3] + + viewpad = np.eye(4) + viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic + + pc_img = np.concatenate([points, np.ones_like(points)[:, :1]], -1) + pc_img = viewpad @ pc_img.T + pc_img = pc_img.T + + cur_pc_in_fov = pc_img[:, 2] > eps + pc_img = pc_img[..., 0:2] / np.maximum(pc_img[..., 2:3], np.ones_like(pc_img[..., 2:3]) * eps) + if image_shape is not None: + img_h, img_w = image_shape + cur_pc_in_fov = ( + cur_pc_in_fov + & (pc_img[:, 0] < (img_w - 1)) + & (pc_img[:, 0] > 0) + & (pc_img[:, 1] < (img_h - 1)) + & (pc_img[:, 1] > 0) + ) + return pc_img, cur_pc_in_fov + + +# def _transform_pcs_to_images( +# lidar_pc: npt.NDArray[np.float32], +# sensor2lidar_rotation: npt.NDArray[np.float32], +# sensor2lidar_translation: npt.NDArray[np.float32], +# intrinsic: npt.NDArray[np.float32], +# img_shape: Optional[Tuple[int, int]] = None, +# eps: float = 1e-3, +# ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: +# """ +# Transforms points in camera frame to image pixel coordinates +# TODO: refactor +# :param lidar_pc: lidar point cloud +# :param sensor2lidar_rotation: camera rotation +# :param sensor2lidar_translation: camera translation +# :param intrinsic: camera intrinsics +# :param img_shape: image shape in pixels, defaults to None +# :param eps: threshold for lidar pc height, defaults to 1e-3 +# :return: lidar pc in pixel coordinates, mask of values in frame +# """ +# pc_xyz = lidar_pc[LidarIndex.POSITION, :].T + +# lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) +# lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T +# lidar2cam_rt = np.eye(4) +# lidar2cam_rt[:3, :3] = lidar2cam_r.T +# lidar2cam_rt[3, :3] = -lidar2cam_t + +# viewpad = np.eye(4) +# viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic +# lidar2img_rt = viewpad @ lidar2cam_rt.T + +# cur_pc_xyz = np.concatenate([pc_xyz, np.ones_like(pc_xyz)[:, :1]], -1) +# cur_pc_cam = lidar2img_rt @ cur_pc_xyz.T +# cur_pc_cam = cur_pc_cam.T +# cur_pc_in_fov = cur_pc_cam[:, 2] > eps +# cur_pc_cam = cur_pc_cam[..., 0:2] / np.maximum(cur_pc_cam[..., 2:3], np.ones_like(cur_pc_cam[..., 2:3]) * eps) + +# if img_shape is not None: +# img_h, img_w = img_shape +# cur_pc_in_fov = ( +# cur_pc_in_fov +# & (cur_pc_cam[:, 0] < (img_w - 1)) +# & (cur_pc_cam[:, 0] > 0) +# & (cur_pc_cam[:, 1] < (img_h - 1)) +# & (cur_pc_cam[:, 1] > 0) +# ) +# return cur_pc_cam, cur_pc_in_fov diff --git a/d123/common/visualization/matplotlib/lidar.py b/d123/common/visualization/matplotlib/lidar.py new file mode 100644 index 00000000..29ffeea8 --- /dev/null +++ b/d123/common/visualization/matplotlib/lidar.py @@ -0,0 +1,70 @@ +# TODO: implement +# from typing import Any, List + +# import matplotlib +# import matplotlib.pyplot as plt +# import numpy as np +# import numpy.typing as npt + +# from navsim.common.enums import LidarIndex +# from navsim.visualization.config import LIDAR_CONFIG + + +# def filter_lidar_pc(lidar_pc: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]: +# """ +# Filter lidar point cloud according to global configuration +# :param lidar_pc: numpy array of shape (6,n) +# :return: filtered point cloud +# """ + +# pc = lidar_pc.T +# mask = ( +# np.ones((len(pc)), dtype=bool) +# & (pc[:, LidarIndex.X] > LIDAR_CONFIG["x_lim"][0]) +# & (pc[:, LidarIndex.X] < LIDAR_CONFIG["x_lim"][1]) +# & (pc[:, LidarIndex.Y] > LIDAR_CONFIG["y_lim"][0]) +# & (pc[:, LidarIndex.Y] < LIDAR_CONFIG["y_lim"][1]) +# & (pc[:, LidarIndex.Z] > LIDAR_CONFIG["z_lim"][0]) +# & (pc[:, LidarIndex.Z] < LIDAR_CONFIG["z_lim"][1]) +# ) +# pc = pc[mask] +# return pc.T + + +# def get_lidar_pc_color(lidar_pc: npt.NDArray[np.float32], as_hex: bool = False) -> List[Any]: +# """ +# Compute color map of lidar point cloud according to global configuration +# :param lidar_pc: numpy array of shape (6,n) +# :param as_hex: whether to return hex values, defaults to False +# :return: list of RGB or hex values +# """ + +# pc = lidar_pc.T +# if LIDAR_CONFIG["color_element"] == "none": +# colors_rgb = np.zeros((len(pc), 3), dtype=np.uin8) + +# else: +# if LIDAR_CONFIG["color_element"] == "distance": +# color_intensities = np.linalg.norm(pc[:, LidarIndex.POSITION], axis=-1) +# else: +# color_element_map = { +# "x": LidarIndex.X, +# "y": LidarIndex.Y, +# "z": LidarIndex.Z, +# "intensity": LidarIndex.INTENSITY, +# "ring": LidarIndex.RING, +# "id": LidarIndex.ID, +# } +# color_intensities = pc[:, color_element_map[LIDAR_CONFIG["color_element"]]] + +# min, max = color_intensities.min(), color_intensities.max() +# norm_intensities = [(value - min) / (max - min) for value in color_intensities] +# colormap = plt.get_cmap("viridis") +# colors_rgb = np.array([colormap(value) for value in norm_intensities]) +# colors_rgb = (colors_rgb[:, :3] * 255).astype(np.uint8) + +# assert len(colors_rgb) == len(pc) +# if as_hex: +# return [matplotlib.colors.to_hex(tuple(c / 255.0 for c in rgb)) for rgb in colors_rgb] + +# return [tuple(value) for value in colors_rgb] diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index 7c13e7f8..c679c3ef 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -6,6 +6,253 @@ "id": "0", "metadata": {}, "outputs": [], + "source": [ + "from typing import Tuple\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "\n", + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "from typing import Dict\n", + "from d123.common.datatypes.sensor.camera import CameraType\n", + "from d123.common.visualization.matplotlib.camera import add_camera_ax\n", + "from d123.common.visualization.matplotlib.camera import add_box_detections_to_camera_ax" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "splits = [\"carla\"]\n", + "log_names = None\n", + "# splits = [\"nuplan_private_test\"]\n", + "# log_names = [\"2021.09.08.14.01.10_veh-39_01898_02179\"]\n", + "\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=20.1,\n", + " history_s=1.0,\n", + " timestamp_threshold_s=20.1,\n", + " shuffle=True,\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "\n", + "def plot_cameras_frame(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 2\n", + " fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L2: ax[1, 0],\n", + " CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R2: ax[1, 2],\n", + " }\n", + "\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_camera_ax(camera_ax, camera)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + " return fig, ax\n", + "\n", + "\n", + "# plot_cameras_frame(scenes[0], iteration=0)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + "# CameraType.CAM_L0: ax[0, 0],\n", + "# CameraType.CAM_F0: ax[0, 1],\n", + "# CameraType.CAM_R0: ax[0, 2],\n", + "# CameraType.CAM_L2: ax[1, 0],\n", + "# CameraType.CAM_B0: ax[1, 1],\n", + "# CameraType.CAM_R2: ax[1, 2],\n", + "# }\n", + "\n", + "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + "# for camera_type, camera_ax in camera_ax_assignment.items():\n", + "# assert camera_type in scene.available_camera_types\n", + "# camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + "# if camera is not None:\n", + "# add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + "# return ax\n", + "\n", + "\n", + "# def plot_cameras_frame(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "# \"\"\"\n", + "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + "# :param scene: navsim scene dataclass\n", + "# :param frame_idx: index of selected frame\n", + "# :return: figure and ax object of matplotlib\n", + "# \"\"\"\n", + "# scale = 2\n", + "# fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", + "# _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + "# fig.tight_layout()\n", + "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "# plot_cameras_frame(scenes[4], iteration=10)\n", + "\n", + "\n", + "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", + "\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + " return ax\n", + "\n", + "\n", + "def plot_cameras_frame(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 2\n", + " fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", + " _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "plot_cameras_frame(scenes[2], iteration=10)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from typing import Optional, Union\n", + "\n", + "from matplotlib import animation\n", + "from tqdm import tqdm\n", + "\n", + "\n", + "def render_scene_animation(\n", + " scene: AbstractScene,\n", + " output_path: Union[str, Path],\n", + " start_idx: int = 0,\n", + " end_idx: Optional[int] = None,\n", + " step: int = 10,\n", + " fps: float = 20.0,\n", + " dpi: int = 300,\n", + " format: str = \"mp4\",\n", + ") -> None:\n", + " assert format in [\"mp4\", \"gif\"], \"Format must be either 'mp4' or 'gif'.\"\n", + " output_path = Path(output_path)\n", + " output_path.mkdir(parents=True, exist_ok=True)\n", + "\n", + " scene.open()\n", + "\n", + " if end_idx is None:\n", + " end_idx = scene.get_number_of_iterations()\n", + " end_idx = min(end_idx, scene.get_number_of_iterations())\n", + "\n", + " scale = 1\n", + " fig, ax = plt.subplots(1, 1, figsize=(scale * 16, scale * 9))\n", + "\n", + " def update(i):\n", + " ax.clear()\n", + " # for a in ax.flatten():\n", + " # a.clear()\n", + " _add_camera_with_detections(ax, scene, i)\n", + " plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)\n", + " pbar.update(1)\n", + "\n", + " frames = list(range(start_idx, end_idx, step))\n", + " pbar = tqdm(total=len(frames), desc=f\"Rendering {scene.log_name} as {format}\")\n", + " ani = animation.FuncAnimation(fig, update, frames=frames, repeat=False)\n", + "\n", + " ani.save(output_path / f\"{scene.log_name}_{scene.token}.{format}\", writer=\"ffmpeg\", fps=fps, dpi=dpi)\n", + " plt.close(fig)\n", + " scene.close()\n", + "\n", + "\n", + "for i in range(10):\n", + " render_scene_animation(scenes[i], output_path=\"output\", start_idx=0, end_idx=None, step=1, fps=20, dpi=300, format=\"mp4\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], "source": [] } ], @@ -16,7 +263,15 @@ "name": "python3" }, "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", "version": "3.12.11" } }, From 2892ced1f9d2ada23e9f26339b0d1bc3fe043738 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sat, 2 Aug 2025 14:12:11 +0200 Subject: [PATCH 12/42] Add option for camera filter in `SceneFilter` --- d123/dataset/scene/scene_builder.py | 5 +++++ d123/dataset/scene/scene_filter.py | 19 +++++++++++++++++++ notebooks/viz/camera_matplotlib.ipynb | 17 ++++++++--------- 3 files changed, 32 insertions(+), 9 deletions(-) diff --git a/d123/dataset/scene/scene_builder.py b/d123/dataset/scene/scene_builder.py index 0ea3c2c3..4ec1ec6c 100644 --- a/d123/dataset/scene/scene_builder.py +++ b/d123/dataset/scene/scene_builder.py @@ -106,6 +106,11 @@ def _get_scene_extraction_info(log_path: Union[str, Path], filter: SceneFilter) if filter.map_names is not None and log_metadata.map_name not in filter.map_names: return scene_extraction_infos + # 2. Filter by camera type if specified in filter + if filter.camera_types is not None: + if not all(camera_type.serialize() in recording_table.column_names for camera_type in filter.camera_types): + return scene_extraction_infos + start_idx = int(filter.history_s / log_metadata.timestep_seconds) # if filter.history_s is not None else 0 end_idx = ( len(recording_table) - int(filter.duration_s / log_metadata.timestep_seconds) diff --git a/d123/dataset/scene/scene_filter.py b/d123/dataset/scene/scene_filter.py index 69223c1b..0a58eba6 100644 --- a/d123/dataset/scene/scene_filter.py +++ b/d123/dataset/scene/scene_filter.py @@ -1,6 +1,8 @@ from dataclasses import dataclass from typing import List, Optional +from d123.common.datatypes.sensor.camera import CameraType + # TODO: Add more filter options (e.g. scene tags, ego movement, or whatever appropriate) @@ -21,5 +23,22 @@ class SceneFilter: duration_s: Optional[float] = 10.0 history_s: Optional[float] = 3.0 + camera_types: Optional[List[CameraType]] = None + max_num_scenes: Optional[int] = None shuffle: bool = False + + def __post_init__(self): + if self.camera_types is not None: + assert isinstance(self.camera_types, list), "camera_types must be a list of CameraType" + camera_types = [] + for camera_type in self.camera_types: + if isinstance(camera_type, str): + camera_type = CameraType.deserialize[camera_type] + camera_types.append(camera_type) + elif isinstance(camera_type, int): + camera_type = CameraType(camera_type) + camera_types.append(camera_type) + else: + raise ValueError(f"Invalid camera type: {camera_type}") + self.camera_types = camera_types diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index c679c3ef..8866ab6d 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -34,9 +34,9 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "splits = [\"carla\"]\n", + "# splits = [\"carla\"]\n", "log_names = None\n", - "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"nuplan_private_test\"]\n", "# log_names = [\"2021.09.08.14.01.10_veh-39_01898_02179\"]\n", "\n", "scene_tokens = None\n", @@ -49,6 +49,7 @@ " history_s=1.0,\n", " timestamp_threshold_s=20.1,\n", " shuffle=True,\n", + " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", @@ -97,9 +98,7 @@ " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", " return fig, ax\n", - "\n", - "\n", - "# plot_cameras_frame(scenes[0], iteration=0)" + "\n" ] }, { @@ -131,7 +130,7 @@ "# return ax\n", "\n", "\n", - "# def plot_cameras_frame(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "# def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", "# \"\"\"\n", "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", "# :param scene: navsim scene dataclass\n", @@ -145,7 +144,7 @@ "# fig.tight_layout()\n", "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "# plot_cameras_frame(scenes[4], iteration=10)\n", + "# plot_cameras_with_detections(scenes[4], iteration=10)\n", "\n", "\n", "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", @@ -163,7 +162,7 @@ " return ax\n", "\n", "\n", - "def plot_cameras_frame(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", " \"\"\"\n", " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", " :param scene: navsim scene dataclass\n", @@ -178,7 +177,7 @@ " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", "\n", - "plot_cameras_frame(scenes[2], iteration=10)" + "plot_cameras_with_detections(scenes[2], iteration=10)" ] }, { From 24ae260d5a859e27558a961aa489e9d29840dfd9 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sat, 2 Aug 2025 14:54:55 +0200 Subject: [PATCH 13/42] Add a few fixes --- .../visualization/matplotlib/camera copy.py | 330 ++++++++++++++++++ d123/common/visualization/viser/utils.py | 15 +- d123/dataset/arrow/conversion.py | 4 +- d123/dataset/scene/scene_builder.py | 17 +- notebooks/viz/camera_matplotlib.ipynb | 41 ++- notebooks/viz/viser_testing_v2_scene.ipynb | 17 +- 6 files changed, 385 insertions(+), 39 deletions(-) create mode 100644 d123/common/visualization/matplotlib/camera copy.py diff --git a/d123/common/visualization/matplotlib/camera copy.py b/d123/common/visualization/matplotlib/camera copy.py new file mode 100644 index 00000000..ea758366 --- /dev/null +++ b/d123/common/visualization/matplotlib/camera copy.py @@ -0,0 +1,330 @@ +# from typing import List, Optional, Tuple + +from typing import List, Optional, Tuple + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +import numpy.typing as npt + +# from PIL import ImageColor +from pyquaternion import Quaternion + +from d123.common.datatypes.detection.detection import BoxDetectionSE3, BoxDetectionWrapper +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import Camera +from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 +from d123.common.geometry.base import StateSE3 +from d123.common.geometry.bounding_box.bounding_box_index import BoundingBoxSE3Index, Corners3DIndex +from d123.common.geometry.transform.se3 import convert_absolute_to_relative_se3_array, get_rotation_matrix +from d123.common.visualization.color.default import BOX_DETECTION_CONFIG + +# from navsim.common.dataclasses import Annotations, Camera, Lidar +# from navsim.common.enums import BoundingBoxIndex, LidarIndex +# from navsim.planning.scenario_builder.navsim_scenario_utils import tracked_object_types +# from navsim.visualization.config import AGENT_CONFIG +# from navsim.visualization.lidar import filter_lidar_pc, get_lidar_pc_color + + +def add_camera_ax(ax: plt.Axes, camera: Camera) -> plt.Axes: + """ + Adds camera image to matplotlib ax object + :param ax: matplotlib ax object + :param camera: navsim camera dataclass + :return: ax object with image + """ + ax.imshow(camera.image) + return ax + + +# FIXME: +# def add_lidar_to_camera_ax(ax: plt.Axes, camera: Camera, lidar: Lidar) -> plt.Axes: +# """ +# Adds camera image with lidar point cloud on matplotlib ax object +# :param ax: matplotlib ax object +# :param camera: navsim camera dataclass +# :param lidar: navsim lidar dataclass +# :return: ax object with image +# """ + +# image, lidar_pc = camera.image.copy(), lidar.lidar_pc.copy() +# image_height, image_width = image.shape[:2] + +# lidar_pc = filter_lidar_pc(lidar_pc) +# lidar_pc_colors = np.array(get_lidar_pc_color(lidar_pc)) + +# pc_in_cam, pc_in_fov_mask = _transform_pcs_to_images( +# lidar_pc, +# camera.sensor2lidar_rotation, +# camera.sensor2lidar_translation, +# camera.intrinsics, +# img_shape=(image_height, image_width), +# ) + +# for (x, y), color in zip(pc_in_cam[pc_in_fov_mask], lidar_pc_colors[pc_in_fov_mask]): +# color = (int(color[0]), int(color[1]), int(color[2])) +# cv2.circle(image, (int(x), int(y)), 5, color, -1) + +# ax.imshow(image) +# return ax + + +def add_box_detections_to_camera_ax( + ax: plt.Axes, + camera: Camera, + box_detections: BoxDetectionWrapper, + ego_state_se3: EgoStateSE3, +) -> plt.Axes: + + box_detection_array = np.zeros((len(box_detections.box_detections), len(BoundingBoxSE3Index)), dtype=np.float64) + detection_types = np.array( + [detection.metadata.detection_type for detection in box_detections.box_detections], dtype=object + ) + for idx, box_detection in enumerate(box_detections.box_detections): + assert isinstance( + box_detection, BoxDetectionSE3 + ), f"Box detection must be of type BoxDetectionSE3, got {type(box_detection)}" + box_detection_array[idx] = box_detection.bounding_box_se3.array + + box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] = convert_absolute_to_relative_se3_array( + ego_state_se3.rear_axle_se3, box_detection_array[..., BoundingBoxSE3Index.STATE_SE3] + ) + detection_positions, detection_extents, detection_yaws = _transform_annotations_to_camera( + box_detection_array, camera.extrinsic + ) + + corners_norm = np.stack(np.unravel_index(np.arange(len(Corners3DIndex)), [2] * 3), axis=1) + corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] + corners_norm = corners_norm - np.array([0.5, 0.5, 0.5]) + corners = detection_extents.reshape([-1, 1, 3]) * corners_norm.reshape([1, 8, 3]) + + corners = _rotation_3d_in_axis(corners, detection_yaws, axis=1) + corners += detection_positions.reshape(-1, 1, 3) + + # Then draw project corners to image. + box_corners, corners_pc_in_fov = _transform_points_to_image(corners.reshape(-1, 3), camera.metadata.intrinsic) + box_corners = box_corners.reshape(-1, 8, 2) + corners_pc_in_fov = corners_pc_in_fov.reshape(-1, 8) + valid_corners = corners_pc_in_fov.any(-1) + + box_corners, detection_types = box_corners[valid_corners], detection_types[valid_corners] + image = _plot_rect_3d_on_img(camera.image.copy(), box_corners, detection_types) + + ax.imshow(image) + return ax + + +def _transform_annotations_to_camera( + boxes: npt.NDArray[np.float32], extrinsic: npt.NDArray[np.float64] +) -> npt.NDArray[np.float32]: + """ + Helper function to transform bounding boxes into camera frame + TODO: Refactor + :param boxes: array representation of bounding boxes + :param sensor2lidar_rotation: camera rotation + :param sensor2lidar_translation: camera translation + :return: bounding boxes in camera coordinates + """ + sensor2ego_rotation = extrinsic[:3, :3] + sensor2ego_translation = extrinsic[:3, 3] + + locs, rots = ( + boxes[:, BoundingBoxSE3Index.XYZ], + boxes[:, BoundingBoxSE3Index.YAW], + ) + dims_cam = boxes[ + :, [BoundingBoxSE3Index.LENGTH, BoundingBoxSE3Index.HEIGHT, BoundingBoxSE3Index.WIDTH] + ] # l, w, h -> l, h, w + + rots_cam = np.zeros_like(rots) + for idx, state_se3_array in enumerate(boxes[:, BoundingBoxSE3Index.STATE_SE3]): + rot = Quaternion(matrix=get_rotation_matrix(StateSE3.from_array(state_se3_array))) + rot = Quaternion(matrix=sensor2ego_rotation).inverse * rot + rots_cam[idx] = -rot.yaw_pitch_roll[0] + + lidar2cam_r = np.linalg.inv(sensor2ego_rotation) + lidar2cam_t = sensor2ego_translation @ lidar2cam_r.T + lidar2cam_rt = np.eye(4) + lidar2cam_rt[:3, :3] = lidar2cam_r.T + lidar2cam_rt[3, :3] = -lidar2cam_t + + locs_cam = np.concatenate([locs, np.ones_like(locs)[:, :1]], -1) # -1, 4 + locs_cam = lidar2cam_rt.T @ locs_cam.T + locs_cam = locs_cam.T + locs_cam = locs_cam[:, :-1] + return locs_cam, dims_cam, rots_cam + + +def _rotation_3d_in_axis(points: npt.NDArray[np.float32], angles: npt.NDArray[np.float32], axis: int = 0): + """ + Rotate 3D points by angles according to axis. + TODO: Refactor + :param points: array of points + :param angles: array of angles + :param axis: axis to perform rotation, defaults to 0 + :raises value: _description_ + :raises ValueError: if axis invalid + :return: rotated points + """ + rot_sin = np.sin(angles) + rot_cos = np.cos(angles) + ones = np.ones_like(rot_cos) + zeros = np.zeros_like(rot_cos) + if axis == 1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, zeros, -rot_sin]), + np.stack([zeros, ones, zeros]), + np.stack([rot_sin, zeros, rot_cos]), + ] + ) + elif axis == 2 or axis == -1: + rot_mat_T = np.stack( + [ + np.stack([rot_cos, -rot_sin, zeros]), + np.stack([rot_sin, rot_cos, zeros]), + np.stack([zeros, zeros, ones]), + ] + ) + elif axis == 0: + rot_mat_T = np.stack( + [ + np.stack([zeros, rot_cos, -rot_sin]), + np.stack([zeros, rot_sin, rot_cos]), + np.stack([ones, zeros, zeros]), + ] + ) + else: + raise ValueError(f"axis should in range [0, 1, 2], got {axis}") + return np.einsum("aij,jka->aik", points, rot_mat_T) + + +def _plot_rect_3d_on_img( + image: npt.NDArray[np.float32], + box_corners: npt.NDArray[np.float32], + detection_types: List[DetectionType], + thickness: int = 1, +) -> npt.NDArray[np.uint8]: + """ + Plot the boundary lines of 3D rectangular on 2D images. + TODO: refactor + :param image: The numpy array of image. + :param box_corners: Coordinates of the corners of 3D, shape of [N, 8, 2]. + :param box_labels: labels of boxes for coloring + :param thickness: pixel width of liens, defaults to 3 + :return: image with 3D bounding boxes + """ + line_indices = ( + (0, 1), + (0, 3), + (0, 4), + (1, 2), + (1, 5), + (3, 2), + (3, 7), + (4, 5), + (4, 7), + (2, 6), + (5, 6), + (6, 7), + ) + for i in range(len(box_corners)): + color = BOX_DETECTION_CONFIG[detection_types[i]].fill_color.rgb + corners = box_corners[i].astype(np.int64) + for start, end in line_indices: + cv2.line( + image, + (corners[start, 0], corners[start, 1]), + (corners[end, 0], corners[end, 1]), + color, + thickness, + cv2.LINE_AA, + ) + return image.astype(np.uint8) + + +def _transform_points_to_image( + points: npt.NDArray[np.float32], + intrinsic: npt.NDArray[np.float32], + image_shape: Optional[Tuple[int, int]] = None, + eps: float = 1e-3, +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: + """ + Transforms points in camera frame to image pixel coordinates + TODO: refactor + :param points: points in camera frame + :param intrinsic: camera intrinsics + :param image_shape: shape of image in pixel + :param eps: lower threshold of points, defaults to 1e-3 + :return: points in pixel coordinates, mask of values in frame + """ + points = points[:, :3] + + viewpad = np.eye(4) + viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic + + pc_img = np.concatenate([points, np.ones_like(points)[:, :1]], -1) + pc_img = viewpad @ pc_img.T + pc_img = pc_img.T + + cur_pc_in_fov = pc_img[:, 2] > eps + pc_img = pc_img[..., 0:2] / np.maximum(pc_img[..., 2:3], np.ones_like(pc_img[..., 2:3]) * eps) + if image_shape is not None: + img_h, img_w = image_shape + cur_pc_in_fov = ( + cur_pc_in_fov + & (pc_img[:, 0] < (img_w - 1)) + & (pc_img[:, 0] > 0) + & (pc_img[:, 1] < (img_h - 1)) + & (pc_img[:, 1] > 0) + ) + return pc_img, cur_pc_in_fov + + +# def _transform_pcs_to_images( +# lidar_pc: npt.NDArray[np.float32], +# sensor2lidar_rotation: npt.NDArray[np.float32], +# sensor2lidar_translation: npt.NDArray[np.float32], +# intrinsic: npt.NDArray[np.float32], +# img_shape: Optional[Tuple[int, int]] = None, +# eps: float = 1e-3, +# ) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.bool_]]: +# """ +# Transforms points in camera frame to image pixel coordinates +# TODO: refactor +# :param lidar_pc: lidar point cloud +# :param sensor2lidar_rotation: camera rotation +# :param sensor2lidar_translation: camera translation +# :param intrinsic: camera intrinsics +# :param img_shape: image shape in pixels, defaults to None +# :param eps: threshold for lidar pc height, defaults to 1e-3 +# :return: lidar pc in pixel coordinates, mask of values in frame +# """ +# pc_xyz = lidar_pc[LidarIndex.POSITION, :].T + +# lidar2cam_r = np.linalg.inv(sensor2lidar_rotation) +# lidar2cam_t = sensor2lidar_translation @ lidar2cam_r.T +# lidar2cam_rt = np.eye(4) +# lidar2cam_rt[:3, :3] = lidar2cam_r.T +# lidar2cam_rt[3, :3] = -lidar2cam_t + +# viewpad = np.eye(4) +# viewpad[: intrinsic.shape[0], : intrinsic.shape[1]] = intrinsic +# lidar2img_rt = viewpad @ lidar2cam_rt.T + +# cur_pc_xyz = np.concatenate([pc_xyz, np.ones_like(pc_xyz)[:, :1]], -1) +# cur_pc_cam = lidar2img_rt @ cur_pc_xyz.T +# cur_pc_cam = cur_pc_cam.T +# cur_pc_in_fov = cur_pc_cam[:, 2] > eps +# cur_pc_cam = cur_pc_cam[..., 0:2] / np.maximum(cur_pc_cam[..., 2:3], np.ones_like(cur_pc_cam[..., 2:3]) * eps) + +# if img_shape is not None: +# img_h, img_w = img_shape +# cur_pc_in_fov = ( +# cur_pc_in_fov +# & (cur_pc_cam[:, 0] < (img_w - 1)) +# & (cur_pc_cam[:, 0] > 0) +# & (cur_pc_cam[:, 1] < (img_h - 1)) +# & (cur_pc_cam[:, 1] > 0) +# ) +# return cur_pc_cam, cur_pc_in_fov diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index f294ba97..dbca9a02 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -185,20 +185,7 @@ def get_camera_values( ego_transform[:3, :3] = get_rotation_matrix(rear_axle) ego_transform[:3, 3] = rear_axle.point_3d.array - CARLA_DEBUG = True - if CARLA_DEBUG: - rotation = StateSE3( - 0.0, - 0.0, - 0.0, - -np.deg2rad(90.0), - np.deg2rad(0.0), - -np.deg2rad(90.0), - ) - camera_to_ego[:3, :3] = get_rotation_matrix(rotation) - camera_transform = ego_transform @ camera_to_ego - else: - camera_transform = ego_transform @ camera_to_ego + camera_transform = ego_transform @ camera_to_ego # Camera transformation in ego frame diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index ada284c4..9525afe2 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -103,9 +103,9 @@ def get_camera_from_arrow_table( table_data = arrow_table[camera_metadata.camera_type.serialize()][index].as_py() extrinsic = arrow_table[f"{camera_metadata.camera_type.serialize()}_extrinsic"][index].as_py() - extrinsic = np.array(extrinsic).reshape((4, 4)) if extrinsic else np.eye(4) + extrinsic = np.array(extrinsic).reshape((4, 4)) if extrinsic else None - if table_data is None: + if table_data is None or extrinsic is None: return None image: Optional[npt.NDArray[np.uint8]] = None diff --git a/d123/dataset/scene/scene_builder.py b/d123/dataset/scene/scene_builder.py index 4ec1ec6c..c9917de4 100644 --- a/d123/dataset/scene/scene_builder.py +++ b/d123/dataset/scene/scene_builder.py @@ -106,11 +106,6 @@ def _get_scene_extraction_info(log_path: Union[str, Path], filter: SceneFilter) if filter.map_names is not None and log_metadata.map_name not in filter.map_names: return scene_extraction_infos - # 2. Filter by camera type if specified in filter - if filter.camera_types is not None: - if not all(camera_type.serialize() in recording_table.column_names for camera_type in filter.camera_types): - return scene_extraction_infos - start_idx = int(filter.history_s / log_metadata.timestep_seconds) # if filter.history_s is not None else 0 end_idx = ( len(recording_table) - int(filter.duration_s / log_metadata.timestep_seconds) @@ -148,12 +143,22 @@ def _get_scene_extraction_info(log_path: Union[str, Path], filter: SceneFilter) ) if scene_extraction_info is not None: - # TODO: add more options + # Check of timestamp threshold exceeded between previous scene, if specified in filter if filter.timestamp_threshold_s is not None and len(scene_extraction_infos) > 0: iteration_delta = idx - scene_extraction_infos[-1].initial_idx if (iteration_delta * log_metadata.timestep_seconds) < filter.timestamp_threshold_s: continue + # Check if camera data is available for the scene, if specified in filter + # NOTE: We only check camera availability at the initial index of the scene. + if filter.camera_types is not None: + cameras_available = [ + recording_table[camera_type.serialize()][start_idx].as_py() is not None + for camera_type in filter.camera_types + ] + if not all(cameras_available): + continue + scene_extraction_infos.append(scene_extraction_info) del recording_table, log_metadata diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index 8866ab6d..87b7b205 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -35,9 +35,9 @@ "\n", "\n", "# splits = [\"carla\"]\n", - "log_names = None\n", + "# log_names = None\n", "splits = [\"nuplan_private_test\"]\n", - "# log_names = [\"2021.09.08.14.01.10_veh-39_01898_02179\"]\n", + "log_names = None\n", "\n", "scene_tokens = None\n", "\n", @@ -48,7 +48,7 @@ " duration_s=20.1,\n", " history_s=1.0,\n", " timestamp_threshold_s=20.1,\n", - " shuffle=True,\n", + " shuffle=False,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", @@ -65,6 +65,16 @@ "id": "2", "metadata": {}, "outputs": [], + "source": [ + "all([])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], "source": [ "\n", "\n", @@ -104,7 +114,7 @@ { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], "source": [ @@ -176,14 +186,25 @@ " fig.tight_layout()\n", " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", + "plot_cameras_with_detections(scenes[70], iteration=0)\n", "\n", - "plot_cameras_with_detections(scenes[2], iteration=10)" + "\n" ] }, { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "scenes[3].log_name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -234,14 +255,14 @@ " scene.close()\n", "\n", "\n", - "for i in range(10):\n", - " render_scene_animation(scenes[i], output_path=\"output\", start_idx=0, end_idx=None, step=1, fps=20, dpi=300, format=\"mp4\")" + "# for i in range(10):\n", + " # render_scene_animation(scenes[i], output_path=\"output\", start_idx=0, end_idx=None, step=1, fps=20, dpi=300, format=\"mp4\")" ] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "7", "metadata": {}, "outputs": [], "source": [] @@ -249,7 +270,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "8", "metadata": {}, "outputs": [], "source": [] diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index d04d3665..6eb07e46 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -10,7 +10,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" ] }, { @@ -23,14 +24,15 @@ "# split = \"nuplan_private_test\"\n", "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", + "\n", "splits = [\"carla\"]\n", - "# log_names = [\"_Rep0_routes_123d_route0_08_01_17_32_01\", \"_Rep0_routes_123d_route1_08_01_17_33_57\"]\n", - "log_names = None\n", + "# log_names = None\n", "\n", - "# splits = [\"nuplan_private_test\"]\n", - "# log_names = [\"2021.07.01.21.22.09_veh-14_00016_00656\"]\n", "\n", "\n", + "# splits = [\"nuplan_private_test\"]\n", + "log_names = None\n", + "\n", "scene_tokens = None\n", "\n", "scene_filter = SceneFilter(\n", @@ -40,7 +42,8 @@ " duration_s=20.1,\n", " history_s=1.0,\n", " timestamp_threshold_s=20.1,\n", - " shuffle=True,\n", + " shuffle=False,\n", + " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", @@ -61,7 +64,7 @@ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=70)" ] }, { From 47c6b8f904373a5395ac7c81092d045543df93fa Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sat, 2 Aug 2025 16:43:37 +0200 Subject: [PATCH 14/42] Restructure some of the notebooks. --- d123/common/datatypes/sensor/lidar.py | 5 + d123/dataset/arrow/conversion.py | 8 +- notebooks/demo_gym.ipynb | 310 ------------------ notebooks/{ => gym}/test_gym.ipynb | 0 notebooks/{ => gym}/test_simulation_2d.ipynb | 0 .../{ => nuplan}/nuplan_sensor_loading.ipynb | 0 notebooks/scene_sensor_loading.ipynb | 8 +- notebooks/waymo_perception/testing.ipynb | 19 ++ 8 files changed, 32 insertions(+), 318 deletions(-) delete mode 100644 notebooks/demo_gym.ipynb rename notebooks/{ => gym}/test_gym.ipynb (100%) rename notebooks/{ => gym}/test_simulation_2d.ipynb (100%) rename notebooks/{ => nuplan}/nuplan_sensor_loading.ipynb (100%) create mode 100644 notebooks/waymo_perception/testing.ipynb diff --git a/d123/common/datatypes/sensor/lidar.py b/d123/common/datatypes/sensor/lidar.py index 7b13f261..a14e6edd 100644 --- a/d123/common/datatypes/sensor/lidar.py +++ b/d123/common/datatypes/sensor/lidar.py @@ -4,6 +4,11 @@ import numpy.typing as npt +@dataclass +class LiDARMetadata: + pass + + @dataclass class LiDAR: diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 9525afe2..e08fb9a9 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -27,8 +27,6 @@ from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.vector import Vector3D -from d123.dataset.dataset_specific.carla.load_sensor import load_carla_lidar_from_path -from d123.dataset.dataset_specific.nuplan.load_sensor import load_nuplan_lidar_from_path from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import List @@ -134,9 +132,15 @@ def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int, log_metadata: sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] full_lidar_path = sensor_root / lidar_data assert full_lidar_path.exists(), f"LiDAR file not found: {full_lidar_path}" + + # NOTE: We move data specific import into if-else block, to avoid data specific import errors if log_metadata.dataset == "nuplan": + from d123.dataset.dataset_specific.nuplan.load_sensor import load_nuplan_lidar_from_path + lidar = load_nuplan_lidar_from_path(full_lidar_path) elif log_metadata.dataset == "carla": + from d123.dataset.dataset_specific.carla.load_sensor import load_carla_lidar_from_path + lidar = load_carla_lidar_from_path(full_lidar_path) else: raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") diff --git a/notebooks/demo_gym.ipynb b/notebooks/demo_gym.ipynb deleted file mode 100644 index 26622558..00000000 --- a/notebooks/demo_gym.ipynb +++ /dev/null @@ -1,310 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", - "from d123.dataset.scene.scene_filter import SceneFilter\n", - "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "import os, psutil\n", - "\n", - "\n", - "def print_memory_usage():\n", - " process = psutil.Process(os.getpid())\n", - " memory_info = process.memory_info()\n", - " print(f\"Memory usage: {memory_info.rss / 1024 ** 2:.2f} MB\")\n", - "\n", - "\n", - "split = \"nuplan_mini_val\"\n", - "\n", - "# log_names = [\"2021.06.07.12.54.00_veh-35_01843_02314\"]\n", - "scene_tokens = None\n", - "scene_tokens = [\"2283aea39bc1505e\"]\n", - "log_names = None\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=[split], log_names=log_names, scene_tokens=scene_tokens, duration_s=15.1, history_s=1.0\n", - ")\n", - "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", - "worker = Sequential()\n", - "# worker = RayDistributed()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "print(len(scenes))\n", - "\n", - "for scene in scenes[:10]:\n", - " print(scene.log_name, scene.token)\n", - "\n", - "\n", - "# 875f407079785268\n", - "\n", - "# _Rep0_longest1_route0_06_13_17_21_21 457da031d28ba67b\n", - "# _Rep0_longest1_route0_06_13_17_21_21 1f6d4a8c9a399b3b\n", - "# _Rep0_longest1_route0_06_13_17_21_21 5b7a3e90922db277\n", - "# _Rep0_longest1_route0_06_13_17_21_21 827d6ac2ff5356d0\n", - "# _Rep0_longest1_route0_06_13_17_21_21 16717aba72175425\n", - "# _Rep0_longest1_route0_06_13_17_21_21 5b7444f86f19b444\n", - "# _Rep0_longest1_route0_06_13_17_21_21 6ee6004b304f1d3c\n", - "# _Rep0_longest1_route0_06_13_17_21_21 c98966ed60a77f7e\n", - "# _Rep0_longest1_route0_06_13_17_21_21 790b09ca88553770\n", - "# _Rep0_longest1_route0_06_13_17_21_21 b6c1a8d385648623" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [ - "from pathlib import Path\n", - "from typing import Optional, Tuple\n", - "\n", - "import matplotlib.animation as animation\n", - "import matplotlib.pyplot as plt\n", - "from tqdm import tqdm\n", - "\n", - "from d123.common.geometry.base import Point2D, StateSE2\n", - "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE2\n", - "from d123.common.visualization.color.default import EGO_VEHICLE_CONFIG\n", - "from d123.common.visualization.matplotlib.observation import (\n", - " add_bounding_box_to_ax,\n", - " add_box_detections_to_ax,\n", - " add_default_map_on_ax,\n", - " add_traffic_lights_to_ax,\n", - ")\n", - "from d123.dataset.arrow.conversion import TrafficLightDetectionWrapper\n", - "from d123.dataset.maps.abstract_map import AbstractMap\n", - "from d123.common.datatypes.detection.detection import BoxDetectionWrapper\n", - "from d123.dataset.scene.abstract_scene import AbstractScene\n", - "from nuplan.common.actor_state.ego_state import EgoState\n", - "import io\n", - "from PIL import Image\n", - "\n", - "\n", - "def add_ego_vehicle_to_ax_(ax: plt.Axes, ego_state: EgoState) -> None:\n", - " bounding_box = BoundingBoxSE2(\n", - " center=StateSE2(*ego_state.center),\n", - " length=ego_state.car_footprint.length,\n", - " width=ego_state.car_footprint.width,\n", - " )\n", - " add_bounding_box_to_ax(ax, bounding_box, EGO_VEHICLE_CONFIG)\n", - "\n", - "\n", - "def _plot_scene_on_ax(\n", - " ax: plt.Axes,\n", - " map_api: AbstractMap,\n", - " ego_state: EgoState,\n", - " initial_ego_state: Optional[EgoState],\n", - " box_detections: BoxDetectionWrapper,\n", - " traffic_light_detections: TrafficLightDetectionWrapper,\n", - " radius: float = 120,\n", - ") -> plt.Axes:\n", - "\n", - " if initial_ego_state is not None:\n", - " point_2d = Point2D(initial_ego_state.center.x, initial_ego_state.center.y)\n", - " else:\n", - "\n", - " point_2d = Point2D(ego_state.center.x, ego_state.center.y)\n", - " add_default_map_on_ax(ax, map_api, point_2d, radius=radius)\n", - " add_traffic_lights_to_ax(ax, traffic_light_detections, map_api)\n", - "\n", - " add_box_detections_to_ax(ax, box_detections)\n", - " add_ego_vehicle_to_ax_(ax, ego_state)\n", - "\n", - " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", - " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", - "\n", - " ax.set_aspect(\"equal\", adjustable=\"box\")\n", - " return ax\n", - "\n", - "\n", - "def plot_scene_to_image(\n", - " map_api: AbstractMap,\n", - " ego_state: EgoState,\n", - " initial_ego_state: Optional[EgoState],\n", - " box_detections: BoxDetectionWrapper,\n", - " traffic_light_detections: TrafficLightDetectionWrapper,\n", - " radius: float = 120,\n", - " figsize: Tuple[int, int] = (8, 8),\n", - ") -> Image:\n", - "\n", - " fig, ax = plt.subplots(figsize=figsize)\n", - " _plot_scene_on_ax(ax, map_api, ego_state, initial_ego_state, box_detections, traffic_light_detections, radius)\n", - " ax.set_aspect(\"equal\", adjustable=\"box\")\n", - " plt.subplots_adjust(left=0.05, right=0.95, top=0.95, bottom=0.05)\n", - " # plt.tight_layout()\n", - "\n", - " buf = io.BytesIO()\n", - " fig.savefig(buf, format=\"png\", bbox_inches=\"tight\")\n", - " plt.close(fig)\n", - " buf.seek(0)\n", - " img = Image.open(buf)\n", - " return img" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.dataset.arrow.conversion import DetectionType\n", - "from d123.simulation.gym.demo_gym_env import DemoGymEnv\n", - "from d123.simulation.observation.agents_observation import _filter_agents_by_type\n", - "\n", - "import time\n", - "\n", - "images = []\n", - "agent_rollouts = []\n", - "plot: bool = True\n", - "action = [1.0, 0.1] # Placeholder action, replace with actual action logic\n", - "env = DemoGymEnv(scenes)\n", - "\n", - "start = time.time()\n", - "\n", - "map_api, ego_state, detection_observation, current_scene = env.reset(None)\n", - "initial_ego_state = ego_state\n", - "cars, _, _ = _filter_agents_by_type(detection_observation.box_detections, detection_types=[DetectionType.VEHICLE])\n", - "agent_rollouts.append(BoxDetectionWrapper(cars))\n", - "if plot:\n", - " images.append(\n", - " plot_scene_to_image(\n", - " map_api,\n", - " ego_state,\n", - " initial_ego_state,\n", - " detection_observation.box_detections,\n", - " detection_observation.traffic_light_detections,\n", - " )\n", - " )\n", - "\n", - "\n", - "for i in range(160):\n", - " ego_state, detection_observation, end = env.step(action)\n", - " cars, _, _ = _filter_agents_by_type(detection_observation.box_detections, detection_types=[DetectionType.VEHICLE])\n", - " agent_rollouts.append(BoxDetectionWrapper(cars))\n", - " if plot:\n", - " images.append(\n", - " plot_scene_to_image(\n", - " map_api,\n", - " ego_state,\n", - " initial_ego_state,\n", - " detection_observation.box_detections,\n", - " detection_observation.traffic_light_detections,\n", - " )\n", - " )\n", - " if end:\n", - " print(\"End of scene reached.\")\n", - " break\n", - "\n", - "print(time.time() - start)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "\n", - "\n", - "def create_gif(images, output_path, duration=100):\n", - " \"\"\"\n", - " Create a GIF from a list of PIL images.\n", - "\n", - " Args:\n", - " images (list): List of PIL.Image objects.\n", - " output_path (str): Path to save the GIF.\n", - " duration (int): Duration between frames in milliseconds.\n", - " \"\"\"\n", - " if images:\n", - " print(len(images))\n", - " images_p = [img.convert(\"P\", palette=Image.ADAPTIVE) for img in images]\n", - " images_p[0].save(output_path, save_all=True, append_images=images_p[1:], duration=duration, loop=0)\n", - "\n", - "\n", - "if plot:\n", - " create_gif(images, f\"{split}_{current_scene.token}.gif\", duration=20)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [ - "from typing import List\n", - "\n", - "from d123.simulation.metrics.sim_agents.sim_agents import get_sim_agents_metrics\n", - "\n", - "\n", - "\n", - "\n", - "result = get_sim_agents_metrics(current_scene, agent_rollouts)\n", - "result\n", - "\n", - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "d123", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/test_gym.ipynb b/notebooks/gym/test_gym.ipynb similarity index 100% rename from notebooks/test_gym.ipynb rename to notebooks/gym/test_gym.ipynb diff --git a/notebooks/test_simulation_2d.ipynb b/notebooks/gym/test_simulation_2d.ipynb similarity index 100% rename from notebooks/test_simulation_2d.ipynb rename to notebooks/gym/test_simulation_2d.ipynb diff --git a/notebooks/nuplan_sensor_loading.ipynb b/notebooks/nuplan/nuplan_sensor_loading.ipynb similarity index 100% rename from notebooks/nuplan_sensor_loading.ipynb rename to notebooks/nuplan/nuplan_sensor_loading.ipynb diff --git a/notebooks/scene_sensor_loading.ipynb b/notebooks/scene_sensor_loading.ipynb index e10a0842..1a4093fe 100644 --- a/notebooks/scene_sensor_loading.ipynb +++ b/notebooks/scene_sensor_loading.ipynb @@ -13,7 +13,7 @@ "\n", "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", "# from nuplan.planning.utils.multithreading.worker_ray impo\n", - "# rt RayDistributed\n", + "# rt RayDistribute\n", "\n", "\n" ] @@ -88,11 +88,7 @@ "id": "4", "metadata": {}, "outputs": [], - "source": [ - "\n", - "\n", - "plt.scatter(pc[0, :], pc[1, :], s=0.1, alpha=0.5)" - ] + "source": [] }, { "cell_type": "code", diff --git a/notebooks/waymo_perception/testing.ipynb b/notebooks/waymo_perception/testing.ipynb new file mode 100644 index 00000000..7ebcef1c --- /dev/null +++ b/notebooks/waymo_perception/testing.ipynb @@ -0,0 +1,19 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From bdef371a26f8b2b592f5dfc4380388eb88dadc51 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 3 Aug 2025 15:52:19 +0200 Subject: [PATCH 15/42] Add conversion for waymo open perception dataset (currently loading ego state and cameras) --- .../vehicle_state/vehicle_parameters.py | 14 + d123/common/geometry/base.py | 6 +- d123/common/visualization/viser/server.py | 15 +- d123/common/visualization/viser/utils.py | 35 +- d123/dataset/arrow/conversion.py | 4 +- d123/dataset/dataset_specific/waymo/.gitkeep | 0 .../wopd/wopd_data_converter.py | 415 ++++++++++++++++++ d123/dataset/scene/arrow_scene.py | 7 +- .../default_dataset_conversion.yaml | 3 +- d123/script/config/datasets/wopd_dataset.yaml | 16 + notebooks/tools/merge_videos.ipynb | 4 +- notebooks/viz/viser_testing_v2_scene.ipynb | 8 +- notebooks/waymo_perception/testing.ipynb | 371 +++++++++++++++- 13 files changed, 878 insertions(+), 20 deletions(-) delete mode 100644 d123/dataset/dataset_specific/waymo/.gitkeep create mode 100644 d123/dataset/dataset_specific/wopd/wopd_data_converter.py create mode 100644 d123/script/config/datasets/wopd_dataset.yaml diff --git a/d123/common/datatypes/vehicle_state/vehicle_parameters.py b/d123/common/datatypes/vehicle_state/vehicle_parameters.py index c09e9b6d..b899d400 100644 --- a/d123/common/datatypes/vehicle_state/vehicle_parameters.py +++ b/d123/common/datatypes/vehicle_state/vehicle_parameters.py @@ -46,6 +46,20 @@ def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: ) +def get_wopd_pacifica_parameters() -> VehicleParameters: + # NOTE: use parameters from nuPlan dataset + # Find better parameters for WOPD ego vehicle + return VehicleParameters( + vehicle_name="wopd_pacifica", + width=2.297, + length=5.176, + height=1.777, + wheel_base=3.089, + rear_axle_to_center_vertical=0.45, + rear_axle_to_center_longitudinal=1.461, + ) + + def center_se3_to_rear_axle_se3(center_se3: StateSE3, vehicle_parameters: VehicleParameters) -> StateSE3: """ Converts a center state to a rear axle state. diff --git a/d123/common/geometry/base.py b/d123/common/geometry/base.py index 28adc163..97abe712 100644 --- a/d123/common/geometry/base.py +++ b/d123/common/geometry/base.py @@ -144,8 +144,10 @@ class Point3D: @classmethod def from_array(cls, array: npt.NDArray[np.float64]) -> "Point3D": - assert array.ndim == 1 - assert array.shape[0] == len(Point3DIndex) + assert array.ndim == 1, f"Array must be 1-dimensional, got shape {array.shape}" + assert array.shape[0] == len( + Point3DIndex + ), f"Array must have the same length as Point3DIndex, got shape {array.shape}" return cls(array[Point3DIndex.X], array[Point3DIndex.Y], array[Point3DIndex.Z]) @property diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 6a401593..f772c096 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -29,14 +29,23 @@ # ] LIDAR_POINT_SIZE: float = 0.05 -MAP_AVAILABLE: bool = True +MAP_AVAILABLE: bool = False BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" LINE_WIDTH: float = 4.0 CAMERA_SCALE: float = 1.0 -VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] + +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ + CameraType.CAM_F0, + CameraType.CAM_L0, + CameraType.CAM_R0, + CameraType.CAM_L1, + CameraType.CAM_R1, +] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0] VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] -LIDAR_AVAILABLE: bool = True +LIDAR_AVAILABLE: bool = False class ViserVisualizationServer: diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index dbca9a02..35b0c39b 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -185,12 +185,39 @@ def get_camera_values( ego_transform[:3, :3] = get_rotation_matrix(rear_axle) ego_transform[:3, 3] = rear_axle.point_3d.array - camera_transform = ego_transform @ camera_to_ego + DEBUG = False + if DEBUG: + print("DEBUG") + + camera_to_ego = camera.extrinsic + rotation = camera_to_ego[:3, :3] + + flip_camera = get_rotation_matrix( + StateSE3( + x=0.0, + y=0.0, + z=0.0, + roll=np.deg2rad(-90.0), + pitch=np.deg2rad(0.0), + yaw=np.deg2rad(-90.0), + ) + ) + + corrected_rotation = camera_to_ego[:3, :3] @ flip_camera + + camera_to_ego[:3, :3] = corrected_rotation + camera_transform = ego_transform @ camera_to_ego + + # Camera transformation in ego frame + camera_position = Point3D(*camera_transform[:3, 3]) + camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) - # Camera transformation in ego frame + else: + camera_transform = ego_transform @ camera_to_ego - camera_position = Point3D(*camera_transform[:3, 3]) - camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) + # Camera transformation in ego frame + camera_position = Point3D(*camera_transform[:3, 3]) + camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) return camera_position, camera_rotation, camera diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index e08fb9a9..a671e2d7 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -1,6 +1,6 @@ # TODO: rename this file and potentially move somewhere more appropriate. - +import io import os from pathlib import Path from typing import Dict, Optional @@ -115,6 +115,8 @@ def get_camera_from_arrow_table( img = Image.open(full_image_path) img.load() image = np.asarray(img, dtype=np.uint8) + elif isinstance(table_data, bytes): + image = np.array(Image.open(io.BytesIO(table_data))) else: raise NotImplementedError("Only string file paths for camera data are supported.") diff --git a/d123/dataset/dataset_specific/waymo/.gitkeep b/d123/dataset/dataset_specific/waymo/.gitkeep deleted file mode 100644 index e69de29b..00000000 diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py new file mode 100644 index 00000000..347a7ac6 --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -0,0 +1,415 @@ +import gc +import hashlib +import json +from dataclasses import asdict +from functools import partial +from pathlib import Path +from typing import Any, Dict, Final, List, Optional, Tuple, Union + +import numpy as np +import numpy.typing as npt +import pyarrow as pa +import tensorflow as tf +from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map +from pyquaternion import Quaternion +from waymo_open_dataset import dataset_pb2 + +from d123.common.datatypes.detection.detection import TrafficLightStatus +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index +from d123.common.datatypes.vehicle_state.vehicle_parameters import get_wopd_pacifica_parameters +from d123.common.geometry.base import Point3D, StateSE3 +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index +from d123.common.geometry.transform.se3 import get_rotation_matrix +from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.logs.log_metadata import LogMetadata + +TARGET_DT: Final[float] = 0.1 +NUPLAN_DT: Final[float] = 0.05 +SORT_BY_TIMESTAMP: Final[bool] = True + +NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { + "green": TrafficLightStatus.GREEN, + "red": TrafficLightStatus.RED, + "unknown": TrafficLightStatus.UNKNOWN, +} +NUPLAN_DETECTION_NAME_DICT = { + "vehicle": DetectionType.VEHICLE, + "bicycle": DetectionType.BICYCLE, + "pedestrian": DetectionType.PEDESTRIAN, + "traffic_cone": DetectionType.TRAFFIC_CONE, + "barrier": DetectionType.BARRIER, + "czone_sign": DetectionType.CZONE_SIGN, + "generic_object": DetectionType.GENERIC_OBJECT, +} + +WOPD_CAMERA_TYPES = { + CameraType.CAM_F0: 1, # front_camera + CameraType.CAM_L0: 2, # front_left_camera + CameraType.CAM_R0: 3, # front_right_camera + CameraType.CAM_L1: 4, # left_camera + CameraType.CAM_R1: 5, # right_camera +} + +WOPD_DATA_ROOT = Path("/media/nvme1/waymo_perception") # TODO: set as environment variable + + +def create_token(input_data: str) -> str: + # TODO: Refactor this function. + # TODO: Add a general function to create tokens from arbitrary data. + if isinstance(input_data, str): + input_data = input_data.encode("utf-8") + + hash_obj = hashlib.sha256(input_data) + return hash_obj.hexdigest()[:16] + + +class WOPDDataConverter(RawDataConverter): + def __init__( + self, + splits: List[str], + log_path: Union[Path, str], + data_converter_config: DataConverterConfig, + ) -> None: + super().__init__(data_converter_config) + for split in splits: + assert ( + split in self.get_available_splits() + ), f"Split {split} is not available. Available splits: {self.available_splits}" + + self._splits: List[str] = splits + self._tf_records_per_split: Dict[str, List[Path]] = self._collect_tf_records() + self._target_dt: float = 0.1 + + def _collect_tf_records(self) -> Dict[str, List[Path]]: + tf_records_per_split: Dict[str, List[Path]] = {} + + for split in self._splits: + if split in ["wopd_train"]: + log_path = WOPD_DATA_ROOT / "training" + else: + raise ValueError(f"Split {split} is not supported.") + + log_paths = [log_file for log_file in log_path.glob("*.tfrecord")] + tf_records_per_split[split] = log_paths + + return tf_records_per_split + + def get_available_splits(self) -> List[str]: + # TODO: Add more splits if available + return [ + "wopd_train", + ] + + def convert_maps(self, worker: WorkerPool) -> None: + # TODO: Implement map conversion + pass + + def convert_logs(self, worker: WorkerPool) -> None: + log_args = [ + { + "tf_record": tf_record, + "split": split, + } + for split, tf_record_paths in self._tf_records_per_split.items() + for tf_record in tf_record_paths + ] + + worker_map( + worker, + partial( + convert_wopd_tfrecord_log_to_arrow, + data_converter_config=self.data_converter_config, + ), + log_args, + ) + + +def convert_wopd_tfrecord_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: + # TODO: Implement map conversion + return [] + + +def convert_wopd_tfrecord_log_to_arrow( + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig +) -> List[Any]: + for log_info in args: + tf_record_path: Path = log_info["tf_record"] + split: str = log_info["split"] + + if not tf_record_path.exists(): + raise FileNotFoundError(f"TFRecord path {tf_record_path} does not exist.") + + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") + for data in dataset: + frame = dataset_pb2.Frame() + frame.ParseFromString(data.numpy()) + break + + log_name = str(frame.context.name) + log_file_path = data_converter_config.output_path / split / f"{log_name}.arrow" + + if data_converter_config.force_log_conversion or not log_file_path.exists(): + log_file_path.unlink(missing_ok=True) + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + if data_converter_config.lidar_store_option == "path": + schema_column_list.append(("lidar", pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + # TODO: Adjust how cameras are added + if data_converter_config.camera_store_option is not None: + for camera_type in WOPD_CAMERA_TYPES.keys(): + if data_converter_config.camera_store_option == "path": + raise NotImplementedError("Path camera storage is not implemented.") + elif data_converter_config.camera_store_option == "binary": + schema_column_list.append((camera_type.serialize(), pa.binary())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) + ) + + initial_frame = next( + (dataset_pb2.Frame().FromString(data.numpy()) or dataset_pb2.Frame() for data in dataset), None + ) + if initial_frame is None: + raise ValueError(f"No frames found in TFRecord {tf_record_path}") + + recording_schema = pa.schema(schema_column_list) + metadata = LogMetadata( + dataset="wopd", + log_name=log_name, + location=None, # TODO: implement map name + timestep_seconds=0.1, # TODO: Check if correct. Maybe not hardcode + map_has_z=True, + ) + vehicle_parameters = get_wopd_pacifica_parameters() + camera_metadata = get_wopd_camera_metadata(initial_frame) + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + } + ) + + _write_recording_table(dataset, recording_schema, log_file_path, tf_record_path, data_converter_config) + + del recording_schema, vehicle_parameters, dataset + gc.collect() + return [] + + +def get_wopd_camera_metadata(initial_frame: dataset_pb2.Frame) -> Dict[str, CameraMetadata]: + + name_to_camera_type = {v: k for k, v in WOPD_CAMERA_TYPES.items()} + cam_metadatas: Dict[str, CameraMetadata] = {} + for calibration in initial_frame.context.camera_calibrations: + camera_type = name_to_camera_type[calibration.name] + + # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L96 + # https://github.com/waymo-research/waymo-open-dataset/issues/834#issuecomment-2134995440 + fx, fy, cx, cy, k1, k2, p1, p2, k3 = calibration.intrinsic + _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) + _distortions = np.array([k1, k2, p1, p2, k3]) + + if camera_type in WOPD_CAMERA_TYPES: + cam_metadatas[camera_type.serialize()] = CameraMetadata( + camera_type=camera_type, + width=calibration.width, + height=calibration.height, + intrinsic=_intrinsics, + distortion=_distortions, + ) + + return cam_metadatas + + +def _write_recording_table( + dataset: tf.data.TFRecordDataset, + recording_schema: pa.schema, + log_file_path: Path, + tf_record_path: Path, + data_converter_config: DataConverterConfig, +) -> None: + + # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: + with pa.OSFile(str(log_file_path), "wb") as sink: + with pa.ipc.new_file(sink, recording_schema) as writer: + + for frame_idx, data in enumerate(dataset): + frame = dataset_pb2.Frame() + frame.ParseFromString(data.numpy()) + + ( + detections_state, + detections_velocity, + detections_token, + detections_types, + ) = _extract_detections(frame) + # traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) + # route_lane_group_ids = [ + # int(roadblock_id) + # for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") + # if len(roadblock_id) > 0 + # ] + + # TODO: Implement traffic light extraction + traffic_light_ids = [] + traffic_light_types = [] + + # TODO: Implement detections + detections_state = [] + detections_velocity = [] + detections_token = [] + detections_types = [] + + row_data = { + "token": [create_token(f"{frame.context.name}_{int(frame.timestamp_micros)}")], + "timestamp": [int(frame.timestamp_micros)], + "detections_state": [detections_state], + "detections_velocity": [detections_velocity], + "detections_token": [detections_token], + "detections_type": [detections_types], + "ego_states": [_extract_ego_state(frame)], + "traffic_light_ids": [traffic_light_ids], + "traffic_light_types": [traffic_light_types], + "scenario_tag": ["unknown"], + "route_lane_group_ids": [None], + } + + # TODO: Implement lidar extraction + if data_converter_config.lidar_store_option is not None: + row_data["lidar"] = [None] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_camera(frame, data_converter_config) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] + else: + row_data[camera_type.serialize()] = [None] + row_data[f"{camera_type.serialize()}_extrinsic"] = [None] + + batch = pa.record_batch(row_data, schema=recording_schema) + writer.write_batch(batch) + del batch, row_data, detections_state, detections_velocity, detections_token, detections_types + + if SORT_BY_TIMESTAMP: + recording_table = open_arrow_table(log_file_path) + recording_table = recording_table.sort_by([("timestamp", "ascending")]) + write_arrow_table(recording_table, log_file_path) + + +def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: + detections_state: List[List[float]] = [] + detections_velocity: List[List[float]] = [] + detections_token: List[str] = [] + detections_types: List[int] = [] + # TODO: implement + + return detections_state, detections_velocity, detections_token, detections_types + + +def _extract_ego_state(frame: dataset_pb2.Frame) -> List[float]: + + ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) + yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll + ego_point_3d = Point3D.from_array(ego_pose_matrix[:3, 3]) + + vehicle_parameters = get_wopd_pacifica_parameters() + + # TODO: figure out if ego frame is given in rear axle or center frame + rear_axle_pose = StateSE3( + x=ego_point_3d.x, + y=ego_point_3d.y, + z=ego_point_3d.z, + roll=roll, + pitch=pitch, + yaw=yaw, + ) + # FIXME: Find dynamic state in waymo open perception dataset + # https://github.com/waymo-research/waymo-open-dataset/issues/55#issuecomment-546152290 + dynamic_state = DynamicStateSE3( + velocity=Vector3D(*np.zeros(3)), + acceleration=Vector3D(*np.zeros(3)), + angular_velocity=Vector3D(*np.zeros(3)), + ) + + return EgoStateSE3.from_rear_axle( + rear_axle_se3=rear_axle_pose, + dynamic_state_se3=dynamic_state, + vehicle_parameters=vehicle_parameters, + time_point=None, + ).array.tolist() + + +def _extract_traffic_lights() -> Tuple[List[int], List[int]]: + pass + + +def _extract_camera( + frame: dataset_pb2.Frame, + data_converter_config: DataConverterConfig, +) -> Dict[CameraType, Union[str, bytes]]: + + name_to_camera_type = {v: k for k, v in WOPD_CAMERA_TYPES.items()} + camera_dict: Dict[str, Union[str, bytes]] = {} # TODO: Fix wrong type hint + ego_global_transform = np.array(frame.pose.transform).reshape(4, 4) + + # NOTE: The extrinsic matrix in frame.context.camera_calibration is fixed to model the ego to camera transformation. + # The poses in frame.images[idx] are the motion compensated ego poses when the camera triggers. + # + + context_extrinsic: Dict[str, npt.NDArray] = {} + for calibration in frame.context.camera_calibrations: + camera_type = name_to_camera_type[calibration.name] + + transform = np.array(calibration.extrinsic.transform).reshape(4, 4) + + # FIXME: This is an ugly hack to convert to uniform camera convention. + flip_camera = get_rotation_matrix( + StateSE3(x=0.0, y=0.0, z=0.0, roll=np.deg2rad(-90.0), pitch=np.deg2rad(0.0), yaw=np.deg2rad(-90.0)) + ) + + transform[:3, :3] = transform[:3, :3] @ flip_camera + context_extrinsic[camera_type] = transform + + for image_proto in frame.images: + camera_type = name_to_camera_type[image_proto.name] + + ego_at_trigger_transform = np.array(image_proto.pose.transform).reshape(4, 4) + camera_bytes = image_proto.image + + # Compute the transform from ego_global_transform to ego_at_camera_transform + # ego_global_transform * T = ego_at_camera_transform => T = ego_global_transform^-1 * ego_at_camera_transform + np.linalg.inv(ego_global_transform) @ ego_at_trigger_transform + + # TODO: figure out the correct transform + camera_dict[camera_type] = camera_bytes, context_extrinsic[camera_type].flatten().tolist() + + return camera_dict + + +def _extract_lidar() -> Optional[str]: + lidar: Optional[str] = None + return lidar diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index 434147be..76c5e530 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -163,8 +163,11 @@ def _lazy_initialize(self) -> None: def open(self) -> None: if self._map_api is None: - self._map_api = get_map_api_from_names(self._metadata.dataset, self._metadata.location) - self._map_api.initialize() + try: + self._map_api = get_map_api_from_names(self._metadata.dataset, self._metadata.location) + self._map_api.initialize() + except Exception as e: + print(f"Error initializing map API: {e}") if self._recording_table is None: self._recording_table = open_arrow_table(self._arrow_log_path) if self._scene_extraction_info is None: diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 7d5fcd7b..a98adec2 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -16,7 +16,8 @@ defaults: - _self_ - datasets: # - nuplan_private_dataset - - carla_dataset + # - carla_dataset + - wopd_dataset force_log_conversion: True force_map_conversion: False diff --git a/d123/script/config/datasets/wopd_dataset.yaml b/d123/script/config/datasets/wopd_dataset.yaml new file mode 100644 index 00000000..9234dbfa --- /dev/null +++ b/d123/script/config/datasets/wopd_dataset.yaml @@ -0,0 +1,16 @@ +wopd_dataset: + _target_: d123.dataset.dataset_specific.wopd.wopd_data_converter.WOPDDataConverter + _convert_: 'all' + + splits: ["wopd_train"] + log_path: null # TODO: implement + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: "binary" + lidar_store_option: null diff --git a/notebooks/tools/merge_videos.ipynb b/notebooks/tools/merge_videos.ipynb index 9e1c9f1b..97d442d6 100644 --- a/notebooks/tools/merge_videos.ipynb +++ b/notebooks/tools/merge_videos.ipynb @@ -89,7 +89,7 @@ " return False\n", "\n", "# List your MP4 files in the order you want them merged\n", - "video_folder = Path(\"/home/daniel/d123_logs_videos/nuplan_mini_val\")\n", + "video_folder = Path(\"/home/daniel/d123_workspace/d123/notebooks/waymo_perception\")\n", "video_files = [str(file) for file in video_folder.glob(\"*.mp4\") if file.is_file()]\n", "\n", "\n", @@ -129,7 +129,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.3" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 6eb07e46..caefb11e 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -25,7 +25,7 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "splits = [\"carla\"]\n", + "splits = [\"wopd_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -39,9 +39,9 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=20.1,\n", + " duration_s=19,\n", " history_s=1.0,\n", - " timestamp_threshold_s=20.1,\n", + " timestamp_threshold_s=20,\n", " shuffle=False,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", @@ -64,7 +64,7 @@ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes, scene_index=70)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" ] }, { diff --git a/notebooks/waymo_perception/testing.ipynb b/notebooks/waymo_perception/testing.ipynb index 7ebcef1c..bd035f8f 100644 --- a/notebooks/waymo_perception/testing.ipynb +++ b/notebooks/waymo_perception/testing.ipynb @@ -6,12 +6,381 @@ "id": "0", "metadata": {}, "outputs": [], + "source": [ + "from waymo_open_dataset import dataset_pb2\n", + "\n", + "import json\n", + "import os\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "from PIL import Image\n", + "from tqdm import tqdm\n", + "from waymo_open_dataset import label_pb2\n", + "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", + "from waymo_open_dataset.utils import box_utils\n", + "from waymo_open_dataset.utils.frame_utils import parse_range_image_and_camera_projection\n", + "\n", + "\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "\n", + "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", + "\n", + "\n", + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from pyquaternion import Quaternion\n", + "\n", + "from d123.common.geometry.base import StateSE3\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", + "\n", + "\n", + "# Frame attributes:\n", + "# context: \n", + "# timestamp_micros: \n", + "# pose: \n", + "# images: List with 5 images\n", + "# lasers: \n", + "# Length: 5\n", + "# laser_labels: \n", + "# Length: 0\n", + "# projected_lidar_labels: \n", + "# Length: 0\n", + "# camera_labels: \n", + "# Length: 0\n", + "# no_label_zones: \n", + "# Length: 0\n", + "# map_features: \n", + "# Length: 0\n", + "# map_pose_offset: \n", + "\n", + "file_idx = 0\n", + "pathname = tfrecords_file_list[file_idx]\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "num_frames = sum(1 for _ in dataset)\n", + "\n", + "\n", + "def read_jpg_image(data: bytes) -> np.ndarray:\n", + " \"\"\"Read a JPEG image from bytes and return it as a numpy array.\"\"\"\n", + " image = Image.open(io.BytesIO(data))\n", + " return np.array(image)\n", + "\n", + "\n", + "\n", + "\n", + "ego_state_se3s = []\n", + "front_images = []\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "\n", + "boxes = []\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + "\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " # print(frame.camera_labels)\n", + " for label in frame.laser_labels:\n", + " boxes.append(\n", + " BoundingBoxSE3(\n", + " center=StateSE3(\n", + " x=label.box.center_x,\n", + " y=label.box.center_y,\n", + " z=label.box.center_z,\n", + " pitch=0.0,\n", + " roll=0.0,\n", + " yaw=label.box.heading,\n", + " ),\n", + " length=label.box.length,\n", + " width=label.box.width,\n", + " height=label.box.height,\n", + " )\n", + " )\n", + "\n", + " print(frame.context)\n", + "\n", + " # Print all attributes of the frame\n", + " # print(\"Frame attributes:\")\n", + " # for field in frame.DESCRIPTOR.fields:\n", + " # field_name = field.name\n", + " # if hasattr(frame, field_name):\n", + " # value = getattr(frame, field_name)\n", + " # if field_name != \"images\": # Don't print the whole image data\n", + " # print(f\" {field_name}: {type(value)}\")\n", + " # if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", + " # print(f\" Length: {len(value)}\")\n", + " # else:\n", + " # print(f\" {field_name}: List with {len(value)} images\")\n", + "\n", + " # Print information about the first image if available\n", + " # if frame.images:\n", + " # print(\"\\nFirst image details:\")\n", + " # first_image = frame.images[0]\n", + " # for field in first_image.DESCRIPTOR.fields:\n", + " # field_name = field.name\n", + " # if hasattr(first_image, field_name):\n", + " # value = getattr(first_image, field_name)\n", + " # if field_name != \"image\": # Don't print the binary data\n", + " # print(f\" {field_name}: {value}\")\n", + " # else:\n", + " # print(f\" {field_name}: binary data of length {len(value)} bytes\")\n", + "\n", + " # for image in frame.images:\n", + " # print(image.name)\n", + "\n", + " # print([image.name for image in frame.images])\n", + " # print(frame.images[0])\n", + "\n", + " # # 1. pose\n", + " pose = np.array(frame.pose.transform).reshape(4, 4)\n", + " yaw_pitch_roll = Quaternion(matrix=pose[:3, :3]).yaw_pitch_roll\n", + " ego_state_se3s.append(\n", + " np.array(\n", + " [\n", + " pose[0, 3], # x\n", + " pose[1, 3], # y\n", + " pose[2, 3], # z\n", + " yaw_pitch_roll[2], # yaw\n", + " yaw_pitch_roll[1], # pitch\n", + " yaw_pitch_roll[0], # roll\n", + " ],\n", + " dtype=np.float64,\n", + " )\n", + " )\n", + "\n", + " # # plt.show()\n", + " if frame_idx == 0:\n", + " break\n", + "\n", + "ego_state_se3s = np.array(ego_state_se3s, dtype=np.float64)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "for frame_idx, data in enumerate(dataset):\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " if frame_idx == 3:\n", + " break\n", + "\n", + "print(\"Ego\")\n", + "ego_transform = np.array(frame.pose.transform).reshape(4, 4)\n", + "print(ego_transform[:3, 3])\n", + "\n", + "print(\"Frame\")\n", + "for image in frame.images:\n", + " image_transform = np.array(image.pose.transform).reshape(4, 4)\n", + " print(image.name, image_transform[:3, 3])\n", + "\n", + "print(\"Context\")\n", + "for image in frame.context.camera_calibrations:\n", + " image_transform = np.array(image.extrinsic.transform).reshape(4, 4)\n", + " print(image.name, image_transform[:3, 3])\n", + "\n", + "# 1 [ 1.5441613 -0.02302364 2.11557864]\n", + "# 2 [1.49672397 0.0954948 2.11616463]\n", + "# 3 [ 1.49442485 -0.09637497 2.11519385]\n", + "# 4 [1.43213651 0.11612398 2.11625087]\n", + "# 5 [ 1.42936162 -0.11545043 2.1150792 ]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "from d123.common.datatypes.time.time_point import TimePoint\n", + "\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " if frame_idx == 4:\n", + " break\n", + " break\n", + "\n", + "\n", + "\n", + "# for calibration in frame.context.camera_calibrations:\n", + "\n", + "frame.timestamp_micros, frame.images[0].pose_timestamp\n", + "# frame.images[0]\n", + "\n", + "frame_timestamp = TimePoint.from_us(frame.timestamp_micros)\n", + "image_timestamp = TimePoint.from_s(frame.images[0].pose_timestamp)\n", + "frame_timestamp.time_s, frame_timestamp.time_s" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "frame = next((dataset_pb2.Frame().FromString(data.numpy()) or dataset_pb2.Frame() for data in dataset), None)\n", + "# if frame is None:\n", + "# raise ValueError(f\"No frames found in TFRecord {tf_record_path}\")\n", + "\n", + "frame" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.datatypes.detection.detection_types import DetectionType\n", + "from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG\n", + "from d123.common.visualization.matplotlib.observation import add_bounding_box_to_ax\n", + "\n", + "\n", + "ego_rear_axle = StateSE3.from_array(ego_state_se3s[0])\n", + "\n", + "ego_rear_axle = StateSE3.from_array(np.zeros_like(ego_state_se3s[0]))\n", + "\n", + "ego_box = BoundingBoxSE3(center=ego_rear_axle, length=4.0, width=1.8, height=1.6)\n", + "\n", + "plot_config = BOX_DETECTION_CONFIG[DetectionType.VEHICLE]\n", + "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n", + "\n", + "\n", + "for box in boxes:\n", + " add_bounding_box_to_ax(ax, box, plot_config)\n", + "\n", + "\n", + "add_bounding_box_to_ax(ax, ego_box, EGO_VEHICLE_CONFIG)\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "ax.plot(ego_state_se3s[:, 0], ego_state_se3s[:, 1])\n", + "ax.set_aspect('equal', adjustable='box')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "import cv2\n", + "\n", + "# Define the output video path\n", + "output_video_path = str(f\"front_camera_video_{file_idx}.mp4\")\n", + "\n", + "# Get the dimensions of the first image\n", + "height, width, channels = front_images[0].shape\n", + "\n", + "# Create the video writer\n", + "fourcc = cv2.VideoWriter_fourcc(*'mp4v') # MP4 codec\n", + "fps = 10 # 10 frames per second as requested\n", + "out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))\n", + "\n", + "# Write each frame to the video\n", + "for img in tqdm(front_images, desc=\"Creating video\"):\n", + " # Convert from RGB to BGR (OpenCV uses BGR)\n", + " img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)\n", + " out.write(img_bgr)\n", + "\n", + "# Release the video writer\n", + "out.release()\n", + "\n", + "print(f\"Video saved to {output_video_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "# import pyarrow as pa\n", + "import pandas as pd\n", + "# parquet_file = \"/home/daniel/Downloads/testing_location_vehicle_pose_10084636266401282188_1120_000_1140_000.parquet\"\n", + "\n", + "parquet_file = \"/home/daniel/Downloads/validation_stats_10203656353524179475_7625_000_7645_000.parquet\"\n", + "\n", + "df = pd.read_parquet(parquet_file)\n", + "\n", + "df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "array" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], "source": [] } ], "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, "language_info": { - "name": "python" + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" } }, "nbformat": 4, From b07af56460fd3758f34a3f984de0472118d9bf1c Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 3 Aug 2025 19:28:49 +0200 Subject: [PATCH 16/42] Add conversion for waymo open perception dataset (currently loading ego state and cameras) --- .../datatypes/detection/detection_types.py | 1 + d123/common/geometry/base.py | 26 ++++ d123/common/geometry/base_index.py | 1 + d123/common/geometry/transform/se3.py | 32 ++++ d123/common/visualization/color/default.py | 10 ++ d123/common/visualization/viser/server.py | 1 + .../wopd/wopd_data_converter.py | 111 ++++++++------ notebooks/viz/bev_matplotlib.ipynb | 139 ++++++++++++++++++ notebooks/viz/camera_matplotlib.ipynb | 19 ++- notebooks/viz/viser_testing_v2_scene.ipynb | 12 +- 10 files changed, 299 insertions(+), 53 deletions(-) create mode 100644 d123/common/geometry/base_index.py create mode 100644 notebooks/viz/bev_matplotlib.ipynb diff --git a/d123/common/datatypes/detection/detection_types.py b/d123/common/datatypes/detection/detection_types.py index 8df27d2f..a22a614c 100644 --- a/d123/common/datatypes/detection/detection_types.py +++ b/d123/common/datatypes/detection/detection_types.py @@ -24,6 +24,7 @@ class DetectionType(SerialIntEnum): GENERIC_OBJECT = 6 # Animals, debris, pushable/pullable objects, permanent poles. EGO = 7 + SIGN = 8 # TODO: Remove or extent DYNAMIC_DETECTION_TYPES: set[DetectionType] = { diff --git a/d123/common/geometry/base.py b/d123/common/geometry/base.py index 97abe712..7fdec5cc 100644 --- a/d123/common/geometry/base.py +++ b/d123/common/geometry/base.py @@ -8,6 +8,7 @@ import numpy.typing as npt import shapely.geometry as geom +# from d123.common.geometry.transform.se3 import get_rotation_matrix from d123.common.utils.enums import classproperty # TODO: Reconsider if 2D/3D or SE2/SE3 structure would be better hierarchical, e.g. inheritance or composition. @@ -228,6 +229,21 @@ def from_array(cls, array: npt.NDArray[np.float64]) -> StateSE3: array[StateSE3Index.YAW], ) + @classmethod + def from_matrix(cls, array: npt.NDArray[np.float64]) -> StateSE3: + assert array.ndim == 2 + assert array.shape == (4, 4) + translation = array[:3, 3] + rotation = array[:3, :3] + return StateSE3( + x=translation[0], + y=translation[1], + z=translation[2], + roll=np.arctan2(rotation[2, 1], rotation[2, 2]), + pitch=np.arctan2(-rotation[2, 0], np.sqrt(rotation[2, 1] ** 2 + rotation[2, 2] ** 2)), + yaw=np.arctan2(rotation[1, 0], rotation[0, 0]), + ) + @property def array(self) -> npt.NDArray[np.float64]: array = np.zeros(len(StateSE3Index), dtype=np.float64) @@ -239,6 +255,16 @@ def array(self) -> npt.NDArray[np.float64]: array[StateSE3Index.YAW] = self.yaw return array + # @property + # def matrix(self) -> npt.NDArray[np.float64]: + # """Convert SE3 state to 4x4 transformation matrix.""" + # R = get_rotation_matrix(self) + # translation = np.array([self.x, self.y, self.z], dtype=np.float64) + # matrix = np.eye(4, dtype=np.float64) + # matrix[:3, :3] = R + # matrix[:3, 3] = translation + # return matrix + @property def state_se2(self) -> StateSE2: return StateSE2(self.x, self.y, self.yaw) diff --git a/d123/common/geometry/base_index.py b/d123/common/geometry/base_index.py new file mode 100644 index 00000000..cb258618 --- /dev/null +++ b/d123/common/geometry/base_index.py @@ -0,0 +1 @@ +# TODO: Move base index here to avoid circular imports. diff --git a/d123/common/geometry/transform/se3.py b/d123/common/geometry/transform/se3.py index dc09ed83..7dc38c5c 100644 --- a/d123/common/geometry/transform/se3.py +++ b/d123/common/geometry/transform/se3.py @@ -102,6 +102,7 @@ def convert_absolute_to_relative_se3_array( origin: StateSE3, se3_array: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: assert se3_array.shape[-1] == len(StateSE3Index) + # TODO: remove transform for-loop, use vectorized operations # Extract rotation and translation of origin R_origin = get_rotation_matrix(origin) @@ -128,6 +129,37 @@ def convert_absolute_to_relative_se3_array( return rel_se3_array +def convert_relative_to_absolute_se3_array( + origin: StateSE3, se3_array: npt.NDArray[np.float64] +) -> npt.NDArray[np.float64]: + assert se3_array.shape[-1] == len(StateSE3Index) + # TODO: remove transform for-loop, use vectorized operations + + # Extract rotation and translation of origin + R_origin = get_rotation_matrix(origin) + t_origin = origin.point_3d.array + + # Prepare output array + abs_se3_array = np.empty_like(se3_array) + + # For each SE3 in the array + for i in range(se3_array.shape[0]): + rel_se3 = se3_array[i] + rel_pos = rel_se3[StateSE3Index.XYZ] + rel_rpy = rel_se3[StateSE3Index.ROLL : StateSE3Index.YAW + 1] + + # Absolute position: rotate and translate + abs_pos = R_origin @ rel_pos + t_origin + + # Absolute orientation: add origin's rpy + abs_rpy = rel_rpy + np.array([origin.roll, origin.pitch, origin.yaw], dtype=np.float64) + + abs_se3_array[i, StateSE3Index.X : StateSE3Index.Z + 1] = abs_pos + abs_se3_array[i, StateSE3Index.ROLL : StateSE3Index.YAW + 1] = abs_rpy + + return abs_se3_array + + def translate_points_3d_along_z( state_se3: StateSE3, points_3d: npt.NDArray[np.float64], diff --git a/d123/common/visualization/color/default.py b/d123/common/visualization/color/default.py index 83921339..f5a3cf94 100644 --- a/d123/common/visualization/color/default.py +++ b/d123/common/visualization/color/default.py @@ -157,6 +157,16 @@ marker_style=None, zorder=2, ), + DetectionType.SIGN: PlotConfig( + fill_color=NEW_TAB_10[8], + fill_color_alpha=1.0, + line_color=BLACK, + line_color_alpha=1.0, + line_width=1.0, + line_style="-", + marker_style=None, + zorder=2, + ), } EGO_VEHICLE_CONFIG: PlotConfig = PlotConfig( diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index f772c096..d1bfba39 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -57,6 +57,7 @@ def __init__( port: int = 8080, label: str = "D123 Viser Server", ): + assert len(scenes) > 0, "At least one scene must be provided." self.scenes = scenes self.scene_index = scene_index diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 347a7ac6..df94cc6f 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -21,7 +21,8 @@ from d123.common.datatypes.vehicle_state.vehicle_parameters import get_wopd_pacifica_parameters from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index -from d123.common.geometry.transform.se3 import get_rotation_matrix +from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array, get_rotation_matrix from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter @@ -36,22 +37,23 @@ "red": TrafficLightStatus.RED, "unknown": TrafficLightStatus.UNKNOWN, } -NUPLAN_DETECTION_NAME_DICT = { - "vehicle": DetectionType.VEHICLE, - "bicycle": DetectionType.BICYCLE, - "pedestrian": DetectionType.PEDESTRIAN, - "traffic_cone": DetectionType.TRAFFIC_CONE, - "barrier": DetectionType.BARRIER, - "czone_sign": DetectionType.CZONE_SIGN, - "generic_object": DetectionType.GENERIC_OBJECT, + +# https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/label.proto#L63 +WOPD_DETECTION_NAME_DICT = { + 0: DetectionType.GENERIC_OBJECT, # TYPE_UNKNOWN + 1: DetectionType.VEHICLE, # TYPE_VEHICLE + 2: DetectionType.PEDESTRIAN, # TYPE_PEDESTRIAN + 3: DetectionType.SIGN, # TYPE_SIGN + 4: DetectionType.BICYCLE, # TYPE_CYCLIST } +# https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L50 WOPD_CAMERA_TYPES = { - CameraType.CAM_F0: 1, # front_camera - CameraType.CAM_L0: 2, # front_left_camera - CameraType.CAM_R0: 3, # front_right_camera - CameraType.CAM_L1: 4, # left_camera - CameraType.CAM_R1: 5, # right_camera + 1: CameraType.CAM_F0, # front_camera + 2: CameraType.CAM_L0, # front_left_camera + 3: CameraType.CAM_R0, # front_right_camera + 4: CameraType.CAM_L1, # left_camera + 5: CameraType.CAM_R1, # right_camera } WOPD_DATA_ROOT = Path("/media/nvme1/waymo_perception") # TODO: set as environment variable @@ -120,10 +122,7 @@ def convert_logs(self, worker: WorkerPool) -> None: worker_map( worker, - partial( - convert_wopd_tfrecord_log_to_arrow, - data_converter_config=self.data_converter_config, - ), + partial(convert_wopd_tfrecord_log_to_arrow, data_converter_config=self.data_converter_config), log_args, ) @@ -178,7 +177,7 @@ def convert_wopd_tfrecord_log_to_arrow( # TODO: Adjust how cameras are added if data_converter_config.camera_store_option is not None: - for camera_type in WOPD_CAMERA_TYPES.keys(): + for camera_type in WOPD_CAMERA_TYPES.values(): if data_converter_config.camera_store_option == "path": raise NotImplementedError("Path camera storage is not implemented.") elif data_converter_config.camera_store_option == "binary": @@ -220,10 +219,9 @@ def convert_wopd_tfrecord_log_to_arrow( def get_wopd_camera_metadata(initial_frame: dataset_pb2.Frame) -> Dict[str, CameraMetadata]: - name_to_camera_type = {v: k for k, v in WOPD_CAMERA_TYPES.items()} cam_metadatas: Dict[str, CameraMetadata] = {} for calibration in initial_frame.context.camera_calibrations: - camera_type = name_to_camera_type[calibration.name] + camera_type = WOPD_CAMERA_TYPES[calibration.name] # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L96 # https://github.com/waymo-research/waymo-open-dataset/issues/834#issuecomment-2134995440 @@ -231,7 +229,7 @@ def get_wopd_camera_metadata(initial_frame: dataset_pb2.Frame) -> Dict[str, Came _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) _distortions = np.array([k1, k2, p1, p2, k3]) - if camera_type in WOPD_CAMERA_TYPES: + if camera_type in WOPD_CAMERA_TYPES.values(): cam_metadatas[camera_type.serialize()] = CameraMetadata( camera_type=camera_type, width=calibration.width, @@ -259,12 +257,7 @@ def _write_recording_table( frame = dataset_pb2.Frame() frame.ParseFromString(data.numpy()) - ( - detections_state, - detections_velocity, - detections_token, - detections_types, - ) = _extract_detections(frame) + (detections_state, detections_velocity, detections_token, detections_types) = _extract_detections(frame) # traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) # route_lane_group_ids = [ # int(roadblock_id) @@ -277,11 +270,6 @@ def _write_recording_table( traffic_light_types = [] # TODO: Implement detections - detections_state = [] - detections_velocity = [] - detections_token = [] - detections_types = [] - row_data = { "token": [create_token(f"{frame.context.name}_{int(frame.timestamp_micros)}")], "timestamp": [int(frame.timestamp_micros)], @@ -321,13 +309,47 @@ def _write_recording_table( def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: - detections_state: List[List[float]] = [] - detections_velocity: List[List[float]] = [] + # TODO: implement + + ego_rear_axle = StateSE3.from_matrix(np.array(frame.pose.transform).reshape(4, 4)) + + num_detections = len(frame.laser_labels) + detections_state = np.zeros((num_detections, len(BoundingBoxSE3Index)), dtype=np.float64) + detections_velocity = np.zeros((num_detections, len(Vector3DIndex)), dtype=np.float64) detections_token: List[str] = [] detections_types: List[int] = [] - # TODO: implement - return detections_state, detections_velocity, detections_token, detections_types + for detection_idx, detection in enumerate(frame.laser_labels): + if detection.type not in WOPD_DETECTION_NAME_DICT: + continue + + # 1. SS3 Bounding Box + detections_state[detection_idx, BoundingBoxSE3Index.X] = detection.box.center_x + detections_state[detection_idx, BoundingBoxSE3Index.Y] = detection.box.center_y + detections_state[detection_idx, BoundingBoxSE3Index.Z] = detection.box.center_z + detections_state[detection_idx, BoundingBoxSE3Index.ROLL] = DEFAULT_ROLL # not provided in WOPD + detections_state[detection_idx, BoundingBoxSE3Index.PITCH] = DEFAULT_PITCH # not provided in WOPD + detections_state[detection_idx, BoundingBoxSE3Index.YAW] = detection.box.heading + detections_state[detection_idx, BoundingBoxSE3Index.LENGTH] = detection.box.length + detections_state[detection_idx, BoundingBoxSE3Index.WIDTH] = detection.box.width + detections_state[detection_idx, BoundingBoxSE3Index.HEIGHT] = detection.box.height + + # 2. Velocity TODO: check if velocity needs to be rotated + detections_velocity[detection_idx] = Vector3D( + x=detection.metadata.speed_x, + y=detection.metadata.speed_y, + z=detection.metadata.speed_z, + ).array + + # 3. Type and track token + detections_token.append(str(detection.id)) + detections_types.append(int(WOPD_DETECTION_NAME_DICT[detection.type])) + + detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( + origin=ego_rear_axle, se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3] + ) + + return detections_state.tolist(), detections_velocity.tolist(), detections_token, detections_types def _extract_ego_state(frame: dataset_pb2.Frame) -> List[float]: @@ -372,9 +394,8 @@ def _extract_camera( data_converter_config: DataConverterConfig, ) -> Dict[CameraType, Union[str, bytes]]: - name_to_camera_type = {v: k for k, v in WOPD_CAMERA_TYPES.items()} camera_dict: Dict[str, Union[str, bytes]] = {} # TODO: Fix wrong type hint - ego_global_transform = np.array(frame.pose.transform).reshape(4, 4) + np.array(frame.pose.transform).reshape(4, 4) # NOTE: The extrinsic matrix in frame.context.camera_calibration is fixed to model the ego to camera transformation. # The poses in frame.images[idx] are the motion compensated ego poses when the camera triggers. @@ -382,7 +403,7 @@ def _extract_camera( context_extrinsic: Dict[str, npt.NDArray] = {} for calibration in frame.context.camera_calibrations: - camera_type = name_to_camera_type[calibration.name] + camera_type = WOPD_CAMERA_TYPES[calibration.name] transform = np.array(calibration.extrinsic.transform).reshape(4, 4) @@ -395,14 +416,14 @@ def _extract_camera( context_extrinsic[camera_type] = transform for image_proto in frame.images: - camera_type = name_to_camera_type[image_proto.name] + camera_type = WOPD_CAMERA_TYPES[image_proto.name] - ego_at_trigger_transform = np.array(image_proto.pose.transform).reshape(4, 4) + np.array(image_proto.pose.transform).reshape(4, 4) camera_bytes = image_proto.image - # Compute the transform from ego_global_transform to ego_at_camera_transform - # ego_global_transform * T = ego_at_camera_transform => T = ego_global_transform^-1 * ego_at_camera_transform - np.linalg.inv(ego_global_transform) @ ego_at_trigger_transform + # # Compute the transform from ego_global_transform to ego_at_camera_transform + # # ego_global_transform * T = ego_at_camera_transform => T = ego_global_transform^-1 * ego_at_camera_transform + # np.linalg.inv(ego_global_transform) @ ego_at_trigger_transform # TODO: figure out the correct transform camera_dict[camera_type] = camera_bytes, context_extrinsic[camera_type].flatten().tolist() diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb new file mode 100644 index 00000000..1bb8217c --- /dev/null +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -0,0 +1,139 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "splits = [\"wopd_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "\n", + "# splits = [\"nuplan_private_test\"]\n", + "log_names = None\n", + "\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=19,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=20,\n", + " shuffle=False,\n", + " camera_types=[CameraType.CAM_F0],\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Tuple\n", + "import matplotlib.pyplot as plt\n", + "from d123.common.visualization.matplotlib.observation import add_box_detections_to_ax, add_default_map_on_ax, add_ego_vehicle_to_ax, add_traffic_lights_to_ax\n", + "from d123.common.visualization.matplotlib.plots import plot_scene_at_iteration\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", + "\n", + " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "\n", + " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", + " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=route_lane_group_ids)\n", + " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", + "\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + "\n", + " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", + " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", + "\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")\n", + " return ax\n", + "\n", + "\n", + "def plot_scene_at_iteration(\n", + " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", + ") -> Tuple[plt.Figure, plt.Axes]:\n", + "\n", + " fig, ax = plt.subplots(figsize=(10, 10))\n", + " _plot_scene_on_ax(ax, scene, iteration, radius)\n", + " return fig, ax\n", + "\n", + "\n", + "scene_index = 0\n", + "plot_scene_at_iteration(scenes[scene_index], iteration=0,radius=80)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index 87b7b205..8259b8c3 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -36,7 +36,16 @@ "\n", "# splits = [\"carla\"]\n", "# log_names = None\n", - "splits = [\"nuplan_private_test\"]\n", + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "splits = [\"wopd_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "\n", + "# splits = [\"nuplan_private_test\"]\n", "log_names = None\n", "\n", "scene_tokens = None\n", @@ -45,9 +54,9 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=20.1,\n", - " history_s=1.0,\n", - " timestamp_threshold_s=20.1,\n", + " duration_s=19,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=20,\n", " shuffle=False,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", @@ -186,7 +195,7 @@ " fig.tight_layout()\n", " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "plot_cameras_with_detections(scenes[70], iteration=0)\n", + "plot_cameras_with_detections(scenes[3], iteration=0)\n", "\n", "\n" ] diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index caefb11e..4ff1a159 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -40,7 +40,7 @@ " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", " duration_s=19,\n", - " history_s=1.0,\n", + " history_s=0.0,\n", " timestamp_threshold_s=20,\n", " shuffle=False,\n", " camera_types=[CameraType.CAM_F0],\n", @@ -67,10 +67,16 @@ "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" ] }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [] + }, { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], "source": [] @@ -78,7 +84,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [] From 1085f74a221f19acb7e798e6ba5bae3ee9cedcff Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Mon, 4 Aug 2025 13:19:58 +0200 Subject: [PATCH 17/42] Fix euler angles rotation order to intrinsic Tait-Bryan angles following the z-y'-x'' convention. --- d123/common/geometry/transform/rotation.py | 27 ++ d123/common/geometry/transform/se3.py | 30 +- d123/common/visualization/viser/server.py | 2 +- d123/common/visualization/viser/utils.py | 11 +- d123/dataset/arrow/conversion.py | 6 +- .../wopd/wopd_data_converter.py | 91 ++++-- .../dataset_specific/wopd/wopd_utils.py | 76 +++++ d123/script/config/datasets/wopd_dataset.yaml | 2 +- notebooks/viz/camera_matplotlib.ipynb | 92 +++--- notebooks/viz/viser_testing_v2_scene.ipynb | 11 +- .../waymo_perception/lidar_testing.ipynb | 301 ++++++++++++++++++ notebooks/waymo_perception/testing.ipynb | 282 ++++++---------- 12 files changed, 661 insertions(+), 270 deletions(-) create mode 100644 d123/common/geometry/transform/rotation.py create mode 100644 d123/dataset/dataset_specific/wopd/wopd_utils.py create mode 100644 notebooks/waymo_perception/lidar_testing.ipynb diff --git a/d123/common/geometry/transform/rotation.py b/d123/common/geometry/transform/rotation.py new file mode 100644 index 00000000..03072609 --- /dev/null +++ b/d123/common/geometry/transform/rotation.py @@ -0,0 +1,27 @@ +# import numpy as np +# import numpy.typing as npt + +# from d123.common.geometry.base import Point3DIndex, StateSE3, StateSE3Index +# from d123.common.geometry.vector import Vector3D + + +# def get_roll_pitch_yaw_from_rotation_matrix( +# rotation_matrix: npt.NDArray[np.float64], +# ) -> Vector3D: +# """Extract roll, pitch, and yaw angles from a rotation matrix.""" +# assert rotation_matrix.shape == (3, 3), "Rotation matrix must be 3x3." + +# sy = np.sqrt(rotation_matrix[0, 0] ** 2 + rotation_matrix[1, 0] ** 2) + +# singular = sy < 1e-6 + +# if not singular: +# x = np.arctan2(rotation_matrix[2, 1], rotation_matrix[2, 2]) +# y = np.arctan2(-rotation_matrix[2, 0], sy) +# z = np.arctan2(rotation_matrix[1, 0], rotation_matrix[0, 0]) +# else: +# x = np.arctan2(-rotation_matrix[1, 2], rotation_matrix[1, 1]) +# y = np.arctan2(-rotation_matrix[2, 0], sy) +# z = 0.0 + +# return Vector3D(x=x, y=y, z=z) diff --git a/d123/common/geometry/transform/se3.py b/d123/common/geometry/transform/se3.py index 7dc38c5c..e61451b4 100644 --- a/d123/common/geometry/transform/se3.py +++ b/d123/common/geometry/transform/se3.py @@ -4,8 +4,36 @@ from d123.common.geometry.base import Point3DIndex, StateSE3, StateSE3Index from d123.common.geometry.vector import Vector3D +# def get_rotation_matrix(state_se3: StateSE3) -> npt.NDArray[np.float64]: +# R_x = np.array( +# [ +# [1, 0, 0], +# [0, np.cos(state_se3.roll), -np.sin(state_se3.roll)], +# [0, np.sin(state_se3.roll), np.cos(state_se3.roll)], +# ], +# dtype=np.float64, +# ) +# R_y = np.array( +# [ +# [np.cos(state_se3.pitch), 0, np.sin(state_se3.pitch)], +# [0, 1, 0], +# [-np.sin(state_se3.pitch), 0, np.cos(state_se3.pitch)], +# ], +# dtype=np.float64, +# ) +# R_z = np.array( +# [ +# [np.cos(state_se3.yaw), -np.sin(state_se3.yaw), 0], +# [np.sin(state_se3.yaw), np.cos(state_se3.yaw), 0], +# [0, 0, 1], +# ], +# dtype=np.float64, +# ) +# return R_z @ R_y @ R_x + def get_rotation_matrix(state_se3: StateSE3) -> npt.NDArray[np.float64]: + # Intrinsic Z-Y'-X'' rotation: R = R_x(roll) @ R_y(pitch) @ R_z(yaw) R_x = np.array( [ [1, 0, 0], @@ -30,7 +58,7 @@ def get_rotation_matrix(state_se3: StateSE3) -> npt.NDArray[np.float64]: ], dtype=np.float64, ) - return R_z @ R_y @ R_x + return R_x @ R_y @ R_z def translate_se3_along_z(state_se3: StateSE3, distance: float) -> StateSE3: diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index d1bfba39..291f076b 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -45,7 +45,7 @@ # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0] VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] -LIDAR_AVAILABLE: bool = False +LIDAR_AVAILABLE: bool = True class ViserVisualizationServer: diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 35b0c39b..a3a3577e 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -185,27 +185,24 @@ def get_camera_values( ego_transform[:3, :3] = get_rotation_matrix(rear_axle) ego_transform[:3, 3] = rear_axle.point_3d.array - DEBUG = False + DEBUG = True if DEBUG: print("DEBUG") camera_to_ego = camera.extrinsic - rotation = camera_to_ego[:3, :3] flip_camera = get_rotation_matrix( StateSE3( x=0.0, y=0.0, z=0.0, - roll=np.deg2rad(-90.0), - pitch=np.deg2rad(0.0), + roll=np.deg2rad(0.0), + pitch=np.deg2rad(90.0), yaw=np.deg2rad(-90.0), ) ) + camera_to_ego[:3, :3] = camera_to_ego[:3, :3] @ flip_camera - corrected_rotation = camera_to_ego[:3, :3] @ flip_camera - - camera_to_ego[:3, :3] = corrected_rotation camera_transform = ego_transform @ camera_to_ego # Camera transformation in ego frame diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index a671e2d7..6e4dbb4a 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -148,6 +148,10 @@ def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int, log_metadata: raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") else: - raise NotImplementedError("Only string file paths for lidar data are supported.") + if log_metadata.dataset == "wopd": + lidar_data = np.array(lidar_data, dtype=np.float64) + lidar = LiDAR(point_cloud=lidar_data[:, 3:].T) + else: + raise NotImplementedError("Only string file paths for lidar data are supported.") return lidar diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index df94cc6f..16aff2c0 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -1,10 +1,11 @@ import gc import hashlib import json +import os from dataclasses import asdict from functools import partial from pathlib import Path -from typing import Any, Dict, Final, List, Optional, Tuple, Union +from typing import Any, Dict, Final, List, Literal, Optional, Tuple, Union import numpy as np import numpy.typing as npt @@ -26,8 +27,11 @@ from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.dataset_specific.wopd.wopd_utils import parse_range_image_and_camera_projection from d123.dataset.logs.log_metadata import LogMetadata +os.environ["CUDA_VISIBLE_DEVICES"] = "-1" + TARGET_DT: Final[float] = 0.1 NUPLAN_DT: Final[float] = 0.05 SORT_BY_TIMESTAMP: Final[bool] = True @@ -58,6 +62,9 @@ WOPD_DATA_ROOT = Path("/media/nvme1/waymo_perception") # TODO: set as environment variable +# Whether to use ego or zero roll and pitch values for bounding box detections (after global conversion) +DETECTION_ROLL_PITCH: Final[Literal["ego", "zero"]] = "zero" + def create_token(input_data: str) -> str: # TODO: Refactor this function. @@ -171,9 +178,9 @@ def convert_wopd_tfrecord_log_to_arrow( ] if data_converter_config.lidar_store_option is not None: if data_converter_config.lidar_store_option == "path": - schema_column_list.append(("lidar", pa.string())) + raise NotImplementedError("Filepath lidar storage is not implemented.") elif data_converter_config.lidar_store_option == "binary": - raise NotImplementedError("Binary lidar storage is not implemented.") + schema_column_list.append(("lidar", pa.list_(pa.list_(pa.float32(), 6)))) # TODO: Adjust how cameras are added if data_converter_config.camera_store_option is not None: @@ -286,7 +293,7 @@ def _write_recording_table( # TODO: Implement lidar extraction if data_converter_config.lidar_store_option is not None: - row_data["lidar"] = [None] + row_data["lidar"] = [_extract_lidar(frame, data_converter_config).tolist()] if data_converter_config.camera_store_option is not None: camera_data_dict = _extract_camera(frame, data_converter_config) @@ -308,10 +315,30 @@ def _write_recording_table( write_arrow_table(recording_table, log_file_path) +# def _get_ego_pose_se3(frame: dataset_pb2.Frame) -> StateSE3: +# ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) +# yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll +# ego_point_3d = Point3D.from_array(ego_pose_matrix[:3, 3]) + +# # TODO: figure out if ego frame is given in rear axle or center frame +# return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=pitch, pitch=-roll, yaw=yaw) + + +def _get_ego_pose_se3(frame: dataset_pb2.Frame) -> StateSE3: + ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) + yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll + ego_point_3d = Point3D.from_array(ego_pose_matrix[:3, 3]) + + return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=roll, pitch=pitch, yaw=yaw) + + # TODO: figure out if ego frame is given in rear axle or center frame + # return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=pitch, pitch=-roll, yaw=yaw) + + def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: # TODO: implement - ego_rear_axle = StateSE3.from_matrix(np.array(frame.pose.transform).reshape(4, 4)) + ego_rear_axle = _get_ego_pose_se3(frame) num_detections = len(frame.laser_labels) detections_state = np.zeros((num_detections, len(BoundingBoxSE3Index)), dtype=np.float64) @@ -348,27 +375,21 @@ def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], Li detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( origin=ego_rear_axle, se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3] ) + if DETECTION_ROLL_PITCH == "ego": + pass + if DETECTION_ROLL_PITCH == "zero": + detections_state[:, BoundingBoxSE3Index.ROLL] = DEFAULT_ROLL + detections_state[:, BoundingBoxSE3Index.PITCH] = DEFAULT_PITCH + else: + raise ValueError(f"Invalid DETECTION_ROLL_PITCH value: {DETECTION_ROLL_PITCH}. Must be 'ego' or 'zero'.") return detections_state.tolist(), detections_velocity.tolist(), detections_token, detections_types def _extract_ego_state(frame: dataset_pb2.Frame) -> List[float]: - - ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) - yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll - ego_point_3d = Point3D.from_array(ego_pose_matrix[:3, 3]) + rear_axle_pose = _get_ego_pose_se3(frame) vehicle_parameters = get_wopd_pacifica_parameters() - - # TODO: figure out if ego frame is given in rear axle or center frame - rear_axle_pose = StateSE3( - x=ego_point_3d.x, - y=ego_point_3d.y, - z=ego_point_3d.z, - roll=roll, - pitch=pitch, - yaw=yaw, - ) # FIXME: Find dynamic state in waymo open perception dataset # https://github.com/waymo-research/waymo-open-dataset/issues/55#issuecomment-546152290 dynamic_state = DynamicStateSE3( @@ -390,8 +411,7 @@ def _extract_traffic_lights() -> Tuple[List[int], List[int]]: def _extract_camera( - frame: dataset_pb2.Frame, - data_converter_config: DataConverterConfig, + frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig ) -> Dict[CameraType, Union[str, bytes]]: camera_dict: Dict[str, Union[str, bytes]] = {} # TODO: Fix wrong type hint @@ -409,7 +429,14 @@ def _extract_camera( # FIXME: This is an ugly hack to convert to uniform camera convention. flip_camera = get_rotation_matrix( - StateSE3(x=0.0, y=0.0, z=0.0, roll=np.deg2rad(-90.0), pitch=np.deg2rad(0.0), yaw=np.deg2rad(-90.0)) + StateSE3( + x=0.0, + y=0.0, + z=0.0, + roll=0.0, + pitch=0.0, + yaw=np.deg2rad(0.0), + ) ) transform[:3, :3] = transform[:3, :3] @ flip_camera @@ -431,6 +458,20 @@ def _extract_camera( return camera_dict -def _extract_lidar() -> Optional[str]: - lidar: Optional[str] = None - return lidar +def _extract_lidar( + frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Optional[npt.NDArray[np.float32]]: + from waymo_open_dataset.utils import frame_utils + + assert data_converter_config.lidar_store_option == "binary", "Lidar store option must be 'binary' for WOPD." + (range_images, camera_projections, _, range_image_top_pose) = parse_range_image_and_camera_projection(frame) + + points, cp_points = frame_utils.convert_range_image_to_point_cloud( + frame=frame, + range_images=range_images, + camera_projections=camera_projections, + range_image_top_pose=range_image_top_pose, + keep_polar_features=True, + ) + points = np.array(points[0], dtype=np.float32) + return points diff --git a/d123/dataset/dataset_specific/wopd/wopd_utils.py b/d123/dataset/dataset_specific/wopd/wopd_utils.py new file mode 100644 index 00000000..c9c302ef --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/wopd_utils.py @@ -0,0 +1,76 @@ +from typing import Dict, List, Tuple + +import tensorflow as tf +from waymo_open_dataset import dataset_pb2 + +RangeImages = Dict["dataset_pb2.LaserName.Name", List[dataset_pb2.MatrixFloat]] +CameraProjections = Dict["dataset_pb2.LaserName.Name", List[dataset_pb2.MatrixInt32]] +SegmentationLabels = Dict["dataset_pb2.LaserName.Name", List[dataset_pb2.MatrixInt32]] +ParsedFrame = Tuple[RangeImages, CameraProjections, SegmentationLabels, dataset_pb2.MatrixFloat] + + +def parse_range_image_and_camera_projection(frame: dataset_pb2.Frame) -> ParsedFrame: + """Parse range images and camera projections given a frame. + + Args: + frame: open dataset frame proto + + Returns: + range_images: A dict of {laser_name, + [range_image_first_return, range_image_second_return]}. + camera_projections: A dict of {laser_name, + [camera_projection_from_first_return, + camera_projection_from_second_return]}. + seg_labels: segmentation labels, a dict of {laser_name, + [seg_label_first_return, seg_label_second_return]} + range_image_top_pose: range image pixel pose for top lidar. + """ + range_images = {} + camera_projections = {} + seg_labels = {} + range_image_top_pose: dataset_pb2.MatrixFloat = dataset_pb2.MatrixFloat() + for laser in frame.lasers: + if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed(laser.ri_return1.range_image_compressed, "ZLIB") + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(range_image_str_tensor.numpy()) + range_images[laser.name] = [ri] + + if laser.name == dataset_pb2.LaserName.TOP: + range_image_top_pose_str_tensor = tf.io.decode_compressed( + laser.ri_return1.range_image_pose_compressed, "ZLIB" + ) + range_image_top_pose = dataset_pb2.MatrixFloat() + range_image_top_pose.ParseFromString(range_image_top_pose_str_tensor.numpy()) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return1.camera_projection_compressed, "ZLIB" + ) + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(camera_projection_str_tensor.numpy()) + camera_projections[laser.name] = [cp] + + if len(laser.ri_return1.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test + seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return1.segmentation_label_compressed, "ZLIB") + seg_label = dataset_pb2.MatrixInt32() + seg_label.ParseFromString(seg_label_str_tensor.numpy()) + seg_labels[laser.name] = [seg_label] + if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test + range_image_str_tensor = tf.io.decode_compressed(laser.ri_return2.range_image_compressed, "ZLIB") + ri = dataset_pb2.MatrixFloat() + ri.ParseFromString(range_image_str_tensor.numpy()) + range_images[laser.name].append(ri) + + camera_projection_str_tensor = tf.io.decode_compressed( + laser.ri_return2.camera_projection_compressed, "ZLIB" + ) + cp = dataset_pb2.MatrixInt32() + cp.ParseFromString(camera_projection_str_tensor.numpy()) + camera_projections[laser.name].append(cp) + + if len(laser.ri_return2.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test + seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return2.segmentation_label_compressed, "ZLIB") + seg_label = dataset_pb2.MatrixInt32() + seg_label.ParseFromString(seg_label_str_tensor.numpy()) + seg_labels[laser.name].append(seg_label) + return range_images, camera_projections, seg_labels, range_image_top_pose diff --git a/d123/script/config/datasets/wopd_dataset.yaml b/d123/script/config/datasets/wopd_dataset.yaml index 9234dbfa..1abb0381 100644 --- a/d123/script/config/datasets/wopd_dataset.yaml +++ b/d123/script/config/datasets/wopd_dataset.yaml @@ -13,4 +13,4 @@ wopd_dataset: force_log_conversion: ${force_log_conversion} force_map_conversion: ${force_map_conversion} camera_store_option: "binary" - lidar_store_option: null + lidar_store_option: "binary" diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index 8259b8c3..4497ce3f 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -34,13 +34,11 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "# splits = [\"carla\"]\n", - "# log_names = None\n", - "# split = \"nuplan_private_test\"\n", - "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "splits = [\"wopd_train\"]\n", + "# splits = [\"carla\"]\n", + "splits = [\"nuplan_private_test\"]\n", + "# splits = [\"wopd_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -127,48 +125,16 @@ "metadata": {}, "outputs": [], "source": [ - "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", - "\n", - "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", - "# CameraType.CAM_L0: ax[0, 0],\n", - "# CameraType.CAM_F0: ax[0, 1],\n", - "# CameraType.CAM_R0: ax[0, 2],\n", - "# CameraType.CAM_L2: ax[1, 0],\n", - "# CameraType.CAM_B0: ax[1, 1],\n", - "# CameraType.CAM_R2: ax[1, 2],\n", - "# }\n", - "\n", - "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", - "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", - "# for camera_type, camera_ax in camera_ax_assignment.items():\n", - "# assert camera_type in scene.available_camera_types\n", - "# camera = scene.get_camera_at_iteration(iteration, camera_type)\n", - "# if camera is not None:\n", - "# add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", - "\n", - "# return ax\n", - "\n", - "\n", - "# def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", - "# \"\"\"\n", - "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", - "# :param scene: navsim scene dataclass\n", - "# :param frame_idx: index of selected frame\n", - "# :return: figure and ax object of matplotlib\n", - "# \"\"\"\n", - "# scale = 2\n", - "# fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", - "# _add_camera_with_detections(ax, scene, iteration)\n", - "\n", - "# fig.tight_layout()\n", - "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", - "\n", - "# plot_cameras_with_detections(scenes[4], iteration=10)\n", - "\n", - "\n", "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", "\n", - " camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L2: ax[1, 0],\n", + " CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R2: ax[1, 2],\n", + " }\n", "\n", " box_detections = scene.get_box_detections_at_iteration(iteration)\n", " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", @@ -189,13 +155,45 @@ " :return: figure and ax object of matplotlib\n", " \"\"\"\n", " scale = 2\n", - " fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", + " fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", " _add_camera_with_detections(ax, scene, iteration)\n", "\n", " fig.tight_layout()\n", " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "plot_cameras_with_detections(scenes[3], iteration=0)\n", + "plot_cameras_with_detections(scenes[4], iteration=10)\n", + "\n", + "\n", + "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", + "\n", + "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + "# for camera_type, camera_ax in camera_ax_assignment.items():\n", + "# assert camera_type in scene.available_camera_types\n", + "# camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + "# if camera is not None:\n", + "# add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + "# return ax\n", + "\n", + "\n", + "# def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "# \"\"\"\n", + "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + "# :param scene: navsim scene dataclass\n", + "# :param frame_idx: index of selected frame\n", + "# :return: figure and ax object of matplotlib\n", + "# \"\"\"\n", + "# scale = 2\n", + "# fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", + "# _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + "# fig.tight_layout()\n", + "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "# plot_cameras_with_detections(scenes[4], iteration=150)\n", "\n", "\n" ] diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 4ff1a159..4a105e1b 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -31,6 +31,7 @@ "\n", "\n", "# splits = [\"nuplan_private_test\"]\n", + "# splits = [\"carla\"]\n", "log_names = None\n", "\n", "scene_tokens = None\n", @@ -60,11 +61,10 @@ "metadata": {}, "outputs": [], "source": [ - "\n", "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=2)" ] }, { @@ -79,7 +79,12 @@ "id": "4", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "ego_state = scenes[2].get_ego_state_at_iteration(0)\n", + "\n", + "\n", + "ego_state.center" + ] }, { "cell_type": "code", diff --git a/notebooks/waymo_perception/lidar_testing.ipynb b/notebooks/waymo_perception/lidar_testing.ipynb new file mode 100644 index 00000000..3cd2a77e --- /dev/null +++ b/notebooks/waymo_perception/lidar_testing.ipynb @@ -0,0 +1,301 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset import dataset_pb2\n", + "\n", + "import json\n", + "import os\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "from PIL import Image\n", + "from tqdm import tqdm\n", + "from waymo_open_dataset import label_pb2\n", + "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", + "from waymo_open_dataset.utils import box_utils\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "\n", + "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", + "\n", + "\n", + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from pyquaternion import Quaternion\n", + "\n", + "from d123.common.geometry.base import StateSE3\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", + "\n", + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "\n", + "# Frame attributes:\n", + "# context: \n", + "# timestamp_micros: \n", + "# pose: \n", + "# images: List with 5 images\n", + "# lasers: \n", + "# Length: 5\n", + "# laser_labels: \n", + "# Length: 0\n", + "# projected_lidar_labels: \n", + "# Length: 0\n", + "# camera_labels: \n", + "# Length: 0\n", + "# no_label_zones: \n", + "# Length: 0\n", + "# map_features: \n", + "# Length: 0\n", + "# map_pose_offset: \n", + "\n", + "file_idx = 0\n", + "pathname = tfrecords_file_list[file_idx]\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "num_frames = sum(1 for _ in dataset)\n", + "\n", + "\n", + "def read_jpg_image(data: bytes) -> np.ndarray:\n", + " \"\"\"Read a JPEG image from bytes and return it as a numpy array.\"\"\"\n", + " image = Image.open(io.BytesIO(data))\n", + " return np.array(image)\n", + "\n", + "\n", + "ego_state_se3s = []\n", + "front_images = []\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "\n", + "boxes = []\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + "\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " # print(frame.camera_labels)\n", + " print(\"Frame attributes:\")\n", + " for field in frame.DESCRIPTOR.fields:\n", + " field_name = field.name\n", + " if hasattr(frame, field_name):\n", + " value = getattr(frame, field_name)\n", + " if field_name != \"images\": # Don't print the whole image data\n", + " print(f\" {field_name}: {type(value)}\")\n", + " if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", + " print(f\" Length: {len(value)}\")\n", + " else:\n", + " print(f\" {field_name}: List with {len(value)} images\")\n", + "\n", + "\n", + " # # 1. pose\n", + " pose = np.array(frame.pose.transform).reshape(4, 4)\n", + " yaw_pitch_roll = Quaternion(matrix=pose[:3, :3]).yaw_pitch_roll\n", + " ego_state_se3s.append(\n", + " np.array(\n", + " [\n", + " pose[0, 3], # x\n", + " pose[1, 3], # y\n", + " pose[2, 3], # z\n", + " yaw_pitch_roll[2], # yaw\n", + " yaw_pitch_roll[1], # pitch\n", + " yaw_pitch_roll[0], # roll\n", + " ],\n", + " dtype=np.float64,\n", + " )\n", + " )\n", + "\n", + " # # plt.show()\n", + " if frame_idx == 0:\n", + " break\n", + "\n", + "ego_state_se3s = np.array(ego_state_se3s, dtype=np.float64)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Dict, List, Tuple\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "\n", + "from waymo_open_dataset import dataset_pb2\n", + "from waymo_open_dataset.utils import range_image_utils\n", + "from waymo_open_dataset.utils import transform_utils\n", + "\n", + "RangeImages = Dict[\"dataset_pb2.LaserName.Name\", List[dataset_pb2.MatrixFloat]]\n", + "CameraProjections = Dict[\"dataset_pb2.LaserName.Name\", List[dataset_pb2.MatrixInt32]]\n", + "SegmentationLabels = Dict[\"dataset_pb2.LaserName.Name\", List[dataset_pb2.MatrixInt32]]\n", + "ParsedFrame = Tuple[RangeImages, CameraProjections, SegmentationLabels, dataset_pb2.MatrixFloat]\n", + "\n", + "\n", + "def parse_range_image_and_camera_projection(frame: dataset_pb2.Frame) -> ParsedFrame:\n", + " \"\"\"Parse range images and camera projections given a frame.\n", + "\n", + " Args:\n", + " frame: open dataset frame proto\n", + "\n", + " Returns:\n", + " range_images: A dict of {laser_name,\n", + " [range_image_first_return, range_image_second_return]}.\n", + " camera_projections: A dict of {laser_name,\n", + " [camera_projection_from_first_return,\n", + " camera_projection_from_second_return]}.\n", + " seg_labels: segmentation labels, a dict of {laser_name,\n", + " [seg_label_first_return, seg_label_second_return]}\n", + " range_image_top_pose: range image pixel pose for top lidar.\n", + " \"\"\"\n", + " range_images = {}\n", + " camera_projections = {}\n", + " seg_labels = {}\n", + " range_image_top_pose: dataset_pb2.MatrixFloat = dataset_pb2.MatrixFloat()\n", + " for laser in frame.lasers:\n", + " if len(laser.ri_return1.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " range_image_str_tensor = tf.io.decode_compressed(laser.ri_return1.range_image_compressed, \"ZLIB\")\n", + " ri = dataset_pb2.MatrixFloat()\n", + " ri.ParseFromString(range_image_str_tensor.numpy())\n", + " range_images[laser.name] = [ri]\n", + "\n", + " if laser.name == dataset_pb2.LaserName.TOP:\n", + " range_image_top_pose_str_tensor = tf.io.decode_compressed(\n", + " laser.ri_return1.range_image_pose_compressed, \"ZLIB\"\n", + " )\n", + " range_image_top_pose = dataset_pb2.MatrixFloat()\n", + " range_image_top_pose.ParseFromString(range_image_top_pose_str_tensor.numpy())\n", + "\n", + " camera_projection_str_tensor = tf.io.decode_compressed(\n", + " laser.ri_return1.camera_projection_compressed, \"ZLIB\"\n", + " )\n", + " cp = dataset_pb2.MatrixInt32()\n", + " cp.ParseFromString(camera_projection_str_tensor.numpy())\n", + " camera_projections[laser.name] = [cp]\n", + "\n", + " if len(laser.ri_return1.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return1.segmentation_label_compressed, \"ZLIB\")\n", + " seg_label = dataset_pb2.MatrixInt32()\n", + " seg_label.ParseFromString(seg_label_str_tensor.numpy())\n", + " seg_labels[laser.name] = [seg_label]\n", + " if len(laser.ri_return2.range_image_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " range_image_str_tensor = tf.io.decode_compressed(laser.ri_return2.range_image_compressed, \"ZLIB\")\n", + " ri = dataset_pb2.MatrixFloat()\n", + " ri.ParseFromString(range_image_str_tensor.numpy())\n", + " range_images[laser.name].append(ri)\n", + "\n", + " camera_projection_str_tensor = tf.io.decode_compressed(\n", + " laser.ri_return2.camera_projection_compressed, \"ZLIB\"\n", + " )\n", + " cp = dataset_pb2.MatrixInt32()\n", + " cp.ParseFromString(camera_projection_str_tensor.numpy())\n", + " camera_projections[laser.name].append(cp)\n", + "\n", + " if len(laser.ri_return2.segmentation_label_compressed) > 0: # pylint: disable=g-explicit-length-test\n", + " seg_label_str_tensor = tf.io.decode_compressed(laser.ri_return2.segmentation_label_compressed, \"ZLIB\")\n", + " seg_label = dataset_pb2.MatrixInt32()\n", + " seg_label.ParseFromString(seg_label_str_tensor.numpy())\n", + " seg_labels[laser.name].append(seg_label)\n", + " return range_images, camera_projections, seg_labels, range_image_top_pose\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "dataset = tf.data.TFRecordDataset(tfrecords_file_list[file_idx], compression_type=\"\")\n", + "for data in dataset:\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy()) # No need for bytearray conversion\n", + " break\n", + "\n", + "(range_images, camera_projections, _, range_image_top_pose) = parse_range_image_and_camera_projection(frame)\n", + "points, cp_points = frame_utils.convert_range_image_to_point_cloud(\n", + " frame=frame,\n", + " range_images=range_images,\n", + " camera_projections=camera_projections,\n", + " range_image_top_pose=range_image_top_pose,\n", + " keep_polar_features=True,\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "pc = points[0]\n", + "\n", + "\n", + "plt.scatter(pc[:, 3], pc[:, 4], s=0.1, c=pc[:, 5], cmap='viridis')" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "cp_points[0].shape, points[0].shape" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/waymo_perception/testing.ipynb b/notebooks/waymo_perception/testing.ipynb index bd035f8f..6ee57cf5 100644 --- a/notebooks/waymo_perception/testing.ipynb +++ b/notebooks/waymo_perception/testing.ipynb @@ -19,7 +19,6 @@ "from waymo_open_dataset import label_pb2\n", "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", "from waymo_open_dataset.utils import box_utils\n", - "from waymo_open_dataset.utils.frame_utils import parse_range_image_and_camera_projection\n", "\n", "\n", "import matplotlib.pyplot as plt" @@ -38,7 +37,7 @@ "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", "\n", "\n", - "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))\n" + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))" ] }, { @@ -54,6 +53,8 @@ "from d123.common.geometry.base import StateSE3\n", "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", "\n", + "from waymo_open_dataset.utils import frame_utils\n", + "\n", "\n", "# Frame attributes:\n", "# context: \n", @@ -86,8 +87,6 @@ " return np.array(image)\n", "\n", "\n", - "\n", - "\n", "ego_state_se3s = []\n", "front_images = []\n", "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", @@ -99,56 +98,18 @@ " frame = dataset_pb2.Frame()\n", " frame.ParseFromString(data.numpy())\n", " # print(frame.camera_labels)\n", - " for label in frame.laser_labels:\n", - " boxes.append(\n", - " BoundingBoxSE3(\n", - " center=StateSE3(\n", - " x=label.box.center_x,\n", - " y=label.box.center_y,\n", - " z=label.box.center_z,\n", - " pitch=0.0,\n", - " roll=0.0,\n", - " yaw=label.box.heading,\n", - " ),\n", - " length=label.box.length,\n", - " width=label.box.width,\n", - " height=label.box.height,\n", - " )\n", - " )\n", + " print(\"Frame attributes:\")\n", + " for field in frame.DESCRIPTOR.fields:\n", + " field_name = field.name\n", + " if hasattr(frame, field_name):\n", + " value = getattr(frame, field_name)\n", + " if field_name != \"images\": # Don't print the whole image data\n", + " print(f\" {field_name}: {type(value)}\")\n", + " if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", + " print(f\" Length: {len(value)}\")\n", + " else:\n", + " print(f\" {field_name}: List with {len(value)} images\")\n", "\n", - " print(frame.context)\n", - "\n", - " # Print all attributes of the frame\n", - " # print(\"Frame attributes:\")\n", - " # for field in frame.DESCRIPTOR.fields:\n", - " # field_name = field.name\n", - " # if hasattr(frame, field_name):\n", - " # value = getattr(frame, field_name)\n", - " # if field_name != \"images\": # Don't print the whole image data\n", - " # print(f\" {field_name}: {type(value)}\")\n", - " # if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", - " # print(f\" Length: {len(value)}\")\n", - " # else:\n", - " # print(f\" {field_name}: List with {len(value)} images\")\n", - "\n", - " # Print information about the first image if available\n", - " # if frame.images:\n", - " # print(\"\\nFirst image details:\")\n", - " # first_image = frame.images[0]\n", - " # for field in first_image.DESCRIPTOR.fields:\n", - " # field_name = field.name\n", - " # if hasattr(first_image, field_name):\n", - " # value = getattr(first_image, field_name)\n", - " # if field_name != \"image\": # Don't print the binary data\n", - " # print(f\" {field_name}: {value}\")\n", - " # else:\n", - " # print(f\" {field_name}: binary data of length {len(value)} bytes\")\n", - "\n", - " # for image in frame.images:\n", - " # print(image.name)\n", - "\n", - " # print([image.name for image in frame.images])\n", - " # print(frame.images[0])\n", "\n", " # # 1. pose\n", " pose = np.array(frame.pose.transform).reshape(4, 4)\n", @@ -184,28 +145,21 @@ "for frame_idx, data in enumerate(dataset):\n", " frame = dataset_pb2.Frame()\n", " frame.ParseFromString(data.numpy())\n", - " if frame_idx == 3:\n", + " if frame_idx == 2:\n", " break\n", "\n", "print(\"Ego\")\n", "ego_transform = np.array(frame.pose.transform).reshape(4, 4)\n", "print(ego_transform[:3, 3])\n", "\n", - "print(\"Frame\")\n", - "for image in frame.images:\n", - " image_transform = np.array(image.pose.transform).reshape(4, 4)\n", - " print(image.name, image_transform[:3, 3])\n", - "\n", - "print(\"Context\")\n", - "for image in frame.context.camera_calibrations:\n", - " image_transform = np.array(image.extrinsic.transform).reshape(4, 4)\n", - " print(image.name, image_transform[:3, 3])\n", - "\n", "# 1 [ 1.5441613 -0.02302364 2.11557864]\n", "# 2 [1.49672397 0.0954948 2.11616463]\n", "# 3 [ 1.49442485 -0.09637497 2.11519385]\n", "# 4 [1.43213651 0.11612398 2.11625087]\n", - "# 5 [ 1.42936162 -0.11545043 2.1150792 ]" + "# 5 [ 1.42936162 -0.11545043 2.1150792 ]\n", + "\n", + "\n", + "# frame.map_pose_offset" ] }, { @@ -215,7 +169,85 @@ "metadata": {}, "outputs": [], "source": [ - "\n", + "plt.figure(figsize=(64, 20))\n", + "\n", + "\n", + "def plot_range_image_helper(data, name, layout, vmin=0, vmax=1, cmap=\"gray\"):\n", + " \"\"\"Plots range image.\n", + "\n", + " Args:\n", + " data: range image data\n", + " name: the image title\n", + " layout: plt layout\n", + " vmin: minimum value of the passed data\n", + " vmax: maximum value of the passed data\n", + " cmap: color map\n", + " \"\"\"\n", + " plt.subplot(*layout)\n", + " plt.imshow(data, cmap=cmap, vmin=vmin, vmax=vmax)\n", + " plt.title(name)\n", + " plt.grid(False)\n", + " plt.axis(\"off\")\n", + "\n", + "\n", + "def get_range_image(laser_name, return_index):\n", + " \"\"\"Returns range image given a laser name and its return index.\"\"\"\n", + " return range_images[laser_name][return_index]\n", + "\n", + "\n", + "# def show_range_image(range_image, layout_index_start=1):\n", + "# \"\"\"Shows range image.\n", + "\n", + "# Args:\n", + "# range_image: the range image data from a given lidar of type MatrixFloat.\n", + "# layout_index_start: layout offset\n", + "# \"\"\"\n", + "# range_image_tensor = tf.convert_to_tensor(range_image.data)\n", + "# range_image_tensor = tf.reshape(range_image_tensor, range_image.shape.dims)\n", + "# lidar_image_mask = tf.greater_equal(range_image_tensor, 0)\n", + "# range_image_tensor = tf.where(lidar_image_mask, range_image_tensor, tf.ones_like(range_image_tensor) * 1e10)\n", + "# range_image_range = range_image_tensor[..., 0]\n", + "# range_image_intensity = range_image_tensor[..., 1]\n", + "# range_image_elongation = range_image_tensor[..., 2]\n", + "# plot_range_image_helper(range_image_range.numpy(), \"range\", [8, 1, layout_index_start], vmax=75, cmap=\"gray\")\n", + "# plot_range_image_helper(\n", + "# range_image_intensity.numpy(), \"intensity\", [8, 1, layout_index_start + 1], vmax=1.5, cmap=\"gray\"\n", + "# )\n", + "# plot_range_image_helper(\n", + "# range_image_elongation.numpy(), \"elongation\", [8, 1, layout_index_start + 2], vmax=1.5, cmap=\"gray\"\n", + "# )\n", + "\n", + "\n", + "def show_range_image(range_image, layout_index_start=1):\n", + " \"\"\"Shows range image.\n", + "\n", + " Args:\n", + " range_image: the range image data from a given lidar of type MatrixFloat.\n", + " layout_index_start: layout offset\n", + " \"\"\"\n", + " range_image_tensor = np.array([data for data in range_image.data]).reshape(range_image.shape.dims)\n", + " lidar_image_mask = np.greater_equal(range_image_tensor, 0)\n", + " range_image_tensor = np.where(lidar_image_mask, range_image_tensor, np.ones_like(range_image_tensor) * 1e10)\n", + " range_image_range = range_image_tensor[..., 0]\n", + " range_image_intensity = range_image_tensor[..., 1]\n", + " range_image_elongation = range_image_tensor[..., 2]\n", + " plot_range_image_helper(range_image_range, \"range\", [8, 1, layout_index_start], vmax=75, cmap=\"gray\")\n", + " plot_range_image_helper(range_image_intensity, \"intensity\", [8, 1, layout_index_start + 1], vmax=1.5, cmap=\"gray\")\n", + " plot_range_image_helper(range_image_elongation, \"elongation\", [8, 1, layout_index_start + 2], vmax=1.5, cmap=\"gray\")\n", + "\n", + "\n", + "frame.lasers.sort(key=lambda laser: laser.name)\n", + "show_range_image(get_range_image(open_dataset.LaserName.TOP, 0), 1)\n", + "# show_range_image(get_range_image(open_dataset.LaserName.TOP, 1), 4)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ "from d123.common.datatypes.time.time_point import TimePoint\n", "\n", "\n", @@ -227,7 +259,6 @@ " break\n", "\n", "\n", - "\n", "# for calibration in frame.context.camera_calibrations:\n", "\n", "frame.timestamp_micros, frame.images[0].pose_timestamp\n", @@ -238,129 +269,12 @@ "frame_timestamp.time_s, frame_timestamp.time_s" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [ - "frame = next((dataset_pb2.Frame().FromString(data.numpy()) or dataset_pb2.Frame() for data in dataset), None)\n", - "# if frame is None:\n", - "# raise ValueError(f\"No frames found in TFRecord {tf_record_path}\")\n", - "\n", - "frame" - ] - }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], - "source": [ - "from d123.common.datatypes.detection.detection_types import DetectionType\n", - "from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG\n", - "from d123.common.visualization.matplotlib.observation import add_bounding_box_to_ax\n", - "\n", - "\n", - "ego_rear_axle = StateSE3.from_array(ego_state_se3s[0])\n", - "\n", - "ego_rear_axle = StateSE3.from_array(np.zeros_like(ego_state_se3s[0]))\n", - "\n", - "ego_box = BoundingBoxSE3(center=ego_rear_axle, length=4.0, width=1.8, height=1.6)\n", - "\n", - "plot_config = BOX_DETECTION_CONFIG[DetectionType.VEHICLE]\n", - "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n", - "\n", - "\n", - "for box in boxes:\n", - " add_bounding_box_to_ax(ax, box, plot_config)\n", - "\n", - "\n", - "add_bounding_box_to_ax(ax, ego_box, EGO_VEHICLE_CONFIG)\n", - "ax.set_aspect(\"equal\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "fig, ax = plt.subplots(figsize=(10, 10))\n", - "ax.plot(ego_state_se3s[:, 0], ego_state_se3s[:, 1])\n", - "ax.set_aspect('equal', adjustable='box')" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8", - "metadata": {}, - "outputs": [], - "source": [ - "import cv2\n", - "\n", - "# Define the output video path\n", - "output_video_path = str(f\"front_camera_video_{file_idx}.mp4\")\n", - "\n", - "# Get the dimensions of the first image\n", - "height, width, channels = front_images[0].shape\n", - "\n", - "# Create the video writer\n", - "fourcc = cv2.VideoWriter_fourcc(*'mp4v') # MP4 codec\n", - "fps = 10 # 10 frames per second as requested\n", - "out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))\n", - "\n", - "# Write each frame to the video\n", - "for img in tqdm(front_images, desc=\"Creating video\"):\n", - " # Convert from RGB to BGR (OpenCV uses BGR)\n", - " img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)\n", - " out.write(img_bgr)\n", - "\n", - "# Release the video writer\n", - "out.release()\n", - "\n", - "print(f\"Video saved to {output_video_path}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "9", - "metadata": {}, - "outputs": [], - "source": [ - "# import pyarrow as pa\n", - "import pandas as pd\n", - "# parquet_file = \"/home/daniel/Downloads/testing_location_vehicle_pose_10084636266401282188_1120_000_1140_000.parquet\"\n", - "\n", - "parquet_file = \"/home/daniel/Downloads/validation_stats_10203656353524179475_7625_000_7645_000.parquet\"\n", - "\n", - "df = pd.read_parquet(parquet_file)\n", - "\n", - "df" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "10", - "metadata": {}, - "outputs": [], - "source": [ - "array" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "11", - "metadata": {}, - "outputs": [], "source": [] } ], From eb0895d8561fed76666a389a69212d508ebfc9af Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Mon, 4 Aug 2025 13:26:41 +0200 Subject: [PATCH 18/42] Fix waymo camera pose. --- d123/common/visualization/viser/utils.py | 2 +- d123/dataset/dataset_specific/wopd/wopd_data_converter.py | 7 +++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index a3a3577e..5525a1bc 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -185,7 +185,7 @@ def get_camera_values( ego_transform[:3, :3] = get_rotation_matrix(rear_axle) ego_transform[:3, 3] = rear_axle.point_3d.array - DEBUG = True + DEBUG = False if DEBUG: print("DEBUG") diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 16aff2c0..11c214ab 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -433,12 +433,11 @@ def _extract_camera( x=0.0, y=0.0, z=0.0, - roll=0.0, - pitch=0.0, - yaw=np.deg2rad(0.0), + roll=np.deg2rad(0.0), + pitch=np.deg2rad(90.0), + yaw=np.deg2rad(-90.0), ) ) - transform[:3, :3] = transform[:3, :3] @ flip_camera context_extrinsic[camera_type] = transform From 826d05ffbd3b6544f4d7efa6620cee2ec14fe6b3 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Tue, 5 Aug 2025 18:39:39 +0200 Subject: [PATCH 19/42] Add initial version of wopd map conversion. --- .../vehicle_state/vehicle_parameters.py | 3 +- d123/common/geometry/line/polylines.py | 9 +- .../common/visualization/matplotlib/camera.py | 2 +- d123/common/visualization/viser/server.py | 38 +- .../wopd/wopd_data_converter.py | 57 ++- .../dataset_specific/wopd/wopd_map_utils.py | 222 +++++++++++ d123/dataset/logs/log_metadata.py | 4 + d123/dataset/maps/gpkg/gpkg_map.py | 10 + d123/dataset/scene/arrow_scene.py | 9 +- .../default_dataset_conversion.yaml | 2 +- notebooks/viz/camera_matplotlib.ipynb | 43 +- notebooks/viz/viser_testing_v2_scene.ipynb | 20 +- notebooks/waymo_perception/map_testing.ipynb | 376 ++++++++++++++++++ 13 files changed, 736 insertions(+), 59 deletions(-) create mode 100644 d123/dataset/dataset_specific/wopd/wopd_map_utils.py create mode 100644 notebooks/waymo_perception/map_testing.ipynb diff --git a/d123/common/datatypes/vehicle_state/vehicle_parameters.py b/d123/common/datatypes/vehicle_state/vehicle_parameters.py index b899d400..714f3628 100644 --- a/d123/common/datatypes/vehicle_state/vehicle_parameters.py +++ b/d123/common/datatypes/vehicle_state/vehicle_parameters.py @@ -55,7 +55,8 @@ def get_wopd_pacifica_parameters() -> VehicleParameters: length=5.176, height=1.777, wheel_base=3.089, - rear_axle_to_center_vertical=0.45, + # rear_axle_to_center_vertical=0.45, + rear_axle_to_center_vertical=0.60, rear_axle_to_center_longitudinal=1.461, ) diff --git a/d123/common/geometry/line/polylines.py b/d123/common/geometry/line/polylines.py index 881eef87..cbd18f56 100644 --- a/d123/common/geometry/line/polylines.py +++ b/d123/common/geometry/line/polylines.py @@ -149,10 +149,15 @@ def from_linestring(cls, linestring: geom.LineString) -> Polyline3D: else Polyline3D(geom_creation.linestrings(*linestring.xy, z=DEFAULT_Z)) ) + @classmethod + def from_array(cls, array: npt.NDArray[np.float64]) -> Polyline3D: + assert array.ndim == 2 and array.shape[1] == 3, "Array must be 2D with shape (N, 3)" + linestring = geom_creation.linestrings(*array.T) + return Polyline3D(linestring) + @property def polyline_2d(self) -> Polyline2D: - Polyline2D(geom_creation.linestrings(*self.linestring.xy)) - return Polyline2D.from_linestring(self.linestring) + return Polyline2D(geom_creation.linestrings(*self.linestring.xy)) @property def polyline_se2(self) -> PolylineSE2: diff --git a/d123/common/visualization/matplotlib/camera.py b/d123/common/visualization/matplotlib/camera.py index 1f3192f0..52d873ae 100644 --- a/d123/common/visualization/matplotlib/camera.py +++ b/d123/common/visualization/matplotlib/camera.py @@ -219,7 +219,7 @@ def _plot_rect_3d_on_img( image: npt.NDArray[np.float32], box_corners: npt.NDArray[np.float32], detection_types: List[DetectionType], - thickness: int = 1, + thickness: int = 3, ) -> npt.NDArray[np.uint8]: """ Plot the boundary lines of 3D rectangular on 2D images. diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 291f076b..f83ea474 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -17,33 +17,33 @@ # TODO: Try to fix performance issues. # TODO: Refactor this file. -# all_camera_types: List[CameraType] = [ -# CameraType.CAM_F0, -# CameraType.CAM_B0, -# CameraType.CAM_L0, -# CameraType.CAM_L1, -# CameraType.CAM_L2, -# CameraType.CAM_R0, -# CameraType.CAM_R1, -# CameraType.CAM_R2, -# ] +all_camera_types: List[CameraType] = [ + CameraType.CAM_F0, + CameraType.CAM_B0, + CameraType.CAM_L0, + CameraType.CAM_L1, + CameraType.CAM_L2, + CameraType.CAM_R0, + CameraType.CAM_R1, + CameraType.CAM_R2, +] LIDAR_POINT_SIZE: float = 0.05 -MAP_AVAILABLE: bool = False +MAP_AVAILABLE: bool = True BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" LINE_WIDTH: float = 4.0 CAMERA_SCALE: float = 1.0 -VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ - CameraType.CAM_F0, - CameraType.CAM_L0, - CameraType.CAM_R0, - CameraType.CAM_L1, - CameraType.CAM_R1, -] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ +# CameraType.CAM_F0, +# CameraType.CAM_L0, +# CameraType.CAM_R0, +# CameraType.CAM_L1, +# CameraType.CAM_R1, +# ] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] -# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0] +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] LIDAR_AVAILABLE: bool = True diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 11c214ab..9db471e3 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -27,10 +27,12 @@ from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.dataset_specific.wopd.wopd_map_utils import convert_wopd_map from d123.dataset.dataset_specific.wopd.wopd_utils import parse_range_image_and_camera_projection from d123.dataset.logs.log_metadata import LogMetadata os.environ["CUDA_VISIBLE_DEVICES"] = "-1" +D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) TARGET_DT: Final[float] = 0.1 NUPLAN_DT: Final[float] = 0.05 @@ -114,8 +116,20 @@ def get_available_splits(self) -> List[str]: ] def convert_maps(self, worker: WorkerPool) -> None: - # TODO: Implement map conversion - pass + log_args = [ + { + "tf_record": tf_record, + "split": split, + } + for split, tf_record_paths in self._tf_records_per_split.items() + for tf_record in tf_record_paths + ] + + worker_map( + worker, + partial(convert_wopd_tfrecord_map_to_gpkg, data_converter_config=self.data_converter_config), + log_args, + ) def convert_logs(self, worker: WorkerPool) -> None: log_args = [ @@ -134,8 +148,27 @@ def convert_logs(self, worker: WorkerPool) -> None: ) -def convert_wopd_tfrecord_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: - # TODO: Implement map conversion +def convert_wopd_tfrecord_map_to_gpkg( + args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig +) -> List[Any]: + + for log_info in args: + tf_record_path: Path = log_info["tf_record"] + split: str = log_info["split"] + + if not tf_record_path.exists(): + raise FileNotFoundError(f"TFRecord path {tf_record_path} does not exist.") + + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") + for data in dataset: + initial_frame = dataset_pb2.Frame() + initial_frame.ParseFromString(data.numpy()) + break + log_name = str(initial_frame.context.name) + map_file_path = D123_MAPS_ROOT / split / f"{log_name}.gpkg" + + if data_converter_config.force_map_conversion or not map_file_path.exists(): + convert_wopd_map(initial_frame, map_file_path) return [] @@ -151,11 +184,11 @@ def convert_wopd_tfrecord_log_to_arrow( dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") for data in dataset: - frame = dataset_pb2.Frame() - frame.ParseFromString(data.numpy()) + initial_frame = dataset_pb2.Frame() + initial_frame.ParseFromString(data.numpy()) break - log_name = str(frame.context.name) + log_name = str(initial_frame.context.name) log_file_path = data_converter_config.output_path / split / f"{log_name}.arrow" if data_converter_config.force_log_conversion or not log_file_path.exists(): @@ -193,12 +226,6 @@ def convert_wopd_tfrecord_log_to_arrow( (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) ) - initial_frame = next( - (dataset_pb2.Frame().FromString(data.numpy()) or dataset_pb2.Frame() for data in dataset), None - ) - if initial_frame is None: - raise ValueError(f"No frames found in TFRecord {tf_record_path}") - recording_schema = pa.schema(schema_column_list) metadata = LogMetadata( dataset="wopd", @@ -219,8 +246,8 @@ def convert_wopd_tfrecord_log_to_arrow( _write_recording_table(dataset, recording_schema, log_file_path, tf_record_path, data_converter_config) - del recording_schema, vehicle_parameters, dataset - gc.collect() + del recording_schema, vehicle_parameters, dataset + gc.collect() return [] diff --git a/d123/dataset/dataset_specific/wopd/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/wopd_map_utils.py new file mode 100644 index 00000000..a8533793 --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/wopd_map_utils.py @@ -0,0 +1,222 @@ +from pathlib import Path +from typing import Dict, Tuple + +import geopandas as gpd +import numpy as np +import numpy.typing as npt +import pandas as pd +import shapely.geometry as geom +from waymo_open_dataset import dataset_pb2 + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.line.polylines import Polyline3D +from d123.dataset.maps.map_datatypes import MapSurfaceType + + +def convert_wopd_map(frame: dataset_pb2.Frame, map_file_path: Path) -> None: + + def _get_polyline(data) -> npt.NDArray[np.float64]: + polyline = np.array([[p.x, p.y, p.z] for p in data.polyline], dtype=np.float64) + return polyline + + def _get_polygon(data) -> npt.NDArray[np.float64]: + polygon = np.array([[p.x, p.y, p.z] for p in data.polygon], dtype=np.float64) + assert polygon.shape[0] >= 3, "Polygon must have at least 3 points" + assert polygon.shape[1] == 3, "Polygon must have 3 coordinates (x, y, z)" + return polygon + + lanes: Dict[int, npt.NDArray[np.float64]] = {} + crosswalks: Dict[int, npt.NDArray[np.float64]] = {} + carparks: Dict[int, npt.NDArray[np.float64]] = {} + for map_feature in frame.map_features: + if map_feature.HasField("lane"): + polyline = _get_polyline(map_feature.lane) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + lanes[map_feature.id] = polyline + elif map_feature.HasField("road_line"): + pass + elif map_feature.HasField("road_edge"): + pass + elif map_feature.HasField("stop_sign"): + # TODO: implement stop signs + pass + elif map_feature.HasField("crosswalk"): + crosswalks[map_feature.id] = _get_polygon(map_feature.crosswalk) + elif map_feature.HasField("speed_bump"): + # TODO: implement speed bumps + pass + elif map_feature.HasField("driveway"): + # NOTE: Determine whether to use a different semantic type for driveways. + carparks[map_feature.id] = _get_polygon(map_feature.driveway) + + lane_df = get_lane_df() + lane_group_df = get_lane_group_df(lanes) + intersection_df = get_intersections_df() + crosswalk_df = get_crosswalk_df(crosswalks) + walkway_df = get_walkway_df() + carpark_df = get_carpark_df(carparks) + generic_drivable_df = get_generic_drivable_df() + + map_file_path.unlink(missing_ok=True) + if not map_file_path.parent.exists(): + map_file_path.parent.mkdir(parents=True, exist_ok=True) + + lane_df.to_file(map_file_path, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") + lane_group_df.to_file(map_file_path, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") + intersection_df.to_file(map_file_path, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") + crosswalk_df.to_file(map_file_path, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") + walkway_df.to_file(map_file_path, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") + carpark_df.to_file(map_file_path, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") + generic_drivable_df.to_file( + map_file_path, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" + ) + + +def get_lane_df() -> gpd.GeoDataFrame: + + ids = [] + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + baseline_paths = [] + geometries = [] + + data = pd.DataFrame( + { + "id": ids, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "baseline_path": baseline_paths, + } + ) + + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_lane_group_df(lanes: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + + ids = [] + lane_ids = [] + intersection_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + for lane_id, lane_array in lanes.items(): + if lane_array.ndim != 2 or lane_array.shape[0] < 2: + continue + + centerline = Polyline3D.from_array(lane_array) + left_boundary, right_boundary = create_lane_boundaries(centerline) + lane_polygon = geom.Polygon(np.vstack([left_boundary.array[:, :2], right_boundary.array[:, :2][::-1]])) + + ids.append(lane_id) + lane_ids.append([lane_id]) + intersection_ids.append([]) + predecessor_lane_group_ids.append([]) + successor_lane_group_ids.append([]) + left_boundaries.append(left_boundary.linestring) + right_boundaries.append(right_boundary.linestring) + geometries.append(lane_polygon) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "intersection_id": intersection_ids, + "predecessor_lane_group_ids": predecessor_lane_group_ids, + "successor_lane_group_ids": successor_lane_group_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_intersections_df() -> gpd.GeoDataFrame: + ids = [] + lane_group_ids = [] + geometries = [] + + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_carpark_df(carparks) -> gpd.GeoDataFrame: + ids = list(carparks.keys()) + outlines = [geom.LineString(outline) for outline in carparks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in carparks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_walkway_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(crosswalks.keys()) + outlines = [geom.LineString(outline) for outline in crosswalks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in crosswalks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_generic_drivable_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def create_lane_boundaries(polyline_3d: Polyline3D, width: float = 4) -> Tuple[Polyline3D, Polyline3D]: + + points = polyline_3d.array + half_width = width / 2.0 + + # Calculate the direction vectors between consecutive points + directions = np.diff(points, axis=0) + + # Normalize the direction vectors + directions_norm = np.linalg.norm(directions, axis=1, keepdims=True) + directions_normalized = directions / directions_norm + + # Calculate perpendicular vectors in the xy plane (z remains 0) + perpendiculars = np.zeros_like(directions) + perpendiculars[:, 0] = -directions_normalized[:, 1] # -dy + perpendiculars[:, 1] = directions_normalized[:, 0] # dx + + # Create boundaries (need to handle the last point separately) + left_boundary = points[:-1] + perpendiculars * half_width + right_boundary = points[:-1] - perpendiculars * half_width + + # Handle the last point based on the last direction + last_perp = perpendiculars[-1] + left_boundary = np.vstack([left_boundary, points[-1] + last_perp * half_width]) + right_boundary = np.vstack([right_boundary, points[-1] - last_perp * half_width]) + + return Polyline3D.from_array(left_boundary), Polyline3D.from_array(right_boundary) diff --git a/d123/dataset/logs/log_metadata.py b/d123/dataset/logs/log_metadata.py index 018156e4..818bc495 100644 --- a/d123/dataset/logs/log_metadata.py +++ b/d123/dataset/logs/log_metadata.py @@ -8,6 +8,10 @@ @dataclass class LogMetadata: + # TODO: add + # - split + # - global/local map + dataset: str log_name: str location: str diff --git a/d123/dataset/maps/gpkg/gpkg_map.py b/d123/dataset/maps/gpkg/gpkg_map.py index d202f2ce..89b42fed 100644 --- a/d123/dataset/maps/gpkg/gpkg_map.py +++ b/d123/dataset/maps/gpkg/gpkg_map.py @@ -359,3 +359,13 @@ def get_map_api_from_names(dataset: str, location: str) -> GPKGMap: map_api = GPKGMap(gpkg_path) map_api.initialize() return map_api + + +def get_local_map_api(split_name: str, log_name: str) -> GPKGMap: + print(split_name, log_name) + D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) + gpkg_path = D123_MAPS_ROOT / split_name / f"{log_name}.gpkg" + assert gpkg_path.is_file(), f"{log_name}.gpkg not found in {str(D123_MAPS_ROOT)}." + map_api = GPKGMap(gpkg_path) + map_api.initialize() + return map_api diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index 76c5e530..676f3831 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -22,7 +22,7 @@ from d123.dataset.arrow.helper import open_arrow_table from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import AbstractMap -from d123.dataset.maps.gpkg.gpkg_map import get_map_api_from_names +from d123.dataset.maps.gpkg.gpkg_map import get_local_map_api, get_map_api_from_names from d123.dataset.scene.abstract_scene import AbstractScene, SceneExtractionInfo # TODO: Remove or improve open/close dynamic of Scene object. @@ -164,7 +164,12 @@ def _lazy_initialize(self) -> None: def open(self) -> None: if self._map_api is None: try: - self._map_api = get_map_api_from_names(self._metadata.dataset, self._metadata.location) + if self._metadata.dataset == "wopd": + # FIXME: + split = str(self._arrow_log_path.parent.name) + self._map_api = get_local_map_api(split, self._metadata.log_name) + else: + self._map_api = get_map_api_from_names(self._metadata.dataset, self._metadata.location) self._map_api.initialize() except Exception as e: print(f"Error initializing map API: {e}") diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index a98adec2..0c5c118b 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -20,4 +20,4 @@ defaults: - wopd_dataset force_log_conversion: True -force_map_conversion: False +force_map_conversion: True diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index 4497ce3f..d7b611b2 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -161,7 +161,6 @@ " fig.tight_layout()\n", " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "plot_cameras_with_detections(scenes[4], iteration=10)\n", "\n", "\n", "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", @@ -193,9 +192,47 @@ "# fig.tight_layout()\n", "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "# plot_cameras_with_detections(scenes[4], iteration=150)\n", "\n", - "\n" + "# plot_cameras_with_detections(scenes[9], iteration=20)\n", + "\n", + "\n", + "\n", + "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + "# CameraType.CAM_L1: ax[0],\n", + "# CameraType.CAM_L0: ax[1],\n", + "# CameraType.CAM_F0: ax[2],\n", + "# CameraType.CAM_R0: ax[3],\n", + "# CameraType.CAM_R1: ax[4],\n", + "# }\n", + "\n", + "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + "# for camera_type, camera_ax in camera_ax_assignment.items():\n", + "# assert camera_type in scene.available_camera_types\n", + "# camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + "# if camera is not None:\n", + "# add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + "# return ax\n", + "\n", + "\n", + "# def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + "# \"\"\"\n", + "# Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + "# :param scene: navsim scene dataclass\n", + "# :param frame_idx: index of selected frame\n", + "# :return: figure and ax object of matplotlib\n", + "# \"\"\"\n", + "# scale = 2\n", + "# fig, ax = plt.subplots(1, 5, figsize=(scale * 11.5, scale * 4.3))\n", + "# _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + "# fig.tight_layout()\n", + "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "plot_cameras_with_detections(scenes[5], iteration=0)\n" ] }, { diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 4a105e1b..953d6caf 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -23,7 +23,7 @@ "source": [ "# split = \"nuplan_private_test\"\n", "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", - "\n", + " \n", "\n", "splits = [\"wopd_train\"]\n", "# log_names = None\n", @@ -32,6 +32,7 @@ "\n", "# splits = [\"nuplan_private_test\"]\n", "# splits = [\"carla\"]\n", + "# log_names = [\"11839652018869852123_2565_000_2585_000\"]\n", "log_names = None\n", "\n", "scene_tokens = None\n", @@ -67,29 +68,18 @@ "visualization_server = ViserVisualizationServer(scenes, scene_index=2)" ] }, - { - "cell_type": "markdown", - "id": "3", - "metadata": {}, - "source": [] - }, { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "3", "metadata": {}, "outputs": [], - "source": [ - "ego_state = scenes[2].get_ego_state_at_iteration(0)\n", - "\n", - "\n", - "ego_state.center" - ] + "source": [] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "4", "metadata": {}, "outputs": [], "source": [] diff --git a/notebooks/waymo_perception/map_testing.ipynb b/notebooks/waymo_perception/map_testing.ipynb new file mode 100644 index 00000000..37948271 --- /dev/null +++ b/notebooks/waymo_perception/map_testing.ipynb @@ -0,0 +1,376 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from waymo_open_dataset import dataset_pb2\n", + "\n", + "import json\n", + "import os\n", + "\n", + "import numpy as np\n", + "import tensorflow as tf\n", + "from PIL import Image\n", + "from tqdm import tqdm\n", + "from waymo_open_dataset import label_pb2\n", + "from waymo_open_dataset.protos import camera_segmentation_pb2 as cs_pb2\n", + "from waymo_open_dataset.utils import box_utils\n", + "\n", + "\n", + "import matplotlib.pyplot as plt" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "\n", + "\n", + "WOPD_DATA_ROOT = Path(\"/media/nvme1/waymo_perception/training\")\n", + "\n", + "\n", + "tfrecords_file_list = list(WOPD_DATA_ROOT.glob(\"*.tfrecord\"))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import io\n", + "from pyquaternion import Quaternion\n", + "\n", + "from d123.common.geometry.base import StateSE3\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3\n", + "\n", + "from waymo_open_dataset.utils import frame_utils\n", + "\n", + "\n", + "# Frame attributes:\n", + "# context: \n", + "# timestamp_micros: \n", + "# pose: \n", + "# images: List with 5 images\n", + "# lasers: \n", + "# Length: 5\n", + "# laser_labels: \n", + "# Length: 0\n", + "# projected_lidar_labels: \n", + "# Length: 0\n", + "# camera_labels: \n", + "# Length: 0\n", + "# no_label_zones: \n", + "# Length: 0\n", + "# map_features: \n", + "# Length: 0\n", + "# map_pose_offset: \n", + "\n", + "file_idx = 1\n", + "pathname = tfrecords_file_list[file_idx]\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "num_frames = sum(1 for _ in dataset)\n", + "\n", + "\n", + "def read_jpg_image(data: bytes) -> np.ndarray:\n", + " \"\"\"Read a JPEG image from bytes and return it as a numpy array.\"\"\"\n", + " image = Image.open(io.BytesIO(data))\n", + " return np.array(image)\n", + "\n", + "\n", + "ego_state_se3s = []\n", + "front_images = []\n", + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "\n", + "boxes = []\n", + "\n", + "for frame_idx, data in enumerate(dataset):\n", + "\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " print(\"Frame attributes:\")\n", + " for field in frame.DESCRIPTOR.fields:\n", + " field_name = field.name\n", + " if hasattr(frame, field_name):\n", + " value = getattr(frame, field_name)\n", + " if field_name != \"images\": # Don't print the whole image data\n", + " print(f\" {field_name}: {type(value)}\")\n", + " if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", + " print(f\" Length: {len(value)}\")\n", + " else:\n", + " print(f\" {field_name}: List with {len(value)} images\")\n", + "\n", + " # # plt.show()\n", + " if frame_idx == 0:\n", + " break\n", + "\n", + "ego_state_se3s = np.array(ego_state_se3s, dtype=np.float64)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "frame.context.stats.location" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", + "for frame_idx, data in enumerate(dataset):\n", + " frame = dataset_pb2.Frame()\n", + " frame.ParseFromString(data.numpy())\n", + " break\n", + "\n", + "driveways = []\n", + "stop_signs = []\n", + "\n", + "road_edges = []\n", + "road_lines = []\n", + "lanes = []\n", + "for map_feature in frame.map_features:\n", + "\n", + " if map_feature.HasField(\"lane\"):\n", + " lanes.append(\n", + " np.array([[p.x, p.y, p.z] for p in map_feature.lane.polyline])\n", + " )\n", + " elif map_feature.HasField(\"road_line\"):\n", + " road_lines.append(\n", + " np.array([[p.x, p.y, p.z] for p in map_feature.road_line.polyline])\n", + " )\n", + " elif map_feature.HasField(\"road_edge\"):\n", + " road_edges.append(\n", + " np.array([[p.x, p.y, p.z] for p in map_feature.road_edge.polyline])\n", + " )\n", + " elif map_feature.HasField(\"stop_sign\"):\n", + " # print(map_feature)\n", + " # outline = np.array(\n", + " # [[p.x, p.y, p.z] for p in map_feature.stop_sign.polygon]\n", + " # )\n", + " # stop_signs.append(outline)\n", + " pass\n", + " elif map_feature.HasField(\"crosswalk\"):\n", + " outline = np.array(\n", + " [[p.x, p.y, p.z] for p in map_feature.crosswalk.polygon]\n", + " )\n", + " elif map_feature.HasField(\"speed_bump\"):\n", + " pass\n", + " elif map_feature.HasField(\"driveway\"):\n", + " # print(map_feature.driveway)\n", + " outline = np.array(\n", + " [[p.x, p.y, p.z] for p in map_feature.driveway.polygon]\n", + " )\n", + " driveways.append(outline)\n", + "\n", + " # print(f\"Roadline: {map_feature.road_line}\")\n", + " # break\n", + "\n", + "# print(frame.map_pose_offset)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "from d123.common.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", + "\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "for road_edge in road_edges:\n", + " # print(len(driveway))\n", + " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"blue\", label=\"road_edge\")\n", + "\n", + "for lane in lanes:\n", + " # print(len(driveway))\n", + " ax.plot(lane[:, 0], lane[:, 1], color=\"gray\", label=\"lane\")\n", + "\n", + "\n", + "for road_line in road_lines:\n", + " # print(len(driveway))\n", + " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "for driveway in driveways:\n", + " # print(len(driveway))\n", + " ax.plot(driveway[:, 0], driveway[:, 1], color=\"green\", label=\"driveway\")\n", + "\n", + "\n", + "for crosswalk in crosswalks:\n", + " # print(len(driveway))\n", + " ax.plot(crosswalk[:, 0], crosswalk[:, 1], color=\"green\", label=\"driveway\")\n", + "\n", + "\n", + "add_non_repeating_legend_to_ax(ax)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.geometry.line.polylines import Polyline3D\n", + "import numpy as np\n", + "\n", + "\n", + "polyline_3d = Polyline3D.from_array(lanes[0])\n", + "\n", + "\n", + "# Create left and right boundaries for the lane\n", + "\n", + "def create_lane_boundaries(polyline_3d, width=2.0):\n", + " \"\"\"\n", + " Create left and right boundaries for a lane by offsetting perpendicular to the heading.\n", + " \n", + " Args:\n", + " polyline_3d: Polyline3D object representing the centerline\n", + " width: Total width of the lane in meters\n", + " \n", + " Returns:\n", + " left_boundary: numpy array representing the left boundary\n", + " right_boundary: numpy array representing the right boundary\n", + " \"\"\"\n", + " points = polyline_3d.array\n", + " half_width = width / 2.0\n", + " \n", + " # Calculate the direction vectors between consecutive points\n", + " directions = np.diff(points, axis=0)\n", + " \n", + " # Normalize the direction vectors\n", + " directions_norm = np.linalg.norm(directions, axis=1, keepdims=True)\n", + " directions_normalized = directions / directions_norm\n", + " \n", + " # Calculate perpendicular vectors in the xy plane (z remains 0)\n", + " perpendiculars = np.zeros_like(directions)\n", + " perpendiculars[:, 0] = -directions_normalized[:, 1] # -dy\n", + " perpendiculars[:, 1] = directions_normalized[:, 0] # dx\n", + " \n", + " # Create boundaries (need to handle the last point separately)\n", + " left_boundary = points[:-1] + perpendiculars * half_width\n", + " right_boundary = points[:-1] - perpendiculars * half_width\n", + " \n", + " # Handle the last point based on the last direction\n", + " last_perp = perpendiculars[-1]\n", + " left_boundary = np.vstack([left_boundary, points[-1] + last_perp * half_width])\n", + " right_boundary = np.vstack([right_boundary, points[-1] - last_perp * half_width])\n", + " \n", + " return left_boundary, right_boundary\n", + "\n", + "# Create the boundaries with a 4m wide lane\n", + "left_boundary, right_boundary = create_lane_boundaries(polyline_3d, width=4.0)\n", + "\n", + "# Plot the centerline and the boundaries\n", + "plt.figure(figsize=(12, 8))\n", + "plt.plot(polyline_3d.array[:, 0], polyline_3d.array[:, 1], 'g-', label='Centerline')\n", + "plt.plot(left_boundary[:, 0], left_boundary[:, 1], 'b-', label='Left Boundary')\n", + "plt.plot(right_boundary[:, 0], right_boundary[:, 1], 'r-', label='Right Boundary')\n", + "plt.axis('equal')\n", + "plt.grid(True)\n", + "plt.legend()\n", + "plt.title('Lane with Left and Right Boundaries')\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "for driveway in driveways:\n", + " # print(len(driveway))\n", + " distance = abs(driveway[:, -1].min() - driveway[:, -1].max())\n", + " if distance > 1:\n", + " print(driveway[:, -1])\n", + " plt.plot(driveway[:, 0], driveway[:, 1], \"r-\", color=\"red\")\n", + " else:\n", + " plt.plot(driveway[:, 0], driveway[:, 1], \"r-\", color=\"blue\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "for polyline in polylines:\n", + " try:\n", + " plt.plot(polyline[:, 0], polyline[:, 1])\n", + " except Exception as e:\n", + " print(f\"Error plotting polyline: {polyline.shape}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 42e1606df42e751f47a216034c18a6dc72f65ce3 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Wed, 6 Aug 2025 23:14:28 +0200 Subject: [PATCH 20/42] Push unfinished state of waymo conversion --- d123/common/geometry/line/polylines.py | 1 + d123/common/visualization/viser/server.py | 5 +- d123/common/visualization/viser/utils.py | 3 +- .../wopd/womp_map_conversion.py | 203 +++++++ .../dataset_specific/wopd/wopd_map_utils.py | 232 +++++++- .../default_dataset_conversion.yaml | 2 +- notebooks/viz/bev_matplotlib.ipynb | 121 +++- notebooks/viz/viser_testing_v2_scene.ipynb | 2 +- notebooks/waymo_perception/map_testing.ipynb | 521 +++++++++++++++--- 9 files changed, 984 insertions(+), 106 deletions(-) create mode 100644 d123/dataset/dataset_specific/wopd/womp_map_conversion.py diff --git a/d123/common/geometry/line/polylines.py b/d123/common/geometry/line/polylines.py index cbd18f56..6adb11fc 100644 --- a/d123/common/geometry/line/polylines.py +++ b/d123/common/geometry/line/polylines.py @@ -79,6 +79,7 @@ def __post_init__(self): if self.linestring is None: self.linestring = geom_creation.linestrings(self.se2_array[..., StateSE2Index.XY]) + self.se2_array[:, StateSE2Index.YAW] = np.unwrap(self.se2_array[:, StateSE2Index.YAW], axis=0) self._progress = get_path_progress(self.se2_array) self._interpolator = interp1d(self._progress, self.se2_array, axis=0, bounds_error=False, fill_value=0.0) diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index f83ea474..ca37a13e 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -43,9 +43,10 @@ # CameraType.CAM_R1, # ] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] -VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] -LIDAR_AVAILABLE: bool = True +LIDAR_AVAILABLE: bool = False class ViserVisualizationServer: diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 5525a1bc..341c629e 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -84,7 +84,8 @@ def get_map_meshes(scene: AbstractScene): initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 map_surface_types = [ - MapSurfaceType.LANE_GROUP, + # MapSurfaceType.LANE_GROUP, + MapSurfaceType.LANE, MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK, MapSurfaceType.CARPARK, diff --git a/d123/dataset/dataset_specific/wopd/womp_map_conversion.py b/d123/dataset/dataset_specific/wopd/womp_map_conversion.py new file mode 100644 index 00000000..6372f19c --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/womp_map_conversion.py @@ -0,0 +1,203 @@ +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +import numpy as np +import numpy.typing as npt +import shapely.geometry as geom + +from d123.common.geometry.base import Point3D, StateSE2 +from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.common.geometry.transform.tranform_2d import translate_along_yaw +from d123.common.geometry.vector import Vector2D + +MAX_LANE_WIDTH = 25.0 # meters +MIN_LANE_WIDTH = 2.0 +DEFAULT_LANE_WIDTH = 4.0 +BOUNDARY_STEP_SIZE = 0.5 # meters +MAX_Z_DISTANCE = 1.0 # meters + + +@dataclass +class BoundaryHit: + distance_along_perp_2d: float + distance_along_polyline_3d: float + polyline_token: str + + +def _extract_lane_boundaries( + lanes: Dict[int, npt.NDArray[np.float64]], + road_lines: Dict[int, npt.NDArray[np.float64]], + road_edges: Dict[int, npt.NDArray[np.float64]], +) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: + polyline_dict: Dict[str, Dict[int, Polyline3D]] = {"lane": {}, "roadline": {}, "roadedge": {}} + + for lane_id, lane_polyline in lanes.items(): + if lane_polyline.ndim == 2 and lane_polyline.shape[1] == 3 and len(lane_polyline) > 0: + polyline_dict["lane"][lane_id] = Polyline3D.from_array(lane_polyline) + + # for road_line_id, road_line_polyline in road_lines.items(): + # if road_line_polyline.ndim == 2 and road_line_polyline.shape[1] == 3 and len(road_line_polyline) > 0: + # polyline_dict["roadline"][road_line_id] = Polyline3D.from_array(road_line_polyline) + + for road_edge_id, road_edge_polyline in road_edges.items(): + if road_edge_polyline.ndim == 2 and road_edge_polyline.shape[1] == 3 and len(road_edge_polyline) > 0: + polyline_dict["roadedge"][road_edge_id] = Polyline3D.from_array(road_edge_polyline) + + geometries = [] + tokens = [] + for line_type, polylines in polyline_dict.items(): + for polyline_id, polyline in polylines.items(): + geometries.append(polyline.polyline_2d.linestring) + tokens.append(f"{line_type}_{polyline_id}") + + occupancy_2d = OccupancyMap2D(geometries, tokens) + + # for each lane + # - sample poses along centerline + # - query left and right perpendicular lines in 2d occupancy map + # - for {left, right} retrieve first perpendicular intersection with length MAX_HALF_WIDTH + # - if intersection with road edge, add intersecting point to {left, right} boundary + # - if intersection with centerline, add middle point to {left, right} boundary + # - if no intersection found, add MAX_HALF_WIDTH to {left, right} boundary + + left_boundaries = {lane_id: [] for lane_id in lanes.keys()} + right_boundaries = {lane_id: [] for lane_id in lanes.keys()} + + def get_type_and_id_from_token(token: str) -> Tuple[str, int]: + """Extract type and id from token.""" + line_type, line_id = token.split("_") + return line_type, int(line_id) + + def get_polyline_from_token(token: str) -> Polyline3D: + """Extract polyline from token.""" + line_type, line_id = get_type_and_id_from_token(token) + return polyline_dict[line_type][line_id] + + def get_intersecting_point_3d(perp_linestring: geom.LineString, intersecting_token: str) -> Optional[Point3D]: + + intersecting_polyline_3d = get_polyline_from_token(intersecting_token) + intersecting_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[intersecting_token]] + intersecting_point_2d = perp_linestring.intersection(intersecting_linestring) + + distance_2d_norm = None + if isinstance(intersecting_point_2d, geom.Point): + distance_2d_norm = intersecting_linestring.project(intersecting_point_2d, normalized=True) + elif isinstance(intersecting_point_2d, geom.MultiPoint): + geom_points = [point for point in intersecting_point_2d.geoms] + starting_perp_point = geom.Point(perp_linestring.coords[0]) + start_perp_point_distance = np.array( + [starting_perp_point.distance(geom_point) for geom_point in geom_points] + ) + start_perp_point_distance[start_perp_point_distance < MIN_LANE_WIDTH / 2.0] = np.inf + closest_index = np.argmin(start_perp_point_distance) + distance_2d_norm = intersecting_linestring.project(geom_points[closest_index], normalized=True) + + intersecting_point_3d = None + if distance_2d_norm is not None: + intersecting_point_3d = intersecting_polyline_3d.interpolate( + distance_2d_norm * intersecting_polyline_3d.length + ) + if np.linalg.norm(intersecting_point_3d.point_2d.array - perp_linestring.coords[0]) < MIN_LANE_WIDTH / 2.0: + return None + + return intersecting_point_3d + + for lane_id, lane_polyline in polyline_dict["lane"].items(): + current_lane_token = f"lane_{lane_id}" + lane_polyline_se2 = lane_polyline.polyline_se2 + + # 0. Find connected or intersecting lanes + connected_lane_tokens = [ + token + for token in occupancy_2d.intersects(lane_polyline_se2.linestring) + if token.startswith("lane_") + if token != current_lane_token + ] + + # 1. sample poses along centerline + distances_se2 = np.linspace( + 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_query_se2 = [ + StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) + ] + distances_3d = np.linspace( + 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_query_3d = [ + Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) + ] + assert len(lane_query_se2) == len(lane_query_3d) + + for lane_query_se2, lane_query_3d in zip(lane_query_se2, lane_query_3d): + for sign in [1.0, -1.0]: + perp_start_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * 0.1)) + perp_end_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)) + perp_linestring = geom.LineString( + [[perp_start_point.x, perp_start_point.y], [perp_end_point.x, perp_end_point.y]] + ) + + # 1. find intersecting lines, compute 3D distance + intersecting_tokens = occupancy_2d.intersects(perp_linestring) + intersecting_points_3d: Dict[str, Point3D] = {} + crosses_centerline = False + + for intersecting_token in intersecting_tokens: + # if intersecting_token in connected_lane_tokens: + # crosses_centerline = True + # continue + # if intersecting_token in connected_lane_tokens: + # continue + + # if intersecting_token == current_lane_token: + # crosses_centerline = True + + intersecting_point_3d = get_intersecting_point_3d(perp_linestring, intersecting_token) + if intersecting_point_3d is None: + continue + + if np.abs(intersecting_point_3d.z - lane_query_3d.z) > MAX_Z_DISTANCE: + continue + + intersecting_points_3d[intersecting_token] = intersecting_point_3d + + boundary_point_3d: Optional[Point3D] = None + if len(intersecting_points_3d) > 0 and not crosses_centerline: + # 2. find closest intersection + intersecting_distances = { + token: np.linalg.norm(point_3d.array - lane_query_3d.array) + for token, point_3d in intersecting_points_3d.items() + } + closest_token = min(intersecting_distances, key=intersecting_distances.get) + closest_type, closest_id = get_type_and_id_from_token(closest_token) + + if closest_type == "lane": + boundary_point_3d = Point3D.from_array( + (intersecting_points_3d[closest_token].array + lane_query_3d.array) / 2.0 + ) + elif closest_type == "roadedge": + boundary_point_3d = intersecting_points_3d[closest_token] + else: + perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 # Default to half the lane width + boundary_point_se2 = translate_along_yaw( + lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance) + ) + boundary_point_3d = Point3D(boundary_point_se2.x, boundary_point_se2.y, lane_query_3d.z) + if sign == 1.0: + left_boundaries[lane_id].append(boundary_point_3d.array) + else: + right_boundaries[lane_id].append(boundary_point_3d.array) + + left_boundaries = { + lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) + for lane_id, boundary_array in left_boundaries.items() + if len(boundary_array) > 1 + } + right_boundaries = { + lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) + for lane_id, boundary_array in right_boundaries.items() + if len(boundary_array) > 1 + } + + return left_boundaries, right_boundaries diff --git a/d123/dataset/dataset_specific/wopd/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/wopd_map_utils.py index a8533793..5c8b0239 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_map_utils.py +++ b/d123/dataset/dataset_specific/wopd/wopd_map_utils.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Dict, Tuple +from typing import Dict, Optional, Tuple import geopandas as gpd import numpy as np @@ -8,8 +8,11 @@ import shapely.geometry as geom from waymo_open_dataset import dataset_pb2 -from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.base import Point3D, Point3DIndex, StateSE2 from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.common.geometry.transform.tranform_2d import translate_along_yaw +from d123.common.geometry.vector import Vector2D from d123.dataset.maps.map_datatypes import MapSurfaceType @@ -26,8 +29,11 @@ def _get_polygon(data) -> npt.NDArray[np.float64]: return polygon lanes: Dict[int, npt.NDArray[np.float64]] = {} + road_lines: Dict[int, npt.NDArray[np.float64]] = {} + road_edges: Dict[int, npt.NDArray[np.float64]] = {} crosswalks: Dict[int, npt.NDArray[np.float64]] = {} carparks: Dict[int, npt.NDArray[np.float64]] = {} + for map_feature in frame.map_features: if map_feature.HasField("lane"): polyline = _get_polyline(map_feature.lane) @@ -35,9 +41,15 @@ def _get_polygon(data) -> npt.NDArray[np.float64]: continue lanes[map_feature.id] = polyline elif map_feature.HasField("road_line"): - pass + polyline = _get_polyline(map_feature.road_line) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + road_lines[map_feature.id] = polyline elif map_feature.HasField("road_edge"): - pass + polyline = _get_polyline(map_feature.road_edge) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + road_edges[map_feature.id] = polyline elif map_feature.HasField("stop_sign"): # TODO: implement stop signs pass @@ -50,7 +62,7 @@ def _get_polygon(data) -> npt.NDArray[np.float64]: # NOTE: Determine whether to use a different semantic type for driveways. carparks[map_feature.id] = _get_polygon(map_feature.driveway) - lane_df = get_lane_df() + lane_df = get_lane_df(lanes, road_lines, road_edges) lane_group_df = get_lane_group_df(lanes) intersection_df = get_intersections_df() crosswalk_df = get_crosswalk_df(crosswalks) @@ -73,7 +85,11 @@ def _get_polygon(data) -> npt.NDArray[np.float64]: ) -def get_lane_df() -> gpd.GeoDataFrame: +def get_lane_df( + lanes: Dict[int, npt.NDArray[np.float64]], + road_lines: Dict[int, npt.NDArray[np.float64]], + road_edges: Dict[int, npt.NDArray[np.float64]], +) -> gpd.GeoDataFrame: ids = [] lane_group_ids = [] @@ -85,6 +101,26 @@ def get_lane_df() -> gpd.GeoDataFrame: baseline_paths = [] geometries = [] + left_boundaries_3d, right_boundaries_3d = _extract_lane_boundaries(lanes, road_lines, road_edges) + for lane_id, lane_centerline_array in lanes.items(): + if lane_id not in left_boundaries_3d or lane_id not in right_boundaries_3d: + continue + lane_centerline = Polyline3D.from_array(lane_centerline_array) + + ids.append(lane_id) + lane_group_ids.append([lane_id]) + speed_limits_mps.append(None) # TODO: Implement speed limits + predecessor_ids.append([]) # TODO: Implement predecessor lanes + successor_ids.append([]) # TODO: Implement successor lanes + left_boundaries.append(left_boundaries_3d[lane_id].linestring) + right_boundaries.append(right_boundaries_3d[lane_id].linestring) + baseline_paths.append(lane_centerline.linestring) + + geometry = geom.Polygon( + np.vstack([left_boundaries_3d[lane_id].array[:, :2], right_boundaries_3d[lane_id].array[:, :2][::-1]]) + ) + geometries.append(geometry) + data = pd.DataFrame( { "id": ids, @@ -220,3 +256,187 @@ def create_lane_boundaries(polyline_3d: Polyline3D, width: float = 4) -> Tuple[P right_boundary = np.vstack([right_boundary, points[-1] - last_perp * half_width]) return Polyline3D.from_array(left_boundary), Polyline3D.from_array(right_boundary) + + +def _extract_lane_boundaries( + lanes: Dict[int, npt.NDArray[np.float64]], + road_lines: Dict[int, npt.NDArray[np.float64]], + road_edges: Dict[int, npt.NDArray[np.float64]], +) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: + polyline_dict: Dict[str, Dict[int, Polyline3D]] = {"lane": {}, "roadline": {}, "roadedge": {}} + + for lane_id, lane_polyline in lanes.items(): + if lane_polyline.ndim == 2 and lane_polyline.shape[1] == 3 and len(lane_polyline) > 0: + polyline_dict["lane"][lane_id] = Polyline3D.from_array(lane_polyline) + + # for road_line_id, road_line_polyline in road_lines.items(): + # if road_line_polyline.ndim == 2 and road_line_polyline.shape[1] == 3 and len(road_line_polyline) > 0: + # polyline_dict["roadline"][road_line_id] = Polyline3D.from_array(road_line_polyline) + + for road_edge_id, road_edge_polyline in road_edges.items(): + if road_edge_polyline.ndim == 2 and road_edge_polyline.shape[1] == 3 and len(road_edge_polyline) > 0: + polyline_dict["roadedge"][road_edge_id] = Polyline3D.from_array(road_edge_polyline) + + geometries = [] + tokens = [] + for line_type, polylines in polyline_dict.items(): + for polyline_id, polyline in polylines.items(): + geometries.append(polyline.polyline_2d.linestring) + tokens.append(f"{line_type}_{polyline_id}") + + occupancy_2d = OccupancyMap2D(geometries, tokens) + + # for each lane + # - sample poses along centerline + # - query left and right perpendicular lines in 2d occupancy map + # - for {left, right} retrieve first perpendicular intersection with length MAX_HALF_WIDTH + # - if intersection with road edge, add intersecting point to {left, right} boundary + # - if intersection with centerline, add middle point to {left, right} boundary + # - if no intersection found, add MAX_HALF_WIDTH to {left, right} boundary + + MAX_LANE_WIDTH = 25.0 # meters + MIN_LANE_WIDTH = 2.0 + DEFAULT_LANE_WIDTH = 4.0 + BOUNDARY_STEP_SIZE = 0.5 # meters + MAX_Z_DISTANCE = 1.0 # meters + + left_boundaries = {lane_id: [] for lane_id in lanes.keys()} + right_boundaries = {lane_id: [] for lane_id in lanes.keys()} + + def get_type_and_id_from_token(token: str) -> Tuple[str, int]: + """Extract type and id from token.""" + line_type, line_id = token.split("_") + return line_type, int(line_id) + + def get_polyline_from_token(token: str) -> Polyline3D: + """Extract polyline from token.""" + line_type, line_id = get_type_and_id_from_token(token) + return polyline_dict[line_type][line_id] + + def get_intersecting_point_3d(perp_linestring: geom.LineString, intersecting_token: str) -> Optional[Point3D]: + + intersecting_polyline_3d = get_polyline_from_token(intersecting_token) + intersecting_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[intersecting_token]] + intersecting_point_2d = perp_linestring.intersection(intersecting_linestring) + + distance_2d_norm = None + if isinstance(intersecting_point_2d, geom.Point): + distance_2d_norm = intersecting_linestring.project(intersecting_point_2d, normalized=True) + elif isinstance(intersecting_point_2d, geom.MultiPoint): + geom_points = [point for point in intersecting_point_2d.geoms] + starting_perp_point = geom.Point(perp_linestring.coords[0]) + start_perp_point_distance = np.array( + [starting_perp_point.distance(geom_point) for geom_point in geom_points] + ) + start_perp_point_distance[start_perp_point_distance < MIN_LANE_WIDTH / 2.0] = np.inf + closest_index = np.argmin(start_perp_point_distance) + distance_2d_norm = intersecting_linestring.project(geom_points[closest_index], normalized=True) + + intersecting_point_3d = None + if distance_2d_norm is not None: + intersecting_point_3d = intersecting_polyline_3d.interpolate( + distance_2d_norm * intersecting_polyline_3d.length + ) + if np.linalg.norm(intersecting_point_3d.point_2d.array - perp_linestring.coords[0]) < MIN_LANE_WIDTH / 2.0: + return None + + return intersecting_point_3d + + for lane_id, lane_polyline in polyline_dict["lane"].items(): + current_lane_token = f"lane_{lane_id}" + lane_polyline_se2 = lane_polyline.polyline_se2 + + # 0. Find connected or intersecting lanes + connected_lane_tokens = [ + token + for token in occupancy_2d.intersects(lane_polyline_se2.linestring) + if token.startswith("lane_") + if token != current_lane_token + ] + + # 1. sample poses along centerline + distances_se2 = np.linspace( + 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_query_se2 = [ + StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) + ] + distances_3d = np.linspace( + 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_query_3d = [ + Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) + ] + assert len(lane_query_se2) == len(lane_query_3d) + + for lane_query_se2, lane_query_3d in zip(lane_query_se2, lane_query_3d): + for sign in [1.0, -1.0]: + perp_start_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * 0.1)) + perp_end_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)) + perp_linestring = geom.LineString( + [[perp_start_point.x, perp_start_point.y], [perp_end_point.x, perp_end_point.y]] + ) + + # 1. find intersecting lines, compute 3D distance + intersecting_tokens = occupancy_2d.intersects(perp_linestring) + intersecting_points_3d: Dict[str, Point3D] = {} + crosses_centerline = False + + for intersecting_token in intersecting_tokens: + # if intersecting_token in connected_lane_tokens: + # crosses_centerline = True + # continue + # if intersecting_token in connected_lane_tokens: + # continue + + # if intersecting_token == current_lane_token: + # crosses_centerline = True + + intersecting_point_3d = get_intersecting_point_3d(perp_linestring, intersecting_token) + if intersecting_point_3d is None: + continue + + if np.abs(intersecting_point_3d.z - lane_query_3d.z) > MAX_Z_DISTANCE: + continue + + intersecting_points_3d[intersecting_token] = intersecting_point_3d + + boundary_point_3d: Optional[Point3D] = None + if len(intersecting_points_3d) > 0 and not crosses_centerline: + # 2. find closest intersection + intersecting_distances = { + token: np.linalg.norm(point_3d.array - lane_query_3d.array) + for token, point_3d in intersecting_points_3d.items() + } + closest_token = min(intersecting_distances, key=intersecting_distances.get) + closest_type, closest_id = get_type_and_id_from_token(closest_token) + + if closest_type == "lane": + boundary_point_3d = Point3D.from_array( + (intersecting_points_3d[closest_token].array + lane_query_3d.array) / 2.0 + ) + elif closest_type == "roadedge": + boundary_point_3d = intersecting_points_3d[closest_token] + else: + perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 # Default to half the lane width + boundary_point_se2 = translate_along_yaw( + lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance) + ) + boundary_point_3d = Point3D(boundary_point_se2.x, boundary_point_se2.y, lane_query_3d.z) + if sign == 1.0: + left_boundaries[lane_id].append(boundary_point_3d.array) + else: + right_boundaries[lane_id].append(boundary_point_3d.array) + + left_boundaries = { + lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) + for lane_id, boundary_array in left_boundaries.items() + if len(boundary_array) > 1 + } + right_boundaries = { + lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) + for lane_id, boundary_array in right_boundaries.items() + if len(boundary_array) > 1 + } + + return left_boundaries, right_boundaries diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 0c5c118b..cceb2911 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -19,5 +19,5 @@ defaults: # - carla_dataset - wopd_dataset -force_log_conversion: True +force_log_conversion: False force_map_conversion: True diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 1bb8217c..fc24de85 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -26,6 +26,7 @@ "\n", "\n", "splits = [\"wopd_train\"]\n", + "# splits = [\"carla\"]\n", "# log_names = None\n", "\n", "\n", @@ -60,23 +61,127 @@ "metadata": {}, "outputs": [], "source": [ - "from typing import Tuple\n", + "from typing import List, Optional, Tuple\n", "import matplotlib.pyplot as plt\n", - "from d123.common.visualization.matplotlib.observation import add_box_detections_to_ax, add_default_map_on_ax, add_ego_vehicle_to_ax, add_traffic_lights_to_ax\n", - "from d123.common.visualization.matplotlib.plots import plot_scene_at_iteration\n", + "from d123.common.geometry.base import Point2D\n", + "from d123.common.visualization.color.color import BLACK, LIGHT_GREY, TAB_10\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", + "from d123.common.visualization.matplotlib.observation import (\n", + " add_box_detections_to_ax,\n", + " add_ego_vehicle_to_ax,\n", + " add_traffic_lights_to_ax,\n", + ")\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", + "from d123.dataset.maps.abstract_map import AbstractMap\n", + "from d123.dataset.maps.abstract_map_objects import AbstractLane\n", + "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", "from d123.dataset.scene.abstract_scene import AbstractScene\n", "\n", + "\n", + "import shapely.geometry as geom\n", + "\n", + "LEFT_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=BLACK,\n", + " fill_color_alpha=1.0,\n", + " line_color=TAB_10[2],\n", + " line_color_alpha=0.5,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "RIGHT_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=BLACK,\n", + " fill_color_alpha=1.0,\n", + " line_color=TAB_10[3],\n", + " line_color_alpha=0.5,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "\n", + "LANE_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=LIGHT_GREY,\n", + " fill_color_alpha=0.5,\n", + " line_color=BLACK,\n", + " line_color_alpha=0.2,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=1,\n", + ")\n", + "\n", + "\n", + "def add_debug_map_on_ax(\n", + " ax: plt.Axes,\n", + " map_api: AbstractMap,\n", + " point_2d: Point2D,\n", + " radius: float,\n", + " route_lane_group_ids: Optional[List[int]] = None,\n", + ") -> None:\n", + " layers: List[MapSurfaceType] = [\n", + " MapSurfaceType.LANE,\n", + " MapSurfaceType.GENERIC_DRIVABLE,\n", + " MapSurfaceType.CARPARK,\n", + " MapSurfaceType.CROSSWALK,\n", + " MapSurfaceType.INTERSECTION,\n", + " MapSurfaceType.WALKWAY,\n", + " ]\n", + " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", + " y_min, y_max = point_2d.y - radius, point_2d.y + radius\n", + " patch = geom.box(x_min, y_min, x_max, y_max)\n", + " map_objects_dict = map_api.query(geometry=patch, layers=layers, predicate=\"intersects\")\n", + "\n", + " for layer, map_objects in map_objects_dict.items():\n", + " for map_object in map_objects:\n", + " try:\n", + " if layer in [\n", + " MapSurfaceType.GENERIC_DRIVABLE,\n", + " MapSurfaceType.CARPARK,\n", + " # MapSurfaceType.CROSSWALK,\n", + " # MapSurfaceType.INTERSECTION,\n", + " # MapSurfaceType.WALKWAY,\n", + " ]:\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + " if layer in [MapSurfaceType.LANE]:\n", + " map_object: AbstractLane\n", + " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", + " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", + " # add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", + "\n", + " # centroid = map_object.shapely_polygon.centroid\n", + " # ax.text(\n", + " # centroid.x,\n", + " # centroid.y,\n", + " # str(map_object.id),\n", + " # horizontalalignment=\"center\",\n", + " # verticalalignment=\"center\",\n", + " # fontsize=8,\n", + " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", + " # )\n", + "\n", + " except Exception:\n", + " import traceback\n", + "\n", + " print(f\"Error adding map object of type {layer.name} and id {map_object.id}\")\n", + " traceback.print_exc()\n", + "\n", + " ax.set_title(f\"Map: {map_api.map_name}\")\n", + "\n", + "\n", "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", "\n", " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", " box_detections = scene.get_box_detections_at_iteration(iteration)\n", "\n", " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", - " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=route_lane_group_ids)\n", + " add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", - " add_box_detections_to_ax(ax, box_detections)\n", - " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + " # add_box_detections_to_ax(ax, box_detections)\n", + " # add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", "\n", " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", @@ -89,13 +194,13 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " fig, ax = plt.subplots(figsize=(10, 10))\n", + " fig, ax = plt.subplots(figsize=(25, 25))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", "scene_index = 0\n", - "plot_scene_at_iteration(scenes[scene_index], iteration=0,radius=80)" + "plot_scene_at_iteration(scenes[scene_index], iteration=20, radius=200)" ] }, { diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 953d6caf..ed6e7fdc 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -65,7 +65,7 @@ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes, scene_index=2)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=4)" ] }, { diff --git a/notebooks/waymo_perception/map_testing.ipynb b/notebooks/waymo_perception/map_testing.ipynb index 37948271..f4da563b 100644 --- a/notebooks/waymo_perception/map_testing.ipynb +++ b/notebooks/waymo_perception/map_testing.ipynb @@ -75,7 +75,7 @@ "# Length: 0\n", "# map_pose_offset: \n", "\n", - "file_idx = 1\n", + "file_idx = 3\n", "pathname = tfrecords_file_list[file_idx]\n", "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", "num_frames = sum(1 for _ in dataset)\n", @@ -133,102 +133,322 @@ "metadata": {}, "outputs": [], "source": [ + "from collections import defaultdict\n", + "\n", + "from d123.common.geometry.units import mph_to_mps\n", + "\n", + "\n", "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", "for frame_idx, data in enumerate(dataset):\n", " frame = dataset_pb2.Frame()\n", " frame.ParseFromString(data.numpy())\n", " break\n", "\n", - "driveways = []\n", - "stop_signs = []\n", + "driveways = {}\n", + "stop_signs = {}\n", + "crosswalks = {}\n", + "road_edges = {}\n", + "road_lines = {}\n", + "\n", + "lanes = {}\n", + "lanes_type = {}\n", + "lanes_speed_limit_mps = {}\n", + "lanes_interpolate = {}\n", + "lanes_successors = defaultdict(list)\n", + "lanes_predecessors = defaultdict(list)\n", + "\n", + "lanes_map_features = {}\n", "\n", - "road_edges = []\n", - "road_lines = []\n", - "lanes = []\n", "for map_feature in frame.map_features:\n", "\n", " if map_feature.HasField(\"lane\"):\n", - " lanes.append(\n", - " np.array([[p.x, p.y, p.z] for p in map_feature.lane.polyline])\n", - " )\n", + " lanes[map_feature.id] = np.array([[p.x, p.y, p.z] for p in map_feature.lane.polyline])\n", + " lanes_map_features[map_feature.id] = map_feature\n", + " lanes_type[map_feature.id] = map_feature.lane.type\n", + " lanes_speed_limit_mps[map_feature.id] = mph_to_mps(map_feature.lane.speed_limit_mph)\n", + " for lane_id_ in map_feature.lane.exit_lanes:\n", + " lanes_successors[map_feature.id].append(lane_id_)\n", + " for lane_id_ in map_feature.lane.exit_lanes:\n", + " lanes_predecessors[map_feature.id].append(lane_id_)\n", + "\n", " elif map_feature.HasField(\"road_line\"):\n", - " road_lines.append(\n", - " np.array([[p.x, p.y, p.z] for p in map_feature.road_line.polyline])\n", - " )\n", + " road_lines[map_feature.id] = np.array([[p.x, p.y, p.z] for p in map_feature.road_line.polyline])\n", " elif map_feature.HasField(\"road_edge\"):\n", - " road_edges.append(\n", - " np.array([[p.x, p.y, p.z] for p in map_feature.road_edge.polyline])\n", - " )\n", + " road_edges[map_feature.id] = np.array([[p.x, p.y, p.z] for p in map_feature.road_edge.polyline])\n", " elif map_feature.HasField(\"stop_sign\"):\n", - " # print(map_feature)\n", - " # outline = np.array(\n", - " # [[p.x, p.y, p.z] for p in map_feature.stop_sign.polygon]\n", - " # )\n", - " # stop_signs.append(outline)\n", " pass\n", " elif map_feature.HasField(\"crosswalk\"):\n", - " outline = np.array(\n", - " [[p.x, p.y, p.z] for p in map_feature.crosswalk.polygon]\n", - " )\n", + " outline = np.array([[p.x, p.y, p.z] for p in map_feature.crosswalk.polygon])\n", + " crosswalks[map_feature.id] = outline\n", " elif map_feature.HasField(\"speed_bump\"):\n", " pass\n", " elif map_feature.HasField(\"driveway\"):\n", " # print(map_feature.driveway)\n", - " outline = np.array(\n", - " [[p.x, p.y, p.z] for p in map_feature.driveway.polygon]\n", + " outline = np.array([[p.x, p.y, p.z] for p in map_feature.driveway.polygon])\n", + " driveways[map_feature.id] = outline\n", + "\n", + " # print(f\"Roadline: {map_feature.road_line}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.dataset_specific.wopd.wopd_map_utils import _extract_lane_boundaries\n", + "\n", + "\n", + "left_boundaries, right_boundaries = _extract_lane_boundaries(\n", + " lanes=lanes,\n", + " road_lines=road_lines,\n", + " road_edges=road_edges,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "size = 20\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "for lane_id, lane_ in lanes.items():\n", + " if lane_id not in left_boundaries or lane_id not in right_boundaries:\n", + " continue\n", + " left_boundary = left_boundaries[lane_id].array\n", + " right_boundary = right_boundaries[lane_id].array\n", + "\n", + " assert len(left_boundary) > 0 and len(right_boundary) > 0\n", + " ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", + " ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "\n", + " ax.plot(lane_[:, 0], lane_[:, 1], color=\"black\")\n", + "\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.geometry.base import StateSE2\n", + "from d123.common.geometry.line.polylines import PolylineSE2\n", + "from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", + "from d123.common.geometry.vector import Vector2D\n", + "\n", + "size = 30\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "lane_id = 114\n", + "lane_ = lanes[lane_id]\n", + "\n", + "lane_polyline_se2 = PolylineSE2.from_array(lane_[:, :2])\n", + "BOUNDARY_STEP_SIZE = 0.5\n", + "\n", + "\n", + "distances_se2 = np.linspace(\n", + " 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True\n", + ")\n", + "lane_query_se2 = lane_polyline_se2.interpolate(distances_se2)\n", + "\n", + "\n", + "left_boundary = left_boundaries[lane_id].array\n", + "right_boundary = right_boundaries[lane_id].array\n", + "\n", + "assert len(left_boundary) > 0 and len(right_boundary) > 0\n", + "ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", + "ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "\n", + "ax.scatter(lane_query_se2[:, 0], lane_query_se2[:, 1], c=lane_query_se2[:, 2] / np.pi, cmap=\"viridis\")\n", + "\n", + "\n", + "MAX_LANE_WIDTH = 25\n", + "for state_se2_array in lane_query_se2:\n", + " for sign in [1.0, -1.0]:\n", + " perp_start_point = translate_along_yaw(StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * 0.1))\n", + " perp_end_point = translate_along_yaw(\n", + " StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)\n", + " )\n", + " ax.plot(\n", + " [perp_start_point.x, perp_end_point.x],\n", + " [perp_start_point.y, perp_end_point.y],\n", + " color=\"lime\" if sign > 0 else \"red\",\n", + " linestyle=\"--\",\n", + " alpha=0.5,\n", " )\n", - " driveways.append(outline)\n", "\n", - " # print(f\"Roadline: {map_feature.road_line}\")\n", - " # break\n", "\n", - "# print(frame.map_pose_offset)\n" + "for road_edge in road_edges.values():\n", + " # print(len(driveway))\n", + " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"black\", label=\"road_edge\", linestyle=\"dashdot\")\n", + "\n", + "\n", + "\n", + "for road_line in road_lines.values():\n", + " # print(len(driveway))\n", + " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "\n", + "ax.set_xlim(lane_[:, 0].min() - 10, lane_[:, 0].max() + 10)\n", + "ax.set_ylim(lane_[:, 1].min() - 10, lane_[:, 1].max() + 10)\n", + "\n", + "lanes_map_features[lane_id]" ] }, { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "8", "metadata": {}, "outputs": [], "source": [ + "from d123.common.geometry.base import StateSE2\n", + "from d123.common.geometry.line.polylines import PolylineSE2\n", + "from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", + "from d123.common.geometry.vector import Vector2D\n", + "from d123.common.visualization.color.color import TAB_10\n", + "\n", + "size = 20\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "lane_id = 114\n", + "lane_ = lanes[lane_id]\n", + "\n", + "lane_polyline_se2 = PolylineSE2.from_array(lane_[:, :2])\n", + "BOUNDARY_STEP_SIZE = 0.5\n", + "\n", + "\n", + "distances_se2 = np.linspace(\n", + " 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True\n", + ")\n", + "lane_query_se2 = lane_polyline_se2.interpolate(distances_se2)\n", + "\n", + "\n", + "left_boundary = left_boundaries[lane_id].array\n", + "right_boundary = right_boundaries[lane_id].array\n", + "\n", + "assert len(left_boundary) > 0 and len(right_boundary) > 0\n", + "# ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", + "# ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "\n", + "# ax.scatter(lane_query_se2[:, 0], lane_query_se2[:, 1], c=lane_query_se2[:, 2] / np.pi, cmap=\"viridis\")\n", + "\n", + "\n", + "# MAX_LANE_WIDTH = 15\n", + "# for state_se2_array in lane_query_se2:\n", + "# for sign in [1.0, -1.0]:\n", + "# perp_end_point = translate_along_yaw(\n", + "# StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)\n", + "# )\n", + "# ax.plot(\n", + "# [state_se2_array[0], perp_end_point.x],\n", + "# [state_se2_array[1], perp_end_point.y],\n", + "# color=\"lime\" if sign > 0 else \"red\",\n", + "# linestyle=\"--\",\n", + "# alpha=0.5,\n", + "# )\n", + "\n", "\n", + "for idx, right_boundary in enumerate(lanes_map_features[lane_id].lane.right_boundaries):\n", + " lane_start_index = right_boundary.lane_start_index\n", + " lane_end_index = right_boundary.lane_end_index\n", "\n", + " road_edge = road_edges[right_boundary.boundary_feature_id]\n", + " if len(lane_[lane_start_index:lane_end_index, 0]) < 10:\n", + " continue\n", + "\n", + " ax.plot(\n", + " lane_[lane_start_index:lane_end_index, 0],\n", + " lane_[lane_start_index:lane_end_index, 1],\n", + " color=TAB_10[idx % len(TAB_10)].hex,\n", + " )\n", + " ax.plot(\n", + " road_edge[:, 0],\n", + " road_edge[:, 1],\n", + " color=TAB_10[idx % len(TAB_10)].hex,\n", + " label=\"road_edge\",\n", + " linestyle=\"dashdot\",\n", + " alpha=0.5,\n", + " )\n", + " print(len(road_edge))\n", + "\n", + "\n", + "# # print(len(driveway))\n", + "# ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\", label=\"right_boundary\", linestyle=\"--\")\n", + "\n", + "\n", + "# for road_edge in road_edges.values():\n", + "# # print(len(driveway))\n", + "# ax.plot(road_edge[:, 0], road_edge[:, 1], label=\"road_edge\", linestyle=\"dashdot\")\n", + "\n", + "\n", + "# for road_line in road_lines.values():\n", + "# # print(len(driveway))\n", + "# ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "\n", + "zoom_out = 100\n", + "ax.set_xlim(lane_[:, 0].min() - zoom_out, lane_[:, 0].max() + zoom_out)\n", + "ax.set_ylim(lane_[:, 1].min() - zoom_out, lane_[:, 1].max() + zoom_out)\n", + "lanes_map_features[lane_id]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ "from d123.common.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", "\n", "\n", - "fig, ax = plt.subplots(figsize=(10, 10))\n", + "fig, ax = plt.subplots(figsize=(30, 30))\n", "\n", - "for road_edge in road_edges:\n", + "for road_edge in road_edges.values():\n", " # print(len(driveway))\n", " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"blue\", label=\"road_edge\")\n", "\n", - "for lane in lanes:\n", + "for lane in lanes.values():\n", " # print(len(driveway))\n", " ax.plot(lane[:, 0], lane[:, 1], color=\"gray\", label=\"lane\")\n", "\n", "\n", - "for road_line in road_lines:\n", + "for road_line in road_lines.values():\n", " # print(len(driveway))\n", " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", "\n", - "for driveway in driveways:\n", + "for driveway in driveways.values():\n", " # print(len(driveway))\n", " ax.plot(driveway[:, 0], driveway[:, 1], color=\"green\", label=\"driveway\")\n", "\n", "\n", - "for crosswalk in crosswalks:\n", + "for crosswalk in crosswalks.values():\n", " # print(len(driveway))\n", - " ax.plot(crosswalk[:, 0], crosswalk[:, 1], color=\"green\", label=\"driveway\")\n", + " ax.plot(crosswalk[:, 0], crosswalk[:, 1], color=\"violet\", label=\"crosswalk\")\n", + "\n", "\n", + "add_non_repeating_legend_to_ax(ax)\n", "\n", - "add_non_repeating_legend_to_ax(ax)" + "ax.set_aspect(\"equal\")" ] }, { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "10", "metadata": {}, "outputs": [], "source": [ @@ -241,115 +461,242 @@ "\n", "# Create left and right boundaries for the lane\n", "\n", + "\n", "def create_lane_boundaries(polyline_3d, width=2.0):\n", " \"\"\"\n", " Create left and right boundaries for a lane by offsetting perpendicular to the heading.\n", - " \n", + "\n", " Args:\n", " polyline_3d: Polyline3D object representing the centerline\n", " width: Total width of the lane in meters\n", - " \n", + "\n", " Returns:\n", " left_boundary: numpy array representing the left boundary\n", " right_boundary: numpy array representing the right boundary\n", " \"\"\"\n", " points = polyline_3d.array\n", " half_width = width / 2.0\n", - " \n", + "\n", " # Calculate the direction vectors between consecutive points\n", " directions = np.diff(points, axis=0)\n", - " \n", + "\n", " # Normalize the direction vectors\n", " directions_norm = np.linalg.norm(directions, axis=1, keepdims=True)\n", " directions_normalized = directions / directions_norm\n", - " \n", + "\n", " # Calculate perpendicular vectors in the xy plane (z remains 0)\n", " perpendiculars = np.zeros_like(directions)\n", " perpendiculars[:, 0] = -directions_normalized[:, 1] # -dy\n", - " perpendiculars[:, 1] = directions_normalized[:, 0] # dx\n", - " \n", + " perpendiculars[:, 1] = directions_normalized[:, 0] # dx\n", + "\n", " # Create boundaries (need to handle the last point separately)\n", " left_boundary = points[:-1] + perpendiculars * half_width\n", " right_boundary = points[:-1] - perpendiculars * half_width\n", - " \n", + "\n", " # Handle the last point based on the last direction\n", " last_perp = perpendiculars[-1]\n", " left_boundary = np.vstack([left_boundary, points[-1] + last_perp * half_width])\n", " right_boundary = np.vstack([right_boundary, points[-1] - last_perp * half_width])\n", - " \n", + "\n", " return left_boundary, right_boundary\n", "\n", + "\n", "# Create the boundaries with a 4m wide lane\n", "left_boundary, right_boundary = create_lane_boundaries(polyline_3d, width=4.0)\n", "\n", "# Plot the centerline and the boundaries\n", "plt.figure(figsize=(12, 8))\n", - "plt.plot(polyline_3d.array[:, 0], polyline_3d.array[:, 1], 'g-', label='Centerline')\n", - "plt.plot(left_boundary[:, 0], left_boundary[:, 1], 'b-', label='Left Boundary')\n", - "plt.plot(right_boundary[:, 0], right_boundary[:, 1], 'r-', label='Right Boundary')\n", - "plt.axis('equal')\n", + "plt.plot(polyline_3d.array[:, 0], polyline_3d.array[:, 1], \"g-\", label=\"Centerline\")\n", + "plt.plot(left_boundary[:, 0], left_boundary[:, 1], \"b-\", label=\"Left Boundary\")\n", + "plt.plot(right_boundary[:, 0], right_boundary[:, 1], \"r-\", label=\"Right Boundary\")\n", + "plt.axis(\"equal\")\n", "plt.grid(True)\n", "plt.legend()\n", - "plt.title('Lane with Left and Right Boundaries')\n", - "\n" + "plt.title(\"Lane with Left and Right Boundaries\")" ] }, { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "11", "metadata": {}, "outputs": [], "source": [ - "for driveway in driveways:\n", - " # print(len(driveway))\n", - " distance = abs(driveway[:, -1].min() - driveway[:, -1].max())\n", - " if distance > 1:\n", - " print(driveway[:, -1])\n", - " plt.plot(driveway[:, 0], driveway[:, 1], \"r-\", color=\"red\")\n", - " else:\n", - " plt.plot(driveway[:, 0], driveway[:, 1], \"r-\", color=\"blue\")\n" + "lane_idx = 59\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "# for right_neighbor in lane.right_neighbors:\n", + "# print(right_neighbor)\n", + "\n", + "# lane\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from torch import le\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.plot(lanes[lane_id][:, 0], lanes[lane_id][:, 1], color=\"blue\", label=\"Lane Centerline\")\n", + "\n", + "for right_neighbor in lane.right_neighbors:\n", + " # print(\"right\", right_neighbor)\n", + " ax.plot(\n", + " lanes[right_neighbor.feature_id][:, 0],\n", + " lanes[right_neighbor.feature_id][:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[right_neighbor.feature_id][right_neighbor.neighbor_start_index: right_neighbor.neighbor_end_index, 0],\n", + " lanes[right_neighbor.feature_id][right_neighbor.neighbor_start_index: right_neighbor.neighbor_end_index, 1],\n", + " color=\"red\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "for left_neighbor in lane.left_neighbors:\n", + " # print(\"left\", left_neighbor)\n", + " ax.plot(\n", + " lanes[left_neighbor.feature_id][:, 0],\n", + " lanes[left_neighbor.feature_id][:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[left_neighbor.feature_id][left_neighbor.neighbor_start_index: left_neighbor.neighbor_end_index, 0],\n", + " lanes[left_neighbor.feature_id][left_neighbor.neighbor_start_index: left_neighbor.neighbor_end_index, 1],\n", + " color=\"lime\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes" ] }, { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "12", "metadata": {}, "outputs": [], "source": [ + "lane_idx = 48\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "\n", + "\n", "import matplotlib.pyplot as plt\n", "\n", - "for polyline in polylines:\n", - " try:\n", - " plt.plot(polyline[:, 0], polyline[:, 1])\n", - " except Exception as e:\n", - " print(f\"Error plotting polyline: {polyline.shape}\")\n" + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.scatter(lanes[lane_id][0, 0], lanes[lane_id][0, 1], color=\"black\", label=\"Lane Centerline\", alpha=0.5)\n", + "ax.scatter(lanes[lane_id][-1, 0], lanes[lane_id][-1, 1], color=\"blue\", label=\"Lane Centerline\", alpha=0.5)\n", + "\n", + "for entry_lane in lane.entry_lanes:\n", + " # print(\"right\", right_neighbor)\n", + " ax.plot(\n", + " lanes[entry_lane][:, 0],\n", + " lanes[entry_lane][:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "for exit_lane in lane.exit_lanes:\n", + " # print(\"left\", left_neighbor)\n", + " ax.plot(\n", + " lanes[exit_lane][:, 0],\n", + " lanes[exit_lane][:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes" ] }, { "cell_type": "code", "execution_count": null, - "id": "9", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "10", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "11", + "id": "13", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "lane_idx = 58\n", + "\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "ax.plot(lanes[lane_id][:, 0], lanes[lane_id][:, 1], color=\"black\", label=\"Lane Centerline\", alpha=0.5)\n", + "\n", + "for left_boundary in lane.left_boundaries:\n", + " print(\"left\", left_boundary)\n", + " try:\n", + " boundary = road_lines[left_boundary.boundary_feature_id]\n", + "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {left_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[left_boundary.boundary_feature_id]\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 0],\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 1],\n", + " color=\"lime\",\n", + " linestyle=\"-\",\n", + " label=\"Left Neighbor\",\n", + " )\n", + "\n", + "\n", + "for right_boundary in lane.right_boundaries:\n", + " print(\"right\", right_boundary)\n", + " try:\n", + " boundary = road_lines[right_boundary.boundary_feature_id]\n", + "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {right_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[right_boundary.boundary_feature_id]\n", + "\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 0],\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 1],\n", + " color=\"red\",\n", + " linestyle=\"-\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + "\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "# lanes" + ] } ], "metadata": { From f257c851f2a0772bd2f901e7c1f263fce9247856 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Thu, 7 Aug 2025 17:50:58 +0200 Subject: [PATCH 21/42] Update waymo map conversion --- d123/common/visualization/viser/server_map.py | 167 +++++++ d123/common/visualization/viser/utils.py | 46 +- .../waymo_map_utils/womp_boundary_utils.py | 296 ++++++++++++ .../wopd/waymo_map_utils/wopd_map_utils.py | 286 ++++++++++++ .../wopd/womp_map_conversion.py | 203 -------- .../dataset_specific/wopd/wopd_map_utils.py | 442 ------------------ notebooks/viz/bev_matplotlib.ipynb | 34 +- notebooks/viz/viser_testing_v2_map.ipynb | 116 +++++ notebooks/waymo_perception/map_testing.ipynb | 366 ++++++++------- 9 files changed, 1106 insertions(+), 850 deletions(-) create mode 100644 d123/common/visualization/viser/server_map.py create mode 100644 d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py create mode 100644 d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py delete mode 100644 d123/dataset/dataset_specific/wopd/womp_map_conversion.py delete mode 100644 d123/dataset/dataset_specific/wopd/wopd_map_utils.py create mode 100644 notebooks/viz/viser_testing_v2_map.ipynb diff --git a/d123/common/visualization/viser/server_map.py b/d123/common/visualization/viser/server_map.py new file mode 100644 index 00000000..551fa2ba --- /dev/null +++ b/d123/common/visualization/viser/server_map.py @@ -0,0 +1,167 @@ +import time +from typing import List, Literal + +import viser + +from d123.common.datatypes.sensor.camera import CameraType +from d123.common.visualization.color.color import BLACK, TAB_10 +from d123.common.visualization.viser.utils import get_map_lines, get_map_meshes +from d123.dataset.scene.abstract_scene import AbstractScene + +# TODO: Try to fix performance issues. +# TODO: Refactor this file. + +all_camera_types: List[CameraType] = [ + CameraType.CAM_F0, + CameraType.CAM_B0, + CameraType.CAM_L0, + CameraType.CAM_L1, + CameraType.CAM_L2, + CameraType.CAM_R0, + CameraType.CAM_R1, + CameraType.CAM_R2, +] + +LIDAR_POINT_SIZE: float = 0.05 +MAP_AVAILABLE: bool = True +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" +LINE_WIDTH: float = 4.0 + +CAMERA_SCALE: float = 1.0 + +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ +# CameraType.CAM_F0, +# CameraType.CAM_L0, +# CameraType.CAM_R0, +# CameraType.CAM_L1, +# CameraType.CAM_R1, +# ] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] +VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] +LIDAR_AVAILABLE: bool = False + + +class ViserMapVisualizationServer: + def __init__( + self, + scenes: List[AbstractScene], + scene_index: int = 0, + host: str = "localhost", + port: int = 8080, + label: str = "D123 Viser Server", + ): + assert len(scenes) > 0, "At least one scene must be provided." + self.scenes = scenes + self.scene_index = scene_index + + self.host = host + self.port = port + self.label = label + + self.server = viser.ViserServer(host=self.host, port=self.port, label=self.label) + self.set_scene(self.scenes[self.scene_index % len(self.scenes)]) + + def next(self) -> None: + self.server.flush() + self.server.gui.reset() + self.server.scene.reset() + self.scene_index = (self.scene_index + 1) % len(self.scenes) + print(f"Viser server started at {self.host}:{self.port}") + self.set_scene(self.scenes[self.scene_index]) + + def set_scene(self, scene: AbstractScene) -> None: + num_frames = scene.get_number_of_iterations() + print(scene.available_camera_types) + + self.server.gui.configure_theme(control_width="large") + with self.server.gui.add_folder("Playback"): + server_playing = True + + gui_timestep = self.server.gui.add_slider( + "Timestep", + min=0, + max=num_frames - 1, + step=1, + initial_value=0, + disabled=True, + ) + gui_next_frame = self.server.gui.add_button("Next Frame", disabled=True) + gui_prev_frame = self.server.gui.add_button("Prev Frame", disabled=True) + gui_next_scene = self.server.gui.add_button("Next Scene", disabled=False) + gui_playing = self.server.gui.add_checkbox("Playing", True) + gui_framerate = self.server.gui.add_slider("FPS", min=1, max=60, step=0.1, initial_value=10) + gui_framerate_options = self.server.gui.add_button_group("FPS options", ("10", "20", "30", "60")) + + # Frame step buttons. + @gui_next_frame.on_click + def _(_) -> None: + gui_timestep.value = (gui_timestep.value + 1) % num_frames + + @gui_prev_frame.on_click + def _(_) -> None: + gui_timestep.value = (gui_timestep.value - 1) % num_frames + + @gui_next_scene.on_click + def _(_) -> None: + nonlocal server_playing + server_playing = False + + # Disable frame controls when we're playing. + @gui_playing.on_update + def _(_) -> None: + gui_timestep.disabled = gui_playing.value + gui_next_frame.disabled = gui_playing.value + gui_prev_frame.disabled = gui_playing.value + + # Set the framerate when we click one of the options. + @gui_framerate_options.on_click + def _(_) -> None: + gui_framerate.value = int(gui_framerate_options.value) + + # Frame step buttons + + # Toggle frame visibility when the timestep slider changes. + @gui_timestep.on_update + def _(_) -> None: + + time.sleep(0.1) + + prev_timestep = gui_timestep.value + + # Load in frames. + if MAP_AVAILABLE: + for name, mesh in get_map_meshes(scene).items(): + self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) + + centerlines, left_boundaries, right_boundaries = get_map_lines(scene) + # for i, centerline in enumerate(centerlines): + self.server.scene.add_line_segments( + "/map/centerlines", + centerlines, + colors=[[BLACK.rgb]], + line_width=LINE_WIDTH, + ) + self.server.scene.add_line_segments( + "/map/left_boundary", + left_boundaries, + colors=[[TAB_10[2].rgb]], + line_width=LINE_WIDTH, + ) + self.server.scene.add_line_segments( + "/map/right_boundary", + right_boundaries, + colors=[[TAB_10[3].rgb]], + line_width=LINE_WIDTH, + ) + + # Playback update loop. + prev_timestep = gui_timestep.value + while server_playing: + # Update the timestep if we're playing. + if gui_playing.value: + gui_timestep.value = (gui_timestep.value + 1) % num_frames + + self.server.flush() + self.next() diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 341c629e..5b765e0c 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -4,6 +4,7 @@ import numpy.typing as npt import trimesh from pyquaternion import Quaternion +from typing_extensions import Final # from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters from d123.common.datatypes.sensor.camera import Camera, CameraType @@ -14,12 +15,14 @@ from d123.common.visualization.color.config import PlotConfig from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG, MAP_SURFACE_CONFIG from d123.dataset.maps.abstract_map import MapSurfaceType -from d123.dataset.maps.abstract_map_objects import AbstractSurfaceMapObject +from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractSurfaceMapObject from d123.dataset.scene.abstract_scene import AbstractScene # TODO: Refactor this file. # TODO: Add general utilities for 3D primitives and mesh support. +MAP_RADIUS: Final[float] = 500 + def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> trimesh.Trimesh: @@ -84,15 +87,16 @@ def get_map_meshes(scene: AbstractScene): initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 map_surface_types = [ - # MapSurfaceType.LANE_GROUP, - MapSurfaceType.LANE, + MapSurfaceType.LANE_GROUP, + # MapSurfaceType.LANE, MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK, MapSurfaceType.CARPARK, ] - radius = 500 - map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=radius, layers=map_surface_types) + map_objects_dict = scene.map_api.get_proximal_map_objects( + center.point_2d, radius=MAP_RADIUS, layers=map_surface_types + ) output = {} for map_surface_type in map_objects_dict.keys(): @@ -101,10 +105,8 @@ def get_map_meshes(scene: AbstractScene): map_surface: AbstractSurfaceMapObject trimesh_mesh = map_surface.trimesh_mesh if map_surface_type in [MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK]: - # trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=1.777 / 2 - 0.05).array trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.05).array else: - # trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=1.777 / 2).array trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z).array if not scene.log_metadata.map_has_z: @@ -119,6 +121,36 @@ def get_map_meshes(scene: AbstractScene): return output +def get_map_lines(scene: AbstractScene): + map_surface_types = [MapSurfaceType.LANE] + initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) + center = initial_ego_vehicle_state.center_se3 + map_objects_dict = scene.map_api.get_proximal_map_objects( + center.point_2d, radius=MAP_RADIUS, layers=map_surface_types + ) + + def polyline_to_segments(polyline: Polyline3D) -> npt.NDArray[np.float64]: + polyline_array = polyline.array - center.point_3d.array + polyline_array = polyline_array.reshape(-1, 1, 3) + polyline_array = np.concatenate([polyline_array[:-1], polyline_array[1:]], axis=1) + return polyline_array + + lanes = map_objects_dict[MapSurfaceType.LANE] + centerlines, right_boundaries, left_boundaries = [], [], [] + for lane in lanes: + lane: AbstractLane + + centerlines.append(polyline_to_segments(lane.centerline)) + right_boundaries.append(polyline_to_segments(lane.right_boundary)) + left_boundaries.append(polyline_to_segments(lane.left_boundary)) + + centerlines = np.concatenate(centerlines, axis=0) + left_boundaries = np.concatenate(left_boundaries, axis=0) + right_boundaries = np.concatenate(right_boundaries, axis=0) + + return centerlines, left_boundaries, right_boundaries + + def get_trimesh_from_boundaries( left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 1.0 ) -> trimesh.Trimesh: diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py new file mode 100644 index 00000000..753ca36a --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py @@ -0,0 +1,296 @@ +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import numpy as np +import numpy.typing as npt +import shapely.geometry as geom + +from d123.common.geometry.base import Point3D, StateSE2 +from d123.common.geometry.line.polylines import Polyline3D, PolylineSE2 +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.common.geometry.transform.tranform_2d import translate_along_yaw +from d123.common.geometry.utils import normalize_angle +from d123.common.geometry.vector import Vector2D + +MAX_LANE_WIDTH = 25.0 # meters +MIN_LANE_WIDTH = 2.0 +DEFAULT_LANE_WIDTH = 3.7 +BOUNDARY_STEP_SIZE = 0.1 # meters +MAX_Z_DISTANCE = 1.0 # meters + +PERP_START_OFFSET = 0.1 # meters + +MIN_HIT_DISTANCE = 0.1 # meters +MIN_AVERAGE_DISTANCE = 1.5 +MAX_AVERAGE_DISTANCE = 6.0 + + +def get_type_and_id_from_token(token: str) -> Tuple[str, int]: + """Extract type and id from token.""" + line_type, line_id = token.split("_") + return line_type, int(line_id) + + +def get_polyline_from_token(polyline_dict: Dict[str, Dict[int, Polyline3D]], token: str) -> Polyline3D: + """Extract polyline from token.""" + line_type, line_id = get_type_and_id_from_token(token) + return polyline_dict[line_type][line_id] + + +@dataclass +class PerpendicularHit: + distance_along_perp_2d: float + hit_point_3d: Point3D + hit_polyline_token: str + centerline_hit_crossing: bool + heading_error: Optional[float] = None + + @property + def hit_polyline_id(self) -> int: + """Extract polyline id from token.""" + return get_type_and_id_from_token(self.hit_polyline_token)[1] + + @property + def hit_polyline_type(self) -> int: + """Extract polyline id from token.""" + return get_type_and_id_from_token(self.hit_polyline_token)[0] + + +def _collect_perpendicular_hits( + lane_query_se2: StateSE2, + lane_token: str, + polyline_dict: Dict[str, Dict[int, Polyline3D]], + lane_polyline_se2_dict: Dict[int, PolylineSE2], + occupancy_2d: OccupancyMap2D, + sign: float, +) -> List[PerpendicularHit]: + assert sign in [1.0, -1.0], "Sign must be either 1.0 (left) or -1.0 (right)" + # perp_start_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * PERP_START_OFFSET)) + perp_start_point = lane_query_se2 + perp_end_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)) + perp_linestring = geom.LineString([[perp_start_point.x, perp_start_point.y], [perp_end_point.x, perp_end_point.y]]) + + lane_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[lane_token]] + + # 1. find intersecting lines, compute 3D distance + intersecting_tokens = occupancy_2d.intersects(perp_linestring) + + perpendicular_hits: List[PerpendicularHit] = [] + for intersecting_token in intersecting_tokens: + intersecting_polyline_3d = get_polyline_from_token(polyline_dict, intersecting_token) + intersecting_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[intersecting_token]] + centerline_hit_crossing: bool = ( + lane_linestring.intersects(intersecting_linestring) if intersecting_token.startswith("lane_") else False + ) + + intersecting_geom_points: List[geom.Point] = [] + intersecting_geometries = perp_linestring.intersection(intersecting_linestring) + if isinstance(intersecting_geometries, geom.Point): + intersecting_geom_points.append(intersecting_geometries) + elif isinstance(intersecting_geometries, geom.MultiPoint): + intersecting_geom_points.extend([geom for geom in intersecting_geometries.geoms]) + + for intersecting_geom_point in intersecting_geom_points: + distance_along_perp_2d = perp_linestring.project(intersecting_geom_point) + + distance_along_intersecting_norm = intersecting_linestring.project(intersecting_geom_point, normalized=True) + intersecting_point_3d = intersecting_polyline_3d.interpolate( + distance_along_intersecting_norm * intersecting_polyline_3d.length + ) + + heading_error = None + if intersecting_token.startswith("lane_"): + # Compute heading error if the intersecting token is a lane + intersecting_polyline_se2 = lane_polyline_se2_dict[intersecting_token] + lane_heading = intersecting_polyline_se2.interpolate( + distance_along_intersecting_norm * intersecting_polyline_se2.length + ) + heading_error = normalize_angle(lane_query_se2.yaw - lane_heading.yaw) + + perpendicular_hits.append( + PerpendicularHit( + distance_along_perp_2d=distance_along_perp_2d, + hit_point_3d=intersecting_point_3d, + hit_polyline_token=intersecting_token, + centerline_hit_crossing=centerline_hit_crossing, + heading_error=heading_error, + ) + ) + + return perpendicular_hits + + +def _filter_perpendicular_hits( + perpendicular_hits: List[PerpendicularHit], + lane_point_3d: Point3D, +) -> List[PerpendicularHit]: + + filtered_hits = [] + for hit in perpendicular_hits: + + # 1. filter hits too far in the vertical direction + z_distance = np.abs(hit.hit_point_3d.z - lane_point_3d.z) + if z_distance > MAX_Z_DISTANCE: + continue + + # 2. filter hits that are too close and not with the road edge (e.g. close lane lines) + if hit.distance_along_perp_2d < MIN_HIT_DISTANCE and hit.hit_polyline_type != "roadedge": + continue + + filtered_hits.append(hit) + + # Sort hits by distance_along_perp_2d + filtered_hits.sort(key=lambda hit: hit.distance_along_perp_2d) + + return filtered_hits + + +def extract_lane_boundaries( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + road_lines: Dict[int, npt.NDArray[np.float64]], + road_edges: Dict[int, npt.NDArray[np.float64]], +) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: + polyline_dict: Dict[str, Dict[int, Polyline3D]] = {"lane": {}, "roadline": {}, "roadedge": {}} + lane_polyline_se2_dict: Dict[int, PolylineSE2] = {} + + for lane_id, lane_polyline in lanes.items(): + if lane_polyline.ndim == 2 and lane_polyline.shape[1] == 3 and len(lane_polyline) > 0: + polyline_dict["lane"][lane_id] = Polyline3D.from_array(lane_polyline) + lane_polyline_se2_dict[f"lane_{lane_id}"] = polyline_dict["lane"][lane_id].polyline_se2 + + # for road_line_id, road_line_polyline in road_lines.items(): + # if road_line_polyline.ndim == 2 and road_line_polyline.shape[1] == 3 and len(road_line_polyline) > 0: + # polyline_dict["roadline"][road_line_id] = Polyline3D.from_array(road_line_polyline) + + for road_edge_id, road_edge_polyline in road_edges.items(): + if road_edge_polyline.ndim == 2 and road_edge_polyline.shape[1] == 3 and len(road_edge_polyline) > 0: + polyline_dict["roadedge"][road_edge_id] = Polyline3D.from_array(road_edge_polyline) + + geometries = [] + tokens = [] + for line_type, polylines in polyline_dict.items(): + for polyline_id, polyline in polylines.items(): + geometries.append(polyline.polyline_2d.linestring) + tokens.append(f"{line_type}_{polyline_id}") + + occupancy_2d = OccupancyMap2D(geometries, tokens) + + left_boundaries = {} + right_boundaries = {} + + for lane_id, lane_polyline in polyline_dict["lane"].items(): + current_lane_token = f"lane_{lane_id}" + lane_polyline_se2 = lane_polyline_se2_dict[current_lane_token] + + # 1. sample poses along centerline + distances_se2 = np.linspace( + 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_queries_se2 = [ + StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) + ] + distances_3d = np.linspace( + 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True + ) + lane_queries_3d = [ + Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) + ] + assert len(lane_queries_se2) == len(lane_queries_3d) + + for sign in [1.0, -1.0]: + boundary_points_3d: List[Optional[Point3D]] = [] + for lane_query_se2, lane_query_3d in zip(lane_queries_se2, lane_queries_3d): + + perpendicular_hits = _collect_perpendicular_hits( + lane_query_se2=lane_query_se2, + lane_token=current_lane_token, + polyline_dict=polyline_dict, + lane_polyline_se2_dict=lane_polyline_se2_dict, + occupancy_2d=occupancy_2d, + sign=sign, + ) + perpendicular_hits = _filter_perpendicular_hits( + perpendicular_hits=perpendicular_hits, lane_point_3d=lane_query_3d + ) + + boundary_point_3d: Optional[Point3D] = None + # 1. First, try to find the boundary point from the perpendicular hits + if len(perpendicular_hits) > 0: + first_hit = perpendicular_hits[0] + + # 1.1. If the first hit is a road edge, use it as the boundary point + if first_hit.hit_polyline_type == "roadedge": + boundary_point_3d = first_hit.hit_point_3d + elif first_hit.hit_polyline_type == "roadline": + boundary_point_3d = first_hit.hit_point_3d + elif first_hit.hit_polyline_type == "lane": + + for hit in perpendicular_hits: + if hit.hit_polyline_type == "roadedge": + continue + if hit.hit_polyline_type == "lane": + + has_same_predecessor = ( + len(set(lanes_predecessors[hit.hit_polyline_id]) & set(lanes_predecessors[lane_id])) + > 0 + ) + has_same_successor = ( + len(set(lanes_successors[hit.hit_polyline_id]) & set(lanes_successors[lane_id])) > 0 + ) + heading_min = np.pi / 8.0 + invalid_heading_error = heading_min < abs(hit.heading_error) < (np.pi - heading_min) + if ( + not has_same_predecessor + and not has_same_successor + and not hit.centerline_hit_crossing + and MAX_AVERAGE_DISTANCE > hit.distance_along_perp_2d + and MIN_AVERAGE_DISTANCE < hit.distance_along_perp_2d + and not invalid_heading_error + ): + # 2. if first hit is lane line, use it as boundary point + boundary_point_3d = Point3D.from_array( + (hit.hit_point_3d.array + lane_query_3d.array) / 2.0 + ) + break + + boundary_points_3d.append(boundary_point_3d) + + no_boundary_ratio = boundary_points_3d.count(None) / len(boundary_points_3d) + final_boundary_points_3d = [] + + def _get_default_boundary_point_3d( + lane_query_se2: StateSE2, lane_query_3d: Point3D, sign: float + ) -> Point3D: + perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 + boundary_point_se2 = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance)) + return Point3D(boundary_point_se2.x, boundary_point_se2.y, lane_query_3d.z) + + if no_boundary_ratio > 0.8: + for lane_query_se2, lane_query_3d in zip(lane_queries_se2, lane_queries_3d): + boundary_point_3d = _get_default_boundary_point_3d(lane_query_se2, lane_query_3d, sign) + final_boundary_points_3d.append(boundary_point_3d.array) + + else: + for boundary_idx, (lane_query_se2, lane_query_3d) in enumerate(zip(lane_queries_se2, lane_queries_3d)): + if boundary_points_3d[boundary_idx] is None: + boundary_point_3d = _get_default_boundary_point_3d(lane_query_se2, lane_query_3d, sign) + else: + boundary_point_3d = boundary_points_3d[boundary_idx] + final_boundary_points_3d.append(boundary_point_3d.array) + + # # 2. If no boundary point was found, use the lane query point as the boundary point + # if boundary_point_3d is None: + + if len(final_boundary_points_3d) > 1: + if sign == 1.0: + left_boundaries[lane_id] = Polyline3D.from_array( + np.array(final_boundary_points_3d, dtype=np.float64) + ) + else: + right_boundaries[lane_id] = Polyline3D.from_array( + np.array(final_boundary_points_3d, dtype=np.float64) + ) + + return left_boundaries, right_boundaries diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py new file mode 100644 index 00000000..d28df0e6 --- /dev/null +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py @@ -0,0 +1,286 @@ +from collections import defaultdict +from pathlib import Path +from typing import Dict, List + +import geopandas as gpd +import numpy as np +import numpy.typing as npt +import pandas as pd +import shapely.geometry as geom +from waymo_open_dataset import dataset_pb2 + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.units import mph_to_mps +from d123.dataset.dataset_specific.wopd.waymo_map_utils.womp_boundary_utils import extract_lane_boundaries +from d123.dataset.maps.map_datatypes import MapSurfaceType + +# TODO: +# - Implement stop signs +# - Implement speed bumps +# - Implement driveways with a different semantic type if needed +# - Implement intersections and lane group logic +# - Handle lane type, e.g. TYPE_UNDEFINED = 0; TYPE_FREEWAY = 1; TYPE_SURFACE_STREET = 2; TYPE_BIKE_LANE = 3; + + +def convert_wopd_map(frame: dataset_pb2.Frame, map_file_path: Path) -> None: + + def _extract_polyline(data) -> npt.NDArray[np.float64]: + polyline = np.array([[p.x, p.y, p.z] for p in data.polyline], dtype=np.float64) + return polyline + + def _extract_polygon(data) -> npt.NDArray[np.float64]: + polygon = np.array([[p.x, p.y, p.z] for p in data.polygon], dtype=np.float64) + assert polygon.shape[0] >= 3, "Polygon must have at least 3 points" + assert polygon.shape[1] == 3, "Polygon must have 3 coordinates (x, y, z)" + return polygon + + lanes: Dict[int, npt.NDArray[np.float64]] = {} + lanes_successors = defaultdict(list) + lanes_predecessors = defaultdict(list) + lanes_speed_limit_mps: Dict[int, float] = {} + lanes_type: Dict[int, int] = {} + + road_lines: Dict[int, npt.NDArray[np.float64]] = {} + road_edges: Dict[int, npt.NDArray[np.float64]] = {} + crosswalks: Dict[int, npt.NDArray[np.float64]] = {} + carparks: Dict[int, npt.NDArray[np.float64]] = {} + + for map_feature in frame.map_features: + if map_feature.HasField("lane"): + polyline = _extract_polyline(map_feature.lane) + # Ignore lanes with less than 2 points or not 2D + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + lanes[map_feature.id] = polyline + for lane_id_ in map_feature.lane.exit_lanes: + lanes_successors[map_feature.id].append(lane_id_) + for lane_id_ in map_feature.lane.exit_lanes: + lanes_predecessors[map_feature.id].append(lane_id_) + lanes_speed_limit_mps[map_feature.id] = mph_to_mps(map_feature.lane.speed_limit_mph) + lanes_type[map_feature.id] = map_feature.lane.type + elif map_feature.HasField("road_line"): + polyline = _extract_polyline(map_feature.road_line) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + road_lines[map_feature.id] = polyline + elif map_feature.HasField("road_edge"): + polyline = _extract_polyline(map_feature.road_edge) + if polyline.ndim != 2 or polyline.shape[0] < 2: + continue + road_edges[map_feature.id] = polyline + elif map_feature.HasField("stop_sign"): + # TODO: implement stop signs + pass + elif map_feature.HasField("crosswalk"): + crosswalks[map_feature.id] = _extract_polygon(map_feature.crosswalk) + elif map_feature.HasField("speed_bump"): + # TODO: implement speed bumps + pass + elif map_feature.HasField("driveway"): + # NOTE: Determine whether to use a different semantic type for driveways. + carparks[map_feature.id] = _extract_polygon(map_feature.driveway) + + lane_left_boundaries_3d, lane_right_boundaries_3d = extract_lane_boundaries( + lanes, lanes_successors, lanes_predecessors, road_lines, road_edges + ) + + lane_df = get_lane_df( + lanes, + lanes_successors, + lanes_predecessors, + lanes_speed_limit_mps, + lane_left_boundaries_3d, + lane_right_boundaries_3d, + ) + lane_group_df = get_lane_group_df( + lanes, + lanes_successors, + lanes_predecessors, + lane_left_boundaries_3d, + lane_right_boundaries_3d, + ) + intersection_df = get_intersections_df() + crosswalk_df = get_crosswalk_df(crosswalks) + walkway_df = get_walkway_df() + carpark_df = get_carpark_df(carparks) + generic_drivable_df = get_generic_drivable_df() + + map_file_path.unlink(missing_ok=True) + if not map_file_path.parent.exists(): + map_file_path.parent.mkdir(parents=True, exist_ok=True) + + lane_df.to_file(map_file_path, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") + lane_group_df.to_file(map_file_path, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") + intersection_df.to_file(map_file_path, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") + crosswalk_df.to_file(map_file_path, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") + walkway_df.to_file(map_file_path, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") + carpark_df.to_file(map_file_path, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") + generic_drivable_df.to_file( + map_file_path, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" + ) + + +def get_lane_df( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + lanes_speed_limit_mps: Dict[int, float], + lanes_left_boundaries_3d: Dict[int, Polyline3D], + lanes_right_boundaries_3d: Dict[int, Polyline3D], +) -> gpd.GeoDataFrame: + + ids = [] + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + baseline_paths = [] + geometries = [] + + for lane_id, lane_centerline_array in lanes.items(): + if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: + continue + lane_centerline = Polyline3D.from_array(lane_centerline_array) + lane_speed_limit_mps = lanes_speed_limit_mps[lane_id] if lanes_speed_limit_mps[lane_id] > 0.0 else None + + ids.append(lane_id) + lane_group_ids.append([lane_id]) + speed_limits_mps.append(lane_speed_limit_mps) + predecessor_ids.append(lanes_predecessors[lane_id]) + successor_ids.append(lanes_successors[lane_id]) + left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) + right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + baseline_paths.append(lane_centerline.linestring) + + geometry = geom.Polygon( + np.vstack( + [ + lanes_left_boundaries_3d[lane_id].array[:, :2], + lanes_right_boundaries_3d[lane_id].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "baseline_path": baseline_paths, + } + ) + + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_lane_group_df( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + lanes_left_boundaries_3d: Dict[int, Polyline3D], + lanes_right_boundaries_3d: Dict[int, Polyline3D], +) -> gpd.GeoDataFrame: + + ids = [] + lane_ids = [] + intersection_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + # NOTE: WOPD does not provide lane groups, so we create a lane group for each lane. + for lane_id in lanes.keys(): + if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: + continue + ids.append(lane_id) + lane_ids.append([lane_id]) + intersection_ids.append(None) # WOPD does not provide intersections + predecessor_lane_group_ids.append(lanes_predecessors[lane_id]) + successor_lane_group_ids.append(lanes_successors[lane_id]) + left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) + right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + geometry = geom.Polygon( + np.vstack( + [ + lanes_left_boundaries_3d[lane_id].array[:, :2], + lanes_right_boundaries_3d[lane_id].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "intersection_id": intersection_ids, + "predecessor_lane_group_ids": predecessor_lane_group_ids, + "successor_lane_group_ids": successor_lane_group_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_intersections_df() -> gpd.GeoDataFrame: + ids = [] + lane_group_ids = [] + geometries = [] + + # NOTE: WOPD does not provide intersections, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_carpark_df(carparks) -> gpd.GeoDataFrame: + ids = list(carparks.keys()) + outlines = [geom.LineString(outline) for outline in carparks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in carparks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_walkway_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + # NOTE: WOPD does not provide walkways, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(crosswalks.keys()) + outlines = [geom.LineString(outline) for outline in crosswalks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in crosswalks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_generic_drivable_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + # NOTE: WOPD does not provide generic drivable areas, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf diff --git a/d123/dataset/dataset_specific/wopd/womp_map_conversion.py b/d123/dataset/dataset_specific/wopd/womp_map_conversion.py deleted file mode 100644 index 6372f19c..00000000 --- a/d123/dataset/dataset_specific/wopd/womp_map_conversion.py +++ /dev/null @@ -1,203 +0,0 @@ -from dataclasses import dataclass -from typing import Dict, Optional, Tuple - -import numpy as np -import numpy.typing as npt -import shapely.geometry as geom - -from d123.common.geometry.base import Point3D, StateSE2 -from d123.common.geometry.line.polylines import Polyline3D -from d123.common.geometry.occupancy_map import OccupancyMap2D -from d123.common.geometry.transform.tranform_2d import translate_along_yaw -from d123.common.geometry.vector import Vector2D - -MAX_LANE_WIDTH = 25.0 # meters -MIN_LANE_WIDTH = 2.0 -DEFAULT_LANE_WIDTH = 4.0 -BOUNDARY_STEP_SIZE = 0.5 # meters -MAX_Z_DISTANCE = 1.0 # meters - - -@dataclass -class BoundaryHit: - distance_along_perp_2d: float - distance_along_polyline_3d: float - polyline_token: str - - -def _extract_lane_boundaries( - lanes: Dict[int, npt.NDArray[np.float64]], - road_lines: Dict[int, npt.NDArray[np.float64]], - road_edges: Dict[int, npt.NDArray[np.float64]], -) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: - polyline_dict: Dict[str, Dict[int, Polyline3D]] = {"lane": {}, "roadline": {}, "roadedge": {}} - - for lane_id, lane_polyline in lanes.items(): - if lane_polyline.ndim == 2 and lane_polyline.shape[1] == 3 and len(lane_polyline) > 0: - polyline_dict["lane"][lane_id] = Polyline3D.from_array(lane_polyline) - - # for road_line_id, road_line_polyline in road_lines.items(): - # if road_line_polyline.ndim == 2 and road_line_polyline.shape[1] == 3 and len(road_line_polyline) > 0: - # polyline_dict["roadline"][road_line_id] = Polyline3D.from_array(road_line_polyline) - - for road_edge_id, road_edge_polyline in road_edges.items(): - if road_edge_polyline.ndim == 2 and road_edge_polyline.shape[1] == 3 and len(road_edge_polyline) > 0: - polyline_dict["roadedge"][road_edge_id] = Polyline3D.from_array(road_edge_polyline) - - geometries = [] - tokens = [] - for line_type, polylines in polyline_dict.items(): - for polyline_id, polyline in polylines.items(): - geometries.append(polyline.polyline_2d.linestring) - tokens.append(f"{line_type}_{polyline_id}") - - occupancy_2d = OccupancyMap2D(geometries, tokens) - - # for each lane - # - sample poses along centerline - # - query left and right perpendicular lines in 2d occupancy map - # - for {left, right} retrieve first perpendicular intersection with length MAX_HALF_WIDTH - # - if intersection with road edge, add intersecting point to {left, right} boundary - # - if intersection with centerline, add middle point to {left, right} boundary - # - if no intersection found, add MAX_HALF_WIDTH to {left, right} boundary - - left_boundaries = {lane_id: [] for lane_id in lanes.keys()} - right_boundaries = {lane_id: [] for lane_id in lanes.keys()} - - def get_type_and_id_from_token(token: str) -> Tuple[str, int]: - """Extract type and id from token.""" - line_type, line_id = token.split("_") - return line_type, int(line_id) - - def get_polyline_from_token(token: str) -> Polyline3D: - """Extract polyline from token.""" - line_type, line_id = get_type_and_id_from_token(token) - return polyline_dict[line_type][line_id] - - def get_intersecting_point_3d(perp_linestring: geom.LineString, intersecting_token: str) -> Optional[Point3D]: - - intersecting_polyline_3d = get_polyline_from_token(intersecting_token) - intersecting_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[intersecting_token]] - intersecting_point_2d = perp_linestring.intersection(intersecting_linestring) - - distance_2d_norm = None - if isinstance(intersecting_point_2d, geom.Point): - distance_2d_norm = intersecting_linestring.project(intersecting_point_2d, normalized=True) - elif isinstance(intersecting_point_2d, geom.MultiPoint): - geom_points = [point for point in intersecting_point_2d.geoms] - starting_perp_point = geom.Point(perp_linestring.coords[0]) - start_perp_point_distance = np.array( - [starting_perp_point.distance(geom_point) for geom_point in geom_points] - ) - start_perp_point_distance[start_perp_point_distance < MIN_LANE_WIDTH / 2.0] = np.inf - closest_index = np.argmin(start_perp_point_distance) - distance_2d_norm = intersecting_linestring.project(geom_points[closest_index], normalized=True) - - intersecting_point_3d = None - if distance_2d_norm is not None: - intersecting_point_3d = intersecting_polyline_3d.interpolate( - distance_2d_norm * intersecting_polyline_3d.length - ) - if np.linalg.norm(intersecting_point_3d.point_2d.array - perp_linestring.coords[0]) < MIN_LANE_WIDTH / 2.0: - return None - - return intersecting_point_3d - - for lane_id, lane_polyline in polyline_dict["lane"].items(): - current_lane_token = f"lane_{lane_id}" - lane_polyline_se2 = lane_polyline.polyline_se2 - - # 0. Find connected or intersecting lanes - connected_lane_tokens = [ - token - for token in occupancy_2d.intersects(lane_polyline_se2.linestring) - if token.startswith("lane_") - if token != current_lane_token - ] - - # 1. sample poses along centerline - distances_se2 = np.linspace( - 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True - ) - lane_query_se2 = [ - StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) - ] - distances_3d = np.linspace( - 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True - ) - lane_query_3d = [ - Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) - ] - assert len(lane_query_se2) == len(lane_query_3d) - - for lane_query_se2, lane_query_3d in zip(lane_query_se2, lane_query_3d): - for sign in [1.0, -1.0]: - perp_start_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * 0.1)) - perp_end_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)) - perp_linestring = geom.LineString( - [[perp_start_point.x, perp_start_point.y], [perp_end_point.x, perp_end_point.y]] - ) - - # 1. find intersecting lines, compute 3D distance - intersecting_tokens = occupancy_2d.intersects(perp_linestring) - intersecting_points_3d: Dict[str, Point3D] = {} - crosses_centerline = False - - for intersecting_token in intersecting_tokens: - # if intersecting_token in connected_lane_tokens: - # crosses_centerline = True - # continue - # if intersecting_token in connected_lane_tokens: - # continue - - # if intersecting_token == current_lane_token: - # crosses_centerline = True - - intersecting_point_3d = get_intersecting_point_3d(perp_linestring, intersecting_token) - if intersecting_point_3d is None: - continue - - if np.abs(intersecting_point_3d.z - lane_query_3d.z) > MAX_Z_DISTANCE: - continue - - intersecting_points_3d[intersecting_token] = intersecting_point_3d - - boundary_point_3d: Optional[Point3D] = None - if len(intersecting_points_3d) > 0 and not crosses_centerline: - # 2. find closest intersection - intersecting_distances = { - token: np.linalg.norm(point_3d.array - lane_query_3d.array) - for token, point_3d in intersecting_points_3d.items() - } - closest_token = min(intersecting_distances, key=intersecting_distances.get) - closest_type, closest_id = get_type_and_id_from_token(closest_token) - - if closest_type == "lane": - boundary_point_3d = Point3D.from_array( - (intersecting_points_3d[closest_token].array + lane_query_3d.array) / 2.0 - ) - elif closest_type == "roadedge": - boundary_point_3d = intersecting_points_3d[closest_token] - else: - perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 # Default to half the lane width - boundary_point_se2 = translate_along_yaw( - lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance) - ) - boundary_point_3d = Point3D(boundary_point_se2.x, boundary_point_se2.y, lane_query_3d.z) - if sign == 1.0: - left_boundaries[lane_id].append(boundary_point_3d.array) - else: - right_boundaries[lane_id].append(boundary_point_3d.array) - - left_boundaries = { - lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) - for lane_id, boundary_array in left_boundaries.items() - if len(boundary_array) > 1 - } - right_boundaries = { - lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) - for lane_id, boundary_array in right_boundaries.items() - if len(boundary_array) > 1 - } - - return left_boundaries, right_boundaries diff --git a/d123/dataset/dataset_specific/wopd/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/wopd_map_utils.py deleted file mode 100644 index 5c8b0239..00000000 --- a/d123/dataset/dataset_specific/wopd/wopd_map_utils.py +++ /dev/null @@ -1,442 +0,0 @@ -from pathlib import Path -from typing import Dict, Optional, Tuple - -import geopandas as gpd -import numpy as np -import numpy.typing as npt -import pandas as pd -import shapely.geometry as geom -from waymo_open_dataset import dataset_pb2 - -from d123.common.geometry.base import Point3D, Point3DIndex, StateSE2 -from d123.common.geometry.line.polylines import Polyline3D -from d123.common.geometry.occupancy_map import OccupancyMap2D -from d123.common.geometry.transform.tranform_2d import translate_along_yaw -from d123.common.geometry.vector import Vector2D -from d123.dataset.maps.map_datatypes import MapSurfaceType - - -def convert_wopd_map(frame: dataset_pb2.Frame, map_file_path: Path) -> None: - - def _get_polyline(data) -> npt.NDArray[np.float64]: - polyline = np.array([[p.x, p.y, p.z] for p in data.polyline], dtype=np.float64) - return polyline - - def _get_polygon(data) -> npt.NDArray[np.float64]: - polygon = np.array([[p.x, p.y, p.z] for p in data.polygon], dtype=np.float64) - assert polygon.shape[0] >= 3, "Polygon must have at least 3 points" - assert polygon.shape[1] == 3, "Polygon must have 3 coordinates (x, y, z)" - return polygon - - lanes: Dict[int, npt.NDArray[np.float64]] = {} - road_lines: Dict[int, npt.NDArray[np.float64]] = {} - road_edges: Dict[int, npt.NDArray[np.float64]] = {} - crosswalks: Dict[int, npt.NDArray[np.float64]] = {} - carparks: Dict[int, npt.NDArray[np.float64]] = {} - - for map_feature in frame.map_features: - if map_feature.HasField("lane"): - polyline = _get_polyline(map_feature.lane) - if polyline.ndim != 2 or polyline.shape[0] < 2: - continue - lanes[map_feature.id] = polyline - elif map_feature.HasField("road_line"): - polyline = _get_polyline(map_feature.road_line) - if polyline.ndim != 2 or polyline.shape[0] < 2: - continue - road_lines[map_feature.id] = polyline - elif map_feature.HasField("road_edge"): - polyline = _get_polyline(map_feature.road_edge) - if polyline.ndim != 2 or polyline.shape[0] < 2: - continue - road_edges[map_feature.id] = polyline - elif map_feature.HasField("stop_sign"): - # TODO: implement stop signs - pass - elif map_feature.HasField("crosswalk"): - crosswalks[map_feature.id] = _get_polygon(map_feature.crosswalk) - elif map_feature.HasField("speed_bump"): - # TODO: implement speed bumps - pass - elif map_feature.HasField("driveway"): - # NOTE: Determine whether to use a different semantic type for driveways. - carparks[map_feature.id] = _get_polygon(map_feature.driveway) - - lane_df = get_lane_df(lanes, road_lines, road_edges) - lane_group_df = get_lane_group_df(lanes) - intersection_df = get_intersections_df() - crosswalk_df = get_crosswalk_df(crosswalks) - walkway_df = get_walkway_df() - carpark_df = get_carpark_df(carparks) - generic_drivable_df = get_generic_drivable_df() - - map_file_path.unlink(missing_ok=True) - if not map_file_path.parent.exists(): - map_file_path.parent.mkdir(parents=True, exist_ok=True) - - lane_df.to_file(map_file_path, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") - lane_group_df.to_file(map_file_path, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") - intersection_df.to_file(map_file_path, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_path, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") - walkway_df.to_file(map_file_path, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_path, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_df.to_file( - map_file_path, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" - ) - - -def get_lane_df( - lanes: Dict[int, npt.NDArray[np.float64]], - road_lines: Dict[int, npt.NDArray[np.float64]], - road_edges: Dict[int, npt.NDArray[np.float64]], -) -> gpd.GeoDataFrame: - - ids = [] - lane_group_ids = [] - speed_limits_mps = [] - predecessor_ids = [] - successor_ids = [] - left_boundaries = [] - right_boundaries = [] - baseline_paths = [] - geometries = [] - - left_boundaries_3d, right_boundaries_3d = _extract_lane_boundaries(lanes, road_lines, road_edges) - for lane_id, lane_centerline_array in lanes.items(): - if lane_id not in left_boundaries_3d or lane_id not in right_boundaries_3d: - continue - lane_centerline = Polyline3D.from_array(lane_centerline_array) - - ids.append(lane_id) - lane_group_ids.append([lane_id]) - speed_limits_mps.append(None) # TODO: Implement speed limits - predecessor_ids.append([]) # TODO: Implement predecessor lanes - successor_ids.append([]) # TODO: Implement successor lanes - left_boundaries.append(left_boundaries_3d[lane_id].linestring) - right_boundaries.append(right_boundaries_3d[lane_id].linestring) - baseline_paths.append(lane_centerline.linestring) - - geometry = geom.Polygon( - np.vstack([left_boundaries_3d[lane_id].array[:, :2], right_boundaries_3d[lane_id].array[:, :2][::-1]]) - ) - geometries.append(geometry) - - data = pd.DataFrame( - { - "id": ids, - "lane_group_id": lane_group_ids, - "speed_limit_mps": speed_limits_mps, - "predecessor_ids": predecessor_ids, - "successor_ids": successor_ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - "baseline_path": baseline_paths, - } - ) - - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def get_lane_group_df(lanes: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: - - ids = [] - lane_ids = [] - intersection_ids = [] - predecessor_lane_group_ids = [] - successor_lane_group_ids = [] - left_boundaries = [] - right_boundaries = [] - geometries = [] - - for lane_id, lane_array in lanes.items(): - if lane_array.ndim != 2 or lane_array.shape[0] < 2: - continue - - centerline = Polyline3D.from_array(lane_array) - left_boundary, right_boundary = create_lane_boundaries(centerline) - lane_polygon = geom.Polygon(np.vstack([left_boundary.array[:, :2], right_boundary.array[:, :2][::-1]])) - - ids.append(lane_id) - lane_ids.append([lane_id]) - intersection_ids.append([]) - predecessor_lane_group_ids.append([]) - successor_lane_group_ids.append([]) - left_boundaries.append(left_boundary.linestring) - right_boundaries.append(right_boundary.linestring) - geometries.append(lane_polygon) - - data = pd.DataFrame( - { - "id": ids, - "lane_ids": lane_ids, - "intersection_id": intersection_ids, - "predecessor_lane_group_ids": predecessor_lane_group_ids, - "successor_lane_group_ids": successor_lane_group_ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - } - ) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def get_intersections_df() -> gpd.GeoDataFrame: - ids = [] - lane_group_ids = [] - geometries = [] - - data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def get_carpark_df(carparks) -> gpd.GeoDataFrame: - ids = list(carparks.keys()) - outlines = [geom.LineString(outline) for outline in carparks.values()] - geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in carparks.values()] - - data = pd.DataFrame({"id": ids, "outline": outlines}) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def get_walkway_df() -> gpd.GeoDataFrame: - ids = [] - geometries = [] - - data = pd.DataFrame({"id": ids}) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: - ids = list(crosswalks.keys()) - outlines = [geom.LineString(outline) for outline in crosswalks.values()] - geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in crosswalks.values()] - - data = pd.DataFrame({"id": ids, "outline": outlines}) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def get_generic_drivable_df() -> gpd.GeoDataFrame: - ids = [] - geometries = [] - - data = pd.DataFrame({"id": ids}) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf - - -def create_lane_boundaries(polyline_3d: Polyline3D, width: float = 4) -> Tuple[Polyline3D, Polyline3D]: - - points = polyline_3d.array - half_width = width / 2.0 - - # Calculate the direction vectors between consecutive points - directions = np.diff(points, axis=0) - - # Normalize the direction vectors - directions_norm = np.linalg.norm(directions, axis=1, keepdims=True) - directions_normalized = directions / directions_norm - - # Calculate perpendicular vectors in the xy plane (z remains 0) - perpendiculars = np.zeros_like(directions) - perpendiculars[:, 0] = -directions_normalized[:, 1] # -dy - perpendiculars[:, 1] = directions_normalized[:, 0] # dx - - # Create boundaries (need to handle the last point separately) - left_boundary = points[:-1] + perpendiculars * half_width - right_boundary = points[:-1] - perpendiculars * half_width - - # Handle the last point based on the last direction - last_perp = perpendiculars[-1] - left_boundary = np.vstack([left_boundary, points[-1] + last_perp * half_width]) - right_boundary = np.vstack([right_boundary, points[-1] - last_perp * half_width]) - - return Polyline3D.from_array(left_boundary), Polyline3D.from_array(right_boundary) - - -def _extract_lane_boundaries( - lanes: Dict[int, npt.NDArray[np.float64]], - road_lines: Dict[int, npt.NDArray[np.float64]], - road_edges: Dict[int, npt.NDArray[np.float64]], -) -> Tuple[Dict[str, Polyline3D], Dict[str, Polyline3D]]: - polyline_dict: Dict[str, Dict[int, Polyline3D]] = {"lane": {}, "roadline": {}, "roadedge": {}} - - for lane_id, lane_polyline in lanes.items(): - if lane_polyline.ndim == 2 and lane_polyline.shape[1] == 3 and len(lane_polyline) > 0: - polyline_dict["lane"][lane_id] = Polyline3D.from_array(lane_polyline) - - # for road_line_id, road_line_polyline in road_lines.items(): - # if road_line_polyline.ndim == 2 and road_line_polyline.shape[1] == 3 and len(road_line_polyline) > 0: - # polyline_dict["roadline"][road_line_id] = Polyline3D.from_array(road_line_polyline) - - for road_edge_id, road_edge_polyline in road_edges.items(): - if road_edge_polyline.ndim == 2 and road_edge_polyline.shape[1] == 3 and len(road_edge_polyline) > 0: - polyline_dict["roadedge"][road_edge_id] = Polyline3D.from_array(road_edge_polyline) - - geometries = [] - tokens = [] - for line_type, polylines in polyline_dict.items(): - for polyline_id, polyline in polylines.items(): - geometries.append(polyline.polyline_2d.linestring) - tokens.append(f"{line_type}_{polyline_id}") - - occupancy_2d = OccupancyMap2D(geometries, tokens) - - # for each lane - # - sample poses along centerline - # - query left and right perpendicular lines in 2d occupancy map - # - for {left, right} retrieve first perpendicular intersection with length MAX_HALF_WIDTH - # - if intersection with road edge, add intersecting point to {left, right} boundary - # - if intersection with centerline, add middle point to {left, right} boundary - # - if no intersection found, add MAX_HALF_WIDTH to {left, right} boundary - - MAX_LANE_WIDTH = 25.0 # meters - MIN_LANE_WIDTH = 2.0 - DEFAULT_LANE_WIDTH = 4.0 - BOUNDARY_STEP_SIZE = 0.5 # meters - MAX_Z_DISTANCE = 1.0 # meters - - left_boundaries = {lane_id: [] for lane_id in lanes.keys()} - right_boundaries = {lane_id: [] for lane_id in lanes.keys()} - - def get_type_and_id_from_token(token: str) -> Tuple[str, int]: - """Extract type and id from token.""" - line_type, line_id = token.split("_") - return line_type, int(line_id) - - def get_polyline_from_token(token: str) -> Polyline3D: - """Extract polyline from token.""" - line_type, line_id = get_type_and_id_from_token(token) - return polyline_dict[line_type][line_id] - - def get_intersecting_point_3d(perp_linestring: geom.LineString, intersecting_token: str) -> Optional[Point3D]: - - intersecting_polyline_3d = get_polyline_from_token(intersecting_token) - intersecting_linestring = occupancy_2d.geometries[occupancy_2d.token_to_idx[intersecting_token]] - intersecting_point_2d = perp_linestring.intersection(intersecting_linestring) - - distance_2d_norm = None - if isinstance(intersecting_point_2d, geom.Point): - distance_2d_norm = intersecting_linestring.project(intersecting_point_2d, normalized=True) - elif isinstance(intersecting_point_2d, geom.MultiPoint): - geom_points = [point for point in intersecting_point_2d.geoms] - starting_perp_point = geom.Point(perp_linestring.coords[0]) - start_perp_point_distance = np.array( - [starting_perp_point.distance(geom_point) for geom_point in geom_points] - ) - start_perp_point_distance[start_perp_point_distance < MIN_LANE_WIDTH / 2.0] = np.inf - closest_index = np.argmin(start_perp_point_distance) - distance_2d_norm = intersecting_linestring.project(geom_points[closest_index], normalized=True) - - intersecting_point_3d = None - if distance_2d_norm is not None: - intersecting_point_3d = intersecting_polyline_3d.interpolate( - distance_2d_norm * intersecting_polyline_3d.length - ) - if np.linalg.norm(intersecting_point_3d.point_2d.array - perp_linestring.coords[0]) < MIN_LANE_WIDTH / 2.0: - return None - - return intersecting_point_3d - - for lane_id, lane_polyline in polyline_dict["lane"].items(): - current_lane_token = f"lane_{lane_id}" - lane_polyline_se2 = lane_polyline.polyline_se2 - - # 0. Find connected or intersecting lanes - connected_lane_tokens = [ - token - for token in occupancy_2d.intersects(lane_polyline_se2.linestring) - if token.startswith("lane_") - if token != current_lane_token - ] - - # 1. sample poses along centerline - distances_se2 = np.linspace( - 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True - ) - lane_query_se2 = [ - StateSE2.from_array(state_se2_array) for state_se2_array in lane_polyline_se2.interpolate(distances_se2) - ] - distances_3d = np.linspace( - 0, lane_polyline.length, int(lane_polyline.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True - ) - lane_query_3d = [ - Point3D.from_array(point_3d_array) for point_3d_array in lane_polyline.interpolate(distances_3d) - ] - assert len(lane_query_se2) == len(lane_query_3d) - - for lane_query_se2, lane_query_3d in zip(lane_query_se2, lane_query_3d): - for sign in [1.0, -1.0]: - perp_start_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * 0.1)) - perp_end_point = translate_along_yaw(lane_query_se2, Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)) - perp_linestring = geom.LineString( - [[perp_start_point.x, perp_start_point.y], [perp_end_point.x, perp_end_point.y]] - ) - - # 1. find intersecting lines, compute 3D distance - intersecting_tokens = occupancy_2d.intersects(perp_linestring) - intersecting_points_3d: Dict[str, Point3D] = {} - crosses_centerline = False - - for intersecting_token in intersecting_tokens: - # if intersecting_token in connected_lane_tokens: - # crosses_centerline = True - # continue - # if intersecting_token in connected_lane_tokens: - # continue - - # if intersecting_token == current_lane_token: - # crosses_centerline = True - - intersecting_point_3d = get_intersecting_point_3d(perp_linestring, intersecting_token) - if intersecting_point_3d is None: - continue - - if np.abs(intersecting_point_3d.z - lane_query_3d.z) > MAX_Z_DISTANCE: - continue - - intersecting_points_3d[intersecting_token] = intersecting_point_3d - - boundary_point_3d: Optional[Point3D] = None - if len(intersecting_points_3d) > 0 and not crosses_centerline: - # 2. find closest intersection - intersecting_distances = { - token: np.linalg.norm(point_3d.array - lane_query_3d.array) - for token, point_3d in intersecting_points_3d.items() - } - closest_token = min(intersecting_distances, key=intersecting_distances.get) - closest_type, closest_id = get_type_and_id_from_token(closest_token) - - if closest_type == "lane": - boundary_point_3d = Point3D.from_array( - (intersecting_points_3d[closest_token].array + lane_query_3d.array) / 2.0 - ) - elif closest_type == "roadedge": - boundary_point_3d = intersecting_points_3d[closest_token] - else: - perp_boundary_distance = DEFAULT_LANE_WIDTH / 2.0 # Default to half the lane width - boundary_point_se2 = translate_along_yaw( - lane_query_se2, Vector2D(0.0, sign * perp_boundary_distance) - ) - boundary_point_3d = Point3D(boundary_point_se2.x, boundary_point_se2.y, lane_query_3d.z) - if sign == 1.0: - left_boundaries[lane_id].append(boundary_point_3d.array) - else: - right_boundaries[lane_id].append(boundary_point_3d.array) - - left_boundaries = { - lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) - for lane_id, boundary_array in left_boundaries.items() - if len(boundary_array) > 1 - } - right_boundaries = { - lane_id: Polyline3D.from_array(np.array(boundary_array, dtype=np.float64)) - for lane_id, boundary_array in right_boundaries.items() - if len(boundary_array) > 1 - } - - return left_boundaries, right_boundaries diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index fc24de85..21614ea0 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -104,10 +104,10 @@ "\n", "LANE_CONFIG: PlotConfig = PlotConfig(\n", " fill_color=LIGHT_GREY,\n", - " fill_color_alpha=0.5,\n", + " fill_color_alpha=1.0,\n", " line_color=BLACK,\n", - " line_color_alpha=0.2,\n", - " line_width=1.0,\n", + " line_color_alpha=0.0,\n", + " line_width=0.0,\n", " line_style=\"-\",\n", " zorder=1,\n", ")\n", @@ -147,20 +147,20 @@ " if layer in [MapSurfaceType.LANE]:\n", " map_object: AbstractLane\n", " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", - " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", - " # add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", + " add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", + " add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", "\n", - " # centroid = map_object.shapely_polygon.centroid\n", - " # ax.text(\n", - " # centroid.x,\n", - " # centroid.y,\n", - " # str(map_object.id),\n", - " # horizontalalignment=\"center\",\n", - " # verticalalignment=\"center\",\n", - " # fontsize=8,\n", - " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", - " # )\n", + " centroid = map_object.shapely_polygon.centroid\n", + " ax.text(\n", + " centroid.x,\n", + " centroid.y,\n", + " str(map_object.id),\n", + " horizontalalignment=\"center\",\n", + " verticalalignment=\"center\",\n", + " fontsize=8,\n", + " bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", + " )\n", "\n", " except Exception:\n", " import traceback\n", @@ -199,8 +199,8 @@ " return fig, ax\n", "\n", "\n", - "scene_index = 0\n", - "plot_scene_at_iteration(scenes[scene_index], iteration=20, radius=200)" + "scene_index = 1\n", + "plot_scene_at_iteration(scenes[scene_index], iteration=70, radius=250)" ] }, { diff --git a/notebooks/viz/viser_testing_v2_map.ipynb b/notebooks/viz/viser_testing_v2_map.ipynb new file mode 100644 index 00000000..38eaaaee --- /dev/null +++ b/notebooks/viz/viser_testing_v2_map.ipynb @@ -0,0 +1,116 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + " \n", + "\n", + "splits = [\"wopd_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "\n", + "# splits = [\"nuplan_private_test\"]\n", + "# splits = [\"carla\"]\n", + "# log_names = [\"11839652018869852123_2565_000_2585_000\"]\n", + "log_names = None\n", + "\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=19,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=20,\n", + " shuffle=False,\n", + " camera_types=[CameraType.CAM_F0],\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "from d123.common.visualization.viser.server_map import ViserMapVisualizationServer\n", + "\n", + "\n", + "visualization_server = ViserMapVisualizationServer(scenes, scene_index=4)" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/waymo_perception/map_testing.ipynb b/notebooks/waymo_perception/map_testing.ipynb index f4da563b..b3388863 100644 --- a/notebooks/waymo_perception/map_testing.ipynb +++ b/notebooks/waymo_perception/map_testing.ipynb @@ -21,7 +21,11 @@ "from waymo_open_dataset.utils import box_utils\n", "\n", "\n", - "import matplotlib.pyplot as plt" + "import matplotlib.pyplot as plt\n", + "\n", + "\n", + "import os\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"" ] }, { @@ -75,7 +79,7 @@ "# Length: 0\n", "# map_pose_offset: \n", "\n", - "file_idx = 3\n", + "file_idx = 0\n", "pathname = tfrecords_file_list[file_idx]\n", "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", "num_frames = sum(1 for _ in dataset)\n", @@ -187,7 +191,9 @@ " outline = np.array([[p.x, p.y, p.z] for p in map_feature.driveway.polygon])\n", " driveways[map_feature.id] = outline\n", "\n", - " # print(f\"Roadline: {map_feature.road_line}\")" + " # print(f\"Roadline: {map_feature.road_line}\")\n", + "\n", + "lanes_speed_limit_mps" ] }, { @@ -197,11 +203,53 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.dataset.dataset_specific.wopd.wopd_map_utils import _extract_lane_boundaries\n", + "from d123.common.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", + "\n", + "\n", + "fig, ax = plt.subplots(figsize=(30, 30))\n", + "\n", + "for road_edge in road_edges.values():\n", + " # print(len(driveway))\n", + " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"blue\", label=\"road_edge\")\n", + "\n", + "for lane in lanes.values():\n", + " # print(len(driveway))\n", + " ax.plot(lane[:, 0], lane[:, 1], color=\"gray\", label=\"lane\")\n", + "\n", + "\n", + "for road_line in road_lines.values():\n", + " # print(len(driveway))\n", + " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "for driveway in driveways.values():\n", + " # print(len(driveway))\n", + " ax.plot(driveway[:, 0], driveway[:, 1], color=\"green\", label=\"driveway\")\n", + "\n", + "\n", + "for crosswalk in crosswalks.values():\n", + " # print(len(driveway))\n", + " ax.plot(crosswalk[:, 0], crosswalk[:, 1], color=\"violet\", label=\"crosswalk\")\n", "\n", "\n", - "left_boundaries, right_boundaries = _extract_lane_boundaries(\n", + "add_non_repeating_legend_to_ax(ax)\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.dataset_specific.wopd.wopd_map_utils import extract_lane_boundaries\n", + "\n", + "\n", + "left_boundaries, right_boundaries = extract_lane_boundaries(\n", " lanes=lanes,\n", + " lanes_successors=lanes_successors,\n", + " lanes_predecessors=lanes_predecessors,\n", " road_lines=road_lines,\n", " road_edges=road_edges,\n", ")" @@ -210,7 +258,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", "metadata": {}, "outputs": [], "source": [ @@ -237,218 +285,102 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [], "source": [ - "from d123.common.geometry.base import StateSE2\n", - "from d123.common.geometry.line.polylines import PolylineSE2\n", - "from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", - "from d123.common.geometry.vector import Vector2D\n", - "\n", - "size = 30\n", - "fig, ax = plt.subplots(figsize=(size, size))\n", - "\n", - "lane_id = 114\n", - "lane_ = lanes[lane_id]\n", - "\n", - "lane_polyline_se2 = PolylineSE2.from_array(lane_[:, :2])\n", - "BOUNDARY_STEP_SIZE = 0.5\n", - "\n", - "\n", - "distances_se2 = np.linspace(\n", - " 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True\n", - ")\n", - "lane_query_se2 = lane_polyline_se2.interpolate(distances_se2)\n", - "\n", - "\n", - "left_boundary = left_boundaries[lane_id].array\n", - "right_boundary = right_boundaries[lane_id].array\n", - "\n", - "assert len(left_boundary) > 0 and len(right_boundary) > 0\n", - "ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", - "ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", - "\n", - "ax.scatter(lane_query_se2[:, 0], lane_query_se2[:, 1], c=lane_query_se2[:, 2] / np.pi, cmap=\"viridis\")\n", - "\n", - "\n", - "MAX_LANE_WIDTH = 25\n", - "for state_se2_array in lane_query_se2:\n", - " for sign in [1.0, -1.0]:\n", - " perp_start_point = translate_along_yaw(StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * 0.1))\n", - " perp_end_point = translate_along_yaw(\n", - " StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)\n", - " )\n", - " ax.plot(\n", - " [perp_start_point.x, perp_end_point.x],\n", - " [perp_start_point.y, perp_end_point.y],\n", - " color=\"lime\" if sign > 0 else \"red\",\n", - " linestyle=\"--\",\n", - " alpha=0.5,\n", - " )\n", - "\n", - "\n", - "for road_edge in road_edges.values():\n", - " # print(len(driveway))\n", - " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"black\", label=\"road_edge\", linestyle=\"dashdot\")\n", - "\n", - "\n", - "\n", - "for road_line in road_lines.values():\n", - " # print(len(driveway))\n", - " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", - "\n", - "ax.set_aspect(\"equal\")\n", - "\n", - "ax.set_xlim(lane_[:, 0].min() - 10, lane_[:, 0].max() + 10)\n", - "ax.set_ylim(lane_[:, 1].min() - 10, lane_[:, 1].max() + 10)\n", - "\n", - "lanes_map_features[lane_id]" + "# " ] }, { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "9", "metadata": {}, "outputs": [], "source": [ - "from d123.common.geometry.base import StateSE2\n", - "from d123.common.geometry.line.polylines import PolylineSE2\n", - "from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", - "from d123.common.geometry.vector import Vector2D\n", - "from d123.common.visualization.color.color import TAB_10\n", - "\n", - "size = 20\n", - "fig, ax = plt.subplots(figsize=(size, size))\n", - "\n", - "lane_id = 114\n", - "lane_ = lanes[lane_id]\n", - "\n", - "lane_polyline_se2 = PolylineSE2.from_array(lane_[:, :2])\n", - "BOUNDARY_STEP_SIZE = 0.5\n", - "\n", - "\n", - "distances_se2 = np.linspace(\n", - " 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True\n", - ")\n", - "lane_query_se2 = lane_polyline_se2.interpolate(distances_se2)\n", - "\n", - "\n", - "left_boundary = left_boundaries[lane_id].array\n", - "right_boundary = right_boundaries[lane_id].array\n", - "\n", - "assert len(left_boundary) > 0 and len(right_boundary) > 0\n", - "# ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", - "# ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "lane_idx = 158\n", "\n", - "# ax.scatter(lane_query_se2[:, 0], lane_query_se2[:, 1], c=lane_query_se2[:, 2] / np.pi, cmap=\"viridis\")\n", + "lane = lanes_map_features[lane_idx].lane\n", + "lane_id = lanes_map_features[lane_idx].id\n", "\n", "\n", - "# MAX_LANE_WIDTH = 15\n", - "# for state_se2_array in lane_query_se2:\n", - "# for sign in [1.0, -1.0]:\n", - "# perp_end_point = translate_along_yaw(\n", - "# StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)\n", - "# )\n", - "# ax.plot(\n", - "# [state_se2_array[0], perp_end_point.x],\n", - "# [state_se2_array[1], perp_end_point.y],\n", - "# color=\"lime\" if sign > 0 else \"red\",\n", - "# linestyle=\"--\",\n", - "# alpha=0.5,\n", - "# )\n", + "import matplotlib.pyplot as plt\n", "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", "\n", - "for idx, right_boundary in enumerate(lanes_map_features[lane_id].lane.right_boundaries):\n", - " lane_start_index = right_boundary.lane_start_index\n", - " lane_end_index = right_boundary.lane_end_index\n", + "ax.plot(lanes[lane_id][:, 0], lanes[lane_id][:, 1], color=\"black\", label=\"Lane Centerline\", alpha=0.5)\n", "\n", - " road_edge = road_edges[right_boundary.boundary_feature_id]\n", - " if len(lane_[lane_start_index:lane_end_index, 0]) < 10:\n", - " continue\n", + "for left_boundary in lane.left_boundaries:\n", + " print(\"left\", left_boundary)\n", + " try:\n", + " boundary = road_lines[left_boundary.boundary_feature_id]\n", "\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {left_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[left_boundary.boundary_feature_id]\n", " ax.plot(\n", - " lane_[lane_start_index:lane_end_index, 0],\n", - " lane_[lane_start_index:lane_end_index, 1],\n", - " color=TAB_10[idx % len(TAB_10)].hex,\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"lime\",\n", + " linestyle=\"--\",\n", + " label=\"Left Neighbor\",\n", " )\n", " ax.plot(\n", - " road_edge[:, 0],\n", - " road_edge[:, 1],\n", - " color=TAB_10[idx % len(TAB_10)].hex,\n", - " label=\"road_edge\",\n", - " linestyle=\"dashdot\",\n", - " alpha=0.5,\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 0],\n", + " lanes[lane_id][left_boundary.lane_start_index : left_boundary.lane_end_index, 1],\n", + " color=\"lime\",\n", + " linestyle=\"-\",\n", + " label=\"Left Neighbor\",\n", " )\n", - " print(len(road_edge))\n", "\n", "\n", - "# # print(len(driveway))\n", - "# ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\", label=\"right_boundary\", linestyle=\"--\")\n", - "\n", + "for right_boundary in lane.right_boundaries:\n", + " print(\"right\", right_boundary)\n", + " try:\n", + " boundary = road_lines[right_boundary.boundary_feature_id]\n", "\n", - "# for road_edge in road_edges.values():\n", - "# # print(len(driveway))\n", - "# ax.plot(road_edge[:, 0], road_edge[:, 1], label=\"road_edge\", linestyle=\"dashdot\")\n", + " except KeyError:\n", + " print(f\"Boundary feature ID {right_boundary.boundary_feature_id} not found.\")\n", + " boundary = road_edges[right_boundary.boundary_feature_id]\n", "\n", + " ax.plot(\n", + " boundary[:, 0],\n", + " boundary[:, 1],\n", + " color=\"red\",\n", + " linestyle=\"--\",\n", + " label=\"Right Neighbor\",\n", + " )\n", + " ax.plot(\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 0],\n", + " lanes[lane_id][right_boundary.lane_start_index : right_boundary.lane_end_index, 1],\n", + " color=\"red\",\n", + " linestyle=\"-\",\n", + " label=\"Right Neighbor\",\n", + " )\n", "\n", - "# for road_line in road_lines.values():\n", - "# # print(len(driveway))\n", - "# ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", "\n", "ax.set_aspect(\"equal\")\n", + "# lanes\n", "\n", - "zoom_out = 100\n", - "ax.set_xlim(lane_[:, 0].min() - zoom_out, lane_[:, 0].max() + zoom_out)\n", - "ax.set_ylim(lane_[:, 1].min() - zoom_out, lane_[:, 1].max() + zoom_out)\n", - "lanes_map_features[lane_id]" + "zoom_out = 10\n", + "lane = lanes[lane_id]\n", + "ax.set_xlim(lane[:, 0].min() - zoom_out, lane[:, 0].max() + zoom_out)\n", + "ax.set_ylim(lane[:, 1].min() - zoom_out, lane[:, 1].max() + zoom_out)" ] }, { "cell_type": "code", "execution_count": null, - "id": "9", + "id": "10", "metadata": {}, "outputs": [], - "source": [ - "from d123.common.visualization.matplotlib.utils import add_non_repeating_legend_to_ax\n", - "\n", - "\n", - "fig, ax = plt.subplots(figsize=(30, 30))\n", - "\n", - "for road_edge in road_edges.values():\n", - " # print(len(driveway))\n", - " ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"blue\", label=\"road_edge\")\n", - "\n", - "for lane in lanes.values():\n", - " # print(len(driveway))\n", - " ax.plot(lane[:, 0], lane[:, 1], color=\"gray\", label=\"lane\")\n", - "\n", - "\n", - "for road_line in road_lines.values():\n", - " # print(len(driveway))\n", - " ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", - "\n", - "for driveway in driveways.values():\n", - " # print(len(driveway))\n", - " ax.plot(driveway[:, 0], driveway[:, 1], color=\"green\", label=\"driveway\")\n", - "\n", - "\n", - "for crosswalk in crosswalks.values():\n", - " # print(len(driveway))\n", - " ax.plot(crosswalk[:, 0], crosswalk[:, 1], color=\"violet\", label=\"crosswalk\")\n", - "\n", - "\n", - "add_non_repeating_legend_to_ax(ax)\n", - "\n", - "ax.set_aspect(\"equal\")" - ] + "source": [] }, { "cell_type": "code", "execution_count": null, - "id": "10", + "id": "11", "metadata": {}, "outputs": [], "source": [ @@ -518,7 +450,7 @@ { "cell_type": "code", "execution_count": null, - "id": "11", + "id": "12", "metadata": {}, "outputs": [], "source": [ @@ -581,7 +513,7 @@ { "cell_type": "code", "execution_count": null, - "id": "12", + "id": "13", "metadata": {}, "outputs": [], "source": [ @@ -629,11 +561,11 @@ { "cell_type": "code", "execution_count": null, - "id": "13", + "id": "14", "metadata": {}, "outputs": [], "source": [ - "lane_idx = 58\n", + "lane_idx = 196\n", "\n", "lane = lanes_map_features[lane_idx].lane\n", "lane_id = lanes_map_features[lane_idx].id\n", @@ -697,6 +629,78 @@ "ax.set_aspect(\"equal\")\n", "# lanes" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.geometry.base import StateSE2\n", + "# from d123.common.geometry.line.polylines import PolylineSE2\n", + "# from d123.common.geometry.transform.tranform_2d import translate_along_yaw\n", + "# from d123.common.geometry.vector import Vector2D\n", + "\n", + "# size = 30\n", + "# fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "# lane_id = 114\n", + "# lane_ = lanes[lane_id]\n", + "\n", + "# lane_polyline_se2 = PolylineSE2.from_array(lane_[:, :2])\n", + "# BOUNDARY_STEP_SIZE = 0.5\n", + "\n", + "\n", + "# distances_se2 = np.linspace(\n", + "# 0, lane_polyline_se2.length, int(lane_polyline_se2.length / BOUNDARY_STEP_SIZE) + 1, endpoint=True\n", + "# )\n", + "# lane_query_se2 = lane_polyline_se2.interpolate(distances_se2)\n", + "\n", + "\n", + "# left_boundary = left_boundaries[lane_id].array\n", + "# right_boundary = right_boundaries[lane_id].array\n", + "\n", + "# assert len(left_boundary) > 0 and len(right_boundary) > 0\n", + "# ax.plot(left_boundary[:, 0], left_boundary[:, 1], color=\"lime\")\n", + "# ax.plot(right_boundary[:, 0], right_boundary[:, 1], color=\"red\")\n", + "\n", + "# ax.scatter(lane_query_se2[:, 0], lane_query_se2[:, 1], c=lane_query_se2[:, 2] / np.pi, cmap=\"viridis\")\n", + "\n", + "\n", + "# MAX_LANE_WIDTH = 25\n", + "# for state_se2_array in lane_query_se2:\n", + "# for sign in [1.0, -1.0]:\n", + "# perp_start_point = translate_along_yaw(StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * 0.1))\n", + "# perp_end_point = translate_along_yaw(\n", + "# StateSE2.from_array(state_se2_array), Vector2D(0.0, sign * MAX_LANE_WIDTH / 2.0)\n", + "# )\n", + "# ax.plot(\n", + "# [perp_start_point.x, perp_end_point.x],\n", + "# [perp_start_point.y, perp_end_point.y],\n", + "# color=\"lime\" if sign > 0 else \"red\",\n", + "# linestyle=\"--\",\n", + "# alpha=0.5,\n", + "# )\n", + "\n", + "\n", + "# for road_edge in road_edges.values():\n", + "# # print(len(driveway))\n", + "# ax.plot(road_edge[:, 0], road_edge[:, 1], color=\"black\", label=\"road_edge\", linestyle=\"dashdot\")\n", + "\n", + "\n", + "\n", + "# for road_line in road_lines.values():\n", + "# # print(len(driveway))\n", + "# ax.plot(road_line[:, 0], road_line[:, 1], color=\"orange\", label=\"road_line\")\n", + "\n", + "# ax.set_aspect(\"equal\")\n", + "\n", + "# ax.set_xlim(lane_[:, 0].min() - 10, lane_[:, 0].max() + 10)\n", + "# ax.set_ylim(lane_[:, 1].min() - 10, lane_[:, 1].max() + 10)\n", + "\n", + "# lanes_map_features[lane_id]" + ] } ], "metadata": { From 5f8c6bc1af3ae11e23a1b50479f613c6c0104298 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Thu, 7 Aug 2025 17:54:36 +0200 Subject: [PATCH 22/42] Fix import issue --- d123/dataset/dataset_specific/wopd/wopd_data_converter.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 9db471e3..713a44c9 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -27,7 +27,7 @@ from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter -from d123.dataset.dataset_specific.wopd.wopd_map_utils import convert_wopd_map +from d123.dataset.dataset_specific.wopd.waymo_map_utils.wopd_map_utils import convert_wopd_map from d123.dataset.dataset_specific.wopd.wopd_utils import parse_range_image_and_camera_projection from d123.dataset.logs.log_metadata import LogMetadata @@ -35,7 +35,6 @@ D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) TARGET_DT: Final[float] = 0.1 -NUPLAN_DT: Final[float] = 0.05 SORT_BY_TIMESTAMP: Final[bool] = True NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { From e24b637db367713cedc9fcf847ecfc6922793041 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Thu, 7 Aug 2025 21:14:13 +0200 Subject: [PATCH 23/42] Change some adjustments in the wopd processing and visualization --- .../vehicle_state/vehicle_parameters.py | 2 +- d123/common/visualization/color/color.py | 17 ++ d123/common/visualization/viser/server.py | 23 ++- d123/common/visualization/viser/server_map.py | 167 ------------------ d123/common/visualization/viser/utils.py | 12 +- d123/common/visualization/viser/utils_v2.py | 9 +- .../waymo_map_utils/womp_boundary_utils.py | 2 +- .../wopd/wopd_data_converter.py | 150 ++++++++-------- d123/dataset/maps/gpkg/gpkg_map.py | 2 + d123/dataset/maps/map_datatypes.py | 2 + .../default_dataset_conversion.yaml | 4 +- notebooks/viz/bev_matplotlib.ipynb | 41 +++-- notebooks/viz/log_rendering.ipynb | 12 +- notebooks/viz/viser_testing_v2_scene.ipynb | 166 ++++++++++++++++- 14 files changed, 338 insertions(+), 271 deletions(-) delete mode 100644 d123/common/visualization/viser/server_map.py diff --git a/d123/common/datatypes/vehicle_state/vehicle_parameters.py b/d123/common/datatypes/vehicle_state/vehicle_parameters.py index 714f3628..8fe4d048 100644 --- a/d123/common/datatypes/vehicle_state/vehicle_parameters.py +++ b/d123/common/datatypes/vehicle_state/vehicle_parameters.py @@ -56,7 +56,7 @@ def get_wopd_pacifica_parameters() -> VehicleParameters: height=1.777, wheel_base=3.089, # rear_axle_to_center_vertical=0.45, - rear_axle_to_center_vertical=0.60, + rear_axle_to_center_vertical=1.777 / 2, rear_axle_to_center_longitudinal=1.461, ) diff --git a/d123/common/visualization/color/color.py b/d123/common/visualization/color/color.py index b151037a..5ba3c866 100644 --- a/d123/common/visualization/color/color.py +++ b/d123/common/visualization/color/color.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from dataclasses import dataclass from typing import Dict, Tuple @@ -9,6 +11,11 @@ class Color: hex: str + @classmethod + def from_rgb(cls, rgb: Tuple[int, int, int]) -> Color: + r, g, b = rgb + return cls(f"#{r:02x}{g:02x}{b:02x}") + @property def rgb(self) -> Tuple[int, int, int]: return ImageColor.getcolor(self.hex, "RGB") @@ -25,6 +32,16 @@ def rgb_norm(self) -> Tuple[float, float, float]: def rgba_norm(self) -> Tuple[float, float, float]: return tuple([c / 255 for c in self.rgba]) + def set_brightness(self, factor: float) -> Color: + r, g, b = self.rgb + return Color.from_rgb( + ( + max(min(int(r * factor), 255), 0), + max(min(int(g * factor), 255), 0), + max(min(int(b * factor), 255), 0), + ) + ) + def __str__(self) -> str: return self.hex diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index ca37a13e..984aa968 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -30,7 +30,7 @@ LIDAR_POINT_SIZE: float = 0.05 MAP_AVAILABLE: bool = True -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" LINE_WIDTH: float = 4.0 CAMERA_SCALE: float = 1.0 @@ -229,6 +229,27 @@ def _(_) -> None: for name, mesh in get_map_meshes(scene).items(): self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) + # centerlines, __, __ = get_map_lines(scene) + # for i, centerline in enumerate(centerlines): + # self.server.scene.add_line_segments( + # "/map/centerlines", + # centerlines, + # colors=[[BLACK.rgb]], + # line_width=LINE_WIDTH, + # ) + # self.server.scene.add_line_segments( + # "/map/left_boundary", + # left_boundaries, + # colors=[[TAB_10[2].rgb]], + # line_width=LINE_WIDTH, + # ) + # self.server.scene.add_line_segments( + # "/map/right_boundary", + # right_boundaries, + # colors=[[TAB_10[3].rgb]], + # line_width=LINE_WIDTH, + # ) + # Playback update loop. prev_timestep = gui_timestep.value while server_playing: diff --git a/d123/common/visualization/viser/server_map.py b/d123/common/visualization/viser/server_map.py deleted file mode 100644 index 551fa2ba..00000000 --- a/d123/common/visualization/viser/server_map.py +++ /dev/null @@ -1,167 +0,0 @@ -import time -from typing import List, Literal - -import viser - -from d123.common.datatypes.sensor.camera import CameraType -from d123.common.visualization.color.color import BLACK, TAB_10 -from d123.common.visualization.viser.utils import get_map_lines, get_map_meshes -from d123.dataset.scene.abstract_scene import AbstractScene - -# TODO: Try to fix performance issues. -# TODO: Refactor this file. - -all_camera_types: List[CameraType] = [ - CameraType.CAM_F0, - CameraType.CAM_B0, - CameraType.CAM_L0, - CameraType.CAM_L1, - CameraType.CAM_L2, - CameraType.CAM_R0, - CameraType.CAM_R1, - CameraType.CAM_R2, -] - -LIDAR_POINT_SIZE: float = 0.05 -MAP_AVAILABLE: bool = True -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" -LINE_WIDTH: float = 4.0 - -CAMERA_SCALE: float = 1.0 - -# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ -# CameraType.CAM_F0, -# CameraType.CAM_L0, -# CameraType.CAM_R0, -# CameraType.CAM_L1, -# CameraType.CAM_R1, -# ] -# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] -# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types -VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] -VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] -LIDAR_AVAILABLE: bool = False - - -class ViserMapVisualizationServer: - def __init__( - self, - scenes: List[AbstractScene], - scene_index: int = 0, - host: str = "localhost", - port: int = 8080, - label: str = "D123 Viser Server", - ): - assert len(scenes) > 0, "At least one scene must be provided." - self.scenes = scenes - self.scene_index = scene_index - - self.host = host - self.port = port - self.label = label - - self.server = viser.ViserServer(host=self.host, port=self.port, label=self.label) - self.set_scene(self.scenes[self.scene_index % len(self.scenes)]) - - def next(self) -> None: - self.server.flush() - self.server.gui.reset() - self.server.scene.reset() - self.scene_index = (self.scene_index + 1) % len(self.scenes) - print(f"Viser server started at {self.host}:{self.port}") - self.set_scene(self.scenes[self.scene_index]) - - def set_scene(self, scene: AbstractScene) -> None: - num_frames = scene.get_number_of_iterations() - print(scene.available_camera_types) - - self.server.gui.configure_theme(control_width="large") - with self.server.gui.add_folder("Playback"): - server_playing = True - - gui_timestep = self.server.gui.add_slider( - "Timestep", - min=0, - max=num_frames - 1, - step=1, - initial_value=0, - disabled=True, - ) - gui_next_frame = self.server.gui.add_button("Next Frame", disabled=True) - gui_prev_frame = self.server.gui.add_button("Prev Frame", disabled=True) - gui_next_scene = self.server.gui.add_button("Next Scene", disabled=False) - gui_playing = self.server.gui.add_checkbox("Playing", True) - gui_framerate = self.server.gui.add_slider("FPS", min=1, max=60, step=0.1, initial_value=10) - gui_framerate_options = self.server.gui.add_button_group("FPS options", ("10", "20", "30", "60")) - - # Frame step buttons. - @gui_next_frame.on_click - def _(_) -> None: - gui_timestep.value = (gui_timestep.value + 1) % num_frames - - @gui_prev_frame.on_click - def _(_) -> None: - gui_timestep.value = (gui_timestep.value - 1) % num_frames - - @gui_next_scene.on_click - def _(_) -> None: - nonlocal server_playing - server_playing = False - - # Disable frame controls when we're playing. - @gui_playing.on_update - def _(_) -> None: - gui_timestep.disabled = gui_playing.value - gui_next_frame.disabled = gui_playing.value - gui_prev_frame.disabled = gui_playing.value - - # Set the framerate when we click one of the options. - @gui_framerate_options.on_click - def _(_) -> None: - gui_framerate.value = int(gui_framerate_options.value) - - # Frame step buttons - - # Toggle frame visibility when the timestep slider changes. - @gui_timestep.on_update - def _(_) -> None: - - time.sleep(0.1) - - prev_timestep = gui_timestep.value - - # Load in frames. - if MAP_AVAILABLE: - for name, mesh in get_map_meshes(scene).items(): - self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) - - centerlines, left_boundaries, right_boundaries = get_map_lines(scene) - # for i, centerline in enumerate(centerlines): - self.server.scene.add_line_segments( - "/map/centerlines", - centerlines, - colors=[[BLACK.rgb]], - line_width=LINE_WIDTH, - ) - self.server.scene.add_line_segments( - "/map/left_boundary", - left_boundaries, - colors=[[TAB_10[2].rgb]], - line_width=LINE_WIDTH, - ) - self.server.scene.add_line_segments( - "/map/right_boundary", - right_boundaries, - colors=[[TAB_10[3].rgb]], - line_width=LINE_WIDTH, - ) - - # Playback update loop. - prev_timestep = gui_timestep.value - while server_playing: - # Update the timestep if we're playing. - if gui_playing.value: - gui_timestep.value = (gui_timestep.value + 1) % num_frames - - self.server.flush() - self.next() diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 5b765e0c..cc8c2bae 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -22,6 +22,7 @@ # TODO: Add general utilities for 3D primitives and mesh support. MAP_RADIUS: Final[float] = 500 +BRIGHTNESS_FACTOR: Final[float] = 0.8 def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> trimesh.Trimesh: @@ -36,8 +37,8 @@ def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> tr # Apply translation box_mesh = box_mesh.apply_translation([bbox.center.x, bbox.center.y, bbox.center.z]) - base_color = [r / 255.0 for r in plot_config.fill_color.rgba] - box_mesh.visual.face_colors = plot_config.fill_color.rgba + base_color = [r / 255.0 for r in plot_config.fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba] + box_mesh.visual.face_colors = plot_config.fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba pbr_material = trimesh.visual.material.PBRMaterial( baseColorFactor=base_color, # Your desired color (RGBA, 0-1 range) @@ -105,6 +106,7 @@ def get_map_meshes(scene: AbstractScene): map_surface: AbstractSurfaceMapObject trimesh_mesh = map_surface.trimesh_mesh if map_surface_type in [MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK]: + # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.05).array else: trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z).array @@ -114,7 +116,9 @@ def get_map_meshes(scene: AbstractScene): x=0, y=0, z=center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 ).array - trimesh_mesh.visual.face_colors = MAP_SURFACE_CONFIG[map_surface_type].fill_color.rgba + trimesh_mesh.visual.face_colors = ( + MAP_SURFACE_CONFIG[map_surface_type].fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba + ) surface_meshes.append(trimesh_mesh) output[f"{map_surface_type.serialize()}"] = trimesh.util.concatenate(surface_meshes) @@ -192,7 +196,7 @@ def _create_lane_mesh_from_boundary_arrays( faces = np.array(faces) mesh = trimesh.Trimesh(vertices=vertices, faces=faces) - mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapSurfaceType.LANE].fill_color.rgba + mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapSurfaceType.LANE].fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba return mesh diff --git a/d123/common/visualization/viser/utils_v2.py b/d123/common/visualization/viser/utils_v2.py index 9fc3a45c..ac5137c1 100644 --- a/d123/common/visualization/viser/utils_v2.py +++ b/d123/common/visualization/viser/utils_v2.py @@ -8,6 +8,7 @@ from d123.common.geometry.transform.se3 import translate_body_frame from d123.common.geometry.vector import Vector3D from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG +from d123.common.visualization.viser.utils import BRIGHTNESS_FACTOR from d123.dataset.scene.abstract_scene import AbstractScene # TODO: Refactor this file. @@ -98,7 +99,11 @@ def get_bounding_box_outlines(scene: AbstractScene, iteration: int): bbox_lines = _get_bounding_box_lines(bbox) bbox_lines = translate_points_3d(bbox_lines, initial_ego_vehicle_state.center_se3.point_3d) bbox_color = np.zeros(bbox_lines.shape, dtype=np.float32) - bbox_color[..., :] = BOX_DETECTION_CONFIG[box_detection.metadata.detection_type].fill_color.rgb_norm + bbox_color[..., :] = ( + BOX_DETECTION_CONFIG[box_detection.metadata.detection_type] + .fill_color.set_brightness(BRIGHTNESS_FACTOR) + .rgb_norm + ) lines.append(bbox_lines) colors.append(bbox_color) @@ -106,7 +111,7 @@ def get_bounding_box_outlines(scene: AbstractScene, iteration: int): ego_bbox_lines = _get_bounding_box_lines(ego_vehicle_state.bounding_box_se3) ego_bbox_lines = translate_points_3d(ego_bbox_lines, initial_ego_vehicle_state.center_se3.point_3d) ego_bbox_color = np.zeros(ego_bbox_lines.shape, dtype=np.float32) - ego_bbox_color[..., :] = EGO_VEHICLE_CONFIG.fill_color.rgb_norm + ego_bbox_color[..., :] = EGO_VEHICLE_CONFIG.fill_color.set_brightness(BRIGHTNESS_FACTOR).rgb_norm lines.append(ego_bbox_lines) colors.append(ego_bbox_color) diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py index 753ca36a..a2ba1ee2 100644 --- a/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/womp_boundary_utils.py @@ -22,7 +22,7 @@ MIN_HIT_DISTANCE = 0.1 # meters MIN_AVERAGE_DISTANCE = 1.5 -MAX_AVERAGE_DISTANCE = 6.0 +MAX_AVERAGE_DISTANCE = 7.0 def get_type_and_id_from_token(token: str) -> Tuple[str, int]: diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 713a44c9..2670024e 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -175,78 +175,85 @@ def convert_wopd_tfrecord_log_to_arrow( args: List[Dict[str, Union[List[str], List[Path]]]], data_converter_config: DataConverterConfig ) -> List[Any]: for log_info in args: - tf_record_path: Path = log_info["tf_record"] - split: str = log_info["split"] - - if not tf_record_path.exists(): - raise FileNotFoundError(f"TFRecord path {tf_record_path} does not exist.") - - dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") - for data in dataset: - initial_frame = dataset_pb2.Frame() - initial_frame.ParseFromString(data.numpy()) - break - - log_name = str(initial_frame.context.name) - log_file_path = data_converter_config.output_path / split / f"{log_name}.arrow" - - if data_converter_config.force_log_conversion or not log_file_path.exists(): - log_file_path.unlink(missing_ok=True) - if not log_file_path.parent.exists(): - log_file_path.parent.mkdir(parents=True, exist_ok=True) - - schema_column_list = [ - ("token", pa.string()), - ("timestamp", pa.int64()), - ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), - ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), - ("detections_token", pa.list_(pa.string())), - ("detections_type", pa.list_(pa.int16())), - ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), - ("traffic_light_ids", pa.list_(pa.int64())), - ("traffic_light_types", pa.list_(pa.int16())), - ("scenario_tag", pa.list_(pa.string())), - ("route_lane_group_ids", pa.list_(pa.int64())), - ] - if data_converter_config.lidar_store_option is not None: - if data_converter_config.lidar_store_option == "path": - raise NotImplementedError("Filepath lidar storage is not implemented.") - elif data_converter_config.lidar_store_option == "binary": - schema_column_list.append(("lidar", pa.list_(pa.list_(pa.float32(), 6)))) - - # TODO: Adjust how cameras are added - if data_converter_config.camera_store_option is not None: - for camera_type in WOPD_CAMERA_TYPES.values(): - if data_converter_config.camera_store_option == "path": - raise NotImplementedError("Path camera storage is not implemented.") - elif data_converter_config.camera_store_option == "binary": - schema_column_list.append((camera_type.serialize(), pa.binary())) - schema_column_list.append( - (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) - ) - - recording_schema = pa.schema(schema_column_list) - metadata = LogMetadata( - dataset="wopd", - log_name=log_name, - location=None, # TODO: implement map name - timestep_seconds=0.1, # TODO: Check if correct. Maybe not hardcode - map_has_z=True, - ) - vehicle_parameters = get_wopd_pacifica_parameters() - camera_metadata = get_wopd_camera_metadata(initial_frame) - recording_schema = recording_schema.with_metadata( - { - "log_metadata": json.dumps(asdict(metadata)), - "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), - "camera_metadata": camera_metadata_dict_to_json(camera_metadata), - } - ) - - _write_recording_table(dataset, recording_schema, log_file_path, tf_record_path, data_converter_config) + try: + + tf_record_path: Path = log_info["tf_record"] + split: str = log_info["split"] + + if not tf_record_path.exists(): + raise FileNotFoundError(f"TFRecord path {tf_record_path} does not exist.") + + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") + for data in dataset: + initial_frame = dataset_pb2.Frame() + initial_frame.ParseFromString(data.numpy()) + break + + log_name = str(initial_frame.context.name) + log_file_path = data_converter_config.output_path / split / f"{log_name}.arrow" + + if data_converter_config.force_log_conversion or not log_file_path.exists(): + log_file_path.unlink(missing_ok=True) + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + if data_converter_config.lidar_store_option == "path": + raise NotImplementedError("Filepath lidar storage is not implemented.") + elif data_converter_config.lidar_store_option == "binary": + schema_column_list.append(("lidar", pa.list_(pa.list_(pa.float32(), 6)))) - del recording_schema, vehicle_parameters, dataset - gc.collect() + # TODO: Adjust how cameras are added + if data_converter_config.camera_store_option is not None: + for camera_type in WOPD_CAMERA_TYPES.values(): + if data_converter_config.camera_store_option == "path": + raise NotImplementedError("Path camera storage is not implemented.") + elif data_converter_config.camera_store_option == "binary": + schema_column_list.append((camera_type.serialize(), pa.binary())) + schema_column_list.append( + (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) + ) + + recording_schema = pa.schema(schema_column_list) + metadata = LogMetadata( + dataset="wopd", + log_name=log_name, + location=None, # TODO: implement map name + timestep_seconds=TARGET_DT, # TODO: Check if correct. Maybe not hardcode + map_has_z=True, + ) + vehicle_parameters = get_wopd_pacifica_parameters() + camera_metadata = get_wopd_camera_metadata(initial_frame) + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + } + ) + + _write_recording_table(dataset, recording_schema, log_file_path, tf_record_path, data_converter_config) + + del recording_schema, vehicle_parameters, dataset + except Exception as e: + import traceback + + print(f"Error processing log {str(tf_record_path)}: {e}") + traceback.print_exc() + gc.collect() return [] @@ -286,6 +293,7 @@ def _write_recording_table( with pa.OSFile(str(log_file_path), "wb") as sink: with pa.ipc.new_file(sink, recording_schema) as writer: + dataset = tf.data.TFRecordDataset(tf_record_path, compression_type="") for frame_idx, data in enumerate(dataset): frame = dataset_pb2.Frame() frame.ParseFromString(data.numpy()) diff --git a/d123/dataset/maps/gpkg/gpkg_map.py b/d123/dataset/maps/gpkg/gpkg_map.py index 89b42fed..f219e516 100644 --- a/d123/dataset/maps/gpkg/gpkg_map.py +++ b/d123/dataset/maps/gpkg/gpkg_map.py @@ -41,6 +41,8 @@ def __init__(self, file_path: Path) -> None: MapSurfaceType.WALKWAY: self._get_walkway, MapSurfaceType.CARPARK: self._get_carpark, MapSurfaceType.GENERIC_DRIVABLE: self._get_generic_drivable, + # MapSurfaceType.ROAD_EDGE: self._get_road_edge, + # MapSurfaceType.ROAD_LINE: self._get_road_line, } # loaded during `.initialize()` diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index ca48497e..b1325202 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -18,3 +18,5 @@ class MapSurfaceType(SerialIntEnum): CARPARK = 5 GENERIC_DRIVABLE = 6 STOP_LINE = 7 + # ROAD_EDGE = 8 + # ROAD_LINE = 9 diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index cceb2911..a98adec2 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -19,5 +19,5 @@ defaults: # - carla_dataset - wopd_dataset -force_log_conversion: False -force_map_conversion: True +force_log_conversion: True +force_map_conversion: False diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 21614ea0..70595dac 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -69,6 +69,7 @@ "from d123.common.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", "from d123.common.visualization.matplotlib.observation import (\n", " add_box_detections_to_ax,\n", + " add_default_map_on_ax,\n", " add_ego_vehicle_to_ax,\n", " add_traffic_lights_to_ax,\n", ")\n", @@ -151,16 +152,16 @@ " add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", "\n", - " centroid = map_object.shapely_polygon.centroid\n", - " ax.text(\n", - " centroid.x,\n", - " centroid.y,\n", - " str(map_object.id),\n", - " horizontalalignment=\"center\",\n", - " verticalalignment=\"center\",\n", - " fontsize=8,\n", - " bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", - " )\n", + " # centroid = map_object.shapely_polygon.centroid\n", + " # ax.text(\n", + " # centroid.x,\n", + " # centroid.y,\n", + " # str(map_object.id),\n", + " # horizontalalignment=\"center\",\n", + " # verticalalignment=\"center\",\n", + " # fontsize=8,\n", + " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", + " # )\n", "\n", " except Exception:\n", " import traceback\n", @@ -178,10 +179,11 @@ "\n", " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", " add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", - " # add_box_detections_to_ax(ax, box_detections)\n", - " # add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", "\n", " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", @@ -194,13 +196,14 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " fig, ax = plt.subplots(figsize=(25, 25))\n", + " size = 16\n", + " fig, ax = plt.subplots(figsize=(size, size))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", - "scene_index = 1\n", - "plot_scene_at_iteration(scenes[scene_index], iteration=70, radius=250)" + "scene_index = 5\n", + "plot_scene_at_iteration(scenes[scene_index], iteration=60, radius=50)" ] }, { @@ -218,6 +221,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/notebooks/viz/log_rendering.ipynb b/notebooks/viz/log_rendering.ipynb index 908d1825..5893308e 100644 --- a/notebooks/viz/log_rendering.ipynb +++ b/notebooks/viz/log_rendering.ipynb @@ -13,8 +13,8 @@ "\n", "\n", "\n", - "log_name = \"2021.07.01.20.35.47_veh-38_00016_00281\"\n", - "log_file = Path(f\"/home/daniel/d123_workspace/data/nuplan_private_test/{log_name}.arrow\")\n", + "log_name = \"1005081002024129653_5313_150_5333_150\"\n", + "log_file = Path(f\"/home/daniel/d123_workspace/data/wopd_train/{log_name}.arrow\")\n", "scene = ArrowScene(log_file)\n", "\n", "fig, ax = plot_scene_at_iteration(scene, iteration=10)\n", @@ -47,13 +47,13 @@ "import traceback\n", "from d123.common.visualization.matplotlib.plots import render_scene_animation\n", "\n", - "\n", - "output_path = Path(\"/home/daniel/d123_logs_videos\")\n", - "log_path = Path(\"/home/daniel/d123_workspace/data/carla\")\n", + "split = \"wopd_train\"\n", + "output_path = Path(f\"/home/daniel/d123_logs_videos/{split}\")\n", + "log_path = Path(f\"/home/daniel/d123_workspace/data/{split}\")\n", "for log_file in log_path.iterdir():\n", " scene = ArrowScene(log_file)\n", " try:\n", - " render_scene_animation(scene, output_path, fps=30, end_idx=None, step=5)\n", + " render_scene_animation(scene, output_path, fps=20, end_idx=None, step=1)\n", " except Exception as e:\n", " traceback.print_exc()\n", " erroneous_file = output_path / f\"{log_name}.mp4\"\n", diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index ed6e7fdc..1e24c584 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -74,7 +74,171 @@ "id": "3", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "from d123.common.visualization.color.color import Color\n", + "\n", + "\n", + "# def adjust_brightness(hex_code, brightness_factor):\n", + "# \"\"\"\n", + "# Adjust the brightness of a hex color code.\n", + "\n", + "# Args:\n", + "# hex_code (str): Hex color code (with or without #)\n", + "# brightness_factor (float): Factor to adjust brightness by\n", + "# - Values > 1.0 make the color brighter\n", + "# - Values < 1.0 make the color darker\n", + "# - 1.0 leaves the color unchanged\n", + "# - 0.0 makes the color black\n", + "\n", + "# Returns:\n", + "# str: Adjusted hex color code with # prefix\n", + "\n", + "# Example:\n", + "# adjust_brightness(\"#FF5733\", 1.5) # Make 50% brighter\n", + "# adjust_brightness(\"FF5733\", 0.7) # Make 30% darker\n", + "# \"\"\"\n", + "# # Remove # if present and convert to uppercase\n", + "# hex_code = hex_code.lstrip(\"#\").upper()\n", + "\n", + "# # Validate hex code\n", + "# if len(hex_code) != 6:\n", + "# raise ValueError(\"Hex code must be 6 characters long\")\n", + "\n", + "# try:\n", + "# # Extract RGB components\n", + "# r = int(hex_code[0:2], 16)\n", + "# g = int(hex_code[2:4], 16)\n", + "# b = int(hex_code[4:6], 16)\n", + "# except ValueError:\n", + "# raise ValueError(\"Invalid hex code format\")\n", + "\n", + "# # Adjust brightness\n", + "# r = min(255, max(0, int(r * brightness_factor)))\n", + "# g = min(255, max(0, int(g * brightness_factor)))\n", + "# b = min(255, max(0, int(b * brightness_factor)))\n", + "\n", + "# # Convert back to hex\n", + "# return f\"#{r:02X}{g:02X}{b:02X}\"\n", + "\n", + "\n", + "# def adjust_brightness_hsl(hex_code, brightness_change):\n", + "# \"\"\"\n", + "# Adjust brightness using HSL color space for more natural results.\n", + "\n", + "# Args:\n", + "# hex_code (str): Hex color code (with or without #)\n", + "# brightness_change (float): Amount to change brightness by\n", + "# - Positive values make brighter\n", + "# - Negative values make darker\n", + "# - Range: -1.0 to 1.0\n", + "\n", + "# Returns:\n", + "# str: Adjusted hex color code with # prefix\n", + "# \"\"\"\n", + "# # Remove # if present\n", + "# hex_code = hex_code.lstrip(\"#\")\n", + "\n", + "# # Convert hex to RGB\n", + "# r = int(hex_code[0:2], 16) / 255.0\n", + "# g = int(hex_code[2:4], 16) / 255.0\n", + "# b = int(hex_code[4:6], 16) / 255.0\n", + "\n", + "# # Convert RGB to HSL\n", + "# max_val = max(r, g, b)\n", + "# min_val = min(r, g, b)\n", + "# diff = max_val - min_val\n", + "\n", + "# # Calculate lightness\n", + "# lightness = (max_val + min_val) / 2\n", + "\n", + "# if diff == 0:\n", + "# hue = saturation = 0\n", + "# else:\n", + "# # Calculate saturation\n", + "# if lightness < 0.5:\n", + "# saturation = diff / (max_val + min_val)\n", + "# else:\n", + "# saturation = diff / (2 - max_val - min_val)\n", + "\n", + "# # Calculate hue\n", + "# if max_val == r:\n", + "# hue = ((g - b) / diff + (6 if g < b else 0)) / 6\n", + "# elif max_val == g:\n", + "# hue = ((b - r) / diff + 2) / 6\n", + "# else:\n", + "# hue = ((r - g) / diff + 4) / 6\n", + "\n", + "# # Adjust lightness\n", + "# lightness = max(0, min(1, lightness + brightness_change))\n", + "\n", + "# # Convert HSL back to RGB\n", + "# if saturation == 0:\n", + "# r = g = b = lightness\n", + "# else:\n", + "\n", + "# def hue_to_rgb(p, q, t):\n", + "# if t < 0:\n", + "# t += 1\n", + "# if t > 1:\n", + "# t -= 1\n", + "# if t < 1 / 6:\n", + "# return p + (q - p) * 6 * t\n", + "# if t < 1 / 2:\n", + "# return q\n", + "# if t < 2 / 3:\n", + "# return p + (q - p) * (2 / 3 - t) * 6\n", + "# return p\n", + "\n", + "# q = lightness * (1 + saturation) if lightness < 0.5 else lightness + saturation - lightness * saturation\n", + "# p = 2 * lightness - q\n", + "\n", + "# r = hue_to_rgb(p, q, hue + 1 / 3)\n", + "# g = hue_to_rgb(p, q, hue)\n", + "# b = hue_to_rgb(p, q, hue - 1 / 3)\n", + "\n", + "# # Convert back to hex\n", + "# r = int(round(r * 255))\n", + "# g = int(round(g * 255))\n", + "# b = int(round(b * 255))\n", + "\n", + "# return f\"#{r:02X}{g:02X}{b:02X}\"\n", + "\n", + "\n", + "# Example usage and testing\n", + "if __name__ == \"__main__\":\n", + " # Test the functions\n", + "\n", + "\n", + " # Test with different colors\n", + " test_colors = [\n", + " Color(\"#3498DB\"),\n", + " Color(\"#E74C3C\"),\n", + " Color(\"#2ECC71\"),\n", + " Color(\"#F39C12\"),\n", + " ]\n", + "\n", + " import matplotlib.pyplot as plt\n", + "\n", + "\n", + " # Visualize the color adjustments\n", + " factor = 0.5\n", + " fig, ax = plt.subplots(len(test_colors), 3, figsize=(8, 12))\n", + " for i, color in enumerate(test_colors):\n", + " ax[i, 0].imshow([[color.rgb]])\n", + " ax[i, 0].set_title(f\"Original: {color}\")\n", + " ax[i, 0].axis(\"off\")\n", + "\n", + " ax[i, 1].imshow([[color.set_brightness(factor).rgb]])\n", + " ax[i, 1].set_title(f\"Brighter: {str(color.set_brightness(factor))}\")\n", + " ax[i, 1].axis(\"off\")\n", + "\n", + " ax[i, 2].imshow([[color.set_brightness(1 / factor).rgb]])\n", + " ax[i, 2].set_title(f\"Darker: {str(color.set_brightness(1 / factor))}\")\n", + " ax[i, 2].axis(\"off\")\n", + "\n", + " plt.tight_layout()\n", + " plt.show()" + ] }, { "cell_type": "code", From 0494ce7aec50f4d1e5a046a59c0f9f6f3744d46a Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sat, 9 Aug 2025 19:20:53 +0200 Subject: [PATCH 24/42] Add road edges to nuPlan and carla map conversion --- d123/common/geometry/occupancy_map.py | 8 + d123/common/visualization/color/config.py | 2 +- d123/common/visualization/viser/server.py | 27 +- d123/common/visualization/viser/utils.py | 21 +- .../opendrive/conversion/group_collections.py | 2 +- .../carla/opendrive/opendrive_converter.py | 247 +++++++++++++-- .../nuplan/nuplan_data_converter.py | 6 +- .../nuplan/nuplan_map_conversion.py | 80 ++++- .../wopd/waymo_map_utils/wopd_map_utils.py | 22 ++ d123/dataset/maps/abstract_map_objects.py | 60 +++- d123/dataset/maps/gpkg/gpkg_map.py | 20 +- d123/dataset/maps/gpkg/gpkg_map_objects.py | 35 ++- d123/dataset/maps/map_datatypes.py | 4 +- .../default_dataset_conversion.yaml | 8 +- notebooks/Untitled-2.ipynb | 286 ++++++++++++++++++ notebooks/viz/bev_matplotlib.ipynb | 67 +++- notebooks/viz/viser_testing_v2_map.ipynb | 116 ------- notebooks/viz/viser_testing_v2_scene.ipynb | 184 +---------- 18 files changed, 852 insertions(+), 343 deletions(-) create mode 100644 notebooks/Untitled-2.ipynb delete mode 100644 notebooks/viz/viser_testing_v2_map.ipynb diff --git a/d123/common/geometry/occupancy_map.py b/d123/common/geometry/occupancy_map.py index ed80ae23..2a8085d6 100644 --- a/d123/common/geometry/occupancy_map.py +++ b/d123/common/geometry/occupancy_map.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Dict, List, Literal, Optional, Sequence, Union import numpy as np @@ -34,6 +36,12 @@ def __init__( self._node_capacity = node_capacity self._str_tree = STRtree(self._geometries, node_capacity) + @classmethod + def from_dict(cls, geometry_dict: Dict[Union[str, int], BaseGeometry], node_capacity: int = 10) -> OccupancyMap2D: + ids = list(geometry_dict.keys()) + geometries = list(geometry_dict.values()) + return cls(geometries=geometries, ids=ids, node_capacity=node_capacity) + def __getitem__(self, id: Union[str, int]) -> BaseGeometry: """ Retrieves geometry of token. diff --git a/d123/common/visualization/color/config.py b/d123/common/visualization/color/config.py index 69eb5f3b..ccc7ecbb 100644 --- a/d123/common/visualization/color/config.py +++ b/d123/common/visualization/color/config.py @@ -9,7 +9,7 @@ class PlotConfig: fill_color: Color = field(default_factory=lambda: BLACK) fill_color_alpha: float = 1.0 line_color: Color = BLACK - line_color_alpha: float = 0.0 + line_color_alpha: float = 1.0 line_width: float = 1.0 line_style: str = "-" marker_style: str = "o" diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 984aa968..97c456c5 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -5,10 +5,12 @@ import viser from d123.common.datatypes.sensor.camera import CameraType +from d123.common.visualization.color.color import BLACK from d123.common.visualization.viser.utils import ( get_bounding_box_meshes, get_camera_values, get_lidar_points, + get_map_lines, get_map_meshes, ) from d123.common.visualization.viser.utils_v2 import get_bounding_box_outlines @@ -28,13 +30,17 @@ CameraType.CAM_R2, ] -LIDAR_POINT_SIZE: float = 0.05 -MAP_AVAILABLE: bool = True -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" +# MISC config: LINE_WIDTH: float = 4.0 -CAMERA_SCALE: float = 1.0 +# Bounding box config: +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" + +# Map config: +MAP_AVAILABLE: bool = True + +# Cameras config: # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ # CameraType.CAM_F0, # CameraType.CAM_L0, @@ -46,7 +52,11 @@ # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] +CAMERA_SCALE: float = 1.0 + +# Lidar config: LIDAR_AVAILABLE: bool = False +LIDAR_POINT_SIZE: float = 0.05 class ViserVisualizationServer: @@ -229,7 +239,7 @@ def _(_) -> None: for name, mesh in get_map_meshes(scene).items(): self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) - # centerlines, __, __ = get_map_lines(scene) + centerlines, __, __, road_edges = get_map_lines(scene) # for i, centerline in enumerate(centerlines): # self.server.scene.add_line_segments( # "/map/centerlines", @@ -249,6 +259,13 @@ def _(_) -> None: # colors=[[TAB_10[3].rgb]], # line_width=LINE_WIDTH, # ) + print(centerlines.shape, road_edges.shape) + self.server.scene.add_line_segments( + "/map/road_edges", + road_edges, + colors=[[BLACK.rgb]], + line_width=LINE_WIDTH, + ) # Playback update loop. prev_timestep = gui_timestep.value diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index cc8c2bae..a2981153 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -93,6 +93,7 @@ def get_map_meshes(scene: AbstractScene): MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK, MapSurfaceType.CARPARK, + MapSurfaceType.GENERIC_DRIVABLE, ] map_objects_dict = scene.map_api.get_proximal_map_objects( @@ -126,7 +127,7 @@ def get_map_meshes(scene: AbstractScene): def get_map_lines(scene: AbstractScene): - map_surface_types = [MapSurfaceType.LANE] + map_surface_types = [MapSurfaceType.LANE, MapSurfaceType.ROAD_EDGE] initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 map_objects_dict = scene.map_api.get_proximal_map_objects( @@ -139,20 +140,30 @@ def polyline_to_segments(polyline: Polyline3D) -> npt.NDArray[np.float64]: polyline_array = np.concatenate([polyline_array[:-1], polyline_array[1:]], axis=1) return polyline_array - lanes = map_objects_dict[MapSurfaceType.LANE] - centerlines, right_boundaries, left_boundaries = [], [], [] - for lane in lanes: + centerlines, right_boundaries, left_boundaries, road_edges = [], [], [], [] + for lane in map_objects_dict[MapSurfaceType.LANE]: lane: AbstractLane centerlines.append(polyline_to_segments(lane.centerline)) right_boundaries.append(polyline_to_segments(lane.right_boundary)) left_boundaries.append(polyline_to_segments(lane.left_boundary)) + for road_edge in map_objects_dict[MapSurfaceType.ROAD_EDGE]: + road_edges.append(polyline_to_segments(road_edge.polyline_3d)) + centerlines = np.concatenate(centerlines, axis=0) left_boundaries = np.concatenate(left_boundaries, axis=0) right_boundaries = np.concatenate(right_boundaries, axis=0) + road_edges = np.concatenate(road_edges, axis=0) + + if not scene.log_metadata.map_has_z: + # If the map does not have a z-coordinate, we set it to the height of the ego vehicle. + centerlines[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + left_boundaries[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + right_boundaries[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 + road_edges[:, :, 2] += center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 - return centerlines, left_boundaries, right_boundaries + return centerlines, left_boundaries, right_boundaries, road_edges def get_trimesh_from_boundaries( diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py b/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py index 845f1cfa..e5082014 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py +++ b/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py @@ -126,7 +126,7 @@ def center_polyline_3d(self) -> npt.NDArray[np.float64]: def outline_polyline_3d(self) -> npt.NDArray[np.float64]: inner_polyline = self.inner_polyline_3d[::-1] outer_polyline = self.outer_polyline_3d - return np.concatenate([inner_polyline, outer_polyline], axis=0, dtype=np.float64) + return np.concatenate([inner_polyline, outer_polyline, inner_polyline[None, 0]], axis=0, dtype=np.float64) @property def shapely_polygon(self) -> shapely.Polygon: diff --git a/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py b/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py index cd22aebe..d3e8b806 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py +++ b/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py @@ -1,14 +1,19 @@ import os import warnings +from collections import defaultdict from pathlib import Path -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Set import geopandas as gpd +import networkx as nx import numpy as np +import numpy.typing as npt import pandas as pd import shapely from shapely.ops import polygonize, unary_union +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.occupancy_map import OccupancyMap2D from d123.dataset.dataset_specific.carla.opendrive.conversion.group_collections import ( OpenDriveLaneGroupHelper, OpenDriveLaneHelper, @@ -27,6 +32,7 @@ from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import Junction, OpenDrive, Road from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border from d123.dataset.dataset_specific.carla.opendrive.id_mapping import IntIDMapping +from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import get_road_edge_linestrings from d123.dataset.maps.map_datatypes import MapSurfaceType ENABLE_WARNING: bool = False @@ -77,6 +83,8 @@ def run(self, map_name: str) -> None: crosswalk_df, ) + road_edge_df = self._extract_road_edge_df(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) + # Store dataframes map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" lane_df.to_file(map_file_name, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") @@ -91,6 +99,7 @@ def run(self, map_name: str) -> None: intersections_df.to_file(map_file_name, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") lane_group_df.to_file(map_file_name, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") crosswalk_df.to_file(map_file_name, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") + road_edge_df.to_file(map_file_name, layer=MapSurfaceType.ROAD_EDGE.serialize(), driver="GPKG", mode="a") def _collect_lane_helpers(self) -> None: for road in self.opendrive.roads: @@ -364,16 +373,16 @@ def _extract_lane_dataframe(self) -> gpd.GeoDataFrame: def _extract_walkways_dataframe(self) -> gpd.GeoDataFrame: ids = [] - # left_boundaries = [] - # right_boundaries = [] + left_boundaries = [] + right_boundaries = [] outlines = [] geometries = [] for lane_helper in self.lane_helper_dict.values(): if lane_helper.type == "sidewalk": ids.append(lane_helper.lane_id) - # left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - # right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) geometries.append(lane_helper.shapely_polygon) @@ -390,24 +399,24 @@ def _extract_walkways_dataframe(self) -> gpd.GeoDataFrame: def _extract_carpark_dataframe(self) -> gpd.GeoDataFrame: ids = [] - # left_boundaries = [] - # right_boundaries = [] + left_boundaries = [] + right_boundaries = [] outlines = [] geometries = [] for lane_helper in self.lane_helper_dict.values(): if lane_helper.type == "parking": ids.append(lane_helper.lane_id) - # left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - # right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) geometries.append(lane_helper.shapely_polygon) data = pd.DataFrame( { "id": ids, - # "left_boundary": left_boundaries, - # "right_boundary": right_boundaries, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, "outline": outlines, } ) @@ -416,24 +425,24 @@ def _extract_carpark_dataframe(self) -> gpd.GeoDataFrame: def _extract_generic_drivable_dataframe(self) -> gpd.GeoDataFrame: ids = [] - # left_boundaries = [] - # right_boundaries = [] + left_boundaries = [] + right_boundaries = [] outlines = [] geometries = [] for lane_helper in self.lane_helper_dict.values(): if lane_helper.type in ["none", "border", "bidirectional"]: ids.append(lane_helper.lane_id) - # left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - # right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) geometries.append(lane_helper.shapely_polygon) data = pd.DataFrame( { "id": ids, - # "left_boundary": left_boundaries, - # "right_boundary": left_boundaries, + "left_boundary": left_boundaries, + "right_boundary": left_boundaries, "outline": outlines, } ) @@ -554,6 +563,19 @@ def _convert_ids_to_int( lambda x: lane_group_id_mapping.map_list(x) ) + def _extract_road_edge_df( + self, + lane_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, + ) -> None: + road_edges = _get_road_edges_from_gdf(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) + + ids = np.arange(len(road_edges), dtype=np.int64).tolist() + geometries = road_edges + return gpd.GeoDataFrame(pd.DataFrame({"id": ids}), geometry=geometries) + # TODO: move this somewhere else and improve def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper]) -> shapely.Polygon: @@ -581,3 +603,194 @@ def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper] else: # Take the largest polygon if there are multiple return max(polygons, key=lambda p: p.area) + + +def _get_road_edges_from_gdf( + lane_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> List[shapely.LineString]: + + # 1. Find conflicting lane groups, e.g. groups of lanes that overlap in 2D but have different Z-values (bridges) + conflicting_lane_groups = _get_conflicting_lane_groups(lane_group_df, lane_df) + + # 2. Extract road edges in 2D (including conflicting lane groups) + drivable_polygons = ( + lane_group_df.geometry.tolist() + carpark_df.geometry.tolist() + generic_drivable_area_df.geometry.tolist() + ) + road_edges_2d = get_road_edge_linestrings(drivable_polygons, max_road_edge_length=None) + + # 3. Collect 3D boundaries of non-conflicting lane groups and other drivable areas + non_conflicting_boundaries: List[shapely.LineString] = [] + for lane_group_id, lane_group_helper in lane_group_df.iterrows(): + if lane_group_id not in conflicting_lane_groups.keys(): + non_conflicting_boundaries.append(lane_group_helper["left_boundary"]) + non_conflicting_boundaries.append(lane_group_helper["right_boundary"]) + for outline in carpark_df.outline.tolist() + generic_drivable_area_df.outline.tolist(): + non_conflicting_boundaries.append(outline) + + # 4. Lift road edges to 3D using the boundaries of non-conflicting elements + non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, non_conflicting_boundaries) + + # 5. Add road edges from conflicting lane groups + resolved_road_edges = _resolve_conflicting_lane_groups(conflicting_lane_groups, lane_group_df) + + all_road_edges = non_conflicting_road_edges + resolved_road_edges + + return all_road_edges + + +def _get_nearest_z_from_points_3d(points_3d: npt.NDArray[np.float64], query_point: npt.NDArray[np.float64]) -> float: + assert points_3d.ndim == 2 and points_3d.shape[1] == len( + Point3DIndex + ), "points_3d must be a 2D array with shape (N, 3)" + distances = np.linalg.norm(points_3d[..., Point3DIndex.XY] - query_point[..., Point3DIndex.XY], axis=1) + closest_point = points_3d[np.argmin(distances)] + return closest_point[2] + + +def _get_conflicting_lane_groups(lane_group_df: gpd.GeoDataFrame, lane_df: gpd.GeoDataFrame) -> Dict[int, List[int]]: + Z_THRESHOLD = 5.0 # [m] Z-value threshold for conflict detection + + ids = lane_group_df.id.tolist() + polygons = lane_group_df.geometry.tolist() + + def _get_centerline_points_3d(lane_group_id: int) -> npt.NDArray[np.float64]: + """Helper function to get the centerline points in 3D.""" + lane_ids = lane_group_df[lane_group_df.id == lane_group_id].lane_ids.values[0] + centerlines: List[npt.NDArray[np.float64]] = [] + for lane_id in lane_ids: + centerline = lane_df[lane_df.id == lane_id].baseline_path.values[0] + assert isinstance(centerline, shapely.LineString) + centerlines.append(np.array(centerline.coords, dtype=np.float64)) + return np.concatenate(centerlines, axis=0) + + occupancy_map = OccupancyMap2D(polygons, ids) + conflicting_lane_groups: Dict[int, List[int]] = defaultdict(list) + + for lane_group_id, lane_group_polygon in zip(ids, polygons): + + # Extract internal centerline points for Z-value check + lane_group_centerlines = _get_centerline_points_3d(lane_group_id) + + # Query lane groups that overlap in 2D + intersecting_lane_group_ids = occupancy_map.intersects(lane_group_polygon) + intersecting_lane_group_ids.remove(lane_group_id) # Remove self from the list + for intersecting_id in intersecting_lane_group_ids: + intersecting_geometry = occupancy_map[intersecting_id] + + # ignore non-polygon geometries + if intersecting_geometry.geom_type != "Polygon": + continue + + # Check if Z-values deviate at intersection centroid + intersection_centroid = np.array(intersecting_geometry.centroid.coords[0], dtype=np.float64) + intersecting_centerlines = _get_centerline_points_3d(intersecting_id) + z_at_intersecting = _get_nearest_z_from_points_3d(intersecting_centerlines, intersection_centroid) + z_at_lane_group = _get_nearest_z_from_points_3d(lane_group_centerlines, intersection_centroid) + if np.abs(z_at_lane_group - z_at_intersecting) < Z_THRESHOLD: + continue + conflicting_lane_groups[lane_group_id].append(intersecting_id) + return conflicting_lane_groups + + +def _resolve_conflicting_lane_groups( + conflicting_lane_groups: Dict[int, List[int]], lane_group_df: gpd.GeoDataFrame +) -> List[shapely.LineString]: + + # Split conflicting lane groups into non-conflicting sets for further merging + non_conflicting_sets = create_non_conflicting_sets(conflicting_lane_groups) + + road_edges_3d: List[shapely.LineString] = [] + for non_conflicting_set in non_conflicting_sets: + + # Collect 2D polygons of non-conflicting lane group set + set_polygons = [ + lane_group_df[lane_group_df.id == lane_group_id].geometry.values[0] for lane_group_id in non_conflicting_set + ] + # Get 2D road edge linestrings for the non-conflicting set + set_road_edges_2d = get_road_edge_linestrings(set_polygons, max_road_edge_length=None) + + # Collect 3D boundaries of non-conflicting lane groups + set_boundaries_3d: List[shapely.LineString] = [] + for lane_group_id in non_conflicting_set: + lane_group_helper = lane_group_df[lane_group_df.id == lane_group_id] + set_boundaries_3d.append(lane_group_helper.left_boundary.values[0]) + set_boundaries_3d.append(lane_group_helper.right_boundary.values[0]) + + # Lift road edges to 3D using the boundaries of non-conflicting lane groups + lifted_road_edges_3d = lift_road_edges_to_3d(set_road_edges_2d, set_boundaries_3d) + road_edges_3d.extend(lifted_road_edges_3d) + + return road_edges_3d + + +def lift_road_edges_to_3d( + road_edges_2d: List[shapely.LineString], boundaries: List[shapely.LineString] +) -> List[shapely.LineString]: + + QUERY_MAX_DISTANCE = 0.01 # [m] Maximum distance for nearest neighbor query + + def _find_continuous_sublists(integers: List[int]) -> List[List[int]]: + """Find continuous sublists in a list of integers.""" + arr = np.array(integers, dtype=np.int64) + breaks = np.where(np.diff(arr) != 1)[0] + 1 + splits = np.split(arr, breaks) + return [sublist.tolist() for sublist in splits] + + occupancy_map = OccupancyMap2D(boundaries) + + road_edges_3d: List[shapely.LineString] = [] + for idx, linestring in enumerate(road_edges_2d): + # print(list(linestring.coords)) + points_3d = np.array(list(linestring.coords), dtype=np.float64) + + results = occupancy_map.query_nearest( + shapely.points(points_3d[..., :2]), max_distance=QUERY_MAX_DISTANCE, exclusive=True + ) + for query_idx, geometry_idx in zip(*results): + intersecting_boundary: shapely.LineString = occupancy_map[occupancy_map.ids[geometry_idx]] + points_3d[query_idx, 2] = _get_nearest_z_from_points_3d( + np.array(list(intersecting_boundary.coords), dtype=np.float64), points_3d[query_idx] + ) + + for continuous_slice in _find_continuous_sublists(results[0]): + if len(continuous_slice) < 2: + continue + lifted_linestring = shapely.LineString(points_3d[continuous_slice]) + road_edges_3d.append(lifted_linestring) + return road_edges_3d + + +def create_non_conflicting_sets(conflicts: Dict[int, List[int]]) -> List[Set[int]]: + """ + Creates sets of non-conflicting indices using NetworkX. + """ + # Create graph from conflicts + G = nx.Graph() + for idx, conflict_list in conflicts.items(): + for conflict_idx in conflict_list: + G.add_edge(idx, conflict_idx) + + result = [] + + # Process each connected component + for component in nx.connected_components(G): + subgraph = G.subgraph(component) + + # Try bipartite coloring first (most common case) + if nx.is_bipartite(subgraph): + sets = nx.bipartite.sets(subgraph) + result.extend([set(s) for s in sets]) + else: + # Fall back to greedy coloring for non-bipartite graphs + coloring = nx.greedy_color(subgraph, strategy="largest_first") + color_groups = {} + for node, color in coloring.items(): + if color not in color_groups: + color_groups[color] = set() + color_groups[color].add(node) + result.extend(color_groups.values()) + + return result diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index 59aea551..012b169c 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -242,9 +242,9 @@ def convert_nuplan_log_to_arrow( _write_recording_table(log_db, recording_schema, log_file_path, log_path, data_converter_config) - log_db.detach_tables() - log_db.remove_ref() - del recording_schema, vehicle_parameters, log_db + log_db.detach_tables() + log_db.remove_ref() + del recording_schema, vehicle_parameters, log_db gc.collect() return [] diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py index f3861aa9..10c338be 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py @@ -2,15 +2,20 @@ import os import warnings +from ast import Tuple from pathlib import Path -from typing import Dict, List, Optional +from typing import Dict, Final, List, Optional import geopandas as gpd import numpy as np +import numpy.typing as npt import pandas as pd import pyogrio +from shapely import union_all from shapely.geometry import LineString +from sympy import Polygon +from d123.common.geometry.constants import DEFAULT_Z from d123.dataset.maps.gpkg.utils import get_all_rows_with_value, get_row_with_value from d123.dataset.maps.map_datatypes import MapSurfaceType @@ -42,6 +47,10 @@ "gen_lane_connectors_scaled_width_polygons", ] +ROAD_EDGE_STEP_SIZE = 0.25 # meters +MAX_ROAD_EDGE_LENGTH = 100.0 # meters, used to filter out very long road edges +ROAD_EDGE_BUFFER: Final[float] = 0.05 + class NuPlanMapConverter: def __init__(self, map_path: Path) -> None: @@ -62,6 +71,7 @@ def convert(self, map_name: str = "us-pa-pittsburgh-hazelwood") -> None: walkway_df = self._extract_walkway_dataframe() carpark_df = self._extract_carpark_dataframe() generic_drivable_df = self._extract_generic_drivable_dataframe() + road_edge_df = self._extract_road_edge_dataframe() if not self._map_path.exists(): self._map_path.mkdir(parents=True, exist_ok=True) @@ -76,6 +86,7 @@ def convert(self, map_name: str = "us-pa-pittsburgh-hazelwood") -> None: generic_drivable_df.to_file( map_file_name, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" ) + road_edge_df.to_file(map_file_name, layer=MapSurfaceType.ROAD_EDGE.serialize(), driver="GPKG", mode="a") def _load_dataframes(self, map_file_path: Path) -> None: @@ -384,6 +395,18 @@ def _extract_generic_drivable_dataframe(self) -> gpd.GeoDataFrame: data = pd.DataFrame({"id": self._gdf["generic_drivable_areas"].fid.to_list()}) return gpd.GeoDataFrame(data, geometry=self._gdf["generic_drivable_areas"].geometry.to_list()) + def _extract_road_edge_dataframe(self) -> gpd.GeoDataFrame: + drivable_polygons = ( + self._gdf["intersections"].geometry.to_list() + + self._gdf["lane_groups_polygons"].geometry.to_list() + + self._gdf["carpark_areas"].geometry.to_list() + + self._gdf["generic_drivable_areas"].geometry.to_list() + ) + road_edge_linestrings = get_road_edge_linestrings(drivable_polygons) + + data = pd.DataFrame({"id": [idx for idx in range(len(road_edge_linestrings))]}) + return gpd.GeoDataFrame(data, geometry=road_edge_linestrings) + def flip_linestring(linestring: LineString) -> LineString: # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. @@ -409,3 +432,58 @@ def align_boundary_direction(centerline: LineString, boundary: LineString) -> Li if not lines_same_direction(centerline, boundary): return flip_linestring(boundary) return boundary + + +def get_road_edge_linestrings( + drivable_polygons: List[Polygon], + step_size: float = ROAD_EDGE_STEP_SIZE, + max_road_edge_length: Optional[float] = MAX_ROAD_EDGE_LENGTH, +) -> List[LineString]: + # TODO: move this function for general usage. + + def _coords_to_points(coords: npt.NDArray[np.float64]) -> List[npt.NDArray[np.float64]]: + coords = np.array(coords).reshape((-1, 2)) + linestring = LineString(coords) + + points_list: List[npt.NDArray[np.float64]] = [] + + segments: List[Tuple[float, float]] = [] # Start and end points of segments + if max_road_edge_length is not None: + num_segments = int(linestring.length / max_road_edge_length) + for i in range(num_segments + 1): + start = np.clip(float(i) * max_road_edge_length, 0, linestring.length) + end = np.clip((float(i) + 1) * max_road_edge_length, 0, linestring.length) + segments.append((start, end)) + else: + segments = [(0.0, linestring.length)] + + for segment in segments: + start, end = segment + distances = np.arange(start, end + step_size, step_size) + points = [linestring.interpolate(distance, normalized=False) for distance in distances] + points_list.append(np.array([[p.x, p.y, DEFAULT_Z] for p in points])) + + return points_list + + def _polygon_to_coords(polygon: Polygon) -> List[npt.NDArray[np.float64]]: + assert polygon.geom_type == "Polygon" + points_list = [] + points_list.extend(_coords_to_points(polygon.exterior.coords)) + for interior in polygon.interiors: + points_list.extend(_coords_to_points(interior.coords)) + return points_list + + union_polygon = union_all([polygon.buffer(ROAD_EDGE_BUFFER, join_style=2) for polygon in drivable_polygons]).buffer( + -ROAD_EDGE_BUFFER, join_style=2 + ) + + linestring_list = [] + if union_polygon.geom_type == "Polygon": + for polyline in _polygon_to_coords(union_polygon): + linestring_list.append(LineString(polyline)) + elif union_polygon.geom_type == "MultiPolygon": + for polygon in union_polygon.geoms: + for polyline in _polygon_to_coords(polygon): + linestring_list.append(LineString(polyline)) + + return linestring_list diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py index d28df0e6..160dae76 100644 --- a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py @@ -105,6 +105,8 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: walkway_df = get_walkway_df() carpark_df = get_carpark_df(carparks) generic_drivable_df = get_generic_drivable_df() + road_edge_df = get_road_edge_df(road_edges) + road_line_df = get_road_line_df(road_lines) map_file_path.unlink(missing_ok=True) if not map_file_path.parent.exists(): @@ -119,6 +121,8 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: generic_drivable_df.to_file( map_file_path, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" ) + road_edge_df.to_file(map_file_path, layer=MapSurfaceType.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + road_line_df.to_file(map_file_path, layer=MapSurfaceType.ROAD_LINE.serialize(), driver="GPKG", mode="a") def get_lane_df( @@ -284,3 +288,21 @@ def get_generic_drivable_df() -> gpd.GeoDataFrame: data = pd.DataFrame({"id": ids}) gdf = gpd.GeoDataFrame(data, geometry=geometries) return gdf + + +def get_road_edge_df(road_edges: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(road_edges.keys()) + geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_edges.values()] + + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_road_line_df(road_lines: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(road_lines.keys()) + geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_lines.values()] + + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf diff --git a/d123/dataset/maps/abstract_map_objects.py b/d123/dataset/maps/abstract_map_objects.py index e8e33f8a..93eaa6aa 100644 --- a/d123/dataset/maps/abstract_map_objects.py +++ b/d123/dataset/maps/abstract_map_objects.py @@ -1,7 +1,7 @@ from __future__ import annotations import abc -from typing import List, Optional +from typing import List, Optional, Tuple import shapely.geometry as geom import trimesh @@ -127,6 +127,14 @@ def lane_group(self) -> AbstractLaneGroup: :return: returns lane group """ + @property + def boundaries(self) -> Tuple[Polyline3D, Polyline3D]: + """ + Property of left and right boundary. + :return: returns tuple of left and right boundary polylines + """ + return self.left_boundary, self.right_boundary + class AbstractLaneGroup(AbstractSurfaceMapObject): """Abstract interface lane groups (nearby lanes going in the same direction).""" @@ -239,3 +247,53 @@ class AbstractStopLine(AbstractSurfaceMapObject): def surface_type(self) -> MapSurfaceType: # return MapSurfaceType.STOP_LINE raise NotImplementedError + + +class AbstractRoadEdge(AbstractMapObject): + """Abstract interface for road edge objects.""" + + @property + def surface_type(self) -> MapSurfaceType: + return MapSurfaceType.ROAD_EDGE + + @property + @abc.abstractmethod + def polyline_3d(self) -> Polyline3D: + """ + Returns the 3D polyline of the road edge. + :return: 3D polyline + """ + raise NotImplementedError + + +class AbstractRoadLine(AbstractMapObject): + """Abstract interface for road line objects.""" + + @property + def surface_type(self) -> MapSurfaceType: + return MapSurfaceType.ROAD_LINE + + @property + @abc.abstractmethod + def polyline_3d(self) -> Polyline3D: + """ + Returns the 3D polyline of the road edge. + :return: 3D polyline + """ + raise NotImplementedError + + @property + def polyline_2d(self) -> Polyline2D: + """ + Returns the 2D polyline of the road line. + :return: 2D polyline + """ + return self.polyline_3d.polyline_2d + + @property + def polyline_se2(self) -> Polyline2D: + """ + Returns the 2D polyline of the road line in SE(2) coordinates. + :return: 2D polyline in SE(2) + """ + return self.polyline_3d.polyline_se2 diff --git a/d123/dataset/maps/gpkg/gpkg_map.py b/d123/dataset/maps/gpkg/gpkg_map.py index f219e516..e193c012 100644 --- a/d123/dataset/maps/gpkg/gpkg_map.py +++ b/d123/dataset/maps/gpkg/gpkg_map.py @@ -21,6 +21,8 @@ GPKGIntersection, GPKGLane, GPKGLaneGroup, + GPKGRoadEdge, + GPKGRoadLine, GPKGWalkway, ) from d123.dataset.maps.gpkg.utils import load_gdf_with_geometry_columns @@ -41,8 +43,8 @@ def __init__(self, file_path: Path) -> None: MapSurfaceType.WALKWAY: self._get_walkway, MapSurfaceType.CARPARK: self._get_carpark, MapSurfaceType.GENERIC_DRIVABLE: self._get_generic_drivable, - # MapSurfaceType.ROAD_EDGE: self._get_road_edge, - # MapSurfaceType.ROAD_LINE: self._get_road_line, + MapSurfaceType.ROAD_EDGE: self._get_road_edge, + MapSurfaceType.ROAD_LINE: self._get_road_line, } # loaded during `.initialize()` @@ -309,6 +311,20 @@ def _get_generic_drivable(self, id: str) -> Optional[GPKGGenericDrivable]: else None ) + def _get_road_edge(self, id: str) -> Optional[GPKGRoadEdge]: + return ( + GPKGRoadEdge(id, self._gpd_dataframes[MapSurfaceType.ROAD_EDGE]) + if id in self._gpd_dataframes[MapSurfaceType.ROAD_EDGE]["id"].tolist() + else None + ) + + def _get_road_line(self, id: str) -> Optional[GPKGRoadLine]: + return ( + GPKGRoadLine(id, self._gpd_dataframes[MapSurfaceType.ROAD_LINE]) + if id in self._gpd_dataframes[MapSurfaceType.ROAD_LINE]["id"].tolist() + else None + ) + # def _query_layer( # self, # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], diff --git a/d123/dataset/maps/gpkg/gpkg_map_objects.py b/d123/dataset/maps/gpkg/gpkg_map_objects.py index 6965b885..876bdf4d 100644 --- a/d123/dataset/maps/gpkg/gpkg_map_objects.py +++ b/d123/dataset/maps/gpkg/gpkg_map_objects.py @@ -6,6 +6,7 @@ import geopandas as gpd import numpy as np +import pandas as pd import shapely.geometry as geom import trimesh @@ -19,6 +20,8 @@ AbstractIntersection, AbstractLane, AbstractLaneGroup, + AbstractRoadEdge, + AbstractRoadLine, AbstractSurfaceMapObject, AbstractWalkway, ) @@ -206,7 +209,7 @@ def intersection(self) -> Optional[GPKGIntersection]: self._lane_df, self._object_df, ) - if intersection_id is not None + if intersection_id is not None and not pd.isna(intersection_id) else None ) @@ -256,3 +259,33 @@ def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): class GPKGGenericDrivable(GPKGSurfaceObject, AbstractGenericDrivable): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): super().__init__(object_id, object_df) + + +class GPKGRoadEdge(AbstractRoadEdge): + def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): + super().__init__(object_id) + self._object_df = object_df + + @cached_property + def _object_row(self) -> gpd.GeoSeries: + return get_row_with_value(self._object_df, "id", self.id) + + @property + def polyline_3d(self) -> Polyline3D: + """Inherited, see superclass.""" + return Polyline3D.from_linestring(self._object_row.geometry) + + +class GPKGRoadLine(AbstractRoadLine): + def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): + super().__init__(object_id) + self._object_df = object_df + + @cached_property + def _object_row(self) -> gpd.GeoSeries: + return get_row_with_value(self._object_df, "id", self.id) + + @property + def polyline_3d(self) -> Polyline3D: + """Inherited, see superclass.""" + return Polyline3D.from_linestring(self._object_row.geometry) diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index b1325202..e16ccb10 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -18,5 +18,5 @@ class MapSurfaceType(SerialIntEnum): CARPARK = 5 GENERIC_DRIVABLE = 6 STOP_LINE = 7 - # ROAD_EDGE = 8 - # ROAD_LINE = 9 + ROAD_EDGE = 8 + ROAD_LINE = 9 diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index a98adec2..37db915f 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -16,8 +16,8 @@ defaults: - _self_ - datasets: # - nuplan_private_dataset - # - carla_dataset - - wopd_dataset + - carla_dataset + # - wopd_dataset -force_log_conversion: True -force_map_conversion: False +force_log_conversion: False +force_map_conversion: True diff --git a/notebooks/Untitled-2.ipynb b/notebooks/Untitled-2.ipynb new file mode 100644 index 00000000..d4de968a --- /dev/null +++ b/notebooks/Untitled-2.ipynb @@ -0,0 +1,286 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.maps.gpkg.gpkg_map import get_map_api_from_names\n", + "\n", + "map_api = get_map_api_from_names(\"carla\", \"town05\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.common.visualization.color.default import MAP_SURFACE_CONFIG\n", + "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "\n", + "from typing import List\n", + "from shapely import LineString, points\n", + "import shapely\n", + "from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import get_road_edge_linestrings\n", + "\n", + "import matplotlib.pyplot as plt\n", + "\n", + "from d123.common.visualization.color.color import BLACK, TAB_10\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", + "\n", + "lane_groups_ = map_api._gpd_dataframes[MapSurfaceType.LANE_GROUP].geometry.tolist()\n", + "car_parks_ = map_api._gpd_dataframes[MapSurfaceType.CARPARK].geometry.tolist()\n", + "generic_drivable_ = map_api._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE].geometry.tolist()\n", + "\n", + "\n", + "s = 30\n", + "fig, ax = plt.subplots(figsize=(s, s))\n", + "for polygon in lane_groups_:\n", + " add_shapely_polygon_to_ax(ax, polygon, MAP_SURFACE_CONFIG[MapSurfaceType.LANE_GROUP])\n", + "for polygon in car_parks_:\n", + " add_shapely_polygon_to_ax(ax, polygon, MAP_SURFACE_CONFIG[MapSurfaceType.CARPARK])\n", + "for polygon in generic_drivable_:\n", + " add_shapely_polygon_to_ax(ax, polygon, MAP_SURFACE_CONFIG[MapSurfaceType.WALKWAY])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "from shapely import union_all\n", + "from d123.common.geometry.occupancy_map import OccupancyMap2D\n", + "from d123.dataset.maps.abstract_map_objects import AbstractLane\n", + "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "\n", + "all_lane_group_ids = map_api._gpd_dataframes[MapSurfaceType.LANE_GROUP].id.tolist()\n", + "\n", + "all_lane_groups = [map_api.get_map_object(id, MapSurfaceType.LANE_GROUP) for id in all_lane_group_ids]\n", + "# lane_groups_no_intersection = [lane_group for lane_group in all_lane_groups if lane_group.intersection is None]\n", + "lane_groups_no_intersection = all_lane_groups\n", + "\n", + "occupancy_map = OccupancyMap2D.from_dict(\n", + " {lane_group.id: lane_group.shapely_polygon for lane_group in all_lane_groups}\n", + ")\n", + "\n", + "\n", + "def get_z_at_point(lane_group: AbstractLane, point: np.ndarray) -> float:\n", + " \"\"\"\n", + " Placeholder function to get the z-coordinate at a point.\n", + " This should be replaced with the actual implementation that retrieves the z-coordinate.\n", + " \"\"\"\n", + " centerlines = np.concatenate([lane.centerline.array for lane in lane_group.lanes], axis=0, dtype=np.float64)\n", + " distances_2d = np.linalg.norm(centerlines[..., :2] - point[:2], axis=-1)\n", + " closest_lane_index = np.argmin(distances_2d)\n", + " return centerlines[closest_lane_index, 2]\n", + "\n", + "\n", + "overlapping_lane_groups = {}\n", + "\n", + "for lane_group in lane_groups_no_intersection:\n", + " lane_centerlines = np.array([lane.centerline.array for lane in lane_group.lanes], dtype=np.float64)\n", + "\n", + " intersecting_lane_groups = occupancy_map.intersects(lane_group.shapely_polygon)\n", + " intersecting_lane_groups.remove(lane_group.id)\n", + " lane_group_polygon = lane_group.shapely_polygon\n", + "\n", + " for intersecting_id in intersecting_lane_groups:\n", + " intersecting_polygon = occupancy_map[intersecting_id]\n", + " intersection = intersecting_polygon.intersection(lane_group_polygon)\n", + " if intersection.geom_type != \"Polygon\":\n", + " continue\n", + "\n", + " intersecting_lane_group = map_api.get_map_object(intersecting_id, MapSurfaceType.LANE_GROUP)\n", + " z_at_query = get_z_at_point(lane_group, intersection.centroid.coords[0])\n", + " z_at_inter = get_z_at_point(intersecting_lane_group, intersection.centroid.coords[0])\n", + "\n", + " if np.abs(z_at_query - z_at_inter) < 5.0:\n", + " continue\n", + "\n", + " overlapping_lane_groups[lane_group.id] = lane_group\n", + " overlapping_lane_groups[intersecting_lane_group.id] = intersecting_lane_group\n", + "\n", + "# occupancy_map" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "non_overlapping_lane_groups = [\n", + " lane_group for lane_group in lane_groups_no_intersection if lane_group.id not in overlapping_lane_groups\n", + "]\n", + "\n", + "\n", + "def get_linestring_z_at_point(linestring: LineString, point: np.ndarray) -> float:\n", + " assert linestring.has_z, \"Linestring must have z-coordinates\"\n", + " linestring_coords = np.array(list(linestring.coords), dtype=np.float64)\n", + " distances_2d = np.linalg.norm(linestring_coords[:, :2] - point[:2], axis=-1)\n", + " closest_point_index = np.argmin(distances_2d)\n", + " return float(linestring_coords[closest_point_index, 2])\n", + "\n", + "\n", + "def lift_road_edges_to_3d(ax: plt.Axes, road_edges_2d: List[LineString]):\n", + " def _find_continuous_sublists(integers: List[int]) -> List[List[int]]:\n", + " arr = np.array(integers, dtype=np.int64)\n", + " breaks = np.where(np.diff(arr) != 1)[0] + 1\n", + " splits = np.split(arr, breaks)\n", + " return [sublist.tolist() for sublist in splits]\n", + "\n", + " boundaries = []\n", + " for lane_group in non_overlapping_lane_groups:\n", + " boundaries.append(lane_group.left_boundary.linestring)\n", + " boundaries.append(lane_group.right_boundary.linestring)\n", + "\n", + " boundaries.extend(\n", + " map_api._gpd_dataframes[MapSurfaceType.CARPARK].outline.tolist()\n", + " + map_api._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE].outline.tolist()\n", + " )\n", + "\n", + " occupancy_map = OccupancyMap2D(boundaries)\n", + "\n", + " road_edges_3d: List[LineString] = []\n", + " for idx, linestring in enumerate(road_edges_2d):\n", + " # print(list(linestring.coords))\n", + " points_3d = np.array(list(linestring.coords), dtype=np.float64)\n", + "\n", + " results = occupancy_map.query_nearest(\n", + " shapely.points(points_3d[..., :2]),\n", + " max_distance=0.01,\n", + " exclusive=False,\n", + " )\n", + " for query_idx, geometry_idx in zip(*results):\n", + " intersecting_boundary = occupancy_map[occupancy_map.ids[geometry_idx]]\n", + " points_3d[query_idx, 2] = get_linestring_z_at_point(intersecting_boundary, points_3d[query_idx])\n", + "\n", + " for continuous_slice in _find_continuous_sublists(results[0]):\n", + " if len(continuous_slice) < 2:\n", + " continue\n", + " lifted_linestring = LineString(points_3d[continuous_slice])\n", + " road_edges_3d.append(lifted_linestring)\n", + " return road_edges_3d\n", + "\n", + "\n", + "drivable_polygons = (\n", + " map_api._gpd_dataframes[MapSurfaceType.LANE_GROUP].geometry.tolist()\n", + " + map_api._gpd_dataframes[MapSurfaceType.CARPARK].geometry.tolist()\n", + " + map_api._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE].geometry.tolist()\n", + ")\n", + "\n", + "road_edges_2d = get_road_edge_linestrings(drivable_polygons, step_size=0.25, max_road_edge_length=None)\n", + "\n", + "\n", + "s = 30\n", + "fig, ax = plt.subplots(figsize=(s, s))\n", + "road_edges_3d = lift_road_edges_to_3d(ax, road_edges_2d)\n", + "\n", + "\n", + "\n", + "for lane_group in overlapping_lane_groups.values():\n", + " road_edges_3d.extend([lane_group.right_boundary.linestring, lane_group.left_boundary.linestring])\n", + "\n", + "\n", + "for idx, linestring in enumerate(road_edges_3d):\n", + " config = PlotConfig(\n", + " fill_color=TAB_10[idx % len(TAB_10)], fill_color_alpha=0.25, line_color=TAB_10[idx % len(TAB_10)]\n", + " )\n", + " # config = PlotConfig(fill_color=BLACK, fill_color_alpha=0.25, line_color=BLACK)\n", + " add_shapely_linestring_to_ax(ax, linestring, config)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "def find_continuous_sublists(arr):\n", + " arr = np.array(arr)\n", + " breaks = np.where(np.diff(arr) != 1)[0] + 1\n", + " splits = np.split(arr, breaks)\n", + " return [sublist.tolist() for sublist in splits]\n", + "\n", + "# Example usage\n", + "arr = [0, 1, 2, 3, 5, 6, 7, 9, 10, 22, 23, 24]\n", + "result = find_continuous_sublists(arr)\n", + "print(result) # [[0, 1, 2, 3], [5, 6, 7]]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "s = 30\n", + "\n", + "fig, ax = plt.subplots(figsize=(s, s))\n", + "\n", + "\n", + "for idx, lane_group in enumerate(overlapping_lane_groups.values()):\n", + " config = PlotConfig(fill_color=TAB_10[idx % len(TAB_10)], fill_color_alpha=0.25, line_color=TAB_10[idx % len(TAB_10)])\n", + " add_shapely_polygon_to_ax(ax, lane_group.shapely_polygon, config)\n", + "\n", + "# for idx, polygon in enumerate(intersections):\n", + "# config = PlotConfig(fill_color=BLACK, line_color=BLACK)\n", + "# add_shapely_polygon_to_ax(ax, polygon, config)\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 70595dac..16d0add2 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -25,8 +25,8 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "splits = [\"wopd_train\"]\n", - "# splits = [\"carla\"]\n", + "# splits = [\"wopd_train\"]\n", + "splits = [\"carla\"]\n", "# log_names = None\n", "\n", "\n", @@ -40,7 +40,7 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=19,\n", + " duration_s=10,\n", " history_s=0.0,\n", " timestamp_threshold_s=20,\n", " shuffle=False,\n", @@ -63,8 +63,9 @@ "source": [ "from typing import List, Optional, Tuple\n", "import matplotlib.pyplot as plt\n", + "import numpy as np\n", "from d123.common.geometry.base import Point2D\n", - "from d123.common.visualization.color.color import BLACK, LIGHT_GREY, TAB_10\n", + "from d123.common.visualization.color.color import BLACK, DARK_GREY, DARKER_GREY, LIGHT_GREY, NEW_TAB_10, TAB_10\n", "from d123.common.visualization.color.config import PlotConfig\n", "from d123.common.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", "from d123.common.visualization.matplotlib.observation import (\n", @@ -113,6 +114,27 @@ " zorder=1,\n", ")\n", "\n", + "ROAD_EDGE_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=DARKER_GREY.set_brightness(0.0),\n", + " fill_color_alpha=1.0,\n", + " line_color=DARKER_GREY.set_brightness(0.0),\n", + " line_color_alpha=1.0,\n", + " line_width=1.0,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "ROAD_LINE_CONFIG: PlotConfig = PlotConfig(\n", + " fill_color=DARKER_GREY,\n", + " fill_color_alpha=1.0,\n", + " line_color=NEW_TAB_10[5],\n", + " line_color_alpha=1.0,\n", + " line_width=1.5,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + ")\n", + "\n", + "\n", "\n", "def add_debug_map_on_ax(\n", " ax: plt.Axes,\n", @@ -123,17 +145,21 @@ ") -> None:\n", " layers: List[MapSurfaceType] = [\n", " MapSurfaceType.LANE,\n", + " MapSurfaceType.LANE_GROUP,\n", " MapSurfaceType.GENERIC_DRIVABLE,\n", " MapSurfaceType.CARPARK,\n", " MapSurfaceType.CROSSWALK,\n", " MapSurfaceType.INTERSECTION,\n", " MapSurfaceType.WALKWAY,\n", + " MapSurfaceType.ROAD_EDGE,\n", + " # MapSurfaceType.ROAD_LINE,\n", " ]\n", " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", " y_min, y_max = point_2d.y - radius, point_2d.y + radius\n", " patch = geom.box(x_min, y_min, x_max, y_max)\n", " map_objects_dict = map_api.query(geometry=patch, layers=layers, predicate=\"intersects\")\n", "\n", + " edge_lengths = []\n", " for layer, map_objects in map_objects_dict.items():\n", " for map_object in map_objects:\n", " try:\n", @@ -141,16 +167,20 @@ " MapSurfaceType.GENERIC_DRIVABLE,\n", " MapSurfaceType.CARPARK,\n", " # MapSurfaceType.CROSSWALK,\n", - " # MapSurfaceType.INTERSECTION,\n", + " MapSurfaceType.INTERSECTION,\n", " # MapSurfaceType.WALKWAY,\n", " ]:\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + "\n", + " if layer in [MapSurfaceType.LANE_GROUP]:\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + "\n", " if layer in [MapSurfaceType.LANE]:\n", " map_object: AbstractLane\n", " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", - " add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", - " add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", - " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", + " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", + " # add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", + " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", "\n", " # centroid = map_object.shapely_polygon.centroid\n", " # ax.text(\n", @@ -162,6 +192,12 @@ " # fontsize=8,\n", " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", " # )\n", + " if layer in [MapSurfaceType.ROAD_EDGE]:\n", + " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", + " edge_lengths.append(map_object.polyline_3d.linestring.length)\n", + "\n", + " # if layer in [MapSurfaceType.ROAD_LINE]:\n", + " # add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_LINE_CONFIG)\n", "\n", " except Exception:\n", " import traceback\n", @@ -169,6 +205,9 @@ " print(f\"Error adding map object of type {layer.name} and id {map_object.id}\")\n", " traceback.print_exc()\n", "\n", + " print(f\"road edge mean: {np.mean(edge_lengths)} m\")\n", + " print(f\"road edge min: {np.min(edge_lengths)} m\")\n", + " print(f\"road edge max: {np.max(edge_lengths)} m\")\n", " ax.set_title(f\"Map: {map_api.map_name}\")\n", "\n", "\n", @@ -185,8 +224,9 @@ " add_box_detections_to_ax(ax, box_detections)\n", " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", "\n", - " ax.set_xlim(point_2d.x - radius, point_2d.x + radius)\n", - " ax.set_ylim(point_2d.y - radius, point_2d.y + radius)\n", + " zoom = 1.0\n", + " ax.set_xlim(point_2d.x - radius * zoom, point_2d.x + radius * zoom)\n", + " ax.set_ylim(point_2d.y - radius * zoom, point_2d.y + radius * zoom)\n", "\n", " ax.set_aspect(\"equal\", adjustable=\"box\")\n", " return ax\n", @@ -196,14 +236,15 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " size = 16\n", + " size = 10\n", + "\n", " fig, ax = plt.subplots(figsize=(size, size))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", - "scene_index = 5\n", - "plot_scene_at_iteration(scenes[scene_index], iteration=60, radius=50)" + "scene_index = 14\n", + "plot_scene_at_iteration(scenes[scene_index], iteration=1, radius=82)" ] }, { diff --git a/notebooks/viz/viser_testing_v2_map.ipynb b/notebooks/viz/viser_testing_v2_map.ipynb deleted file mode 100644 index 38eaaaee..00000000 --- a/notebooks/viz/viser_testing_v2_map.ipynb +++ /dev/null @@ -1,116 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", - "from d123.dataset.scene.scene_filter import SceneFilter\n", - "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "from d123.common.datatypes.sensor.camera import CameraType" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# split = \"nuplan_private_test\"\n", - "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", - " \n", - "\n", - "splits = [\"wopd_train\"]\n", - "# log_names = None\n", - "\n", - "\n", - "\n", - "# splits = [\"nuplan_private_test\"]\n", - "# splits = [\"carla\"]\n", - "# log_names = [\"11839652018869852123_2565_000_2585_000\"]\n", - "log_names = None\n", - "\n", - "scene_tokens = None\n", - "\n", - "scene_filter = SceneFilter(\n", - " split_names=splits,\n", - " log_names=log_names,\n", - " scene_tokens=scene_tokens,\n", - " duration_s=19,\n", - " history_s=0.0,\n", - " timestamp_threshold_s=20,\n", - " shuffle=False,\n", - " camera_types=[CameraType.CAM_F0],\n", - ")\n", - "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", - "worker = Sequential()\n", - "# worker = RayDistributed()\n", - "scenes = scene_builder.get_scenes(scene_filter, worker)\n", - "\n", - "print(f\"Found {len(scenes)} scenes\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "from d123.common.visualization.viser.server_map import ViserMapVisualizationServer\n", - "\n", - "\n", - "visualization_server = ViserMapVisualizationServer(scenes, scene_index=4)" - ] - }, - { - "cell_type": "markdown", - "id": "3", - "metadata": {}, - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "d123", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 1e24c584..75f9de83 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -25,12 +25,12 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", " \n", "\n", - "splits = [\"wopd_train\"]\n", + "# splits = [\"wopd_train\"]\n", "# log_names = None\n", "\n", "\n", "\n", - "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"nuplan_private_test\"]\n", "# splits = [\"carla\"]\n", "# log_names = [\"11839652018869852123_2565_000_2585_000\"]\n", "log_names = None\n", @@ -44,7 +44,7 @@ " duration_s=19,\n", " history_s=0.0,\n", " timestamp_threshold_s=20,\n", - " shuffle=False,\n", + " shuffle=True,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", @@ -65,185 +65,27 @@ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes, scene_index=4)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" ] }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [] + }, { "cell_type": "code", "execution_count": null, - "id": "3", + "id": "4", "metadata": {}, "outputs": [], - "source": [ - "from d123.common.visualization.color.color import Color\n", - "\n", - "\n", - "# def adjust_brightness(hex_code, brightness_factor):\n", - "# \"\"\"\n", - "# Adjust the brightness of a hex color code.\n", - "\n", - "# Args:\n", - "# hex_code (str): Hex color code (with or without #)\n", - "# brightness_factor (float): Factor to adjust brightness by\n", - "# - Values > 1.0 make the color brighter\n", - "# - Values < 1.0 make the color darker\n", - "# - 1.0 leaves the color unchanged\n", - "# - 0.0 makes the color black\n", - "\n", - "# Returns:\n", - "# str: Adjusted hex color code with # prefix\n", - "\n", - "# Example:\n", - "# adjust_brightness(\"#FF5733\", 1.5) # Make 50% brighter\n", - "# adjust_brightness(\"FF5733\", 0.7) # Make 30% darker\n", - "# \"\"\"\n", - "# # Remove # if present and convert to uppercase\n", - "# hex_code = hex_code.lstrip(\"#\").upper()\n", - "\n", - "# # Validate hex code\n", - "# if len(hex_code) != 6:\n", - "# raise ValueError(\"Hex code must be 6 characters long\")\n", - "\n", - "# try:\n", - "# # Extract RGB components\n", - "# r = int(hex_code[0:2], 16)\n", - "# g = int(hex_code[2:4], 16)\n", - "# b = int(hex_code[4:6], 16)\n", - "# except ValueError:\n", - "# raise ValueError(\"Invalid hex code format\")\n", - "\n", - "# # Adjust brightness\n", - "# r = min(255, max(0, int(r * brightness_factor)))\n", - "# g = min(255, max(0, int(g * brightness_factor)))\n", - "# b = min(255, max(0, int(b * brightness_factor)))\n", - "\n", - "# # Convert back to hex\n", - "# return f\"#{r:02X}{g:02X}{b:02X}\"\n", - "\n", - "\n", - "# def adjust_brightness_hsl(hex_code, brightness_change):\n", - "# \"\"\"\n", - "# Adjust brightness using HSL color space for more natural results.\n", - "\n", - "# Args:\n", - "# hex_code (str): Hex color code (with or without #)\n", - "# brightness_change (float): Amount to change brightness by\n", - "# - Positive values make brighter\n", - "# - Negative values make darker\n", - "# - Range: -1.0 to 1.0\n", - "\n", - "# Returns:\n", - "# str: Adjusted hex color code with # prefix\n", - "# \"\"\"\n", - "# # Remove # if present\n", - "# hex_code = hex_code.lstrip(\"#\")\n", - "\n", - "# # Convert hex to RGB\n", - "# r = int(hex_code[0:2], 16) / 255.0\n", - "# g = int(hex_code[2:4], 16) / 255.0\n", - "# b = int(hex_code[4:6], 16) / 255.0\n", - "\n", - "# # Convert RGB to HSL\n", - "# max_val = max(r, g, b)\n", - "# min_val = min(r, g, b)\n", - "# diff = max_val - min_val\n", - "\n", - "# # Calculate lightness\n", - "# lightness = (max_val + min_val) / 2\n", - "\n", - "# if diff == 0:\n", - "# hue = saturation = 0\n", - "# else:\n", - "# # Calculate saturation\n", - "# if lightness < 0.5:\n", - "# saturation = diff / (max_val + min_val)\n", - "# else:\n", - "# saturation = diff / (2 - max_val - min_val)\n", - "\n", - "# # Calculate hue\n", - "# if max_val == r:\n", - "# hue = ((g - b) / diff + (6 if g < b else 0)) / 6\n", - "# elif max_val == g:\n", - "# hue = ((b - r) / diff + 2) / 6\n", - "# else:\n", - "# hue = ((r - g) / diff + 4) / 6\n", - "\n", - "# # Adjust lightness\n", - "# lightness = max(0, min(1, lightness + brightness_change))\n", - "\n", - "# # Convert HSL back to RGB\n", - "# if saturation == 0:\n", - "# r = g = b = lightness\n", - "# else:\n", - "\n", - "# def hue_to_rgb(p, q, t):\n", - "# if t < 0:\n", - "# t += 1\n", - "# if t > 1:\n", - "# t -= 1\n", - "# if t < 1 / 6:\n", - "# return p + (q - p) * 6 * t\n", - "# if t < 1 / 2:\n", - "# return q\n", - "# if t < 2 / 3:\n", - "# return p + (q - p) * (2 / 3 - t) * 6\n", - "# return p\n", - "\n", - "# q = lightness * (1 + saturation) if lightness < 0.5 else lightness + saturation - lightness * saturation\n", - "# p = 2 * lightness - q\n", - "\n", - "# r = hue_to_rgb(p, q, hue + 1 / 3)\n", - "# g = hue_to_rgb(p, q, hue)\n", - "# b = hue_to_rgb(p, q, hue - 1 / 3)\n", - "\n", - "# # Convert back to hex\n", - "# r = int(round(r * 255))\n", - "# g = int(round(g * 255))\n", - "# b = int(round(b * 255))\n", - "\n", - "# return f\"#{r:02X}{g:02X}{b:02X}\"\n", - "\n", - "\n", - "# Example usage and testing\n", - "if __name__ == \"__main__\":\n", - " # Test the functions\n", - "\n", - "\n", - " # Test with different colors\n", - " test_colors = [\n", - " Color(\"#3498DB\"),\n", - " Color(\"#E74C3C\"),\n", - " Color(\"#2ECC71\"),\n", - " Color(\"#F39C12\"),\n", - " ]\n", - "\n", - " import matplotlib.pyplot as plt\n", - "\n", - "\n", - " # Visualize the color adjustments\n", - " factor = 0.5\n", - " fig, ax = plt.subplots(len(test_colors), 3, figsize=(8, 12))\n", - " for i, color in enumerate(test_colors):\n", - " ax[i, 0].imshow([[color.rgb]])\n", - " ax[i, 0].set_title(f\"Original: {color}\")\n", - " ax[i, 0].axis(\"off\")\n", - "\n", - " ax[i, 1].imshow([[color.set_brightness(factor).rgb]])\n", - " ax[i, 1].set_title(f\"Brighter: {str(color.set_brightness(factor))}\")\n", - " ax[i, 1].axis(\"off\")\n", - "\n", - " ax[i, 2].imshow([[color.set_brightness(1 / factor).rgb]])\n", - " ax[i, 2].set_title(f\"Darker: {str(color.set_brightness(1 / factor))}\")\n", - " ax[i, 2].axis(\"off\")\n", - "\n", - " plt.tight_layout()\n", - " plt.show()" - ] + "source": [] }, { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [] From d8f1438890036497525fb0e5acd98151dc4c2287 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sat, 9 Aug 2025 19:33:53 +0200 Subject: [PATCH 25/42] Improve matplotlib function for rendering polygons (with transparent interiors.). --- d123/common/visualization/color/config.py | 2 - d123/common/visualization/matplotlib/utils.py | 67 +++++++++---------- 2 files changed, 31 insertions(+), 38 deletions(-) diff --git a/d123/common/visualization/color/config.py b/d123/common/visualization/color/config.py index ccc7ecbb..6d610844 100644 --- a/d123/common/visualization/color/config.py +++ b/d123/common/visualization/color/config.py @@ -18,5 +18,3 @@ class PlotConfig: zorder: int = 0 smoothing_radius: Optional[float] = None - - shadow: bool = False diff --git a/d123/common/visualization/matplotlib/utils.py b/d123/common/visualization/matplotlib/utils.py index 1af6f9f5..144530ab 100644 --- a/d123/common/visualization/matplotlib/utils.py +++ b/d123/common/visualization/matplotlib/utils.py @@ -1,13 +1,14 @@ from functools import lru_cache from typing import Union +import matplotlib.patches as patches import matplotlib.pyplot as plt import numpy as np import shapely.affinity as affinity import shapely.geometry as geom +from matplotlib.path import Path from d123.common.geometry.base import StateSE2, StateSE3 -from d123.common.visualization.color.color import WHITE from d123.common.visualization.color.config import PlotConfig @@ -18,54 +19,48 @@ def add_shapely_polygon_to_ax( disable_smoothing: bool = False, ) -> plt.Axes: """ - Adds shapely polygon to birds-eye-view visualization + Adds shapely polygon to birds-eye-view visualization with proper hole handling :param ax: matplotlib ax object :param polygon: shapely Polygon - :param config: dictionary containing plot parameters + :param plot_config: dictionary containing plot parameters + :param disable_smoothing: whether to overwrite smoothing of the polygon :return: ax with plot """ def _add_element_helper(element: geom.Polygon): - """Helper to add single polygon to ax""" + """Helper to add single polygon to ax with proper holes""" if plot_config.smoothing_radius is not None and not disable_smoothing: element = element.buffer(-plot_config.smoothing_radius).buffer(plot_config.smoothing_radius) - exterior_x, exterior_y = element.exterior.xy - - if plot_config.shadow: - shadow_offset = 0.5 - shadow_x = [x + shadow_offset for x in exterior_x] - shadow_y = [y - shadow_offset for y in exterior_y] - ax.fill( - shadow_x, - shadow_y, - color="gray", - alpha=1.0, - edgecolor=None, - linewidth=0, - zorder=plot_config.zorder, - ) - - ax.fill( - exterior_x, - exterior_y, - color=plot_config.fill_color.hex, + + # Create path with exterior and interior rings + def create_polygon_path(polygon): + # Get exterior coordinates + exterior_coords = list(polygon.exterior.coords) + + # Start with exterior ring + vertices = exterior_coords + codes = [Path.MOVETO] + [Path.LINETO] * (len(exterior_coords) - 2) + [Path.CLOSEPOLY] + + # Add interior rings (holes) + for interior in polygon.interiors: + interior_coords = list(interior.coords) + vertices.extend(interior_coords) + codes.extend([Path.MOVETO] + [Path.LINETO] * (len(interior_coords) - 2) + [Path.CLOSEPOLY]) + + return Path(vertices, codes) + + path = create_polygon_path(element) + + # Add main polygon with holes + patch = patches.PathPatch( + path, + facecolor=plot_config.fill_color.hex, alpha=plot_config.fill_color_alpha, edgecolor=plot_config.line_color.hex, linewidth=plot_config.line_width, zorder=plot_config.zorder, ) - - for interior in element.interiors: - x_interior, y_interior = interior.xy - ax.fill( - x_interior, - y_interior, - color=WHITE.hex, - alpha=plot_config.fill_color_alpha, - edgecolor=plot_config.line_color.hex, - linewidth=plot_config.line_width, - zorder=plot_config.zorder, - ) + ax.add_patch(patch) if isinstance(polygon, geom.Polygon): _add_element_helper(polygon) From 5cde0aca6f0355a972a8a0bee83d61e17778d8fa Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 10 Aug 2025 14:23:07 +0200 Subject: [PATCH 26/42] Refactor and improve road edge extraction for opendrive (3D) and nuplan (2D) --- d123/common/visualization/viser/utils.py | 17 +- d123/dataset/conversion/__init__.py | 0 d123/dataset/conversion/map/__init__.py | 0 .../conversion/map/opendrive/__init__.py | 0 .../opendrive/conversion/group_collections.py | 8 +- .../map}/opendrive/conversion/id_system.py | 0 .../conversion/objects_collections.py | 4 +- .../map}/opendrive/elements/elevation.py | 0 .../map}/opendrive/elements/enums.py | 0 .../map}/opendrive/elements/geometry.py | 0 .../map}/opendrive/elements/lane.py | 0 .../map}/opendrive/elements/objects.py | 0 .../map}/opendrive/elements/opendrive.py | 2 +- .../map}/opendrive/elements/reference.py | 6 +- .../map}/opendrive/elements/road.py | 8 +- .../map}/opendrive/id_mapping.py | 0 .../map}/opendrive/opendrive_converter.py | 276 +++-------------- .../conversion/map/road_edge/__init__.py | 0 .../map/road_edge/road_edge_2d_utils.py | 60 ++++ .../map/road_edge/road_edge_3d_utils.py | 279 ++++++++++++++++++ .../carla/carla_data_converter.py | 4 +- .../nuplan/nuplan_map_conversion.py | 106 ++----- d123/dataset/maps/gpkg/gpkg_map_objects.py | 36 ++- .../default_dataset_conversion.yaml | 4 +- .../test_nuplan_conversion.ipynb | 31 +- .../test_opendrive_conversion.ipynb | 8 +- notebooks/viz/bev_matplotlib.ipynb | 19 +- notebooks/viz/viser_testing_v2_scene.ipynb | 4 +- scripts/dataset/run_log_caching.sh | 2 +- 29 files changed, 472 insertions(+), 402 deletions(-) create mode 100644 d123/dataset/conversion/__init__.py create mode 100644 d123/dataset/conversion/map/__init__.py create mode 100644 d123/dataset/conversion/map/opendrive/__init__.py rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/conversion/group_collections.py (96%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/conversion/id_system.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/conversion/objects_collections.py (94%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/elevation.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/enums.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/geometry.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/lane.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/objects.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/opendrive.py (98%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/reference.py (96%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/elements/road.py (92%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/id_mapping.py (100%) rename d123/dataset/{dataset_specific/carla => conversion/map}/opendrive/opendrive_converter.py (68%) create mode 100644 d123/dataset/conversion/map/road_edge/__init__.py create mode 100644 d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py create mode 100644 d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index a2981153..151e820c 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -21,7 +21,7 @@ # TODO: Refactor this file. # TODO: Add general utilities for 3D primitives and mesh support. -MAP_RADIUS: Final[float] = 500 +MAP_RADIUS: Final[float] = 80 BRIGHTNESS_FACTOR: Final[float] = 0.8 @@ -106,7 +106,16 @@ def get_map_meshes(scene: AbstractScene): for map_surface in map_objects_dict[map_surface_type]: map_surface: AbstractSurfaceMapObject trimesh_mesh = map_surface.trimesh_mesh - if map_surface_type in [MapSurfaceType.WALKWAY, MapSurfaceType.CROSSWALK]: + if map_surface_type == MapSurfaceType.GENERIC_DRIVABLE: + print("Generic Drivable Surface Mesh:") + print(trimesh_mesh) + output[f"{map_surface_type.serialize()}_{map_surface.id}"] = trimesh_mesh + if map_surface_type in [ + MapSurfaceType.WALKWAY, + MapSurfaceType.CROSSWALK, + # MapSurfaceType.GENERIC_DRIVABLE, + # MapSurfaceType.CARPARK, + ]: # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.05).array else: @@ -117,9 +126,7 @@ def get_map_meshes(scene: AbstractScene): x=0, y=0, z=center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 ).array - trimesh_mesh.visual.face_colors = ( - MAP_SURFACE_CONFIG[map_surface_type].fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba - ) + trimesh_mesh.visual.face_colors = MAP_SURFACE_CONFIG[map_surface_type].fill_color.rgba surface_meshes.append(trimesh_mesh) output[f"{map_surface_type.serialize()}"] = trimesh.util.concatenate(surface_meshes) diff --git a/d123/dataset/conversion/__init__.py b/d123/dataset/conversion/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/conversion/map/__init__.py b/d123/dataset/conversion/map/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/conversion/map/opendrive/__init__.py b/d123/dataset/conversion/map/opendrive/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py b/d123/dataset/conversion/map/opendrive/conversion/group_collections.py similarity index 96% rename from d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py rename to d123/dataset/conversion/map/opendrive/conversion/group_collections.py index e5082014..5e509859 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/conversion/group_collections.py +++ b/d123/dataset/conversion/map/opendrive/conversion/group_collections.py @@ -8,14 +8,14 @@ from d123.common.geometry.base import StateSE2Index from d123.common.geometry.units import kmph_to_mps, mph_to_mps -from d123.dataset.dataset_specific.carla.opendrive.conversion.id_system import ( +from d123.dataset.conversion.map.opendrive.conversion.id_system import ( derive_lane_group_id, derive_lane_id, lane_group_id_from_lane_id, ) -from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lane, LaneSection -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border -from d123.dataset.dataset_specific.carla.opendrive.elements.road import RoadType +from d123.dataset.conversion.map.opendrive.elements.lane import Lane, LaneSection +from d123.dataset.conversion.map.opendrive.elements.reference import Border +from d123.dataset.conversion.map.opendrive.elements.road import RoadType # TODO: Add to config STEP_SIZE = 1.0 diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/id_system.py b/d123/dataset/conversion/map/opendrive/conversion/id_system.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/conversion/id_system.py rename to d123/dataset/conversion/map/opendrive/conversion/id_system.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/conversion/objects_collections.py b/d123/dataset/conversion/map/opendrive/conversion/objects_collections.py similarity index 94% rename from d123/dataset/dataset_specific/carla/opendrive/conversion/objects_collections.py rename to d123/dataset/conversion/map/opendrive/conversion/objects_collections.py index d94d3eaf..1eb52f77 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/conversion/objects_collections.py +++ b/d123/dataset/conversion/map/opendrive/conversion/objects_collections.py @@ -8,8 +8,8 @@ from d123.common.geometry.base import Point2D, Point3D, Point3DIndex, StateSE2 from d123.common.geometry.transform.tranform_2d import translate_along_yaw from d123.common.geometry.utils import normalize_angle -from d123.dataset.dataset_specific.carla.opendrive.elements.objects import Object -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border +from d123.dataset.conversion.map.opendrive.elements.objects import Object +from d123.dataset.conversion.map.opendrive.elements.reference import Border # TODO: make naming consistent with group_collections.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/elevation.py b/d123/dataset/conversion/map/opendrive/elements/elevation.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/elements/elevation.py rename to d123/dataset/conversion/map/opendrive/elements/elevation.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/enums.py b/d123/dataset/conversion/map/opendrive/elements/enums.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/elements/enums.py rename to d123/dataset/conversion/map/opendrive/elements/enums.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/geometry.py b/d123/dataset/conversion/map/opendrive/elements/geometry.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/elements/geometry.py rename to d123/dataset/conversion/map/opendrive/elements/geometry.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/lane.py b/d123/dataset/conversion/map/opendrive/elements/lane.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/elements/lane.py rename to d123/dataset/conversion/map/opendrive/elements/lane.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/objects.py b/d123/dataset/conversion/map/opendrive/elements/objects.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/elements/objects.py rename to d123/dataset/conversion/map/opendrive/elements/objects.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/opendrive.py b/d123/dataset/conversion/map/opendrive/elements/opendrive.py similarity index 98% rename from d123/dataset/dataset_specific/carla/opendrive/elements/opendrive.py rename to d123/dataset/conversion/map/opendrive/elements/opendrive.py index f06876ef..de73ce87 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/opendrive.py +++ b/d123/dataset/conversion/map/opendrive/elements/opendrive.py @@ -6,7 +6,7 @@ from typing import List, Optional from xml.etree.ElementTree import Element, parse -from d123.dataset.dataset_specific.carla.opendrive.elements.road import Road +from d123.dataset.conversion.map.opendrive.elements.road import Road @dataclass diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/reference.py b/d123/dataset/conversion/map/opendrive/elements/reference.py similarity index 96% rename from d123/dataset/dataset_specific/carla/opendrive/elements/reference.py rename to d123/dataset/conversion/map/opendrive/elements/reference.py index ebca6ecb..cb521d4e 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/reference.py +++ b/d123/dataset/conversion/map/opendrive/elements/reference.py @@ -11,9 +11,9 @@ from d123.common.geometry.base import Point3DIndex, StateSE2Index from d123.common.geometry.utils import normalize_angle -from d123.dataset.dataset_specific.carla.opendrive.elements.elevation import ElevationProfile -from d123.dataset.dataset_specific.carla.opendrive.elements.geometry import Arc, Geometry, Line, Spiral -from d123.dataset.dataset_specific.carla.opendrive.elements.lane import LaneOffset +from d123.dataset.conversion.map.opendrive.elements.elevation import ElevationProfile +from d123.dataset.conversion.map.opendrive.elements.geometry import Arc, Geometry, Line, Spiral +from d123.dataset.conversion.map.opendrive.elements.lane import LaneOffset @dataclass diff --git a/d123/dataset/dataset_specific/carla/opendrive/elements/road.py b/d123/dataset/conversion/map/opendrive/elements/road.py similarity index 92% rename from d123/dataset/dataset_specific/carla/opendrive/elements/road.py rename to d123/dataset/conversion/map/opendrive/elements/road.py index b22c77ed..cdef3472 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/elements/road.py +++ b/d123/dataset/conversion/map/opendrive/elements/road.py @@ -4,10 +4,10 @@ from typing import List, Optional from xml.etree.ElementTree import Element -from d123.dataset.dataset_specific.carla.opendrive.elements.elevation import ElevationProfile, LateralProfile -from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lanes -from d123.dataset.dataset_specific.carla.opendrive.elements.objects import Object -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import PlanView +from d123.dataset.conversion.map.opendrive.elements.elevation import ElevationProfile, LateralProfile +from d123.dataset.conversion.map.opendrive.elements.lane import Lanes +from d123.dataset.conversion.map.opendrive.elements.objects import Object +from d123.dataset.conversion.map.opendrive.elements.reference import PlanView @dataclass diff --git a/d123/dataset/dataset_specific/carla/opendrive/id_mapping.py b/d123/dataset/conversion/map/opendrive/id_mapping.py similarity index 100% rename from d123/dataset/dataset_specific/carla/opendrive/id_mapping.py rename to d123/dataset/conversion/map/opendrive/id_mapping.py diff --git a/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py b/d123/dataset/conversion/map/opendrive/opendrive_converter.py similarity index 68% rename from d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py rename to d123/dataset/conversion/map/opendrive/opendrive_converter.py index d3e8b806..3ad5bea3 100644 --- a/d123/dataset/dataset_specific/carla/opendrive/opendrive_converter.py +++ b/d123/dataset/conversion/map/opendrive/opendrive_converter.py @@ -1,44 +1,42 @@ import os import warnings -from collections import defaultdict from pathlib import Path -from typing import Dict, List, Optional, Set +from typing import Dict, Final, List, Optional import geopandas as gpd -import networkx as nx import numpy as np -import numpy.typing as npt import pandas as pd import shapely from shapely.ops import polygonize, unary_union -from d123.common.geometry.base import Point3DIndex -from d123.common.geometry.occupancy_map import OccupancyMap2D -from d123.dataset.dataset_specific.carla.opendrive.conversion.group_collections import ( +from d123.dataset.conversion.map.opendrive.conversion.group_collections import ( OpenDriveLaneGroupHelper, OpenDriveLaneHelper, lane_section_to_lane_helpers, ) -from d123.dataset.dataset_specific.carla.opendrive.conversion.id_system import ( +from d123.dataset.conversion.map.opendrive.conversion.id_system import ( build_lane_id, derive_lane_section_id, lane_group_id_from_lane_id, road_id_from_lane_group_id, ) -from d123.dataset.dataset_specific.carla.opendrive.conversion.objects_collections import ( +from d123.dataset.conversion.map.opendrive.conversion.objects_collections import ( OpenDriveObjectHelper, get_object_helper, ) -from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import Junction, OpenDrive, Road -from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border -from d123.dataset.dataset_specific.carla.opendrive.id_mapping import IntIDMapping -from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import get_road_edge_linestrings +from d123.dataset.conversion.map.opendrive.elements.opendrive import Junction, OpenDrive, Road +from d123.dataset.conversion.map.opendrive.elements.reference import Border +from d123.dataset.conversion.map.opendrive.id_mapping import IntIDMapping +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length +from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import get_road_edges_3d_from_gdf from d123.dataset.maps.map_datatypes import MapSurfaceType ENABLE_WARNING: bool = False CONNECTION_DISTANCE_THRESHOLD: float = 0.1 # [m] D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # [m] + class OpenDriveConverter: def __init__(self, opendrive: OpenDrive): @@ -65,41 +63,37 @@ def run(self, map_name: str) -> None: self._collect_crosswalks() # Collect data frames and store - lane_df = self._extract_lane_dataframe() - walkways_df = self._extract_walkways_dataframe() - carpark_df = self._extract_carpark_dataframe() - generic_drivable_area_df = self._extract_generic_drivable_dataframe() - intersections_df = self._extract_intersections_dataframe() - lane_group_df = self._extract_lane_group_dataframe() - crosswalk_df = self._extract_crosswalk_dataframe() + dataframes: Dict[MapSurfaceType, gpd.GeoDataFrame] = {} + dataframes[MapSurfaceType.LANE] = self._extract_lane_dataframe() + dataframes[MapSurfaceType.WALKWAY] = self._extract_walkways_dataframe() + dataframes[MapSurfaceType.CARPARK] = self._extract_carpark_dataframe() + dataframes[MapSurfaceType.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() + dataframes[MapSurfaceType.INTERSECTION] = self._extract_intersections_dataframe() + dataframes[MapSurfaceType.LANE_GROUP] = self._extract_lane_group_dataframe() + dataframes[MapSurfaceType.CROSSWALK] = self._extract_crosswalk_dataframe() self._convert_ids_to_int( - lane_df, - walkways_df, - carpark_df, - generic_drivable_area_df, - lane_group_df, - intersections_df, - crosswalk_df, + dataframes[MapSurfaceType.LANE], + dataframes[MapSurfaceType.WALKWAY], + dataframes[MapSurfaceType.CARPARK], + dataframes[MapSurfaceType.GENERIC_DRIVABLE], + dataframes[MapSurfaceType.LANE_GROUP], + dataframes[MapSurfaceType.INTERSECTION], + dataframes[MapSurfaceType.CROSSWALK], + ) + dataframes[MapSurfaceType.ROAD_EDGE] = self._extract_road_edge_df( + dataframes[MapSurfaceType.LANE], + dataframes[MapSurfaceType.CARPARK], + dataframes[MapSurfaceType.GENERIC_DRIVABLE], + dataframes[MapSurfaceType.LANE_GROUP], ) - - road_edge_df = self._extract_road_edge_df(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) # Store dataframes map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" - lane_df.to_file(map_file_name, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") - walkways_df.to_file(map_file_name, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_name, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_area_df.to_file( - map_file_name, - layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), - driver="GPKG", - mode="a", - ) - intersections_df.to_file(map_file_name, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") - lane_group_df.to_file(map_file_name, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_name, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") - road_edge_df.to_file(map_file_name, layer=MapSurfaceType.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_name, layer=layer.serialize(), driver="GPKG", mode="a") def _collect_lane_helpers(self) -> None: for road in self.opendrive.roads: @@ -389,8 +383,8 @@ def _extract_walkways_dataframe(self) -> gpd.GeoDataFrame: data = pd.DataFrame( { "id": ids, - # "left_boundary": left_boundaries, - # "right_boundary": left_boundaries, + "left_boundary": left_boundaries, + "right_boundary": left_boundaries, "outline": outlines, } ) @@ -570,7 +564,8 @@ def _extract_road_edge_df( generic_drivable_area_df: gpd.GeoDataFrame, lane_group_df: gpd.GeoDataFrame, ) -> None: - road_edges = _get_road_edges_from_gdf(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) + road_edges = get_road_edges_3d_from_gdf(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) + road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) ids = np.arange(len(road_edges), dtype=np.int64).tolist() geometries = road_edges @@ -603,194 +598,3 @@ def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper] else: # Take the largest polygon if there are multiple return max(polygons, key=lambda p: p.area) - - -def _get_road_edges_from_gdf( - lane_df: gpd.GeoDataFrame, - carpark_df: gpd.GeoDataFrame, - generic_drivable_area_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, -) -> List[shapely.LineString]: - - # 1. Find conflicting lane groups, e.g. groups of lanes that overlap in 2D but have different Z-values (bridges) - conflicting_lane_groups = _get_conflicting_lane_groups(lane_group_df, lane_df) - - # 2. Extract road edges in 2D (including conflicting lane groups) - drivable_polygons = ( - lane_group_df.geometry.tolist() + carpark_df.geometry.tolist() + generic_drivable_area_df.geometry.tolist() - ) - road_edges_2d = get_road_edge_linestrings(drivable_polygons, max_road_edge_length=None) - - # 3. Collect 3D boundaries of non-conflicting lane groups and other drivable areas - non_conflicting_boundaries: List[shapely.LineString] = [] - for lane_group_id, lane_group_helper in lane_group_df.iterrows(): - if lane_group_id not in conflicting_lane_groups.keys(): - non_conflicting_boundaries.append(lane_group_helper["left_boundary"]) - non_conflicting_boundaries.append(lane_group_helper["right_boundary"]) - for outline in carpark_df.outline.tolist() + generic_drivable_area_df.outline.tolist(): - non_conflicting_boundaries.append(outline) - - # 4. Lift road edges to 3D using the boundaries of non-conflicting elements - non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, non_conflicting_boundaries) - - # 5. Add road edges from conflicting lane groups - resolved_road_edges = _resolve_conflicting_lane_groups(conflicting_lane_groups, lane_group_df) - - all_road_edges = non_conflicting_road_edges + resolved_road_edges - - return all_road_edges - - -def _get_nearest_z_from_points_3d(points_3d: npt.NDArray[np.float64], query_point: npt.NDArray[np.float64]) -> float: - assert points_3d.ndim == 2 and points_3d.shape[1] == len( - Point3DIndex - ), "points_3d must be a 2D array with shape (N, 3)" - distances = np.linalg.norm(points_3d[..., Point3DIndex.XY] - query_point[..., Point3DIndex.XY], axis=1) - closest_point = points_3d[np.argmin(distances)] - return closest_point[2] - - -def _get_conflicting_lane_groups(lane_group_df: gpd.GeoDataFrame, lane_df: gpd.GeoDataFrame) -> Dict[int, List[int]]: - Z_THRESHOLD = 5.0 # [m] Z-value threshold for conflict detection - - ids = lane_group_df.id.tolist() - polygons = lane_group_df.geometry.tolist() - - def _get_centerline_points_3d(lane_group_id: int) -> npt.NDArray[np.float64]: - """Helper function to get the centerline points in 3D.""" - lane_ids = lane_group_df[lane_group_df.id == lane_group_id].lane_ids.values[0] - centerlines: List[npt.NDArray[np.float64]] = [] - for lane_id in lane_ids: - centerline = lane_df[lane_df.id == lane_id].baseline_path.values[0] - assert isinstance(centerline, shapely.LineString) - centerlines.append(np.array(centerline.coords, dtype=np.float64)) - return np.concatenate(centerlines, axis=0) - - occupancy_map = OccupancyMap2D(polygons, ids) - conflicting_lane_groups: Dict[int, List[int]] = defaultdict(list) - - for lane_group_id, lane_group_polygon in zip(ids, polygons): - - # Extract internal centerline points for Z-value check - lane_group_centerlines = _get_centerline_points_3d(lane_group_id) - - # Query lane groups that overlap in 2D - intersecting_lane_group_ids = occupancy_map.intersects(lane_group_polygon) - intersecting_lane_group_ids.remove(lane_group_id) # Remove self from the list - for intersecting_id in intersecting_lane_group_ids: - intersecting_geometry = occupancy_map[intersecting_id] - - # ignore non-polygon geometries - if intersecting_geometry.geom_type != "Polygon": - continue - - # Check if Z-values deviate at intersection centroid - intersection_centroid = np.array(intersecting_geometry.centroid.coords[0], dtype=np.float64) - intersecting_centerlines = _get_centerline_points_3d(intersecting_id) - z_at_intersecting = _get_nearest_z_from_points_3d(intersecting_centerlines, intersection_centroid) - z_at_lane_group = _get_nearest_z_from_points_3d(lane_group_centerlines, intersection_centroid) - if np.abs(z_at_lane_group - z_at_intersecting) < Z_THRESHOLD: - continue - conflicting_lane_groups[lane_group_id].append(intersecting_id) - return conflicting_lane_groups - - -def _resolve_conflicting_lane_groups( - conflicting_lane_groups: Dict[int, List[int]], lane_group_df: gpd.GeoDataFrame -) -> List[shapely.LineString]: - - # Split conflicting lane groups into non-conflicting sets for further merging - non_conflicting_sets = create_non_conflicting_sets(conflicting_lane_groups) - - road_edges_3d: List[shapely.LineString] = [] - for non_conflicting_set in non_conflicting_sets: - - # Collect 2D polygons of non-conflicting lane group set - set_polygons = [ - lane_group_df[lane_group_df.id == lane_group_id].geometry.values[0] for lane_group_id in non_conflicting_set - ] - # Get 2D road edge linestrings for the non-conflicting set - set_road_edges_2d = get_road_edge_linestrings(set_polygons, max_road_edge_length=None) - - # Collect 3D boundaries of non-conflicting lane groups - set_boundaries_3d: List[shapely.LineString] = [] - for lane_group_id in non_conflicting_set: - lane_group_helper = lane_group_df[lane_group_df.id == lane_group_id] - set_boundaries_3d.append(lane_group_helper.left_boundary.values[0]) - set_boundaries_3d.append(lane_group_helper.right_boundary.values[0]) - - # Lift road edges to 3D using the boundaries of non-conflicting lane groups - lifted_road_edges_3d = lift_road_edges_to_3d(set_road_edges_2d, set_boundaries_3d) - road_edges_3d.extend(lifted_road_edges_3d) - - return road_edges_3d - - -def lift_road_edges_to_3d( - road_edges_2d: List[shapely.LineString], boundaries: List[shapely.LineString] -) -> List[shapely.LineString]: - - QUERY_MAX_DISTANCE = 0.01 # [m] Maximum distance for nearest neighbor query - - def _find_continuous_sublists(integers: List[int]) -> List[List[int]]: - """Find continuous sublists in a list of integers.""" - arr = np.array(integers, dtype=np.int64) - breaks = np.where(np.diff(arr) != 1)[0] + 1 - splits = np.split(arr, breaks) - return [sublist.tolist() for sublist in splits] - - occupancy_map = OccupancyMap2D(boundaries) - - road_edges_3d: List[shapely.LineString] = [] - for idx, linestring in enumerate(road_edges_2d): - # print(list(linestring.coords)) - points_3d = np.array(list(linestring.coords), dtype=np.float64) - - results = occupancy_map.query_nearest( - shapely.points(points_3d[..., :2]), max_distance=QUERY_MAX_DISTANCE, exclusive=True - ) - for query_idx, geometry_idx in zip(*results): - intersecting_boundary: shapely.LineString = occupancy_map[occupancy_map.ids[geometry_idx]] - points_3d[query_idx, 2] = _get_nearest_z_from_points_3d( - np.array(list(intersecting_boundary.coords), dtype=np.float64), points_3d[query_idx] - ) - - for continuous_slice in _find_continuous_sublists(results[0]): - if len(continuous_slice) < 2: - continue - lifted_linestring = shapely.LineString(points_3d[continuous_slice]) - road_edges_3d.append(lifted_linestring) - return road_edges_3d - - -def create_non_conflicting_sets(conflicts: Dict[int, List[int]]) -> List[Set[int]]: - """ - Creates sets of non-conflicting indices using NetworkX. - """ - # Create graph from conflicts - G = nx.Graph() - for idx, conflict_list in conflicts.items(): - for conflict_idx in conflict_list: - G.add_edge(idx, conflict_idx) - - result = [] - - # Process each connected component - for component in nx.connected_components(G): - subgraph = G.subgraph(component) - - # Try bipartite coloring first (most common case) - if nx.is_bipartite(subgraph): - sets = nx.bipartite.sets(subgraph) - result.extend([set(s) for s in sets]) - else: - # Fall back to greedy coloring for non-bipartite graphs - coloring = nx.greedy_color(subgraph, strategy="largest_first") - color_groups = {} - for node, color in coloring.items(): - if color not in color_groups: - color_groups[color] = set() - color_groups[color].add(node) - result.extend(color_groups.values()) - - return result diff --git a/d123/dataset/conversion/map/road_edge/__init__.py b/d123/dataset/conversion/map/road_edge/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py b/d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py new file mode 100644 index 00000000..f0eb580d --- /dev/null +++ b/d123/dataset/conversion/map/road_edge/road_edge_2d_utils.py @@ -0,0 +1,60 @@ +from typing import Final, List, Union + +import numpy as np +import shapely +from shapely import LinearRing, LineString, Polygon, union_all + +ROAD_EDGE_BUFFER: Final[float] = 0.05 + + +def get_road_edge_linear_rings(drivable_polygons: List[Polygon]) -> List[LinearRing]: + + def _polygon_to_linear_rings(polygon: Polygon) -> List[LinearRing]: + assert polygon.geom_type == "Polygon" + linear_ring_list = [] + linear_ring_list.append(polygon.exterior) + for interior in polygon.interiors: + linear_ring_list.append(interior) + return linear_ring_list + + union_polygon = union_all([polygon.buffer(ROAD_EDGE_BUFFER, join_style=2) for polygon in drivable_polygons]).buffer( + -ROAD_EDGE_BUFFER, join_style=2 + ) + + linear_ring_list = [] + if union_polygon.geom_type == "Polygon": + for polyline in _polygon_to_linear_rings(union_polygon): + linear_ring_list.append(LinearRing(polyline)) + elif union_polygon.geom_type == "MultiPolygon": + for polygon in union_polygon.geoms: + for polyline in _polygon_to_linear_rings(polygon): + linear_ring_list.append(LinearRing(polyline)) + + return linear_ring_list + + +def split_line_geometry_by_max_length( + geometries: List[Union[LineString, LinearRing]], + max_length_meters: float, +) -> List[LineString]: + # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. + + if not isinstance(geometries, list): + geometries = [geometries] + + all_segments = [] + for geom in geometries: + if geom.length <= max_length_meters: + all_segments.append(LineString(geom.coords)) + continue + + num_segments = int(np.ceil(geom.length / max_length_meters)) + segment_length = geom.length / num_segments + + for i in range(num_segments): + start_dist = i * segment_length + end_dist = min((i + 1) * segment_length, geom.length) + segment = shapely.ops.substring(geom, start_dist, end_dist) + all_segments.append(segment) + + return all_segments diff --git a/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py b/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py new file mode 100644 index 00000000..91e9dc47 --- /dev/null +++ b/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py @@ -0,0 +1,279 @@ +from collections import defaultdict +from typing import Dict, List, Set + +import geopandas as gpd +import networkx as nx +import numpy as np +import numpy.typing as npt +import shapely +from shapely.geometry import LineString + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.occupancy_map import OccupancyMap2D +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings + + +def get_road_edges_3d_from_gdf( + lane_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> List[LineString]: + + # 1. Find conflicting lane groups, e.g. groups of lanes that overlap in 2D but have different Z-values (bridges) + conflicting_lane_groups = _get_conflicting_lane_groups(lane_group_df, lane_df) + + # 2. Extract road edges in 2D (including conflicting lane groups) + drivable_polygons = ( + lane_group_df.geometry.tolist() + carpark_df.geometry.tolist() + generic_drivable_area_df.geometry.tolist() + ) + road_edges_2d = get_road_edge_linear_rings(drivable_polygons) + + # 3. Collect 3D boundaries of non-conflicting lane groups and other drivable areas + non_conflicting_boundaries: List[LineString] = [] + for lane_group_id, lane_group_helper in lane_group_df.iterrows(): + if lane_group_id not in conflicting_lane_groups.keys(): + non_conflicting_boundaries.append(lane_group_helper["left_boundary"]) + non_conflicting_boundaries.append(lane_group_helper["right_boundary"]) + for outline in carpark_df.outline.tolist() + generic_drivable_area_df.outline.tolist(): + non_conflicting_boundaries.append(outline) + + # 4. Lift road edges to 3D using the boundaries of non-conflicting elements + non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, non_conflicting_boundaries) + + # 5. Add road edges from conflicting lane groups + resolved_road_edges = _resolve_conflicting_lane_groups(conflicting_lane_groups, lane_group_df) + + all_road_edges = non_conflicting_road_edges + resolved_road_edges + + return all_road_edges + + +def _get_conflicting_lane_groups(lane_group_df: gpd.GeoDataFrame, lane_df: gpd.GeoDataFrame) -> Dict[int, List[int]]: + """ + Even more optimized version using vectorized operations where possible. + """ + Z_THRESHOLD = 5.0 # [m] Z-value threshold for conflict detection + + # Convert to regular dictionaries for faster access + lane_group_dict = lane_group_df.set_index("id").to_dict("index") + lane_baseline_dict = dict(zip(lane_df.id.values, lane_df.baseline_path.values)) + + # Pre-compute all centerlines + centerlines_cache = {} + polygons = [] + ids = [] + + for lane_group_id, data in lane_group_dict.items(): + geometry = data["geometry"] + lane_ids = data["lane_ids"] + + # Vectorized centerline computation + centerlines = [np.array(lane_baseline_dict[lane_id].coords, dtype=np.float64) for lane_id in lane_ids] + centerlines_3d = np.concatenate(centerlines, axis=0) + + centerlines_cache[lane_group_id] = centerlines_3d + polygons.append(geometry) + ids.append(lane_group_id) + + occupancy_map = OccupancyMap2D(polygons, ids) + conflicting_lane_groups: Dict[int, List[int]] = defaultdict(list) + processed_pairs = set() + + for i, lane_group_id in enumerate(ids): + lane_group_polygon = polygons[i] + lane_group_centerlines = centerlines_cache[lane_group_id] + + # Get all intersecting geometries at once + intersecting_ids = occupancy_map.intersects(lane_group_polygon) + intersecting_ids.remove(lane_group_id) + + for intersecting_id in intersecting_ids: + pair_key = tuple(sorted([lane_group_id, intersecting_id])) + if pair_key in processed_pairs: + continue + processed_pairs.add(pair_key) + + intersecting_geometry = occupancy_map[intersecting_id] + if intersecting_geometry.geom_type != "Polygon": + continue + + # Compute actual intersection for better centroid + intersection = lane_group_polygon.intersection(intersecting_geometry) + if intersection.is_empty: + continue + + intersection_centroid = np.array(intersection.centroid.coords[0], dtype=np.float64) + intersecting_centerlines = centerlines_cache[intersecting_id] + + z_at_intersecting = _get_nearest_z_from_points_3d(intersecting_centerlines, intersection_centroid) + z_at_lane_group = _get_nearest_z_from_points_3d(lane_group_centerlines, intersection_centroid) + + if np.abs(z_at_lane_group - z_at_intersecting) >= Z_THRESHOLD: + conflicting_lane_groups[lane_group_id].append(intersecting_id) + conflicting_lane_groups[intersecting_id].append(lane_group_id) + + return conflicting_lane_groups + + +def lift_road_edges_to_3d( + road_edges_2d: List[shapely.LinearRing], + boundaries: List[LineString], + max_distance: float = 0.01, +) -> List[LineString]: + """ + Even faster version using batch processing and optimized data structures. + """ + if not road_edges_2d or not boundaries: + return [] + + # 1. Build comprehensive spatial index with all boundary segments + boundary_segments = [] + + for boundary_idx, boundary in enumerate(boundaries): + coords = np.array(boundary.coords, dtype=np.float64).reshape(-1, 1, 3) + segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1) + boundary_segments.append(segment_coords_boundary) + + boundary_segments = np.concatenate(boundary_segments, axis=0) + boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments) + + occupancy_map = OccupancyMap2D(boundary_segment_linestrings) + + road_edges_3d = [] + for linear_ring in road_edges_2d: + points_2d = np.array(linear_ring.coords, dtype=np.float64) + points_3d = np.zeros((len(points_2d), 3), dtype=np.float64) + points_3d[:, :2] = points_2d + + # 3. Batch query for all points + query_points = shapely.creation.points(points_2d) + results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True) + + for query_idx, geometry_idx in zip(*results): + query_point = query_points[query_idx] + segment_coords = boundary_segments[geometry_idx] + best_z = _interpolate_z_on_segment(query_point, segment_coords) + points_3d[query_idx, 2] = best_z + + continuous_segments = _find_continuous_segments(np.array(results[0])) + + for segment_indices in continuous_segments: + if len(segment_indices) >= 2: + segment_points = points_3d[segment_indices] + road_edges_3d.append(LineString(segment_points)) + + return road_edges_3d + + +def _get_nearest_z_from_points_3d(points_3d: npt.NDArray[np.float64], query_point: npt.NDArray[np.float64]) -> float: + assert points_3d.ndim == 2 and points_3d.shape[1] == len( + Point3DIndex + ), "points_3d must be a 2D array with shape (N, 3)" + distances = np.linalg.norm(points_3d[..., Point3DIndex.XY] - query_point[..., Point3DIndex.XY], axis=1) + closest_point = points_3d[np.argmin(distances)] + return closest_point[2] + + +def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float: + """Interpolate Z coordinate along a 3D line segment.""" + p1, p2 = segment_coords[0], segment_coords[1] + + # Project point onto segment + segment_vec = p2[:2] - p1[:2] + point_vec = np.array([point.x, point.y]) - p1[:2] + + # Handle degenerate case + segment_length_sq = np.dot(segment_vec, segment_vec) + if segment_length_sq == 0: + return p1[2] + + # Calculate projection parameter + t = np.dot(point_vec, segment_vec) / segment_length_sq + t = np.clip(t, 0, 1) # Clamp to segment bounds + + # Interpolate Z + return p1[2] + t * (p2[2] - p1[2]) + + +def _find_continuous_segments(indices: np.ndarray) -> List[np.ndarray]: + """Vectorized version of finding continuous segments.""" + if len(indices) == 0: + return [] + + # Find breaks in continuity + breaks = np.where(np.diff(indices) != 1)[0] + 1 + segments = np.split(indices, breaks) + + # Filter segments with at least 2 points + return [seg for seg in segments if len(seg) >= 2] + + +def _resolve_conflicting_lane_groups( + conflicting_lane_groups: Dict[int, List[int]], lane_group_df: gpd.GeoDataFrame +) -> List[LineString]: + + # Split conflicting lane groups into non-conflicting sets for further merging + non_conflicting_sets = _create_non_conflicting_sets(conflicting_lane_groups) + + road_edges_3d: List[LineString] = [] + for non_conflicting_set in non_conflicting_sets: + + # Collect 2D polygons of non-conflicting lane group set + set_lane_group_rows = lane_group_df[lane_group_df.id.isin(non_conflicting_set)] + connected_lane_group = [] + for row in set_lane_group_rows.itertuples(): + connected_lane_group.extend(row.predecessor_ids) + connected_lane_group.extend(row.successor_ids) + connected_lane_group_rows = lane_group_df[lane_group_df.id.isin(connected_lane_group)] + + set_polygons = set_lane_group_rows.geometry.tolist() + connected_lane_group_rows.geometry.tolist() + + # Get 2D road edge linestrings for the non-conflicting set + set_road_edges_2d = get_road_edge_linear_rings(set_polygons) + + # Collect 3D boundaries of non-conflicting lane groups + set_boundaries_3d: List[LineString] = [] + for lane_group_id in non_conflicting_set: + lane_group_helper = lane_group_df[lane_group_df.id == lane_group_id] + set_boundaries_3d.append(lane_group_helper.left_boundary.values[0]) + set_boundaries_3d.append(lane_group_helper.right_boundary.values[0]) + + # Lift road edges to 3D using the boundaries of non-conflicting lane groups + lifted_road_edges_3d = lift_road_edges_to_3d(set_road_edges_2d, set_boundaries_3d) + road_edges_3d.extend(lifted_road_edges_3d) + + return road_edges_3d + + +def _create_non_conflicting_sets(conflicts: Dict[int, List[int]]) -> List[Set[int]]: + """ + Creates sets of non-conflicting indices using NetworkX. + """ + # Create graph from conflicts + G = nx.Graph() + for idx, conflict_list in conflicts.items(): + for conflict_idx in conflict_list: + G.add_edge(idx, conflict_idx) + + result = [] + + # Process each connected component + for component in nx.connected_components(G): + subgraph = G.subgraph(component) + + # Try bipartite coloring first (most common case) + if nx.is_bipartite(subgraph): + sets = nx.bipartite.sets(subgraph) + result.extend([set(s) for s in sets]) + else: + # Fall back to greedy coloring for non-bipartite graphs + coloring = nx.greedy_color(subgraph, strategy="largest_first") + color_groups = {} + for node, color in coloring.items(): + if color not in color_groups: + color_groups[color] = set() + color_groups[color].add(node) + result.extend(color_groups.values()) + + return result diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index fc0fddbe..2cf3a857 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -19,8 +19,8 @@ from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index from d123.common.geometry.vector import Vector3DIndex from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive -from d123.dataset.dataset_specific.carla.opendrive.opendrive_converter import OpenDriveConverter +from d123.dataset.conversion.map.opendrive.elements.opendrive import OpenDrive +from d123.dataset.conversion.map.opendrive.opendrive_converter import OpenDriveConverter from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import AbstractMap, MapSurfaceType diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py index 10c338be..a5a236ec 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py @@ -2,20 +2,19 @@ import os import warnings -from ast import Tuple from pathlib import Path -from typing import Dict, Final, List, Optional +from typing import Dict, List, Optional import geopandas as gpd import numpy as np -import numpy.typing as npt import pandas as pd import pyogrio -from shapely import union_all from shapely.geometry import LineString -from sympy import Polygon -from d123.common.geometry.constants import DEFAULT_Z +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import ( + get_road_edge_linear_rings, + split_line_geometry_by_max_length, +) from d123.dataset.maps.gpkg.utils import get_all_rows_with_value, get_row_with_value from d123.dataset.maps.map_datatypes import MapSurfaceType @@ -47,9 +46,7 @@ "gen_lane_connectors_scaled_width_polygons", ] -ROAD_EDGE_STEP_SIZE = 0.25 # meters MAX_ROAD_EDGE_LENGTH = 100.0 # meters, used to filter out very long road edges -ROAD_EDGE_BUFFER: Final[float] = 0.05 class NuPlanMapConverter: @@ -64,29 +61,24 @@ def convert(self, map_name: str = "us-pa-pittsburgh-hazelwood") -> None: map_file_path = Path(NUPLAN_MAPS_ROOT) / MAP_FILES[map_name] self._load_dataframes(map_file_path) - lane_df = self._extract_lane_dataframe() - lane_group_df = self._extract_lane_group_dataframe() - intersection_df = self._extract_intersection_dataframe() - crosswalk_df = self._extract_crosswalk_dataframe() - walkway_df = self._extract_walkway_dataframe() - carpark_df = self._extract_carpark_dataframe() - generic_drivable_df = self._extract_generic_drivable_dataframe() - road_edge_df = self._extract_road_edge_dataframe() + dataframes: Dict[MapSurfaceType, gpd.GeoDataFrame] = {} + dataframes[MapSurfaceType.LANE] = self._extract_lane_dataframe() + dataframes[MapSurfaceType.LANE_GROUP] = self._extract_lane_group_dataframe() + dataframes[MapSurfaceType.INTERSECTION] = self._extract_intersection_dataframe() + dataframes[MapSurfaceType.CROSSWALK] = self._extract_crosswalk_dataframe() + dataframes[MapSurfaceType.WALKWAY] = self._extract_walkway_dataframe() + dataframes[MapSurfaceType.CARPARK] = self._extract_carpark_dataframe() + dataframes[MapSurfaceType.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() + dataframes[MapSurfaceType.ROAD_EDGE] = self._extract_road_edge_dataframe() if not self._map_path.exists(): self._map_path.mkdir(parents=True, exist_ok=True) map_file_name = self._map_path / f"nuplan_{map_name}.gpkg" - lane_df.to_file(map_file_name, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") - lane_group_df.to_file(map_file_name, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") - intersection_df.to_file(map_file_name, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_name, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") - walkway_df.to_file(map_file_name, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_name, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_df.to_file( - map_file_name, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" - ) - road_edge_df.to_file(map_file_name, layer=MapSurfaceType.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_name, layer=layer.serialize(), driver="GPKG", mode="a") def _load_dataframes(self, map_file_path: Path) -> None: @@ -402,10 +394,11 @@ def _extract_road_edge_dataframe(self) -> gpd.GeoDataFrame: + self._gdf["carpark_areas"].geometry.to_list() + self._gdf["generic_drivable_areas"].geometry.to_list() ) - road_edge_linestrings = get_road_edge_linestrings(drivable_polygons) + road_edge_linear_rings = get_road_edge_linear_rings(drivable_polygons) + road_edges = split_line_geometry_by_max_length(road_edge_linear_rings, MAX_ROAD_EDGE_LENGTH) - data = pd.DataFrame({"id": [idx for idx in range(len(road_edge_linestrings))]}) - return gpd.GeoDataFrame(data, geometry=road_edge_linestrings) + data = pd.DataFrame({"id": [idx for idx in range(len(road_edges))]}) + return gpd.GeoDataFrame(data, geometry=road_edges) def flip_linestring(linestring: LineString) -> LineString: @@ -432,58 +425,3 @@ def align_boundary_direction(centerline: LineString, boundary: LineString) -> Li if not lines_same_direction(centerline, boundary): return flip_linestring(boundary) return boundary - - -def get_road_edge_linestrings( - drivable_polygons: List[Polygon], - step_size: float = ROAD_EDGE_STEP_SIZE, - max_road_edge_length: Optional[float] = MAX_ROAD_EDGE_LENGTH, -) -> List[LineString]: - # TODO: move this function for general usage. - - def _coords_to_points(coords: npt.NDArray[np.float64]) -> List[npt.NDArray[np.float64]]: - coords = np.array(coords).reshape((-1, 2)) - linestring = LineString(coords) - - points_list: List[npt.NDArray[np.float64]] = [] - - segments: List[Tuple[float, float]] = [] # Start and end points of segments - if max_road_edge_length is not None: - num_segments = int(linestring.length / max_road_edge_length) - for i in range(num_segments + 1): - start = np.clip(float(i) * max_road_edge_length, 0, linestring.length) - end = np.clip((float(i) + 1) * max_road_edge_length, 0, linestring.length) - segments.append((start, end)) - else: - segments = [(0.0, linestring.length)] - - for segment in segments: - start, end = segment - distances = np.arange(start, end + step_size, step_size) - points = [linestring.interpolate(distance, normalized=False) for distance in distances] - points_list.append(np.array([[p.x, p.y, DEFAULT_Z] for p in points])) - - return points_list - - def _polygon_to_coords(polygon: Polygon) -> List[npt.NDArray[np.float64]]: - assert polygon.geom_type == "Polygon" - points_list = [] - points_list.extend(_coords_to_points(polygon.exterior.coords)) - for interior in polygon.interiors: - points_list.extend(_coords_to_points(interior.coords)) - return points_list - - union_polygon = union_all([polygon.buffer(ROAD_EDGE_BUFFER, join_style=2) for polygon in drivable_polygons]).buffer( - -ROAD_EDGE_BUFFER, join_style=2 - ) - - linestring_list = [] - if union_polygon.geom_type == "Polygon": - for polyline in _polygon_to_coords(union_polygon): - linestring_list.append(LineString(polyline)) - elif union_polygon.geom_type == "MultiPolygon": - for polygon in union_polygon.geoms: - for polyline in _polygon_to_coords(polygon): - linestring_list.append(LineString(polyline)) - - return linestring_list diff --git a/d123/dataset/maps/gpkg/gpkg_map_objects.py b/d123/dataset/maps/gpkg/gpkg_map_objects.py index 876bdf4d..24b0fdbb 100644 --- a/d123/dataset/maps/gpkg/gpkg_map_objects.py +++ b/d123/dataset/maps/gpkg/gpkg_map_objects.py @@ -65,10 +65,18 @@ def outline_3d(self) -> Polyline3D: @property def trimesh_mesh(self) -> trimesh.Trimesh: """Inherited, see superclass.""" - outline_3d_array = self.outline_3d.array - _, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) - # NOTE: Optional add color information to the mesh - return trimesh.Trimesh(vertices=outline_3d_array, faces=faces) + + trimesh_mesh: Optional[trimesh.Trimesh] = None + if "right_boundary" in self._object_df.columns and "left_boundary" in self._object_df.columns: + left_boundary = Polyline3D.from_linestring(self._object_row.left_boundary) + right_boundary = Polyline3D.from_linestring(self._object_row.right_boundary) + trimesh_mesh = get_trimesh_from_boundaries(left_boundary, right_boundary) + elif "geometry" in self._object_df.columns: + # Fallback to geometry if no boundaries are available + outline_3d_array = self.outline_3d.array + _, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) + trimesh_mesh = trimesh.Trimesh(vertices=outline_3d_array, faces=faces) + return trimesh_mesh class GPKGLane(GPKGSurfaceObject, AbstractLane): @@ -122,11 +130,6 @@ def outline_3d(self) -> Polyline3D: outline_array = np.vstack((outline_array, outline_array[0])) return Polyline3D.from_linestring(geom.LineString(outline_array)) - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - return get_trimesh_from_boundaries(self.left_boundary, self.right_boundary) - @property def lane_group(self) -> GPKGLaneGroup: """Inherited, see superclass.""" @@ -179,11 +182,6 @@ def outline_3d(self) -> Polyline3D: outline_array = np.vstack((self.left_boundary.array, self.right_boundary.array[::-1])) return Polyline3D.from_linestring(geom.LineString(outline_array)) - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - return get_trimesh_from_boundaries(self.left_boundary, self.right_boundary) - @property def lanes(self) -> List[GPKGLane]: """Inherited, see superclass.""" @@ -260,6 +258,16 @@ class GPKGGenericDrivable(GPKGSurfaceObject, AbstractGenericDrivable): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): super().__init__(object_id, object_df) + @property + def trimesh_mesh(self) -> trimesh.Trimesh: + """Inherited, see superclass.""" + + # Fallback to geometry if no boundaries are available + outline_3d_array = self.outline_3d.array + _, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) + trimesh_mesh = trimesh.Trimesh(vertices=outline_3d_array, faces=faces) + return trimesh_mesh + class GPKGRoadEdge(AbstractRoadEdge): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 37db915f..20540413 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,8 +15,8 @@ defaults: - default_dataset_paths - _self_ - datasets: - # - nuplan_private_dataset - - carla_dataset + - nuplan_private_dataset + # - carla_dataset # - wopd_dataset force_log_conversion: False diff --git a/notebooks/map_conversion/test_nuplan_conversion.ipynb b/notebooks/map_conversion/test_nuplan_conversion.ipynb index 869e55f3..32805c97 100644 --- a/notebooks/map_conversion/test_nuplan_conversion.ipynb +++ b/notebooks/map_conversion/test_nuplan_conversion.ipynb @@ -6,37 +6,10 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", - "from shapely.geometry import LineString, Polygon, Point\n", - "import numpy as np\n", - "\n", - "from typing import List\n", - "\n", - "\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lane, LaneSection\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", "from pathlib import Path\n", - "from d123.dataset.dataset_specific.nuplan.map_conversion import NuPlanMapConverter, MAP_LOCATIONS\n", + "from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import NuPlanMapConverter, MAP_LOCATIONS\n", "\n", - "import pyogrio\n", "\n", - "import geopandas as gpd\n", "\n", "\n", "# file_to_delete = Path(\"nuplan_us-pa-pittsburgh-hazelwood.gpkg\")\n", @@ -53,7 +26,7 @@ " file_to_delete = root / f\"nuplan_{map_name}.gpkg\"\n", " if file_to_delete.exists():\n", " file_to_delete.unlink()\n", - " \n", + "\n", " NuPlanMapConverter(root).convert(map_name=map_name)\n", " print(f\"Converting {map_name} map... done\")\n", "\n", diff --git a/notebooks/map_conversion/test_opendrive_conversion.ipynb b/notebooks/map_conversion/test_opendrive_conversion.ipynb index 0de615ec..9bc80a2c 100644 --- a/notebooks/map_conversion/test_opendrive_conversion.ipynb +++ b/notebooks/map_conversion/test_opendrive_conversion.ipynb @@ -12,8 +12,8 @@ "\n", "from typing import List\n", "\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.lane import Lane, LaneSection\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.reference import Border" + "from d123.dataset.conversion.map.opendrive.elements.lane import Lane, LaneSection\n", + "from d123.dataset.conversion.map.opendrive.elements.reference import Border" ] }, { @@ -23,7 +23,7 @@ "outputs": [], "source": [ "from pathlib import Path\n", - "from d123.dataset.dataset_specific.carla.opendrive.elements.opendrive import OpenDrive\n", + "from d123.dataset.conversion.map.opendrive.elements.opendrive import OpenDrive\n", "\n", "CARLA_MAP_LOCATIONS = [\n", " \"Town01\", # A small, simple town with a river and several bridges.\n", @@ -79,7 +79,7 @@ "metadata": {}, "outputs": [], "source": [ - "from d123.dataset.dataset_specific.carla.opendrive.opendrive_converter import OpenDriveConverter\n", + "from d123.dataset.conversion.map.opendrive.opendrive_converter import OpenDriveConverter\n", "\n", "\n", "\n", diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 16d0add2..f2a64092 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -25,25 +25,24 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "# splits = [\"wopd_train\"]\n", - "splits = [\"carla\"]\n", + "splits = [\"wopd_train\"]\n", + "# splits = [\"carla\"]\n", + "# splits = [\"nuplan_private_test\"]\n", "# log_names = None\n", "\n", "\n", "\n", - "# splits = [\"nuplan_private_test\"]\n", "log_names = None\n", - "\n", "scene_tokens = None\n", "\n", "scene_filter = SceneFilter(\n", " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=10,\n", + " duration_s=19,\n", " history_s=0.0,\n", " timestamp_threshold_s=20,\n", - " shuffle=False,\n", + " shuffle=True,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", @@ -236,15 +235,17 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " size = 10\n", + " size = 15\n", "\n", " fig, ax = plt.subplots(figsize=(size, size))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", - "scene_index = 14\n", - "plot_scene_at_iteration(scenes[scene_index], iteration=1, radius=82)" + "scene_index = 8\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=100, radius=40)\n", + "\n", + "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] }, { diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 75f9de83..e7584e06 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -30,8 +30,8 @@ "\n", "\n", "\n", - "splits = [\"nuplan_private_test\"]\n", - "# splits = [\"carla\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"carla\"]\n", "# log_names = [\"11839652018869852123_2565_000_2585_000\"]\n", "log_names = None\n", "\n", diff --git a/scripts/dataset/run_log_caching.sh b/scripts/dataset/run_log_caching.sh index 4d278d15..34ff923c 100644 --- a/scripts/dataset/run_log_caching.sh +++ b/scripts/dataset/run_log_caching.sh @@ -3,5 +3,5 @@ python $D123_DEVKIT_ROOT/d123/script/run_dataset_conversion.py \ experiment_name="caching" \ -worker.threads_per_node=10 +worker.threads_per_node=32 # worker=single_machine_thread_pool From ada4442a5cfc45db84d6324010c396ece9eb4830 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 10 Aug 2025 15:08:50 +0200 Subject: [PATCH 27/42] Rename MapSurfaceType to MapLayer, as map types were expanded to road lines and edges, which are polylines (not surfaces). Also adapted the map object structure accordingly. --- d123/common/visualization/color/default.py | 18 +- .../visualization/matplotlib/observation.py | 34 +-- d123/common/visualization/viser/utils.py | 56 ++-- .../map/opendrive/opendrive_converter.py | 42 +-- .../carla/carla_data_converter.py | 10 +- .../nuplan/nuplan_map_conversion.py | 20 +- .../wopd/waymo_map_utils/wopd_map_utils.py | 22 +- d123/dataset/maps/abstract_map.py | 26 +- d123/dataset/maps/abstract_map_objects.py | 117 +++---- d123/dataset/maps/gpkg/gpkg_map.py | 120 ++++---- d123/dataset/maps/gpkg/gpkg_map_objects.py | 28 +- d123/dataset/maps/map_datatypes.py | 2 +- .../environment/helper/environment_cache.py | 16 +- .../metrics/sim_agents/map_based.py | 16 +- .../feature_builder/smart_feature_builder.py | 18 +- notebooks/Untitled-2.ipynb | 286 ------------------ notebooks/deprecated/arrow.ipynb | 4 +- .../test_opendrive_conversion.ipynb | 12 +- notebooks/viz/bev_matplotlib.ipynb | 58 ++-- notebooks/viz/viser_testing_v2.ipynb | 87 ------ 20 files changed, 316 insertions(+), 676 deletions(-) delete mode 100644 notebooks/Untitled-2.ipynb delete mode 100644 notebooks/viz/viser_testing_v2.ipynb diff --git a/d123/common/visualization/color/default.py b/d123/common/visualization/color/default.py index f5a3cf94..076a1c81 100644 --- a/d123/common/visualization/color/default.py +++ b/d123/common/visualization/color/default.py @@ -13,12 +13,12 @@ Color, ) from d123.common.visualization.color.config import PlotConfig -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer HEADING_MARKER_STYLE: str = "^" # "^": triangle, "-": line -MAP_SURFACE_CONFIG: Dict[MapSurfaceType, PlotConfig] = { - MapSurfaceType.LANE: PlotConfig( +MAP_SURFACE_CONFIG: Dict[MapLayer, PlotConfig] = { + MapLayer.LANE: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -27,7 +27,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.LANE_GROUP: PlotConfig( + MapLayer.LANE_GROUP: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -36,7 +36,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.INTERSECTION: PlotConfig( + MapLayer.INTERSECTION: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, @@ -45,7 +45,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.CROSSWALK: PlotConfig( + MapLayer.CROSSWALK: PlotConfig( fill_color=Color("#c69fbb"), fill_color_alpha=1.0, line_color=Color("#c69fbb"), @@ -54,7 +54,7 @@ line_style="-", zorder=2, ), - MapSurfaceType.WALKWAY: PlotConfig( + MapLayer.WALKWAY: PlotConfig( fill_color=Color("#d4d19e"), fill_color_alpha=1.0, line_color=Color("#d4d19e"), @@ -63,7 +63,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.CARPARK: PlotConfig( + MapLayer.CARPARK: PlotConfig( fill_color=Color("#b9d3b4"), fill_color_alpha=1.0, line_color=Color("#b9d3b4"), @@ -72,7 +72,7 @@ line_style="-", zorder=1, ), - MapSurfaceType.GENERIC_DRIVABLE: PlotConfig( + MapLayer.GENERIC_DRIVABLE: PlotConfig( fill_color=LIGHT_GREY, fill_color_alpha=1.0, line_color=LIGHT_GREY, diff --git a/d123/common/visualization/matplotlib/observation.py b/d123/common/visualization/matplotlib/observation.py index 26cb65b6..0d46ca93 100644 --- a/d123/common/visualization/matplotlib/observation.py +++ b/d123/common/visualization/matplotlib/observation.py @@ -26,7 +26,7 @@ ) from d123.dataset.maps.abstract_map import AbstractMap from d123.dataset.maps.abstract_map_objects import AbstractLane -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer def add_default_map_on_ax( @@ -36,14 +36,14 @@ def add_default_map_on_ax( radius: float, route_lane_group_ids: Optional[List[int]] = None, ) -> None: - layers: List[MapSurfaceType] = [ - MapSurfaceType.LANE, - MapSurfaceType.LANE_GROUP, - MapSurfaceType.GENERIC_DRIVABLE, - MapSurfaceType.CARPARK, - MapSurfaceType.CROSSWALK, - MapSurfaceType.INTERSECTION, - MapSurfaceType.WALKWAY, + layers: List[MapLayer] = [ + MapLayer.LANE, + MapLayer.LANE_GROUP, + MapLayer.GENERIC_DRIVABLE, + MapLayer.CARPARK, + MapLayer.CROSSWALK, + MapLayer.INTERSECTION, + MapLayer.WALKWAY, ] x_min, x_max = point_2d.x - radius, point_2d.x + radius y_min, y_max = point_2d.y - radius, point_2d.y + radius @@ -53,20 +53,20 @@ def add_default_map_on_ax( for layer, map_objects in map_objects_dict.items(): for map_object in map_objects: try: - if layer in [MapSurfaceType.LANE_GROUP]: + if layer in [MapLayer.LANE_GROUP]: if route_lane_group_ids is not None and int(map_object.id) in route_lane_group_ids: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, ROUTE_CONFIG) else: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer]) if layer in [ - MapSurfaceType.GENERIC_DRIVABLE, - MapSurfaceType.CARPARK, - MapSurfaceType.CROSSWALK, - MapSurfaceType.INTERSECTION, - MapSurfaceType.WALKWAY, + MapLayer.GENERIC_DRIVABLE, + MapLayer.CARPARK, + MapLayer.CROSSWALK, + MapLayer.INTERSECTION, + MapLayer.WALKWAY, ]: add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer]) - if layer in [MapSurfaceType.LANE]: + if layer in [MapLayer.LANE]: map_object: AbstractLane add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG) except Exception: @@ -95,7 +95,7 @@ def add_traffic_lights_to_ax( ax: plt.Axes, traffic_light_detections: TrafficLightDetectionWrapper, map_api: AbstractMap ) -> None: for traffic_light_detection in traffic_light_detections: - lane: AbstractLane = map_api.get_map_object(str(traffic_light_detection.lane_id), MapSurfaceType.LANE) + lane: AbstractLane = map_api.get_map_object(str(traffic_light_detection.lane_id), MapLayer.LANE) if lane is not None: add_shapely_linestring_to_ax( ax, diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 151e820c..188100eb 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -14,14 +14,14 @@ from d123.common.geometry.transform.se3 import convert_relative_to_absolute_points_3d_array from d123.common.visualization.color.config import PlotConfig from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG, MAP_SURFACE_CONFIG -from d123.dataset.maps.abstract_map import MapSurfaceType +from d123.dataset.maps.abstract_map import MapLayer from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractSurfaceMapObject from d123.dataset.scene.abstract_scene import AbstractScene # TODO: Refactor this file. # TODO: Add general utilities for 3D primitives and mesh support. -MAP_RADIUS: Final[float] = 80 +MAP_RADIUS: Final[float] = 500 BRIGHTNESS_FACTOR: Final[float] = 0.8 @@ -87,34 +87,28 @@ def get_bounding_box_meshes(scene: AbstractScene, iteration: int): def get_map_meshes(scene: AbstractScene): initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 - map_surface_types = [ - MapSurfaceType.LANE_GROUP, - # MapSurfaceType.LANE, - MapSurfaceType.WALKWAY, - MapSurfaceType.CROSSWALK, - MapSurfaceType.CARPARK, - MapSurfaceType.GENERIC_DRIVABLE, + map_layers = [ + MapLayer.LANE_GROUP, + # MapLayer.LANE, + MapLayer.WALKWAY, + MapLayer.CROSSWALK, + MapLayer.CARPARK, + MapLayer.GENERIC_DRIVABLE, ] - map_objects_dict = scene.map_api.get_proximal_map_objects( - center.point_2d, radius=MAP_RADIUS, layers=map_surface_types - ) + map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) output = {} - for map_surface_type in map_objects_dict.keys(): + for map_layer in map_objects_dict.keys(): surface_meshes = [] - for map_surface in map_objects_dict[map_surface_type]: + for map_surface in map_objects_dict[map_layer]: map_surface: AbstractSurfaceMapObject trimesh_mesh = map_surface.trimesh_mesh - if map_surface_type == MapSurfaceType.GENERIC_DRIVABLE: - print("Generic Drivable Surface Mesh:") - print(trimesh_mesh) - output[f"{map_surface_type.serialize()}_{map_surface.id}"] = trimesh_mesh - if map_surface_type in [ - MapSurfaceType.WALKWAY, - MapSurfaceType.CROSSWALK, - # MapSurfaceType.GENERIC_DRIVABLE, - # MapSurfaceType.CARPARK, + if map_layer in [ + MapLayer.WALKWAY, + MapLayer.CROSSWALK, + MapLayer.GENERIC_DRIVABLE, + # MapLayer.CARPARK, ]: # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.05).array @@ -126,20 +120,18 @@ def get_map_meshes(scene: AbstractScene): x=0, y=0, z=center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 ).array - trimesh_mesh.visual.face_colors = MAP_SURFACE_CONFIG[map_surface_type].fill_color.rgba + trimesh_mesh.visual.face_colors = MAP_SURFACE_CONFIG[map_layer].fill_color.set_brightness(0.8).rgba surface_meshes.append(trimesh_mesh) - output[f"{map_surface_type.serialize()}"] = trimesh.util.concatenate(surface_meshes) + output[f"{map_layer.serialize()}"] = trimesh.util.concatenate(surface_meshes) return output def get_map_lines(scene: AbstractScene): - map_surface_types = [MapSurfaceType.LANE, MapSurfaceType.ROAD_EDGE] + map_layers = [MapLayer.LANE, MapLayer.ROAD_EDGE] initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 - map_objects_dict = scene.map_api.get_proximal_map_objects( - center.point_2d, radius=MAP_RADIUS, layers=map_surface_types - ) + map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) def polyline_to_segments(polyline: Polyline3D) -> npt.NDArray[np.float64]: polyline_array = polyline.array - center.point_3d.array @@ -148,14 +140,14 @@ def polyline_to_segments(polyline: Polyline3D) -> npt.NDArray[np.float64]: return polyline_array centerlines, right_boundaries, left_boundaries, road_edges = [], [], [], [] - for lane in map_objects_dict[MapSurfaceType.LANE]: + for lane in map_objects_dict[MapLayer.LANE]: lane: AbstractLane centerlines.append(polyline_to_segments(lane.centerline)) right_boundaries.append(polyline_to_segments(lane.right_boundary)) left_boundaries.append(polyline_to_segments(lane.left_boundary)) - for road_edge in map_objects_dict[MapSurfaceType.ROAD_EDGE]: + for road_edge in map_objects_dict[MapLayer.ROAD_EDGE]: road_edges.append(polyline_to_segments(road_edge.polyline_3d)) centerlines = np.concatenate(centerlines, axis=0) @@ -214,7 +206,7 @@ def _create_lane_mesh_from_boundary_arrays( faces = np.array(faces) mesh = trimesh.Trimesh(vertices=vertices, faces=faces) - mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapSurfaceType.LANE].fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba + mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapLayer.LANE].fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba return mesh diff --git a/d123/dataset/conversion/map/opendrive/opendrive_converter.py b/d123/dataset/conversion/map/opendrive/opendrive_converter.py index 3ad5bea3..f536b7fc 100644 --- a/d123/dataset/conversion/map/opendrive/opendrive_converter.py +++ b/d123/dataset/conversion/map/opendrive/opendrive_converter.py @@ -29,7 +29,7 @@ from d123.dataset.conversion.map.opendrive.id_mapping import IntIDMapping from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import get_road_edges_3d_from_gdf -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer ENABLE_WARNING: bool = False CONNECTION_DISTANCE_THRESHOLD: float = 0.1 # [m] @@ -63,29 +63,29 @@ def run(self, map_name: str) -> None: self._collect_crosswalks() # Collect data frames and store - dataframes: Dict[MapSurfaceType, gpd.GeoDataFrame] = {} - dataframes[MapSurfaceType.LANE] = self._extract_lane_dataframe() - dataframes[MapSurfaceType.WALKWAY] = self._extract_walkways_dataframe() - dataframes[MapSurfaceType.CARPARK] = self._extract_carpark_dataframe() - dataframes[MapSurfaceType.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() - dataframes[MapSurfaceType.INTERSECTION] = self._extract_intersections_dataframe() - dataframes[MapSurfaceType.LANE_GROUP] = self._extract_lane_group_dataframe() - dataframes[MapSurfaceType.CROSSWALK] = self._extract_crosswalk_dataframe() + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + dataframes[MapLayer.LANE] = self._extract_lane_dataframe() + dataframes[MapLayer.WALKWAY] = self._extract_walkways_dataframe() + dataframes[MapLayer.CARPARK] = self._extract_carpark_dataframe() + dataframes[MapLayer.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() + dataframes[MapLayer.INTERSECTION] = self._extract_intersections_dataframe() + dataframes[MapLayer.LANE_GROUP] = self._extract_lane_group_dataframe() + dataframes[MapLayer.CROSSWALK] = self._extract_crosswalk_dataframe() self._convert_ids_to_int( - dataframes[MapSurfaceType.LANE], - dataframes[MapSurfaceType.WALKWAY], - dataframes[MapSurfaceType.CARPARK], - dataframes[MapSurfaceType.GENERIC_DRIVABLE], - dataframes[MapSurfaceType.LANE_GROUP], - dataframes[MapSurfaceType.INTERSECTION], - dataframes[MapSurfaceType.CROSSWALK], + dataframes[MapLayer.LANE], + dataframes[MapLayer.WALKWAY], + dataframes[MapLayer.CARPARK], + dataframes[MapLayer.GENERIC_DRIVABLE], + dataframes[MapLayer.LANE_GROUP], + dataframes[MapLayer.INTERSECTION], + dataframes[MapLayer.CROSSWALK], ) - dataframes[MapSurfaceType.ROAD_EDGE] = self._extract_road_edge_df( - dataframes[MapSurfaceType.LANE], - dataframes[MapSurfaceType.CARPARK], - dataframes[MapSurfaceType.GENERIC_DRIVABLE], - dataframes[MapSurfaceType.LANE_GROUP], + dataframes[MapLayer.ROAD_EDGE] = self._extract_road_edge_df( + dataframes[MapLayer.LANE], + dataframes[MapLayer.CARPARK], + dataframes[MapLayer.GENERIC_DRIVABLE], + dataframes[MapLayer.LANE_GROUP], ) # Store dataframes diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 2cf3a857..eb37ff86 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -23,7 +23,7 @@ from d123.dataset.conversion.map.opendrive.opendrive_converter import OpenDriveConverter from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata -from d123.dataset.maps.abstract_map import AbstractMap, MapSurfaceType +from d123.dataset.maps.abstract_map import AbstractMap, MapLayer from d123.dataset.maps.abstract_map_objects import AbstractLane from d123.dataset.scene.arrow_scene import get_map_api_from_names @@ -342,8 +342,8 @@ def _extract_traffic_light_data( for traffic_light_waypoint in traffic_light_waypoints: point_3d = Point3D(*traffic_light_waypoint) nearby_lanes = map_api.get_proximal_map_objects( - point_3d, TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE, [MapSurfaceType.LANE] - )[MapSurfaceType.LANE] + point_3d, TRAFFIC_LIGHT_ASSIGNMENT_DISTANCE, [MapLayer.LANE] + )[MapLayer.LANE] for lane in nearby_lanes: lane: AbstractLane @@ -366,8 +366,8 @@ def _extract_route_lane_group_ids(route: List[List[float]], map_api: AbstractMap for point in route[:200]: point_2d = Point2D(point[0], point[1]) - nearby_lane_groups = map_api.query(point_2d.shapely_point, [MapSurfaceType.LANE_GROUP], predicate="intersects")[ - MapSurfaceType.LANE_GROUP + nearby_lane_groups = map_api.query(point_2d.shapely_point, [MapLayer.LANE_GROUP], predicate="intersects")[ + MapLayer.LANE_GROUP ] if len(nearby_lane_groups) == 0: continue diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py index a5a236ec..7f4c683b 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py @@ -16,7 +16,7 @@ split_line_geometry_by_max_length, ) from d123.dataset.maps.gpkg.utils import get_all_rows_with_value, get_row_with_value -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer MAP_FILES = { "sg-one-north": "sg-one-north/9.17.1964/map.gpkg", @@ -61,15 +61,15 @@ def convert(self, map_name: str = "us-pa-pittsburgh-hazelwood") -> None: map_file_path = Path(NUPLAN_MAPS_ROOT) / MAP_FILES[map_name] self._load_dataframes(map_file_path) - dataframes: Dict[MapSurfaceType, gpd.GeoDataFrame] = {} - dataframes[MapSurfaceType.LANE] = self._extract_lane_dataframe() - dataframes[MapSurfaceType.LANE_GROUP] = self._extract_lane_group_dataframe() - dataframes[MapSurfaceType.INTERSECTION] = self._extract_intersection_dataframe() - dataframes[MapSurfaceType.CROSSWALK] = self._extract_crosswalk_dataframe() - dataframes[MapSurfaceType.WALKWAY] = self._extract_walkway_dataframe() - dataframes[MapSurfaceType.CARPARK] = self._extract_carpark_dataframe() - dataframes[MapSurfaceType.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() - dataframes[MapSurfaceType.ROAD_EDGE] = self._extract_road_edge_dataframe() + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + dataframes[MapLayer.LANE] = self._extract_lane_dataframe() + dataframes[MapLayer.LANE_GROUP] = self._extract_lane_group_dataframe() + dataframes[MapLayer.INTERSECTION] = self._extract_intersection_dataframe() + dataframes[MapLayer.CROSSWALK] = self._extract_crosswalk_dataframe() + dataframes[MapLayer.WALKWAY] = self._extract_walkway_dataframe() + dataframes[MapLayer.CARPARK] = self._extract_carpark_dataframe() + dataframes[MapLayer.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() + dataframes[MapLayer.ROAD_EDGE] = self._extract_road_edge_dataframe() if not self._map_path.exists(): self._map_path.mkdir(parents=True, exist_ok=True) diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py index 160dae76..ed023eaa 100644 --- a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py @@ -13,7 +13,7 @@ from d123.common.geometry.line.polylines import Polyline3D from d123.common.geometry.units import mph_to_mps from d123.dataset.dataset_specific.wopd.waymo_map_utils.womp_boundary_utils import extract_lane_boundaries -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer # TODO: # - Implement stop signs @@ -112,17 +112,15 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: if not map_file_path.parent.exists(): map_file_path.parent.mkdir(parents=True, exist_ok=True) - lane_df.to_file(map_file_path, layer=MapSurfaceType.LANE.serialize(), driver="GPKG") - lane_group_df.to_file(map_file_path, layer=MapSurfaceType.LANE_GROUP.serialize(), driver="GPKG", mode="a") - intersection_df.to_file(map_file_path, layer=MapSurfaceType.INTERSECTION.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_path, layer=MapSurfaceType.CROSSWALK.serialize(), driver="GPKG", mode="a") - walkway_df.to_file(map_file_path, layer=MapSurfaceType.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_path, layer=MapSurfaceType.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_df.to_file( - map_file_path, layer=MapSurfaceType.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a" - ) - road_edge_df.to_file(map_file_path, layer=MapSurfaceType.ROAD_EDGE.serialize(), driver="GPKG", mode="a") - road_line_df.to_file(map_file_path, layer=MapSurfaceType.ROAD_LINE.serialize(), driver="GPKG", mode="a") + lane_df.to_file(map_file_path, layer=MapLayer.LANE.serialize(), driver="GPKG") + lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") + intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") + crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") + walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") + carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") + generic_drivable_df.to_file(map_file_path, layer=MapLayer.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a") + road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + road_line_df.to_file(map_file_path, layer=MapLayer.ROAD_LINE.serialize(), driver="GPKG", mode="a") def get_lane_df( diff --git a/d123/dataset/maps/abstract_map.py b/d123/dataset/maps/abstract_map.py index 5db687be..0dc483bb 100644 --- a/d123/dataset/maps/abstract_map.py +++ b/d123/dataset/maps/abstract_map.py @@ -7,7 +7,7 @@ from d123.common.geometry.base import Point2D from d123.dataset.maps.abstract_map_objects import AbstractMapObject -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer # TODO: # - add docstrings @@ -27,57 +27,57 @@ def initialize(self) -> None: pass @abc.abstractmethod - def get_available_map_objects(self) -> List[MapSurfaceType]: + def get_available_map_objects(self) -> List[MapLayer]: pass @abc.abstractmethod - def get_map_object(self, object_id: str, layer: MapSurfaceType) -> Optional[AbstractMapObject]: + def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[AbstractMapObject]: pass @abc.abstractmethod - def get_all_map_objects(self, point_2d: Point2D, layer: MapSurfaceType) -> List[AbstractMapObject]: + def get_all_map_objects(self, point_2d: Point2D, layer: MapLayer) -> List[AbstractMapObject]: pass @abc.abstractmethod - def is_in_layer(self, point: Point2D, layer: MapSurfaceType) -> bool: + def is_in_layer(self, point: Point2D, layer: MapLayer) -> bool: pass @abc.abstractmethod def get_proximal_map_objects( - self, point: Point2D, radius: float, layers: List[MapSurfaceType] - ) -> Dict[MapSurfaceType, List[AbstractMapObject]]: + self, point: Point2D, radius: float, layers: List[MapLayer] + ) -> Dict[MapLayer, List[AbstractMapObject]]: pass @abc.abstractmethod def query( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: + ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: pass @abc.abstractmethod def query_object_ids( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]]: + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: pass @abc.abstractmethod def query_nearest( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], return_all: bool = True, max_distance: Optional[float] = None, return_distance: bool = False, exclusive: bool = False, - ) -> Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: + ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: pass diff --git a/d123/dataset/maps/abstract_map_objects.py b/d123/dataset/maps/abstract_map_objects.py index 93eaa6aa..e19eb679 100644 --- a/d123/dataset/maps/abstract_map_objects.py +++ b/d123/dataset/maps/abstract_map_objects.py @@ -6,8 +6,8 @@ import shapely.geometry as geom import trimesh -from d123.common.geometry.line.polylines import Polyline2D, Polyline3D -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.common.geometry.line.polylines import Polyline2D, Polyline3D, PolylineSE2 +from d123.dataset.maps.map_datatypes import MapLayer class AbstractMapObject(abc.ABC): @@ -22,6 +22,14 @@ def __init__(self, object_id: str): """ self.id = str(object_id) + @property + @abc.abstractmethod + def layer(self) -> MapLayer: + """ + Returns map layer type, e.g. LANE, ROAD_EDGE. + :return: map layer type + """ + class AbstractSurfaceMapObject(AbstractMapObject): """ @@ -36,14 +44,6 @@ def shapely_polygon(self) -> geom.Polygon: :return: shapely polygon """ - @property - @abc.abstractmethod - def surface_type(self) -> MapSurfaceType: - """ - Returns map surface type, e.g. LANE. - :return: map surface type - """ - @property @abc.abstractmethod def outline_3d(self) -> Polyline3D: @@ -64,12 +64,39 @@ def outline_2d(self) -> Polyline2D: return self.outline_3d.polyline_2d +class AbstractLineMapObject(AbstractMapObject): + + @property + @abc.abstractmethod + def polyline_3d(self) -> Polyline3D: + """ + Returns the 3D polyline of the road edge. + :return: 3D polyline + """ + + @property + def polyline_2d(self) -> Polyline2D: + """ + Returns the 2D polyline of the road line. + :return: 2D polyline + """ + return self.polyline_3d.polyline_2d + + @property + def polyline_se2(self) -> PolylineSE2: + """ + Returns the 2D polyline of the road line in SE(2) coordinates. + :return: 2D polyline in SE(2) + """ + return self.polyline_3d.polyline_se2 + + class AbstractLane(AbstractSurfaceMapObject): """Abstract interface for lane objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.LANE + def layer(self) -> MapLayer: + return MapLayer.LANE @property @abc.abstractmethod @@ -140,8 +167,8 @@ class AbstractLaneGroup(AbstractSurfaceMapObject): """Abstract interface lane groups (nearby lanes going in the same direction).""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.LANE_GROUP + def layer(self) -> MapLayer: + return MapLayer.LANE_GROUP @property @abc.abstractmethod @@ -196,8 +223,8 @@ class AbstractIntersection(AbstractSurfaceMapObject): """Abstract interface for intersection objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.INTERSECTION + def layer(self) -> MapLayer: + return MapLayer.INTERSECTION @property @abc.abstractmethod @@ -212,49 +239,48 @@ class AbstractCrosswalk(AbstractSurfaceMapObject): """Abstract interface for crosswalk objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.CROSSWALK + def layer(self) -> MapLayer: + return MapLayer.CROSSWALK class AbstractWalkway(AbstractSurfaceMapObject): """Abstract interface for walkway objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.WALKWAY + def layer(self) -> MapLayer: + return MapLayer.WALKWAY class AbstractCarpark(AbstractSurfaceMapObject): """Abstract interface for carpark objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.CARPARK + def layer(self) -> MapLayer: + return MapLayer.CARPARK class AbstractGenericDrivable(AbstractSurfaceMapObject): """Abstract interface for generic drivable objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.GENERIC_DRIVABLE + def layer(self) -> MapLayer: + return MapLayer.GENERIC_DRIVABLE class AbstractStopLine(AbstractSurfaceMapObject): """Abstract interface for stop line objects.""" @property - def surface_type(self) -> MapSurfaceType: - # return MapSurfaceType.STOP_LINE - raise NotImplementedError + def layer(self) -> MapLayer: + return MapLayer.STOP_LINE -class AbstractRoadEdge(AbstractMapObject): +class AbstractRoadEdge(AbstractLineMapObject): """Abstract interface for road edge objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.ROAD_EDGE + def layer(self) -> MapLayer: + return MapLayer.ROAD_EDGE @property @abc.abstractmethod @@ -266,34 +292,9 @@ def polyline_3d(self) -> Polyline3D: raise NotImplementedError -class AbstractRoadLine(AbstractMapObject): +class AbstractRoadLine(AbstractLineMapObject): """Abstract interface for road line objects.""" @property - def surface_type(self) -> MapSurfaceType: - return MapSurfaceType.ROAD_LINE - - @property - @abc.abstractmethod - def polyline_3d(self) -> Polyline3D: - """ - Returns the 3D polyline of the road edge. - :return: 3D polyline - """ - raise NotImplementedError - - @property - def polyline_2d(self) -> Polyline2D: - """ - Returns the 2D polyline of the road line. - :return: 2D polyline - """ - return self.polyline_3d.polyline_2d - - @property - def polyline_se2(self) -> Polyline2D: - """ - Returns the 2D polyline of the road line in SE(2) coordinates. - :return: 2D polyline in SE(2) - """ - return self.polyline_3d.polyline_se2 + def layer(self) -> MapLayer: + return MapLayer.ROAD_LINE diff --git a/d123/dataset/maps/gpkg/gpkg_map.py b/d123/dataset/maps/gpkg/gpkg_map.py index e193c012..29df246d 100644 --- a/d123/dataset/maps/gpkg/gpkg_map.py +++ b/d123/dataset/maps/gpkg/gpkg_map.py @@ -26,7 +26,7 @@ GPKGWalkway, ) from d123.dataset.maps.gpkg.utils import load_gdf_with_geometry_columns -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer USE_ARROW: bool = True @@ -35,20 +35,20 @@ class GPKGMap(AbstractMap): def __init__(self, file_path: Path) -> None: self._file_path = file_path - self._map_object_getter: Dict[MapSurfaceType, Callable[[str], Optional[AbstractMapObject]]] = { - MapSurfaceType.LANE: self._get_lane, - MapSurfaceType.LANE_GROUP: self._get_lane_group, - MapSurfaceType.INTERSECTION: self._get_intersection, - MapSurfaceType.CROSSWALK: self._get_crosswalk, - MapSurfaceType.WALKWAY: self._get_walkway, - MapSurfaceType.CARPARK: self._get_carpark, - MapSurfaceType.GENERIC_DRIVABLE: self._get_generic_drivable, - MapSurfaceType.ROAD_EDGE: self._get_road_edge, - MapSurfaceType.ROAD_LINE: self._get_road_line, + self._map_object_getter: Dict[MapLayer, Callable[[str], Optional[AbstractMapObject]]] = { + MapLayer.LANE: self._get_lane, + MapLayer.LANE_GROUP: self._get_lane_group, + MapLayer.INTERSECTION: self._get_intersection, + MapLayer.CROSSWALK: self._get_crosswalk, + MapLayer.WALKWAY: self._get_walkway, + MapLayer.CARPARK: self._get_carpark, + MapLayer.GENERIC_DRIVABLE: self._get_generic_drivable, + MapLayer.ROAD_EDGE: self._get_road_edge, + MapLayer.ROAD_LINE: self._get_road_line, } # loaded during `.initialize()` - self._gpd_dataframes: Dict[MapSurfaceType, gpd.GeoDataFrame] = {} + self._gpd_dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} @property def map_name(self) -> str: @@ -59,7 +59,7 @@ def initialize(self) -> None: """Inherited, see superclass.""" if len(self._gpd_dataframes) == 0: available_layers = list(gpd.list_layers(self._file_path).name) - for map_layer in list(MapSurfaceType): + for map_layer in list(MapLayer): map_layer_name = map_layer.serialize() if map_layer_name in available_layers: self._gpd_dataframes[map_layer] = gpd.read_file( @@ -79,16 +79,16 @@ def _assert_initialize(self) -> None: "Checks if `.initialize()` was called, before retrieving data." assert len(self._gpd_dataframes) > 0, "GPKGMap: Call `.initialize()` before retrieving data!" - def _assert_layer_available(self, layer: MapSurfaceType) -> None: + def _assert_layer_available(self, layer: MapLayer) -> None: "Checks if layer is available." - assert layer in self.get_available_map_objects(), f"GPKGMap: MapSurfaceType {layer.name} is unavailable." + assert layer in self.get_available_map_objects(), f"GPKGMap: MapLayer {layer.name} is unavailable." - def get_available_map_objects(self) -> List[MapSurfaceType]: + def get_available_map_objects(self) -> List[MapLayer]: """Inherited, see superclass.""" self._assert_initialize() return list(self._gpd_dataframes.keys()) - def get_map_object(self, object_id: str, layer: MapSurfaceType) -> Optional[AbstractMapObject]: + def get_map_object(self, object_id: str, layer: MapLayer) -> Optional[AbstractMapObject]: """Inherited, see superclass.""" self._assert_initialize() @@ -98,17 +98,17 @@ def get_map_object(self, object_id: str, layer: MapSurfaceType) -> Optional[Abst except KeyError: raise ValueError(f"Object representation for layer: {layer.name} object: {object_id} is unavailable") - def get_all_map_objects(self, point_2d: Point2D, layer: MapSurfaceType) -> List[AbstractMapObject]: + def get_all_map_objects(self, point_2d: Point2D, layer: MapLayer) -> List[AbstractMapObject]: """Inherited, see superclass.""" raise NotImplementedError - def is_in_layer(self, point: Point2D, layer: MapSurfaceType) -> bool: + def is_in_layer(self, point: Point2D, layer: MapLayer) -> bool: """Inherited, see superclass.""" raise NotImplementedError def get_proximal_map_objects( - self, point_2d: Point2D, radius: float, layers: List[MapSurfaceType] - ) -> Dict[MapSurfaceType, List[AbstractMapObject]]: + self, point_2d: Point2D, radius: float, layers: List[MapLayer] + ) -> Dict[MapLayer, List[AbstractMapObject]]: """Inherited, see superclass.""" center_point = geom.Point(point_2d.x, point_2d.y) patch = center_point.buffer(radius) @@ -117,16 +117,16 @@ def get_proximal_map_objects( def query( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: + ) -> Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]]: supported_layers = self.get_available_map_objects() unsupported_layers = [layer for layer in layers if layer not in supported_layers] assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapSurfaceType, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]] = ( - defaultdict(list) + object_map: Dict[MapLayer, Union[List[AbstractMapObject], Dict[int, List[AbstractMapObject]]]] = defaultdict( + list ) for layer in layers: object_map[layer] = self._query_layer(geometry, layer, predicate, sort, distance) @@ -135,15 +135,15 @@ def query( def query_object_ids( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, - ) -> Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]]: + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]]]]: supported_layers = self.get_available_map_objects() unsupported_layers = [layer for layer in layers if layer not in supported_layers] assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) + object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) for layer in layers: object_map[layer] = self._query_layer_objects_ids(geometry, layer, predicate, sort, distance) return object_map @@ -151,16 +151,16 @@ def query_object_ids( def query_nearest( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layers: List[MapSurfaceType], + layers: List[MapLayer], return_all: bool = True, max_distance: Optional[float] = None, return_distance: bool = False, exclusive: bool = False, - ) -> Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]]: + ) -> Dict[MapLayer, Union[List[str], Dict[int, List[str]], Dict[int, List[Tuple[str, float]]]]]: supported_layers = self.get_available_map_objects() unsupported_layers = [layer for layer in layers if layer not in supported_layers] assert len(unsupported_layers) == 0, f"Object representation for layer(s): {unsupported_layers} is unavailable" - object_map: Dict[MapSurfaceType, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) + object_map: Dict[MapLayer, Union[List[str], Dict[int, List[str]]]] = defaultdict(list) for layer in layers: object_map[layer] = self._query_layer_nearest( geometry, layer, return_all, max_distance, return_distance, exclusive @@ -170,7 +170,7 @@ def query_nearest( def _query_layer( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapSurfaceType, + layer: MapLayer, predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, @@ -192,7 +192,7 @@ def _query_layer( def _query_layer_objects_ids( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapSurfaceType, + layer: MapLayer, predicate: Optional[str] = None, sort: bool = False, distance: Optional[float] = None, @@ -215,7 +215,7 @@ def _query_layer_objects_ids( def _query_layer_nearest( self, geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - layer: MapSurfaceType, + layer: MapLayer, return_all: bool = True, max_distance: Optional[float] = None, return_distance: bool = False, @@ -251,11 +251,11 @@ def _get_lane(self, id: str) -> Optional[GPKGLane]: return ( GPKGLane( id, - self._gpd_dataframes[MapSurfaceType.LANE], - self._gpd_dataframes[MapSurfaceType.LANE_GROUP], - self._gpd_dataframes[MapSurfaceType.INTERSECTION], + self._gpd_dataframes[MapLayer.LANE], + self._gpd_dataframes[MapLayer.LANE_GROUP], + self._gpd_dataframes[MapLayer.INTERSECTION], ) - if id in self._gpd_dataframes[MapSurfaceType.LANE]["id"].tolist() + if id in self._gpd_dataframes[MapLayer.LANE]["id"].tolist() else None ) @@ -263,11 +263,11 @@ def _get_lane_group(self, id: str) -> Optional[GPKGLaneGroup]: return ( GPKGLaneGroup( id, - self._gpd_dataframes[MapSurfaceType.LANE_GROUP], - self._gpd_dataframes[MapSurfaceType.LANE], - self._gpd_dataframes[MapSurfaceType.INTERSECTION], + self._gpd_dataframes[MapLayer.LANE_GROUP], + self._gpd_dataframes[MapLayer.LANE], + self._gpd_dataframes[MapLayer.INTERSECTION], ) - if id in self._gpd_dataframes[MapSurfaceType.LANE_GROUP]["id"].tolist() + if id in self._gpd_dataframes[MapLayer.LANE_GROUP]["id"].tolist() else None ) @@ -275,60 +275,60 @@ def _get_intersection(self, id: str) -> Optional[GPKGIntersection]: return ( GPKGIntersection( id, - self._gpd_dataframes[MapSurfaceType.INTERSECTION], - self._gpd_dataframes[MapSurfaceType.LANE], - self._gpd_dataframes[MapSurfaceType.LANE_GROUP], + self._gpd_dataframes[MapLayer.INTERSECTION], + self._gpd_dataframes[MapLayer.LANE], + self._gpd_dataframes[MapLayer.LANE_GROUP], ) - if id in self._gpd_dataframes[MapSurfaceType.INTERSECTION]["id"].tolist() + if id in self._gpd_dataframes[MapLayer.INTERSECTION]["id"].tolist() else None ) def _get_crosswalk(self, id: str) -> Optional[GPKGCrosswalk]: return ( - GPKGCrosswalk(id, self._gpd_dataframes[MapSurfaceType.CROSSWALK]) - if id in self._gpd_dataframes[MapSurfaceType.CROSSWALK]["id"].tolist() + GPKGCrosswalk(id, self._gpd_dataframes[MapLayer.CROSSWALK]) + if id in self._gpd_dataframes[MapLayer.CROSSWALK]["id"].tolist() else None ) def _get_walkway(self, id: str) -> Optional[GPKGWalkway]: return ( - GPKGWalkway(id, self._gpd_dataframes[MapSurfaceType.WALKWAY]) - if id in self._gpd_dataframes[MapSurfaceType.WALKWAY]["id"].tolist() + GPKGWalkway(id, self._gpd_dataframes[MapLayer.WALKWAY]) + if id in self._gpd_dataframes[MapLayer.WALKWAY]["id"].tolist() else None ) def _get_carpark(self, id: str) -> Optional[GPKGCarpark]: return ( - GPKGCarpark(id, self._gpd_dataframes[MapSurfaceType.CARPARK]) - if id in self._gpd_dataframes[MapSurfaceType.CARPARK]["id"].tolist() + GPKGCarpark(id, self._gpd_dataframes[MapLayer.CARPARK]) + if id in self._gpd_dataframes[MapLayer.CARPARK]["id"].tolist() else None ) def _get_generic_drivable(self, id: str) -> Optional[GPKGGenericDrivable]: return ( - GPKGGenericDrivable(id, self._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE]) - if id in self._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE]["id"].tolist() + GPKGGenericDrivable(id, self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE]) + if id in self._gpd_dataframes[MapLayer.GENERIC_DRIVABLE]["id"].tolist() else None ) def _get_road_edge(self, id: str) -> Optional[GPKGRoadEdge]: return ( - GPKGRoadEdge(id, self._gpd_dataframes[MapSurfaceType.ROAD_EDGE]) - if id in self._gpd_dataframes[MapSurfaceType.ROAD_EDGE]["id"].tolist() + GPKGRoadEdge(id, self._gpd_dataframes[MapLayer.ROAD_EDGE]) + if id in self._gpd_dataframes[MapLayer.ROAD_EDGE]["id"].tolist() else None ) def _get_road_line(self, id: str) -> Optional[GPKGRoadLine]: return ( - GPKGRoadLine(id, self._gpd_dataframes[MapSurfaceType.ROAD_LINE]) - if id in self._gpd_dataframes[MapSurfaceType.ROAD_LINE]["id"].tolist() + GPKGRoadLine(id, self._gpd_dataframes[MapLayer.ROAD_LINE]) + if id in self._gpd_dataframes[MapLayer.ROAD_LINE]["id"].tolist() else None ) # def _query_layer( # self, # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - # layer: MapSurfaceType, + # layer: MapLayer, # predicate: Optional[str] = None, # sort: bool = False, # distance: Optional[float] = None, @@ -350,7 +350,7 @@ def _get_road_line(self, id: str) -> Optional[GPKGRoadLine]: # def _query_layer_objects_ids( # self, # geometry: Union[shapely.Geometry, Iterable[shapely.Geometry]], - # layer: MapSurfaceType, + # layer: MapLayer, # predicate: Optional[str] = None, # sort: bool = False, # distance: Optional[float] = None, diff --git a/d123/dataset/maps/gpkg/gpkg_map_objects.py b/d123/dataset/maps/gpkg/gpkg_map_objects.py index 24b0fdbb..5d1a4d5b 100644 --- a/d123/dataset/maps/gpkg/gpkg_map_objects.py +++ b/d123/dataset/maps/gpkg/gpkg_map_objects.py @@ -20,6 +20,7 @@ AbstractIntersection, AbstractLane, AbstractLaneGroup, + AbstractLineMapObject, AbstractRoadEdge, AbstractRoadLine, AbstractSurfaceMapObject, @@ -33,7 +34,6 @@ class GPKGSurfaceObject(AbstractSurfaceMapObject): Base interface representation of all map objects. """ - # TODO: Extend for 3D outline def __init__(self, object_id: str, surface_df: gpd.GeoDataFrame) -> None: """ Constructor of the base surface map object type. @@ -79,6 +79,22 @@ def trimesh_mesh(self) -> trimesh.Trimesh: return trimesh_mesh +class GPKGLineObject(AbstractLineMapObject): + + def __init__(self, object_id: str, line_df: gpd.GeoDataFrame) -> None: + """ + Constructor of the base line map object type. + :param object_id: unique identifier of a line map object. + """ + super().__init__(object_id) + # TODO: add assertion if columns are available + self._object_df = line_df + + @cached_property + def _object_row(self) -> gpd.GeoSeries: + return get_row_with_value(self._object_df, "id", self.id) + + class GPKGLane(GPKGSurfaceObject, AbstractLane): def __init__( self, @@ -269,10 +285,9 @@ def trimesh_mesh(self) -> trimesh.Trimesh: return trimesh_mesh -class GPKGRoadEdge(AbstractRoadEdge): +class GPKGRoadEdge(GPKGLineObject, AbstractRoadEdge): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): - super().__init__(object_id) - self._object_df = object_df + super().__init__(object_id, object_df) @cached_property def _object_row(self) -> gpd.GeoSeries: @@ -284,10 +299,9 @@ def polyline_3d(self) -> Polyline3D: return Polyline3D.from_linestring(self._object_row.geometry) -class GPKGRoadLine(AbstractRoadLine): +class GPKGRoadLine(GPKGLineObject, AbstractRoadLine): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): - super().__init__(object_id) - self._object_df = object_df + super().__init__(object_id, object_df) @cached_property def _object_row(self) -> gpd.GeoSeries: diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index e16ccb10..0add2d7d 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -5,7 +5,7 @@ # TODO: Add stop pads or stop lines. -class MapSurfaceType(SerialIntEnum): +class MapLayer(SerialIntEnum): """ Enum for AbstractMapSurface. """ diff --git a/d123/simulation/gym/environment/helper/environment_cache.py b/d123/simulation/gym/environment/helper/environment_cache.py index 8a46df52..55fbea92 100644 --- a/d123/simulation/gym/environment/helper/environment_cache.py +++ b/d123/simulation/gym/environment/helper/environment_cache.py @@ -25,7 +25,7 @@ AbstractLaneGroup, AbstractStopLine, ) -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer from d123.simulation.gym.environment.helper.environment_area import AbstractEnvironmentArea from d123.simulation.planning.abstract_planner import PlannerInitialization, PlannerInput @@ -77,12 +77,12 @@ def __init__( def _load_cache(self) -> None: - query_map_layers = [MapSurfaceType.LANE_GROUP, MapSurfaceType.CARPARK] + query_map_layers = [MapLayer.LANE_GROUP, MapLayer.CARPARK] # FIXME: Add stop lines and crosswalks to the map layers if needed # if self.load_crosswalks: - # query_map_layers.append(MapSurfaceType.CROSSWALK) + # query_map_layers.append(MapLayer.CROSSWALK) # if self.load_stop_lines: - # query_map_layers.append(MapSurfaceType.STOP_LINE) + # query_map_layers.append(MapLayer.STOP_LINE) map_object_dict = self.map_api.query( geometry=self.environment_area.get_global_polygon(self.ego_state.center), @@ -91,7 +91,7 @@ def _load_cache(self) -> None: ) # 1. load (1.1) lane groups, (1.2) lanes, (1.3) intersections - for lane_group in map_object_dict[MapSurfaceType.LANE_GROUP]: + for lane_group in map_object_dict[MapLayer.LANE_GROUP]: lane_group: AbstractLaneGroup self.lane_groups[lane_group.id] = lane_group for lane in lane_group.lanes: @@ -101,18 +101,18 @@ def _load_cache(self) -> None: self.intersections[optional_intersection.id] = optional_intersection # 2. load car parks - for car_park in map_object_dict[MapSurfaceType.CARPARK]: + for car_park in map_object_dict[MapLayer.CARPARK]: car_park: AbstractCarpark self.car_parks[car_park.id] = car_park # FIXME: Add stop lines and crosswalks to the map layers if needed # if self.load_crosswalks: - # for crosswalk in map_object_dict[MapSurfaceType.CROSSWALK]: + # for crosswalk in map_object_dict[MapLayer.CROSSWALK]: # crosswalk: AbstractCarpark # self.crosswalks[crosswalk.id] = crosswalk # if self.load_stop_lines: - # for stop_line in map_object_dict[MapSurfaceType.STOP_LINE]: + # for stop_line in map_object_dict[MapLayer.STOP_LINE]: # stop_line: AbstractStopLine # self.stop_lines[stop_line.id] = stop_line diff --git a/d123/simulation/metrics/sim_agents/map_based.py b/d123/simulation/metrics/sim_agents/map_based.py index 513aca42..dbd9a16d 100644 --- a/d123/simulation/metrics/sim_agents/map_based.py +++ b/d123/simulation/metrics/sim_agents/map_based.py @@ -10,7 +10,7 @@ from d123.common.geometry.utils import normalize_angle from d123.dataset.maps.abstract_map import AbstractMap from d123.dataset.maps.abstract_map_objects import AbstractLane -from d123.dataset.maps.map_datatypes import MapSurfaceType +from d123.dataset.maps.map_datatypes import MapLayer MAX_LANE_CENTER_DISTANCE: Final[float] = 10.0 @@ -32,10 +32,10 @@ def _get_offroad_feature( output = map_api.query_object_ids( agent_shapely_corners, layers=[ - MapSurfaceType.INTERSECTION, - MapSurfaceType.LANE_GROUP, - MapSurfaceType.CARPARK, - MapSurfaceType.GENERIC_DRIVABLE, + MapLayer.INTERSECTION, + MapLayer.LANE_GROUP, + MapLayer.CARPARK, + MapLayer.GENERIC_DRIVABLE, ], predicate="within", ) @@ -61,7 +61,7 @@ def _get_road_center_distance_feature( def get_lane_by_id(lane_id: str, lane_dict: Dict[str, AbstractLane]) -> AbstractLane: if lane_id not in lane_dict.keys(): - lane_dict[lane_id] = map_api.get_map_object(lane_id, MapSurfaceType.LANE) + lane_dict[lane_id] = map_api.get_map_object(lane_id, MapLayer.LANE) return lane_dict[lane_id] assert agents_array.shape[-1] == len(BoundingBoxSE2Index) @@ -74,12 +74,12 @@ def get_lane_by_id(lane_id: str, lane_dict: Dict[str, AbstractLane]) -> Abstract nearest_query_output = map_api.query_nearest( agent_shapely_centers, - layers=[MapSurfaceType.LANE], + layers=[MapLayer.LANE], max_distance=MAX_LANE_CENTER_DISTANCE, return_all=True, return_distance=False, exclusive=False, - )[MapSurfaceType.LANE] + )[MapLayer.LANE] for object_idx in range(n_objects): for iteration in range(n_iterations): diff --git a/d123/training/feature_builder/smart_feature_builder.py b/d123/training/feature_builder/smart_feature_builder.py index 8a2af0ba..829f0498 100644 --- a/d123/training/feature_builder/smart_feature_builder.py +++ b/d123/training/feature_builder/smart_feature_builder.py @@ -12,7 +12,7 @@ from d123.common.geometry.line.polylines import PolylineSE2 from d123.common.geometry.transform.se2_array import convert_absolute_to_relative_se2_array from d123.common.visualization.color.default import TrafficLightStatus -from d123.dataset.maps.abstract_map import MapSurfaceType +from d123.dataset.maps.abstract_map import MapLayer from d123.dataset.maps.abstract_map_objects import ( AbstractCarpark, AbstractCrosswalk, @@ -74,10 +74,10 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. map_objects = map_api.query( map_bounding_box.shapely_polygon, layers=[ - MapSurfaceType.LANE_GROUP, - MapSurfaceType.CROSSWALK, - MapSurfaceType.CARPARK, - MapSurfaceType.GENERIC_DRIVABLE, + MapLayer.LANE_GROUP, + MapLayer.CROSSWALK, + MapLayer.CARPARK, + MapLayer.GENERIC_DRIVABLE, ], predicate="intersects", ) @@ -91,7 +91,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types: List[int] = [] # 1. Add lane - for lane_group in map_objects[MapSurfaceType.LANE_GROUP]: + for lane_group in map_objects[MapLayer.LANE_GROUP]: lane_group: AbstractLaneGroup is_intersection = lane_group.intersection is not None @@ -129,7 +129,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types.extend([int(lane_traffic_light.status)] * len(lane_traj_se2)) # 2. Crosswalks - for crosswalk in map_objects[MapSurfaceType.CROSSWALK]: + for crosswalk in map_objects[MapLayer.CROSSWALK]: crosswalk: AbstractCrosswalk crosswalk_traj_se2 = _split_segments( crosswalk.outline_3d.polyline_se2, @@ -143,7 +143,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types.extend([int(TrafficLightStatus.OFF)] * len(crosswalk_traj_se2)) # 3. Parking - for carpark in map_objects[MapSurfaceType.CARPARK]: + for carpark in map_objects[MapLayer.CARPARK]: carpark: AbstractCarpark carpark_traj_se2 = _split_segments( carpark.outline_3d.polyline_se2, @@ -157,7 +157,7 @@ def _build_map_features(scene: AbstractScene, origin: StateSE2) -> Dict[str, np. pl_light_types.extend([int(TrafficLightStatus.OFF)] * len(carpark_traj_se2)) # 4. Generic drivable - for generic_drivable in map_objects[MapSurfaceType.GENERIC_DRIVABLE]: + for generic_drivable in map_objects[MapLayer.GENERIC_DRIVABLE]: generic_drivable: AbstractGenericDrivable drivable_traj_se2 = _split_segments( generic_drivable.outline_3d.polyline_se2, diff --git a/notebooks/Untitled-2.ipynb b/notebooks/Untitled-2.ipynb deleted file mode 100644 index d4de968a..00000000 --- a/notebooks/Untitled-2.ipynb +++ /dev/null @@ -1,286 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.dataset.maps.gpkg.gpkg_map import get_map_api_from_names\n", - "\n", - "map_api = get_map_api_from_names(\"carla\", \"town05\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.common.visualization.color.default import MAP_SURFACE_CONFIG\n", - "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", - "\n", - "from typing import List\n", - "from shapely import LineString, points\n", - "import shapely\n", - "from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import get_road_edge_linestrings\n", - "\n", - "import matplotlib.pyplot as plt\n", - "\n", - "from d123.common.visualization.color.color import BLACK, TAB_10\n", - "from d123.common.visualization.color.config import PlotConfig\n", - "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", - "\n", - "lane_groups_ = map_api._gpd_dataframes[MapSurfaceType.LANE_GROUP].geometry.tolist()\n", - "car_parks_ = map_api._gpd_dataframes[MapSurfaceType.CARPARK].geometry.tolist()\n", - "generic_drivable_ = map_api._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE].geometry.tolist()\n", - "\n", - "\n", - "s = 30\n", - "fig, ax = plt.subplots(figsize=(s, s))\n", - "for polygon in lane_groups_:\n", - " add_shapely_polygon_to_ax(ax, polygon, MAP_SURFACE_CONFIG[MapSurfaceType.LANE_GROUP])\n", - "for polygon in car_parks_:\n", - " add_shapely_polygon_to_ax(ax, polygon, MAP_SURFACE_CONFIG[MapSurfaceType.CARPARK])\n", - "for polygon in generic_drivable_:\n", - " add_shapely_polygon_to_ax(ax, polygon, MAP_SURFACE_CONFIG[MapSurfaceType.WALKWAY])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "from shapely import union_all\n", - "from d123.common.geometry.occupancy_map import OccupancyMap2D\n", - "from d123.dataset.maps.abstract_map_objects import AbstractLane\n", - "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", - "\n", - "all_lane_group_ids = map_api._gpd_dataframes[MapSurfaceType.LANE_GROUP].id.tolist()\n", - "\n", - "all_lane_groups = [map_api.get_map_object(id, MapSurfaceType.LANE_GROUP) for id in all_lane_group_ids]\n", - "# lane_groups_no_intersection = [lane_group for lane_group in all_lane_groups if lane_group.intersection is None]\n", - "lane_groups_no_intersection = all_lane_groups\n", - "\n", - "occupancy_map = OccupancyMap2D.from_dict(\n", - " {lane_group.id: lane_group.shapely_polygon for lane_group in all_lane_groups}\n", - ")\n", - "\n", - "\n", - "def get_z_at_point(lane_group: AbstractLane, point: np.ndarray) -> float:\n", - " \"\"\"\n", - " Placeholder function to get the z-coordinate at a point.\n", - " This should be replaced with the actual implementation that retrieves the z-coordinate.\n", - " \"\"\"\n", - " centerlines = np.concatenate([lane.centerline.array for lane in lane_group.lanes], axis=0, dtype=np.float64)\n", - " distances_2d = np.linalg.norm(centerlines[..., :2] - point[:2], axis=-1)\n", - " closest_lane_index = np.argmin(distances_2d)\n", - " return centerlines[closest_lane_index, 2]\n", - "\n", - "\n", - "overlapping_lane_groups = {}\n", - "\n", - "for lane_group in lane_groups_no_intersection:\n", - " lane_centerlines = np.array([lane.centerline.array for lane in lane_group.lanes], dtype=np.float64)\n", - "\n", - " intersecting_lane_groups = occupancy_map.intersects(lane_group.shapely_polygon)\n", - " intersecting_lane_groups.remove(lane_group.id)\n", - " lane_group_polygon = lane_group.shapely_polygon\n", - "\n", - " for intersecting_id in intersecting_lane_groups:\n", - " intersecting_polygon = occupancy_map[intersecting_id]\n", - " intersection = intersecting_polygon.intersection(lane_group_polygon)\n", - " if intersection.geom_type != \"Polygon\":\n", - " continue\n", - "\n", - " intersecting_lane_group = map_api.get_map_object(intersecting_id, MapSurfaceType.LANE_GROUP)\n", - " z_at_query = get_z_at_point(lane_group, intersection.centroid.coords[0])\n", - " z_at_inter = get_z_at_point(intersecting_lane_group, intersection.centroid.coords[0])\n", - "\n", - " if np.abs(z_at_query - z_at_inter) < 5.0:\n", - " continue\n", - "\n", - " overlapping_lane_groups[lane_group.id] = lane_group\n", - " overlapping_lane_groups[intersecting_lane_group.id] = intersecting_lane_group\n", - "\n", - "# occupancy_map" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [ - "non_overlapping_lane_groups = [\n", - " lane_group for lane_group in lane_groups_no_intersection if lane_group.id not in overlapping_lane_groups\n", - "]\n", - "\n", - "\n", - "def get_linestring_z_at_point(linestring: LineString, point: np.ndarray) -> float:\n", - " assert linestring.has_z, \"Linestring must have z-coordinates\"\n", - " linestring_coords = np.array(list(linestring.coords), dtype=np.float64)\n", - " distances_2d = np.linalg.norm(linestring_coords[:, :2] - point[:2], axis=-1)\n", - " closest_point_index = np.argmin(distances_2d)\n", - " return float(linestring_coords[closest_point_index, 2])\n", - "\n", - "\n", - "def lift_road_edges_to_3d(ax: plt.Axes, road_edges_2d: List[LineString]):\n", - " def _find_continuous_sublists(integers: List[int]) -> List[List[int]]:\n", - " arr = np.array(integers, dtype=np.int64)\n", - " breaks = np.where(np.diff(arr) != 1)[0] + 1\n", - " splits = np.split(arr, breaks)\n", - " return [sublist.tolist() for sublist in splits]\n", - "\n", - " boundaries = []\n", - " for lane_group in non_overlapping_lane_groups:\n", - " boundaries.append(lane_group.left_boundary.linestring)\n", - " boundaries.append(lane_group.right_boundary.linestring)\n", - "\n", - " boundaries.extend(\n", - " map_api._gpd_dataframes[MapSurfaceType.CARPARK].outline.tolist()\n", - " + map_api._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE].outline.tolist()\n", - " )\n", - "\n", - " occupancy_map = OccupancyMap2D(boundaries)\n", - "\n", - " road_edges_3d: List[LineString] = []\n", - " for idx, linestring in enumerate(road_edges_2d):\n", - " # print(list(linestring.coords))\n", - " points_3d = np.array(list(linestring.coords), dtype=np.float64)\n", - "\n", - " results = occupancy_map.query_nearest(\n", - " shapely.points(points_3d[..., :2]),\n", - " max_distance=0.01,\n", - " exclusive=False,\n", - " )\n", - " for query_idx, geometry_idx in zip(*results):\n", - " intersecting_boundary = occupancy_map[occupancy_map.ids[geometry_idx]]\n", - " points_3d[query_idx, 2] = get_linestring_z_at_point(intersecting_boundary, points_3d[query_idx])\n", - "\n", - " for continuous_slice in _find_continuous_sublists(results[0]):\n", - " if len(continuous_slice) < 2:\n", - " continue\n", - " lifted_linestring = LineString(points_3d[continuous_slice])\n", - " road_edges_3d.append(lifted_linestring)\n", - " return road_edges_3d\n", - "\n", - "\n", - "drivable_polygons = (\n", - " map_api._gpd_dataframes[MapSurfaceType.LANE_GROUP].geometry.tolist()\n", - " + map_api._gpd_dataframes[MapSurfaceType.CARPARK].geometry.tolist()\n", - " + map_api._gpd_dataframes[MapSurfaceType.GENERIC_DRIVABLE].geometry.tolist()\n", - ")\n", - "\n", - "road_edges_2d = get_road_edge_linestrings(drivable_polygons, step_size=0.25, max_road_edge_length=None)\n", - "\n", - "\n", - "s = 30\n", - "fig, ax = plt.subplots(figsize=(s, s))\n", - "road_edges_3d = lift_road_edges_to_3d(ax, road_edges_2d)\n", - "\n", - "\n", - "\n", - "for lane_group in overlapping_lane_groups.values():\n", - " road_edges_3d.extend([lane_group.right_boundary.linestring, lane_group.left_boundary.linestring])\n", - "\n", - "\n", - "for idx, linestring in enumerate(road_edges_3d):\n", - " config = PlotConfig(\n", - " fill_color=TAB_10[idx % len(TAB_10)], fill_color_alpha=0.25, line_color=TAB_10[idx % len(TAB_10)]\n", - " )\n", - " # config = PlotConfig(fill_color=BLACK, fill_color_alpha=0.25, line_color=BLACK)\n", - " add_shapely_linestring_to_ax(ax, linestring, config)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "5", - "metadata": {}, - "outputs": [], - "source": [ - "def find_continuous_sublists(arr):\n", - " arr = np.array(arr)\n", - " breaks = np.where(np.diff(arr) != 1)[0] + 1\n", - " splits = np.split(arr, breaks)\n", - " return [sublist.tolist() for sublist in splits]\n", - "\n", - "# Example usage\n", - "arr = [0, 1, 2, 3, 5, 6, 7, 9, 10, 22, 23, 24]\n", - "result = find_continuous_sublists(arr)\n", - "print(result) # [[0, 1, 2, 3], [5, 6, 7]]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "6", - "metadata": {}, - "outputs": [], - "source": [ - "\n", - "s = 30\n", - "\n", - "fig, ax = plt.subplots(figsize=(s, s))\n", - "\n", - "\n", - "for idx, lane_group in enumerate(overlapping_lane_groups.values()):\n", - " config = PlotConfig(fill_color=TAB_10[idx % len(TAB_10)], fill_color_alpha=0.25, line_color=TAB_10[idx % len(TAB_10)])\n", - " add_shapely_polygon_to_ax(ax, lane_group.shapely_polygon, config)\n", - "\n", - "# for idx, polygon in enumerate(intersections):\n", - "# config = PlotConfig(fill_color=BLACK, line_color=BLACK)\n", - "# add_shapely_polygon_to_ax(ax, polygon, config)\n", - "\n", - "ax.set_aspect(\"equal\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "7", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "d123", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/notebooks/deprecated/arrow.ipynb b/notebooks/deprecated/arrow.ipynb index 0a2b6337..8fa5d9bb 100644 --- a/notebooks/deprecated/arrow.ipynb +++ b/notebooks/deprecated/arrow.ipynb @@ -32,7 +32,7 @@ "\n", "import shapely.geometry as geom\n", "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", - "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", "\n", "\n", "temp_gpkg = \"/home/daniel/d123_workspace/data/maps/carla_town01.gpkg\"\n", @@ -44,7 +44,7 @@ "print_memory_usage()\n", "\n", "\n", - "table = map_api._gpd_dataframes[MapSurfaceType.LANE]\n", + "table = map_api._gpd_dataframes[MapLayer.LANE]\n", "\n", "# import pyarrow as pa\n", "# pa.table(table).schema\n", diff --git a/notebooks/map_conversion/test_opendrive_conversion.ipynb b/notebooks/map_conversion/test_opendrive_conversion.ipynb index 9bc80a2c..6eb9ef23 100644 --- a/notebooks/map_conversion/test_opendrive_conversion.ipynb +++ b/notebooks/map_conversion/test_opendrive_conversion.ipynb @@ -129,21 +129,21 @@ "metadata": {}, "outputs": [], "source": [ - "# from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "# from d123.dataset.maps.map_datatypes import MapLayer\n", "\n", "\n", - "# str(MapSurfaceType.GENERIC_DRIVABLE).split(\".\")[-1].lower()\n", + "# str(MapLayer.GENERIC_DRIVABLE).split(\".\")[-1].lower()\n", "\n", "\n", - "# MapSurfaceType.GENERIC_DRIVABLE.name\n", + "# MapLayer.GENERIC_DRIVABLE.name\n", "\n", - "# MapSurfaceType.deserialize(MapSurfaceType.GENERIC_DRIVABLE.name)\n", + "# MapLayer.deserialize(MapLayer.GENERIC_DRIVABLE.name)\n", "\n", "\n", - "# MapSurfaceType.GENERIC_DRIVABLE.name.lower().islower()\n", + "# MapLayer.GENERIC_DRIVABLE.name.lower().islower()\n", "\n", "\n", - "# AVAILABLE_MAP_LAYERS = list(MapSurfaceType)\n", + "# AVAILABLE_MAP_LAYERS = list(MapLayer)\n", "# AVAILABLE_MAP_LAYERS\n", "\n", "# pyogrio.read_dataframe()." diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index f2a64092..0f7e077b 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -25,9 +25,9 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "splits = [\"wopd_train\"]\n", + "# splits = [\"wopd_train\"]\n", "# splits = [\"carla\"]\n", - "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"nuplan_private_test\"]\n", "# log_names = None\n", "\n", "\n", @@ -76,7 +76,7 @@ "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", "from d123.dataset.maps.abstract_map import AbstractMap\n", "from d123.dataset.maps.abstract_map_objects import AbstractLane\n", - "from d123.dataset.maps.map_datatypes import MapSurfaceType\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", "from d123.dataset.scene.abstract_scene import AbstractScene\n", "\n", "\n", @@ -142,16 +142,16 @@ " radius: float,\n", " route_lane_group_ids: Optional[List[int]] = None,\n", ") -> None:\n", - " layers: List[MapSurfaceType] = [\n", - " MapSurfaceType.LANE,\n", - " MapSurfaceType.LANE_GROUP,\n", - " MapSurfaceType.GENERIC_DRIVABLE,\n", - " MapSurfaceType.CARPARK,\n", - " MapSurfaceType.CROSSWALK,\n", - " MapSurfaceType.INTERSECTION,\n", - " MapSurfaceType.WALKWAY,\n", - " MapSurfaceType.ROAD_EDGE,\n", - " # MapSurfaceType.ROAD_LINE,\n", + " layers: List[MapLayer] = [\n", + " MapLayer.LANE,\n", + " MapLayer.LANE_GROUP,\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.CARPARK,\n", + " MapLayer.CROSSWALK,\n", + " MapLayer.INTERSECTION,\n", + " MapLayer.WALKWAY,\n", + " MapLayer.ROAD_EDGE,\n", + " # MapLayer.ROAD_LINE,\n", " ]\n", " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", " y_min, y_max = point_2d.y - radius, point_2d.y + radius\n", @@ -163,18 +163,18 @@ " for map_object in map_objects:\n", " try:\n", " if layer in [\n", - " MapSurfaceType.GENERIC_DRIVABLE,\n", - " MapSurfaceType.CARPARK,\n", - " # MapSurfaceType.CROSSWALK,\n", - " MapSurfaceType.INTERSECTION,\n", - " # MapSurfaceType.WALKWAY,\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.CARPARK,\n", + " # MapLayer.CROSSWALK,\n", + " MapLayer.INTERSECTION,\n", + " # MapLayer.WALKWAY,\n", " ]:\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", "\n", - " if layer in [MapSurfaceType.LANE_GROUP]:\n", + " if layer in [MapLayer.LANE_GROUP]:\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", "\n", - " if layer in [MapSurfaceType.LANE]:\n", + " if layer in [MapLayer.LANE]:\n", " map_object: AbstractLane\n", " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", @@ -191,12 +191,12 @@ " # fontsize=8,\n", " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", " # )\n", - " if layer in [MapSurfaceType.ROAD_EDGE]:\n", + " if layer in [MapLayer.ROAD_EDGE]:\n", " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", " edge_lengths.append(map_object.polyline_3d.linestring.length)\n", "\n", - " # if layer in [MapSurfaceType.ROAD_LINE]:\n", - " # add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_LINE_CONFIG)\n", + " if layer in [MapLayer.ROAD_LINE]:\n", + " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_LINE_CONFIG)\n", "\n", " except Exception:\n", " import traceback\n", @@ -242,8 +242,8 @@ " return fig, ax\n", "\n", "\n", - "scene_index = 8\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=100, radius=40)\n", + "scene_index = 15\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=100, radius=90)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] @@ -271,6 +271,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/notebooks/viz/viser_testing_v2.ipynb b/notebooks/viz/viser_testing_v2.ipynb deleted file mode 100644 index 60b88c98..00000000 --- a/notebooks/viz/viser_testing_v2.ipynb +++ /dev/null @@ -1,87 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": null, - "id": "0", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.common.visualization.viser.server import ViserVisualizationServer\n", - "\n", - "from pathlib import Path\n", - "from d123.dataset.scene.arrow_scene import ArrowScene" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1", - "metadata": {}, - "outputs": [], - "source": [ - "# log_name = \"Town12_Rep0_longest6_route17_06_03_21_40_12\" # Town3\n", - "log_name = \"2021.09.29.17.35.58_veh-44_01671_01819\" # Town4\n", - "log_file = Path(f\"/home/daniel/d123_workspace/data/nuplan_private_test/{log_name}.arrow\")\n", - "\n", - "# log_name = \"_Rep0_longest1_route0_07_04_10_18_47\"\n", - "# log_file = Path(f\"/home/daniel/d123_workspace/data/carla/{log_name}.arrow\")\n", - "scene = ArrowScene(log_file)\n", - "\n", - "visualization_server = ViserVisualizationServer()\n", - "\n", - "visualization_server.set_scene(scene)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "2", - "metadata": {}, - "outputs": [], - "source": [ - "ego_vehicle_state = scene.get_ego_state_at_iteration(0)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "3", - "metadata": {}, - "outputs": [], - "source": [ - "ego_vehicle_state.vehicle_parameters.height/2 - ego_vehicle_state.vehicle_parameters.height/3\n", - "\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "4", - "metadata": {}, - "outputs": [], - "source": [] - } - ], - "metadata": { - "kernelspec": { - "display_name": "d123", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.11" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} From c4cf2a3d70b039948d4afb8562e4ec8858cf5c94 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Sun, 10 Aug 2025 21:44:43 +0200 Subject: [PATCH 28/42] Add lane neighbors and more specific map subtypes. --- .../opendrive/conversion/group_collections.py | 2 - .../map/opendrive/opendrive_converter.py | 76 ++++- .../dataset_specific/kitti_360/.gitkeep | 0 .../nuplan/nuplan_map_conversion.py | 62 ++++- .../wopd/waymo_map_utils/wopd_map_utils.py | 99 ++++++- .../wopd/wopd_data_converter.py | 1 + d123/dataset/maps/abstract_map_objects.py | 33 ++- d123/dataset/maps/gpkg/gpkg_map_objects.py | 34 ++- d123/dataset/maps/map_datatypes.py | 43 +++ .../default_dataset_conversion.yaml | 4 +- .../test_nuplan_conversion.ipynb | 0 .../test_opendrive_conversion.ipynb | 0 .../deprecated/nuplan_map_dataframe.ipynb | 263 +++++++++++++++++- notebooks/tools/plot_map_sizes.ipynb | 2 +- notebooks/viz/bev_matplotlib.ipynb | 106 ++++--- notebooks/waymo_perception/map_testing.ipynb | 4 +- 16 files changed, 643 insertions(+), 86 deletions(-) create mode 100644 d123/dataset/dataset_specific/kitti_360/.gitkeep rename notebooks/{ => deprecated}/map_conversion/test_nuplan_conversion.ipynb (100%) rename notebooks/{ => deprecated}/map_conversion/test_opendrive_conversion.ipynb (100%) diff --git a/d123/dataset/conversion/map/opendrive/conversion/group_collections.py b/d123/dataset/conversion/map/opendrive/conversion/group_collections.py index 5e509859..1e978693 100644 --- a/d123/dataset/conversion/map/opendrive/conversion/group_collections.py +++ b/d123/dataset/conversion/map/opendrive/conversion/group_collections.py @@ -33,7 +33,6 @@ class OpenDriveLaneHelper: inner_border: Border outer_border: Border speed_limit_mps: Optional[float] - reverse: bool = False # lazy loaded predecessor_lane_ids: Optional[List[str]] = None @@ -213,7 +212,6 @@ def lane_section_to_lane_helpers( coeff_factor = -1.0 if side == "right" else 1.0 lane_borders = [reference_border] - for lane in lanes: lane_id = derive_lane_id(lane_group_id, lane.id) s_inner_offset = lane_section.s if len(lane_borders) == 1 else 0.0 diff --git a/d123/dataset/conversion/map/opendrive/opendrive_converter.py b/d123/dataset/conversion/map/opendrive/opendrive_converter.py index f536b7fc..0a212e9f 100644 --- a/d123/dataset/conversion/map/opendrive/opendrive_converter.py +++ b/d123/dataset/conversion/map/opendrive/opendrive_converter.py @@ -29,7 +29,7 @@ from d123.dataset.conversion.map.opendrive.id_mapping import IntIDMapping from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import get_road_edges_3d_from_gdf -from d123.dataset.maps.map_datatypes import MapLayer +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType ENABLE_WARNING: bool = False CONNECTION_DISTANCE_THRESHOLD: float = 0.1 # [m] @@ -87,6 +87,10 @@ def run(self, map_name: str) -> None: dataframes[MapLayer.GENERIC_DRIVABLE], dataframes[MapLayer.LANE_GROUP], ) + dataframes[MapLayer.ROAD_LINE] = self._extract_road_line_df( + dataframes[MapLayer.LANE], + dataframes[MapLayer.LANE_GROUP], + ) # Store dataframes map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" @@ -334,14 +338,19 @@ def _extract_lane_dataframe(self) -> gpd.GeoDataFrame: successor_ids = [] left_boundaries = [] right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] baseline_paths = [] geometries = [] - # TODO: Extract speed limit and convert to mps - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type == "driving": + for lane_group_helper in self.lane_group_helper_dict.values(): + lane_group_id = lane_group_helper.lane_group_id + lane_helpers = lane_group_helper.lane_helpers + num_lanes = len(lane_helpers) + # NOTE: Lanes are going left to right + for lane_idx, lane_helper in enumerate(lane_helpers): ids.append(lane_helper.lane_id) - lane_group_ids.append(lane_group_id_from_lane_id(lane_helper.lane_id)) + lane_group_ids.append(lane_group_id) speed_limits_mps.append(lane_helper.speed_limit_mps) predecessor_ids.append(lane_helper.predecessor_lane_ids) successor_ids.append(lane_helper.successor_lane_ids) @@ -349,6 +358,10 @@ def _extract_lane_dataframe(self) -> gpd.GeoDataFrame: right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) baseline_paths.append(shapely.LineString(lane_helper.center_polyline_3d)) geometries.append(lane_helper.shapely_polygon) + left_lane_id = lane_helpers[lane_idx - 1].lane_id if lane_idx > 0 else None + right_lane_id = lane_helpers[lane_idx + 1].lane_id if lane_idx < num_lanes - 1 else None + left_lane_ids.append(left_lane_id) + right_lane_ids.append(right_lane_id) data = pd.DataFrame( { @@ -359,6 +372,8 @@ def _extract_lane_dataframe(self) -> gpd.GeoDataFrame: "successor_ids": successor_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, "baseline_path": baseline_paths, } ) @@ -547,6 +562,11 @@ def _convert_ids_to_int( lane_df[column] = lane_df[column].apply(lambda x: lane_id_mapping.map_list(x)) lane_group_df[column] = lane_group_df[column].apply(lambda x: lane_group_id_mapping.map_list(x)) + for column in ["left_lane_id", "right_lane_id"]: + lane_df[column] = lane_df[column].apply( + lambda x: str(lane_id_mapping.str_to_int[x]) if pd.notna(x) and x is not None else x + ) + lane_df["id"] = lane_df["id"].map(lane_id_mapping.str_to_int) walkways_df["id"] = walkways_df["id"].map(walkway_id_mapping.str_to_int) carpark_df["id"] = carpark_df["id"].map(carpark_id_mapping.str_to_int) @@ -568,8 +588,52 @@ def _extract_road_edge_df( road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) ids = np.arange(len(road_edges), dtype=np.int64).tolist() + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types = [int(RoadEdgeType.ROAD_EDGE_BOUNDARY)] * len(road_edges) geometries = road_edges - return gpd.GeoDataFrame(pd.DataFrame({"id": ids}), geometry=geometries) + return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) + + def _extract_road_line_df( + self, + lane_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, + ) -> None: + + lane_group_on_intersection = { + lane_group_id: str(intersection_id) != "nan" + for lane_group_id, intersection_id in zip(lane_group_df.id.tolist(), lane_group_df.intersection_id.tolist()) + } + ids = [] + road_line_types = [] + geometries = [] + + running_id = 0 + for lane_row in lane_df.itertuples(): + on_intersection = lane_group_on_intersection.get(lane_row.lane_group_id, False) + if on_intersection: + # Skip road lines on intersections + continue + if str(lane_row.right_lane_id) == "nan": + # This is a boundary lane, e.g. a border or sidewalk + ids.append(running_id) + road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) + geometries.append(lane_row.right_boundary) + running_id += 1 + else: + # This is a regular lane + ids.append(running_id) + road_line_types.append(int(RoadLineType.BROKEN_SINGLE_WHITE)) + geometries.append(lane_row.right_boundary) + running_id += 1 + if str(lane_row.left_lane_id) == "nan": + # This is a boundary lane, e.g. a border or sidewalk + ids.append(running_id) + road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) + geometries.append(lane_row.left_boundary) + running_id += 1 + + data = pd.DataFrame({"id": ids, "road_line_type": road_line_types}) + return gpd.GeoDataFrame(data, geometry=geometries) # TODO: move this somewhere else and improve diff --git a/d123/dataset/dataset_specific/kitti_360/.gitkeep b/d123/dataset/dataset_specific/kitti_360/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py index 7f4c683b..2f0050fd 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py @@ -16,7 +16,7 @@ split_line_geometry_by_max_length, ) from d123.dataset.maps.gpkg.utils import get_all_rows_with_value, get_row_with_value -from d123.dataset.maps.map_datatypes import MapLayer +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType MAP_FILES = { "sg-one-north": "sg-one-north/9.17.1964/map.gpkg", @@ -46,8 +46,17 @@ "gen_lane_connectors_scaled_width_polygons", ] +# 0: generic lane I guess. +# 1: ending? +# 3: bike lanes. + MAX_ROAD_EDGE_LENGTH = 100.0 # meters, used to filter out very long road edges +NUPLAN_ROAD_LINE_CONVERSION = { + 0: RoadLineType.BROKEN_SINGLE_WHITE, + 2: RoadLineType.SOLID_SINGLE_WHITE, +} + class NuPlanMapConverter: def __init__(self, map_path: Path) -> None: @@ -70,6 +79,7 @@ def convert(self, map_name: str = "us-pa-pittsburgh-hazelwood") -> None: dataframes[MapLayer.CARPARK] = self._extract_carpark_dataframe() dataframes[MapLayer.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() dataframes[MapLayer.ROAD_EDGE] = self._extract_road_edge_dataframe() + dataframes[MapLayer.ROAD_LINE] = self._extract_road_line_dataframe() if not self._map_path.exists(): self._map_path.mkdir(parents=True, exist_ok=True) @@ -121,6 +131,8 @@ def _extract_nuplan_lane_dataframe(self) -> gpd.GeoDataFrame: successor_ids = [] left_boundaries = [] right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] baseline_paths = [] geometries = self._gdf["lanes_polygons"].geometry.to_list() @@ -148,6 +160,16 @@ def _extract_nuplan_lane_dataframe(self) -> gpd.GeoDataFrame: right_boundary_fid = lane_series["right_boundary_fid"] right_boundary = get_row_with_value(self._gdf["boundaries"], "fid", str(right_boundary_fid))["geometry"] + # 3. left_lane_ids, right_lane_ids + lane_index = lane_series["lane_index"] + all_group_lanes = get_all_rows_with_value( + self._gdf["lanes_polygons"], "lane_group_fid", lane_series["lane_group_fid"] + ) + left_lane_id = all_group_lanes[all_group_lanes["lane_index"] == int(lane_index) - 1]["fid"] + right_lane_id = all_group_lanes[all_group_lanes["lane_index"] == int(lane_index) + 1]["fid"] + left_lane_ids.append(left_lane_id.item() if not left_lane_id.empty else None) + right_lane_ids.append(right_lane_id.item() if not right_lane_id.empty else None) + # 3. baseline_paths baseline_path = get_row_with_value(self._gdf["baseline_paths"], "lane_fid", float(lane_id))["geometry"] @@ -167,6 +189,8 @@ def _extract_nuplan_lane_dataframe(self) -> gpd.GeoDataFrame: "successor_ids": successor_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, "baseline_path": baseline_paths, } ) @@ -228,6 +252,8 @@ def _extract_nuplan_lane_connector_dataframe(self) -> None: "successor_ids": successor_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, + "left_lane_id": [None] * len(ids), + "right_lane_id": [None] * len(ids), "baseline_path": baseline_paths, } ) @@ -397,9 +423,41 @@ def _extract_road_edge_dataframe(self) -> gpd.GeoDataFrame: road_edge_linear_rings = get_road_edge_linear_rings(drivable_polygons) road_edges = split_line_geometry_by_max_length(road_edge_linear_rings, MAX_ROAD_EDGE_LENGTH) - data = pd.DataFrame({"id": [idx for idx in range(len(road_edges))]}) + ids = [] + road_edge_types = [] + for idx in range(len(road_edges)): + ids.append(idx) + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types.append(int(RoadEdgeType.ROAD_EDGE_BOUNDARY)) + + data = pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}) return gpd.GeoDataFrame(data, geometry=road_edges) + def _extract_road_line_dataframe(self) -> gpd.GeoDataFrame: + boundaries = self._gdf["boundaries"].geometry.to_list() + fids = self._gdf["boundaries"].fid.to_list() + boundary_types = self._gdf["boundaries"].boundary_type_fid.to_list() + + ids = [] + road_line_types = [] + geometries = [] + + for idx in range(len(boundary_types)): + # NOTE @DanielDauner: We ignore boundaries of nuPlan-type, which are on intersections. + if boundary_types[idx] not in NUPLAN_ROAD_LINE_CONVERSION.keys(): + continue + ids.append(fids[idx]) + road_line_types.append(int(NUPLAN_ROAD_LINE_CONVERSION[boundary_types[idx]])) + geometries.append(boundaries[idx]) + + data = pd.DataFrame( + { + "id": ids, + "road_line_type": road_line_types, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + def flip_linestring(linestring: LineString) -> LineString: # TODO: move somewhere more appropriate or implement in Polyline2D, PolylineSE2, etc. diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py index ed023eaa..a2de1233 100644 --- a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py @@ -1,6 +1,6 @@ from collections import defaultdict from pathlib import Path -from typing import Dict, List +from typing import Dict, List, Optional import geopandas as gpd import numpy as np @@ -13,7 +13,7 @@ from d123.common.geometry.line.polylines import Polyline3D from d123.common.geometry.units import mph_to_mps from d123.dataset.dataset_specific.wopd.waymo_map_utils.womp_boundary_utils import extract_lane_boundaries -from d123.dataset.maps.map_datatypes import MapLayer +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType # TODO: # - Implement stop signs @@ -22,6 +22,24 @@ # - Implement intersections and lane group logic # - Handle lane type, e.g. TYPE_UNDEFINED = 0; TYPE_FREEWAY = 1; TYPE_SURFACE_STREET = 2; TYPE_BIKE_LANE = 3; +WAYMO_ROAD_LINE_CONVERSION = { + 0: RoadLineType.UNKNOWN, + 1: RoadLineType.BROKEN_SINGLE_WHITE, + 2: RoadLineType.SOLID_SINGLE_WHITE, + 3: RoadLineType.SOLID_DOUBLE_WHITE, + 4: RoadLineType.BROKEN_SINGLE_YELLOW, + 5: RoadLineType.BROKEN_DOUBLE_YELLOW, + 6: RoadLineType.SOLID_SINGLE_YELLOW, + 7: RoadLineType.SOLID_DOUBLE_YELLOW, + 8: RoadLineType.PASSING_DOUBLE_YELLOW, +} + +WAYMO_ROAD_EDGE_CONVERSION = { + 0: RoadEdgeType.UNKNOWN, + 1: RoadEdgeType.ROAD_EDGE_BOUNDARY, + 2: RoadEdgeType.ROAD_EDGE_MEDIAN, +} + def convert_wopd_map(frame: dataset_pb2.Frame, map_file_path: Path) -> None: @@ -35,14 +53,34 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: assert polygon.shape[1] == 3, "Polygon must have 3 coordinates (x, y, z)" return polygon + def _extract_neighbors(data) -> List[Dict[str, int]]: + neighbors = [] + for neighbor in data: + neighbors.append( + { + "lane_id": neighbor.feature_id, + "self_start_index": neighbor.self_start_index, + "self_end_index": neighbor.self_end_index, + "neighbor_start_index": neighbor.neighbor_start_index, + "neighbor_end_index": neighbor.neighbor_end_index, + } + ) + return neighbors + lanes: Dict[int, npt.NDArray[np.float64]] = {} lanes_successors = defaultdict(list) lanes_predecessors = defaultdict(list) lanes_speed_limit_mps: Dict[int, float] = {} lanes_type: Dict[int, int] = {} + lanes_left_neighbors: Dict[int, List[Dict[str, int]]] = {} + lanes_right_neighbors: Dict[int, List[Dict[str, int]]] = {} road_lines: Dict[int, npt.NDArray[np.float64]] = {} + road_lines_type: Dict[int, RoadLineType] = {} + road_edges: Dict[int, npt.NDArray[np.float64]] = {} + road_edges_type: Dict[int, int] = {} + crosswalks: Dict[int, npt.NDArray[np.float64]] = {} carparks: Dict[int, npt.NDArray[np.float64]] = {} @@ -59,16 +97,24 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: lanes_predecessors[map_feature.id].append(lane_id_) lanes_speed_limit_mps[map_feature.id] = mph_to_mps(map_feature.lane.speed_limit_mph) lanes_type[map_feature.id] = map_feature.lane.type + lanes_left_neighbors[map_feature.id] = _extract_neighbors(map_feature.lane.left_neighbors) + lanes_right_neighbors[map_feature.id] = _extract_neighbors(map_feature.lane.right_neighbors) elif map_feature.HasField("road_line"): polyline = _extract_polyline(map_feature.road_line) if polyline.ndim != 2 or polyline.shape[0] < 2: continue road_lines[map_feature.id] = polyline + road_lines_type[map_feature.id] = WAYMO_ROAD_LINE_CONVERSION.get( + map_feature.road_line.type, RoadLineType.UNKNOWN + ) elif map_feature.HasField("road_edge"): polyline = _extract_polyline(map_feature.road_edge) if polyline.ndim != 2 or polyline.shape[0] < 2: continue road_edges[map_feature.id] = polyline + road_edges_type[map_feature.id] = WAYMO_ROAD_EDGE_CONVERSION.get( + map_feature.road_edge.type, RoadEdgeType.UNKNOWN + ) elif map_feature.HasField("stop_sign"): # TODO: implement stop signs pass @@ -92,6 +138,9 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: lanes_speed_limit_mps, lane_left_boundaries_3d, lane_right_boundaries_3d, + lanes_type, + lanes_left_neighbors, + lanes_right_neighbors, ) lane_group_df = get_lane_group_df( lanes, @@ -105,8 +154,8 @@ def _extract_polygon(data) -> npt.NDArray[np.float64]: walkway_df = get_walkway_df() carpark_df = get_carpark_df(carparks) generic_drivable_df = get_generic_drivable_df() - road_edge_df = get_road_edge_df(road_edges) - road_line_df = get_road_line_df(road_lines) + road_edge_df = get_road_edge_df(road_edges, road_edges_type) + road_line_df = get_road_line_df(road_lines, road_lines_type) map_file_path.unlink(missing_ok=True) if not map_file_path.parent.exists(): @@ -130,18 +179,32 @@ def get_lane_df( lanes_speed_limit_mps: Dict[int, float], lanes_left_boundaries_3d: Dict[int, Polyline3D], lanes_right_boundaries_3d: Dict[int, Polyline3D], + lanes_type: Dict[int, int], + lanes_left_neighbors: Dict[int, List[Dict[str, int]]], + lanes_right_neighbors: Dict[int, List[Dict[str, int]]], ) -> gpd.GeoDataFrame: ids = [] + lane_types = [] lane_group_ids = [] speed_limits_mps = [] predecessor_ids = [] successor_ids = [] left_boundaries = [] right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] baseline_paths = [] geometries = [] + def _get_majority_neighbor(neighbors: List[Dict[str, int]]) -> Optional[int]: + if len(neighbors) == 0: + return None + length = { + neighbor["lane_id"]: neighbor["self_end_index"] - neighbor["self_start_index"] for neighbor in neighbors + } + return str(max(length, key=length.get)) + for lane_id, lane_centerline_array in lanes.items(): if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: continue @@ -149,12 +212,15 @@ def get_lane_df( lane_speed_limit_mps = lanes_speed_limit_mps[lane_id] if lanes_speed_limit_mps[lane_id] > 0.0 else None ids.append(lane_id) + lane_types.append(lanes_type[lane_id]) lane_group_ids.append([lane_id]) speed_limits_mps.append(lane_speed_limit_mps) predecessor_ids.append(lanes_predecessors[lane_id]) successor_ids.append(lanes_successors[lane_id]) left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + left_lane_ids.append(_get_majority_neighbor(lanes_left_neighbors[lane_id])) + right_lane_ids.append(_get_majority_neighbor(lanes_right_neighbors[lane_id])) baseline_paths.append(lane_centerline.linestring) geometry = geom.Polygon( @@ -170,12 +236,15 @@ def get_lane_df( data = pd.DataFrame( { "id": ids, + "lane_type": lane_types, "lane_group_id": lane_group_ids, "speed_limit_mps": speed_limits_mps, "predecessor_ids": predecessor_ids, "successor_ids": successor_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, "baseline_path": baseline_paths, } ) @@ -288,19 +357,33 @@ def get_generic_drivable_df() -> gpd.GeoDataFrame: return gdf -def get_road_edge_df(road_edges: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: +def get_road_edge_df( + road_edges: Dict[int, npt.NDArray[np.float64]], road_edges_type: Dict[int, RoadEdgeType] +) -> gpd.GeoDataFrame: ids = list(road_edges.keys()) geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_edges.values()] - data = pd.DataFrame({"id": ids}) + data = pd.DataFrame( + { + "id": ids, + "road_edge_type": [int(road_edge_type) for road_edge_type in road_edges_type.values()], + } + ) gdf = gpd.GeoDataFrame(data, geometry=geometries) return gdf -def get_road_line_df(road_lines: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: +def get_road_line_df( + road_lines: Dict[int, npt.NDArray[np.float64]], road_lines_type: Dict[int, RoadLineType] +) -> gpd.GeoDataFrame: ids = list(road_lines.keys()) geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_lines.values()] - data = pd.DataFrame({"id": ids}) + data = pd.DataFrame( + { + "id": ids, + "road_line_type": [int(road_line_type) for road_line_type in road_lines_type.values()], + } + ) gdf = gpd.GeoDataFrame(data, geometry=geometries) return gdf diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 2670024e..94523d89 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -167,6 +167,7 @@ def convert_wopd_tfrecord_map_to_gpkg( map_file_path = D123_MAPS_ROOT / split / f"{log_name}.gpkg" if data_converter_config.force_map_conversion or not map_file_path.exists(): + map_file_path.unlink(missing_ok=True) convert_wopd_map(initial_frame, map_file_path) return [] diff --git a/d123/dataset/maps/abstract_map_objects.py b/d123/dataset/maps/abstract_map_objects.py index e19eb679..96ce79de 100644 --- a/d123/dataset/maps/abstract_map_objects.py +++ b/d123/dataset/maps/abstract_map_objects.py @@ -7,7 +7,7 @@ import trimesh from d123.common.geometry.line.polylines import Polyline2D, Polyline3D, PolylineSE2 -from d123.dataset.maps.map_datatypes import MapLayer +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType class AbstractMapObject(abc.ABC): @@ -138,6 +138,22 @@ def right_boundary(self) -> Polyline3D: :return: returns 3D polyline """ + @property + @abc.abstractmethod + def left_lane(self) -> Optional[AbstractLane]: + """ + Property of left lane of lane. + :return: returns left lane or none, if no left lane + """ + + @property + @abc.abstractmethod + def right_lane(self) -> Optional[AbstractLane]: + """ + Property of right lane of lane. + :return: returns right lane or none, if no right lane + """ + @property @abc.abstractmethod def centerline(self) -> Polyline3D: @@ -284,12 +300,11 @@ def layer(self) -> MapLayer: @property @abc.abstractmethod - def polyline_3d(self) -> Polyline3D: + def road_edge_type(self) -> RoadEdgeType: """ - Returns the 3D polyline of the road edge. - :return: 3D polyline + Returns the road edge type. + :return: RoadEdgeType """ - raise NotImplementedError class AbstractRoadLine(AbstractLineMapObject): @@ -298,3 +313,11 @@ class AbstractRoadLine(AbstractLineMapObject): @property def layer(self) -> MapLayer: return MapLayer.ROAD_LINE + + @property + @abc.abstractmethod + def road_line_type(self) -> RoadLineType: + """ + Returns the road line type. + :return: RoadLineType + """ diff --git a/d123/dataset/maps/gpkg/gpkg_map_objects.py b/d123/dataset/maps/gpkg/gpkg_map_objects.py index 5d1a4d5b..0b6164f0 100644 --- a/d123/dataset/maps/gpkg/gpkg_map_objects.py +++ b/d123/dataset/maps/gpkg/gpkg_map_objects.py @@ -27,6 +27,7 @@ AbstractWalkway, ) from d123.dataset.maps.gpkg.utils import get_row_with_value +from d123.dataset.maps.map_datatypes import RoadEdgeType, RoadLineType class GPKGSurfaceObject(AbstractSurfaceMapObject): @@ -94,6 +95,11 @@ def __init__(self, object_id: str, line_df: gpd.GeoDataFrame) -> None: def _object_row(self) -> gpd.GeoSeries: return get_row_with_value(self._object_df, "id", self.id) + @property + def polyline_3d(self) -> Polyline3D: + """Inherited, see superclass.""" + return Polyline3D.from_linestring(self._object_row.geometry) + class GPKGLane(GPKGSurfaceObject, AbstractLane): def __init__( @@ -134,6 +140,26 @@ def right_boundary(self) -> Polyline3D: """Inherited, see superclass.""" return Polyline3D.from_linestring(self._object_row.right_boundary) + @property + def left_lane(self) -> Optional[GPKGLane]: + """Inherited, see superclass.""" + left_lane_id = self._object_row.left_lane_id + return ( + GPKGLane(left_lane_id, self._object_df, self._lane_group_df, self._intersection_df) + if left_lane_id is not None and not pd.isna(left_lane_id) + else None + ) + + @property + def right_lane(self) -> Optional[GPKGLane]: + """Inherited, see superclass.""" + right_lane_id = self._object_row.right_lane_id + return ( + GPKGLane(right_lane_id, self._object_df, self._lane_group_df, self._intersection_df) + if right_lane_id is not None and not pd.isna(right_lane_id) + else None + ) + @property def centerline(self) -> Polyline3D: """Inherited, see superclass.""" @@ -294,9 +320,9 @@ def _object_row(self) -> gpd.GeoSeries: return get_row_with_value(self._object_df, "id", self.id) @property - def polyline_3d(self) -> Polyline3D: + def road_edge_type(self) -> RoadEdgeType: """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.geometry) + return RoadEdgeType(int(self._object_row.road_edge_type)) class GPKGRoadLine(GPKGLineObject, AbstractRoadLine): @@ -308,6 +334,6 @@ def _object_row(self) -> gpd.GeoSeries: return get_row_with_value(self._object_df, "id", self.id) @property - def polyline_3d(self) -> Polyline3D: + def road_line_type(self) -> RoadLineType: """Inherited, see superclass.""" - return Polyline3D.from_linestring(self._object_row.geometry) + return RoadLineType(int(self._object_row.road_line_type)) diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index 0add2d7d..55a61486 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -20,3 +20,46 @@ class MapLayer(SerialIntEnum): STOP_LINE = 7 ROAD_EDGE = 8 ROAD_LINE = 9 + + +class LaneType(SerialIntEnum): + """ + Enum for LaneType. + NOTE: We use the lane types from Waymo. + https://github.com/waymo-research/waymo-open-dataset/blob/99a4cb3ff07e2fe06c2ce73da001f850f628e45a/src/waymo_open_dataset/protos/map.proto#L147 + """ + + UNDEFINED = 0 + FREEWAY = 1 + SURFACE_STREET = 2 + BIKE_LANE = 3 + + +class RoadEdgeType(SerialIntEnum): + """ + Enum for RoadEdgeType. + NOTE: We use the road line types from Waymo. + https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L188 + """ + + UNKNOWN = 0 + ROAD_EDGE_BOUNDARY = 1 + ROAD_EDGE_MEDIAN = 2 + + +class RoadLineType(SerialIntEnum): + """ + Enum for RoadLineType. + NOTE: We use the road line types from Waymo. + https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L208 + """ + + UNKNOWN = 0 + BROKEN_SINGLE_WHITE = 1 + SOLID_SINGLE_WHITE = 2 + SOLID_DOUBLE_WHITE = 3 + BROKEN_SINGLE_YELLOW = 4 + BROKEN_DOUBLE_YELLOW = 5 + SOLID_SINGLE_YELLOW = 6 + SOLID_DOUBLE_YELLOW = 7 + PASSING_DOUBLE_YELLOW = 8 diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 20540413..cceb2911 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,9 +15,9 @@ defaults: - default_dataset_paths - _self_ - datasets: - - nuplan_private_dataset + # - nuplan_private_dataset # - carla_dataset - # - wopd_dataset + - wopd_dataset force_log_conversion: False force_map_conversion: True diff --git a/notebooks/map_conversion/test_nuplan_conversion.ipynb b/notebooks/deprecated/map_conversion/test_nuplan_conversion.ipynb similarity index 100% rename from notebooks/map_conversion/test_nuplan_conversion.ipynb rename to notebooks/deprecated/map_conversion/test_nuplan_conversion.ipynb diff --git a/notebooks/map_conversion/test_opendrive_conversion.ipynb b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb similarity index 100% rename from notebooks/map_conversion/test_opendrive_conversion.ipynb rename to notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb diff --git a/notebooks/deprecated/nuplan_map_dataframe.ipynb b/notebooks/deprecated/nuplan_map_dataframe.ipynb index 94c372ff..113d8391 100644 --- a/notebooks/deprecated/nuplan_map_dataframe.ipynb +++ b/notebooks/deprecated/nuplan_map_dataframe.ipynb @@ -11,6 +11,16 @@ "MAP_LOCATIONS = {\"sg-one-north\", \"us-ma-boston\", \"us-nv-las-vegas-strip\", \"us-pa-pittsburgh-hazelwood\"}" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "length = {0: 1, 1: 2, 2: 3, 3: 4, 4: 0}\n", + "max(length, key=length.get)" + ] + }, { "cell_type": "code", "execution_count": null, @@ -49,7 +59,6 @@ "\n", "def load_layer(layer_name: str) -> gpd.geodataframe:\n", "\n", - "\n", " map_meta = gpd.read_file(path_to_gpkg, layer=\"meta\", engine=\"pyogrio\")\n", " projection_system = map_meta[map_meta[\"key\"] == \"projectedCoordSystem\"][\"value\"].iloc[0]\n", "\n", @@ -63,6 +72,7 @@ "\n", " return gdf_in_utm_coords\n", "\n", + "\n", "pyogrio.read_dataframe(path_to_gpkg)" ] }, @@ -72,7 +82,11 @@ "metadata": {}, "outputs": [], "source": [ - "[\n", + "import pandas as pd\n", + "\n", + "\n", + "layers = [\n", + " \"baseline_paths\",\n", " \"carpark_areas\",\n", " \"generic_drivable_areas\",\n", " \"dubins_nodes\",\n", @@ -91,18 +105,25 @@ " \"meta\",\n", "]\n", "\n", - "lanes = load_layer(\"stop_polygons\")\n", + "\n", + "def non_nan_set(series: pd.Series) -> set:\n", + " return set(series.dropna().values)\n", + "\n", + "\n", + "lanes = load_layer(\"intersections\")\n", "boundaries = load_layer(\"boundaries\")\n", "\n", "\n", "def get_lane(id: int):\n", " return lanes[lanes[\"fid\"] == str(id)]\n", "\n", + "\n", "def get_right_boundary(lane):\n", " return boundaries[boundaries[\"fid\"] == str(lane[\"right_boundary_fid\"].iloc[0])]\n", "\n", + "\n", "def get_left_boundary(lane):\n", - " return boundaries[boundaries[\"fid\"] == str(lane[\"left_boundary_fid\"].iloc[0])]\n", + " return boundaries[boundaries[\"fid\"] == str(lane[\"left_boundary_fid\"].iloc[0])]\n", "\n", "\n", "# boundaries\n", @@ -114,11 +135,12 @@ "# error_lane_df = get_lane(error_lane_fid)\n", "\n", "\n", - "\n", "# 46552,46553\n", "# test_id = 46553\n", "\n", - "lanes" + "print(lanes.keys())\n", + "non_nan_set(lanes.intersection_type_fid)\n", + "\n" ] }, { @@ -127,7 +149,114 @@ "metadata": {}, "outputs": [], "source": [ - "lanes[lanes[\"right_boundary_fid\"] == test_id]" + "baseline_paths = load_layer(\"baseline_paths\")\n", + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from networkx import center\n", + "\n", + "# Create a figure\n", + "fig, ax = plt.subplots(figsize=(15, 12))\n", + "\n", + "# Get unique path_type_fid values\n", + "path_types = baseline_paths[\"path_type_fid\"].unique()\n", + "\n", + "# Define colors for different path types\n", + "colors = [\"blue\", \"red\", \"green\", \"brown\", \"purple\", \"orange\", \"pink\"]\n", + "\n", + "# Plot each path type with a different color\n", + "target_label = 0\n", + "for i, path_type in enumerate(path_types):\n", + " # Filter baseline_paths by path_type\n", + " paths_of_type = baseline_paths[baseline_paths[\"path_type_fid\"] == path_type]\n", + "\n", + " if path_type == target_label:\n", + " alpha = 1.0\n", + " target_fids = paths_of_type.fid.tolist()\n", + " else:\n", + " alpha = 0.5\n", + "\n", + " # Plot these paths with a specific color\n", + " paths_of_type.plot(ax=ax, color=colors[i % len(colors)], label=f\"Path Type {path_type}\", alpha=alpha)\n", + "\n", + "# Add title and legend\n", + "ax.set_title(\"Baseline Paths by Path Type\", fontsize=16)\n", + "ax.legend()\n", + "ax.set_xlabel(\"X Coordinate\")\n", + "ax.set_ylabel(\"Y Coordinate\")\n", + "ax.axis(\"equal\") # Maintain aspect ratio\n", + "\n", + "center_zoom = False # Set to True to zoom into the center of the data\n", + "fid_zoom = target_fids[25]\n", + "\n", + "if center_zoom:\n", + " # Get the bounds of the data for better focusing\n", + " x_min, y_min, x_max, y_max = baseline_paths.total_bounds\n", + "\n", + " # Calculate center coordinates\n", + " center_x = (x_min + x_max) / 2\n", + " center_y = (y_min + y_max) / 2\n", + "\n", + " # Set axis limits to zoom into the center (using 30% of the total range)\n", + " range_x = x_max - x_min\n", + " range_y = y_max - y_min\n", + " zoom_factor = 2.0\n", + "\n", + " ax.set_xlim(center_x - range_x * zoom_factor / 2, center_x + range_x * zoom_factor / 2)\n", + " ax.set_ylim(center_y - range_y * zoom_factor / 2, center_y + range_y * zoom_factor / 2)\n", + "\n", + "elif fid_zoom:\n", + " # Filter to get the specific element with the given FID\n", + " specific_path = baseline_paths[baseline_paths[\"fid\"] == str(fid_zoom)]\n", + "\n", + " if not specific_path.empty:\n", + " # Get the bounds of the specific path\n", + " x_min, y_min, x_max, y_max = specific_path.total_bounds\n", + "\n", + " # Add some padding around the element\n", + " padding = 100 # meters\n", + " ax.set_xlim(x_min - padding, x_max + padding)\n", + " ax.set_ylim(y_min - padding, y_max + padding)\n", + "\n", + " # Highlight the specific element\n", + " # specific_path.plot(ax=ax, color=\"yellow\", linewidth=3, zorder=10)\n", + "\n", + " # Update title to show we're zoomed to a specific FID\n", + " ax.set_title(f\"Baseline Path - FID {fid_zoom}\", fontsize=16)\n", + " else:\n", + " print(f\"FID {fid_zoom} not found in baseline_paths\")\n", + "\n", + "\n", + "# Add a title to indicate we're looking at the center\n", + "ax.set_title(\"Baseline Paths by Path Type (Center Zoom)\", fontsize=16)\n", + "\n", + "\n", + "# Convert the specific path coordinates to WGS84 (latitude/longitude) for Google Maps\n", + "if fid_zoom and not specific_path.empty:\n", + " # Create a copy to avoid modifying the original\n", + " wgs84_path = specific_path.copy()\n", + "\n", + " # Convert from the current projection to WGS84 (EPSG:4326)\n", + " wgs84_path = wgs84_path.to_crs(\"EPSG:4326\")\n", + "\n", + " # Get the centroid of the path for easier lookup\n", + " centroid = wgs84_path.geometry.iloc[0].centroid\n", + "\n", + " # Display the coordinates\n", + " print(f\"\\nGoogle Maps coordinates for FID {fid_zoom}:\")\n", + " print(f\"Latitude: {centroid.y}, Longitude: {centroid.x}\")\n", + " print(f\"Google Maps link: https://www.google.com/maps?q={centroid.y},{centroid.x}\")\n", + "\n", + " # Add a text annotation showing coordinates on the plot\n", + " ax.annotate(\n", + " f\"Lat: {centroid.y:.6f}, Lon: {centroid.x:.6f}\",\n", + " xy=(0.05, 0.05),\n", + " xycoords=\"axes fraction\",\n", + " bbox=dict(boxstyle=\"round,pad=0.5\", fc=\"white\", alpha=0.8),\n", + " fontsize=10,\n", + " )\n", + "plt.tight_layout()\n", + "plt.show()" ] }, { @@ -136,7 +265,113 @@ "metadata": {}, "outputs": [], "source": [ - "boundaries[boundaries[\"fid\"] == str(test_id)]" + "baseline_paths = load_layer(\"baseline_paths\")\n", + "intersections = load_layer(\"intersections\") # Load the intersections layer\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from networkx import center\n", + "\n", + "# Create a figure\n", + "fig, ax = plt.subplots(figsize=(15, 12))\n", + "\n", + "# Get unique intersection types\n", + "intersection_types = intersections[\"intersection_type_fid\"].unique()\n", + "\n", + "# Define colors for different intersection types\n", + "colors = [\"blue\", \"red\", \"green\", \"brown\", \"purple\", \"orange\", \"pink\"]\n", + "\n", + "# Plot each intersection type with a different color\n", + "target_label = 2 # Target intersection type to highlight\n", + "for i, intersection_type in enumerate(intersection_types):\n", + " # Filter intersections by type\n", + " intersections_of_type = intersections[intersections[\"intersection_type_fid\"] == intersection_type]\n", + "\n", + " if intersection_type == target_label:\n", + " alpha = 1.0\n", + " target_fids = intersections_of_type.fid.tolist()\n", + " else:\n", + " alpha = 0.5\n", + "\n", + " # Plot these intersections with a specific color\n", + " intersections_of_type.plot(ax=ax, color=colors[i % len(colors)], \n", + " label=f\"Intersection Type {intersection_type}\", alpha=alpha)\n", + "\n", + "# Add title and legend\n", + "ax.set_title(\"Intersections by Type\", fontsize=16)\n", + "\n", + "ax.set_xlabel(\"X Coordinate\")\n", + "ax.set_ylabel(\"Y Coordinate\")\n", + "ax.axis(\"equal\") # Maintain aspect ratio\n", + "\n", + "center_zoom = False # Set to True to zoom into the center of the data\n", + "fid_zoom = target_fids[0] if 'target_fids' in locals() and len(target_fids) > 0 else None\n", + "\n", + "if center_zoom:\n", + " # Get the bounds of the data for better focusing\n", + " x_min, y_min, x_max, y_max = intersections.total_bounds\n", + "\n", + " # Calculate center coordinates\n", + " center_x = (x_min + x_max) / 2\n", + " center_y = (y_min + y_max) / 2\n", + "\n", + " # Set axis limits to zoom into the center (using 30% of the total range)\n", + " range_x = x_max - x_min\n", + " range_y = y_max - y_min\n", + " zoom_factor = 2.0\n", + "\n", + " ax.set_xlim(center_x - range_x * zoom_factor / 2, center_x + range_x * zoom_factor / 2)\n", + " ax.set_ylim(center_y - range_y * zoom_factor / 2, center_y + range_y * zoom_factor / 2)\n", + "\n", + "elif fid_zoom:\n", + " # Filter to get the specific intersection with the given FID\n", + " specific_intersection = intersections[intersections[\"fid\"] == fid_zoom]\n", + "\n", + " if not specific_intersection.empty:\n", + " # Get the bounds of the specific intersection\n", + " x_min, y_min, x_max, y_max = specific_intersection.total_bounds\n", + "\n", + " # Add some padding around the element\n", + " padding = 100 # meters\n", + " ax.set_xlim(x_min - padding, x_max + padding)\n", + " ax.set_ylim(y_min - padding, y_max + padding)\n", + "\n", + " # Update title to show we're zoomed to a specific FID\n", + " ax.set_title(f\"Intersection - FID {fid_zoom} (Type {specific_intersection['intersection_type_fid'].iloc[0]})\", fontsize=16)\n", + " else:\n", + " print(f\"FID {fid_zoom} not found in intersections\")\n", + "\n", + "\n", + "# Add a title if not zooming to a specific FID\n", + "if not fid_zoom:\n", + " ax.set_title(\"Intersections by Type\", fontsize=16)\n", + "\n", + "# Convert the specific intersection coordinates to WGS84 (latitude/longitude) for Google Maps\n", + "if fid_zoom and 'specific_intersection' in locals() and not specific_intersection.empty:\n", + " # Create a copy to avoid modifying the original\n", + " wgs84_intersection = specific_intersection.copy()\n", + "\n", + " # Convert from the current projection to WGS84 (EPSG:4326)\n", + " wgs84_intersection = wgs84_intersection.to_crs(\"EPSG:4326\")\n", + "\n", + " # Get the centroid of the intersection for easier lookup\n", + " centroid = wgs84_intersection.geometry.iloc[0].centroid\n", + "\n", + " # Display the coordinates\n", + " print(f\"\\nGoogle Maps coordinates for Intersection FID {fid_zoom}:\")\n", + " print(f\"Latitude: {centroid.y}, Longitude: {centroid.x}\")\n", + " print(f\"Google Maps link: https://www.google.com/maps?q={centroid.y},{centroid.x}\")\n", + "\n", + " # Add a text annotation showing coordinates on the plot\n", + " ax.annotate(\n", + " f\"Lat: {centroid.y:.6f}, Lon: {centroid.x:.6f}\",\n", + " xy=(0.05, 0.05),\n", + " xycoords=\"axes fraction\",\n", + " bbox=dict(boxstyle=\"round,pad=0.5\", fc=\"white\", alpha=0.8),\n", + " fontsize=10,\n", + " )\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "plt.show()\n" ] }, { @@ -212,9 +447,9 @@ "metadata": {}, "outputs": [], "source": [ - "\n", "import pandas as pd\n", - "lane_df = load_layer(\"generic_drivable_areas\")\n", + "\n", + "lane_df = load_layer(\"generic_drivable_areas\")\n", "lane_df" ] }, @@ -234,7 +469,6 @@ "metadata": {}, "outputs": [], "source": [ - "\n", "gen_lane_connectors_scaled_width_polygons_df = load_layer(\"boundaries\")\n", "\n", "gen_lane_connectors_scaled_width_polygons_df[gen_lane_connectors_scaled_width_polygons_df[\"fid\"] == \"17950\"]" @@ -256,7 +490,8 @@ "outputs": [], "source": [ "import pandas as pd\n", - "lane_df = load_layer(\"baseline_paths\")\n", + "\n", + "lane_df = load_layer(\"baseline_paths\")\n", "lane_df" ] }, @@ -273,6 +508,8 @@ " :return: meters per second [m/s]\n", " \"\"\"\n", " return mph / 0.44704\n", + "\n", + "\n", "mps_to_mph(6.705409029950827)" ] }, @@ -287,7 +524,7 @@ "# type = 3\n", "\n", "# for i in np.random.choice(len(geoms[type]), 10):\n", - "# ax.plot(*geoms[type][i].coords.xy, color=\"blue\")\n" + "# ax.plot(*geoms[type][i].coords.xy, color=\"blue\")" ] }, { diff --git a/notebooks/tools/plot_map_sizes.ipynb b/notebooks/tools/plot_map_sizes.ipynb index 029a63ef..4bea65c7 100644 --- a/notebooks/tools/plot_map_sizes.ipynb +++ b/notebooks/tools/plot_map_sizes.ipynb @@ -59,7 +59,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.13.3" + "version": "3.12.11" } }, "nbformat": 4, diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 0f7e077b..b3284c9e 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -25,13 +25,12 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "# splits = [\"wopd_train\"]\n", + "splits = [\"wopd_train\"]\n", "# splits = [\"carla\"]\n", - "splits = [\"nuplan_private_test\"]\n", + "# splits = [\"nuplan_private_test\"]\n", "# log_names = None\n", "\n", "\n", - "\n", "log_names = None\n", "scene_tokens = None\n", "\n", @@ -83,7 +82,7 @@ "import shapely.geometry as geom\n", "\n", "LEFT_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=BLACK,\n", + " fill_color=TAB_10[2],\n", " fill_color_alpha=1.0,\n", " line_color=TAB_10[2],\n", " line_color_alpha=0.5,\n", @@ -93,7 +92,7 @@ ")\n", "\n", "RIGHT_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=BLACK,\n", + " fill_color=TAB_10[3],\n", " fill_color_alpha=1.0,\n", " line_color=TAB_10[3],\n", " line_color_alpha=0.5,\n", @@ -104,13 +103,13 @@ "\n", "\n", "LANE_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=LIGHT_GREY,\n", + " fill_color=BLACK,\n", " fill_color_alpha=1.0,\n", " line_color=BLACK,\n", " line_color_alpha=0.0,\n", " line_width=0.0,\n", " line_style=\"-\",\n", - " zorder=1,\n", + " zorder=5,\n", ")\n", "\n", "ROAD_EDGE_CONFIG: PlotConfig = PlotConfig(\n", @@ -134,7 +133,6 @@ ")\n", "\n", "\n", - "\n", "def add_debug_map_on_ax(\n", " ax: plt.Axes,\n", " map_api: AbstractMap,\n", @@ -151,52 +149,70 @@ " 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", "\n", - " edge_lengths = []\n", + " done = False\n", " for layer, map_objects in map_objects_dict.items():\n", " for map_object in map_objects:\n", " try:\n", " if layer in [\n", - " MapLayer.GENERIC_DRIVABLE,\n", - " MapLayer.CARPARK,\n", + " # MapLayer.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", "\n", - " if layer in [MapLayer.LANE_GROUP]:\n", - " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + " # if layer in [MapLayer.LANE_GROUP]:\n", + " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", "\n", " if layer in [MapLayer.LANE]:\n", " map_object: AbstractLane\n", - " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", - " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", - " # add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", - " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", - "\n", - " # centroid = map_object.shapely_polygon.centroid\n", - " # ax.text(\n", - " # centroid.x,\n", - " # centroid.y,\n", - " # str(map_object.id),\n", - " # horizontalalignment=\"center\",\n", - " # verticalalignment=\"center\",\n", - " # fontsize=8,\n", - " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", - " # )\n", - " if layer in [MapLayer.ROAD_EDGE]:\n", - " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", - " edge_lengths.append(map_object.polyline_3d.linestring.length)\n", + " if map_object.right_lane is not None and map_object.left_lane is not None and not done:\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", + " add_shapely_polygon_to_ax(ax, map_object.right_lane.shapely_polygon, RIGHT_CONFIG)\n", + " add_shapely_polygon_to_ax(ax, map_object.left_lane.shapely_polygon, LEFT_CONFIG)\n", + " done = True\n", + " else:\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + "\n", + "\n", + " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", + " # add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", + " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", + "\n", + " # centroid = map_object.shapely_polygon.centroid\n", + " # ax.text(\n", + " # centroid.x,\n", + " # centroid.y,\n", + " # str(map_object.id),\n", + " # horizontalalignment=\"center\",\n", + " # verticalalignment=\"center\",\n", + " # fontsize=8,\n", + " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", + " # )\n", + " # if layer in [MapLayer.ROAD_EDGE]:\n", + " # add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", + " # edge_lengths.append(map_object.polyline_3d.linestring.length)\n", "\n", " if layer in [MapLayer.ROAD_LINE]:\n", - " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_LINE_CONFIG)\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)],\n", + " fill_color_alpha=1.0,\n", + " line_color=NEW_TAB_10[line_type % len(NEW_TAB_10)],\n", + " line_color_alpha=1.0,\n", + " line_width=1.5,\n", + " line_style=\"-\",\n", + " zorder=3,\n", + " )\n", + " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, plt_config)\n", "\n", " except Exception:\n", " import traceback\n", @@ -204,9 +220,6 @@ " print(f\"Error adding map object of type {layer.name} and id {map_object.id}\")\n", " traceback.print_exc()\n", "\n", - " print(f\"road edge mean: {np.mean(edge_lengths)} m\")\n", - " print(f\"road edge min: {np.min(edge_lengths)} m\")\n", - " print(f\"road edge max: {np.max(edge_lengths)} m\")\n", " ax.set_title(f\"Map: {map_api.map_name}\")\n", "\n", "\n", @@ -242,8 +255,8 @@ " return fig, ax\n", "\n", "\n", - "scene_index = 15\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=100, radius=90)\n", + "scene_index = 1\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=100, radius=100)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] @@ -254,7 +267,12 @@ "id": "3", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "map_api = scenes[1].map_api\n", + "\n", + "\n", + "map_api._gpd_dataframes[MapLayer.LANE].left_lane_id.tolist()" + ] }, { "cell_type": "code", @@ -279,6 +297,14 @@ "metadata": {}, "outputs": [], "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/notebooks/waymo_perception/map_testing.ipynb b/notebooks/waymo_perception/map_testing.ipynb index b3388863..c29fc212 100644 --- a/notebooks/waymo_perception/map_testing.ipynb +++ b/notebooks/waymo_perception/map_testing.ipynb @@ -191,9 +191,7 @@ " outline = np.array([[p.x, p.y, p.z] for p in map_feature.driveway.polygon])\n", " driveways[map_feature.id] = outline\n", "\n", - " # print(f\"Roadline: {map_feature.road_line}\")\n", - "\n", - "lanes_speed_limit_mps" + " # print(f\"Roadline: {map_feature.road_line}\")" ] }, { From de515a0c3f16f2d190eaf695ac898e6c225e1945 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Tue, 12 Aug 2025 11:54:40 +0200 Subject: [PATCH 29/42] Add lidar metadata option and supper multiple lidar sensors in dataset. --- d123/common/datatypes/sensor/camera.py | 4 +- d123/common/datatypes/sensor/lidar.py | 106 +++++++++++- d123/common/datatypes/sensor/lidar_index.py | 62 +++++++ d123/common/visualization/viser/server.py | 58 ++++--- d123/common/visualization/viser/utils.py | 91 +++++----- d123/common/visualization/viser/utils_v2.py | 47 ++---- d123/dataset/arrow/conversion.py | 23 ++- .../map/opendrive/opendrive_converter.py | 4 +- .../carla/carla_data_converter.py | 50 ++++-- .../dataset_specific/carla/load_sensor.py | 6 +- .../dataset_specific/nuplan/load_sensor.py | 6 +- .../nuplan/nuplan_data_converter.py | 60 ++++--- .../wopd/wopd_data_converter.py | 156 +++++++++++------- d123/dataset/scene/abstract_scene.py | 20 ++- d123/dataset/scene/arrow_scene.py | 32 +++- d123/dataset/scene/scene_builder.py | 6 +- .../default_dataset_conversion.yaml | 8 +- notebooks/nuplan/nuplan_sensor_loading.ipynb | 6 +- notebooks/tools/merge_videos.ipynb | 2 +- notebooks/viz/bev_matplotlib.ipynb | 88 +++++----- notebooks/viz/camera_matplotlib.ipynb | 90 +++++----- notebooks/viz/log_rendering.ipynb | 4 +- notebooks/viz/viser_testing_v2_scene.ipynb | 29 ++-- notebooks/waymo_perception/testing.ipynb | 59 ++----- 24 files changed, 643 insertions(+), 374 deletions(-) create mode 100644 d123/common/datatypes/sensor/lidar_index.py diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py index f03917f2..900f0ab3 100644 --- a/d123/common/datatypes/sensor/camera.py +++ b/d123/common/datatypes/sensor/camera.py @@ -84,7 +84,9 @@ def camera_metadata_dict_to_json(camera_metadata: Dict[CameraType, CameraMetadat :param camera_metadata: Dictionary of CameraMetadata. :return: JSON-serializable dictionary. """ - camera_metadata_dict = {str(camera_type): metadata.to_dict() for camera_type, metadata in camera_metadata.items()} + camera_metadata_dict = { + camera_type.serialize(): metadata.to_dict() for camera_type, metadata in camera_metadata.items() + } return json.dumps(camera_metadata_dict) diff --git a/d123/common/datatypes/sensor/lidar.py b/d123/common/datatypes/sensor/lidar.py index a14e6edd..fe178760 100644 --- a/d123/common/datatypes/sensor/lidar.py +++ b/d123/common/datatypes/sensor/lidar.py @@ -1,17 +1,82 @@ +from __future__ import annotations + +import json from dataclasses import dataclass +from typing import Dict, Optional, Type import numpy as np import numpy.typing as npt +from d123.common.datatypes.sensor.lidar_index import LIDAR_INDEX_REGISTRY, LiDARIndex +from d123.common.utils.enums import SerialIntEnum + + +class LiDARType(SerialIntEnum): + + LIDAR_UNKNOWN = 0 + LIDAR_MERGED = 1 + LIDAR_TOP = 2 + LIDAR_FRONT = 3 + LIDAR_SIDE_LEFT = 4 + LIDAR_SIDE_RIGHT = 5 + LIDAR_BACK = 6 + @dataclass class LiDARMetadata: - pass + + lidar_type: LiDARType + lidar_index: Type[LiDARIndex] + extrinsic: Optional[npt.NDArray[np.float64]] = None # 4x4 matrix + + # TODO: add identifier if point cloud is returned in lidar or ego frame. + + def to_dict(self) -> dict: + return { + "lidar_type": self.lidar_type.name, + "lidar_index": self.lidar_index.__name__, + "extrinsic": self.extrinsic.tolist() if self.extrinsic is not None else None, + } + + @classmethod + def from_dict(cls, json_dict: dict) -> LiDARMetadata: + lidar_type = LiDARType[json_dict["lidar_type"]] + if json_dict["lidar_index"] not in LIDAR_INDEX_REGISTRY: + raise ValueError(f"Unknown lidar index: {json_dict['lidar_index']}") + lidar_index_class = LIDAR_INDEX_REGISTRY[json_dict["lidar_index"]] + extrinsic = np.array(json_dict["extrinsic"]) if json_dict["extrinsic"] is not None else None + return cls(lidar_type=lidar_type, lidar_index=lidar_index_class, extrinsic=extrinsic) + + +def lidar_metadata_dict_to_json(lidar_metadata: Dict[LiDARType, LiDARMetadata]) -> str: + """ + Converts a dictionary of LiDARMetadata to a JSON-serializable format. + :param lidar_metadata: Dictionary of LiDARMetadata. + :return: JSON string. + """ + lidar_metadata_dict = { + lidar_type.serialize(): metadata.to_dict() for lidar_type, metadata in lidar_metadata.items() + } + return json.dumps(lidar_metadata_dict) + + +def lidar_metadata_dict_from_json(json_str: str) -> Dict[LiDARType, LiDARMetadata]: + """ + Converts a JSON string back to a dictionary of LiDARMetadata. + :param json_str: JSON string. + :return: Dictionary of LiDARMetadata. + """ + lidar_metadata_dict = json.loads(json_str) + return { + LiDARType.deserialize(lidar_type): LiDARMetadata.from_dict(metadata) + for lidar_type, metadata in lidar_metadata_dict.items() + } @dataclass class LiDAR: + metadata: LiDARMetadata point_cloud: npt.NDArray[np.float32] @property @@ -19,4 +84,41 @@ def xyz(self) -> npt.NDArray[np.float32]: """ Returns the point cloud as an Nx3 array of x, y, z coordinates. """ - return self.point_cloud[:3].T + return self.point_cloud[self.metadata.lidar_index.XYZ].T + + @property + def xy(self) -> npt.NDArray[np.float32]: + """ + Returns the point cloud as an Nx2 array of x, y coordinates. + """ + return self.point_cloud[self.metadata.lidar_index.XY].T + + @property + def intensity(self) -> Optional[npt.NDArray[np.float32]]: + """ + Returns the intensity values of the LiDAR point cloud if available. + Returns None if intensity is not part of the point cloud. + """ + if hasattr(self.metadata.lidar_index, "INTENSITY"): + return self.point_cloud[self.metadata.lidar_index.INTENSITY] + return None + + @property + def range(self) -> Optional[npt.NDArray[np.float32]]: + """ + Returns the range values of the LiDAR point cloud if available. + Returns None if range is not part of the point cloud. + """ + if hasattr(self.metadata.lidar_index, "RANGE"): + return self.point_cloud[self.metadata.lidar_index.RANGE] + return None + + @property + def elongation(self) -> Optional[npt.NDArray[np.float32]]: + """ + Returns the elongation values of the LiDAR point cloud if available. + Returns None if elongation is not part of the point cloud. + """ + if hasattr(self.metadata.lidar_index, "ELONGATION"): + return self.point_cloud[self.metadata.lidar_index.ELONGATION] + return None diff --git a/d123/common/datatypes/sensor/lidar_index.py b/d123/common/datatypes/sensor/lidar_index.py new file mode 100644 index 00000000..0df92cff --- /dev/null +++ b/d123/common/datatypes/sensor/lidar_index.py @@ -0,0 +1,62 @@ +from enum import IntEnum + +from d123.common.utils.enums import classproperty + +LIDAR_INDEX_REGISTRY = {} + + +def register_lidar_index(enum_class): + LIDAR_INDEX_REGISTRY[enum_class.__name__] = enum_class + return enum_class + + +class LiDARIndex(IntEnum): + + @classproperty + def XY(self) -> slice: + """ + Returns a slice for the XY coordinates of the LiDAR point cloud. + """ + return slice(self.X, self.Y + 1) + + @classproperty + def XYZ(self) -> slice: + """ + Returns a slice for the XYZ coordinates of the LiDAR point cloud. + """ + return slice(self.X, self.Z + 1) + + +@register_lidar_index +class DefaultLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + + +@register_lidar_index +class NuplanLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + RING = 4 + ID = 5 + + +@register_lidar_index +class CarlaLidarIndex(LiDARIndex): + X = 0 + Y = 1 + Z = 2 + INTENSITY = 3 + + +@register_lidar_index +class WopdLidarIndex(LiDARIndex): + RANGE = 0 + INTENSITY = 1 + ELONGATION = 2 + X = 3 + Y = 4 + Z = 5 diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 97c456c5..f2242b64 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -5,12 +5,11 @@ import viser from d123.common.datatypes.sensor.camera import CameraType -from d123.common.visualization.color.color import BLACK +from d123.common.datatypes.sensor.lidar import LiDARType from d123.common.visualization.viser.utils import ( get_bounding_box_meshes, get_camera_values, get_lidar_points, - get_map_lines, get_map_meshes, ) from d123.common.visualization.viser.utils_v2 import get_bounding_box_outlines @@ -34,28 +33,34 @@ LINE_WIDTH: float = 4.0 # Bounding box config: -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" # Map config: MAP_AVAILABLE: bool = True # Cameras config: -# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [ -# CameraType.CAM_F0, -# CameraType.CAM_L0, -# CameraType.CAM_R0, -# CameraType.CAM_L1, -# CameraType.CAM_R1, -# ] -# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] + +VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types -VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] CAMERA_SCALE: float = 1.0 # Lidar config: -LIDAR_AVAILABLE: bool = False +LIDAR_AVAILABLE: bool = True + +LIDAR_TYPES: List[LiDARType] = [ + LiDARType.LIDAR_MERGED, + LiDARType.LIDAR_TOP, + LiDARType.LIDAR_FRONT, + LiDARType.LIDAR_SIDE_LEFT, + LiDARType.LIDAR_SIDE_RIGHT, + LiDARType.LIDAR_BACK, +] +# LIDAR_TYPES: List[LiDARType] = [ +# LiDARType.LIDAR_TOP, +# ] LIDAR_POINT_SIZE: float = 0.05 @@ -91,7 +96,7 @@ def set_scene(self, scene: AbstractScene) -> None: num_frames = scene.get_number_of_iterations() print(scene.available_camera_types) - self.server.gui.configure_theme(control_width="large") + self.server.gui.configure_theme(dark_mode=False, control_width="large") with self.server.gui.add_folder("Playback"): server_playing = True @@ -184,7 +189,9 @@ def _(_) -> None: camera_frustum_handles[camera_type].image = camera.image if LIDAR_AVAILABLE: - gui_lidar.points = get_lidar_points(scene, gui_timestep.value) + points, colors = get_lidar_points(scene, gui_timestep.value, LIDAR_TYPES) + gui_lidar.points = points + gui_lidar.colors = colors prev_timestep = current_timestep @@ -227,10 +234,11 @@ def _(_) -> None: ) if LIDAR_AVAILABLE: + points, colors = get_lidar_points(scene, gui_timestep.value, LIDAR_TYPES) gui_lidar = self.server.scene.add_point_cloud( name="LiDAR", - points=get_lidar_points(scene, gui_timestep.value), - colors=(0.0, 0.0, 0.0), + points=points, + colors=colors, point_size=LIDAR_POINT_SIZE, point_shape="circle", ) @@ -239,7 +247,7 @@ def _(_) -> None: for name, mesh in get_map_meshes(scene).items(): self.server.scene.add_mesh_trimesh(f"/map/{name}", mesh, visible=True) - centerlines, __, __, road_edges = get_map_lines(scene) + # centerlines, __, __, road_edges = get_map_lines(scene) # for i, centerline in enumerate(centerlines): # self.server.scene.add_line_segments( # "/map/centerlines", @@ -259,13 +267,13 @@ def _(_) -> None: # colors=[[TAB_10[3].rgb]], # line_width=LINE_WIDTH, # ) - print(centerlines.shape, road_edges.shape) - self.server.scene.add_line_segments( - "/map/road_edges", - road_edges, - colors=[[BLACK.rgb]], - line_width=LINE_WIDTH, - ) + # print(centerlines.shape, road_edges.shape) + # self.server.scene.add_line_segments( + # "/map/road_edges", + # road_edges, + # colors=[[BLACK.rgb]], + # line_width=LINE_WIDTH, + # ) # Playback update loop. prev_timestep = gui_timestep.value diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 188100eb..ab87a1e3 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import List, Tuple import numpy as np import numpy.typing as npt @@ -8,6 +8,7 @@ # from d123.common.datatypes.sensor.camera_parameters import get_nuplan_camera_parameters from d123.common.datatypes.sensor.camera import Camera, CameraType +from d123.common.datatypes.sensor.lidar import LiDARType from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D @@ -232,36 +233,11 @@ def get_camera_values( ego_transform[:3, :3] = get_rotation_matrix(rear_axle) ego_transform[:3, 3] = rear_axle.point_3d.array - DEBUG = False - if DEBUG: - print("DEBUG") - - camera_to_ego = camera.extrinsic - - flip_camera = get_rotation_matrix( - StateSE3( - x=0.0, - y=0.0, - z=0.0, - roll=np.deg2rad(0.0), - pitch=np.deg2rad(90.0), - yaw=np.deg2rad(-90.0), - ) - ) - camera_to_ego[:3, :3] = camera_to_ego[:3, :3] @ flip_camera - - camera_transform = ego_transform @ camera_to_ego - - # Camera transformation in ego frame - camera_position = Point3D(*camera_transform[:3, 3]) - camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) + camera_transform = ego_transform @ camera_to_ego - else: - camera_transform = ego_transform @ camera_to_ego - - # Camera transformation in ego frame - camera_position = Point3D(*camera_transform[:3, 3]) - camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) + # Camera transformation in ego frame + camera_position = Point3D(*camera_transform[:3, 3]) + camera_rotation = Quaternion(matrix=camera_transform[:3, :3]) return camera_position, camera_rotation, camera @@ -286,19 +262,50 @@ def euler_to_quaternion_scipy(roll: float, pitch: float, yaw: float) -> npt.NDAr return quat -def get_lidar_points(scene: AbstractScene, iteration: int) -> npt.NDArray[np.float32]: - - pass +def get_lidar_points( + scene: AbstractScene, iteration: int, lidar_types: List[LiDARType] +) -> Tuple[npt.NDArray[np.float32], npt.NDArray[np.float32]]: initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) current_ego_vehicle_state = scene.get_ego_state_at_iteration(iteration) - lidar = scene.get_lidar_at_iteration(iteration) - # if scene.log_metadata.dataset == "nuplan": - # NOTE: nuPlan uses the rear axle as origin. - origin = current_ego_vehicle_state.rear_axle_se3 - points = lidar.xyz - points = convert_relative_to_absolute_points_3d_array(origin, points) - points = points - initial_ego_vehicle_state.center_se3.point_3d.array - - return points + def float_to_rgb(values: npt.NDArray[np.float32], cmap_name: str = "viridis") -> npt.NDArray[np.float32]: + """ + Converts an array of float values to RGB colors using a matplotlib colormap. + Normalizes values to [0, 1] using min-max scaling. + Returns an array of shape (N, 3) with RGB values in [0, 1]. + """ + import matplotlib.pyplot as plt + + vmin = np.min(values) + vmax = np.max(values) + if vmax > vmin: + normed = (values - vmin) / (vmax - vmin) + else: + normed = np.zeros_like(values) + cmap = plt.get_cmap(cmap_name) + rgb = cmap(normed)[:, :3] # Ignore alpha channel + return rgb.astype(np.float32) + + points_ = [] + colors_ = [] + for lidar_idx, lidar_type in enumerate(lidar_types): + if lidar_type not in scene.available_lidar_types: + continue + lidar = scene.get_lidar_at_iteration(iteration, lidar_type) + + # 1. convert points to ego frame + points = lidar.xyz + + # 2. convert points to world frame + origin = current_ego_vehicle_state.rear_axle_se3 + points = convert_relative_to_absolute_points_3d_array(origin, points) + points = points - initial_ego_vehicle_state.center_se3.point_3d.array + points_.append(points) + # colors_.append([TAB_10[lidar_idx % len(TAB_10)].rgb] * points.shape[0]) + colors_.append(float_to_rgb(lidar.intensity, cmap_name="viridis")) + + points_ = np.concatenate(points_, axis=0) if points_ else np.empty((0, 3), dtype=np.float32) + colors_ = np.concatenate(colors_, axis=0) if colors_ else np.empty((0, 3), dtype=np.float32) + + return points_, colors_ diff --git a/d123/common/visualization/viser/utils_v2.py b/d123/common/visualization/viser/utils_v2.py index ac5137c1..5db06ab8 100644 --- a/d123/common/visualization/viser/utils_v2.py +++ b/d123/common/visualization/viser/utils_v2.py @@ -19,40 +19,21 @@ def _get_bounding_box_corners(bounding_box: BoundingBoxSE3) -> npt.NDArray[np.fl """ Get the vertices of a bounding box in 3D space. """ - # TODO: apply transform over array batch instead + corner_extent_factors = { + Corners3DIndex.FRONT_LEFT_BOTTOM: Vector3D(+0.5, -0.5, -0.5), + Corners3DIndex.FRONT_RIGHT_BOTTOM: Vector3D(+0.5, +0.5, -0.5), + Corners3DIndex.BACK_RIGHT_BOTTOM: Vector3D(-0.5, +0.5, -0.5), + Corners3DIndex.BACK_LEFT_BOTTOM: Vector3D(-0.5, -0.5, -0.5), + Corners3DIndex.FRONT_LEFT_TOP: Vector3D(+0.5, -0.5, +0.5), + Corners3DIndex.FRONT_RIGHT_TOP: Vector3D(+0.5, +0.5, +0.5), + Corners3DIndex.BACK_RIGHT_TOP: Vector3D(-0.5, +0.5, +0.5), + Corners3DIndex.BACK_LEFT_TOP: Vector3D(-0.5, -0.5, +0.5), + } corners = np.zeros((len(Corners3DIndex), len(Point3DIndex)), dtype=np.float64) - corners[Corners3DIndex.FRONT_LEFT_BOTTOM] = translate_body_frame( - bounding_box.center, - Vector3D(bounding_box.length / 2, -bounding_box.width / 2, -bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.FRONT_RIGHT_BOTTOM] = translate_body_frame( - bounding_box.center, - Vector3D(bounding_box.length / 2, bounding_box.width / 2, -bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.BACK_RIGHT_BOTTOM] = translate_body_frame( - bounding_box.center, - Vector3D(-bounding_box.length / 2, bounding_box.width / 2, -bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.BACK_LEFT_BOTTOM] = translate_body_frame( - bounding_box.center, - Vector3D(-bounding_box.length / 2, -bounding_box.width / 2, -bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.FRONT_LEFT_TOP] = translate_body_frame( - bounding_box.center, - Vector3D(bounding_box.length / 2, -bounding_box.width / 2, bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.FRONT_RIGHT_TOP] = translate_body_frame( - bounding_box.center, - Vector3D(bounding_box.length / 2, bounding_box.width / 2, bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.BACK_RIGHT_TOP] = translate_body_frame( - bounding_box.center, - Vector3D(-bounding_box.length / 2, bounding_box.width / 2, bounding_box.height / 2), - ).point_3d.array - corners[Corners3DIndex.BACK_LEFT_TOP] = translate_body_frame( - bounding_box.center, - Vector3D(-bounding_box.length / 2, -bounding_box.width / 2, bounding_box.height / 2), - ).point_3d.array + bounding_box_extent = np.array([bounding_box.length, bounding_box.width, bounding_box.height], dtype=np.float64) + for idx, vec in corner_extent_factors.items(): + vector_3d = Vector3D.from_array(bounding_box_extent * vec.array) + corners[idx] = translate_body_frame(bounding_box.center, vector_3d).point_3d.array return corners diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index 6e4dbb4a..d9afba6f 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -21,7 +21,7 @@ ) from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.sensor.camera import Camera, CameraMetadata -from d123.common.datatypes.sensor.lidar import LiDAR +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters @@ -127,9 +127,16 @@ def get_camera_from_arrow_table( ) -def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int, log_metadata: LogMetadata) -> LiDAR: - assert "lidar" in arrow_table.schema.names, '"lidar" field not found in Arrow table schema.' - lidar_data = arrow_table["lidar"][index].as_py() +def get_lidar_from_arrow_table( + arrow_table: pa.Table, + index: int, + lidar_metadata: LiDARMetadata, + log_metadata: LogMetadata, +) -> LiDAR: + assert ( + lidar_metadata.lidar_type.serialize() in arrow_table.schema.names + ), f'"{lidar_metadata.lidar_type.serialize()}" field not found in Arrow table schema.' + lidar_data = arrow_table[lidar_metadata.lidar_type.serialize()][index].as_py() if isinstance(lidar_data, str): sensor_root = DATASET_SENSOR_ROOT[log_metadata.dataset] full_lidar_path = sensor_root / lidar_data @@ -139,18 +146,20 @@ def get_lidar_from_arrow_table(arrow_table: pa.Table, index: int, log_metadata: if log_metadata.dataset == "nuplan": from d123.dataset.dataset_specific.nuplan.load_sensor import load_nuplan_lidar_from_path - lidar = load_nuplan_lidar_from_path(full_lidar_path) + lidar = load_nuplan_lidar_from_path(full_lidar_path, lidar_metadata) elif log_metadata.dataset == "carla": from d123.dataset.dataset_specific.carla.load_sensor import load_carla_lidar_from_path - lidar = load_carla_lidar_from_path(full_lidar_path) + lidar = load_carla_lidar_from_path(full_lidar_path, lidar_metadata) + elif log_metadata.dataset == "wopd": + raise NotImplementedError else: raise NotImplementedError(f"Loading LiDAR data for dataset {log_metadata.dataset} is not implemented.") else: if log_metadata.dataset == "wopd": lidar_data = np.array(lidar_data, dtype=np.float64) - lidar = LiDAR(point_cloud=lidar_data[:, 3:].T) + lidar = LiDAR(metadata=lidar_metadata, point_cloud=lidar_data.T) else: raise NotImplementedError("Only string file paths for lidar data are supported.") diff --git a/d123/dataset/conversion/map/opendrive/opendrive_converter.py b/d123/dataset/conversion/map/opendrive/opendrive_converter.py index 0a212e9f..b9dd58c2 100644 --- a/d123/dataset/conversion/map/opendrive/opendrive_converter.py +++ b/d123/dataset/conversion/map/opendrive/opendrive_converter.py @@ -613,7 +613,7 @@ def _extract_road_line_df( if on_intersection: # Skip road lines on intersections continue - if str(lane_row.right_lane_id) == "nan": + if str(lane_row.right_lane_id) in ["nan", "None"]: # This is a boundary lane, e.g. a border or sidewalk ids.append(running_id) road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) @@ -625,7 +625,7 @@ def _extract_road_line_df( road_line_types.append(int(RoadLineType.BROKEN_SINGLE_WHITE)) geometries.append(lane_row.right_boundary) running_id += 1 - if str(lane_row.left_lane_id) == "nan": + if str(lane_row.left_lane_id) in ["nan", "None"]: # This is a boundary lane, e.g. a border or sidewalk ids.append(running_id) road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index eb37ff86..0bdc33d7 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -13,6 +13,8 @@ from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar_index import CarlaLidarIndex from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import get_carla_lincoln_mkz_2020_parameters from d123.common.geometry.base import Point2D, Point3D @@ -169,6 +171,11 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> map_name = first_log_dict["location"] map_api = get_map_api_from_names("carla", map_name) + metadata = _get_metadata(map_name, str(log_path.stem)) + vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() + camera_metadata = get_carla_camera_metadata(first_log_dict) + lidar_metadata = get_carla_lidar_metadata(first_log_dict) + schema_column_list = [ ("token", pa.string()), ("timestamp", pa.int64()), @@ -183,14 +190,15 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> ("route_lane_group_ids", pa.list_(pa.int64())), ] if data_converter_config.lidar_store_option is not None: - if data_converter_config.lidar_store_option == "path": - schema_column_list.append(("lidar", pa.string())) - elif data_converter_config.lidar_store_option == "binary": - raise NotImplementedError("Binary lidar storage is not implemented.") + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + schema_column_list.append((lidar_type.serialize(), pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") # TODO: Adjust how cameras are added if data_converter_config.camera_store_option is not None: - for camera_type in CARLA_CAMERA_TYPES: + for camera_type in camera_metadata.keys(): if data_converter_config.camera_store_option == "path": schema_column_list.append((camera_type.serialize(), pa.string())) schema_column_list.append( @@ -200,14 +208,12 @@ def convert_log_internal(args: List[Dict[str, Union[List[str], List[Path]]]]) -> raise NotImplementedError("Binary camera storage is not implemented.") recording_schema = pa.schema(schema_column_list) - metadata = _get_metadata(map_name, str(log_path.stem)) - vehicle_parameters = get_carla_lincoln_mkz_2020_parameters() - camera_metadata = get_carla_camera_metadata(first_log_dict) recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), } ) @@ -236,7 +242,7 @@ def _get_metadata(location: str, log_name: str) -> LogMetadata: ) -def get_carla_camera_metadata(first_log_dict: Dict[str, Any]) -> Dict[str, CameraMetadata]: +def get_carla_camera_metadata(first_log_dict: Dict[str, Any]) -> Dict[CameraType, CameraMetadata]: # FIXME: This is a placeholder function to return camera metadata. @@ -245,7 +251,7 @@ def get_carla_camera_metadata(first_log_dict: Dict[str, Any]) -> Dict[str, Camer dtype=np.float64, ) camera_metadata = { - CameraType.CAM_F0.serialize(): CameraMetadata( + CameraType.CAM_F0: CameraMetadata( camera_type=CameraType.CAM_F0, width=1024, height=512, @@ -256,6 +262,19 @@ def get_carla_camera_metadata(first_log_dict: Dict[str, Any]) -> Dict[str, Camer return camera_metadata +def get_carla_lidar_metadata(first_log_dict: Dict[str, Any]) -> Dict[LiDARType, LiDARMetadata]: + + # TODO: add lidar extrinsic + lidar_metadata = { + LiDARType.LIDAR_TOP: LiDARMetadata( + lidar_type=LiDARType.LIDAR_TOP, + lidar_index=CarlaLidarIndex, + extrinsic=None, + ) + } + return lidar_metadata + + def _write_recording_table( bounding_box_paths: List[Path], map_api: AbstractMap, @@ -297,7 +316,9 @@ def _write_recording_table( "route_lane_group_ids": [route_lane_group_ids], } if data_converter_config.lidar_store_option is not None: - row_data["lidar"] = [_extract_lidar(log_name, sample_name, data_converter_config)] + lidar_data_dict = _extract_lidar(log_name, sample_name, data_converter_config) + for lidar_type, lidar_data in lidar_data_dict.items(): + row_data[lidar_type.serialize()] = [lidar_data] if data_converter_config.camera_store_option is not None: camera_data_dict = _extract_cameras(data, log_name, sample_name, data_converter_config) @@ -410,7 +431,9 @@ def _extract_cameras( return camera_dict -def _extract_lidar(log_name: str, sample_name: str, data_converter_config: DataConverterConfig) -> Optional[str]: +def _extract_lidar( + log_name: str, sample_name: str, data_converter_config: DataConverterConfig +) -> Dict[LiDARType, Optional[str]]: lidar: Optional[str] = None lidar_full_path = CARLA_DATA_ROOT / "sensor_blobs" / log_name / "lidar" / f"{sample_name}.npy" @@ -421,4 +444,5 @@ def _extract_lidar(log_name: str, sample_name: str, data_converter_config: DataC raise NotImplementedError("Binary lidar storage is not implemented.") else: raise FileNotFoundError(f"LiDAR file not found: {lidar_full_path}") - return lidar + + return {LiDARType.LIDAR_TOP: lidar} if lidar else None diff --git a/d123/dataset/dataset_specific/carla/load_sensor.py b/d123/dataset/dataset_specific/carla/load_sensor.py index 63f5ee35..56ff68fa 100644 --- a/d123/dataset/dataset_specific/carla/load_sensor.py +++ b/d123/dataset/dataset_specific/carla/load_sensor.py @@ -2,9 +2,9 @@ import numpy as np -from d123.common.datatypes.sensor.lidar import LiDAR +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata -def load_carla_lidar_from_path(filepath: Path) -> LiDAR: +def load_carla_lidar_from_path(filepath: Path, lidar_metadata: LiDARMetadata) -> LiDAR: assert filepath.exists(), f"LiDAR file not found: {filepath}" - return LiDAR(point_cloud=np.load(filepath)) + return LiDAR(metadata=lidar_metadata, point_cloud=np.load(filepath)) diff --git a/d123/dataset/dataset_specific/nuplan/load_sensor.py b/d123/dataset/dataset_specific/nuplan/load_sensor.py index dc699b2c..0bbdc406 100644 --- a/d123/dataset/dataset_specific/nuplan/load_sensor.py +++ b/d123/dataset/dataset_specific/nuplan/load_sensor.py @@ -3,14 +3,14 @@ from nuplan.database.utils.pointclouds.lidar import LidarPointCloud -from d123.common.datatypes.sensor.lidar import LiDAR +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata -def load_nuplan_lidar_from_path(filepath: Path) -> LiDAR: +def load_nuplan_lidar_from_path(filepath: Path, lidar_metadata: LiDARMetadata) -> LiDAR: assert filepath.exists(), f"LiDAR file not found: {filepath}" with open(filepath, "rb") as fp: buffer = io.BytesIO(fp.read()) - return LiDAR(point_cloud=LidarPointCloud.from_buffer(buffer, "pcd").points) + return LiDAR(metadata=lidar_metadata, point_cloud=LidarPointCloud.from_buffer(buffer, "pcd").points) # def load_camera_from_path(filename: str, metadata: CameraMetadata) -> Camera: diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index 012b169c..b4b9844a 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -24,6 +24,8 @@ from d123.common.datatypes.detection.detection import TrafficLightStatus from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar_index import NuplanLidarIndex from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( @@ -191,6 +193,17 @@ def convert_nuplan_log_to_arrow( if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) + metadata = LogMetadata( + dataset="nuplan", + log_name=log_db.log_name, + location=log_db.log.map_version, + timestep_seconds=TARGET_DT, + map_has_z=False, + ) + vehicle_parameters = get_nuplan_pacifica_parameters() + camera_metadata = get_nuplan_camera_metadata(log_path) + lidar_metadata = get_nuplan_lidar_metadata(log_db) + schema_column_list = [ ("token", pa.string()), ("timestamp", pa.int64()), @@ -205,14 +218,14 @@ def convert_nuplan_log_to_arrow( ("route_lane_group_ids", pa.list_(pa.int64())), ] if data_converter_config.lidar_store_option is not None: - if data_converter_config.lidar_store_option == "path": - schema_column_list.append(("lidar", pa.string())) - elif data_converter_config.lidar_store_option == "binary": - raise NotImplementedError("Binary lidar storage is not implemented.") + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + schema_column_list.append((lidar_type.serialize(), pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") - # TODO: Adjust how cameras are added if data_converter_config.camera_store_option is not None: - for camera_type in NUPLAN_CAMERA_TYPES.keys(): + for camera_type in camera_metadata.keys(): if data_converter_config.camera_store_option == "path": schema_column_list.append((camera_type.serialize(), pa.string())) schema_column_list.append( @@ -223,20 +236,12 @@ def convert_nuplan_log_to_arrow( raise NotImplementedError("Binary camera storage is not implemented.") recording_schema = pa.schema(schema_column_list) - metadata = LogMetadata( - dataset="nuplan", - log_name=log_db.log_name, - location=log_db.log.map_version, - timestep_seconds=TARGET_DT, - map_has_z=False, - ) - vehicle_parameters = get_nuplan_pacifica_parameters() - camera_metadata = get_nuplan_camera_metadata(log_path) recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), } ) @@ -249,7 +254,7 @@ def convert_nuplan_log_to_arrow( return [] -def get_nuplan_camera_metadata(log_path: Path) -> Dict[str, CameraMetadata]: +def get_nuplan_camera_metadata(log_path: Path) -> Dict[CameraType, CameraMetadata]: def _get_camera_metadata(camera_type: CameraType) -> CameraMetadata: cam = list(get_cameras(log_path, [str(NUPLAN_CAMERA_TYPES[camera_type].value)]))[0] @@ -267,11 +272,21 @@ def _get_camera_metadata(camera_type: CameraType) -> CameraMetadata: log_cam_infos: Dict[str, CameraMetadata] = {} for camera_type in NUPLAN_CAMERA_TYPES.keys(): - log_cam_infos[camera_type.serialize()] = _get_camera_metadata(camera_type) + log_cam_infos[camera_type] = _get_camera_metadata(camera_type) return log_cam_infos +def get_nuplan_lidar_metadata(log_db: NuPlanDB) -> Dict[LiDARType, LiDARMetadata]: + metadata: Dict[LiDARType, LiDARMetadata] = {} + metadata[LiDARType.LIDAR_MERGED] = LiDARMetadata( + lidar_type=LiDARType.LIDAR_MERGED, + lidar_index=NuplanLidarIndex, + extrinsic=None, # NOTE: LiDAR extrinsic are unknown + ) + return metadata + + def _write_recording_table( log_db: NuPlanDB, recording_schema: pa.schema, @@ -314,7 +329,12 @@ def _write_recording_table( } if data_converter_config.lidar_store_option is not None: - row_data["lidar"] = [_extract_lidar(lidar_pc, data_converter_config)] + lidar_data_dict = _extract_lidar(lidar_pc, data_converter_config) + for lidar_type, lidar_data in lidar_data_dict.items(): + if lidar_data is not None: + row_data[lidar_type.serialize()] = [lidar_data] + else: + row_data[lidar_type.serialize()] = [None] if data_converter_config.camera_store_option is not None: camera_data_dict = _extract_camera(log_db, lidar_pc, source_log_path, data_converter_config) @@ -470,11 +490,11 @@ def _extract_camera( return camera_dict -def _extract_lidar(lidar_pc: LidarPc, data_converter_config: DataConverterConfig) -> Optional[str]: +def _extract_lidar(lidar_pc: LidarPc, data_converter_config: DataConverterConfig) -> Dict[LiDARType, Optional[str]]: lidar: Optional[str] = None lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / lidar_pc.filename if lidar_full_path.exists(): lidar = lidar_pc.filename - return lidar + return {LiDARType.LIDAR_MERGED: lidar} diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 94523d89..35c517de 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -5,7 +5,7 @@ from dataclasses import asdict from functools import partial from pathlib import Path -from typing import Any, Dict, Final, List, Literal, Optional, Tuple, Union +from typing import Any, Dict, Final, List, Literal, Tuple, Union import numpy as np import numpy.typing as npt @@ -18,6 +18,8 @@ from d123.common.datatypes.detection.detection import TrafficLightStatus from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar_index import WopdLidarIndex from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import get_wopd_pacifica_parameters from d123.common.geometry.base import Point3D, StateSE3 @@ -35,7 +37,7 @@ D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) TARGET_DT: Final[float] = 0.1 -SORT_BY_TIMESTAMP: Final[bool] = True +SORT_BY_TIMESTAMP: Final[bool] = False NUPLAN_TRAFFIC_STATUS_DICT: Final[Dict[str, TrafficLightStatus]] = { "green": TrafficLightStatus.GREEN, @@ -44,7 +46,7 @@ } # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/label.proto#L63 -WOPD_DETECTION_NAME_DICT = { +WOPD_DETECTION_NAME_DICT: Dict[int, DetectionType] = { 0: DetectionType.GENERIC_OBJECT, # TYPE_UNKNOWN 1: DetectionType.VEHICLE, # TYPE_VEHICLE 2: DetectionType.PEDESTRIAN, # TYPE_PEDESTRIAN @@ -53,7 +55,7 @@ } # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L50 -WOPD_CAMERA_TYPES = { +WOPD_CAMERA_TYPES: Dict[int, CameraType] = { 1: CameraType.CAM_F0, # front_camera 2: CameraType.CAM_L0, # front_left_camera 3: CameraType.CAM_R0, # front_right_camera @@ -61,6 +63,16 @@ 5: CameraType.CAM_R1, # right_camera } +# https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L66 +WOPD_LIDAR_TYPES: Dict[int, LiDARType] = { + 0: LiDARType.LIDAR_UNKNOWN, # UNKNOWN + 1: LiDARType.LIDAR_TOP, # TOP + 2: LiDARType.LIDAR_FRONT, # FRONT + 3: LiDARType.LIDAR_SIDE_LEFT, # SIDE_LEFT + 4: LiDARType.LIDAR_SIDE_RIGHT, # SIDE_RIGHT + 5: LiDARType.LIDAR_BACK, # REAR +} + WOPD_DATA_ROOT = Path("/media/nvme1/waymo_perception") # TODO: set as environment variable # Whether to use ego or zero roll and pitch values for bounding box detections (after global conversion) @@ -198,6 +210,17 @@ def convert_wopd_tfrecord_log_to_arrow( if not log_file_path.parent.exists(): log_file_path.parent.mkdir(parents=True, exist_ok=True) + metadata = LogMetadata( + dataset="wopd", + log_name=log_name, + location=None, # TODO: implement map name + timestep_seconds=TARGET_DT, # TODO: Check if correct. Maybe not hardcode + map_has_z=True, + ) + vehicle_parameters = get_wopd_pacifica_parameters() + camera_metadata = get_wopd_camera_metadata(initial_frame, data_converter_config) + lidar_metadata = get_wopd_lidar_metadata(initial_frame, data_converter_config) + schema_column_list = [ ("token", pa.string()), ("timestamp", pa.int64()), @@ -211,15 +234,9 @@ def convert_wopd_tfrecord_log_to_arrow( ("scenario_tag", pa.list_(pa.string())), ("route_lane_group_ids", pa.list_(pa.int64())), ] - if data_converter_config.lidar_store_option is not None: - if data_converter_config.lidar_store_option == "path": - raise NotImplementedError("Filepath lidar storage is not implemented.") - elif data_converter_config.lidar_store_option == "binary": - schema_column_list.append(("lidar", pa.list_(pa.list_(pa.float32(), 6)))) - # TODO: Adjust how cameras are added if data_converter_config.camera_store_option is not None: - for camera_type in WOPD_CAMERA_TYPES.values(): + for camera_type in camera_metadata.keys(): if data_converter_config.camera_store_option == "path": raise NotImplementedError("Path camera storage is not implemented.") elif data_converter_config.camera_store_option == "binary": @@ -228,21 +245,22 @@ def convert_wopd_tfrecord_log_to_arrow( (f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4)) ) + if data_converter_config.lidar_store_option is not None: + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + raise NotImplementedError("Filepath lidar storage is not implemented.") + elif data_converter_config.lidar_store_option == "binary": + schema_column_list.append( + (lidar_type.serialize(), pa.list_(pa.list_(pa.float32(), len(WopdLidarIndex)))) + ) + recording_schema = pa.schema(schema_column_list) - metadata = LogMetadata( - dataset="wopd", - log_name=log_name, - location=None, # TODO: implement map name - timestep_seconds=TARGET_DT, # TODO: Check if correct. Maybe not hardcode - map_has_z=True, - ) - vehicle_parameters = get_wopd_pacifica_parameters() - camera_metadata = get_wopd_camera_metadata(initial_frame) recording_schema = recording_schema.with_metadata( { "log_metadata": json.dumps(asdict(metadata)), "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), } ) @@ -258,28 +276,53 @@ def convert_wopd_tfrecord_log_to_arrow( return [] -def get_wopd_camera_metadata(initial_frame: dataset_pb2.Frame) -> Dict[str, CameraMetadata]: +def get_wopd_camera_metadata( + initial_frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Dict[CameraType, CameraMetadata]: + + cam_metadatas: Dict[CameraType, CameraMetadata] = {} + if data_converter_config.camera_store_option is not None: + for calibration in initial_frame.context.camera_calibrations: + camera_type = WOPD_CAMERA_TYPES[calibration.name] + + # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L96 + # https://github.com/waymo-research/waymo-open-dataset/issues/834#issuecomment-2134995440 + fx, fy, cx, cy, k1, k2, p1, p2, k3 = calibration.intrinsic + _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) + _distortions = np.array([k1, k2, p1, p2, k3]) + + if camera_type in WOPD_CAMERA_TYPES.values(): + cam_metadatas[camera_type] = CameraMetadata( + camera_type=camera_type, + width=calibration.width, + height=calibration.height, + intrinsic=_intrinsics, + distortion=_distortions, + ) - cam_metadatas: Dict[str, CameraMetadata] = {} - for calibration in initial_frame.context.camera_calibrations: - camera_type = WOPD_CAMERA_TYPES[calibration.name] + return cam_metadatas + + +def get_wopd_lidar_metadata( + initial_frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig +) -> Dict[LiDARType, LiDARMetadata]: - # https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/dataset.proto#L96 - # https://github.com/waymo-research/waymo-open-dataset/issues/834#issuecomment-2134995440 - fx, fy, cx, cy, k1, k2, p1, p2, k3 = calibration.intrinsic - _intrinsics = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]]) - _distortions = np.array([k1, k2, p1, p2, k3]) - - if camera_type in WOPD_CAMERA_TYPES.values(): - cam_metadatas[camera_type.serialize()] = CameraMetadata( - camera_type=camera_type, - width=calibration.width, - height=calibration.height, - intrinsic=_intrinsics, - distortion=_distortions, + laser_metadatas: Dict[LiDARType, LiDARMetadata] = {} + if data_converter_config.lidar_store_option is not None: + for laser_calibration in initial_frame.context.laser_calibrations: + lidar_type = WOPD_LIDAR_TYPES[laser_calibration.name] + extrinsic = ( + np.array(laser_calibration.extrinsic.transform, dtype=np.float64).reshape(4, 4) + if laser_calibration.extrinsic + else None + ) + laser_metadatas[lidar_type] = LiDARMetadata( + lidar_type=lidar_type, + lidar_index=WopdLidarIndex, + extrinsic=extrinsic, ) - return cam_metadatas + return laser_metadatas def _write_recording_table( @@ -300,12 +343,6 @@ def _write_recording_table( frame.ParseFromString(data.numpy()) (detections_state, detections_velocity, detections_token, detections_types) = _extract_detections(frame) - # traffic_light_ids, traffic_light_types = _extract_traffic_lights(log_db, lidar_pc_token) - # route_lane_group_ids = [ - # int(roadblock_id) - # for roadblock_id in str(lidar_pc.scene.roadblock_ids).split(" ") - # if len(roadblock_id) > 0 - # ] # TODO: Implement traffic light extraction traffic_light_ids = [] @@ -328,7 +365,12 @@ def _write_recording_table( # TODO: Implement lidar extraction if data_converter_config.lidar_store_option is not None: - row_data["lidar"] = [_extract_lidar(frame, data_converter_config).tolist()] + lidar_data_dict = _extract_lidar(frame, data_converter_config) + for lidar_type, lidar_data in lidar_data_dict.items(): + if lidar_data is not None: + row_data[lidar_type.serialize()] = [lidar_data.tolist()] + else: + row_data[lidar_type.serialize()] = [None] if data_converter_config.camera_store_option is not None: camera_data_dict = _extract_camera(frame, data_converter_config) @@ -350,15 +392,6 @@ def _write_recording_table( write_arrow_table(recording_table, log_file_path) -# def _get_ego_pose_se3(frame: dataset_pb2.Frame) -> StateSE3: -# ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) -# yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll -# ego_point_3d = Point3D.from_array(ego_pose_matrix[:3, 3]) - -# # TODO: figure out if ego frame is given in rear axle or center frame -# return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=pitch, pitch=-roll, yaw=yaw) - - def _get_ego_pose_se3(frame: dataset_pb2.Frame) -> StateSE3: ego_pose_matrix = np.array(frame.pose.transform).reshape(4, 4) yaw, pitch, roll = Quaternion(matrix=ego_pose_matrix[:3, :3]).yaw_pitch_roll @@ -366,9 +399,6 @@ def _get_ego_pose_se3(frame: dataset_pb2.Frame) -> StateSE3: return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=roll, pitch=pitch, yaw=yaw) - # TODO: figure out if ego frame is given in rear axle or center frame - # return StateSE3(x=ego_point_3d.x, y=ego_point_3d.y, z=ego_point_3d.z, roll=pitch, pitch=-roll, yaw=yaw) - def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: # TODO: implement @@ -494,7 +524,7 @@ def _extract_camera( def _extract_lidar( frame: dataset_pb2.Frame, data_converter_config: DataConverterConfig -) -> Optional[npt.NDArray[np.float32]]: +) -> Dict[LiDARType, npt.NDArray[np.float32]]: from waymo_open_dataset.utils import frame_utils assert data_converter_config.lidar_store_option == "binary", "Lidar store option must be 'binary' for WOPD." @@ -507,5 +537,11 @@ def _extract_lidar( range_image_top_pose=range_image_top_pose, keep_polar_features=True, ) - points = np.array(points[0], dtype=np.float32) - return points + print(f"Extracted {len(points)} points from {len(frame.lasers)} lasers.") + + lidar_data: Dict[LiDARType, npt.NDArray[np.float32]] = {} + for lidar_idx, frame_lidar in enumerate(frame.lasers): + lidar_type = WOPD_LIDAR_TYPES[frame_lidar.name] + lidar_data[lidar_type] = np.array(points[lidar_idx], dtype=np.float32) + + return lidar_data diff --git a/d123/dataset/scene/abstract_scene.py b/d123/dataset/scene/abstract_scene.py index 84d26dd8..4877fc83 100644 --- a/d123/dataset/scene/abstract_scene.py +++ b/d123/dataset/scene/abstract_scene.py @@ -2,12 +2,12 @@ import abc from dataclasses import dataclass -from typing import List +from typing import List, Optional from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording from d123.common.datatypes.sensor.camera import Camera, CameraType -from d123.common.datatypes.sensor.lidar import LiDAR +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARType from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.dataset.logs.log_metadata import LogMetadata @@ -17,10 +17,6 @@ class AbstractScene(abc.ABC): - @property - @abc.abstractmethod - def map_api(self) -> AbstractMap: - raise NotImplementedError @property @abc.abstractmethod @@ -42,6 +38,16 @@ def log_metadata(self) -> LogMetadata: def available_camera_types(self) -> List[CameraType]: raise NotImplementedError + @property + @abc.abstractmethod + def available_lidar_types(self) -> List[LiDARType]: + raise NotImplementedError + + @property + @abc.abstractmethod + def map_api(self) -> Optional[AbstractMap]: + raise NotImplementedError + @abc.abstractmethod def get_number_of_iterations(self) -> int: raise NotImplementedError @@ -79,7 +85,7 @@ def get_camera_at_iteration(self, iteration: int, camera_type: CameraType) -> Ca raise NotImplementedError @abc.abstractmethod - def get_lidar_at_iteration(self, iteration: int) -> LiDAR: + def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> LiDAR: raise NotImplementedError def open(self) -> None: diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index 676f3831..6856d102 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -7,7 +7,7 @@ from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper from d123.common.datatypes.recording.detection_recording import DetectionRecording from d123.common.datatypes.sensor.camera import Camera, CameraMetadata, CameraType, camera_metadata_dict_from_json -from d123.common.datatypes.sensor.lidar import LiDAR +from d123.common.datatypes.sensor.lidar import LiDAR, LiDARMetadata, LiDARType, lidar_metadata_dict_from_json from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE3 from d123.common.datatypes.vehicle_state.vehicle_parameters import VehicleParameters @@ -44,8 +44,13 @@ def _get_scene_data( else: camera_metadata = {} + if b"lidar_metadata" in table.schema.metadata: + lidar_metadata = lidar_metadata_dict_from_json(table.schema.metadata[b"lidar_metadata"].decode()) + else: + lidar_metadata = {} + del table - return metadata, vehicle_parameters, camera_metadata + return metadata, vehicle_parameters, camera_metadata, lidar_metadata class ArrowScene(AbstractScene): @@ -57,10 +62,16 @@ def __init__( self._recording_table: pa.Table = None - _metadata, _vehicle_parameters, _camera_metadata = _get_scene_data(arrow_file_path) + ( + _metadata, + _vehicle_parameters, + _camera_metadata, + _lidar_metadata, + ) = _get_scene_data(arrow_file_path) self._metadata: LogMetadata = _metadata self._vehicle_parameters: VehicleParameters = _vehicle_parameters self._camera_metadata: Dict[CameraType, CameraMetadata] = _camera_metadata + self._lidar_metadata: Dict[LiDARType, LiDARMetadata] = _lidar_metadata self._map_api: Optional[AbstractMap] = None @@ -98,6 +109,10 @@ def log_metadata(self) -> LogMetadata: def available_camera_types(self) -> List[CameraType]: return list(self._camera_metadata.keys()) + @property + def available_lidar_types(self) -> List[LiDARType]: + return list(self._lidar_metadata.keys()) + def _get_table_index(self, iteration: int) -> int: self._lazy_initialize() assert ( @@ -154,9 +169,16 @@ def get_camera_at_iteration(self, iteration: int, camera_type: CameraType) -> Ca self.log_metadata, ) - def get_lidar_at_iteration(self, iteration: int) -> LiDAR: + def get_lidar_at_iteration(self, iteration: int, lidar_type: LiDARType) -> LiDAR: self._lazy_initialize() - return get_lidar_from_arrow_table(self._recording_table, self._get_table_index(iteration), self.log_metadata) + assert lidar_type in self._lidar_metadata, f"LiDAR type {lidar_type} not found in metadata." + table_index = self._get_table_index(iteration) + return get_lidar_from_arrow_table( + self._recording_table, + table_index, + self._lidar_metadata[lidar_type], + self.log_metadata, + ) def _lazy_initialize(self) -> None: self.open() diff --git a/d123/dataset/scene/scene_builder.py b/d123/dataset/scene/scene_builder.py index c9917de4..948d47d3 100644 --- a/d123/dataset/scene/scene_builder.py +++ b/d123/dataset/scene/scene_builder.py @@ -85,7 +85,11 @@ def _discover_log_paths(dataset_path: Path, split_names: Set[str], log_names: Op def _extract_scenes_from_logs(log_paths: List[Path], filter: SceneFilter) -> List[AbstractScene]: scenes: List[AbstractScene] = [] for log_path in log_paths: - scene_extraction_infos = _get_scene_extraction_info(log_path, filter) + try: + scene_extraction_infos = _get_scene_extraction_info(log_path, filter) + except Exception as e: + print(f"Error extracting scenes from {log_path}: {e}") + continue for scene_extraction_info in scene_extraction_infos: scenes.append( ArrowScene( diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index cceb2911..1eb51d46 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -16,8 +16,8 @@ defaults: - _self_ - datasets: # - nuplan_private_dataset - # - carla_dataset - - wopd_dataset + - carla_dataset + # - wopd_dataset -force_log_conversion: False -force_map_conversion: True +force_log_conversion: True +force_map_conversion: False diff --git a/notebooks/nuplan/nuplan_sensor_loading.ipynb b/notebooks/nuplan/nuplan_sensor_loading.ipynb index 0dd69b4e..3fc654d8 100644 --- a/notebooks/nuplan/nuplan_sensor_loading.ipynb +++ b/notebooks/nuplan/nuplan_sensor_loading.ipynb @@ -216,7 +216,7 @@ "metadata": {}, "outputs": [], "source": [ - "# process_images_to_arrow(\"/mnt/nvme/nuplan/dataset/nuplan-v1.1/sensor/2021.07.01.20.35.47_veh-38_00016_00281/CAM_F0\", \"test.arrow\", 10)" + "# process_images_to_arrow(\"/mnt/nvme/nuplan/dataset/nuplan-v1.1/sensor/2021.07.01.20.35.47_veh-38_00016_00281/F0\", \"test.arrow\", 10)" ] }, { @@ -310,7 +310,7 @@ " log_file = os.path.join(NUPLAN_DB_PATH, log_name + \".db\")\n", "\n", " log_cam_infos = {}\n", - " for cam in get_cameras(log_file, [str(CameraChannel.CAM_F0.value)]):\n", + " for cam in get_cameras(log_file, [str(CameraChannel.F0.value)]):\n", " intrinsics = np.array(pickle.loads(cam.intrinsic))\n", " translation = np.array(pickle.loads(cam.translation))\n", " rotation = np.array(pickle.loads(cam.rotation))\n", @@ -330,7 +330,7 @@ "images = []\n", "for lidar_pc in log_db.lidar_pc[::2]:\n", "\n", - " # front_image = get_images_from_lidar_tokens(log_path, [lidar_pc.token], [str(CameraChannel.CAM_F0.value)])\n", + " # front_image = get_images_from_lidar_tokens(log_path, [lidar_pc.token], [str(CameraChannel.F0.value)])\n", " # parameters = get_log_cam_info(log_db.log)\n", " # print(parameters)\n", "\n", diff --git a/notebooks/tools/merge_videos.ipynb b/notebooks/tools/merge_videos.ipynb index 97d442d6..257bb8ad 100644 --- a/notebooks/tools/merge_videos.ipynb +++ b/notebooks/tools/merge_videos.ipynb @@ -89,7 +89,7 @@ " return False\n", "\n", "# List your MP4 files in the order you want them merged\n", - "video_folder = Path(\"/home/daniel/d123_workspace/d123/notebooks/waymo_perception\")\n", + "video_folder = Path(\"/home/daniel/d123_logs_videos/wopd_train\")\n", "video_files = [str(file) for file in video_folder.glob(\"*.mp4\") if file.is_file()]\n", "\n", "\n", diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index b3284c9e..c24f6797 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -161,31 +161,19 @@ " for map_object in map_objects:\n", " try:\n", " if layer in [\n", - " # MapLayer.GENERIC_DRIVABLE,\n", - " # MapLayer.CARPARK,\n", - " # MapLayer.CROSSWALK,\n", - " # MapLayer.INTERSECTION,\n", - " # MapLayer.WALKWAY,\n", + " MapLayer.GENERIC_DRIVABLE,\n", + " MapLayer.CARPARK,\n", + " MapLayer.CROSSWALK,\n", + " MapLayer.INTERSECTION,\n", + " MapLayer.WALKWAY,\n", " ]:\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", "\n", - " # if layer in [MapLayer.LANE_GROUP]:\n", - " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + " if layer in [MapLayer.LANE_GROUP]:\n", + " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", "\n", " if layer in [MapLayer.LANE]:\n", - " map_object: AbstractLane\n", - " if map_object.right_lane is not None and map_object.left_lane is not None and not done:\n", - " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", - " add_shapely_polygon_to_ax(ax, map_object.right_lane.shapely_polygon, RIGHT_CONFIG)\n", - " add_shapely_polygon_to_ax(ax, map_object.left_lane.shapely_polygon, LEFT_CONFIG)\n", - " done = True\n", - " else:\n", - " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", - "\n", - "\n", - " # add_shapely_linestring_to_ax(ax, map_object.right_boundary.linestring, RIGHT_CONFIG)\n", - " # add_shapely_linestring_to_ax(ax, map_object.left_boundary.linestring, LEFT_CONFIG)\n", - " # add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, LANE_CONFIG)\n", + " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", "\n", " # centroid = map_object.shapely_polygon.centroid\n", " # ax.text(\n", @@ -197,22 +185,21 @@ " # fontsize=8,\n", " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", " # )\n", - " # if layer in [MapLayer.ROAD_EDGE]:\n", - " # add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", - " # edge_lengths.append(map_object.polyline_3d.linestring.length)\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)],\n", - " fill_color_alpha=1.0,\n", - " line_color=NEW_TAB_10[line_type % len(NEW_TAB_10)],\n", - " line_color_alpha=1.0,\n", - " line_width=1.5,\n", - " line_style=\"-\",\n", - " zorder=3,\n", - " )\n", - " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, plt_config)\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, plt_config)\n", "\n", " except Exception:\n", " import traceback\n", @@ -256,7 +243,7 @@ "\n", "\n", "scene_index = 1\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=100, radius=100)\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=10, radius=100)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] @@ -268,10 +255,26 @@ "metadata": {}, "outputs": [], "source": [ - "map_api = scenes[1].map_api\n", + "# from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", + "# from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", + "\n", + "\n", + "# map_api: GPKGMap = scenes[scene_index].map_api\n", + "\n", + "# drivable_polygons = map_api._gpd_dataframes[MapLayer.LANE]\n", + "\n", + "\n", "\n", + "# linear_rings = get_road_edge_linear_rings(drivable_polygons.geometry.tolist())\n", "\n", - "map_api._gpd_dataframes[MapLayer.LANE].left_lane_id.tolist()" + "# size = 16\n", + "# fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "# for ring in linear_rings:\n", + "# ax.plot(*ring.xy)\n", + "\n", + "\n", + "# ax.set_aspect(\"equal\", adjustable=\"box\")" ] }, { @@ -304,7 +307,12 @@ "id": "7", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "from d123.common.datatypes.sensor.lidar_index import NuplanLidarIndex\n", + "\n", + "\n", + "NuplanLidarIndex.Y" + ] } ], "metadata": { diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index d7b611b2..049fff6d 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -37,8 +37,8 @@ "\n", "\n", "# splits = [\"carla\"]\n", - "splits = [\"nuplan_private_test\"]\n", - "# splits = [\"wopd_train\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"wopd_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -125,47 +125,16 @@ "metadata": {}, "outputs": [], "source": [ - "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", - "\n", - " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", - " CameraType.CAM_L0: ax[0, 0],\n", - " CameraType.CAM_F0: ax[0, 1],\n", - " CameraType.CAM_R0: ax[0, 2],\n", - " CameraType.CAM_L2: ax[1, 0],\n", - " CameraType.CAM_B0: ax[1, 1],\n", - " CameraType.CAM_R2: ax[1, 2],\n", - " }\n", - "\n", - " box_detections = scene.get_box_detections_at_iteration(iteration)\n", - " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", - " for camera_type, camera_ax in camera_ax_assignment.items():\n", - " assert camera_type in scene.available_camera_types\n", - " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", - " if camera is not None:\n", - " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", - "\n", - " return ax\n", - "\n", - "\n", - "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", - " \"\"\"\n", - " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", - " :param scene: navsim scene dataclass\n", - " :param frame_idx: index of selected frame\n", - " :return: figure and ax object of matplotlib\n", - " \"\"\"\n", - " scale = 2\n", - " fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", - " _add_camera_with_detections(ax, scene, iteration)\n", - "\n", - " fig.tight_layout()\n", - " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", - "\n", - "\n", - "\n", "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", "\n", - "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + "# CameraType.CAM_L0: ax[0, 0],\n", + "# CameraType.CAM_F0: ax[0, 1],\n", + "# CameraType.CAM_R0: ax[0, 2],\n", + "# CameraType.CAM_L2: ax[1, 0],\n", + "# CameraType.CAM_B0: ax[1, 1],\n", + "# CameraType.CAM_R2: ax[1, 2],\n", + "# }\n", "\n", "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", @@ -186,14 +155,45 @@ "# :return: figure and ax object of matplotlib\n", "# \"\"\"\n", "# scale = 2\n", - "# fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", + "# fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", "# _add_camera_with_detections(ax, scene, iteration)\n", "\n", "# fig.tight_layout()\n", "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", "\n", - "# plot_cameras_with_detections(scenes[9], iteration=20)\n", + "\n", + "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", + "\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + " return ax\n", + "\n", + "\n", + "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 2\n", + " fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", + " _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "plot_cameras_with_detections(scenes[9], iteration=20)\n", "\n", "\n", "\n", @@ -202,7 +202,7 @@ "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", "# CameraType.CAM_L1: ax[0],\n", "# CameraType.CAM_L0: ax[1],\n", - "# CameraType.CAM_F0: ax[2],\n", + "# CameraType.F0: ax[2],\n", "# CameraType.CAM_R0: ax[3],\n", "# CameraType.CAM_R1: ax[4],\n", "# }\n", @@ -232,7 +232,7 @@ "# fig.tight_layout()\n", "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "plot_cameras_with_detections(scenes[5], iteration=0)\n" + "plot_cameras_with_detections(scenes[0], iteration=0)\n" ] }, { diff --git a/notebooks/viz/log_rendering.ipynb b/notebooks/viz/log_rendering.ipynb index 5893308e..e88c1012 100644 --- a/notebooks/viz/log_rendering.ipynb +++ b/notebooks/viz/log_rendering.ipynb @@ -51,15 +51,15 @@ "output_path = Path(f\"/home/daniel/d123_logs_videos/{split}\")\n", "log_path = Path(f\"/home/daniel/d123_workspace/data/{split}\")\n", "for log_file in log_path.iterdir():\n", - " scene = ArrowScene(log_file)\n", " try:\n", + " scene = ArrowScene(log_file)\n", " render_scene_animation(scene, output_path, fps=20, end_idx=None, step=1)\n", + " del scene\n", " except Exception as e:\n", " traceback.print_exc()\n", " erroneous_file = output_path / f\"{log_name}.mp4\"\n", " if erroneous_file.exists():\n", " erroneous_file.unlink()\n", - " del scene\n", " # break\n", " # break" ] diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index e7584e06..b6cb640b 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -21,18 +21,11 @@ "metadata": {}, "outputs": [], "source": [ - "# split = \"nuplan_private_test\"\n", - "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", - " \n", - "\n", - "# splits = [\"wopd_train\"]\n", - "# log_names = None\n", - "\n", "\n", "\n", "# splits = [\"nuplan_private_test\"]\n", - "splits = [\"carla\"]\n", - "# log_names = [\"11839652018869852123_2565_000_2585_000\"]\n", + "# splits = [\"carla\"]\n", + "splits = [\"wopd_train\"]\n", "log_names = None\n", "\n", "scene_tokens = None\n", @@ -41,7 +34,7 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=19,\n", + " duration_s=15,\n", " history_s=0.0,\n", " timestamp_threshold_s=20,\n", " shuffle=True,\n", @@ -61,23 +54,31 @@ "id": "2", "metadata": {}, "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], "source": [ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=1)" ] }, { "cell_type": "markdown", - "id": "3", + "id": "4", "metadata": {}, "source": [] }, { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [] @@ -85,7 +86,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [], "source": [] diff --git a/notebooks/waymo_perception/testing.ipynb b/notebooks/waymo_perception/testing.ipynb index 6ee57cf5..19745f8f 100644 --- a/notebooks/waymo_perception/testing.ipynb +++ b/notebooks/waymo_perception/testing.ipynb @@ -21,7 +21,8 @@ "from waymo_open_dataset.utils import box_utils\n", "\n", "\n", - "import matplotlib.pyplot as plt" + "import matplotlib.pyplot as plt\n", + "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"-1\"" ] }, { @@ -86,53 +87,18 @@ " image = Image.open(io.BytesIO(data))\n", " return np.array(image)\n", "\n", - "\n", - "ego_state_se3s = []\n", - "front_images = []\n", "dataset = tf.data.TFRecordDataset(pathname, compression_type=\"\")\n", "\n", - "boxes = []\n", "\n", "for frame_idx, data in enumerate(dataset):\n", "\n", " frame = dataset_pb2.Frame()\n", " frame.ParseFromString(data.numpy())\n", - " # print(frame.camera_labels)\n", - " print(\"Frame attributes:\")\n", - " for field in frame.DESCRIPTOR.fields:\n", - " field_name = field.name\n", - " if hasattr(frame, field_name):\n", - " value = getattr(frame, field_name)\n", - " if field_name != \"images\": # Don't print the whole image data\n", - " print(f\" {field_name}: {type(value)}\")\n", - " if hasattr(value, \"__len__\") and not isinstance(value, (str, bytes)):\n", - " print(f\" Length: {len(value)}\")\n", - " else:\n", - " print(f\" {field_name}: List with {len(value)} images\")\n", - "\n", - "\n", - " # # 1. pose\n", - " pose = np.array(frame.pose.transform).reshape(4, 4)\n", - " yaw_pitch_roll = Quaternion(matrix=pose[:3, :3]).yaw_pitch_roll\n", - " ego_state_se3s.append(\n", - " np.array(\n", - " [\n", - " pose[0, 3], # x\n", - " pose[1, 3], # y\n", - " pose[2, 3], # z\n", - " yaw_pitch_roll[2], # yaw\n", - " yaw_pitch_roll[1], # pitch\n", - " yaw_pitch_roll[0], # roll\n", - " ],\n", - " dtype=np.float64,\n", - " )\n", - " )\n", + " print(frame.context)\n", "\n", " # # plt.show()\n", " if frame_idx == 0:\n", - " break\n", - "\n", - "ego_state_se3s = np.array(ego_state_se3s, dtype=np.float64)" + " break\n" ] }, { @@ -141,6 +107,17 @@ "id": "3", "metadata": {}, "outputs": [], + "source": [ + "for laser_calibration in frame.context.laser_calibrations:\n", + " print(np.array(laser_calibration.extrinsic.transform).reshape(4, 4))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], "source": [ "for frame_idx, data in enumerate(dataset):\n", " frame = dataset_pb2.Frame()\n", @@ -165,7 +142,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [ @@ -244,7 +221,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [], "source": [ @@ -272,7 +249,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", "metadata": {}, "outputs": [], "source": [] From 51dcdd2e827e743ea4db9321319eef3006e2428c Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Wed, 13 Aug 2025 16:19:05 +0200 Subject: [PATCH 30/42] Include scripts to run viser with hydra. --- d123/common/visualization/viser/server.py | 21 ++++-- d123/common/visualization/viser/utils.py | 47 ++++++++------ d123/dataset/scene/scene_builder.py | 2 +- .../common/scene_filter/log_scenes.yaml | 20 ++++++ d123/script/config/viser/__init__.py | 0 d123/script/config/viser/default_viser.yaml | 16 +++++ d123/script/run_viser.py | 29 +++++++++ notebooks/viz/bev_matplotlib.ipynb | 64 +++++++++++-------- scripts/viz/run_viser.sh | 6 ++ 9 files changed, 151 insertions(+), 54 deletions(-) create mode 100644 d123/script/config/common/scene_filter/log_scenes.yaml create mode 100644 d123/script/config/viser/__init__.py create mode 100644 d123/script/config/viser/default_viser.yaml create mode 100644 d123/script/run_viser.py create mode 100644 scripts/viz/run_viser.sh diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index f2242b64..7db2351e 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -33,7 +33,7 @@ LINE_WIDTH: float = 4.0 # Bounding box config: -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" # Map config: MAP_AVAILABLE: bool = True @@ -44,11 +44,11 @@ VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] -VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] -CAMERA_SCALE: float = 1.0 +VISUALIZE_CAMERA_GUI: List[CameraType] = [] +CAMERA_SCALE: float = 2.0 # Lidar config: -LIDAR_AVAILABLE: bool = True +LIDAR_AVAILABLE: bool = False LIDAR_TYPES: List[LiDARType] = [ LiDARType.LIDAR_MERGED, @@ -94,9 +94,18 @@ def next(self) -> None: def set_scene(self, scene: AbstractScene) -> None: num_frames = scene.get_number_of_iterations() - print(scene.available_camera_types) + # print(scene.available_camera_types) self.server.gui.configure_theme(dark_mode=False, control_width="large") + + # TODO: Fix lighting. Environment map can help, but cannot be freely configured. + # self.server.scene.configure_environment_map( + # hdri="warehouse", + # background=False, + # background_intensity=0.25, + # environment_intensity=0.5, + # ) + with self.server.gui.add_folder("Playback"): server_playing = True @@ -262,7 +271,7 @@ def _(_) -> None: # line_width=LINE_WIDTH, # ) # self.server.scene.add_line_segments( - # "/map/right_boundary", + # "/map/right_boundary",clear # right_boundaries, # colors=[[TAB_10[3].rgb]], # line_width=LINE_WIDTH, diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index ab87a1e3..d527b05e 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -13,6 +13,7 @@ from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3 from d123.common.geometry.line.polylines import Polyline3D from d123.common.geometry.transform.se3 import convert_relative_to_absolute_points_3d_array +from d123.common.visualization.color.color import TAB_10, Color from d123.common.visualization.color.config import PlotConfig from d123.common.visualization.color.default import BOX_DETECTION_CONFIG, EGO_VEHICLE_CONFIG, MAP_SURFACE_CONFIG from d123.dataset.maps.abstract_map import MapLayer @@ -23,7 +24,25 @@ # TODO: Add general utilities for 3D primitives and mesh support. MAP_RADIUS: Final[float] = 500 -BRIGHTNESS_FACTOR: Final[float] = 0.8 +BRIGHTNESS_FACTOR: Final[float] = 1.0 + + +def configure_trimesh(mesh: trimesh.Trimesh, color: Color): + # base_color = [r / 255.0 for r in color.rgba] + mesh.visual.face_colors = color.rgba + + # pbr_material = trimesh.visual.material.PBRMaterial( + # baseColorFactor=base_color, # Your desired color (RGBA, 0-1 range) + # metallicFactor=0.0, # 0.0 = non-metallic (more matte) + # roughnessFactor=1.0, # 0.8 = quite rough (less shiny, 0=mirror, 1=completely rough) + # emissiveFactor=[0.0, 0.0, 0.0], # No emission + # alphaCutoff=0.9, # Alpha threshold for transparency + # doubleSided=True, # Single-sided material + # ) + # mesh.visual.material = pbr_material + # mesh.visual = mesh.visual.to_texture() + + return mesh def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> trimesh.Trimesh: @@ -38,20 +57,8 @@ def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> tr # Apply translation box_mesh = box_mesh.apply_translation([bbox.center.x, bbox.center.y, bbox.center.z]) - base_color = [r / 255.0 for r in plot_config.fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba] - box_mesh.visual.face_colors = plot_config.fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba - - pbr_material = trimesh.visual.material.PBRMaterial( - baseColorFactor=base_color, # Your desired color (RGBA, 0-1 range) - metallicFactor=1.0, # 0.0 = non-metallic (more matte) - roughnessFactor=0.9, # 0.8 = quite rough (less shiny, 0=mirror, 1=completely rough) - emissiveFactor=[0.0, 0.0, 0.0], # No emission - alphaCutoff=0.75, # Alpha threshold for transparency - doubleSided=False, # Single-sided material - ) - box_mesh.visual.material = pbr_material - return box_mesh + return configure_trimesh(box_mesh, plot_config.fill_color) def translate_bounding_box_se3(bounding_box_se3: BoundingBoxSE3, point_3d: Point3D) -> BoundingBoxSE3: @@ -109,10 +116,10 @@ def get_map_meshes(scene: AbstractScene): MapLayer.WALKWAY, MapLayer.CROSSWALK, MapLayer.GENERIC_DRIVABLE, - # MapLayer.CARPARK, + MapLayer.CARPARK, ]: # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. - trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.05).array + trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.1).array else: trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z).array @@ -121,7 +128,7 @@ def get_map_meshes(scene: AbstractScene): x=0, y=0, z=center.z - initial_ego_vehicle_state.vehicle_parameters.height / 2 ).array - trimesh_mesh.visual.face_colors = MAP_SURFACE_CONFIG[map_layer].fill_color.set_brightness(0.8).rgba + trimesh_mesh = configure_trimesh(trimesh_mesh, MAP_SURFACE_CONFIG[map_layer].fill_color) surface_meshes.append(trimesh_mesh) output[f"{map_layer.serialize()}"] = trimesh.util.concatenate(surface_meshes) @@ -207,7 +214,7 @@ def _create_lane_mesh_from_boundary_arrays( faces = np.array(faces) mesh = trimesh.Trimesh(vertices=vertices, faces=faces) - mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapLayer.LANE].fill_color.set_brightness(BRIGHTNESS_FACTOR).rgba + mesh.visual.face_colors = MAP_SURFACE_CONFIG[MapLayer.LANE].fill_color.rgba return mesh @@ -302,8 +309,8 @@ def float_to_rgb(values: npt.NDArray[np.float32], cmap_name: str = "viridis") -> points = convert_relative_to_absolute_points_3d_array(origin, points) points = points - initial_ego_vehicle_state.center_se3.point_3d.array points_.append(points) - # colors_.append([TAB_10[lidar_idx % len(TAB_10)].rgb] * points.shape[0]) - colors_.append(float_to_rgb(lidar.intensity, cmap_name="viridis")) + colors_.append([TAB_10[lidar_idx % len(TAB_10)].rgb] * points.shape[0]) + # colors_.append(float_to_rgb(lidar.intensity, cmap_name="viridis")) points_ = np.concatenate(points_, axis=0) if points_ else np.empty((0, 3), dtype=np.float32) colors_ = np.concatenate(colors_, axis=0) if colors_ else np.empty((0, 3), dtype=np.float32) diff --git a/d123/dataset/scene/scene_builder.py b/d123/dataset/scene/scene_builder.py index 948d47d3..3c97e0a4 100644 --- a/d123/dataset/scene/scene_builder.py +++ b/d123/dataset/scene/scene_builder.py @@ -110,7 +110,7 @@ def _get_scene_extraction_info(log_path: Union[str, Path], filter: SceneFilter) if filter.map_names is not None and log_metadata.map_name not in filter.map_names: return scene_extraction_infos - start_idx = int(filter.history_s / log_metadata.timestep_seconds) # if filter.history_s is not None else 0 + start_idx = int(filter.history_s / log_metadata.timestep_seconds) if filter.history_s is not None else 0 end_idx = ( len(recording_table) - int(filter.duration_s / log_metadata.timestep_seconds) if filter.duration_s is not None diff --git a/d123/script/config/common/scene_filter/log_scenes.yaml b/d123/script/config/common/scene_filter/log_scenes.yaml new file mode 100644 index 00000000..6c3bcee7 --- /dev/null +++ b/d123/script/config/common/scene_filter/log_scenes.yaml @@ -0,0 +1,20 @@ +_target_: d123.dataset.scene.scene_filter.SceneFilter +_convert_: 'all' + +split_types: null +split_names: null +log_names: null + + +map_names: null +scene_tokens: null +timestamp_threshold_s: null +ego_displacement_minimum_m: null + +duration_s: null +history_s: null + +camera_types: null + +max_num_scenes: null +shuffle: false diff --git a/d123/script/config/viser/__init__.py b/d123/script/config/viser/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/script/config/viser/default_viser.yaml b/d123/script/config/viser/default_viser.yaml new file mode 100644 index 00000000..c42164dc --- /dev/null +++ b/d123/script/config/viser/default_viser.yaml @@ -0,0 +1,16 @@ +hydra: + run: + dir: . + output_subdir: null + searchpath: # Only in these paths are discoverable + - pkg://d123.script.config + - pkg://d123.script.config.common + - pkg://d123.script.config.preprocessing + job: + chdir: False +# +defaults: + - default_common + - default_dataset_paths + +port_number: 8080 diff --git a/d123/script/run_viser.py b/d123/script/run_viser.py new file mode 100644 index 00000000..e682a96e --- /dev/null +++ b/d123/script/run_viser.py @@ -0,0 +1,29 @@ +import logging + +import hydra +from omegaconf import DictConfig + +from d123.common.visualization.viser.server import ViserVisualizationServer +from d123.script.builders.scene_builder_builder import build_scene_builder +from d123.script.builders.scene_filter_builder import build_scene_filter +from d123.script.run_dataset_conversion import build_worker + +logger = logging.getLogger(__name__) + +CONFIG_PATH = "config/viser" +CONFIG_NAME = "default_viser" + + +@hydra.main(config_path=CONFIG_PATH, config_name=CONFIG_NAME, version_base=None) +def main(cfg: DictConfig) -> None: + + worker = build_worker(cfg) + scene_filter = build_scene_filter(cfg.scene_filter) + scene_builder = build_scene_builder(cfg.scene_builder) + scenes = scene_builder.get_scenes(scene_filter, worker=worker) + + ViserVisualizationServer(scenes=scenes) + + +if __name__ == "__main__": + main() diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index c24f6797..79b1e8e9 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -123,7 +123,7 @@ ")\n", "\n", "ROAD_LINE_CONFIG: PlotConfig = PlotConfig(\n", - " fill_color=DARKER_GREY,\n", + " fill_color=NEW_TAB_10[5],\n", " fill_color_alpha=1.0,\n", " line_color=NEW_TAB_10[5],\n", " line_color_alpha=1.0,\n", @@ -188,18 +188,18 @@ " 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, plt_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", @@ -220,8 +220,8 @@ " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", - " add_box_detections_to_ax(ax, box_detections)\n", - " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + " # 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", @@ -242,8 +242,8 @@ " return fig, ax\n", "\n", "\n", - "scene_index = 1\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=10, radius=100)\n", + "scene_index = 5\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=150, radius=300)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] @@ -255,26 +255,36 @@ "metadata": {}, "outputs": [], "source": [ - "# from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", - "# from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", + "from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", + "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", "\n", "\n", - "# map_api: GPKGMap = scenes[scene_index].map_api\n", + "map_api: GPKGMap = scenes[scene_index].map_api\n", "\n", - "# drivable_polygons = map_api._gpd_dataframes[MapLayer.LANE]\n", + "drivable_polygons = map_api._gpd_dataframes[MapLayer.LANE]\n", "\n", "\n", "\n", - "# linear_rings = get_road_edge_linear_rings(drivable_polygons.geometry.tolist())\n", + "linear_rings = get_road_edge_linear_rings(drivable_polygons.geometry.tolist())\n", + "rings_lengths = [ring.length for ring in linear_rings]\n", + "max_length_idx = np.argmax(rings_lengths)\n", "\n", - "# size = 16\n", - "# fig, ax = plt.subplots(figsize=(size, size))\n", "\n", - "# for ring in linear_rings:\n", - "# ax.plot(*ring.xy)\n", "\n", "\n", - "# ax.set_aspect(\"equal\", adjustable=\"box\")" + "\n", + "\n", + "size = 16\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "for idx, ring in enumerate(linear_rings):\n", + " if idx == max_length_idx:\n", + " ax.plot(*ring.xy, color=\"black\", linewidth=2, label=\"Longest Road Edge\")\n", + " else:\n", + " ax.plot(*ring.xy)\n", + "\n", + "\n", + "ax.set_aspect(\"equal\", adjustable=\"box\")" ] }, { diff --git a/scripts/viz/run_viser.sh b/scripts/viz/run_viser.sh new file mode 100644 index 00000000..0dafc0a3 --- /dev/null +++ b/scripts/viz/run_viser.sh @@ -0,0 +1,6 @@ + + +python $D123_DEVKIT_ROOT/d123/script/run_viser.py \ +scene_filter=log_scenes \ +scene_filter.shuffle=True \ +worker=sequential From 65c4754d2859981591df153a900b70a143b2954c Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Wed, 13 Aug 2025 18:03:23 +0200 Subject: [PATCH 31/42] Add initial spphinx documentation. --- .gitignore | 7 ++++ docs/Makefile | 20 ++++++++++ docs/__init__.py | 0 docs/conf.py | 69 +++++++++++++++++++++++++++++++++ docs/{camera.md => datasets.md} | 14 ++++--- docs/index.rst | 20 ++++++++++ docs/installation.md | 8 ++-- docs/make.bat | 35 +++++++++++++++++ docs/requirements.in | 2 + docs/requirements.txt | 60 ++++++++++++++++++++++++++++ docs/schema.md | 7 ++-- 11 files changed, 229 insertions(+), 13 deletions(-) create mode 100644 docs/Makefile create mode 100644 docs/__init__.py create mode 100644 docs/conf.py rename docs/{camera.md => datasets.md} (97%) create mode 100644 docs/index.rst create mode 100644 docs/make.bat create mode 100644 docs/requirements.in create mode 100644 docs/requirements.txt diff --git a/.gitignore b/.gitignore index 22cfdee9..bdc0a21f 100644 --- a/.gitignore +++ b/.gitignore @@ -23,3 +23,10 @@ *.csv *.log *.mp4 + + +# Sphinx documentation +docs/_build/ +docs/build/ +_build/ +.doctrees/ diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 00000000..d4bb2cbb --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/__init__.py b/docs/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 00000000..c9379f0d --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,69 @@ +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +project = "d123" +copyright = "2025, DanielDauner" +author = "DanielDauner" +release = "v0.0.6" + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + "sphinx.ext.duration", + "sphinx.ext.doctest", + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.intersphinx", + "myst_parser", +] + +intersphinx_mapping = { + "rtd": ("https://docs.readthedocs.io/en/stable/", None), + "python": ("https://docs.python.org/3/", None), + "sphinx": ("https://www.sphinx-doc.org/en/master/", None), +} +intersphinx_disabled_domains = ["std"] + +templates_path = ["_templates"] + +# -- Options for EPUB output +epub_show_urls = "footnote" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path. +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "sphinx_rtd_theme" + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ["_static"] + +html_theme_options = { + "analytics_anonymize_ip": False, + "logo_only": False, + "display_version": True, + "prev_next_buttons_location": "bottom", + "style_external_links": False, + "vcs_pageview_mode": "", + "style_nav_header_background": "#FFE100", + # Toc options + "collapse_navigation": True, + "sticky_navigation": True, + "navigation_depth": 3, + "includehidden": True, + "titles_only": False, +} diff --git a/docs/camera.md b/docs/datasets.md similarity index 97% rename from docs/camera.md rename to docs/datasets.md index 9a016aa7..31ebda18 100644 --- a/docs/camera.md +++ b/docs/datasets.md @@ -1,6 +1,8 @@ -# nuPlan +# Datasets + +## nuPlan ``` 0: { "camera_name": "CAM_F0", @@ -45,7 +47,7 @@ ``` -# nuScenes +## nuScenes ``` 0: { "camera_name": "CAM_FRONT", @@ -80,7 +82,7 @@ ``` -# Waymo +## Waymo ``` 0: { "camera_name": "front_camera", @@ -109,7 +111,7 @@ }, ``` -# Panda set +## Panda set ``` 0: { @@ -145,7 +147,7 @@ ``` -# Argoverse +## Argoverse ``` 0: { @@ -186,7 +188,7 @@ ``` -# KITTI +## KITTI ``` 0: { "camera_name": "CAM_LEFT", diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 00000000..53aedd4d --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,20 @@ +.. 123d documentation master file, created by + sphinx-quickstart on Wed Aug 13 16:57:48 2025. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +123d documentation +================== + +Add your content using ``reStructuredText`` syntax. See the +`reStructuredText `_ +documentation for details. + + +.. toctree:: + :hidden: + :caption: Contents: + + installation + datasets + schema diff --git a/docs/installation.md b/docs/installation.md index 6fe76e59..95d09a04 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -1,5 +1,5 @@ -# Install Code +# Installation Note, the following installation assumes the following folder structure: ``` @@ -19,13 +19,13 @@ Note, the following installation assumes the following folder structure: │ ├── ... │ └── 2021.10.06.07.26.10_veh-52_00006_00398.arrow ├── nuplan_mini_train - │ └── ... + │ └── ... └── nuplan_mini_test └── ... ``` -First you need to create a new conda environment and install `d123` as editable pip package. +First you need to create a new conda environment and install `d123` as editable pip package. ```bash conda env create -f environment.yml conda activate d123 @@ -45,4 +45,4 @@ export CARLA_SIMULATOR_ROOT="$HOME/carla_workspace/carla_garage/carla" # nuPlan export NUPLAN_DATA_ROOT="/path/to/nuplan/dataset" export NUPLAN_MAPS_ROOT="/path/to/nuplan/dataset/maps" -``` \ No newline at end of file +``` diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 00000000..32bb2452 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=. +set BUILDDIR=_build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/requirements.in b/docs/requirements.in new file mode 100644 index 00000000..acbc25d5 --- /dev/null +++ b/docs/requirements.in @@ -0,0 +1,2 @@ +Sphinx>=5,<6 +sphinx_rtd_theme diff --git a/docs/requirements.txt b/docs/requirements.txt new file mode 100644 index 00000000..32230734 --- /dev/null +++ b/docs/requirements.txt @@ -0,0 +1,60 @@ +# +# This file is autogenerated by pip-compile with python 3.10 +# To update, run: +# +# pip-compile docs/requirements.in +# +alabaster==0.7.12 + # via sphinx +babel==2.10.3 + # via sphinx +certifi==2022.6.15 + # via requests +charset-normalizer==2.1.0 + # via requests +docutils==0.17.1 + # via + # sphinx + # sphinx-rtd-theme +idna==3.3 + # via requests +imagesize==1.4.1 + # via sphinx +jinja2==3.1.2 +m2r2 + # via sphinx +markupsafe==2.1.1 + # via requests +myst-parser + # via jinja2 +packaging==21.3 + # via sphinx +pygments==2.12.0 + # via sphinx +pyparsing==3.0.9 + # via packaging +pytz==2022.1 + # via babel +requests==2.28.1 + # via sphinx +snowballstemmer==2.2.0 + # via sphinx +sphinx==5.0.2 + # via + # -r docs/requirements.in + # sphinx-rtd-theme +sphinx-rtd-theme==1.0.0 + # via -r docs/requirements.in +sphinxcontrib-applehelp==1.0.2 + # via sphinx +sphinxcontrib-devhelp==1.0.2 + # via sphinx +sphinxcontrib-htmlhelp==2.0.0 + # via sphinx +sphinxcontrib-jsmath==1.0.1 + # via sphinx +sphinxcontrib-qthelp==1.0.3 + # via sphinx +sphinxcontrib-serializinghtml==1.1.5 + # via sphinx +urllib3==1.26.9 diff --git a/docs/schema.md b/docs/schema.md index 2669585a..b3d44a6e 100644 --- a/docs/schema.md +++ b/docs/schema.md @@ -1,7 +1,8 @@ +# Schema - -token +TODO + From 95dc0f5de8d2a56e6321e9e4dc36ba1d1176fbfe Mon Sep 17 00:00:00 2001 From: DanielDauner Date: Fri, 15 Aug 2025 11:10:02 +0200 Subject: [PATCH 32/42] Add multithreading tools from nuPlan. Change pip install from `setup.py` to `toml` with option for module dependencies. Add download scripts for raw data. --- d123/common/multithreading/ray_execution.py | 122 +++++++++++++ d123/common/multithreading/worker_parallel.py | 60 +++++++ d123/common/multithreading/worker_pool.py | 156 ++++++++++++++++ d123/common/multithreading/worker_ray.py | 166 ++++++++++++++++++ .../multithreading/worker_sequential.py | 46 +++++ d123/common/multithreading/worker_utils.py | 36 ++++ .../carla/carla_data_converter.py | 2 +- .../nuplan/nuplan_data_converter.py | 2 +- .../dataset_specific/raw_data_converter.py | 2 +- .../wopd/wopd_data_converter.py | 2 +- d123/dataset/maps/abstract_map.py | 1 + d123/dataset/scene/scene_builder.py | 3 +- .../script/builders/data_converter_builder.py | 2 +- d123/script/builders/utils/utils_type.py | 73 ++++++++ d123/script/builders/worker_pool_builder.py | 14 +- .../config/common/worker/ray_distributed.yaml | 2 +- .../config/common/worker/sequential.yaml | 2 +- .../worker/single_machine_thread_pool.yaml | 2 +- .../default_dataset_conversion.yaml | 4 +- d123/script/run_preprocessing.py | 2 +- d123/script/run_simulation.py | 2 +- .../gym/environment/environment_wrapper.py | 6 +- .../abstract_output_converter.py | 1 - .../sim_agent/smart/layers/attention_layer.py | 2 +- docs/conf.py | 5 +- docs/index.rst | 2 +- docs/requirements.in | 2 - docs/requirements.txt | 60 ------- environment.yml | 2 - .../deprecated/collect_sim_metrics_gym.ipynb | 2 +- notebooks/deprecated/extraction_testing.ipynb | 2 +- notebooks/deprecated/test_scene_builder.ipynb | 2 +- notebooks/deprecated/test_waypoints.ipynb | 2 +- notebooks/gym/test_gym.ipynb | 4 +- notebooks/gym/test_simulation_2d.ipynb | 4 +- notebooks/scene_rendering.ipynb | 4 +- notebooks/scene_sensor_loading.ipynb | 4 +- notebooks/smarty/smart_feature_testing.ipynb | 2 +- notebooks/smarty/smart_rollout.ipynb | 2 +- notebooks/viz/bev_matplotlib.ipynb | 40 ++--- notebooks/viz/camera_matplotlib.ipynb | 2 +- notebooks/viz/viser_testing_v2_scene.ipynb | 2 +- pyproject.toml | 91 ++++++++++ requirements.txt | 109 ++++++------ scripts/download/download_av2.sh | 15 ++ ...load_nuplan.sh => download_nuplan_logs.sh} | 0 scripts/download/download_nuplan_sensor.sh | 14 ++ setup.py | 30 ---- 48 files changed, 894 insertions(+), 218 deletions(-) create mode 100644 d123/common/multithreading/ray_execution.py create mode 100644 d123/common/multithreading/worker_parallel.py create mode 100644 d123/common/multithreading/worker_pool.py create mode 100644 d123/common/multithreading/worker_ray.py create mode 100644 d123/common/multithreading/worker_sequential.py create mode 100644 d123/common/multithreading/worker_utils.py create mode 100644 d123/script/builders/utils/utils_type.py delete mode 100644 docs/requirements.in delete mode 100644 docs/requirements.txt create mode 100644 pyproject.toml create mode 100644 scripts/download/download_av2.sh rename scripts/download/{download_nuplan.sh => download_nuplan_logs.sh} (100%) create mode 100644 scripts/download/download_nuplan_sensor.sh delete mode 100644 setup.py diff --git a/d123/common/multithreading/ray_execution.py b/d123/common/multithreading/ray_execution.py new file mode 100644 index 00000000..4e44c56f --- /dev/null +++ b/d123/common/multithreading/ray_execution.py @@ -0,0 +1,122 @@ +import logging +import traceback +from functools import partial +from pathlib import Path +from typing import Any, Callable, Iterable, Iterator, List, Optional, Tuple, cast +from uuid import uuid1 + +import ray +from ray.exceptions import RayTaskError +from ray.remote_function import RemoteFunction +from tqdm import tqdm + +from d123.common.multithreading.worker_pool import Task + + +def _ray_object_iterator(initial_ids: List[ray.ObjectRef]) -> Iterator[Tuple[ray.ObjectRef, Any]]: + """ + Iterator that waits for each ray object in the input object list to be completed and fetches the result. + :param initial_ids: list of ray object ids + :yield: result of worker + """ + next_ids = initial_ids + + while next_ids: + ready_ids, not_ready_ids = ray.wait(next_ids) + next_id = ready_ids[0] + + yield next_id, ray.get(next_id) + + next_ids = not_ready_ids + + +def wrap_function(fn: Callable[..., Any], log_dir: Optional[Path] = None) -> Callable[..., Any]: + """ + Wraps a function to save its logs to a unique file inside the log directory. + :param fn: function to be wrapped. + :param log_dir: directory to store logs (wrapper function does nothing if it's not set). + :return: wrapped function which changes logging settings while it runs. + """ + + def wrapped_fn(*args: Any, **kwargs: Any) -> Any: + if log_dir is None: + return fn(*args, **kwargs) + + log_dir.mkdir(parents=True, exist_ok=True) + log_path = log_dir / f"{uuid1().hex}__{fn.__name__}.log" + + logging.basicConfig() + logger = logging.getLogger() + fh = logging.FileHandler(log_path, delay=True) + fh.setLevel(logging.INFO) + logger.addHandler(fh) + logger.setLevel(logging.INFO) + + # Silent botocore which is polluting the terminal because of serialization and deserialization + # with following message: INFO:botocore.credentials:Credentials found in config file: ~/.aws/config + logging.getLogger("botocore").setLevel(logging.WARNING) + + result = fn(*args, **kwargs) + + fh.flush() + fh.close() + logger.removeHandler(fh) + + return result + + return wrapped_fn + + +def _ray_map_items(task: Task, *item_lists: Iterable[List[Any]], log_dir: Optional[Path] = None) -> List[Any]: + """ + Map each item of a list of arguments to a callable and executes in parallel. + :param fn: callable to be run + :param item_list: items to be parallelized + :param log_dir: directory to store worker logs + :return: list of outputs + """ + assert len(item_lists) > 0, "No map arguments received for mapping" + assert all(isinstance(items, list) for items in item_lists), "All map arguments must be lists" + assert all( + len(cast(List, items)) == len(item_lists[0]) for items in item_lists # type: ignore + ), "All lists must have equal size" + fn = task.fn + # Wrap function in remote decorator and create ray objects + if isinstance(fn, partial): + _, _, pack = fn.__reduce__() # type: ignore + fn, _, args, _ = pack + fn = wrap_function(fn, log_dir=log_dir) + remote_fn: RemoteFunction = ray.remote(fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus) + object_ids = [remote_fn.remote(*items, **args) for items in zip(*item_lists)] + else: + fn = wrap_function(fn, log_dir=log_dir) + remote_fn = ray.remote(fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus) + object_ids = [remote_fn.remote(*items) for items in zip(*item_lists)] + + # Create ordered map to track order of objects inserted in the queue + object_result_map = dict.fromkeys(object_ids, None) + + # Asynchronously iterate through the object and track progress + for object_id, output in tqdm(_ray_object_iterator(object_ids), total=len(object_ids), desc="Ray objects"): + object_result_map[object_id] = output + + results = list(object_result_map.values()) + + return results + + +def ray_map(task: Task, *item_lists: Iterable[List[Any]], log_dir: Optional[Path] = None) -> List[Any]: + """ + Initialize ray, align item lists and map each item of a list of arguments to a callable and executes in parallel. + :param task: callable to be run + :param item_lists: items to be parallelized + :param log_dir: directory to store worker logs + :return: list of outputs + """ + try: + results = _ray_map_items(task, *item_lists, log_dir=log_dir) + return results + except (RayTaskError, Exception) as exc: + ray.shutdown() + traceback.print_exc() + raise RuntimeError(exc) diff --git a/d123/common/multithreading/worker_parallel.py b/d123/common/multithreading/worker_parallel.py new file mode 100644 index 00000000..737bd486 --- /dev/null +++ b/d123/common/multithreading/worker_parallel.py @@ -0,0 +1,60 @@ +import concurrent +import concurrent.futures +import logging +from concurrent.futures import Future +from typing import Any, Iterable, List, Optional + +from tqdm import tqdm + +from d123.common.multithreading.worker_pool import ( + Task, + WorkerPool, + WorkerResources, + get_max_size_of_arguments, +) + +logger = logging.getLogger(__name__) + + +class SingleMachineParallelExecutor(WorkerPool): + """ + This worker distributes all tasks across multiple threads on this machine. + """ + + def __init__(self, use_process_pool: bool = False, max_workers: Optional[int] = None): + """ + Create worker with limited threads. + :param use_process_pool: if true, ProcessPoolExecutor will be used as executor, otherwise ThreadPoolExecutor. + :param max_workers: if available, use this number as used number of threads. + """ + # Set the number of available threads + number_of_cpus_per_node = max_workers if max_workers else WorkerResources.current_node_cpu_count() + + super().__init__( + WorkerResources( + number_of_nodes=1, number_of_cpus_per_node=number_of_cpus_per_node, number_of_gpus_per_node=0 + ) + ) + + # Create executor + self._executor = ( + concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) + if use_process_pool + else concurrent.futures.ThreadPoolExecutor(max_workers=number_of_cpus_per_node) + ) + + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """Inherited, see superclass.""" + return list( + tqdm( + self._executor.map(task.fn, *item_lists), + leave=False, + total=get_max_size_of_arguments(*item_lists), + desc="SingleMachineParallelExecutor", + disable=not verbose, + ) + ) + + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """Inherited, see superclass.""" + return self._executor.submit(task.fn, *args, **kwargs) diff --git a/d123/common/multithreading/worker_pool.py b/d123/common/multithreading/worker_pool.py new file mode 100644 index 00000000..8ffea8ed --- /dev/null +++ b/d123/common/multithreading/worker_pool.py @@ -0,0 +1,156 @@ +import abc +import logging +from concurrent.futures import Future +from dataclasses import dataclass +from typing import Any, Callable, Iterable, List, Optional, Tuple, Union + +from psutil import cpu_count + +logger = logging.getLogger(__name__) + + +def get_max_size_of_arguments(*item_lists: Iterable[List[Any]]) -> int: + """ + Find the argument with most elements. + e.g. [db, [arg1, arg2] -> 2. + :param item_lists: arguments where some of the arguments is a list. + :return: size of largest list. + """ + lengths = [len(items) for items in item_lists if isinstance(items, list)] + if len(list(set(lengths))) > 1: + raise RuntimeError(f"There exists lists with different element size = {lengths}!") + return max(lengths) if len(lengths) != 0 else 1 + + +def align_size_of_arguments(*item_lists: Iterable[List[Any]]) -> Tuple[int, Iterable[List[Any]]]: + """ + Align item lists by repeating elements in order to achieve the same size. + eg. [db, [arg1, arg2] -> [[db, db], [arg1, arg2]]. + :param item_lists: multiple arguments which will be used to call a function. + :return: arguments with same dimension, e.g., [[db, db], [arg1, arg1]]. + """ + max_size = get_max_size_of_arguments(*item_lists) + aligned_item_lists = [items if isinstance(items, list) else [items] * max_size for items in item_lists] + return max_size, aligned_item_lists + + +@dataclass(frozen=True) +class Task: + """This class represents a task that can be submitted to a worker with specific resource requirements.""" + + fn: Callable[..., Any] # Function that should be called with arguments + + # Number of CPUs required for this task + # if None, the default number of CPUs will be used that were specified in initialization + num_cpus: Optional[int] = None + + # Fraction of GPUs required for this task, this number can be also smaller than 1 + # It num_gpus is smaller than 1, only a part of a GPU is allocated to this task. + # NOTE: it is the responsibility of a user to make sure that a model fits into num_gpus + # if None, no GPU will be used + num_gpus: Optional[Union[int, float]] = None + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + """ + Call function with args. + :return: output from fn. + """ + return self.fn(*args, **kwargs) + + +@dataclass(frozen=True) +class WorkerResources: + """Data class to indicate resources used by workers.""" + + number_of_nodes: int # Number of available independent nodes + number_of_gpus_per_node: int # Number of GPUs per node + # Number of CPU logical cores per node, this can be smaller than + # current_node_cpu_count if user will decide to limit it + number_of_cpus_per_node: int + + @property + def number_of_threads(self) -> int: + """ + :return: the number of available threads across all nodes. + """ + return self.number_of_nodes * self.number_of_cpus_per_node + + @staticmethod + def current_node_cpu_count() -> int: + """ + :return: the number of logical cores on the current machine. + """ + return cpu_count(logical=True) # type: ignore + + +class WorkerPool(abc.ABC): + """ + This class executed function on list of arguments. This can either be distributed/parallel or sequential worker. + """ + + def __init__(self, config: WorkerResources): + """ + Initialize worker with resource description. + :param config: setup of this worker. + """ + self.config = config + + if self.config.number_of_threads < 1: + raise RuntimeError(f"Number of threads can not be 0, and it is {self.config.number_of_threads}!") + + logger.info(f"Worker: {self.__class__.__name__}") + logger.info(f"{self}") + + def map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """ + Run function with arguments from item_lists, this function will make sure all arguments have the same + number of elements. + :param task: function to be run. + :param item_lists: arguments to the function. + :param verbose: Whether to increase logger verbosity. + :return: type from the fn. + """ + max_size, aligned_item_lists = align_size_of_arguments(*item_lists) + + if verbose: + logger.info(f"Submitting {max_size} tasks!") + return self._map(task, *aligned_item_lists, verbose=verbose) + + @abc.abstractmethod + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """ + Run function with arguments from item_lists. This function can assume that all the args in item_lists have + the same number of elements. + :param fn: function to be run. + :param item_lists: arguments to the function. + :param number_of_elements: number of calls to the function. + :return: type from the fn. + """ + + @abc.abstractmethod + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """ + Submit a task to the worker. + :param task: to be submitted. + :param args: arguments for the task. + :param kwargs: keyword arguments for the task. + :return: future. + """ + + @property + def number_of_threads(self) -> int: + """ + :return: the number of available threads across all nodes. + """ + return self.config.number_of_threads + + def __str__(self) -> str: + """ + :return: string with information about this worker. + """ + return ( + f"Number of nodes: {self.config.number_of_nodes}\n" + f"Number of CPUs per node: {self.config.number_of_cpus_per_node}\n" + f"Number of GPUs per node: {self.config.number_of_gpus_per_node}\n" + f"Number of threads across all nodes: {self.config.number_of_threads}" + ) diff --git a/d123/common/multithreading/worker_ray.py b/d123/common/multithreading/worker_ray.py new file mode 100644 index 00000000..8bb12b04 --- /dev/null +++ b/d123/common/multithreading/worker_ray.py @@ -0,0 +1,166 @@ +import logging +import os +from concurrent.futures import Future +from pathlib import Path +from typing import Any, Iterable, List, Optional, Union + +import ray +from psutil import cpu_count + +from d123.common.multithreading.ray_execution import ray_map +from d123.common.multithreading.worker_pool import Task, WorkerPool, WorkerResources + +logger = logging.getLogger(__name__) + +# Silent botocore which is polluting the terminal because of serialization and deserialization +# with following message: INFO:botocore.credentials:Credentials found in config file: ~/.aws/config +logging.getLogger("botocore").setLevel(logging.WARNING) + + +def initialize_ray( + master_node_ip: Optional[str] = None, + threads_per_node: Optional[int] = None, + local_mode: bool = False, + log_to_driver: bool = True, + use_distributed: bool = False, +) -> WorkerResources: + """ + Initialize ray worker. + ENV_VAR_MASTER_NODE_IP="master node IP". + ENV_VAR_MASTER_NODE_PASSWORD="password to the master node". + ENV_VAR_NUM_NODES="number of nodes available". + :param master_node_ip: if available, ray will connect to remote cluster. + :param threads_per_node: Number of threads to use per node. + :param log_to_driver: If true, the output from all of the worker + processes on all nodes will be directed to the driver. + :param local_mode: If true, the code will be executed serially. This + is useful for debugging. + :param use_distributed: If true, and the env vars are available, + ray will launch in distributed mode + :return: created WorkerResources. + """ + # Env variables which are set through SLURM script + env_var_master_node_ip = "ip_head" + env_var_master_node_password = "redis_password" + env_var_num_nodes = "num_nodes" + + # Read number of CPU cores on current machine + number_of_cpus_per_node = threads_per_node if threads_per_node else cpu_count(logical=True) + + try: + import torch + + number_of_gpus_per_node = torch.cuda.device_count() if torch.cuda.is_available() else 0 + except ImportError: + number_of_gpus_per_node = 0 + if not number_of_gpus_per_node: + logger.info("Not using GPU in ray") + + # Find a way in how the ray should be initialized + if master_node_ip and use_distributed: + # Connect to ray remotely to node ip + logger.info(f"Connecting to cluster at: {master_node_ip}!") + ray.init(address=f"ray://{master_node_ip}:10001", local_mode=local_mode, log_to_driver=log_to_driver) + number_of_nodes = 1 + elif env_var_master_node_ip in os.environ and use_distributed: + # In this way, we started ray on the current machine which generated password and master node ip: + # It was started with "ray start --head" + number_of_nodes = int(os.environ[env_var_num_nodes]) + master_node_ip = os.environ[env_var_master_node_ip].split(":")[0] + redis_password = os.environ[env_var_master_node_password].split(":")[0] + logger.info(f"Connecting as part of a cluster at: {master_node_ip} with password: {redis_password}!") + # Connect to cluster, follow to https://docs.ray.io/en/latest/package-ref.html for more info + ray.init( + address="auto", + _node_ip_address=master_node_ip, + _redis_password=redis_password, + log_to_driver=log_to_driver, + local_mode=local_mode, + ) + else: + # In this case, we will just start ray directly from this script + number_of_nodes = 1 + logger.info("Starting ray local!") + ray.init( + num_cpus=number_of_cpus_per_node, + dashboard_host="0.0.0.0", + local_mode=local_mode, + log_to_driver=log_to_driver, + ) + + return WorkerResources( + number_of_nodes=number_of_nodes, + number_of_cpus_per_node=number_of_cpus_per_node, + number_of_gpus_per_node=number_of_gpus_per_node, + ) + + +class RayDistributed(WorkerPool): + """ + This worker uses ray to distribute work across all available threads. + """ + + def __init__( + self, + master_node_ip: Optional[str] = None, + threads_per_node: Optional[int] = None, + debug_mode: bool = False, + log_to_driver: bool = True, + output_dir: Optional[Union[str, Path]] = None, + logs_subdir: Optional[str] = "logs", + use_distributed: bool = False, + ): + """ + Initialize ray worker. + :param master_node_ip: if available, ray will connect to remote cluster. + :param threads_per_node: Number of threads to use per node. + :param debug_mode: If true, the code will be executed serially. This + is useful for debugging. + :param log_to_driver: If true, the output from all of the worker + processes on all nodes will be directed to the driver. + :param output_dir: Experiment output directory. + :param logs_subdir: Subdirectory inside experiment dir to store worker logs. + :param use_distributed: Boolean flag to explicitly enable/disable distributed computation + """ + self._master_node_ip = master_node_ip + self._threads_per_node = threads_per_node + self._local_mode = debug_mode + self._log_to_driver = log_to_driver + self._log_dir: Optional[Path] = Path(output_dir) / (logs_subdir or "") if output_dir is not None else None + self._use_distributed = use_distributed + super().__init__(self.initialize()) + + def initialize(self) -> WorkerResources: + """ + Initialize ray. + :return: created WorkerResources. + """ + # In case ray was already running, shut it down. This occurs mainly in tests + if ray.is_initialized(): + logger.warning("Ray is running, we will shut it down before starting again!") + ray.shutdown() + + return initialize_ray( + master_node_ip=self._master_node_ip, + threads_per_node=self._threads_per_node, + local_mode=self._local_mode, + log_to_driver=self._log_to_driver, + use_distributed=self._use_distributed, + ) + + def shutdown(self) -> None: + """ + Shutdown the worker and clear memory. + """ + ray.shutdown() + + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """Inherited, see superclass.""" + del verbose + return ray_map(task, *item_lists, log_dir=self._log_dir) # type: ignore + + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """Inherited, see superclass.""" + remote_fn = ray.remote(task.fn).options(num_gpus=task.num_gpus, num_cpus=task.num_cpus) + object_ids: ray._raylet.ObjectRef = remote_fn.remote(*args, **kwargs) + return object_ids.future() # type: ignore diff --git a/d123/common/multithreading/worker_sequential.py b/d123/common/multithreading/worker_sequential.py new file mode 100644 index 00000000..e3d436cb --- /dev/null +++ b/d123/common/multithreading/worker_sequential.py @@ -0,0 +1,46 @@ +import logging +from concurrent.futures import Future +from typing import Any, Iterable, List + +from tqdm import tqdm + +from d123.common.multithreading.worker_pool import ( + Task, + WorkerPool, + WorkerResources, + get_max_size_of_arguments, +) + +logger = logging.getLogger(__name__) + + +class Sequential(WorkerPool): + """ + This function does execute all functions sequentially. + """ + + def __init__(self) -> None: + """ + Initialize simple sequential worker. + """ + super().__init__(WorkerResources(number_of_nodes=1, number_of_cpus_per_node=1, number_of_gpus_per_node=0)) + + def _map(self, task: Task, *item_lists: Iterable[List[Any]], verbose: bool = False) -> List[Any]: + """Inherited, see superclass.""" + if task.num_cpus not in [None, 1]: + raise ValueError(f"Expected num_cpus to be 1 or unset for Sequential worker, got {task.num_cpus}") + output = [ + task.fn(*args) + for args in tqdm( + zip(*item_lists), + leave=False, + total=get_max_size_of_arguments(*item_lists), + desc="Sequential", + disable=not verbose, + ) + ] + return output + + def submit(self, task: Task, *args: Any, **kwargs: Any) -> Future[Any]: + """Inherited, see superclass.""" + raise NotImplementedError diff --git a/d123/common/multithreading/worker_utils.py b/d123/common/multithreading/worker_utils.py new file mode 100644 index 00000000..fbe0a753 --- /dev/null +++ b/d123/common/multithreading/worker_utils.py @@ -0,0 +1,36 @@ +from typing import Any, Callable, List, Optional + +import numpy as np +from psutil import cpu_count + +from d123.common.multithreading.worker_pool import Task, WorkerPool + + +def chunk_list(input_list: List[Any], num_chunks: Optional[int] = None) -> List[List[Any]]: + """ + Chunks a list to equal sized lists. The size of the last list might be truncated. + :param input_list: List to be chunked. + :param num_chunks: Number of chunks, equals to the number of cores if set to None. + :return: List of equal sized lists. + """ + num_chunks = num_chunks if num_chunks else cpu_count(logical=True) + chunks = np.array_split(input_list, num_chunks) # type: ignore + return [chunk.tolist() for chunk in chunks if len(chunk) != 0] + + +def worker_map(worker: WorkerPool, fn: Callable[..., List[Any]], input_objects: List[Any]) -> List[Any]: + """ + Map a list of objects through a worker. + :param worker: Worker pool to use for parallelization. + :param fn: Function to use when mapping. + :param input_objects: List of objects to map. + :return: List of mapped objects. + """ + if worker.number_of_threads == 0: + return fn(input_objects) + + object_chunks = chunk_list(input_objects, worker.number_of_threads) + scattered_objects = worker.map(Task(fn=fn), object_chunks) + output_objects = [result for results in scattered_objects for result in results] + + return output_objects diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 0bdc33d7..81989008 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -10,7 +10,6 @@ import numpy as np import pyarrow as pa -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json @@ -20,6 +19,7 @@ from d123.common.geometry.base import Point2D, Point3D from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index from d123.common.geometry.vector import Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.conversion.map.opendrive.elements.opendrive import OpenDrive from d123.dataset.conversion.map.opendrive.opendrive_converter import OpenDriveConverter diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index b4b9844a..59c791d7 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -16,7 +16,6 @@ from nuplan.database.nuplan_db_orm.lidar_pc import LidarPc from nuplan.database.nuplan_db_orm.nuplandb import NuPlanDB from nuplan.planning.simulation.observation.observation_type import CameraChannel -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map from pyquaternion import Quaternion from sqlalchemy import func @@ -36,6 +35,7 @@ from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3, BoundingBoxSE3Index from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import MAP_LOCATIONS, NuPlanMapConverter from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter diff --git a/d123/dataset/dataset_specific/raw_data_converter.py b/d123/dataset/dataset_specific/raw_data_converter.py index 3a7f826c..179f9ef5 100644 --- a/d123/dataset/dataset_specific/raw_data_converter.py +++ b/d123/dataset/dataset_specific/raw_data_converter.py @@ -3,7 +3,7 @@ from pathlib import Path from typing import List, Literal, Optional, Union -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool +from d123.common.multithreading.worker_utils import WorkerPool @dataclass diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 35c517de..dd6512fe 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -11,7 +11,6 @@ import numpy.typing as npt import pyarrow as pa import tensorflow as tf -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map from pyquaternion import Quaternion from waymo_open_dataset import dataset_pb2 @@ -27,6 +26,7 @@ from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array, get_rotation_matrix from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.dataset_specific.wopd.waymo_map_utils.wopd_map_utils import convert_wopd_map diff --git a/d123/dataset/maps/abstract_map.py b/d123/dataset/maps/abstract_map.py index 0dc483bb..681c90cf 100644 --- a/d123/dataset/maps/abstract_map.py +++ b/d123/dataset/maps/abstract_map.py @@ -17,6 +17,7 @@ class AbstractMap(abc.ABC): + @property @abc.abstractmethod def map_name(self) -> str: diff --git a/d123/dataset/scene/scene_builder.py b/d123/dataset/scene/scene_builder.py index 3c97e0a4..228fcb13 100644 --- a/d123/dataset/scene/scene_builder.py +++ b/d123/dataset/scene/scene_builder.py @@ -5,8 +5,7 @@ from pathlib import Path from typing import Iterator, List, Optional, Set, Union -from nuplan.planning.utils.multithreading.worker_utils import WorkerPool, worker_map - +from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.scene.abstract_scene import AbstractScene diff --git a/d123/script/builders/data_converter_builder.py b/d123/script/builders/data_converter_builder.py index 6047b008..d9c54004 100644 --- a/d123/script/builders/data_converter_builder.py +++ b/d123/script/builders/data_converter_builder.py @@ -2,10 +2,10 @@ from typing import List from hydra.utils import instantiate -from nuplan.planning.script.builders.utils.utils_type import validate_type from omegaconf import DictConfig from d123.dataset.dataset_specific.raw_data_converter import RawDataConverter +from d123.script.builders.utils.utils_type import validate_type logger = logging.getLogger(__name__) diff --git a/d123/script/builders/utils/utils_type.py b/d123/script/builders/utils/utils_type.py new file mode 100644 index 00000000..0fe5e9fd --- /dev/null +++ b/d123/script/builders/utils/utils_type.py @@ -0,0 +1,73 @@ +from typing import Any, Callable, Dict, Type, Union, cast + +from hydra._internal.utils import _locate +from omegaconf import DictConfig + + +def is_TorchModuleWrapper_config(cfg: DictConfig) -> bool: + """ + Check whether the config is meant for a TorchModuleWrapper + :param cfg: config + :return: True if model_config and checkpoint_path is in the cfg, False otherwise + """ + return "model_config" in cfg and "checkpoint_path" in cfg + + +def is_target_type(cfg: DictConfig, target_type: Union[Type[Any], Callable[..., Any]]) -> bool: + """ + Check whether the config's resolved type matches the target type or callable. + :param cfg: config + :param target_type: Type or callable to check against. + :return: Whether cfg._target_ matches the target_type. + """ + return bool(_locate(cfg._target_) == target_type) + + +def validate_type(instantiated_class: Any, desired_type: Type[Any]) -> None: + """ + Validate that constructed type is indeed the desired one + :param instantiated_class: class that was created + :param desired_type: type that the created class should have + """ + assert isinstance( + instantiated_class, desired_type + ), f"Class to be of type {desired_type}, but is {type(instantiated_class)}!" + + +def are_the_same_type(lhs: Any, rhs: Any) -> None: + """ + Validate that lhs and rhs are of the same type + :param lhs: left argument + :param rhs: right argument + """ + lhs_type = type(lhs) + rhs_type = type(rhs) + assert lhs_type == rhs_type, f"Lhs and Rhs are not of the same type! {lhs_type} != {rhs_type}!" + + +def validate_dict_type(instantiated_dict: Dict[str, Any], desired_type: Type[Any]) -> None: + """ + Validate that all entries in dict is indeed the desired one + :param instantiated_dict: dictionary that was created + :param desired_type: type that the created class should have + """ + for value in instantiated_dict.values(): + if isinstance(value, dict): + validate_dict_type(value, desired_type) + else: + validate_type(value, desired_type) + + +def find_builder_in_config(cfg: DictConfig, desired_type: Type[Any]) -> DictConfig: + """ + Find the corresponding config for the desired builder + :param cfg: config structured as a dictionary + :param desired_type: desired builder type + :return: found config + @raise ValueError if the config cannot be found for the builder + """ + for cfg_builder in cfg.values(): + if is_target_type(cfg_builder, desired_type): + return cast(DictConfig, cfg_builder) + + raise ValueError(f"Config does not exist for builder type: {desired_type}!") diff --git a/d123/script/builders/worker_pool_builder.py b/d123/script/builders/worker_pool_builder.py index 4d13877d..123ec3c3 100644 --- a/d123/script/builders/worker_pool_builder.py +++ b/d123/script/builders/worker_pool_builder.py @@ -1,12 +1,11 @@ import logging from hydra.utils import instantiate -from nuplan.planning.script.builders.utils.utils_type import is_target_type, validate_type -from nuplan.planning.utils.multithreading.worker_parallel import SingleMachineParallelExecutor -from nuplan.planning.utils.multithreading.worker_pool import WorkerPool -from nuplan.planning.utils.multithreading.worker_sequential import Sequential from omegaconf import DictConfig +from d123.common.multithreading.worker_pool import WorkerPool +from d123.script.builders.utils.utils_type import validate_type + logger = logging.getLogger(__name__) @@ -17,12 +16,7 @@ def build_worker(cfg: DictConfig) -> WorkerPool: :return: Instance of WorkerPool. """ logger.info("Building WorkerPool...") - worker: WorkerPool = ( - instantiate(cfg.worker) - if (is_target_type(cfg.worker, SingleMachineParallelExecutor) or is_target_type(cfg.worker, Sequential)) - else instantiate(cfg.worker, output_dir=cfg.output_dir) - ) + worker: WorkerPool = instantiate(cfg.worker) validate_type(worker, WorkerPool) - logger.info("Building WorkerPool...DONE!") return worker diff --git a/d123/script/config/common/worker/ray_distributed.yaml b/d123/script/config/common/worker/ray_distributed.yaml index c51de1d3..2c101f66 100644 --- a/d123/script/config/common/worker/ray_distributed.yaml +++ b/d123/script/config/common/worker/ray_distributed.yaml @@ -1,4 +1,4 @@ -_target_: nuplan.planning.utils.multithreading.worker_ray.RayDistributed +_target_: d123.common.multithreading.worker_ray.RayDistributed _convert_: 'all' master_node_ip: null # Set to a master node IP if you desire to connect to cluster remotely threads_per_node: null # Number of CPU threads to use per node, "null" means all threads available diff --git a/d123/script/config/common/worker/sequential.yaml b/d123/script/config/common/worker/sequential.yaml index 0b7955e9..ea004415 100644 --- a/d123/script/config/common/worker/sequential.yaml +++ b/d123/script/config/common/worker/sequential.yaml @@ -1,2 +1,2 @@ -_target_: nuplan.planning.utils.multithreading.worker_sequential.Sequential +_target_: d123.common.multithreading.worker_sequential.Sequential _convert_: 'all' diff --git a/d123/script/config/common/worker/single_machine_thread_pool.yaml b/d123/script/config/common/worker/single_machine_thread_pool.yaml index ba1b9657..ace763f6 100644 --- a/d123/script/config/common/worker/single_machine_thread_pool.yaml +++ b/d123/script/config/common/worker/single_machine_thread_pool.yaml @@ -1,4 +1,4 @@ -_target_: nuplan.planning.utils.multithreading.worker_parallel.SingleMachineParallelExecutor +_target_: d123.common.multithreading.worker_parallel.SingleMachineParallelExecutor _convert_: 'all' use_process_pool: True # If true, use ProcessPoolExecutor as the backend, otherwise uses ThreadPoolExecutor max_workers: 16 # Number of CPU workers (threads/processes) to use per node, "null" means all available diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 1eb51d46..6357115c 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,8 +15,8 @@ defaults: - default_dataset_paths - _self_ - datasets: - # - nuplan_private_dataset - - carla_dataset + - nuplan_private_dataset + # - carla_dataset # - wopd_dataset force_log_conversion: True diff --git a/d123/script/run_preprocessing.py b/d123/script/run_preprocessing.py index d8c426d7..4ed6cf97 100644 --- a/d123/script/run_preprocessing.py +++ b/d123/script/run_preprocessing.py @@ -6,9 +6,9 @@ import hydra import lightning as L -from nuplan.planning.utils.multithreading.worker_utils import worker_map from omegaconf import DictConfig +from d123.common.multithreading.worker_utils import worker_map from d123.dataset.scene.abstract_scene import AbstractScene from d123.script.builders.scene_builder_builder import build_scene_builder from d123.script.builders.scene_filter_builder import build_scene_filter diff --git a/d123/script/run_simulation.py b/d123/script/run_simulation.py index a1be003c..393e4367 100644 --- a/d123/script/run_simulation.py +++ b/d123/script/run_simulation.py @@ -6,10 +6,10 @@ import hydra import lightning as L import pandas as pd -from nuplan.planning.utils.multithreading.worker_utils import worker_map from omegaconf import DictConfig from tqdm import tqdm +from d123.common.multithreading.worker_utils import worker_map from d123.dataset.scene.abstract_scene import AbstractScene from d123.script.builders.scene_builder_builder import build_scene_builder from d123.script.builders.scene_filter_builder import build_scene_filter diff --git a/d123/simulation/gym/environment/environment_wrapper.py b/d123/simulation/gym/environment/environment_wrapper.py index ef0723fe..fe3cd5b9 100644 --- a/d123/simulation/gym/environment/environment_wrapper.py +++ b/d123/simulation/gym/environment/environment_wrapper.py @@ -11,15 +11,15 @@ from d123.simulation.gym.environment.gym_observation.abstract_gym_observation import ( AbstractGymObservation, ) +from d123.simulation.gym.environment.output_converter.abstract_output_converter import ( + AbstractOutputConverter, +) from d123.simulation.gym.environment.reward_builder.abstract_reward_builder import AbstractRewardBuilder from d123.simulation.gym.environment.scenario_sampler.abstract_scenario_sampler import AbstractScenarioSampler from d123.simulation.gym.environment.simulation_builder.abstract_simulation_builder import ( AbstractSimulationBuilder, ) from d123.simulation.gym.environment.simulation_wrapper import SimulationWrapper -from d123.simulation.gym.environment.output_converter.abstract_output_converter import ( - AbstractOutputConverter, -) logger = logging.getLogger(__name__) diff --git a/d123/simulation/gym/environment/output_converter/abstract_output_converter.py b/d123/simulation/gym/environment/output_converter/abstract_output_converter.py index c045d7bf..94f3ef0c 100644 --- a/d123/simulation/gym/environment/output_converter/abstract_output_converter.py +++ b/d123/simulation/gym/environment/output_converter/abstract_output_converter.py @@ -5,7 +5,6 @@ import numpy.typing as npt from gymnasium import spaces - from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2 from d123.simulation.planning.planner_output.abstract_planner_output import AbstractPlannerOutput diff --git a/d123/training/models/sim_agent/smart/layers/attention_layer.py b/d123/training/models/sim_agent/smart/layers/attention_layer.py index e5190552..9452662e 100644 --- a/d123/training/models/sim_agent/smart/layers/attention_layer.py +++ b/d123/training/models/sim_agent/smart/layers/attention_layer.py @@ -17,7 +17,7 @@ def __init__( dropout: float, bipartite: bool, has_pos_emb: bool, - **kwargs + **kwargs, ) -> None: super(AttentionLayer, self).__init__(aggr="add", node_dim=0, **kwargs) self.num_heads = num_heads diff --git a/docs/conf.py b/docs/conf.py index c9379f0d..9d7e8c5f 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -7,8 +7,8 @@ # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information project = "d123" -copyright = "2025, DanielDauner" -author = "DanielDauner" +copyright = "2025, 123D Contributors" +author = "123D Contributors" release = "v0.0.6" # -- General configuration --------------------------------------------------- @@ -59,7 +59,6 @@ "prev_next_buttons_location": "bottom", "style_external_links": False, "vcs_pageview_mode": "", - "style_nav_header_background": "#FFE100", # Toc options "collapse_navigation": True, "sticky_navigation": True, diff --git a/docs/index.rst b/docs/index.rst index 53aedd4d..4ea69ff9 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -12,7 +12,7 @@ documentation for details. .. toctree:: - :hidden: + :maxdepth: 2 :caption: Contents: installation diff --git a/docs/requirements.in b/docs/requirements.in deleted file mode 100644 index acbc25d5..00000000 --- a/docs/requirements.in +++ /dev/null @@ -1,2 +0,0 @@ -Sphinx>=5,<6 -sphinx_rtd_theme diff --git a/docs/requirements.txt b/docs/requirements.txt deleted file mode 100644 index 32230734..00000000 --- a/docs/requirements.txt +++ /dev/null @@ -1,60 +0,0 @@ -# -# This file is autogenerated by pip-compile with python 3.10 -# To update, run: -# -# pip-compile docs/requirements.in -# -alabaster==0.7.12 - # via sphinx -babel==2.10.3 - # via sphinx -certifi==2022.6.15 - # via requests -charset-normalizer==2.1.0 - # via requests -docutils==0.17.1 - # via - # sphinx - # sphinx-rtd-theme -idna==3.3 - # via requests -imagesize==1.4.1 - # via sphinx -jinja2==3.1.2 -m2r2 - # via sphinx -markupsafe==2.1.1 - # via requests -myst-parser - # via jinja2 -packaging==21.3 - # via sphinx -pygments==2.12.0 - # via sphinx -pyparsing==3.0.9 - # via packaging -pytz==2022.1 - # via babel -requests==2.28.1 - # via sphinx -snowballstemmer==2.2.0 - # via sphinx -sphinx==5.0.2 - # via - # -r docs/requirements.in - # sphinx-rtd-theme -sphinx-rtd-theme==1.0.0 - # via -r docs/requirements.in -sphinxcontrib-applehelp==1.0.2 - # via sphinx -sphinxcontrib-devhelp==1.0.2 - # via sphinx -sphinxcontrib-htmlhelp==2.0.0 - # via sphinx -sphinxcontrib-jsmath==1.0.1 - # via sphinx -sphinxcontrib-qthelp==1.0.3 - # via sphinx -sphinxcontrib-serializinghtml==1.1.5 - # via sphinx -urllib3==1.26.9 diff --git a/environment.yml b/environment.yml index eeb40a17..53e458bc 100644 --- a/environment.yml +++ b/environment.yml @@ -5,5 +5,3 @@ dependencies: - python=3.12 - pip=25.1.1 - nb_conda_kernels - - pip: - - -r requirements.txt diff --git a/notebooks/deprecated/collect_sim_metrics_gym.ipynb b/notebooks/deprecated/collect_sim_metrics_gym.ipynb index 2e5b0f82..40100cd6 100644 --- a/notebooks/deprecated/collect_sim_metrics_gym.ipynb +++ b/notebooks/deprecated/collect_sim_metrics_gym.ipynb @@ -9,7 +9,7 @@ "source": [ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n" + "from d123.common.multithreading.worker_sequential import Sequential\n" ] }, { diff --git a/notebooks/deprecated/extraction_testing.ipynb b/notebooks/deprecated/extraction_testing.ipynb index 652601ad..d782d7ba 100644 --- a/notebooks/deprecated/extraction_testing.ipynb +++ b/notebooks/deprecated/extraction_testing.ipynb @@ -14,7 +14,7 @@ "\n", "\n", "import pyarrow as pa\n", - "from nuplan.planning.utils.multithreading.worker_pool import WorkerPool\n", + "from d123.common.multithreading.worker_pool import WorkerPool\n", "\n", "from d123.dataset.arrow.helper import open_arrow_arrow_table\n", "from d123.dataset.dataset_specific.nuplan.nuplan_data_processor import worker_map\n", diff --git a/notebooks/deprecated/test_scene_builder.ipynb b/notebooks/deprecated/test_scene_builder.ipynb index 2f4a608d..00de24f9 100644 --- a/notebooks/deprecated/test_scene_builder.ipynb +++ b/notebooks/deprecated/test_scene_builder.ipynb @@ -10,7 +10,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n" + "from d123.common.multithreading.worker_sequential import Sequential\n" ] }, { diff --git a/notebooks/deprecated/test_waypoints.ipynb b/notebooks/deprecated/test_waypoints.ipynb index a2ce5642..86d2ea9c 100644 --- a/notebooks/deprecated/test_waypoints.ipynb +++ b/notebooks/deprecated/test_waypoints.ipynb @@ -18,7 +18,7 @@ "metadata": {}, "outputs": [], "source": [ - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from d123.common.multithreading.worker_sequential import Sequential" ] }, { diff --git a/notebooks/gym/test_gym.ipynb b/notebooks/gym/test_gym.ipynb index 5a5f40f2..bb6cf7dc 100644 --- a/notebooks/gym/test_gym.ipynb +++ b/notebooks/gym/test_gym.ipynb @@ -10,8 +10,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" + "from d123.common.multithreading.worker_sequential import Sequential\n", + "# from d123.common.multithreading.worker_ray import RayDistributed" ] }, { diff --git a/notebooks/gym/test_simulation_2d.ipynb b/notebooks/gym/test_simulation_2d.ipynb index 26dad379..60c22ceb 100644 --- a/notebooks/gym/test_simulation_2d.ipynb +++ b/notebooks/gym/test_simulation_2d.ipynb @@ -10,8 +10,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" + "from d123.common.multithreading.worker_sequential import Sequential\n", + "# from d123.common.multithreading.worker_ray import RayDistributed" ] }, { diff --git a/notebooks/scene_rendering.ipynb b/notebooks/scene_rendering.ipynb index e575365e..dad53954 100644 --- a/notebooks/scene_rendering.ipynb +++ b/notebooks/scene_rendering.ipynb @@ -10,9 +10,9 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "\n", - "# from nuplan.planning.utils.multithreading.worker_ray import RayDistributed" + "# from d123.common.multithreading.worker_ray import RayDistributed" ] }, { diff --git a/notebooks/scene_sensor_loading.ipynb b/notebooks/scene_sensor_loading.ipynb index 1a4093fe..ed143de1 100644 --- a/notebooks/scene_sensor_loading.ipynb +++ b/notebooks/scene_sensor_loading.ipynb @@ -11,8 +11,8 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", - "# from nuplan.planning.utils.multithreading.worker_ray impo\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", + "# from d123.common.multithreading.worker_ray impo\n", "# rt RayDistribute\n", "\n", "\n" diff --git a/notebooks/smarty/smart_feature_testing.ipynb b/notebooks/smarty/smart_feature_testing.ipynb index 2f35f7de..c134baf9 100644 --- a/notebooks/smarty/smart_feature_testing.ipynb +++ b/notebooks/smarty/smart_feature_testing.ipynb @@ -36,7 +36,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "\n", "\n", "log_names = None\n", diff --git a/notebooks/smarty/smart_rollout.ipynb b/notebooks/smarty/smart_rollout.ipynb index ae38d252..3f3aac67 100644 --- a/notebooks/smarty/smart_rollout.ipynb +++ b/notebooks/smarty/smart_rollout.ipynb @@ -17,7 +17,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential" + "from d123.common.multithreading.worker_sequential import Sequential" ] }, { diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 79b1e8e9..6ead233b 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -10,7 +10,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "from d123.common.datatypes.sensor.camera import CameraType" ] }, @@ -25,9 +25,9 @@ "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", "\n", "\n", - "splits = [\"wopd_train\"]\n", + "# splits = [\"wopd_train\"]\n", "# splits = [\"carla\"]\n", - "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"nuplan_private_test\"]\n", "# log_names = None\n", "\n", "\n", @@ -42,7 +42,7 @@ " history_s=0.0,\n", " timestamp_threshold_s=20,\n", " shuffle=True,\n", - " camera_types=[CameraType.CAM_F0],\n", + " # camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", "worker = Sequential()\n", @@ -188,18 +188,18 @@ " 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", @@ -220,8 +220,8 @@ " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", - " # add_box_detections_to_ax(ax, box_detections)\n", - " # add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + " 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", @@ -243,7 +243,7 @@ "\n", "\n", "scene_index = 5\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=150, radius=300)\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=150, radius=40)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] @@ -327,7 +327,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123", + "display_name": "d123_dev", "language": "python", "name": "python3" }, diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index 049fff6d..aa3ed869 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -11,7 +11,7 @@ "\n", "import matplotlib.pyplot as plt\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "\n", "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index b6cb640b..740ce2e2 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -10,7 +10,7 @@ "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", "from d123.dataset.scene.scene_filter import SceneFilter\n", "\n", - "from nuplan.planning.utils.multithreading.worker_sequential import Sequential\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", "from d123.common.datatypes.sensor.camera import CameraType" ] }, diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..6fb9c0e5 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,91 @@ +[tool.isort] +profile = "black" + +[build-system] +requires = ["setuptools>=58", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Programming Language :: Python :: 3.9", + "License :: OSI Approved :: Apache Software License", +] +name = "d123" +version = "v0.0.6" +authors = [{ name = "Daniel Dauner", email = "daniel.dauner@gmail.com" }] +description = "TODO" +readme = "README.md" +requires-python = ">=3.9" +license = {text = "Apache-2.0"} +dependencies = [ + "nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2", + "aioboto3", + "aiofiles", + "bokeh", + "casadi", + "control", + "Fiona", + "geopandas", + "guppy3", + "joblib", + "matplotlib", + "nest_asyncio", + "numpy", + "opencv-python", + "pandas", + "Pillow", + "psutil", + "pyarrow", + "pyinstrument", + "pyogrio", + "pyquaternion", + "pytest", + "rasterio", + "ray", + "retry", + "rtree", + "scipy", + "selenium", + "setuptools", + "shapely>=2.0.0", + "SQLAlchemy==1.4.27", + "sympy", + "tornado", + "tqdm", + "ujson", + "notebook", + "pre-commit", + "cachetools", + "hydra_colorlog", + "hydra-core", + "lxml", + "trimesh", + "viser", +] + +[project.optional-dependencies] +dev = ["black", "isort", "pre-commit"] +docs = [ + "Sphinx", + "sphinx-rtd-theme", + "myst-parser", +] +waymo = [ + "tensorflow==2.11.0", + "waymo-open-dataset-tf-2-11-0", + "intervaltree", +] +av2 = [ + "av2==0.2.1", +] + +[tool.setuptools.packages.find] +where = ["."] +include = ["d123*"] # Only include d123 package +exclude = ["notebooks*", "tests*", "docs*"] # Explicitly exclude notebooks + +[project.urls] +"Homepage" = "https://github.com/DanielDauner/d123" +"Bug Tracker" = "https://github.com/DanielDauner/d123/issues" diff --git a/requirements.txt b/requirements.txt index f0c697e2..4c17c209 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,59 +1,58 @@ -nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2 +# nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2 -# nuplan requirements -aioboto3 -aiofiles -bokeh -casadi -control -Fiona -geopandas -guppy3 -joblib -matplotlib -nest_asyncio -numpy -opencv-python -pandas -Pillow -psutil -pyarrow -pyinstrument -pyogrio -pyquaternion -pytest -rasterio -ray -retry -rtree -scipy -selenium -setuptools -shapely==2.1.1 -SQLAlchemy==1.4.27 -sympy -tornado -tqdm -ujson -notebook -pre-commit -cachetools +# # # nuplan requirements +# aioboto3 +# aiofiles +# bokeh +# casadi +# control +# Fiona +# geopandas +# guppy3 +# joblib +# matplotlib +# nest_asyncio +# numpy +# opencv-python +# pandas +# Pillow +# psutil +# pyarrow +# pyinstrument +# pyogrio +# pyquaternion +# pytest +# rasterio +# ray +# retry +# rtree +# scipy +# selenium +# setuptools +# shapely>=2.0.0 +# SQLAlchemy==1.4.27 +# sympy +# tornado +# tqdm +# ujson +# notebook +# pre-commit +# cachetools -# hydra -hydra_colorlog -hydra-core +# # hydra +# hydra_colorlog +# hydra-core -# d123 -lxml -trimesh -viser -gym==0.17.2 -gymnasium==0.26.3 +# # d123 +# lxml +# trimesh +# viser +# gym==0.17.2 +# gymnasium==0.26.3 -# torch & lighting -torch==2.6.0 -torchvision==0.21.0 -lightning -tensorboard -protobuf==4.25.3 -pre-commit +# # torch & lighting +# torch==2.6.0 +# torchvision==0.21.0 +# lightning +# tensorboard +# protobuf==4.25.3 diff --git a/scripts/download/download_av2.sh b/scripts/download/download_av2.sh new file mode 100644 index 00000000..7c63854a --- /dev/null +++ b/scripts/download/download_av2.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash + +# Dataset URIs +# s3://argoverse/datasets/av2/sensor/ +# s3://argoverse/datasets/av2/lidar/ +# s3://argoverse/datasets/av2/motion-forecasting/ +# s3://argoverse/datasets/av2/tbv/ + +DATASET_NAMES=("sensor" "lidar" "motion_forecasting" "tbv") +TARGET_DIR="/path/to/argoverse" + +for DATASET_NAME in "${DATASET_NAMES[@]}"; do + mkdir -p "$TARGET_DIR/$DATASET_NAME" + s5cmd --no-sign-request cp "s3://argoverse/datasets/av2/$DATASET_NAME/*" "$TARGET_DIR/$DATASET_NAME" +done diff --git a/scripts/download/download_nuplan.sh b/scripts/download/download_nuplan_logs.sh similarity index 100% rename from scripts/download/download_nuplan.sh rename to scripts/download/download_nuplan_logs.sh diff --git a/scripts/download/download_nuplan_sensor.sh b/scripts/download/download_nuplan_sensor.sh new file mode 100644 index 00000000..6cd3d246 --- /dev/null +++ b/scripts/download/download_nuplan_sensor.sh @@ -0,0 +1,14 @@ +# NOTE: Please check the LICENSE file when downloading the nuPlan dataset +# wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/LICENSE + +# train: nuplan_train + +# val: nuplan_val + +# test: nuplan_test + +# mini: nuplan_mini_train, nuplan_mini_val, nuplan_mini_test +for split in {0..8}; do + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_camera_${split}.zip + wget https://motional-nuplan.s3-ap-northeast-1.amazonaws.com/public/nuplan-v1.1/sensor_blobs/mini_set/nuplan-v1.1_mini_lidar_${split}.zip +done diff --git a/setup.py b/setup.py deleted file mode 100644 index 8c516b45..00000000 --- a/setup.py +++ /dev/null @@ -1,30 +0,0 @@ -import os - -import setuptools - -# Change directory to allow installation from anywhere -script_folder = os.path.dirname(os.path.realpath(__file__)) -os.chdir(script_folder) - -with open("requirements.txt") as f: - requirements = f.read().splitlines() - -# Installs -setuptools.setup( - name="d123", - version="0.0.5", - author="Daniel Dauner", - author_email="daniel.dauner@gmail.com", - description="TODO", - url="https://github.com/autonomousvision/d123", - python_requires=">=3.9", - packages=["d123"], - package_dir={"": "."}, - classifiers=[ - "Programming Language :: Python :: 3.9", - "Operating System :: OS Independent", - "License :: Free for non-commercial use", - ], - license="apache-2.0", - install_requires=requirements, -) From 54f327d0d8bb6d143832cfbda9d8e6875544f693 Mon Sep 17 00:00:00 2001 From: DanielDauner Date: Sat, 16 Aug 2025 15:12:07 +0200 Subject: [PATCH 33/42] Start writing the documentation including contribution guidelines. --- docs/contributing.md | 118 ++++++++++++++++++++++++++++++++++++++++++ docs/index.rst | 4 +- docs/visualization.md | 11 ++++ pyproject.toml | 12 ++++- 4 files changed, 142 insertions(+), 3 deletions(-) create mode 100644 docs/contributing.md create mode 100644 docs/visualization.md diff --git a/docs/contributing.md b/docs/contributing.md new file mode 100644 index 00000000..e447d876 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1,118 @@ + +# Contributing to 123D + +Thank you for your interest in contributing to 123D! This guide will help you get started with the development process. + +## Getting Started + +### 1. Clone the Repository + +```sh +git clone git@github.com:DanielDauner/d123.git +cd d123 +``` + +### 2. Install the pip-package + +```sh +conda env create -f environment.yml --name d123_dev # Optional +conda activate d123_dev +pip install -e .[dev] +pre-commit install +``` + +.. note:: + We might remove the conda environment in the future, but leave the file in the repo during development. + + +### 3. Managing dependencies + +One principal of 123D is to keep *minimal dependencies*. However, various datasets require problematic (or even outdated) dependencies in order to load or preprocess the dataset. In this case, you can add optional dependencies to the `pyproject.toml` install file. You can follow examples of Waymo or nuPlan. These optional dependencies can be install with + +```sh +pip install -e .[dev,waymo,nuplan] +``` +where you can combined the different optional dependencies. + +The optional dependencies should only be required for data pre-processing. If a dataset allows to load sensor data dynamically from the original dataset, please encapsule the import accordingly, e.g. + +```python +import numpy as np +import numpy.typing as npt + +def load_camera_from_outdated_dataset(file_path: str) -> npt.NDArray[np.uint8]: + try: + from outdated_dataset import load_camera_image + except ImportError: + raise ImportError( + "Optional dependency 'outdated_dataset' is required to load camera images from this dataset. " + "Please install it using: pip install .[outdated_dataset]" + ) + return load_camera_image(file_path) +``` + + +## Code Style and Formatting + +We maintain consistent code quality using the following tools: +- **[Black](https://black.readthedocs.io/)** - Code formatter +- **[isort](https://pycqa.github.io/isort/)** - Import statement formatter +- **[flake8](https://flake8.pycqa.github.io/)** - Style guide enforcement +- **[pytest](https://docs.pytest.org/)** - Testing framework for unit and integration tests +- **[pre-commit](https://pre-commit.com/)** - Framework for managing and running Git hooks to automate code quality checks + + +.. note:: + If you're using VSCode, it is recommended to install the [Black](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter), [isort](https://marketplace.visualstudio.com/items?itemName=ms-python.isort), and [Flake8](https://marketplace.visualstudio.com/items?itemName=ms-python.flake8) plugins. + + + +### Editor Setup + +**VS Code Users:** +If you're using VSCode, it is recommended to install the following plguins: +- [Black](https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter) - see above. +- [isort](https://marketplace.visualstudio.com/items?itemName=ms-python.isort) - see above. +- [Flake8](https://marketplace.visualstudio.com/items?itemName=ms-python.flake8) - see above. +- [autodocstring](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) - Creating docstrings (please set `"autoDocstring.docstringFormat": "sphinx-notypes"`). +- [Code Spell Checker](https://marketplace.visualstudio.com/items?itemName=streetsidesoftware.code-spell-checker) - A basic spell checker. + + +**Other Editors:** +Similar plugins are available for most popular editors including PyCharm, Vim, Emacs, and Sublime Text. + + +## Documentation Requirements + +### Docstrings +- **Development:** Docstrings are encouraged but not strictly required during active development +- **Release:** All public functions, classes, and modules must have comprehensive docstrings before release +- **Format:** Use [Sphinx-style docstrings](https://sphinx-rtd-tutorial.readthedocs.io/en/latest/docstrings.html) + +**VS Code Users:** The [autoDocstring extension](https://marketplace.visualstudio.com/items?itemName=njpwerner.autodocstring) can help generate properly formatted docstrings. + +### Type Hints +- **Required:** All function parameters and return values must include type hints +- **Style:** Follow [PEP 484](https://peps.python.org/pep-0484/) conventions + +### Sphinx documentation + +All datasets should be included in the `/docs/datasets.md` documentation. Please follow the documentation format of other datasets. + +You can install relevant dependencies for editing the public documentation via: +```sh +pip install -e .[docs] +``` + +It is recommended to uses [sphinx-autobuild](https://github.com/sphinx-doc/sphinx-autobuild) (installed above) to edit and view the documentation. You can run: +```sh +sphinx-autobuild docs docs/_build/html +``` + +## Adding new datasets +TODO + + +## Questions? + +If you have any questions about contributing, please open an issue or reach out to the maintainers. diff --git a/docs/index.rst b/docs/index.rst index 4ea69ff9..9f6c3518 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -3,7 +3,7 @@ You can adapt this file completely to your liking, but it should at least contain the root `toctree` directive. -123d documentation +123D Documentation ================== Add your content using ``reStructuredText`` syntax. See the @@ -18,3 +18,5 @@ documentation for details. installation datasets schema + visualization + contributing diff --git a/docs/visualization.md b/docs/visualization.md new file mode 100644 index 00000000..9fe2e3cf --- /dev/null +++ b/docs/visualization.md @@ -0,0 +1,11 @@ + +# Visualization + + +## Matplotlib + + +## Viser + + +## Bokeh diff --git a/pyproject.toml b/pyproject.toml index 6fb9c0e5..d5419010 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,7 +20,6 @@ readme = "README.md" requires-python = ">=3.9" license = {text = "Apache-2.0"} dependencies = [ - "nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2", "aioboto3", "aiofiles", "bokeh", @@ -66,12 +65,21 @@ dependencies = [ ] [project.optional-dependencies] -dev = ["black", "isort", "pre-commit"] +dev = [ + "black", + "isort", + "flake8", + "pre-commit", +] docs = [ "Sphinx", "sphinx-rtd-theme", + "sphinx-autobuild", "myst-parser", ] +nuplan = [ + "nuplan-devkit @ git+https://github.com/motional/nuplan-devkit/@nuplan-devkit-v1.2", +] waymo = [ "tensorflow==2.11.0", "waymo-open-dataset-tf-2-11-0", From 10f0482989d8524b8b81a3da0f3d60d3b7492610 Mon Sep 17 00:00:00 2001 From: DanielDauner Date: Sat, 16 Aug 2025 17:49:34 +0200 Subject: [PATCH 34/42] Add dataset and contributing documentation. --- docs/datasets.md | 203 ------------------------------------ docs/datasets/carla.rst | 90 ++++++++++++++++ docs/datasets/index.rst | 16 +++ docs/datasets/kitti-360.rst | 87 ++++++++++++++++ docs/datasets/nuplan.rst | 93 +++++++++++++++++ docs/datasets/nuscenes.rst | 90 ++++++++++++++++ docs/datasets/template.rst | 90 ++++++++++++++++ docs/datasets/wopd.rst | 88 ++++++++++++++++ docs/index.rst | 4 +- 9 files changed, 556 insertions(+), 205 deletions(-) delete mode 100644 docs/datasets.md create mode 100644 docs/datasets/carla.rst create mode 100644 docs/datasets/index.rst create mode 100644 docs/datasets/kitti-360.rst create mode 100644 docs/datasets/nuplan.rst create mode 100644 docs/datasets/nuscenes.rst create mode 100644 docs/datasets/template.rst create mode 100644 docs/datasets/wopd.rst diff --git a/docs/datasets.md b/docs/datasets.md deleted file mode 100644 index 31ebda18..00000000 --- a/docs/datasets.md +++ /dev/null @@ -1,203 +0,0 @@ - -# Datasets - - -## nuPlan -``` -0: { - "camera_name": "CAM_F0", - "original_size": (1080, 1920), - "egocar_visible": False -}, -1: { - "camera_name": "CAM_L0", - "original_size": (1080, 1920), - "egocar_visible": True -}, -2: { - "camera_name": "CAM_R0", - "original_size": (1080, 1920), - "egocar_visible": True -}, -3: { - "camera_name": "CAM_L1", - "original_size": (1080, 1920), - "egocar_visible": False -}, -4: { - "camera_name": "CAM_R1", - "original_size": (1080, 1920), - "egocar_visible": False -}, -5: { - "camera_name": "CAM_L2", - "original_size": (1080, 1920), - "egocar_visible": True -}, -6: { - "camera_name": "CAM_R2", - "original_size": (1080, 1920), - "egocar_visible": True -}, -7: { - "camera_name": "CAM_B0", - "original_size": (1080, 1920), - "egocar_visible": False -}, -``` - - -## nuScenes -``` -0: { - "camera_name": "CAM_FRONT", - "original_size": (900, 1600), - "egocar_visible": False -}, -1: { - "camera_name": "CAM_FRONT_LEFT", - "original_size": (900, 1600), - "egocar_visible": False -}, -2: { - "camera_name": "CAM_FRONT_RIGHT", - "original_size": (900, 1600), - "egocar_visible": False -}, -3: { - "camera_name": "CAM_BACK_LEFT", - "original_size": (900, 1600), - "egocar_visible": False -}, -4: { - "camera_name": "CAM_BACK_RIGHT", - "original_size": (900, 1600), - "egocar_visible": False -}, -5: { - "camera_name": "CAM_BACK", - "original_size": (900, 1600), - "egocar_visible": True -} -``` - - -## Waymo -``` -0: { - "camera_name": "front_camera", - "original_size": (1280, 1920), - "egocar_visible": False -}, -1: { - "camera_name": "front_left_camera", - "original_size": (1280, 1920), - "egocar_visible": False -}, -2: { - "camera_name": "front_right_camera", - "original_size": (1280, 1920), - "egocar_visible": False -}, -3: { - "camera_name": "left_camera", - "original_size": (866, 1920), - "egocar_visible": False -}, -4: { - "camera_name": "right_camera", - "original_size": (866, 1920), - "egocar_visible": False -}, -``` - -## Panda set - -``` -0: { - "camera_name": "front_camera", - "original_size": (1080, 1920), - "egocar_visible": False -}, -1: { - "camera_name": "front_left_camera", - "original_size": (1080, 1920), - "egocar_visible": False -}, -2: { - "camera_name": "front_right_camera", - "original_size": (1080, 1920), - "egocar_visible": False -}, -3: { - "camera_name": "left_camera", - "original_size": (1080, 1920), - "egocar_visible": False -}, -4: { - "camera_name": "right_camera", - "original_size": (1080, 1920), - "egocar_visible": False -}, -5: { - "camera_name": "back_camera", - "original_size": (1080, 1920), - "egocar_visible": True -}, -``` - - -## Argoverse - -``` -0: { - "camera_name": "ring_front_center", - "original_size": (2048, 1550), - "egocar_visible": True -}, -1: { - "camera_name": "ring_front_left", - "original_size": (1550, 2048), - "egocar_visible": False -}, -2: { - "camera_name": "ring_front_right", - "original_size": (1550, 2048), - "egocar_visible": False -}, -3: { - "camera_name": "ring_side_left", - "original_size": (1550, 2048), - "egocar_visible": False -}, -4: { - "camera_name": "ring_side_right", - "original_size": (1550, 2048), - "egocar_visible": False -}, -5: { - "camera_name": "ring_rear_left", - "original_size": (1550, 2048), - "egocar_visible": True -}, -6: { - "camera_name": "ring_rear_right", - "original_size": (1550, 2048), - "egocar_visible": True -}, -``` - - -## KITTI -``` -0: { - "camera_name": "CAM_LEFT", - "original_size": (375, 1242), - "egocar_visible": False -}, -1: { - "camera_name": "CAM_RIGHT", - "original_size": (375, 1242), - "egocar_visible": False -} -``` diff --git a/docs/datasets/carla.rst b/docs/datasets/carla.rst new file mode 100644 index 00000000..3ca7ff09 --- /dev/null +++ b/docs/datasets/carla.rst @@ -0,0 +1,90 @@ +CARLA +----- + +.. sidebar:: Dataset Name + + .. image:: https://carla.org/img/services/getty_center_400_400.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/index.rst b/docs/datasets/index.rst new file mode 100644 index 00000000..8bad761c --- /dev/null +++ b/docs/datasets/index.rst @@ -0,0 +1,16 @@ +Datasets +======== + +Brief overview of the datasets section... + +This section provides comprehensive documentation for various autonomous driving and computer vision datasets. Each dataset entry includes installation instructions, available data types, known issues, and proper citation formats. + +.. toctree:: + :maxdepth: 0 + + nuplan + nuscenes + carla + kitti-360 + wopd + template diff --git a/docs/datasets/kitti-360.rst b/docs/datasets/kitti-360.rst new file mode 100644 index 00000000..76100d27 --- /dev/null +++ b/docs/datasets/kitti-360.rst @@ -0,0 +1,87 @@ +KiTTI-360 +--------- + +.. sidebar:: Dataset Name + + .. image:: https://www.cvlibs.net/datasets/kitti-360/images/example/3d/semantic/02400.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use KiTTI-360 in your research, please cite: + +.. code-block:: bibtex + + @article{Liao2022PAMI, + title = {{KITTI}-360: A Novel Dataset and Benchmarks for Urban Scene Understanding in 2D and 3D}, + author = {Yiyi Liao and Jun Xie and Andreas Geiger}, + journal = {Pattern Analysis and Machine Intelligence (PAMI)}, + year = {2022}, + } diff --git a/docs/datasets/nuplan.rst b/docs/datasets/nuplan.rst new file mode 100644 index 00000000..065c26a5 --- /dev/null +++ b/docs/datasets/nuplan.rst @@ -0,0 +1,93 @@ +nuPlan +----- + +.. sidebar:: nuPlan + + .. image:: https://www.nuplan.org/static/media/nuPlan_final.3fde7586.png + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Towards learning-based planning:The nuPlan benchmark for real-world autonomous driving `_ + | **Download:** `www.nuscenes.org/nuplan `_ + | **Code:** `www.github.com/motional/nuplan-devkit `_ + | **Documentation:** `nuPlan Documentation `_ + | **License:** `CC BY-NC-SA 4.0 `_, `nuPlan Dataset License `_ + | **Duration:** 1282 hours (120 hours of sensor data) + | **Supported Versions:** [TODO] + | **Redistribution:** [TODO] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - ✓ + - [Description of ego vehicle data] + * - Map + - ✓ + - [Description of map data] + * - Bounding Boxes + - X + - [Description of bounding boxes data] + * - Traffic Lights + - X + - [Description of traffic lights data] + * - Cameras + - X + - [Description of cameras data] + * - LiDARs + - X + - [Description of LiDARs data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{Karnchanachari2024ICRA, + title={Towards learning-based planning: The nuplan benchmark for real-world autonomous driving}, + author={Karnchanachari, Napat and Geromichalos, Dimitris and Tan, Kok Seang and Li, Nanxiang and Eriksen, Christopher and Yaghoubi, Shakiba and Mehdipour, Noushin and Bernasconi, Gianmarco and Fong, Whye Kit and Guo, Yiluan and others}, + booktitle={2024 IEEE International Conference on Robotics and Automation (ICRA)}, + year={2024}, + } + @article{Caesar2021CVPRW, + title={nuplan: A closed-loop ml-based planning benchmark for autonomous vehicles}, + author={Caesar, Holger and Kabzan, Juraj and Tan, Kok Seang and Fong, Whye Kit and Wolff, Eric and Lang, Alex and Fletcher, Luke and Beijbom, Oscar and Omari, Sammy}, + booktitle={Proc. IEEE Conf. on Computer Vision and Pattern Recognition (CVPR) Workshops}, + year={2021} + } diff --git a/docs/datasets/nuscenes.rst b/docs/datasets/nuscenes.rst new file mode 100644 index 00000000..1f4e1621 --- /dev/null +++ b/docs/datasets/nuscenes.rst @@ -0,0 +1,90 @@ +nuScenes +-------- + +.. sidebar:: nuScenes + + .. image:: https://ar5iv.labs.arxiv.org/html/1903.11027/assets/figures/sensors.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/template.rst b/docs/datasets/template.rst new file mode 100644 index 00000000..d38723ed --- /dev/null +++ b/docs/datasets/template.rst @@ -0,0 +1,90 @@ +Template +-------- + +.. sidebar:: Dataset Name + + .. image:: https://www.nuplan.org/static/media/nuPlan_final.3fde7586.png + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/wopd.rst b/docs/datasets/wopd.rst new file mode 100644 index 00000000..8dc276a8 --- /dev/null +++ b/docs/datasets/wopd.rst @@ -0,0 +1,88 @@ +Waymo Open Perception Dataset (WOPD) +------------------------------------ + +.. sidebar:: WOPD + + .. image:: https://images.ctfassets.net/e6t5diu0txbw/4LpraC18sHNvS87OFnEGKB/63de105d4ce623d91cfdbc23f77d6a37/Open_Dataset_Download_Hero.jpg?fm=webp&q=90 + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @inproceedings{Sun2020CVPR, + title={Scalability in perception for autonomous driving: Waymo open dataset}, + author={Sun, Pei and Kretzschmar, Henrik and Dotiwalla, Xerxes and Chouard, Aurelien and Patnaik, Vijaysai and Tsui, Paul and Guo, James and Zhou, Yin and Chai, Yuning and Caine, Benjamin and others}, + booktitle={Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, + pages={2446--2454}, + year={2020} + } diff --git a/docs/index.rst b/docs/index.rst index 9f6c3518..b169f012 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -12,11 +12,11 @@ documentation for details. .. toctree:: - :maxdepth: 2 + :maxdepth: 1 :caption: Contents: installation - datasets + datasets/index schema visualization contributing From 00f01e51beb3fec4c14a081204583f7661dad7e4 Mon Sep 17 00:00:00 2001 From: DanielDauner Date: Sun, 17 Aug 2025 21:52:49 +0200 Subject: [PATCH 35/42] Add initial version of argoverse2 sensor dataset --- d123/common/datatypes/sensor/camera.py | 6 +- d123/common/datatypes/time/time_point.py | 10 + d123/common/utils/enums.py | 2 +- d123/common/visualization/viser/server.py | 4 +- .../dataset_specific/av2/av2_constants.py | 87 +++ .../av2/av2_data_converter.py | 528 ++++++++++++++++++ .../dataset_specific/av2/av2_helper.py | 179 ++++++ .../av2/av2_map_conversion.py | 0 .../default_dataset_conversion.yaml | 3 +- .../config/datasets/av2_sensor_dataset.yaml | 16 + docs/datasets/av2.rst | 90 +++ docs/datasets/index.rst | 1 + notebooks/av2/delete_me.ipynb | 480 ++++++++++++++++ notebooks/viz/bev_matplotlib.ipynb | 37 +- notebooks/viz/camera_matplotlib.ipynb | 157 ++++-- notebooks/viz/log_rendering.ipynb | 4 +- notebooks/viz/viser_testing_v2_scene.ipynb | 9 +- scripts/download/download_av2.sh | 6 +- 18 files changed, 1546 insertions(+), 73 deletions(-) create mode 100644 d123/dataset/dataset_specific/av2/av2_constants.py create mode 100644 d123/dataset/dataset_specific/av2/av2_data_converter.py create mode 100644 d123/dataset/dataset_specific/av2/av2_helper.py create mode 100644 d123/dataset/dataset_specific/av2/av2_map_conversion.py create mode 100644 d123/script/config/datasets/av2_sensor_dataset.yaml create mode 100644 docs/datasets/av2.rst create mode 100644 notebooks/av2/delete_me.ipynb diff --git a/d123/common/datatypes/sensor/camera.py b/d123/common/datatypes/sensor/camera.py index 900f0ab3..56fe6f07 100644 --- a/d123/common/datatypes/sensor/camera.py +++ b/d123/common/datatypes/sensor/camera.py @@ -23,6 +23,8 @@ class CameraType(SerialIntEnum): CAM_R0 = 5 CAM_R1 = 6 CAM_R2 = 7 + CAM_STEREO_L = 8 + CAM_STEREO_R = 9 @dataclass @@ -31,8 +33,8 @@ class CameraMetadata: camera_type: CameraType width: int height: int - intrinsic: npt.NDArray[np.float64] # 3x3 matrix - distortion: npt.NDArray[np.float64] # 5x1 vector + intrinsic: npt.NDArray[np.float64] # 3x3 matrix # TODO: don't store matrix but values. + distortion: npt.NDArray[np.float64] # 5x1 vector # TODO: don't store matrix but values. def to_dict(self) -> Dict[str, Any]: # TODO: remove None types. Only a placeholder for now. diff --git a/d123/common/datatypes/time/time_point.py b/d123/common/datatypes/time/time_point.py index 9460f78b..437bf967 100644 --- a/d123/common/datatypes/time/time_point.py +++ b/d123/common/datatypes/time/time_point.py @@ -206,6 +206,16 @@ def __post_init__(self) -> None: """ assert self.time_us >= 0, "Time point has to be positive!" + @classmethod + def from_ns(cls, t_ns: int) -> TimePoint: + """ + Constructs a TimePoint from a value in nanoseconds. + :param t_ns: Time in nanoseconds. + :return: TimePoint. + """ + assert isinstance(t_ns, int), "Nanoseconds must be an integer!" + return TimePoint(time_us=t_ns // 1000) + @classmethod def from_us(cls, t_us: int) -> TimePoint: """ diff --git a/d123/common/utils/enums.py b/d123/common/utils/enums.py index e09103e1..33300c00 100644 --- a/d123/common/utils/enums.py +++ b/d123/common/utils/enums.py @@ -18,7 +18,7 @@ def serialize(self, lower: bool = True) -> str: return self.name.lower() if lower else self.name @classmethod - def deserialize(cls, key: str) -> type[SerialIntEnum]: + def deserialize(cls, key: str) -> SerialIntEnum: """Deserialize the type when loading from a string.""" # Allow for lower/upper case letters during deserialize return cls.__members__[key.upper()] if key.islower() else cls.__members__[key] diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index 7db2351e..bf0acc26 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -33,10 +33,10 @@ LINE_WIDTH: float = 4.0 # Bounding box config: -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" # Map config: -MAP_AVAILABLE: bool = True +MAP_AVAILABLE: bool = False # Cameras config: diff --git a/d123/dataset/dataset_specific/av2/av2_constants.py b/d123/dataset/dataset_specific/av2/av2_constants.py new file mode 100644 index 00000000..65799aae --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_constants.py @@ -0,0 +1,87 @@ +from d123.common.datatypes.detection.detection_types import DetectionType +from d123.common.datatypes.sensor.camera import CameraType +from d123.common.utils.enums import SerialIntEnum + + +class AV2SensorBoxDetectionType(SerialIntEnum): + """Sensor dataset annotation categories.""" + + ANIMAL = 1 + ARTICULATED_BUS = 2 + BICYCLE = 3 + BICYCLIST = 4 + BOLLARD = 5 + BOX_TRUCK = 6 + BUS = 7 + CONSTRUCTION_BARREL = 8 + CONSTRUCTION_CONE = 9 + DOG = 10 + LARGE_VEHICLE = 11 + MESSAGE_BOARD_TRAILER = 12 + MOBILE_PEDESTRIAN_CROSSING_SIGN = 13 + MOTORCYCLE = 14 + MOTORCYCLIST = 15 + OFFICIAL_SIGNALER = 16 + PEDESTRIAN = 17 + RAILED_VEHICLE = 18 + REGULAR_VEHICLE = 19 + SCHOOL_BUS = 20 + SIGN = 21 + STOP_SIGN = 22 + STROLLER = 23 + TRAFFIC_LIGHT_TRAILER = 24 + TRUCK = 25 + TRUCK_CAB = 26 + VEHICULAR_TRAILER = 27 + WHEELCHAIR = 28 + WHEELED_DEVICE = 29 + WHEELED_RIDER = 30 + + +# Mapping from AV2SensorBoxDetectionType to general DetectionType +# TODO: Change the detection types. Multiple mistakes, e.g. animals/dogs are not generic objects. +AV2_TO_DETECTION_TYPE = { + AV2SensorBoxDetectionType.ANIMAL: DetectionType.GENERIC_OBJECT, + AV2SensorBoxDetectionType.ARTICULATED_BUS: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.BICYCLE: DetectionType.BICYCLE, + AV2SensorBoxDetectionType.BICYCLIST: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.BOLLARD: DetectionType.BARRIER, + AV2SensorBoxDetectionType.BOX_TRUCK: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.BUS: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.CONSTRUCTION_BARREL: DetectionType.BARRIER, + AV2SensorBoxDetectionType.CONSTRUCTION_CONE: DetectionType.TRAFFIC_CONE, + AV2SensorBoxDetectionType.DOG: DetectionType.GENERIC_OBJECT, + AV2SensorBoxDetectionType.LARGE_VEHICLE: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.MESSAGE_BOARD_TRAILER: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.MOBILE_PEDESTRIAN_CROSSING_SIGN: DetectionType.CZONE_SIGN, + AV2SensorBoxDetectionType.MOTORCYCLE: DetectionType.BICYCLE, + AV2SensorBoxDetectionType.MOTORCYCLIST: DetectionType.BICYCLE, + AV2SensorBoxDetectionType.OFFICIAL_SIGNALER: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.PEDESTRIAN: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.RAILED_VEHICLE: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.REGULAR_VEHICLE: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.SCHOOL_BUS: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.SIGN: DetectionType.SIGN, + AV2SensorBoxDetectionType.STOP_SIGN: DetectionType.SIGN, + AV2SensorBoxDetectionType.STROLLER: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.TRAFFIC_LIGHT_TRAILER: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.TRUCK: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.TRUCK_CAB: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.VEHICULAR_TRAILER: DetectionType.VEHICLE, + AV2SensorBoxDetectionType.WHEELCHAIR: DetectionType.PEDESTRIAN, + AV2SensorBoxDetectionType.WHEELED_DEVICE: DetectionType.GENERIC_OBJECT, + AV2SensorBoxDetectionType.WHEELED_RIDER: DetectionType.BICYCLE, +} + + +AV2_CAMERA_TYPE_MAPPING = { + "ring_front_center": CameraType.CAM_F0, + "ring_front_left": CameraType.CAM_L0, + "ring_front_right": CameraType.CAM_R0, + "ring_side_left": CameraType.CAM_L1, + "ring_side_right": CameraType.CAM_R1, + "ring_rear_left": CameraType.CAM_L2, + "ring_rear_right": CameraType.CAM_R2, + "stereo_front_left": CameraType.CAM_STEREO_L, + "stereo_front_right": CameraType.CAM_STEREO_R, +} diff --git a/d123/dataset/dataset_specific/av2/av2_data_converter.py b/d123/dataset/dataset_specific/av2/av2_data_converter.py new file mode 100644 index 00000000..36dc143e --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_data_converter.py @@ -0,0 +1,528 @@ +import gc +import hashlib +import json +from dataclasses import asdict +from functools import partial +from pathlib import Path +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import pandas as pd +import pyarrow as pa +from pyquaternion import Quaternion + +from d123.common.datatypes.sensor.camera import CameraMetadata, CameraType, camera_metadata_dict_to_json +from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json +from d123.common.datatypes.time.time_point import TimePoint +from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index +from d123.common.datatypes.vehicle_state.vehicle_parameters import ( + get_nuplan_pacifica_parameters, + rear_axle_se3_to_center_se3, +) +from d123.common.geometry.base import StateSE3 +from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index +from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array +from d123.common.geometry.vector import Vector3D, Vector3DIndex +from d123.common.multithreading.worker_utils import WorkerPool, worker_map +from d123.dataset.dataset_specific.av2.av2_constants import ( + AV2_CAMERA_TYPE_MAPPING, + AV2_TO_DETECTION_TYPE, + AV2SensorBoxDetectionType, +) +from d123.dataset.dataset_specific.av2.av2_helper import ( + build_sensor_dataframe, + build_synchronization_dataframe, + find_closest_target_fpath, + get_slice_with_timestamp_ns, +) +from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import MAP_LOCATIONS +from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter +from d123.dataset.logs.log_metadata import LogMetadata + + +def create_token(input_data: str) -> str: + # TODO: Refactor this function. + # TODO: Add a general function to create tokens from arbitrary data. + if isinstance(input_data, str): + input_data = input_data.encode("utf-8") + + hash_obj = hashlib.sha256(input_data) + return hash_obj.hexdigest()[:16] + + +class AV2SensorDataConverter(RawDataConverter): + def __init__( + self, + splits: List[str], + log_path: Union[Path, str], + data_converter_config: DataConverterConfig, + ) -> None: + super().__init__(data_converter_config) + for split in splits: + assert ( + split in self.get_available_splits() + ), f"Split {split} is not available. Available splits: {self.available_splits}" + + self._splits: List[str] = splits + self._data_root: Path = Path(log_path) + self._log_paths_per_split: Dict[str, List[Path]] = self._collect_log_paths() + self._target_dt: float = 0.1 + + def _collect_log_paths(self) -> Dict[str, List[Path]]: + log_paths_per_split: Dict[str, List[Path]] = {} + + for split in self._splits: + subsplit = split.split("_")[-1] + assert subsplit in ["train", "val", "test"] + + if "av2_sensor" in split: + log_folder = self._data_root / "sensor" / subsplit + elif "av2_lidar" in split: + log_folder = self._data_root / "lidar" / subsplit + elif "av2_motion" in split: + log_folder = self._data_root / "motion-forecasting" / subsplit + elif "av2-sensor-mini" in split: + log_folder = self._data_root / "sensor_mini" / subsplit + + log_paths_per_split[split] = list(log_folder.iterdir()) + + return log_paths_per_split + + def get_available_splits(self) -> List[str]: + return [ + "av2-sensor_train", + "av2-sensor_val", + "av2-sensor_test", + "av2-sensor-mini_train", + "av2-sensor-mini_val", + "av2-sensor-mini_test", + ] + + def convert_maps(self, worker: WorkerPool) -> None: + worker_map( + worker, + partial(convert_av2_map_to_gpkg, data_converter_config=self.data_converter_config), + list(MAP_LOCATIONS), + ) + + def convert_logs(self, worker: WorkerPool) -> None: + log_args = [ + { + "log_path": log_path, + "split": split, + } + for split, log_paths in self._log_paths_per_split.items() + for log_path in log_paths + ] + + worker_map( + worker, + partial( + convert_av2_log_to_arrow, + data_converter_config=self.data_converter_config, + ), + log_args, + ) + + +def convert_av2_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: + # TODO: Implement map + # for map_name in map_names: + # map_path = data_converter_config.output_path / "maps" / f"av2_{map_name}.gpkg" + # if data_converter_config.force_map_conversion or not map_path.exists(): + # map_path.unlink(missing_ok=True) + # AV2MapConverter(data_converter_config.output_path / "maps").convert(map_name=map_name) + return [] + + +def convert_av2_log_to_arrow( + args: List[Dict[str, Union[List[str], List[Path]]]], + data_converter_config: DataConverterConfig, +) -> List[Any]: + for log_info in args: + log_path: Path = log_info["log_path"] + split: str = log_info["split"] + + if not log_path.exists(): + raise FileNotFoundError(f"Log path {log_path} does not exist.") + + log_file_path = data_converter_config.output_path / split / f"{log_path.stem}.arrow" + + if data_converter_config.force_log_conversion or not log_file_path.exists(): + log_file_path.unlink(missing_ok=True) + if not log_file_path.parent.exists(): + log_file_path.parent.mkdir(parents=True, exist_ok=True) + + sensor_df = build_sensor_dataframe(log_path) + synchronization_df = build_synchronization_dataframe(sensor_df) + + metadata = LogMetadata( + dataset="av2-sensor", + log_name=log_path.name, + location=None, + timestep_seconds=0.1, # TODO: verify this + map_has_z=True, + ) + vehicle_parameters = get_nuplan_pacifica_parameters() # TODO: Add av2 vehicle parameters + camera_metadata = get_av2_camera_metadata(log_path) + lidar_metadata = get_av2_lidar_metadata(log_path) + + schema_column_list = [ + ("token", pa.string()), + ("timestamp", pa.int64()), + ("detections_state", pa.list_(pa.list_(pa.float64(), len(BoundingBoxSE3Index)))), + ("detections_velocity", pa.list_(pa.list_(pa.float64(), len(Vector3DIndex)))), + ("detections_token", pa.list_(pa.string())), + ("detections_type", pa.list_(pa.int16())), + ("ego_states", pa.list_(pa.float64(), len(EgoStateSE3Index))), + ("traffic_light_ids", pa.list_(pa.int64())), + ("traffic_light_types", pa.list_(pa.int16())), + ("scenario_tag", pa.list_(pa.string())), + ("route_lane_group_ids", pa.list_(pa.int64())), + ] + if data_converter_config.lidar_store_option is not None: + for lidar_type in lidar_metadata.keys(): + if data_converter_config.lidar_store_option == "path": + schema_column_list.append((lidar_type.serialize(), pa.string())) + elif data_converter_config.lidar_store_option == "binary": + raise NotImplementedError("Binary lidar storage is not implemented.") + + if data_converter_config.camera_store_option is not None: + for camera_type in camera_metadata.keys(): + if data_converter_config.camera_store_option == "path": + schema_column_list.append((camera_type.serialize(), pa.string())) + + elif data_converter_config.camera_store_option == "binary": + schema_column_list.append((camera_type.serialize(), pa.binary())) + + schema_column_list.append((f"{camera_type.serialize()}_extrinsic", pa.list_(pa.float64(), 4 * 4))) + + recording_schema = pa.schema(schema_column_list) + recording_schema = recording_schema.with_metadata( + { + "log_metadata": json.dumps(asdict(metadata)), + "vehicle_parameters": json.dumps(asdict(vehicle_parameters)), + "camera_metadata": camera_metadata_dict_to_json(camera_metadata), + "lidar_metadata": lidar_metadata_dict_to_json(lidar_metadata), + } + ) + + _write_recording_table( + sensor_df, + synchronization_df, + recording_schema, + log_file_path, + log_path, + data_converter_config, + ) + del recording_schema, vehicle_parameters + gc.collect() + return [] + + +def get_av2_camera_metadata(log_path: Path) -> Dict[CameraType, CameraMetadata]: + + intrinsics_file = log_path / "calibration" / "intrinsics.feather" + intrinsics_df = pd.read_feather(intrinsics_file) + + camera_metadata: Dict[CameraType, CameraMetadata] = {} + for _, row in intrinsics_df.iterrows(): + row = row.to_dict() + + camera_type = AV2_CAMERA_TYPE_MAPPING[row["sensor_name"]] + camera_metadata[camera_type] = CameraMetadata( + camera_type=camera_type, + width=row["width_px"], + height=row["height_px"], + intrinsic=np.array( + [[row["fx_px"], 0, row["cx_px"]], [0, row["fy_px"], row["cy_px"]], [0, 0, 1]], dtype=np.float64 + ), + distortion=np.array([row["k1"], row["k2"], row["k3"], 0, 0], dtype=np.float64), + ) + + return camera_metadata + + +def get_av2_lidar_metadata(log_path: Path) -> Dict[LiDARType, LiDARMetadata]: + # metadata: Dict[LiDARType, LiDARMetadata] = {} + # metadata[LiDARType.LIDAR_MERGED] = LiDARMetadata( + # lidar_type=LiDARType.LIDAR_MERGED, + # lidar_index=NuplanLidarIndex, + # extrinsic=None, # NOTE: LiDAR extrinsic are unknown + # ) + # return metadata + return {} + + +def _write_recording_table( + sensor_df: pd.DataFrame, + synchronization_df: pd.DataFrame, + recording_schema: pa.schema, + log_file_path: Path, + source_log_path: Path, + data_converter_config: DataConverterConfig, +) -> None: + + # NOTE: Similar to other datasets, we use the lidar timestamps as reference timestamps. + lidar_sensor = sensor_df.xs(key="lidar", level=2) + lidar_timestamps_ns = np.sort([int(idx_tuple[2]) for idx_tuple in lidar_sensor.index]) + + # NOTE: The annotation dataframe is not available for the test split. + annotations_df = ( + pd.read_feather(source_log_path / "annotations.feather") + if (source_log_path / "annotations.feather").exists() + else None + ) + + city_se3_egovehicle_df = pd.read_feather(source_log_path / "city_SE3_egovehicle.feather") + + egovehicle_se3_sensor_df = ( + pd.read_feather(source_log_path / "calibration" / "egovehicle_SE3_sensor.feather") + if data_converter_config.camera_store_option is not None + else None + ) + + # with pa.ipc.new_stream(str(log_file_path), recording_schema) as writer: + with pa.OSFile(str(log_file_path), "wb") as sink: + with pa.ipc.new_file(sink, recording_schema) as writer: + + for lidar_timestamp_ns in lidar_timestamps_ns: + + ego_state_se3 = _extract_ego_state(city_se3_egovehicle_df, lidar_timestamp_ns) + ( + detections_state, + detections_velocity, + detections_token, + detections_types, + ) = _extract_box_detections(annotations_df, lidar_timestamp_ns, ego_state_se3) + traffic_light_ids, traffic_light_types = _extract_traffic_lights() + route_lane_group_ids = None # TODO: Add route lane group ids extraction ? + row_data = { + "token": [create_token(str(lidar_timestamp_ns))], + "timestamp": [TimePoint.from_ns(int(lidar_timestamp_ns)).time_us], + "detections_state": [detections_state], + "detections_velocity": [detections_velocity], + "detections_token": [detections_token], + "detections_type": [detections_types], + "ego_states": [ego_state_se3.array.tolist()], + "traffic_light_ids": [traffic_light_ids], + "traffic_light_types": [traffic_light_types], + "scenario_tag": [_extract_scenario_tag()], + "route_lane_group_ids": [route_lane_group_ids], + } + + # TODO: add lidar data + + # if data_converter_config.lidar_store_option is not None: + # lidar_data_dict = _extract_lidar(lidar_pc, data_converter_config) + # for lidar_type, lidar_data in lidar_data_dict.items(): + # if lidar_data is not None: + # row_data[lidar_type.serialize()] = [lidar_data] + # else: + # row_data[lidar_type.serialize()] = [None] + + if data_converter_config.camera_store_option is not None: + camera_data_dict = _extract_camera( + lidar_timestamp_ns, + city_se3_egovehicle_df, + egovehicle_se3_sensor_df, + synchronization_df, + source_log_path, + data_converter_config, + ) + for camera_type, camera_data in camera_data_dict.items(): + if camera_data is not None: + row_data[camera_type.serialize()] = [camera_data[0]] + row_data[f"{camera_type.serialize()}_extrinsic"] = [camera_data[1]] + else: + row_data[camera_type.serialize()] = [None] + row_data[f"{camera_type.serialize()}_extrinsic"] = [None] + + batch = pa.record_batch(row_data, schema=recording_schema) + writer.write_batch(batch) + del batch, row_data, detections_state, detections_velocity, detections_token, detections_types + + +def _extract_box_detections( + annotations_df: Optional[pd.DataFrame], + lidar_timestamp_ns: int, + ego_state_se3: EgoStateSE3, +) -> Tuple[List[List[float]], List[List[float]], List[str], List[int]]: + + # TODO: Extract velocity from annotations_df if available. + + if annotations_df is None: + return [], [], [], [] + + annotations_slice = get_slice_with_timestamp_ns(annotations_df, lidar_timestamp_ns) + num_detections = len(annotations_slice) + + detections_state = np.zeros((num_detections, len(BoundingBoxSE3Index)), dtype=np.float64) + detections_velocity = np.zeros((num_detections, len(Vector3DIndex)), dtype=np.float64) + detections_token: List[str] = annotations_slice["track_uuid"].tolist() + detections_types: List[int] = [] + + for detection_idx, (_, row) in enumerate(annotations_slice.iterrows()): + row = row.to_dict() + yaw, pitch, roll = Quaternion( + w=row["qw"], + x=row["qx"], + y=row["qy"], + z=row["qz"], + ).yaw_pitch_roll + + detections_state[detection_idx, BoundingBoxSE3Index.X] = row["tx_m"] + detections_state[detection_idx, BoundingBoxSE3Index.Y] = row["ty_m"] + detections_state[detection_idx, BoundingBoxSE3Index.Z] = row["tz_m"] + detections_state[detection_idx, BoundingBoxSE3Index.ROLL] = roll + detections_state[detection_idx, BoundingBoxSE3Index.PITCH] = pitch + detections_state[detection_idx, BoundingBoxSE3Index.YAW] = yaw + detections_state[detection_idx, BoundingBoxSE3Index.LENGTH] = row["length_m"] + detections_state[detection_idx, BoundingBoxSE3Index.WIDTH] = row["width_m"] + detections_state[detection_idx, BoundingBoxSE3Index.HEIGHT] = row["height_m"] + + av2_detection_type = AV2SensorBoxDetectionType.deserialize(row["category"]) + detections_types.append(int(AV2_TO_DETECTION_TYPE[av2_detection_type])) + + detections_state[:, BoundingBoxSE3Index.STATE_SE3] = convert_relative_to_absolute_se3_array( + origin=ego_state_se3.rear_axle_se3, se3_array=detections_state[:, BoundingBoxSE3Index.STATE_SE3] + ) + + ZERO_BOX_ROLL_PITCH = False # TODO: Add config option or remove + if ZERO_BOX_ROLL_PITCH: + detections_state[:, BoundingBoxSE3Index.ROLL] = DEFAULT_ROLL + detections_state[:, BoundingBoxSE3Index.PITCH] = DEFAULT_PITCH + + return detections_state.tolist(), detections_velocity.tolist(), detections_token, detections_types + + +def _extract_ego_state(city_se3_egovehicle_df: pd.DataFrame, lidar_timestamp_ns: int) -> EgoStateSE3: + ego_state_slice = get_slice_with_timestamp_ns(city_se3_egovehicle_df, lidar_timestamp_ns) + assert ( + len(ego_state_slice) == 1 + ), f"Expected exactly one ego state for timestamp {lidar_timestamp_ns}, got {len(ego_state_slice)}." + + ego_pose_dict = ego_state_slice.iloc[0].to_dict() + + ego_pose_quat = Quaternion( + w=ego_pose_dict["qw"], + x=ego_pose_dict["qx"], + y=ego_pose_dict["qy"], + z=ego_pose_dict["qz"], + ) + + yaw, pitch, roll = ego_pose_quat.yaw_pitch_roll + + rear_axle_pose = StateSE3( + x=ego_pose_dict["tx_m"], + y=ego_pose_dict["ty_m"], + z=ego_pose_dict["tz_m"], + roll=roll, + pitch=pitch, + yaw=yaw, + ) + vehicle_parameters = get_nuplan_pacifica_parameters() # TODO: Add av2 vehicle parameters + center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) + # TODO: Add script to calculate the dynamic state from log sequence. + dynamic_state = DynamicStateSE3( + velocity=Vector3D( + x=0.0, + y=0.0, + z=0.0, + ), + acceleration=Vector3D( + x=0.0, + y=0.0, + z=0.0, + ), + angular_velocity=Vector3D( + x=0.0, + y=0.0, + z=0.0, + ), + ) + + return EgoStateSE3( + center_se3=center, + dynamic_state_se3=dynamic_state, + vehicle_parameters=vehicle_parameters, + timepoint=None, + ) + + +def _extract_traffic_lights() -> Tuple[List[int], List[int]]: + return [], [] + + +def _extract_scenario_tag() -> List[str]: + return ["unknown"] + + +def _extract_camera( + lidar_timestamp_ns: int, + city_se3_egovehicle_df: pd.DataFrame, + egovehicle_se3_sensor_df: pd.DataFrame, + synchronization_df: pd.DataFrame, + source_log_path: Path, + data_converter_config: DataConverterConfig, +) -> Dict[CameraType, Union[str, bytes]]: + + camera_dict: Dict[CameraType, Union[str, bytes]] = { + camera_type: None for camera_type in AV2_CAMERA_TYPE_MAPPING.values() + } + split = source_log_path.parent.name + log_id = source_log_path.name + + source_dataset_dir = source_log_path.parent.parent + + for _, row in egovehicle_se3_sensor_df.iterrows(): + row = row.to_dict() + if row["sensor_name"] not in AV2_CAMERA_TYPE_MAPPING: + continue + + camera_name = row["sensor_name"] + camera_type = AV2_CAMERA_TYPE_MAPPING[camera_name] + + relative_image_path = find_closest_target_fpath( + split=split, + log_id=log_id, + src_sensor_name="lidar", + src_timestamp_ns=lidar_timestamp_ns, + target_sensor_name=camera_name, + synchronization_df=synchronization_df, + ) + if relative_image_path is None: + camera_dict[camera_type] = None + else: + absolute_image_path = source_dataset_dir / relative_image_path + assert absolute_image_path.exists() + # TODO: Adjust for finer IMU timestamps to correct the camera extrinsic. + camera_extrinsic = np.zeros((4, 4), dtype=np.float64) + camera_extrinsic[:3, :3] = Quaternion( + w=row["qw"], + x=row["qx"], + y=row["qy"], + z=row["qz"], + ).rotation_matrix + camera_extrinsic[:3, 3] = np.array([row["tx_m"], row["ty_m"], row["tz_m"]], dtype=np.float64) + camera_extrinsic = camera_extrinsic.flatten().tolist() + + if data_converter_config.camera_store_option == "path": + camera_dict[camera_type] = (str(relative_image_path), camera_extrinsic) + elif data_converter_config.camera_store_option == "binary": + with open(absolute_image_path, "rb") as f: + camera_dict[camera_type] = (f.read(), camera_extrinsic) + + return camera_dict + + +def _extract_lidar(lidar_pc, data_converter_config: DataConverterConfig) -> Dict[LiDARType, Optional[str]]: + + # lidar: Optional[str] = None + # lidar_full_path = NUPLAN_DATA_ROOT / "nuplan-v1.1" / "sensor_blobs" / lidar_pc.filename + # if lidar_full_path.exists(): + # lidar = lidar_pc.filename + + # return {LiDARType.LIDAR_MERGED: lidar} + return {} diff --git a/d123/dataset/dataset_specific/av2/av2_helper.py b/d123/dataset/dataset_specific/av2/av2_helper.py new file mode 100644 index 00000000..8ad881cc --- /dev/null +++ b/d123/dataset/dataset_specific/av2/av2_helper.py @@ -0,0 +1,179 @@ +from pathlib import Path +from typing import Final, List, Literal, Optional + +import pandas as pd + +from d123.dataset.dataset_specific.av2.av2_constants import AV2_CAMERA_TYPE_MAPPING + +AV2_SENSOR_CAM_SHUTTER_INTERVAL_MS: Final[float] = 50.0 +AV2_SENSOR_LIDAR_SWEEP_INTERVAL_W_BUFFER_NS: Final[float] = 102000000.0 + + +AV2_SENSOR_CAMERA_NAMES = [ + "ring_front_center", + "ring_front_left", + "ring_front_right", + "ring_rear_left", + "ring_rear_right", + "ring_side_left", + "ring_side_right", + "stereo_front_left", + "stereo_front_right", +] + + +def get_dataframe_from_file(file_path: Path) -> pd.DataFrame: + if file_path.suffix == ".parquet": + import pyarrow.parquet as pq + + return pq.read_table(file_path) + elif file_path.suffix == ".feather": + import pyarrow.feather as feather + + return feather.read_feather(file_path) + else: + raise ValueError(f"Unsupported file type: {file_path.suffix}") + + +def get_slice_with_timestamp_ns(dataframe: pd.DataFrame, timestamp_ns: int): + """Get the index of the closest timestamp to the target timestamp.""" + return dataframe[dataframe["timestamp_ns"] == timestamp_ns] + + +def build_sensor_dataframe(raw_log_path: Path) -> pd.DataFrame: + + # https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/sensor_dataloader.py#L209 + + split = raw_log_path.parent.name + log_id = raw_log_path.name + + lidar_path = raw_log_path / "sensors" / "lidar" + cameras_path = raw_log_path / "sensors" / "cameras" + + # Find all the lidar records and timestamps from file names. + lidar_records = populate_sensor_records(lidar_path, split, log_id) + + # Find all the camera records and timestamps from file names. + camera_records = [] + for camera_folder in cameras_path.iterdir(): + assert camera_folder.name in AV2_SENSOR_CAMERA_NAMES + camera_record = populate_sensor_records(camera_folder, split, log_id) + camera_records.append(camera_record) + + # Concatenate all the camera records into a single DataFrame. + sensor_records = [lidar_records] + camera_records + sensor_dataframe = pd.concat(sensor_records) + + # Set index as tuples of the form: (split, log_id, sensor_name, timestamp_ns) and sort the index. + # sorts by split, log_id, and then by sensor name, and then by timestamp. + sensor_dataframe.set_index(["split", "log_id", "sensor_name", "timestamp_ns"], inplace=True) + sensor_dataframe.sort_index(inplace=True) + + return sensor_dataframe + + +def build_synchronization_dataframe( + sensor_dataframe: pd.DataFrame, + matching_criterion: Literal["nearest", "forward"] = "nearest", +) -> pd.DataFrame: + + # https://github.com/argoverse/av2-api/blob/main/src/av2/datasets/sensor/sensor_dataloader.py#L382 + + # Create list to store synchronized data frames. + sync_list: List[pd.DataFrame] = [] + unique_sensor_names: List[str] = sensor_dataframe.index.unique(level=2).tolist() + + # Associate a 'source' sensor to a 'target' sensor for all available sensors. + # For example, we associate the lidar sensor with each ring camera which + # produces a mapping from lidar -> all-other-sensors. + for src_sensor_name in unique_sensor_names: + src_records = sensor_dataframe.xs(src_sensor_name, level=2, drop_level=False).reset_index() + src_records = src_records.rename({"timestamp_ns": src_sensor_name}, axis=1).sort_values(src_sensor_name) + + # _Very_ important to convert to timedelta. Tolerance below causes precision loss otherwise. + src_records[src_sensor_name] = pd.to_timedelta(src_records[src_sensor_name]) + for target_sensor_name in unique_sensor_names: + if src_sensor_name == target_sensor_name: + continue + target_records = sensor_dataframe.xs(target_sensor_name, level=2).reset_index() + target_records = target_records.rename({"timestamp_ns": target_sensor_name}, axis=1).sort_values( + target_sensor_name + ) + + # Merge based on matching criterion. + # _Very_ important to convert to timedelta. Tolerance below causes precision loss otherwise. + target_records[target_sensor_name] = pd.to_timedelta(target_records[target_sensor_name]) + tolerance = pd.to_timedelta(AV2_SENSOR_CAM_SHUTTER_INTERVAL_MS / 2 * 1e6) + if "ring" in src_sensor_name: + tolerance = pd.to_timedelta(AV2_SENSOR_LIDAR_SWEEP_INTERVAL_W_BUFFER_NS / 2) + src_records = pd.merge_asof( + src_records, + target_records, + left_on=src_sensor_name, + right_on=target_sensor_name, + by=["split", "log_id"], + direction=matching_criterion, + tolerance=tolerance, + ) + sync_list.append(src_records) + + sync_records = pd.concat(sync_list).reset_index(drop=True) + sync_records.set_index(keys=["split", "log_id", "sensor_name"], inplace=True) + sync_records.sort_index(inplace=True) + + return sync_records + + +def populate_sensor_records(sensor_path: Path, split: str, log_id: str) -> pd.DataFrame: + + sensor_name = sensor_path.name + sensor_files = list(sensor_path.iterdir()) + sensor_records = [] + + for sensor_file in sensor_files: + assert sensor_file.suffix in [ + ".feather", + ".jpg", + ], f"Unsupported file type: {sensor_file.suffix} for {str(sensor_file)}" + row = {} + row["split"] = split + row["log_id"] = log_id + row["sensor_name"] = sensor_name + row["timestamp_ns"] = int(sensor_file.stem) + sensor_records.append(row) + + return pd.DataFrame(sensor_records) + + +def find_closest_target_fpath( + split: str, + log_id: str, + src_sensor_name: str, + src_timestamp_ns: int, + target_sensor_name: str, + synchronization_df: pd.DataFrame, +) -> Optional[Path]: + """Find the file path to the target sensor from a source sensor.""" + # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/datasets/sensor/sensor_dataloader.py#L448 + + src_timedelta_ns = pd.Timedelta(src_timestamp_ns) + src_to_target_records = synchronization_df.loc[(split, log_id, src_sensor_name)].set_index(src_sensor_name) + index = src_to_target_records.index + if src_timedelta_ns not in index: + # This timestamp does not correspond to any lidar sweep. + return None + + # Grab the synchronization record. + target_timestamp_ns = src_to_target_records.loc[src_timedelta_ns, target_sensor_name] + if pd.isna(target_timestamp_ns): + # No match was found within tolerance. + return None + + sensor_dir = Path(split) / log_id / "sensors" + valid_cameras = list(AV2_CAMERA_TYPE_MAPPING.keys()) + timestamp_ns_str = str(target_timestamp_ns.asm8.item()) + if target_sensor_name in valid_cameras: + target_path = sensor_dir / "cameras" / target_sensor_name / f"{timestamp_ns_str}.jpg" + else: + target_path = sensor_dir / target_sensor_name / f"{timestamp_ns_str}.feather" + return target_path diff --git a/d123/dataset/dataset_specific/av2/av2_map_conversion.py b/d123/dataset/dataset_specific/av2/av2_map_conversion.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 6357115c..7287f850 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -15,9 +15,10 @@ defaults: - default_dataset_paths - _self_ - datasets: - - nuplan_private_dataset + # - nuplan_private_dataset # - carla_dataset # - wopd_dataset + - av2_sensor_dataset force_log_conversion: True force_map_conversion: False diff --git a/d123/script/config/datasets/av2_sensor_dataset.yaml b/d123/script/config/datasets/av2_sensor_dataset.yaml new file mode 100644 index 00000000..d5b4643c --- /dev/null +++ b/d123/script/config/datasets/av2_sensor_dataset.yaml @@ -0,0 +1,16 @@ +av2_sensor_dataset: + _target_: d123.dataset.dataset_specific.av2.av2_data_converter.AV2SensorDataConverter + _convert_: 'all' + + splits: ["av2-sensor-mini_train"] + log_path: "/mnt/elements_0/argoverse" + + data_converter_config: + _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig + _convert_: 'all' + + output_path: ${d123_data_root} + force_log_conversion: ${force_log_conversion} + force_map_conversion: ${force_map_conversion} + camera_store_option: binary + lidar_store_option: null diff --git a/docs/datasets/av2.rst b/docs/datasets/av2.rst new file mode 100644 index 00000000..ab01edfb --- /dev/null +++ b/docs/datasets/av2.rst @@ -0,0 +1,90 @@ +Argoverse 2 +----------- + +.. sidebar:: Dataset Name + + .. image:: https://www.argoverse.org/assets/images/reference_images/av2_vehicle.jpg + :alt: Dataset sample image + :width: 290px + + | **Paper:** `Name of Paper `_ + | **Download:** `Documentation `_ + | **Code:** [Code] + | **Documentation:** [License type] + | **License:** [License type] + | **Duration:** [Duration here] + | **Supported Versions:** [Yes/No/Conditions] + | **Redistribution:** [Yes/No/Conditions] + +Description +~~~~~~~~~~~ + +[Provide a detailed description of the dataset here, including its purpose, collection methodology, and key characteristics.] + +Installation +~~~~~~~~~~~~ + +[Instructions for installing or accessing the dataset] + +.. code-block:: bash + + # Example installation commands + pip install d123[dataset_name] + # or + wget https://example.com/dataset.zip + +Available Data +~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + :widths: 30 5 70 + + + * - **Name** + - **Available** + - **Description** + * - Ego Vehicle + - X + - [Description of ego vehicle data] + * - Map + - X + - [Description of ego vehicle data] + * - Bounding Boxes + - X + - [Description of ego vehicle data] + * - Traffic Lights + - X + - [Description of ego vehicle data] + * - Cameras + - X + - [Description of ego vehicle data] + * - LiDARs + - X + - [Description of ego vehicle data] + +Dataset Specific Issues +~~~~~~~~~~~~~~~~~~~~~~~ + +[Document any known issues, limitations, or considerations when using this dataset] + +* Issue 1: Description +* Issue 2: Description +* Issue 3: Description + +Citation +~~~~~~~~ + +If you use this dataset in your research, please cite: + +.. code-block:: bibtex + + @article{AuthorYearConference, + title={Dataset Title}, + author={Author, First and Author, Second}, + journal={Journal Name}, + year={2023}, + volume={1}, + pages={1-10}, + doi={10.1000/example} + } diff --git a/docs/datasets/index.rst b/docs/datasets/index.rst index 8bad761c..ead97169 100644 --- a/docs/datasets/index.rst +++ b/docs/datasets/index.rst @@ -8,6 +8,7 @@ This section provides comprehensive documentation for various autonomous driving .. toctree:: :maxdepth: 0 + av2 nuplan nuscenes carla diff --git a/notebooks/av2/delete_me.ipynb b/notebooks/av2/delete_me.ipynb new file mode 100644 index 00000000..d66c98c8 --- /dev/null +++ b/notebooks/av2/delete_me.ipynb @@ -0,0 +1,480 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import io\n", + "\n", + "from PIL import Image\n", + "\n", + "import pandas as pd" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"test\"\n", + "\n", + "\n", + "split = \"train\"\n", + "# split = \"val\"\n", + "\n", + "split_folder = Path(f\"/mnt/elements_0/argoverse/sensor_mini/{split}\")\n", + "log_folder = sorted(list(split_folder.iterdir()))[0]\n", + "\n", + "log_folder.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "def _ls(path: Path):\n", + " \"\"\"List all files in the given path.\"\"\"\n", + " return [f.name for f in path.iterdir()]\n", + "\n", + "\n", + "def get_arrow_from_file(file_path: Path):\n", + " if file_path.suffix == \".parquet\":\n", + " import pyarrow.parquet as pq\n", + " return pq.read_table(file_path)\n", + " elif file_path.suffix == \".feather\":\n", + " import pyarrow.feather as feather\n", + " return feather.read_feather(file_path)\n", + " else:\n", + " raise ValueError(f\"Unsupported file type: {file_path.suffix}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "\n", + "# 1. Calibration \n", + "calibration_folder = log_folder / \"calibration\"\n", + "\n", + "# 1.1 -> ego to sensor transformation\n", + "egovehicle_SE3_sensor_file = log_folder / \"calibration\" / \"egovehicle_SE3_sensor.feather\"\n", + "egovehicle_se3_sensor_table = get_arrow_from_file(egovehicle_SE3_sensor_file)\n", + "\n", + "egovehicle_se3_sensor_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "# 1.2 -> intrinsic parameters\n", + "intrinsics_file = log_folder / \"calibration\" / \"intrinsics.feather\"\n", + "intrinsics_table = get_arrow_from_file(intrinsics_file)\n", + "\n", + "\n", + "intrinsics_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "# 2. Ego Vehicle\n", + "city_SE3_egovehicle_file = log_folder / \"city_SE3_egovehicle.feather\"\n", + "city_se3_egovehicle_table = get_arrow_from_file(city_SE3_egovehicle_file)\n", + "\n", + "city_se3_egovehicle_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# # 3. Map\n", + "# # map_folder = log_folder / \"map\"\n", + "# # print(_ls(map_folder))\n", + "\n", + "# # # 4. sensors\n", + "# # print(_ls(log_folder))\n", + "\n", + "# # from d123.dataset.dataset_specific.av2.av2_data_converter import AV2SensorDataConverter\n", + "# from d123.dataset.dataset_specific.av2.av2_data_converter import AV2SensorDataConverter\n", + "\n", + "# # AV2SensorDataConverter([])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "# 5. Annotations\n", + "annotations_file = log_folder / \"annotations.feather\"\n", + "annotations_table = get_arrow_from_file(annotations_file)\n", + "\n", + "# print(_ls(annotations_folder))\n", + "annotations_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "camera_name = \"ring_side_left\"\n", + "\n", + "camera_folder = log_folder / \"sensors\" / \"cameras\" / camera_name\n", + "camera_files = sorted(list(camera_folder.iterdir()))\n", + "\n", + "\n", + "def jpg_to_array(file_path):\n", + " image = np.array(Image.open(io.BytesIO(file_path.read_bytes())))\n", + " return image\n", + "\n", + "plt.imshow(jpg_to_array(camera_files[1]))\n", + "print(len(camera_files))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "lidar_folder = log_folder / \"sensors\" / \"lidar\" \n", + "lidar_files = sorted(list(lidar_folder.iterdir()))\n", + "\n", + "\n", + "get_arrow_from_file(lidar_files[0])\n", + "\n", + "print(lidar_files)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "# Testing sensor syn dataframes\n", + "\n", + "from typing import Optional\n", + "from d123.dataset.dataset_specific.av2.av2_constants import AV2_CAMERA_TYPE_MAPPING\n", + "from d123.dataset.dataset_specific.av2.av2_helper import build_sensor_dataframe, build_synchronization_dataframe\n", + "\n", + "\n", + "sensor_df = build_sensor_dataframe(log_folder)\n", + "synchronization_df = build_synchronization_dataframe(sensor_df)\n", + "dataset_dir = split_folder.parent\n", + "\n", + "\n", + "def find_closest_target_fpath(\n", + " split: str,\n", + " log_id: str,\n", + " src_sensor_name: str,\n", + " src_timestamp_ns: int,\n", + " target_sensor_name: str,\n", + ") -> Optional[Path]:\n", + " \"\"\"Find the file path to the target sensor from a source sensor.\"\"\"\n", + " if synchronization_df is None:\n", + " raise RuntimeError(\"Requested synchronized data, but the synchronization database has not been created.\")\n", + "\n", + " src_timedelta_ns = pd.Timedelta(src_timestamp_ns)\n", + " src_to_target_records = synchronization_df.loc[(split, log_id, src_sensor_name)].set_index(src_sensor_name)\n", + " index = src_to_target_records.index\n", + " if src_timedelta_ns not in index:\n", + " # This timestamp does not correspond to any lidar sweep.\n", + " return None\n", + "\n", + " # Grab the synchronization record.\n", + " target_timestamp_ns = src_to_target_records.loc[src_timedelta_ns, target_sensor_name]\n", + " if pd.isna(target_timestamp_ns):\n", + " # No match was found within tolerance.\n", + " return None\n", + "\n", + " sensor_dir = dataset_dir / split / log_id / \"sensors\"\n", + " valid_cameras = list(AV2_CAMERA_TYPE_MAPPING.keys())\n", + " timestamp_ns_str = str(target_timestamp_ns.asm8.item())\n", + " if target_sensor_name in valid_cameras:\n", + " target_path = sensor_dir / \"cameras\" / target_sensor_name / f\"{timestamp_ns_str}.jpg\"\n", + " else:\n", + " target_path = sensor_dir / target_sensor_name / f\"{timestamp_ns_str}.feather\"\n", + " return target_path" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "# split=\"train\"\n", + "# log_id=\"00a6ffc1-6ce9-3bc3-a060-6006e9893a1a\"\n", + "# src_sensor_name=\"ring_front_center\"\n", + "# src_timestamp_ns=315967376959702000\n", + "# target_sensor_name=\"lidar\"\n", + "\n", + "# src_to_target_records = synchronization_df.loc[(\"train\", \"\", src_sensor_name)]\n", + "# # synchronization_df" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "lidar_timestamp_ns_list = [int(path.stem) for path in lidar_files]\n", + "\n", + "\n", + "\n", + "\n", + "for lidar_timestamp_ns in lidar_timestamp_ns_list:\n", + "\n", + " fpath = find_closest_target_fpath(\n", + " split=\"train\",\n", + " log_id=log_folder.name,\n", + " src_sensor_name=\"lidar\",\n", + " src_timestamp_ns=lidar_timestamp_ns,\n", + " target_sensor_name=\"ring_front_center\",\n", + " )\n", + " if fpath is None:\n", + " continue\n", + " # print(fpath)\n", + "\n", + " egovehicle_se3_sensor_table" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "for _, row in egovehicle_se3_sensor_table.iterrows():\n", + " row = row.to_dict()\n", + " print(row)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "14", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "15", + "metadata": {}, + "outputs": [], + "source": [ + "import cv2\n", + "from pyquaternion import Quaternion\n", + "\n", + "lidar_timestamps = [int(f.stem) for f in lidar_files]\n", + "camera_timestamps = [int(f.stem) for f in camera_files]\n", + "\n", + "\n", + "def get_slice_with_timestamp_ns(dataframe: pd.DataFrame, timestamp_ns: int):\n", + " \"\"\"Get the index of the closest timestamp to the target timestamp.\"\"\"\n", + " return dataframe[dataframe[\"timestamp_ns\"] == timestamp_ns]\n", + "\n", + "\n", + "def find_nearest_timestamp(target_ns, timestamp_list):\n", + " timestamp_array = np.array(timestamp_list, dtype=np.int64)\n", + " idx = np.argmin(np.abs(timestamp_array - np.int64(target_ns)))\n", + " return int(timestamp_array[idx])\n", + "\n", + "# for lidar_timestamp in lidar_timestamps:\n", + "# slice = get_slice_with_timestamp_ns(annotations_table, lidar_timestamp)\n", + "# assert len(slice) >= 1\n", + "\n", + "\n", + "\n", + "# ego_pose = city_SE3_egovehicle_table[city_SE3_egovehicle_table[\"timestamp_ns\"] == lidar_timestamps[10]]\n", + "# ego_pose_dict = ego_pose.iloc[0].to_dict()\n", + "\n", + "\n", + "annotations_slice = get_slice_with_timestamp_ns(annotations_table, lidar_timestamps[10])\n", + "for _, row in annotations_slice.iterrows():\n", + "# qw = row[\"qw\"]\n", + "# qx = row[\"qx\"]\n", + "# qy = row[\"qy\"]\n", + "# qz = row[\"qz\"]\n", + "# tx_m = row[\"tx_m\"]\n", + "# ty_m = row[\"ty_m\"]\n", + "# tz_m = row[\"tz_m\"]\n", + " print(row.to_dict())\n", + "\n", + "annotations_slice\n", + "\n", + "# qw\tqx\tqy\tqz\ttx_m\tty_m\ttz_m\n", + "# # def jpg_to_array(file_path):\n", + "\n", + "# camera_frames = []\n", + "# for lidar_timestamp in lidar_timestamps:\n", + "# camera_stamp_at_lidar = find_nearest_timestamp(lidar_timestamp, camera_timestamps)\n", + "# image = jpg_to_array(camera_folder / f\"{camera_stamp_at_lidar}.jpg\")\n", + "# camera_frames.append(image)\n", + " \n", + "# print(len(camera_frames))\n", + "# height, width, _ = camera_frames[0].shape\n", + "# video_path = \"camera_frames_video.mp4\"\n", + "# fourcc = cv2.VideoWriter_fourcc(*'mp4v')\n", + "# out = cv2.VideoWriter(video_path, fourcc, 10, (width, height))\n", + "\n", + "# for frame in camera_frames:\n", + "# out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))\n", + "\n", + "# out.release()\n", + "# print(f\"Saved video to {video_path}\")\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "\n", + "from pyquaternion import Quaternion\n", + "from d123.common.datatypes.detection.detection_types import DetectionType\n", + "from d123.common.geometry.base import StateSE2\n", + "from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE2\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.color.default import BOX_DETECTION_CONFIG\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_polygon_to_ax\n", + "\n", + "\n", + "\n", + "fig, ax = plt.subplots(1, 1, figsize=(10, 10))\n", + "\n", + "for cuboid in sensor_data.annotations:\n", + " yaw, pitch, roll = Quaternion(matrix=cuboid.dst_SE3_object.rotation).yaw_pitch_roll\n", + " center = StateSE2(cuboid.dst_SE3_object.translation[0], cuboid.dst_SE3_object.translation[1], yaw)\n", + " bounding_box = BoundingBoxSE2(center, cuboid.length_m, cuboid.width_m)\n", + " add_shapely_polygon_to_ax(ax, bounding_box.shapely_polygon, BOX_DETECTION_CONFIG[DetectionType.VEHICLE])\n", + "\n", + "ax.set_aspect(\"equal\")\n", + "\n", + "radius = 200\n", + "ax.set_xlim(-radius, radius)\n", + "ax.set_ylim(-radius, radius)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17", + "metadata": {}, + "outputs": [], + "source": [ + "bounding_box.shapely_polygon\n", + "\n", + "bounding_box.corners_array" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "sensor_cache = \"/home/daniel/.cache/av2/sensor_cache.feather\"\n", + "get_arrow_from_file(Path(sensor_cache))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "19", + "metadata": {}, + "outputs": [], + "source": [ + "synchronization_cache = \"/home/daniel/.cache/av2/synchronization_cache.feather\"\n", + "synchronization_cache = get_arrow_from_file(Path(synchronization_cache))\n", + "\n", + "synchronization_cache[\"sensor_name\"].unique()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20", + "metadata": {}, + "outputs": [], + "source": [ + "CAM_SHUTTER_INTERVAL_MS" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123_dev", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 6ead233b..fe16f19b 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -27,7 +27,8 @@ "\n", "# splits = [\"wopd_train\"]\n", "# splits = [\"carla\"]\n", - "splits = [\"nuplan_private_test\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -38,7 +39,7 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=19,\n", + " duration_s=10,\n", " history_s=0.0,\n", " timestamp_threshold_s=20,\n", " shuffle=True,\n", @@ -216,7 +217,7 @@ " box_detections = scene.get_box_detections_at_iteration(iteration)\n", "\n", " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", - " add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", @@ -235,15 +236,15 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " size = 15\n", + " size = 2\n", "\n", " fig, ax = plt.subplots(figsize=(size, size))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", - "scene_index = 5\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=150, radius=40)\n", + "scene_index = 0\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=10, radius=100)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" ] @@ -254,6 +255,22 @@ "id": "3", "metadata": {}, "outputs": [], + "source": [ + "scene = scenes[scene_index]\n", + "\n", + "\n", + "scene.get_camera_at_iteration(camera_type=CameraType.CAM_F0, iteration=0)\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], "source": [ "from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", @@ -290,7 +307,7 @@ { "cell_type": "code", "execution_count": null, - "id": "4", + "id": "5", "metadata": {}, "outputs": [], "source": [] @@ -298,7 +315,7 @@ { "cell_type": "code", "execution_count": null, - "id": "5", + "id": "6", "metadata": {}, "outputs": [], "source": [] @@ -306,7 +323,7 @@ { "cell_type": "code", "execution_count": null, - "id": "6", + "id": "7", "metadata": {}, "outputs": [], "source": [] @@ -314,7 +331,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [], "source": [ diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index aa3ed869..d9e21ec3 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -38,7 +38,8 @@ "\n", "# splits = [\"carla\"]\n", "# splits = [\"nuplan_private_test\"]\n", - "splits = [\"wopd_train\"]\n", + "# splits = [\"wopd_train\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -52,9 +53,9 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=19,\n", + " duration_s=15,\n", " history_s=0.0,\n", - " timestamp_threshold_s=20,\n", + " timestamp_threshold_s=15,\n", " shuffle=False,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", @@ -72,9 +73,7 @@ "id": "2", "metadata": {}, "outputs": [], - "source": [ - "all([])" - ] + "source": [] }, { "cell_type": "code", @@ -125,16 +124,51 @@ "metadata": {}, "outputs": [], "source": [ + "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L2: ax[1, 0],\n", + " CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R2: ax[1, 2],\n", + " }\n", + "\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " if camera_type not in scene.available_camera_types:\n", + " continue\n", + " # assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + " return ax\n", + "\n", + "\n", + "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 1\n", + " fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", + " _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "plot_cameras_with_detections(scenes[9], iteration=20)\n", + "\n", + "\n", "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", "\n", - "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", - "# CameraType.CAM_L0: ax[0, 0],\n", - "# CameraType.CAM_F0: ax[0, 1],\n", - "# CameraType.CAM_R0: ax[0, 2],\n", - "# CameraType.CAM_L2: ax[1, 0],\n", - "# CameraType.CAM_B0: ax[1, 1],\n", - "# CameraType.CAM_R2: ax[1, 2],\n", - "# }\n", + "# camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", "\n", "# box_detections = scene.get_box_detections_at_iteration(iteration)\n", "# ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", @@ -155,7 +189,7 @@ "# :return: figure and ax object of matplotlib\n", "# \"\"\"\n", "# scale = 2\n", - "# fig, ax = plt.subplots(2, 3, figsize=(scale * 11.5, scale * 4.3))\n", + "# fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", "# _add_camera_with_detections(ax, scene, iteration)\n", "\n", "# fig.tight_layout()\n", @@ -163,38 +197,6 @@ "\n", "\n", "\n", - "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", - "\n", - " camera_ax_assignment: Dict[CameraType, plt.Axes] = {CameraType.CAM_F0: ax}\n", - "\n", - " box_detections = scene.get_box_detections_at_iteration(iteration)\n", - " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", - " for camera_type, camera_ax in camera_ax_assignment.items():\n", - " assert camera_type in scene.available_camera_types\n", - " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", - " if camera is not None:\n", - " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", - "\n", - " return ax\n", - "\n", - "\n", - "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", - " \"\"\"\n", - " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", - " :param scene: navsim scene dataclass\n", - " :param frame_idx: index of selected frame\n", - " :return: figure and ax object of matplotlib\n", - " \"\"\"\n", - " scale = 2\n", - " fig, ax = plt.subplots(1, 1, figsize=(scale * 11.5, scale * 4.3))\n", - " _add_camera_with_detections(ax, scene, iteration)\n", - "\n", - " fig.tight_layout()\n", - " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", - "\n", - "\n", - "plot_cameras_with_detections(scenes[9], iteration=20)\n", - "\n", "\n", "\n", "# def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", @@ -232,7 +234,7 @@ "# fig.tight_layout()\n", "# fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", "\n", - "plot_cameras_with_detections(scenes[0], iteration=0)\n" + "# plot_cameras_with_detections(scenes[0], iteration=0)\n" ] }, { @@ -242,7 +244,48 @@ "metadata": {}, "outputs": [], "source": [ - "scenes[3].log_name" + "def _add_camera_with_detections(ax: plt.Axes, scene: AbstractScene, iteration: int) -> plt.Axes:\n", + "\n", + " camera_ax_assignment: Dict[CameraType, plt.Axes] = {\n", + " CameraType.CAM_L0: ax[0, 0],\n", + " CameraType.CAM_F0: ax[0, 1],\n", + " CameraType.CAM_R0: ax[0, 2],\n", + " CameraType.CAM_L1: ax[1, 0],\n", + " # CameraType.CAM_B0: ax[1, 1],\n", + " CameraType.CAM_R1: ax[1, 2],\n", + " CameraType.CAM_L2: ax[2, 0],\n", + " CameraType.CAM_R2: ax[2, 2],\n", + " }\n", + "\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + " ego_state_se3 = scene.get_ego_state_at_iteration(iteration)\n", + " for camera_type, camera_ax in camera_ax_assignment.items():\n", + " if camera_type not in scene.available_camera_types:\n", + " continue\n", + " # assert camera_type in scene.available_camera_types\n", + " camera = scene.get_camera_at_iteration(iteration, camera_type)\n", + " if camera is not None:\n", + " add_box_detections_to_camera_ax(camera_ax, camera, box_detections, ego_state_se3)\n", + "\n", + " return ax\n", + "\n", + "\n", + "def plot_cameras_with_detections(scene: AbstractScene, iteration: int) -> Tuple[plt.Figure, plt.Axes]:\n", + " \"\"\"\n", + " Plots 8x cameras and birds-eye-view visualization in 3x3 grid\n", + " :param scene: navsim scene dataclass\n", + " :param frame_idx: index of selected frame\n", + " :return: figure and ax object of matplotlib\n", + " \"\"\"\n", + " scale = 1\n", + " fig, ax = plt.subplots(3, 3, figsize=(scale * 10, scale * 10))\n", + " _add_camera_with_detections(ax, scene, iteration)\n", + "\n", + " fig.tight_layout()\n", + " fig.subplots_adjust(wspace=0.01, hspace=0.01, left=0.01, right=0.99, top=0.99, bottom=0.01)\n", + "\n", + "\n", + "plot_cameras_with_detections(scenes[9], iteration=20)\n" ] }, { @@ -251,6 +294,16 @@ "id": "6", "metadata": {}, "outputs": [], + "source": [ + "scenes[3].log_name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], "source": [ "from pathlib import Path\n", "from typing import Optional, Union\n", @@ -306,7 +359,7 @@ { "cell_type": "code", "execution_count": null, - "id": "7", + "id": "8", "metadata": {}, "outputs": [], "source": [] @@ -314,7 +367,7 @@ { "cell_type": "code", "execution_count": null, - "id": "8", + "id": "9", "metadata": {}, "outputs": [], "source": [] @@ -322,7 +375,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123", + "display_name": "d123_dev", "language": "python", "name": "python3" }, diff --git a/notebooks/viz/log_rendering.ipynb b/notebooks/viz/log_rendering.ipynb index e88c1012..fe88e006 100644 --- a/notebooks/viz/log_rendering.ipynb +++ b/notebooks/viz/log_rendering.ipynb @@ -47,7 +47,7 @@ "import traceback\n", "from d123.common.visualization.matplotlib.plots import render_scene_animation\n", "\n", - "split = \"wopd_train\"\n", + "split = \"av2-sensor-mini_train\"\n", "output_path = Path(f\"/home/daniel/d123_logs_videos/{split}\")\n", "log_path = Path(f\"/home/daniel/d123_workspace/data/{split}\")\n", "for log_file in log_path.iterdir():\n", @@ -83,7 +83,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123", + "display_name": "d123_dev", "language": "python", "name": "python3" }, diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 740ce2e2..567edd06 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -21,11 +21,16 @@ "metadata": {}, "outputs": [], "source": [ + "import requests\n", + "from PIL import Image\n", + "from io import BytesIO\n", + "\n", "\n", "\n", "# splits = [\"nuplan_private_test\"]\n", "# splits = [\"carla\"]\n", - "splits = [\"wopd_train\"]\n", + "# splits = [\"wopd_train\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", "log_names = None\n", "\n", "scene_tokens = None\n", @@ -94,7 +99,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123", + "display_name": "d123_dev", "language": "python", "name": "python3" }, diff --git a/scripts/download/download_av2.sh b/scripts/download/download_av2.sh index 7c63854a..039de648 100644 --- a/scripts/download/download_av2.sh +++ b/scripts/download/download_av2.sh @@ -6,10 +6,14 @@ # s3://argoverse/datasets/av2/motion-forecasting/ # s3://argoverse/datasets/av2/tbv/ -DATASET_NAMES=("sensor" "lidar" "motion_forecasting" "tbv") +DATASET_NAMES=("sensor" "lidar" "motion-forecasting" "tbv") TARGET_DIR="/path/to/argoverse" for DATASET_NAME in "${DATASET_NAMES[@]}"; do mkdir -p "$TARGET_DIR/$DATASET_NAME" s5cmd --no-sign-request cp "s3://argoverse/datasets/av2/$DATASET_NAME/*" "$TARGET_DIR/$DATASET_NAME" done + + +# wget -r s3://argoverse/datasets/av2/sensor/test/0f0cdd79-bc6c-35cd-9d99-7ae2fc7e165c/sensors/cameras/ring_front_center/315965893599927217.jpg +# wget http://argoverse.s3.amazonaws.com/datasets/av2/sensor/test/0f0cdd79-bc6c-35cd-9d99-7ae2fc7e165c/sensors/cameras/ring_front_center/315965893599927217.jpg From acc6d9e6e9f1b94b29b5e89416b17fba3a473d32 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Mon, 18 Aug 2025 12:01:13 +0200 Subject: [PATCH 36/42] Add initial argoverse 2 sensor map. --- .../vehicle_state/vehicle_parameters.py | 13 + d123/common/visualization/viser/server.py | 9 +- d123/common/visualization/viser/utils.py | 13 +- d123/dataset/arrow/conversion.py | 1 + .../av2/av2_data_converter.py | 51 ++- .../av2/av2_map_conversion.py | 312 ++++++++++++++++++ d123/dataset/maps/gpkg/gpkg_map.py | 1 + d123/dataset/scene/arrow_scene.py | 2 +- .../common/scene_filter/log_scenes.yaml | 4 +- .../default_dataset_conversion.yaml | 2 +- .../config/datasets/av2_sensor_dataset.yaml | 4 +- notebooks/viz/bev_matplotlib.ipynb | 22 +- 12 files changed, 392 insertions(+), 42 deletions(-) diff --git a/d123/common/datatypes/vehicle_state/vehicle_parameters.py b/d123/common/datatypes/vehicle_state/vehicle_parameters.py index 8fe4d048..b6d6b330 100644 --- a/d123/common/datatypes/vehicle_state/vehicle_parameters.py +++ b/d123/common/datatypes/vehicle_state/vehicle_parameters.py @@ -61,6 +61,19 @@ def get_wopd_pacifica_parameters() -> VehicleParameters: ) +def get_av2_fusion_hybrid_parameters() -> VehicleParameters: + # NOTE: use parameters from nuPlan dataset + return VehicleParameters( + vehicle_name="av2_fusion_hybrid", + width=2.297, + length=5.176, + height=1.777, + wheel_base=3.089, + rear_axle_to_center_vertical=0.5885, + rear_axle_to_center_longitudinal=1.461, + ) + + def center_se3_to_rear_axle_se3(center_se3: StateSE3, vehicle_parameters: VehicleParameters) -> StateSE3: """ Converts a center state to a rear axle state. diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index bf0acc26..cdca86c4 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -33,19 +33,20 @@ LINE_WIDTH: float = 4.0 # Bounding box config: -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" # Map config: -MAP_AVAILABLE: bool = False +MAP_AVAILABLE: bool = True # Cameras config: VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_F0, CameraType.CAM_L0, CameraType.CAM_R0] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = all_camera_types +# VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [CameraType.CAM_STEREO_L, CameraType.CAM_STEREO_R] # VISUALIZE_CAMERA_FRUSTUM: List[CameraType] = [] -VISUALIZE_CAMERA_GUI: List[CameraType] = [] -CAMERA_SCALE: float = 2.0 +VISUALIZE_CAMERA_GUI: List[CameraType] = [CameraType.CAM_F0] +CAMERA_SCALE: float = 1.0 # Lidar config: LIDAR_AVAILABLE: bool = False diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index d527b05e..faee4725 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -96,12 +96,12 @@ def get_map_meshes(scene: AbstractScene): initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 map_layers = [ - MapLayer.LANE_GROUP, - # MapLayer.LANE, - MapLayer.WALKWAY, - MapLayer.CROSSWALK, - MapLayer.CARPARK, - MapLayer.GENERIC_DRIVABLE, + # MapLayer.LANE_GROUP, + MapLayer.LANE, + # MapLayer.WALKWAY, + # MapLayer.CROSSWALK, + # MapLayer.CARPARK, + # MapLayer.GENERIC_DRIVABLE, ] map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) @@ -231,7 +231,6 @@ def get_camera_values( rear_axle = StateSE3.from_array(rear_axle_array) camera_to_ego = camera.extrinsic # 4x4 transformation from camera to ego frame - camera.image # Get the rotation matrix of the rear axle pose from d123.common.geometry.transform.se3 import get_rotation_matrix diff --git a/d123/dataset/arrow/conversion.py b/d123/dataset/arrow/conversion.py index d9afba6f..68251f82 100644 --- a/d123/dataset/arrow/conversion.py +++ b/d123/dataset/arrow/conversion.py @@ -33,6 +33,7 @@ DATASET_SENSOR_ROOT: Dict[str, Path] = { "nuplan": Path(os.environ["NUPLAN_DATA_ROOT"]) / "nuplan-v1.1" / "sensor_blobs", "carla": Path(os.environ["CARLA_DATA_ROOT"]) / "sensor_blobs", + # "av2-sensor": Path(os.environ["AV2_SENSOR_DATA_ROOT"]) / "sensor", } diff --git a/d123/dataset/dataset_specific/av2/av2_data_converter.py b/d123/dataset/dataset_specific/av2/av2_data_converter.py index 36dc143e..984d59c7 100644 --- a/d123/dataset/dataset_specific/av2/av2_data_converter.py +++ b/d123/dataset/dataset_specific/av2/av2_data_converter.py @@ -16,13 +16,13 @@ from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( - get_nuplan_pacifica_parameters, + get_av2_fusion_hybrid_parameters, rear_axle_se3_to_center_se3, ) from d123.common.geometry.base import StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL -from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array +from d123.common.geometry.transform.se3 import convert_relative_to_absolute_se3_array, get_rotation_matrix from d123.common.geometry.vector import Vector3D, Vector3DIndex from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.dataset_specific.av2.av2_constants import ( @@ -36,7 +36,7 @@ find_closest_target_fpath, get_slice_with_timestamp_ns, ) -from d123.dataset.dataset_specific.nuplan.nuplan_map_conversion import MAP_LOCATIONS +from d123.dataset.dataset_specific.av2.av2_map_conversion import convert_av2_map from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata @@ -100,10 +100,18 @@ def get_available_splits(self) -> List[str]: ] def convert_maps(self, worker: WorkerPool) -> None: + log_args = [ + { + "log_path": log_path, + "split": split, + } + for split, log_paths in self._log_paths_per_split.items() + for log_path in log_paths + ] worker_map( worker, partial(convert_av2_map_to_gpkg, data_converter_config=self.data_converter_config), - list(MAP_LOCATIONS), + log_args, ) def convert_logs(self, worker: WorkerPool) -> None: @@ -126,13 +134,20 @@ def convert_logs(self, worker: WorkerPool) -> None: ) -def convert_av2_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: - # TODO: Implement map - # for map_name in map_names: - # map_path = data_converter_config.output_path / "maps" / f"av2_{map_name}.gpkg" - # if data_converter_config.force_map_conversion or not map_path.exists(): - # map_path.unlink(missing_ok=True) - # AV2MapConverter(data_converter_config.output_path / "maps").convert(map_name=map_name) +def convert_av2_map_to_gpkg( + args: List[Dict[str, Union[List[str], List[Path]]]], + data_converter_config: DataConverterConfig, +) -> List[Any]: + for log_info in args: + source_log_path: Path = log_info["log_path"] + split: str = log_info["split"] + + source_log_name = source_log_path.name + + map_path = data_converter_config.output_path / "maps" / split / f"{source_log_name}.gpkg" + if data_converter_config.force_map_conversion or not map_path.exists(): + map_path.unlink(missing_ok=True) + convert_av2_map(source_log_path, map_path) return [] @@ -164,7 +179,7 @@ def convert_av2_log_to_arrow( timestep_seconds=0.1, # TODO: verify this map_has_z=True, ) - vehicle_parameters = get_nuplan_pacifica_parameters() # TODO: Add av2 vehicle parameters + vehicle_parameters = get_av2_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters camera_metadata = get_av2_camera_metadata(log_path) lidar_metadata = get_av2_lidar_metadata(log_path) @@ -327,6 +342,7 @@ def _write_recording_table( lidar_timestamp_ns, city_se3_egovehicle_df, egovehicle_se3_sensor_df, + ego_state_se3, synchronization_df, source_log_path, data_converter_config, @@ -422,7 +438,7 @@ def _extract_ego_state(city_se3_egovehicle_df: pd.DataFrame, lidar_timestamp_ns: pitch=pitch, yaw=yaw, ) - vehicle_parameters = get_nuplan_pacifica_parameters() # TODO: Add av2 vehicle parameters + vehicle_parameters = get_av2_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) # TODO: Add script to calculate the dynamic state from log sequence. dynamic_state = DynamicStateSE3( @@ -463,6 +479,7 @@ def _extract_camera( lidar_timestamp_ns: int, city_se3_egovehicle_df: pd.DataFrame, egovehicle_se3_sensor_df: pd.DataFrame, + ego_state_se3: EgoStateSE3, synchronization_df: pd.DataFrame, source_log_path: Path, data_converter_config: DataConverterConfig, @@ -476,6 +493,11 @@ def _extract_camera( source_dataset_dir = source_log_path.parent.parent + rear_axle_se3 = ego_state_se3.rear_axle_se3 + ego_transform = np.zeros((4, 4), dtype=np.float64) + ego_transform[:3, :3] = get_rotation_matrix(ego_state_se3.rear_axle_se3) + ego_transform[:3, 3] = rear_axle_se3.point_3d.array + for _, row in egovehicle_se3_sensor_df.iterrows(): row = row.to_dict() if row["sensor_name"] not in AV2_CAMERA_TYPE_MAPPING: @@ -498,7 +520,7 @@ def _extract_camera( absolute_image_path = source_dataset_dir / relative_image_path assert absolute_image_path.exists() # TODO: Adjust for finer IMU timestamps to correct the camera extrinsic. - camera_extrinsic = np.zeros((4, 4), dtype=np.float64) + camera_extrinsic = np.eye(4, dtype=np.float64) camera_extrinsic[:3, :3] = Quaternion( w=row["qw"], x=row["qx"], @@ -506,6 +528,7 @@ def _extract_camera( z=row["qz"], ).rotation_matrix camera_extrinsic[:3, 3] = np.array([row["tx_m"], row["ty_m"], row["tz_m"]], dtype=np.float64) + # camera_extrinsic = camera_extrinsic @ ego_transform camera_extrinsic = camera_extrinsic.flatten().tolist() if data_converter_config.camera_store_option == "path": diff --git a/d123/dataset/dataset_specific/av2/av2_map_conversion.py b/d123/dataset/dataset_specific/av2/av2_map_conversion.py index e69de29b..368edda8 100644 --- a/d123/dataset/dataset_specific/av2/av2_map_conversion.py +++ b/d123/dataset/dataset_specific/av2/av2_map_conversion.py @@ -0,0 +1,312 @@ +from collections import defaultdict +from pathlib import Path +from typing import Any, Dict, List + +import geopandas as gpd +import numpy as np +import numpy.typing as npt +import pandas as pd +import shapely.geometry as geom +from flask import json + +from d123.common.geometry.base import Point3DIndex +from d123.common.geometry.line.polylines import Polyline3D +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType + +# TODO: +# - TODO + + +def convert_av2_map(source_log_path: Path, map_file_path: Path) -> None: + + def _extract_polyline(data: List[Dict[str, float]]) -> Polyline3D: + polyline = np.array([[p["x"], p["y"], p["z"]] for p in data], dtype=np.float64) + return Polyline3D.from_array(polyline) + + map_folder = source_log_path / "map" + + next(map_folder.glob("*.npy")) + log_map_archive_path = next(map_folder.glob("log_map_archive_*.json")) + next(map_folder.glob("*img_Sim2_city.json*")) + + with open(log_map_archive_path, "r") as f: + log_map_archive = json.load(f) + + defaultdict(list) + drivable_areas: Dict[int, Polyline3D] = {} + + for drivable_area_id, drivable_area_dict in log_map_archive["drivable_areas"].items(): + # keys: ["area_boundary", "id"] + drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict["area_boundary"]) + + for lane_segment_id, lane_segment_dict in log_map_archive["lane_segments"].items(): + # keys = [ + # "id", + # "is_intersection", + # "lane_type", + # "left_lane_boundary", + # "left_lane_mark_type", + # "right_lane_boundary", + # "right_lane_mark_type", + # "successors", + # "predecessors", + # "right_neighbor_id", + # "left_neighbor_id", + # ] + lane_segment_dict["left_lane_boundary"] = _extract_polyline(lane_segment_dict["left_lane_boundary"]) + lane_segment_dict["right_lane_boundary"] = _extract_polyline(lane_segment_dict["right_lane_boundary"]) + + lane_df = get_lane_df(log_map_archive["lane_segments"]) + get_empty_gdf() + get_empty_gdf() + get_empty_gdf() + get_empty_gdf() + get_empty_gdf() + generic_drivable_df = get_generic_drivable_df(drivable_areas) + get_empty_gdf() + get_empty_gdf() + + map_file_path.unlink(missing_ok=True) + if not map_file_path.parent.exists(): + map_file_path.parent.mkdir(parents=True, exist_ok=True) + + lane_df.to_file(map_file_path, layer=MapLayer.LANE.serialize(), driver="GPKG") + # lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") + # intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") + # crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") + # walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") + # carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") + generic_drivable_df.to_file(map_file_path, layer=MapLayer.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a") + # road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + # road_line_df.to_file(map_file_path, layer=MapLayer.ROAD_LINE.serialize(), driver="GPKG", mode="a") + + +def get_lane_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: + + ids = list(lanes.keys()) + lane_types = [0] * len(ids) # TODO: Add lane types + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] + baseline_paths = [] + geometries = [] + + def _get_centerline_from_boundaries( + left_boundary: Polyline3D, right_boundary: Polyline3D, resolution: float = 0.1 + ) -> Polyline3D: + + points_per_meter = 1 / resolution + num_points = int(np.ceil(max([right_boundary.length, left_boundary.length]) * points_per_meter)) + right_array = right_boundary.interpolate(np.linspace(0, right_boundary.length, num_points, endpoint=True)) + left_array = left_boundary.interpolate(np.linspace(0, left_boundary.length, num_points, endpoint=True)) + + return Polyline3D.from_array(np.mean([right_array, left_array], axis=0)) + + for lane_id, lane_dict in lanes.items(): + # keys = [ + # "id", + # "is_intersection", + # "lane_type", + # "left_lane_boundary", + # "left_lane_mark_type", + # "right_lane_boundary", + # "right_lane_mark_type", + # "successors", + # "predecessors", + # "right_neighbor_id", + # "left_neighbor_id", + # ] + lane_centerline = _get_centerline_from_boundaries( + left_boundary=lane_dict["left_lane_boundary"], right_boundary=lane_dict["right_lane_boundary"] + ) + lane_speed_limit_mps = None + + # ids.append(lane_id) + # lane_types.append(lanes_type[lane_id]) + lane_group_ids.append(lane_id) + speed_limits_mps.append(lane_speed_limit_mps) + predecessor_ids.append(lane_dict["predecessors"]) + successor_ids.append(lane_dict["successors"]) + left_boundaries.append(lane_dict["left_lane_boundary"].linestring) + right_boundaries.append(lane_dict["right_lane_boundary"].linestring) + left_lane_ids.append(lane_dict["left_neighbor_id"]) + right_lane_ids.append(lane_dict["right_neighbor_id"]) + baseline_paths.append(lane_centerline.linestring) + + geometry = geom.Polygon( + np.vstack( + [ + lane_dict["left_lane_boundary"].array[:, :2], + lane_dict["right_lane_boundary"].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_type": lane_types, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, + "baseline_path": baseline_paths, + } + ) + + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_empty_gdf() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_lane_group_df( + lanes: Dict[int, npt.NDArray[np.float64]], + lanes_successors: Dict[int, List[int]], + lanes_predecessors: Dict[int, List[int]], + lanes_left_boundaries_3d: Dict[int, Polyline3D], + lanes_right_boundaries_3d: Dict[int, Polyline3D], +) -> gpd.GeoDataFrame: + + ids = [] + lane_ids = [] + intersection_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + # NOTE: WOPD does not provide lane groups, so we create a lane group for each lane. + for lane_id in lanes.keys(): + if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: + continue + ids.append(lane_id) + lane_ids.append([lane_id]) + intersection_ids.append(None) # WOPD does not provide intersections + predecessor_lane_group_ids.append(lanes_predecessors[lane_id]) + successor_lane_group_ids.append(lanes_successors[lane_id]) + left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) + right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + geometry = geom.Polygon( + np.vstack( + [ + lanes_left_boundaries_3d[lane_id].array[:, :2], + lanes_right_boundaries_3d[lane_id].array[:, :2][::-1], + ] + ) + ) + geometries.append(geometry) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "intersection_id": intersection_ids, + "predecessor_lane_group_ids": predecessor_lane_group_ids, + "successor_lane_group_ids": successor_lane_group_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_intersections_df() -> gpd.GeoDataFrame: + ids = [] + lane_group_ids = [] + geometries = [] + + # NOTE: WOPD does not provide intersections, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_carpark_df(carparks) -> gpd.GeoDataFrame: + ids = list(carparks.keys()) + outlines = [geom.LineString(outline) for outline in carparks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in carparks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_walkway_df() -> gpd.GeoDataFrame: + ids = [] + geometries = [] + + # NOTE: WOPD does not provide walkways, so we create an empty DataFrame. + data = pd.DataFrame({"id": ids}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: + ids = list(crosswalks.keys()) + outlines = [geom.LineString(outline) for outline in crosswalks.values()] + geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in crosswalks.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_generic_drivable_df(drivable_areas: Dict[int, Polyline3D]) -> gpd.GeoDataFrame: + ids = list(drivable_areas.keys()) + outlines = [drivable_area.linestring for drivable_area in drivable_areas.values()] + geometries = [geom.Polygon(drivable_area.array[:, Point3DIndex.XY]) for drivable_area in drivable_areas.values()] + + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_road_edge_df( + road_edges: Dict[int, npt.NDArray[np.float64]], road_edges_type: Dict[int, RoadEdgeType] +) -> gpd.GeoDataFrame: + ids = list(road_edges.keys()) + geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_edges.values()] + + data = pd.DataFrame( + { + "id": ids, + "road_edge_type": [int(road_edge_type) for road_edge_type in road_edges_type.values()], + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + +def get_road_line_df( + road_lines: Dict[int, npt.NDArray[np.float64]], road_lines_type: Dict[int, RoadLineType] +) -> gpd.GeoDataFrame: + ids = list(road_lines.keys()) + geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_lines.values()] + + data = pd.DataFrame( + { + "id": ids, + "road_line_type": [int(road_line_type) for road_line_type in road_lines_type.values()], + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf diff --git a/d123/dataset/maps/gpkg/gpkg_map.py b/d123/dataset/maps/gpkg/gpkg_map.py index 29df246d..32b02d41 100644 --- a/d123/dataset/maps/gpkg/gpkg_map.py +++ b/d123/dataset/maps/gpkg/gpkg_map.py @@ -74,6 +74,7 @@ def initialize(self) -> None: self._gpd_dataframes[map_layer]["id"] = self._gpd_dataframes[map_layer]["id"].astype(str) else: warnings.warn(f"GPKGMap: {map_layer_name} not available in {str(self._file_path)}") + self._gpd_dataframes[map_layer] = None def _assert_initialize(self) -> None: "Checks if `.initialize()` was called, before retrieving data." diff --git a/d123/dataset/scene/arrow_scene.py b/d123/dataset/scene/arrow_scene.py index 6856d102..ecd68111 100644 --- a/d123/dataset/scene/arrow_scene.py +++ b/d123/dataset/scene/arrow_scene.py @@ -186,7 +186,7 @@ def _lazy_initialize(self) -> None: def open(self) -> None: if self._map_api is None: try: - if self._metadata.dataset == "wopd": + if self._metadata.dataset in ["wopd", "av2-sensor"]: # FIXME: split = str(self._arrow_log_path.parent.name) self._map_api = get_local_map_api(split, self._metadata.log_name) diff --git a/d123/script/config/common/scene_filter/log_scenes.yaml b/d123/script/config/common/scene_filter/log_scenes.yaml index 6c3bcee7..68114df9 100644 --- a/d123/script/config/common/scene_filter/log_scenes.yaml +++ b/d123/script/config/common/scene_filter/log_scenes.yaml @@ -2,7 +2,7 @@ _target_: d123.dataset.scene.scene_filter.SceneFilter _convert_: 'all' split_types: null -split_names: null +split_names: ["av2-sensor-mini_train"] log_names: null @@ -12,7 +12,7 @@ timestamp_threshold_s: null ego_displacement_minimum_m: null duration_s: null -history_s: null +history_s: 1.0 camera_types: null diff --git a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml index 7287f850..0a4544da 100644 --- a/d123/script/config/dataset_conversion/default_dataset_conversion.yaml +++ b/d123/script/config/dataset_conversion/default_dataset_conversion.yaml @@ -21,4 +21,4 @@ defaults: - av2_sensor_dataset force_log_conversion: True -force_map_conversion: False +force_map_conversion: True diff --git a/d123/script/config/datasets/av2_sensor_dataset.yaml b/d123/script/config/datasets/av2_sensor_dataset.yaml index d5b4643c..4f9a95ec 100644 --- a/d123/script/config/datasets/av2_sensor_dataset.yaml +++ b/d123/script/config/datasets/av2_sensor_dataset.yaml @@ -3,7 +3,7 @@ av2_sensor_dataset: _convert_: 'all' splits: ["av2-sensor-mini_train"] - log_path: "/mnt/elements_0/argoverse" + log_path: "/media/nvme1/argoverse" data_converter_config: _target_: d123.dataset.dataset_specific.raw_data_converter.DataConverterConfig @@ -12,5 +12,5 @@ av2_sensor_dataset: output_path: ${d123_data_root} force_log_conversion: ${force_log_conversion} force_map_conversion: ${force_map_conversion} - camera_store_option: binary + camera_store_option: "binary" lidar_store_option: null diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index fe16f19b..36a7800e 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -143,14 +143,14 @@ ") -> None:\n", " layers: List[MapLayer] = [\n", " MapLayer.LANE,\n", - " MapLayer.LANE_GROUP,\n", + " # MapLayer.LANE_GROUP,\n", " MapLayer.GENERIC_DRIVABLE,\n", - " MapLayer.CARPARK,\n", - " MapLayer.CROSSWALK,\n", - " MapLayer.INTERSECTION,\n", - " MapLayer.WALKWAY,\n", - " MapLayer.ROAD_EDGE,\n", - " MapLayer.ROAD_LINE,\n", + " # MapLayer.CARPARK,\n", + " # MapLayer.CROSSWALK,\n", + " # MapLayer.INTERSECTION,\n", + " # MapLayer.WALKWAY,\n", + " # MapLayer.ROAD_EDGE,\n", + " # MapLayer.ROAD_LINE,\n", " ]\n", " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", " y_min, y_max = point_2d.y - radius, point_2d.y + radius\n", @@ -217,7 +217,7 @@ " box_detections = scene.get_box_detections_at_iteration(iteration)\n", "\n", " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", - " # add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", @@ -236,14 +236,14 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " size = 2\n", + " size = 20\n", "\n", " fig, ax = plt.subplots(figsize=(size, size))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", - "scene_index = 0\n", + "scene_index = 2\n", "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=10, radius=100)\n", "\n", "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" @@ -344,7 +344,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123_dev", + "display_name": "d123", "language": "python", "name": "python3" }, From e0228fb284722521c90af5a4142769191a654a86 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Tue, 19 Aug 2025 20:32:28 +0200 Subject: [PATCH 37/42] Extent Argoverse 2 (sensor) map implementation. --- .../vehicle_state/vehicle_parameters.py | 29 +-- d123/common/visualization/viser/utils.py | 16 +- .../map/road_edge/road_edge_3d_utils.py | 23 +- .../av2/av2_data_converter.py | 6 +- .../av2/av2_map_conversion.py | 197 ++++++++++++------ .../nuplan/nuplan_data_converter.py | 6 +- .../wopd/wopd_data_converter.py | 6 +- d123/dataset/maps/gpkg/gpkg_map_objects.py | 36 ++-- d123/dataset/maps/map_datatypes.py | 1 + notebooks/viz/bev_matplotlib.ipynb | 62 +++--- notebooks/viz/camera_matplotlib.ipynb | 2 +- 11 files changed, 246 insertions(+), 138 deletions(-) diff --git a/d123/common/datatypes/vehicle_state/vehicle_parameters.py b/d123/common/datatypes/vehicle_state/vehicle_parameters.py index b6d6b330..3b4f04b9 100644 --- a/d123/common/datatypes/vehicle_state/vehicle_parameters.py +++ b/d123/common/datatypes/vehicle_state/vehicle_parameters.py @@ -20,10 +20,10 @@ class VehicleParameters: rear_axle_to_center_longitudinal: float -def get_nuplan_pacifica_parameters() -> VehicleParameters: +def get_nuplan_chrysler_pacifica_parameters() -> VehicleParameters: # NOTE: use parameters from nuPlan dataset return VehicleParameters( - vehicle_name="nuplan_pacifica", + vehicle_name="nuplan_chrysler_pacifica", width=2.297, length=5.176, height=1.777, @@ -46,31 +46,32 @@ def get_carla_lincoln_mkz_2020_parameters() -> VehicleParameters: ) -def get_wopd_pacifica_parameters() -> VehicleParameters: +def get_wopd_chrysler_pacifica_parameters() -> VehicleParameters: # NOTE: use parameters from nuPlan dataset # Find better parameters for WOPD ego vehicle return VehicleParameters( - vehicle_name="wopd_pacifica", + vehicle_name="wopd_chrysler_pacifica", width=2.297, length=5.176, height=1.777, wheel_base=3.089, - # rear_axle_to_center_vertical=0.45, rear_axle_to_center_vertical=1.777 / 2, rear_axle_to_center_longitudinal=1.461, ) -def get_av2_fusion_hybrid_parameters() -> VehicleParameters: - # NOTE: use parameters from nuPlan dataset +def get_av2_ford_fusion_hybrid_parameters() -> VehicleParameters: + # NOTE: Parameters are estimated from the vehicle model. + # https://en.wikipedia.org/wiki/Ford_Fusion_Hybrid#Second_generation + # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/tests/unit/map/test_map_api.py#L375 return VehicleParameters( - vehicle_name="av2_fusion_hybrid", - width=2.297, - length=5.176, - height=1.777, - wheel_base=3.089, - rear_axle_to_center_vertical=0.5885, - rear_axle_to_center_longitudinal=1.461, + vehicle_name="av2_ford_fusion_hybrid", + width=1.852 + 0.275, # 0.275 is the estimated width of the side mirrors + length=4.869, + height=1.476, + wheel_base=2.850, + rear_axle_to_center_vertical=0.438, + rear_axle_to_center_longitudinal=1.339, ) diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index faee4725..43345c0c 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -51,9 +51,9 @@ def bounding_box_to_trimesh(bbox: BoundingBoxSE3, plot_config: PlotConfig) -> tr box_mesh = trimesh.creation.box(extents=[bbox.length, bbox.width, bbox.height]) # Apply rotations in order: roll, pitch, yaw - box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.roll, [1, 0, 0])) - box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.pitch, [0, 1, 0])) box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.yaw, [0, 0, 1])) + box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.pitch, [0, 1, 0])) + box_mesh = box_mesh.apply_transform(trimesh.transformations.rotation_matrix(bbox.center.roll, [1, 0, 0])) # Apply translation box_mesh = box_mesh.apply_translation([bbox.center.x, bbox.center.y, bbox.center.z]) @@ -96,12 +96,12 @@ def get_map_meshes(scene: AbstractScene): initial_ego_vehicle_state = scene.get_ego_state_at_iteration(0) center = initial_ego_vehicle_state.center_se3 map_layers = [ - # MapLayer.LANE_GROUP, - MapLayer.LANE, + MapLayer.LANE_GROUP, + # MapLayer.LANE, # MapLayer.WALKWAY, - # MapLayer.CROSSWALK, + MapLayer.CROSSWALK, # MapLayer.CARPARK, - # MapLayer.GENERIC_DRIVABLE, + MapLayer.GENERIC_DRIVABLE, ] map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) @@ -115,13 +115,13 @@ def get_map_meshes(scene: AbstractScene): if map_layer in [ MapLayer.WALKWAY, MapLayer.CROSSWALK, - MapLayer.GENERIC_DRIVABLE, + # MapLayer.GENERIC_DRIVABLE, MapLayer.CARPARK, ]: # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.1).array else: - trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z).array + trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z + 0.1).array if not scene.log_metadata.map_has_z: trimesh_mesh.vertices += Point3D( diff --git a/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py b/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py index 91e9dc47..93a34526 100644 --- a/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py +++ b/d123/dataset/conversion/map/road_edge/road_edge_3d_utils.py @@ -1,3 +1,4 @@ +import logging from collections import defaultdict from typing import Dict, List, Set @@ -12,6 +13,20 @@ from d123.common.geometry.occupancy_map import OccupancyMap2D from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings +logger = logging.getLogger(__name__) + + +def get_road_edges_3d_from_generic_drivable_area_df(generic_drivable_area_df: gpd.GeoDataFrame) -> List[LineString]: + """ + Extracts 3D road edges from the generic drivable area GeoDataFrame. + """ + # NOTE: this is a simplified version that assumes the generic drivable area covers all areas. + # This is the case for argoverse 2. + road_edges_2d = get_road_edge_linear_rings(generic_drivable_area_df.geometry.tolist()) + outlines = generic_drivable_area_df.outline.tolist() + non_conflicting_road_edges = lift_road_edges_to_3d(road_edges_2d, outlines) + return non_conflicting_road_edges + def get_road_edges_3d_from_gdf( lane_df: gpd.GeoDataFrame, @@ -97,9 +112,13 @@ def _get_conflicting_lane_groups(lane_group_df: gpd.GeoDataFrame, lane_df: gpd.G intersecting_geometry = occupancy_map[intersecting_id] if intersecting_geometry.geom_type != "Polygon": continue + try: + # Compute actual intersection for better centroid + intersection = lane_group_polygon.intersection(intersecting_geometry) + except shapely.errors.GEOSException as e: + logger.debug(f"Error computing intersection for {pair_key}: {e}") + continue - # Compute actual intersection for better centroid - intersection = lane_group_polygon.intersection(intersecting_geometry) if intersection.is_empty: continue diff --git a/d123/dataset/dataset_specific/av2/av2_data_converter.py b/d123/dataset/dataset_specific/av2/av2_data_converter.py index 984d59c7..0ff84a7e 100644 --- a/d123/dataset/dataset_specific/av2/av2_data_converter.py +++ b/d123/dataset/dataset_specific/av2/av2_data_converter.py @@ -16,7 +16,7 @@ from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( - get_av2_fusion_hybrid_parameters, + get_av2_ford_fusion_hybrid_parameters, rear_axle_se3_to_center_se3, ) from d123.common.geometry.base import StateSE3 @@ -179,7 +179,7 @@ def convert_av2_log_to_arrow( timestep_seconds=0.1, # TODO: verify this map_has_z=True, ) - vehicle_parameters = get_av2_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters + vehicle_parameters = get_av2_ford_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters camera_metadata = get_av2_camera_metadata(log_path) lidar_metadata = get_av2_lidar_metadata(log_path) @@ -438,7 +438,7 @@ def _extract_ego_state(city_se3_egovehicle_df: pd.DataFrame, lidar_timestamp_ns: pitch=pitch, yaw=yaw, ) - vehicle_parameters = get_av2_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters + vehicle_parameters = get_av2_ford_fusion_hybrid_parameters() # TODO: Add av2 vehicle parameters center = rear_axle_se3_to_center_se3(rear_axle_se3=rear_axle_pose, vehicle_parameters=vehicle_parameters) # TODO: Add script to calculate the dynamic state from log sequence. dynamic_state = DynamicStateSE3( diff --git a/d123/dataset/dataset_specific/av2/av2_map_conversion.py b/d123/dataset/dataset_specific/av2/av2_map_conversion.py index 368edda8..b294573c 100644 --- a/d123/dataset/dataset_specific/av2/av2_map_conversion.py +++ b/d123/dataset/dataset_specific/av2/av2_map_conversion.py @@ -1,6 +1,5 @@ -from collections import defaultdict from pathlib import Path -from typing import Any, Dict, List +from typing import Any, Dict, Final, List import geopandas as gpd import numpy as np @@ -11,16 +10,25 @@ from d123.common.geometry.base import Point3DIndex from d123.common.geometry.line.polylines import Polyline3D +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length +from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import ( + get_road_edges_3d_from_generic_drivable_area_df, +) from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType # TODO: # - TODO +LANE_GROUP_MARK_TYPES = ["DASHED_WHITE", "DOUBLE_DASH_WHITE", "DASH_SOLID_WHITE", "SOLID_DASH_WHITE", "SOLID_WHITE"] +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # TODO: Add to config def convert_av2_map(source_log_path: Path, map_file_path: Path) -> None: - def _extract_polyline(data: List[Dict[str, float]]) -> Polyline3D: + def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Polyline3D: polyline = np.array([[p["x"], p["y"], p["z"]] for p in data], dtype=np.float64) + if close: + polyline = np.vstack([polyline, polyline[0]]) + return Polyline3D.from_array(polyline) map_folder = source_log_path / "map" @@ -32,12 +40,11 @@ def _extract_polyline(data: List[Dict[str, float]]) -> Polyline3D: with open(log_map_archive_path, "r") as f: log_map_archive = json.load(f) - defaultdict(list) drivable_areas: Dict[int, Polyline3D] = {} for drivable_area_id, drivable_area_dict in log_map_archive["drivable_areas"].items(): # keys: ["area_boundary", "id"] - drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict["area_boundary"]) + drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict["area_boundary"], close=True) for lane_segment_id, lane_segment_dict in log_map_archive["lane_segments"].items(): # keys = [ @@ -56,14 +63,22 @@ def _extract_polyline(data: List[Dict[str, float]]) -> Polyline3D: lane_segment_dict["left_lane_boundary"] = _extract_polyline(lane_segment_dict["left_lane_boundary"]) lane_segment_dict["right_lane_boundary"] = _extract_polyline(lane_segment_dict["right_lane_boundary"]) + for crosswalk_id, crosswalk_dict in log_map_archive["pedestrian_crossings"].items(): + # keys = ["id", "outline"] + # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/map/pedestrian_crossing.py + + p1, p2 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge1"]], dtype=np.float64) + p3, p4 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge2"]], dtype=np.float64) + crosswalk_dict["outline"] = Polyline3D.from_array(np.array([p1, p2, p4, p3, p1], dtype=np.float64)) + lane_df = get_lane_df(log_map_archive["lane_segments"]) + lane_group_df = get_lane_group_df(log_map_archive["lane_segments"]) get_empty_gdf() - get_empty_gdf() - get_empty_gdf() + crosswalk_df = get_crosswalk_df(log_map_archive["pedestrian_crossings"]) get_empty_gdf() get_empty_gdf() generic_drivable_df = get_generic_drivable_df(drivable_areas) - get_empty_gdf() + road_edge_df = get_road_edge_df(generic_drivable_df) get_empty_gdf() map_file_path.unlink(missing_ok=True) @@ -71,19 +86,28 @@ def _extract_polyline(data: List[Dict[str, float]]) -> Polyline3D: map_file_path.parent.mkdir(parents=True, exist_ok=True) lane_df.to_file(map_file_path, layer=MapLayer.LANE.serialize(), driver="GPKG") - # lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") + lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") # intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") - # crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") + crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") # walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") # carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") generic_drivable_df.to_file(map_file_path, layer=MapLayer.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a") - # road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") + road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") # road_line_df.to_file(map_file_path, layer=MapLayer.ROAD_LINE.serialize(), driver="GPKG", mode="a") +def get_empty_gdf() -> gpd.GeoDataFrame: + ids = [] + outlines = [] + geometries = [] + data = pd.DataFrame({"id": ids, "outline": outlines}) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + return gdf + + def get_lane_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: - ids = list(lanes.keys()) + ids = [int(lane_id) for lane_id in lanes.keys()] lane_types = [0] * len(ids) # TODO: Add lane types lane_group_ids = [] speed_limits_mps = [] @@ -122,7 +146,8 @@ def _get_centerline_from_boundaries( # "left_neighbor_id", # ] lane_centerline = _get_centerline_from_boundaries( - left_boundary=lane_dict["left_lane_boundary"], right_boundary=lane_dict["right_lane_boundary"] + left_boundary=lane_dict["left_lane_boundary"], + right_boundary=lane_dict["right_lane_boundary"], ) lane_speed_limit_mps = None @@ -168,23 +193,12 @@ def _get_centerline_from_boundaries( return gdf -def get_empty_gdf() -> gpd.GeoDataFrame: - ids = [] - geometries = [] - data = pd.DataFrame({"id": ids}) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf +def get_lane_group_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: + lane_group_sets = find_lane_groups(lanes) + lane_group_set_dict = {i: lane_group for i, lane_group in enumerate(lane_group_sets)} -def get_lane_group_df( - lanes: Dict[int, npt.NDArray[np.float64]], - lanes_successors: Dict[int, List[int]], - lanes_predecessors: Dict[int, List[int]], - lanes_left_boundaries_3d: Dict[int, Polyline3D], - lanes_right_boundaries_3d: Dict[int, Polyline3D], -) -> gpd.GeoDataFrame: - - ids = [] + ids = list(lane_group_set_dict.keys()) lane_ids = [] intersection_ids = [] predecessor_lane_group_ids = [] @@ -193,22 +207,38 @@ def get_lane_group_df( right_boundaries = [] geometries = [] - # NOTE: WOPD does not provide lane groups, so we create a lane group for each lane. - for lane_id in lanes.keys(): - if lane_id not in lanes_left_boundaries_3d or lane_id not in lanes_right_boundaries_3d: - continue - ids.append(lane_id) - lane_ids.append([lane_id]) - intersection_ids.append(None) # WOPD does not provide intersections - predecessor_lane_group_ids.append(lanes_predecessors[lane_id]) - successor_lane_group_ids.append(lanes_successors[lane_id]) - left_boundaries.append(lanes_left_boundaries_3d[lane_id].linestring) - right_boundaries.append(lanes_right_boundaries_3d[lane_id].linestring) + def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: + """Helper to find lane group ids that contain any of the given lane ids.""" + lane_group_ids_ = [] + for lane_group_id_, lane_group_set_ in lane_group_set_dict.items(): + if any(str(lane_id) in lane_group_set_ for lane_id in lane_ids): + lane_group_ids_.append(lane_group_id_) + return list(set(lane_group_ids_)) + + for lane_group_id, lane_group_set in lane_group_set_dict.items(): + + lane_ids.append([int(lane_id) for lane_id in lane_group_set]) + intersection_ids.append(None) # NOTE: AV2 doesn't have explicit intersection objects. + + successor_lanes = [] + predecessor_lanes = [] + for lane_id in lane_group_set: + lane_dict = lanes[str(lane_id)] + successor_lanes.extend(lane_dict["successors"]) + predecessor_lanes.extend(lane_dict["predecessors"]) + + left_boundary = lanes[lane_group_set[0]]["left_lane_boundary"] + right_boundary = lanes[lane_group_set[-1]]["right_lane_boundary"] + + predecessor_lane_group_ids.append(_get_lane_group_ids_of_lanes_ids(predecessor_lanes)) + successor_lane_group_ids.append(_get_lane_group_ids_of_lanes_ids(successor_lanes)) + left_boundaries.append(left_boundary.linestring) + right_boundaries.append(right_boundary.linestring) geometry = geom.Polygon( np.vstack( [ - lanes_left_boundaries_3d[lane_id].array[:, :2], - lanes_right_boundaries_3d[lane_id].array[:, :2][::-1], + left_boundary.array[:, :2], + right_boundary.array[:, :2][::-1], ] ) ) @@ -219,8 +249,8 @@ def get_lane_group_df( "id": ids, "lane_ids": lane_ids, "intersection_id": intersection_ids, - "predecessor_lane_group_ids": predecessor_lane_group_ids, - "successor_lane_group_ids": successor_lane_group_ids, + "predecessor_ids": predecessor_lane_group_ids, + "successor_ids": successor_lane_group_ids, "left_boundary": left_boundaries, "right_boundary": right_boundaries, } @@ -262,8 +292,12 @@ def get_walkway_df() -> gpd.GeoDataFrame: def get_crosswalk_df(crosswalks: Dict[int, npt.NDArray[np.float64]]) -> gpd.GeoDataFrame: ids = list(crosswalks.keys()) - outlines = [geom.LineString(outline) for outline in crosswalks.values()] - geometries = [geom.Polygon(outline[..., Point3DIndex.XY]) for outline in crosswalks.values()] + outlines = [] + geometries = [] + for crosswalk_dict in crosswalks.values(): + outline = crosswalk_dict["outline"] + outlines.append(outline.linestring) + geometries.append(geom.Polygon(outline.array[:, Point3DIndex.XY])) data = pd.DataFrame({"id": ids, "outline": outlines}) gdf = gpd.GeoDataFrame(data, geometry=geometries) @@ -280,20 +314,15 @@ def get_generic_drivable_df(drivable_areas: Dict[int, Polyline3D]) -> gpd.GeoDat return gdf -def get_road_edge_df( - road_edges: Dict[int, npt.NDArray[np.float64]], road_edges_type: Dict[int, RoadEdgeType] -) -> gpd.GeoDataFrame: - ids = list(road_edges.keys()) - geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_edges.values()] +def get_road_edge_df(generic_drivable_area_df: gpd.GeoDataFrame) -> gpd.GeoDataFrame: + road_edges = get_road_edges_3d_from_generic_drivable_area_df(generic_drivable_area_df) + road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) - data = pd.DataFrame( - { - "id": ids, - "road_edge_type": [int(road_edge_type) for road_edge_type in road_edges_type.values()], - } - ) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - return gdf + ids = np.arange(len(road_edges), dtype=np.int64).tolist() + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types = [int(RoadEdgeType.ROAD_EDGE_BOUNDARY)] * len(road_edges) + geometries = road_edges + return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) def get_road_line_df( @@ -310,3 +339,55 @@ def get_road_line_df( ) gdf = gpd.GeoDataFrame(data, geometry=geometries) return gdf + + +def find_lane_groups(lanes) -> List[List[str]]: + + visited = set() + lane_groups = [] + + def _get_valid_neighbor_id(lane_data, direction): + """Helper function to safely get neighbor ID""" + neighbor_key = f"{direction}_neighbor_id" + neighbor_id = str(lane_data.get(neighbor_key)) + mark_type = lane_data.get(f"{direction}_lane_mark_type", None) + + if (neighbor_id is not None) and (neighbor_id in lanes) and (mark_type in LANE_GROUP_MARK_TYPES): + return neighbor_id + return None + + def _traverse_group(start_lane_id): + """ + Traverse left and right from a starting lane to find all connected parallel lanes + """ + group = [start_lane_id] + queue = [start_lane_id] + + while queue: + current_id = queue.pop(0) + if current_id in visited: + continue + + visited.add(current_id) + + # Check left neighbor + left_neighbor = _get_valid_neighbor_id(lanes[current_id], "left") + if left_neighbor is not None and left_neighbor not in visited: + queue.append(left_neighbor) + group = [left_neighbor] + group + + # Check right neighbor + right_neighbor = _get_valid_neighbor_id(lanes[current_id], "right") + if right_neighbor is not None and right_neighbor not in visited: + queue.append(right_neighbor) + group = group + [right_neighbor] + + return group + + # Find all lane groups + for lane_id in lanes: + if lane_id not in visited: + group = _traverse_group(lane_id) + lane_groups.append(group) + + return lane_groups diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py index 59c791d7..23a4d702 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_data_converter.py @@ -28,7 +28,7 @@ from d123.common.datatypes.time.time_point import TimePoint from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index from d123.common.datatypes.vehicle_state.vehicle_parameters import ( - get_nuplan_pacifica_parameters, + get_nuplan_chrysler_pacifica_parameters, rear_axle_se3_to_center_se3, ) from d123.common.geometry.base import StateSE3 @@ -200,7 +200,7 @@ def convert_nuplan_log_to_arrow( timestep_seconds=TARGET_DT, map_has_z=False, ) - vehicle_parameters = get_nuplan_pacifica_parameters() + vehicle_parameters = get_nuplan_chrysler_pacifica_parameters() camera_metadata = get_nuplan_camera_metadata(log_path) lidar_metadata = get_nuplan_lidar_metadata(log_db) @@ -385,7 +385,7 @@ def _extract_detections(lidar_pc: LidarPc) -> Tuple[List[List[float]], List[List def _extract_ego_state(lidar_pc: LidarPc) -> List[float]: yaw, pitch, roll = lidar_pc.ego_pose.quaternion.yaw_pitch_roll - vehicle_parameters = get_nuplan_pacifica_parameters() + vehicle_parameters = get_nuplan_chrysler_pacifica_parameters() # vehicle_parameters = get_pacifica_parameters() rear_axle_pose = StateSE3( diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index dd6512fe..1295b2b0 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -20,7 +20,7 @@ from d123.common.datatypes.sensor.lidar import LiDARMetadata, LiDARType, lidar_metadata_dict_to_json from d123.common.datatypes.sensor.lidar_index import WopdLidarIndex from d123.common.datatypes.vehicle_state.ego_state import DynamicStateSE3, EgoStateSE3, EgoStateSE3Index -from d123.common.datatypes.vehicle_state.vehicle_parameters import get_wopd_pacifica_parameters +from d123.common.datatypes.vehicle_state.vehicle_parameters import get_wopd_chrysler_pacifica_parameters from d123.common.geometry.base import Point3D, StateSE3 from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE3Index from d123.common.geometry.constants import DEFAULT_PITCH, DEFAULT_ROLL @@ -217,7 +217,7 @@ def convert_wopd_tfrecord_log_to_arrow( timestep_seconds=TARGET_DT, # TODO: Check if correct. Maybe not hardcode map_has_z=True, ) - vehicle_parameters = get_wopd_pacifica_parameters() + vehicle_parameters = get_wopd_chrysler_pacifica_parameters() camera_metadata = get_wopd_camera_metadata(initial_frame, data_converter_config) lidar_metadata = get_wopd_lidar_metadata(initial_frame, data_converter_config) @@ -454,7 +454,7 @@ def _extract_detections(frame: dataset_pb2.Frame) -> Tuple[List[List[float]], Li def _extract_ego_state(frame: dataset_pb2.Frame) -> List[float]: rear_axle_pose = _get_ego_pose_se3(frame) - vehicle_parameters = get_wopd_pacifica_parameters() + vehicle_parameters = get_wopd_chrysler_pacifica_parameters() # FIXME: Find dynamic state in waymo open perception dataset # https://github.com/waymo-research/waymo-open-dataset/issues/55#issuecomment-546152290 dynamic_state = DynamicStateSE3( diff --git a/d123/dataset/maps/gpkg/gpkg_map_objects.py b/d123/dataset/maps/gpkg/gpkg_map_objects.py index 0b6164f0..1fa6de6a 100644 --- a/d123/dataset/maps/gpkg/gpkg_map_objects.py +++ b/d123/dataset/maps/gpkg/gpkg_map_objects.py @@ -72,11 +72,21 @@ def trimesh_mesh(self) -> trimesh.Trimesh: left_boundary = Polyline3D.from_linestring(self._object_row.left_boundary) right_boundary = Polyline3D.from_linestring(self._object_row.right_boundary) trimesh_mesh = get_trimesh_from_boundaries(left_boundary, right_boundary) - elif "geometry" in self._object_df.columns: + else: # Fallback to geometry if no boundaries are available outline_3d_array = self.outline_3d.array - _, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) - trimesh_mesh = trimesh.Trimesh(vertices=outline_3d_array, faces=faces) + vertices_2d, faces = trimesh.creation.triangulate_polygon( + geom.Polygon(outline_3d_array[:, Point3DIndex.XY]) + ) + if len(vertices_2d) == len(outline_3d_array): + # Regular case, where vertices match outline_3d_array + vertices_3d = outline_3d_array + elif len(vertices_2d) == len(outline_3d_array) + 1: + # outline array was not closed, so we need to add the first vertex again + vertices_3d = np.vstack((outline_3d_array, outline_3d_array[0])) + else: + raise ValueError("No vertices found for triangulation.") + trimesh_mesh = trimesh.Trimesh(vertices=vertices_3d, faces=faces) return trimesh_mesh @@ -200,13 +210,19 @@ def __init__( def successors(self) -> List[GPKGLaneGroup]: """Inherited, see superclass.""" successor_ids = ast.literal_eval(self._object_row.successor_ids) - return [GPKGLaneGroup(lane_group_id, self._object_df) for lane_group_id in successor_ids] + return [ + GPKGLaneGroup(lane_group_id, self._object_df, self._lane_df, self._intersection_df) + for lane_group_id in successor_ids + ] @property def predecessors(self) -> List[GPKGLaneGroup]: """Inherited, see superclass.""" predecessor_ids = ast.literal_eval(self._object_row.predecessor_ids) - return [GPKGLaneGroup(lane_group_id, self._object_df) for lane_group_id in predecessor_ids] + return [ + GPKGLaneGroup(lane_group_id, self._object_df, self._lane_df, self._intersection_df) + for lane_group_id in predecessor_ids + ] @property def left_boundary(self) -> Polyline3D: @@ -300,16 +316,6 @@ class GPKGGenericDrivable(GPKGSurfaceObject, AbstractGenericDrivable): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): super().__init__(object_id, object_df) - @property - def trimesh_mesh(self) -> trimesh.Trimesh: - """Inherited, see superclass.""" - - # Fallback to geometry if no boundaries are available - outline_3d_array = self.outline_3d.array - _, faces = trimesh.creation.triangulate_polygon(geom.Polygon(outline_3d_array[:, Point3DIndex.XY])) - trimesh_mesh = trimesh.Trimesh(vertices=outline_3d_array, faces=faces) - return trimesh_mesh - class GPKGRoadEdge(GPKGLineObject, AbstractRoadEdge): def __init__(self, object_id: str, object_df: gpd.GeoDataFrame): diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index 55a61486..ab29cbf5 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -51,6 +51,7 @@ class RoadLineType(SerialIntEnum): """ Enum for RoadLineType. NOTE: We use the road line types from Waymo. + TODO: Use the Argoverse 2 road line types instead. https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L208 """ diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 36a7800e..c9e58b07 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -75,7 +75,7 @@ ")\n", "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", "from d123.dataset.maps.abstract_map import AbstractMap\n", - "from d123.dataset.maps.abstract_map_objects import AbstractLane\n", + "from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractLaneGroup\n", "from d123.dataset.maps.map_datatypes import MapLayer\n", "from d123.dataset.scene.abstract_scene import AbstractScene\n", "\n", @@ -143,13 +143,13 @@ ") -> None:\n", " layers: List[MapLayer] = [\n", " MapLayer.LANE,\n", - " # MapLayer.LANE_GROUP,\n", + " MapLayer.LANE_GROUP,\n", " MapLayer.GENERIC_DRIVABLE,\n", " # MapLayer.CARPARK,\n", - " # MapLayer.CROSSWALK,\n", + " MapLayer.CROSSWALK,\n", " # MapLayer.INTERSECTION,\n", " # MapLayer.WALKWAY,\n", - " # MapLayer.ROAD_EDGE,\n", + " MapLayer.ROAD_EDGE,\n", " # MapLayer.ROAD_LINE,\n", " ]\n", " x_min, x_max = point_2d.x - radius, point_2d.x + radius\n", @@ -171,21 +171,25 @@ " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", "\n", " if layer in [MapLayer.LANE_GROUP]:\n", + " map_object: AbstractLaneGroup\n", + "\n", + " successor_ids = [successor.id for successor in map_object.successors]\n", + "\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", + " # centroid = map_object.shapely_polygon.centroid\n", + " # ax.text(\n", + " # centroid.x,\n", + " # centroid.y,\n", + " # f\"{map_object.id} -> {successor_ids}\",\n", + " # horizontalalignment=\"center\",\n", + " # verticalalignment=\"center\",\n", + " # fontsize=8,\n", + " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", + " # )\n", "\n", " if layer in [MapLayer.LANE]:\n", " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", "\n", - " # centroid = map_object.shapely_polygon.centroid\n", - " # ax.text(\n", - " # centroid.x,\n", - " # centroid.y,\n", - " # str(map_object.id),\n", - " # horizontalalignment=\"center\",\n", - " # verticalalignment=\"center\",\n", - " # fontsize=8,\n", - " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\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", @@ -236,17 +240,26 @@ " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", ") -> Tuple[plt.Figure, plt.Axes]:\n", "\n", - " size = 20\n", + " size = 10\n", "\n", " fig, ax = plt.subplots(figsize=(size, size))\n", " _plot_scene_on_ax(ax, scene, iteration, radius)\n", " return fig, ax\n", "\n", "\n", - "scene_index = 2\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=10, radius=100)\n", + "scene_index = 3\n", + "iteration = 1\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=iteration, radius=100)\n", + "plt.show()\n", + "\n", + "# camera = scenes[scene_index].get_camera_at_iteration(\n", + "# iteration=iteration, camera_type=CameraType.CAM_F0\n", + "# )\n", + "\n", + "# plt.imshow(camera.image, cmap=\"gray\", vmin=0, vmax=255)\n", + "# # fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")\n", "\n", - "# fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")" + "# scenes[scene_index].log_name" ] }, { @@ -327,19 +340,6 @@ "metadata": {}, "outputs": [], "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "8", - "metadata": {}, - "outputs": [], - "source": [ - "from d123.common.datatypes.sensor.lidar_index import NuplanLidarIndex\n", - "\n", - "\n", - "NuplanLidarIndex.Y" - ] } ], "metadata": { diff --git a/notebooks/viz/camera_matplotlib.ipynb b/notebooks/viz/camera_matplotlib.ipynb index d9e21ec3..bbd6221d 100644 --- a/notebooks/viz/camera_matplotlib.ipynb +++ b/notebooks/viz/camera_matplotlib.ipynb @@ -375,7 +375,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123_dev", + "display_name": "d123", "language": "python", "name": "python3" }, From 957e82da56ae1745a4798988e8d8df7ce951ee0d Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Wed, 20 Aug 2025 15:50:08 +0200 Subject: [PATCH 38/42] Add some opendrive comments. --- .../conversion/map/opendrive/conversion/group_collections.py | 1 + d123/dataset/conversion/map/opendrive/elements/reference.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/d123/dataset/conversion/map/opendrive/conversion/group_collections.py b/d123/dataset/conversion/map/opendrive/conversion/group_collections.py index 1e978693..6e577305 100644 --- a/d123/dataset/conversion/map/opendrive/conversion/group_collections.py +++ b/d123/dataset/conversion/map/opendrive/conversion/group_collections.py @@ -236,6 +236,7 @@ def _create_outer_lane_border( lane: Lane, coeff_factor: float, ) -> Border: + # NOTE: This function needs to be re-written because of licensing issues. args = {} if len(lane_borders) == 1: diff --git a/d123/dataset/conversion/map/opendrive/elements/reference.py b/d123/dataset/conversion/map/opendrive/elements/reference.py index cb521d4e..30d18bb3 100644 --- a/d123/dataset/conversion/map/opendrive/elements/reference.py +++ b/d123/dataset/conversion/map/opendrive/elements/reference.py @@ -15,6 +15,8 @@ from d123.dataset.conversion.map.opendrive.elements.geometry import Arc, Geometry, Line, Spiral from d123.dataset.conversion.map.opendrive.elements.lane import LaneOffset +# NOTE: This file needs to be re-written because of licensing issues. + @dataclass class PlanView: From e49f0d15bdafe965b95648483abd244d05a8bcd3 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Thu, 21 Aug 2025 21:06:26 +0200 Subject: [PATCH 39/42] Update OpenDrive conversion --- .flake8 | 3 + .../visualization/matplotlib/observation.py | 35 + d123/common/visualization/viser/server.py | 7 +- d123/common/visualization/viser/utils.py | 12 +- .../map/opendrive/elements/enums.py | 45 -- .../map/opendrive/elements/reference.py | 211 ------ .../map/opendrive/opendrive_converter.py | 664 ------------------ .../map/opendrive/opendrive_map_conversion.py | 422 +++++++++++ .../map/opendrive/parser/__init__.py | 0 .../{elements => parser}/elevation.py | 70 +- .../{elements => parser}/geometry.py | 68 +- .../opendrive/{elements => parser}/lane.py | 65 +- .../opendrive/{elements => parser}/objects.py | 7 + .../{elements => parser}/opendrive.py | 35 +- .../map/opendrive/parser/polynomial.py | 32 + .../map/opendrive/parser/reference.py | 153 ++++ .../opendrive/{elements => parser}/road.py | 29 +- .../map/opendrive/utils/collection.py | 325 +++++++++ .../map/opendrive/{ => utils}/id_mapping.py | 0 .../{conversion => utils}/id_system.py | 0 .../lane_helper.py} | 108 ++- .../objects_helper.py} | 10 +- .../carla/carla_data_converter.py | 15 +- .../test_opendrive_conversion.ipynb | 30 +- notebooks/viz/bev_matplotlib.ipynb | 15 +- notebooks/viz/bev_matplotlib_prediction.ipynb | 246 +++++++ notebooks/viz/viser_testing_v2_scene.ipynb | 14 +- 27 files changed, 1479 insertions(+), 1142 deletions(-) delete mode 100644 d123/dataset/conversion/map/opendrive/elements/enums.py delete mode 100644 d123/dataset/conversion/map/opendrive/elements/reference.py create mode 100644 d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py create mode 100644 d123/dataset/conversion/map/opendrive/parser/__init__.py rename d123/dataset/conversion/map/opendrive/{elements => parser}/elevation.py (56%) rename d123/dataset/conversion/map/opendrive/{elements => parser}/geometry.py (76%) rename d123/dataset/conversion/map/opendrive/{elements => parser}/lane.py (77%) rename d123/dataset/conversion/map/opendrive/{elements => parser}/objects.py (85%) rename d123/dataset/conversion/map/opendrive/{elements => parser}/opendrive.py (79%) create mode 100644 d123/dataset/conversion/map/opendrive/parser/polynomial.py create mode 100644 d123/dataset/conversion/map/opendrive/parser/reference.py rename d123/dataset/conversion/map/opendrive/{elements => parser}/road.py (77%) create mode 100644 d123/dataset/conversion/map/opendrive/utils/collection.py rename d123/dataset/conversion/map/opendrive/{ => utils}/id_mapping.py (100%) rename d123/dataset/conversion/map/opendrive/{conversion => utils}/id_system.py (100%) rename d123/dataset/conversion/map/opendrive/{conversion/group_collections.py => utils/lane_helper.py} (70%) rename d123/dataset/conversion/map/opendrive/{conversion/objects_collections.py => utils/objects_helper.py} (85%) create mode 100644 notebooks/viz/bev_matplotlib_prediction.ipynb diff --git a/.flake8 b/.flake8 index f2b974ae..c77f18b6 100644 --- a/.flake8 +++ b/.flake8 @@ -16,4 +16,7 @@ ignore = E731 # Local variable name is assigned to but never used F841 +per-file-ignores = + # imported but unused + __init__.py: F401 max-line-length = 120 diff --git a/d123/common/visualization/matplotlib/observation.py b/d123/common/visualization/matplotlib/observation.py index 0d46ca93..7e225cc7 100644 --- a/d123/common/visualization/matplotlib/observation.py +++ b/d123/common/visualization/matplotlib/observation.py @@ -5,6 +5,7 @@ import shapely.geometry as geom from d123.common.datatypes.detection.detection import BoxDetectionWrapper, TrafficLightDetectionWrapper +from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.vehicle_state.ego_state import EgoStateSE2, EgoStateSE3 from d123.common.geometry.base import Point2D from d123.common.geometry.bounding_box.bounding_box import BoundingBoxSE2, BoundingBoxSE3 @@ -27,6 +28,7 @@ from d123.dataset.maps.abstract_map import AbstractMap from d123.dataset.maps.abstract_map_objects import AbstractLane from d123.dataset.maps.map_datatypes import MapLayer +from d123.dataset.scene.abstract_scene import AbstractScene def add_default_map_on_ax( @@ -87,6 +89,39 @@ def add_box_detections_to_ax(ax: plt.Axes, box_detections: BoxDetectionWrapper) add_bounding_box_to_ax(ax, box_detection.bounding_box, plot_config) +def add_box_future_detections_to_ax(ax: plt.Axes, scene: AbstractScene, iteration: int) -> None: + + # TODO: Refactor this function + initial_agents = scene.get_box_detections_at_iteration(iteration) + agents_poses = { + agent.metadata.track_token: [agent.center_se3] + for agent in initial_agents + if agent.metadata.detection_type == DetectionType.VEHICLE + } + frequency = 1 + for iteration in range(iteration + frequency, scene.get_number_of_iterations(), frequency): + agents = scene.get_box_detections_at_iteration(iteration) + for agent in agents: + if agent.metadata.track_token in agents_poses: + agents_poses[agent.metadata.track_token].append(agent.center_se3) + + for track_token, poses in agents_poses.items(): + if len(poses) < 2: + continue + poses = np.array([pose.point_2d.array for pose in poses]) + num_poses = poses.shape[0] + alphas = 1 - np.linspace(0.2, 1.0, num_poses) # Start low, end high + for i in range(num_poses - 1): + ax.plot( + poses[i : i + 2, 0], + poses[i : i + 2, 1], + color=BOX_DETECTION_CONFIG[DetectionType.VEHICLE].fill_color.hex, + alpha=alphas[i + 1], + linewidth=BOX_DETECTION_CONFIG[DetectionType.VEHICLE].line_width * 5, + zorder=BOX_DETECTION_CONFIG[DetectionType.VEHICLE].zorder, + ) + + def add_ego_vehicle_to_ax(ax: plt.Axes, ego_vehicle_state: Union[EgoStateSE3, EgoStateSE2]) -> None: add_bounding_box_to_ax(ax, ego_vehicle_state.bounding_box, EGO_VEHICLE_CONFIG) diff --git a/d123/common/visualization/viser/server.py b/d123/common/visualization/viser/server.py index cdca86c4..087a7e8b 100644 --- a/d123/common/visualization/viser/server.py +++ b/d123/common/visualization/viser/server.py @@ -33,7 +33,7 @@ LINE_WIDTH: float = 4.0 # Bounding box config: -BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "mesh" +BOUNDING_BOX_TYPE: Literal["mesh", "lines"] = "lines" # Map config: MAP_AVAILABLE: bool = True @@ -182,6 +182,9 @@ def _(_) -> None: else: raise ValueError(f"Unknown bounding box type: {BOUNDING_BOX_TYPE}") + current_frame_handle.remove() + current_frame_handle = mew_frame_handle + for camera_type in VISUALIZE_CAMERA_GUI: if camera_type in scene.available_camera_types: camera_gui_handles[camera_type].image = scene.get_camera_at_iteration( @@ -208,8 +211,6 @@ def _(_) -> None: rendering_time = time.time() - start sleep_time = 1.0 / gui_framerate.value - rendering_time time.sleep(max(sleep_time, 0.0)) - current_frame_handle.remove() - current_frame_handle = mew_frame_handle self.server.flush() # Optional! # Load in frames. diff --git a/d123/common/visualization/viser/utils.py b/d123/common/visualization/viser/utils.py index 43345c0c..561c3676 100644 --- a/d123/common/visualization/viser/utils.py +++ b/d123/common/visualization/viser/utils.py @@ -97,14 +97,15 @@ def get_map_meshes(scene: AbstractScene): center = initial_ego_vehicle_state.center_se3 map_layers = [ MapLayer.LANE_GROUP, - # MapLayer.LANE, - # MapLayer.WALKWAY, + MapLayer.LANE, + MapLayer.WALKWAY, MapLayer.CROSSWALK, - # MapLayer.CARPARK, + MapLayer.CARPARK, MapLayer.GENERIC_DRIVABLE, ] map_objects_dict = scene.map_api.get_proximal_map_objects(center.point_2d, radius=MAP_RADIUS, layers=map_layers) + print(map_objects_dict.keys()) output = {} for map_layer in map_objects_dict.keys(): @@ -115,13 +116,13 @@ def get_map_meshes(scene: AbstractScene): if map_layer in [ MapLayer.WALKWAY, MapLayer.CROSSWALK, - # MapLayer.GENERIC_DRIVABLE, + MapLayer.GENERIC_DRIVABLE, MapLayer.CARPARK, ]: # Push meshes up by a few centimeters to avoid overlap with the ground in the visualization. trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z - 0.1).array else: - trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z + 0.1).array + trimesh_mesh.vertices -= Point3D(x=center.x, y=center.y, z=center.z).array if not scene.log_metadata.map_has_z: trimesh_mesh.vertices += Point3D( @@ -131,7 +132,6 @@ def get_map_meshes(scene: AbstractScene): trimesh_mesh = configure_trimesh(trimesh_mesh, MAP_SURFACE_CONFIG[map_layer].fill_color) surface_meshes.append(trimesh_mesh) output[f"{map_layer.serialize()}"] = trimesh.util.concatenate(surface_meshes) - return output diff --git a/d123/dataset/conversion/map/opendrive/elements/enums.py b/d123/dataset/conversion/map/opendrive/elements/enums.py deleted file mode 100644 index cd8ebe60..00000000 --- a/d123/dataset/conversion/map/opendrive/elements/enums.py +++ /dev/null @@ -1,45 +0,0 @@ -from enum import Enum, IntEnum - - -class NamedEnum(Enum): - def __init__(self, index, names): - self._index = index - self._names = names - - @property - def index(self): - return self._index - - @property - def names(self): - return self._names - - @classmethod - def from_name(cls, name): - for member in cls: - if name in member.names: - return member - raise ValueError(f"No enum member with name '{name}'") - - -class OpenDriveLaneType(IntEnum): - - BIKING = (0, ("biking",)) - BORDER = (1, ("border",)) - CONNECTING_RAMP = (2, ("connectingRamp",)) - CURB = (3, ("curb",)) - DRIVING = (4, ("driving",)) - ENTRY = (5, ("entry",)) - EXIT = (6, ("exit",)) - MEDIAN = (7, ("median",)) - NONE = (8, ("none",)) - OFF_RAMP = (9, ("offRamp",)) - ON_RAMP = (10, ("onRamp",)) - PARKING = (11, ("parking",)) - RAIL = (12, ("rail",)) - RESTRICTED = (13, ("restricted",)) - SHOULDER = (14, ("shoulder",)) - SIDEWALK = (15, ("sidewalk", "walking")) - SLIP_LANE = (16, ("slipLane",)) - STOP = (17, ("stop",)) - TRAM = (18, ("tram",)) diff --git a/d123/dataset/conversion/map/opendrive/elements/reference.py b/d123/dataset/conversion/map/opendrive/elements/reference.py deleted file mode 100644 index 30d18bb3..00000000 --- a/d123/dataset/conversion/map/opendrive/elements/reference.py +++ /dev/null @@ -1,211 +0,0 @@ -from __future__ import annotations - -import xml.etree.ElementTree as ET -from dataclasses import dataclass -from functools import cached_property -from typing import List, Optional, Union -from xml.etree.ElementTree import Element - -import numpy as np -import numpy.typing as npt - -from d123.common.geometry.base import Point3DIndex, StateSE2Index -from d123.common.geometry.utils import normalize_angle -from d123.dataset.conversion.map.opendrive.elements.elevation import ElevationProfile -from d123.dataset.conversion.map.opendrive.elements.geometry import Arc, Geometry, Line, Spiral -from d123.dataset.conversion.map.opendrive.elements.lane import LaneOffset - -# NOTE: This file needs to be re-written because of licensing issues. - - -@dataclass -class PlanView: - - geometries: List[Geometry] - - def __post_init__(self): - # NOTE: added assertion/filtering to check for element type or consistency - self.geometries.sort(key=lambda x: x.s, reverse=False) - - @classmethod - def parse(cls, plan_view_element: Optional[Element]) -> PlanView: - - args = {} - geometries: List[Geometry] = [] - for geometry_element in plan_view_element.findall("geometry"): - if geometry_element.find("line") is not None: - geometry = Line.parse(geometry_element) - elif geometry_element.find("arc") is not None: - geometry = Arc.parse(geometry_element) - elif geometry_element.find("spiral") is not None: - geometry = Spiral.parse(geometry_element) - else: - geometry_str = ET.tostring(geometry_element, encoding="unicode") - raise NotImplementedError(f"Geometry not implemented: {geometry_str}") - geometries.append(geometry) - args["geometries"] = geometries - return PlanView(**args) - - @cached_property - def geometry_lengths(self) -> npt.NDArray[np.float64]: - return np.cumsum([0.0] + [geo.length for geo in self.geometries], dtype=np.float64) - - @property - def length(self) -> float: - return float(self.geometry_lengths[-1]) - - def interpolate_se2(self, s: float, t: float = 0.0, is_last_pos: bool = False) -> npt.NDArray[np.float64]: - - try: - # get index of geometry which is at s_pos - mask = self.geometry_lengths > s - sub_idx = np.argmin(self.geometry_lengths[mask] - s) - geo_idx = np.arange(self.geometry_lengths.shape[0])[mask][sub_idx] - 1 - except ValueError: - # s_pos is after last geometry because of rounding error - if np.isclose(s, self.geometry_lengths[-1], 0.01, 0.01): # todo parameter - geo_idx = self.geometry_lengths.size - 2 - else: - raise Exception( - f"Tried to calculate a position outside of the borders of the reference path at s={s}" - f", but path has only length of l={self.length}" - ) - - return self.geometries[geo_idx].interpolate_se2(s - self.geometry_lengths[geo_idx], t) - - -@dataclass -class Border: - - reference: Union[Border, PlanView] - elevation_profile: ElevationProfile - - width_coefficient_offsets: List[float] - width_coefficients: List[List[float]] - - s_offset: float = 0.0 - - # NOTE: loaded in __post_init__ - length: Optional[float] = None - - def __post_init__(self): - # NOTE: added assertion/filtering to check for element type or consistency - self.length = float(self.reference.length) - - @classmethod - def from_plan_view( - cls, - plan_view: PlanView, - lane_offsets: List[LaneOffset], - elevation_profile: ElevationProfile, - ) -> Border: - args = {} - args["reference"] = plan_view - args["elevation_profile"] = elevation_profile - - width_coefficient_offsets = [] - width_coefficients = [] - - # Lane offsets will be coeffs - # this has to be done if the reference path has the laneoffset attribute - # and thus is different to the geometry described in the plan_view - # openDRIVE lets multiple laneOffsets start at the same position - # but only the last one counts -> delete all previous ones - if any(lane_offsets): - for lane_offset in lane_offsets: - if lane_offset.s in width_coefficient_offsets: - # offset is already there, delete previous entries - idx = width_coefficient_offsets.index(lane_offset.s) - del width_coefficient_offsets[idx] - del width_coefficients[idx] - width_coefficient_offsets.append(lane_offset.s) - width_coefficients.append(lane_offset.polynomial_coefficients) - else: - width_coefficient_offsets.append(0.0) - width_coefficients.append([0.0]) - - args["width_coefficient_offsets"] = width_coefficient_offsets - args["width_coefficients"] = width_coefficients - - return Border(**args) - - def _get_width_index(self, s: float, is_last_pos: bool) -> float: - """Get the index of the width which applies at position s_pos. - - :param s_pos: Position on border in curve_parameter ds - :param is_last_pos: Whether s_pos is the last position - :return: Index of width that applies at position s_pos - """ - - return next( - ( - self.width_coefficient_offsets.index(n) - for n in self.width_coefficient_offsets[::-1] - if ((n <= s and (not is_last_pos or s == 0)) or (n < s and is_last_pos)) - ), - len(self.width_coefficient_offsets) - 1, - ) - - def _get_height_index(self, s: float, is_last_pos: bool) -> float: - height_offsets = [elevation.s for elevation in self.elevation_profile.elevations] - - return next( - ( - height_offsets.index(n) - for n in height_offsets[::-1] - if ((n <= s and (not is_last_pos or s == 0)) or (n < s and is_last_pos)) - ), - len(height_offsets) - 1, - ) - - def interpolate_se2(self, s: float, t: float = 0.0, is_last_pos: bool = False) -> npt.NDArray[np.float64]: - # Last reference has to be a reference geometry (PlanView) - # Offset of all inner lanes (Border) - # calculate position of reference border - if np.isclose(s, 0): - s = 0 - - try: - se2 = self.reference.interpolate_se2(self.s_offset + s, is_last_pos=is_last_pos) - except TypeError: - se2 = self.reference.interpolate_se2(np.round(self.s_offset + s, 3), is_last_pos=is_last_pos) - - if len(self.width_coefficients) == 0 or len(self.width_coefficient_offsets) == 0: - raise Exception("No entries for width definitions.") - - # Find correct coefficients - # find which width segment is at s_pos - width_idx = self._get_width_index(s, is_last_pos=is_last_pos) - # width_idx = min(width_idx, len(self.width_coefficient_offsets)-1) - # Calculate width at s_pos - distance = ( - np.polynomial.polynomial.polyval( - s - self.width_coefficient_offsets[width_idx], - self.width_coefficients[width_idx], - ) - + t - ) - ortho = normalize_angle(se2[StateSE2Index.YAW] + np.pi / 2) - se2[StateSE2Index.X] += distance * np.cos(ortho) - se2[StateSE2Index.Y] += distance * np.sin(ortho) - se2[StateSE2Index.YAW] = normalize_angle(se2[StateSE2Index.YAW]) - - return se2 - - def interpolate_3d(self, s: float, t: float = 0.0, is_last_pos: bool = False) -> npt.NDArray[np.float64]: - - point_3d = np.zeros(len(Point3DIndex), dtype=np.float64) - - se2 = self.interpolate_se2(s, t, is_last_pos=is_last_pos) - point_3d[Point3DIndex.X] = se2[StateSE2Index.X] - point_3d[Point3DIndex.Y] = se2[StateSE2Index.Y] - - if len(self.elevation_profile.elevations) > 0: - height_index = self._get_height_index(s, is_last_pos=is_last_pos) - elevation = self.elevation_profile.elevations[height_index] - - point_3d[Point3DIndex.Z] = np.polynomial.polynomial.polyval( - s - elevation.s, elevation.polynomial_coefficients - ) - - return point_3d diff --git a/d123/dataset/conversion/map/opendrive/opendrive_converter.py b/d123/dataset/conversion/map/opendrive/opendrive_converter.py index b9dd58c2..e69de29b 100644 --- a/d123/dataset/conversion/map/opendrive/opendrive_converter.py +++ b/d123/dataset/conversion/map/opendrive/opendrive_converter.py @@ -1,664 +0,0 @@ -import os -import warnings -from pathlib import Path -from typing import Dict, Final, List, Optional - -import geopandas as gpd -import numpy as np -import pandas as pd -import shapely -from shapely.ops import polygonize, unary_union - -from d123.dataset.conversion.map.opendrive.conversion.group_collections import ( - OpenDriveLaneGroupHelper, - OpenDriveLaneHelper, - lane_section_to_lane_helpers, -) -from d123.dataset.conversion.map.opendrive.conversion.id_system import ( - build_lane_id, - derive_lane_section_id, - lane_group_id_from_lane_id, - road_id_from_lane_group_id, -) -from d123.dataset.conversion.map.opendrive.conversion.objects_collections import ( - OpenDriveObjectHelper, - get_object_helper, -) -from d123.dataset.conversion.map.opendrive.elements.opendrive import Junction, OpenDrive, Road -from d123.dataset.conversion.map.opendrive.elements.reference import Border -from d123.dataset.conversion.map.opendrive.id_mapping import IntIDMapping -from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length -from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import get_road_edges_3d_from_gdf -from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType - -ENABLE_WARNING: bool = False -CONNECTION_DISTANCE_THRESHOLD: float = 0.1 # [m] -D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) - -MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # [m] - - -class OpenDriveConverter: - def __init__(self, opendrive: OpenDrive): - - self.opendrive: OpenDrive = opendrive - - self.road_dict: Dict[int, Road] = {road.id: road for road in opendrive.roads} - self.junction_dict: Dict[int, Junction] = {junction.id: junction for junction in opendrive.junctions} - - # loaded during conversion - self.lane_helper_dict: Dict[str, OpenDriveLaneHelper] = {} - self.lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = {} - self.object_helper_dict: Dict[str, OpenDriveObjectHelper] = {} - - def run(self, map_name: str) -> None: - - # Run processing for map elements - self._collect_lane_helpers() - self._update_connection_from_links() - self._update_connection_from_junctions() - self._flip_and_set_connections() - self._post_process_connections() - self._collect_lane_groups() - self._collect_crosswalks() - - # Collect data frames and store - dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} - dataframes[MapLayer.LANE] = self._extract_lane_dataframe() - dataframes[MapLayer.WALKWAY] = self._extract_walkways_dataframe() - dataframes[MapLayer.CARPARK] = self._extract_carpark_dataframe() - dataframes[MapLayer.GENERIC_DRIVABLE] = self._extract_generic_drivable_dataframe() - dataframes[MapLayer.INTERSECTION] = self._extract_intersections_dataframe() - dataframes[MapLayer.LANE_GROUP] = self._extract_lane_group_dataframe() - dataframes[MapLayer.CROSSWALK] = self._extract_crosswalk_dataframe() - - self._convert_ids_to_int( - dataframes[MapLayer.LANE], - dataframes[MapLayer.WALKWAY], - dataframes[MapLayer.CARPARK], - dataframes[MapLayer.GENERIC_DRIVABLE], - dataframes[MapLayer.LANE_GROUP], - dataframes[MapLayer.INTERSECTION], - dataframes[MapLayer.CROSSWALK], - ) - dataframes[MapLayer.ROAD_EDGE] = self._extract_road_edge_df( - dataframes[MapLayer.LANE], - dataframes[MapLayer.CARPARK], - dataframes[MapLayer.GENERIC_DRIVABLE], - dataframes[MapLayer.LANE_GROUP], - ) - dataframes[MapLayer.ROAD_LINE] = self._extract_road_line_df( - dataframes[MapLayer.LANE], - dataframes[MapLayer.LANE_GROUP], - ) - - # Store dataframes - map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" - with warnings.catch_warnings(): - warnings.filterwarnings("ignore", message="'crs' was not provided") - for layer, gdf in dataframes.items(): - gdf.to_file(map_file_name, layer=layer.serialize(), driver="GPKG", mode="a") - - def _collect_lane_helpers(self) -> None: - for road in self.opendrive.roads: - reference_border = Border.from_plan_view(road.plan_view, road.lanes.lane_offsets, road.elevation_profile) - lane_section_lengths: List[float] = [ls.s for ls in road.lanes.lane_sections] + [road.length] - for idx, lane_section in enumerate(road.lanes.lane_sections): - lane_section_id = derive_lane_section_id(road.id, idx) - lane_helpers_ = lane_section_to_lane_helpers( - lane_section_id, - lane_section, - reference_border, - lane_section_lengths[idx], - lane_section_lengths[idx + 1], - road.road_types, - ) - self.lane_helper_dict.update(lane_helpers_) - - def _update_connection_from_links(self) -> None: - - for lane_id in self.lane_helper_dict.keys(): - road_idx, lane_section_idx, side, lane_idx = lane_id.split("_") - road_idx, lane_section_idx, lane_idx = int(road_idx), int(lane_section_idx), int(lane_idx) - - road = self.road_dict[road_idx] - - successor_lane_idx = self.lane_helper_dict[lane_id].open_drive_lane.successor - if successor_lane_idx is not None: - successor_lane_id: Optional[str] = None - - # Not the last lane section -> Next lane section in same road - if lane_section_idx < road.lanes.last_lane_section_idx: - successor_lane_id = build_lane_id( - road_idx, - lane_section_idx + 1, - successor_lane_idx, - ) - - # Last lane section -> Next road in first lane section - # Try to get next road - elif road.link.successor is not None and road.link.successor.element_type != "junction": - successor_road = self.road_dict[road.link.successor.element_id] - - if road.link.successor.contact_point == "start": - successor_lane_section_idx = 0 - else: - successor_lane_section_idx = successor_road.lanes.last_lane_section_idx - - successor_lane_id = build_lane_id( - successor_road.id, - successor_lane_section_idx, - successor_lane_idx, - ) - - # assert successor_lane_id in self.lane_helper_dict.keys() - if successor_lane_id is None or successor_lane_id not in self.lane_helper_dict.keys(): - continue - self.lane_helper_dict[lane_id].successor_lane_ids.append(successor_lane_id) - self.lane_helper_dict[successor_lane_id].predecessor_lane_ids.append(lane_id) - - predecessor_lane_idx = self.lane_helper_dict[lane_id].open_drive_lane.predecessor - if predecessor_lane_idx is not None: - predecessor_lane_id: Optional[str] = None - - # Not the first lane section -> Previous lane section in same road - if lane_section_idx > 0: - predecessor_lane_id = build_lane_id( - road_idx, - lane_section_idx - 1, - predecessor_lane_idx, - ) - - # First lane section -> Previous road - # Try to get previous road - elif road.link.predecessor is not None and road.link.predecessor.element_type != "junction": - predecessor_road = self.road_dict[road.link.predecessor.element_id] - - if road.link.predecessor.contact_point == "start": - predecessor_lane_section_idx = 0 - else: - predecessor_lane_section_idx = predecessor_road.lanes.last_lane_section_idx - - predecessor_lane_id = build_lane_id( - predecessor_road.id, - predecessor_lane_section_idx, - predecessor_lane_idx, - ) - - # assert predecessor_lane_id in self.lane_helper_dict.keys() - if predecessor_lane_id is None or predecessor_lane_id not in self.lane_helper_dict.keys(): - continue - self.lane_helper_dict[lane_id].predecessor_lane_ids.append(predecessor_lane_id) - self.lane_helper_dict[predecessor_lane_id].successor_lane_ids.append(lane_id) - - def _update_connection_from_junctions(self) -> None: - - # add junctions to link_index - # if contact_point is start, and laneId from connecting_road is negative - # the connecting_road is the successor - # if contact_point is start, and laneId from connecting_road is positive - # the connecting_road is the predecessor - # for contact_point == end it's exactly the other way - for junction_idx, junction in self.junction_dict.items(): - for connection in junction.connections: - incoming_road = self.road_dict[connection.incoming_road] - connecting_road = self.road_dict[connection.connecting_road] - contact_point = connection.contact_point - for lane_link in connection.lane_links: - - incoming_lane_id: Optional[str] = None - connecting_lane_id: Optional[str] = None - - if contact_point == "start": - if lane_link.start < 0: - incoming_lane_section_idx = incoming_road.lanes.last_lane_section_idx - else: - incoming_lane_section_idx = 0 - incoming_lane_id = build_lane_id(incoming_road.id, incoming_lane_section_idx, lane_link.start) - connecting_lane_id = build_lane_id(connecting_road.id, 0, lane_link.end) - else: - incoming_lane_id = build_lane_id(incoming_road.id, 0, lane_link.start) - connecting_lane_id = build_lane_id( - connecting_road.id, connecting_road.lanes.last_lane_section_idx, lane_link.end - ) - - if incoming_lane_id is None or connecting_lane_id is None: - continue - if ( - incoming_lane_id not in self.lane_helper_dict.keys() - or connecting_lane_id not in self.lane_helper_dict.keys() - ): - if ENABLE_WARNING: - warnings.warn( - f"Warning..... Lane connection {incoming_lane_id} -> {connecting_lane_id} not found in lane_helper_dict" - ) - continue - # assert incoming_lane_id in self.lane_helper_dict.keys() - # assert connecting_lane_id in self.lane_helper_dict.keys() - self.lane_helper_dict[incoming_lane_id].successor_lane_ids.append(connecting_lane_id) - self.lane_helper_dict[connecting_lane_id].predecessor_lane_ids.append(incoming_lane_id) - - def _flip_and_set_connections(self) -> None: - - for lane_id in self.lane_helper_dict.keys(): - if self.lane_helper_dict[lane_id].id > 0: - successors_temp = self.lane_helper_dict[lane_id].successor_lane_ids - self.lane_helper_dict[lane_id].successor_lane_ids = self.lane_helper_dict[lane_id].predecessor_lane_ids - self.lane_helper_dict[lane_id].predecessor_lane_ids = successors_temp - self.lane_helper_dict[lane_id].successor_lane_ids = list( - set(self.lane_helper_dict[lane_id].successor_lane_ids) - ) - self.lane_helper_dict[lane_id].predecessor_lane_ids = list( - set(self.lane_helper_dict[lane_id].predecessor_lane_ids) - ) - - def _post_process_connections(self) -> None: - - for lane_id in self.lane_helper_dict.keys(): - self.lane_helper_dict[lane_id] - - centerline = self.lane_helper_dict[lane_id].center_polyline_se2 - - valid_successor_lane_ids: List[str] = [] - for successor_lane_id in self.lane_helper_dict[lane_id].successor_lane_ids: - successor_centerline = self.lane_helper_dict[successor_lane_id].center_polyline_se2 - distance = np.linalg.norm(centerline[-1, :2] - successor_centerline[0, :2]) - if distance > CONNECTION_DISTANCE_THRESHOLD: - if ENABLE_WARNING: - warnings.warn( - f"Warning..... Removing connection {lane_id} -> {successor_lane_id} with distance {distance}" - ) - else: - valid_successor_lane_ids.append(successor_lane_id) - self.lane_helper_dict[lane_id].successor_lane_ids = valid_successor_lane_ids - - valid_predecessor_lane_ids: List[str] = [] - for predecessor_lane_id in self.lane_helper_dict[lane_id].predecessor_lane_ids: - predecessor_centerline = self.lane_helper_dict[predecessor_lane_id].center_polyline_se2 - distance = np.linalg.norm(centerline[0, :2] - predecessor_centerline[-1, :2]) - if distance > CONNECTION_DISTANCE_THRESHOLD: - if ENABLE_WARNING: - warnings.warn( - f"Warning..... Removing connection {predecessor_lane_id} -> {successor_lane_id} with distance {distance}" - ) - else: - valid_predecessor_lane_ids.append(predecessor_lane_id) - self.lane_helper_dict[lane_id].predecessor_lane_ids = valid_predecessor_lane_ids - - def _collect_lane_groups(self) -> None: - def _collect_lane_helper_of_id(lane_group_id: str) -> List[OpenDriveLaneHelper]: - lane_helpers: List[OpenDriveLaneHelper] = [] - for lane_id, lane_helper in self.lane_helper_dict.items(): - if (lane_helper.type in ["driving"]) and (lane_group_id_from_lane_id(lane_id) == lane_group_id): - lane_helpers.append(lane_helper) - return lane_helpers - - all_lane_group_ids = list( - set([lane_group_id_from_lane_id(lane_id) for lane_id in self.lane_helper_dict.keys()]) - ) - - for lane_group_id in all_lane_group_ids: - lane_group_lane_helper = _collect_lane_helper_of_id(lane_group_id) - if len(lane_group_lane_helper) >= 1: - self.lane_group_helper_dict[lane_group_id] = OpenDriveLaneGroupHelper( - lane_group_id, lane_group_lane_helper - ) - - def _collect_lane_group_ids_of_road(road_id: int) -> List[str]: - lane_group_ids: List[str] = [] - for lane_group_id in self.lane_group_helper_dict.keys(): - if int(road_id_from_lane_group_id(lane_group_id)) == road_id: - lane_group_ids.append(lane_group_id) - return lane_group_ids - - for junction in self.junction_dict.values(): - for connection in junction.connections: - connecting_road = self.road_dict[connection.connecting_road] - connecting_lane_group_ids = _collect_lane_group_ids_of_road(connecting_road.id) - for connecting_lane_group_id in connecting_lane_group_ids: - if connecting_lane_group_id in self.lane_group_helper_dict.keys(): - self.lane_group_helper_dict[connecting_lane_group_id].junction_id = junction.id - - def _collect_crosswalks(self) -> None: - for road in self.opendrive.roads: - if len(road.objects) == 0: - continue - reference_border = Border.from_plan_view(road.plan_view, road.lanes.lane_offsets, road.elevation_profile) - for object in road.objects: - if object.type in ["crosswalk"]: - object_helper = get_object_helper(object, reference_border) - self.object_helper_dict[object_helper.object_id] = object_helper - - def _extract_lane_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - lane_group_ids = [] - speed_limits_mps = [] - predecessor_ids = [] - successor_ids = [] - left_boundaries = [] - right_boundaries = [] - left_lane_ids = [] - right_lane_ids = [] - baseline_paths = [] - geometries = [] - - for lane_group_helper in self.lane_group_helper_dict.values(): - lane_group_id = lane_group_helper.lane_group_id - lane_helpers = lane_group_helper.lane_helpers - num_lanes = len(lane_helpers) - # NOTE: Lanes are going left to right - for lane_idx, lane_helper in enumerate(lane_helpers): - ids.append(lane_helper.lane_id) - lane_group_ids.append(lane_group_id) - speed_limits_mps.append(lane_helper.speed_limit_mps) - predecessor_ids.append(lane_helper.predecessor_lane_ids) - successor_ids.append(lane_helper.successor_lane_ids) - left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - baseline_paths.append(shapely.LineString(lane_helper.center_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - left_lane_id = lane_helpers[lane_idx - 1].lane_id if lane_idx > 0 else None - right_lane_id = lane_helpers[lane_idx + 1].lane_id if lane_idx < num_lanes - 1 else None - left_lane_ids.append(left_lane_id) - right_lane_ids.append(right_lane_id) - - data = pd.DataFrame( - { - "id": ids, - "lane_group_id": lane_group_ids, - "speed_limit_mps": speed_limits_mps, - "predecessor_ids": predecessor_ids, - "successor_ids": successor_ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - "left_lane_id": left_lane_ids, - "right_lane_id": right_lane_ids, - "baseline_path": baseline_paths, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_walkways_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - left_boundaries = [] - right_boundaries = [] - outlines = [] - geometries = [] - - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type == "sidewalk": - ids.append(lane_helper.lane_id) - left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - "left_boundary": left_boundaries, - "right_boundary": left_boundaries, - "outline": outlines, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_carpark_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - left_boundaries = [] - right_boundaries = [] - outlines = [] - geometries = [] - - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type == "parking": - ids.append(lane_helper.lane_id) - left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - "outline": outlines, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_generic_drivable_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - left_boundaries = [] - right_boundaries = [] - outlines = [] - geometries = [] - - for lane_helper in self.lane_helper_dict.values(): - if lane_helper.type in ["none", "border", "bidirectional"]: - ids.append(lane_helper.lane_id) - left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) - outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) - geometries.append(lane_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - "left_boundary": left_boundaries, - "right_boundary": left_boundaries, - "outline": outlines, - } - ) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_intersections_dataframe(self) -> gpd.GeoDataFrame: - def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriveLaneGroupHelper]: - return [ - lane_group_helper - for lane_group_helper in self.lane_group_helper_dict.values() - if lane_group_helper.junction_id == junction_id - ] - - ids = [] - lane_group_ids = [] - geometries = [] - for junction in self.junction_dict.values(): - lane_group_helpers = _find_lane_group_helpers_with_junction_id(junction.id) - lane_group_ids_ = [lane_group_helper.lane_group_id for lane_group_helper in lane_group_helpers] - if len(lane_group_ids_) == 0: - warnings.warn(f"Skipped Junction {junction.id} without drivable lanes!") - continue - - polygon = extract_exteriors_polygon(lane_group_helpers) - ids.append(junction.id) - lane_group_ids.append(lane_group_ids_) - geometries.append(polygon) - - data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) - return gpd.GeoDataFrame(data, geometry=geometries) - - def _extract_lane_group_dataframe(self) -> gpd.GeoDataFrame: - - ids = [] - lane_ids = [] - predecessor_lane_group_ids = [] - successor_lane_group_ids = [] - intersection_ids = [] - left_boundaries = [] - right_boundaries = [] - geometries = [] - - for lane_group_helper in self.lane_group_helper_dict.values(): - lane_group_helper: OpenDriveLaneGroupHelper - ids.append(lane_group_helper.lane_group_id) - lane_ids.append([lane_helper.lane_id for lane_helper in lane_group_helper.lane_helpers]) - predecessor_lane_group_ids.append(lane_group_helper.predecessor_lane_group_ids) - successor_lane_group_ids.append(lane_group_helper.successor_lane_group_ids) - intersection_ids.append(lane_group_helper.junction_id) - left_boundaries.append(shapely.LineString(lane_group_helper.inner_polyline_3d)) - right_boundaries.append(shapely.LineString(lane_group_helper.outer_polyline_3d)) - geometries.append(lane_group_helper.shapely_polygon) - - data = pd.DataFrame( - { - "id": ids, - "lane_ids": lane_ids, - "predecessor_ids": predecessor_lane_group_ids, - "successor_ids": successor_lane_group_ids, - "intersection_id": intersection_ids, - "left_boundary": left_boundaries, - "right_boundary": right_boundaries, - } - ) - gdf = gpd.GeoDataFrame(data, geometry=geometries) - - return gdf - - def _extract_crosswalk_dataframe(self) -> gpd.GeoDataFrame: - ids = [] - outlines = [] - geometries = [] - for object_helper in self.object_helper_dict.values(): - ids.append(object_helper.object_id) - outlines.append(shapely.LineString(object_helper.outline_3d)) - geometries.append(object_helper.shapely_polygon) - - data = pd.DataFrame({"id": ids, "outline": outlines}) - return gpd.GeoDataFrame(data, geometry=geometries) - - @staticmethod - def _convert_ids_to_int( - lane_df: gpd.GeoDataFrame, - walkways_df: gpd.GeoDataFrame, - carpark_df: gpd.GeoDataFrame, - generic_drivable_area_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, - intersections_df: gpd.GeoDataFrame, - crosswalk_df: gpd.GeoDataFrame, - ) -> None: - - # NOTE: intersection and crosswalk ids are already integers - - # initialize id mappings - lane_id_mapping = IntIDMapping.from_series(lane_df["id"]) - walkway_id_mapping = IntIDMapping.from_series(walkways_df["id"]) - carpark_id_mapping = IntIDMapping.from_series(carpark_df["id"]) - # TODO: add id mapping for intersections - generic_drivable_id_mapping = IntIDMapping.from_series(generic_drivable_area_df["id"]) - lane_group_id_mapping = IntIDMapping.from_series(lane_group_df["id"]) - - # Adjust cross reference in lane_df and lane_group_df - lane_df["lane_group_id"] = lane_df["lane_group_id"].map(lane_group_id_mapping.str_to_int) - lane_group_df["lane_ids"] = lane_group_df["lane_ids"].apply(lambda x: lane_id_mapping.map_list(x)) - - # Adjust predecessor/successor in lane_df and lane_group_df - for column in ["predecessor_ids", "successor_ids"]: - lane_df[column] = lane_df[column].apply(lambda x: lane_id_mapping.map_list(x)) - lane_group_df[column] = lane_group_df[column].apply(lambda x: lane_group_id_mapping.map_list(x)) - - for column in ["left_lane_id", "right_lane_id"]: - lane_df[column] = lane_df[column].apply( - lambda x: str(lane_id_mapping.str_to_int[x]) if pd.notna(x) and x is not None else x - ) - - lane_df["id"] = lane_df["id"].map(lane_id_mapping.str_to_int) - walkways_df["id"] = walkways_df["id"].map(walkway_id_mapping.str_to_int) - carpark_df["id"] = carpark_df["id"].map(carpark_id_mapping.str_to_int) - generic_drivable_area_df["id"] = generic_drivable_area_df["id"].map(generic_drivable_id_mapping.str_to_int) - lane_group_df["id"] = lane_group_df["id"].map(lane_group_id_mapping.str_to_int) - - intersections_df["lane_group_ids"] = intersections_df["lane_group_ids"].apply( - lambda x: lane_group_id_mapping.map_list(x) - ) - - def _extract_road_edge_df( - self, - lane_df: gpd.GeoDataFrame, - carpark_df: gpd.GeoDataFrame, - generic_drivable_area_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, - ) -> None: - road_edges = get_road_edges_3d_from_gdf(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) - road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) - - ids = np.arange(len(road_edges), dtype=np.int64).tolist() - # TODO @DanielDauner: Figure out if other types should/could be assigned here. - road_edge_types = [int(RoadEdgeType.ROAD_EDGE_BOUNDARY)] * len(road_edges) - geometries = road_edges - return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) - - def _extract_road_line_df( - self, - lane_df: gpd.GeoDataFrame, - lane_group_df: gpd.GeoDataFrame, - ) -> None: - - lane_group_on_intersection = { - lane_group_id: str(intersection_id) != "nan" - for lane_group_id, intersection_id in zip(lane_group_df.id.tolist(), lane_group_df.intersection_id.tolist()) - } - ids = [] - road_line_types = [] - geometries = [] - - running_id = 0 - for lane_row in lane_df.itertuples(): - on_intersection = lane_group_on_intersection.get(lane_row.lane_group_id, False) - if on_intersection: - # Skip road lines on intersections - continue - if str(lane_row.right_lane_id) in ["nan", "None"]: - # This is a boundary lane, e.g. a border or sidewalk - ids.append(running_id) - road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) - geometries.append(lane_row.right_boundary) - running_id += 1 - else: - # This is a regular lane - ids.append(running_id) - road_line_types.append(int(RoadLineType.BROKEN_SINGLE_WHITE)) - geometries.append(lane_row.right_boundary) - running_id += 1 - if str(lane_row.left_lane_id) in ["nan", "None"]: - # This is a boundary lane, e.g. a border or sidewalk - ids.append(running_id) - road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) - geometries.append(lane_row.left_boundary) - running_id += 1 - - data = pd.DataFrame({"id": ids, "road_line_type": road_line_types}) - return gpd.GeoDataFrame(data, geometry=geometries) - - -# TODO: move this somewhere else and improve -def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper]) -> shapely.Polygon: - - # Step 1: Extract all boundary line segments - all_polygons = [] - for lane_group_helper in lane_group_helpers: - all_polygons.append(lane_group_helper.shapely_polygon) - - # Step 2: Merge all boundaries and extract the enclosed polygons - # try: - merged_boundaries = unary_union(all_polygons) - # except Exception as e: - # warnings.warn(f"Topological error during polygon union: {e}") - # print([(helper.lane_group_id, poly.is_valid) for poly, helper in zip(all_polygons, lane_group_helpers)]) - # merged_boundaries = unary_union([poly for poly in all_polygons if poly.is_valid]) - - # Step 3: Generate polygons from the merged lines - polygons = list(polygonize(merged_boundaries)) - - # Step 4: Select the polygon that represents the intersection - # Usually it's the largest polygon - if len(polygons) == 1: - return polygons[0] - else: - # Take the largest polygon if there are multiple - return max(polygons, key=lambda p: p.area) diff --git a/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py b/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py new file mode 100644 index 00000000..c66c8613 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py @@ -0,0 +1,422 @@ +import logging +import os +import warnings +from pathlib import Path +from typing import Dict, Final, List + +import geopandas as gpd +import numpy as np +import pandas as pd +import shapely +from shapely.ops import polygonize, unary_union + +from d123.dataset.conversion.map.opendrive.parser.opendrive import Junction, OpenDrive +from d123.dataset.conversion.map.opendrive.utils.collection import collect_element_helpers +from d123.dataset.conversion.map.opendrive.utils.id_mapping import IntIDMapping +from d123.dataset.conversion.map.opendrive.utils.lane_helper import ( + OpenDriveLaneGroupHelper, + OpenDriveLaneHelper, +) +from d123.dataset.conversion.map.opendrive.utils.objects_helper import ( + OpenDriveObjectHelper, +) +from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length +from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import get_road_edges_3d_from_gdf +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType + +logger = logging.getLogger(__name__) +D123_MAPS_ROOT = Path(os.environ.get("D123_MAPS_ROOT")) + +MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # [m] + + +def convert_from_xodr( + xordr_file: Path, + map_name: str, + interpolation_step_size: float, + connection_distance_threshold: float, +) -> None: + + opendrive = OpenDrive.parse_from_file(xordr_file) + + _, junction_dict, lane_helper_dict, lane_group_helper_dict, object_helper_dict = collect_element_helpers( + opendrive, interpolation_step_size, connection_distance_threshold + ) + + # Collect data frames and store + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + dataframes[MapLayer.LANE] = _extract_lane_dataframe(lane_group_helper_dict) + dataframes[MapLayer.LANE_GROUP] = _extract_lane_group_dataframe(lane_group_helper_dict) + dataframes[MapLayer.WALKWAY] = _extract_walkways_dataframe(lane_helper_dict) + dataframes[MapLayer.CARPARK] = _extract_carpark_dataframe(lane_helper_dict) + dataframes[MapLayer.GENERIC_DRIVABLE] = _extract_generic_drivable_dataframe(lane_helper_dict) + dataframes[MapLayer.INTERSECTION] = _extract_intersections_dataframe(junction_dict, lane_group_helper_dict) + dataframes[MapLayer.CROSSWALK] = _extract_crosswalk_dataframe(object_helper_dict) + + _convert_ids_to_int( + dataframes[MapLayer.LANE], + dataframes[MapLayer.WALKWAY], + dataframes[MapLayer.CARPARK], + dataframes[MapLayer.GENERIC_DRIVABLE], + dataframes[MapLayer.LANE_GROUP], + dataframes[MapLayer.INTERSECTION], + dataframes[MapLayer.CROSSWALK], + ) + dataframes[MapLayer.ROAD_EDGE] = _extract_road_edge_df( + dataframes[MapLayer.LANE], + dataframes[MapLayer.CARPARK], + dataframes[MapLayer.GENERIC_DRIVABLE], + dataframes[MapLayer.LANE_GROUP], + ) + dataframes[MapLayer.ROAD_LINE] = _extract_road_line_df( + dataframes[MapLayer.LANE], + dataframes[MapLayer.LANE_GROUP], + ) + map_file_name = D123_MAPS_ROOT / f"{map_name}.gpkg" + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_name, layer=layer.serialize(), driver="GPKG", mode="a") + + +def _extract_lane_dataframe(lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper]) -> gpd.GeoDataFrame: + + ids = [] + lane_group_ids = [] + speed_limits_mps = [] + predecessor_ids = [] + successor_ids = [] + left_boundaries = [] + right_boundaries = [] + left_lane_ids = [] + right_lane_ids = [] + baseline_paths = [] + geometries = [] + + for lane_group_helper in lane_group_helper_dict.values(): + lane_group_id = lane_group_helper.lane_group_id + lane_helpers = lane_group_helper.lane_helpers + num_lanes = len(lane_helpers) + # NOTE: Lanes are going left to right, ie. inner to outer + for lane_idx, lane_helper in enumerate(lane_helpers): + ids.append(lane_helper.lane_id) + lane_group_ids.append(lane_group_id) + speed_limits_mps.append(lane_helper.speed_limit_mps) + predecessor_ids.append(lane_helper.predecessor_lane_ids) + successor_ids.append(lane_helper.successor_lane_ids) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + baseline_paths.append(shapely.LineString(lane_helper.center_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + left_lane_id = lane_helpers[lane_idx - 1].lane_id if lane_idx > 0 else None + right_lane_id = lane_helpers[lane_idx + 1].lane_id if lane_idx < num_lanes - 1 else None + left_lane_ids.append(left_lane_id) + right_lane_ids.append(right_lane_id) + + data = pd.DataFrame( + { + "id": ids, + "lane_group_id": lane_group_ids, + "speed_limit_mps": speed_limits_mps, + "predecessor_ids": predecessor_ids, + "successor_ids": successor_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "left_lane_id": left_lane_ids, + "right_lane_id": right_lane_ids, + "baseline_path": baseline_paths, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_lane_group_dataframe(lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper]) -> gpd.GeoDataFrame: + + ids = [] + lane_ids = [] + predecessor_lane_group_ids = [] + successor_lane_group_ids = [] + intersection_ids = [] + left_boundaries = [] + right_boundaries = [] + geometries = [] + + for lane_group_helper in lane_group_helper_dict.values(): + lane_group_helper: OpenDriveLaneGroupHelper + ids.append(lane_group_helper.lane_group_id) + lane_ids.append([lane_helper.lane_id for lane_helper in lane_group_helper.lane_helpers]) + predecessor_lane_group_ids.append(lane_group_helper.predecessor_lane_group_ids) + successor_lane_group_ids.append(lane_group_helper.successor_lane_group_ids) + intersection_ids.append(lane_group_helper.junction_id) + left_boundaries.append(shapely.LineString(lane_group_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_group_helper.outer_polyline_3d)) + geometries.append(lane_group_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "lane_ids": lane_ids, + "predecessor_ids": predecessor_lane_group_ids, + "successor_ids": successor_lane_group_ids, + "intersection_id": intersection_ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + } + ) + gdf = gpd.GeoDataFrame(data, geometry=geometries) + + return gdf + + +def _extract_walkways_dataframe(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> gpd.GeoDataFrame: + + ids = [] + left_boundaries = [] + right_boundaries = [] + outlines = [] + geometries = [] + + for lane_helper in lane_helper_dict.values(): + if lane_helper.type == "sidewalk": + ids.append(lane_helper.lane_id) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "outline": outlines, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_carpark_dataframe(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> gpd.GeoDataFrame: + + ids = [] + left_boundaries = [] + right_boundaries = [] + outlines = [] + geometries = [] + + for lane_helper in lane_helper_dict.values(): + if lane_helper.type == "parking": + ids.append(lane_helper.lane_id) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "left_boundary": left_boundaries, + "right_boundary": right_boundaries, + "outline": outlines, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_generic_drivable_dataframe(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> gpd.GeoDataFrame: + + ids = [] + left_boundaries = [] + right_boundaries = [] + outlines = [] + geometries = [] + + for lane_helper in lane_helper_dict.values(): + if lane_helper.type in ["none", "border", "bidirectional"]: + ids.append(lane_helper.lane_id) + left_boundaries.append(shapely.LineString(lane_helper.inner_polyline_3d)) + right_boundaries.append(shapely.LineString(lane_helper.outer_polyline_3d)) + outlines.append(shapely.LineString(lane_helper.outline_polyline_3d)) + geometries.append(lane_helper.shapely_polygon) + + data = pd.DataFrame( + { + "id": ids, + "left_boundary": left_boundaries, + "right_boundary": left_boundaries, + "outline": outlines, + } + ) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_intersections_dataframe( + junction_dict: Dict[str, Junction], + lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper], +) -> gpd.GeoDataFrame: + def _find_lane_group_helpers_with_junction_id(junction_id: int) -> List[OpenDriveLaneGroupHelper]: + return [ + lane_group_helper + for lane_group_helper in lane_group_helper_dict.values() + if lane_group_helper.junction_id == junction_id + ] + + ids = [] + lane_group_ids = [] + geometries = [] + for junction in junction_dict.values(): + lane_group_helpers = _find_lane_group_helpers_with_junction_id(junction.id) + lane_group_ids_ = [lane_group_helper.lane_group_id for lane_group_helper in lane_group_helpers] + if len(lane_group_ids_) == 0: + logger.debug(f"Skipping Junction {junction.id} without lane groups!") + continue + + polygon = extract_exteriors_polygon(lane_group_helpers) + ids.append(junction.id) + lane_group_ids.append(lane_group_ids_) + geometries.append(polygon) + + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_crosswalk_dataframe(object_helper_dict: Dict[int, OpenDriveObjectHelper]) -> gpd.GeoDataFrame: + ids = [] + outlines = [] + geometries = [] + for object_helper in object_helper_dict.values(): + ids.append(object_helper.object_id) + outlines.append(shapely.LineString(object_helper.outline_3d)) + geometries.append(object_helper.shapely_polygon) + + data = pd.DataFrame({"id": ids, "outline": outlines}) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _convert_ids_to_int( + lane_df: gpd.GeoDataFrame, + walkways_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, + intersections_df: gpd.GeoDataFrame, + crosswalk_df: gpd.GeoDataFrame, +) -> None: + + # NOTE: intersection and crosswalk ids are already integers + + # initialize id mappings + lane_id_mapping = IntIDMapping.from_series(lane_df["id"]) + walkway_id_mapping = IntIDMapping.from_series(walkways_df["id"]) + carpark_id_mapping = IntIDMapping.from_series(carpark_df["id"]) + generic_drivable_id_mapping = IntIDMapping.from_series(generic_drivable_area_df["id"]) + lane_group_id_mapping = IntIDMapping.from_series(lane_group_df["id"]) + + # Adjust cross reference in lane_df and lane_group_df + lane_df["lane_group_id"] = lane_df["lane_group_id"].map(lane_group_id_mapping.str_to_int) + lane_group_df["lane_ids"] = lane_group_df["lane_ids"].apply(lambda x: lane_id_mapping.map_list(x)) + + # Adjust predecessor/successor in lane_df and lane_group_df + for column in ["predecessor_ids", "successor_ids"]: + lane_df[column] = lane_df[column].apply(lambda x: lane_id_mapping.map_list(x)) + lane_group_df[column] = lane_group_df[column].apply(lambda x: lane_group_id_mapping.map_list(x)) + + for column in ["left_lane_id", "right_lane_id"]: + lane_df[column] = lane_df[column].apply( + lambda x: str(lane_id_mapping.str_to_int[x]) if pd.notna(x) and x is not None else x + ) + + lane_df["id"] = lane_df["id"].map(lane_id_mapping.str_to_int) + walkways_df["id"] = walkways_df["id"].map(walkway_id_mapping.str_to_int) + carpark_df["id"] = carpark_df["id"].map(carpark_id_mapping.str_to_int) + generic_drivable_area_df["id"] = generic_drivable_area_df["id"].map(generic_drivable_id_mapping.str_to_int) + lane_group_df["id"] = lane_group_df["id"].map(lane_group_id_mapping.str_to_int) + + intersections_df["lane_group_ids"] = intersections_df["lane_group_ids"].apply( + lambda x: lane_group_id_mapping.map_list(x) + ) + + +def _extract_road_line_df( + lane_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> None: + + lane_group_on_intersection = { + lane_group_id: str(intersection_id) != "nan" + for lane_group_id, intersection_id in zip(lane_group_df.id.tolist(), lane_group_df.intersection_id.tolist()) + } + ids = [] + road_line_types = [] + geometries = [] + + running_id = 0 + for lane_row in lane_df.itertuples(): + on_intersection = lane_group_on_intersection.get(lane_row.lane_group_id, False) + if on_intersection: + # Skip road lines on intersections + continue + if str(lane_row.right_lane_id) in ["nan", "None"]: + # This is a boundary lane, e.g. a border or sidewalk + ids.append(running_id) + road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) + geometries.append(lane_row.right_boundary) + running_id += 1 + else: + # This is a regular lane + ids.append(running_id) + road_line_types.append(int(RoadLineType.BROKEN_SINGLE_WHITE)) + geometries.append(lane_row.right_boundary) + running_id += 1 + if str(lane_row.left_lane_id) in ["nan", "None"]: + # This is a boundary lane, e.g. a border or sidewalk + ids.append(running_id) + road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) + geometries.append(lane_row.left_boundary) + running_id += 1 + + data = pd.DataFrame({"id": ids, "road_line_type": road_line_types}) + return gpd.GeoDataFrame(data, geometry=geometries) + + +def _extract_road_edge_df( + lane_df: gpd.GeoDataFrame, + carpark_df: gpd.GeoDataFrame, + generic_drivable_area_df: gpd.GeoDataFrame, + lane_group_df: gpd.GeoDataFrame, +) -> None: + road_edges = get_road_edges_3d_from_gdf(lane_df, carpark_df, generic_drivable_area_df, lane_group_df) + road_edges = split_line_geometry_by_max_length(road_edges, MAX_ROAD_EDGE_LENGTH) + + ids = np.arange(len(road_edges), dtype=np.int64).tolist() + # TODO @DanielDauner: Figure out if other types should/could be assigned here. + road_edge_types = [int(RoadEdgeType.ROAD_EDGE_BOUNDARY)] * len(road_edges) + geometries = road_edges + return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) + + +# TODO: move this somewhere else and improve +def extract_exteriors_polygon(lane_group_helpers: List[OpenDriveLaneGroupHelper]) -> shapely.Polygon: + + # Step 1: Extract all boundary line segments + all_polygons = [] + for lane_group_helper in lane_group_helpers: + all_polygons.append(lane_group_helper.shapely_polygon) + + # Step 2: Merge all boundaries and extract the enclosed polygons + # try: + merged_boundaries = unary_union(all_polygons) + # except Exception as e: + # warnings.warn(f"Topological error during polygon union: {e}") + # print([(helper.lane_group_id, poly.is_valid) for poly, helper in zip(all_polygons, lane_group_helpers)]) + # merged_boundaries = unary_union([poly for poly in all_polygons if poly.is_valid]) + + # Step 3: Generate polygons from the merged lines + polygons = list(polygonize(merged_boundaries)) + + # Step 4: Select the polygon that represents the intersection + # Usually it's the largest polygon + if len(polygons) == 1: + return polygons[0] + else: + # Take the largest polygon if there are multiple + return max(polygons, key=lambda p: p.area) diff --git a/d123/dataset/conversion/map/opendrive/parser/__init__.py b/d123/dataset/conversion/map/opendrive/parser/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/d123/dataset/conversion/map/opendrive/elements/elevation.py b/d123/dataset/conversion/map/opendrive/parser/elevation.py similarity index 56% rename from d123/dataset/conversion/map/opendrive/elements/elevation.py rename to d123/dataset/conversion/map/opendrive/parser/elevation.py index 782c4311..7ee7aa67 100644 --- a/d123/dataset/conversion/map/opendrive/elements/elevation.py +++ b/d123/dataset/conversion/map/opendrive/parser/elevation.py @@ -4,12 +4,17 @@ from typing import List, Optional from xml.etree.ElementTree import Element -import numpy as np -import numpy.typing as npt +from d123.dataset.conversion.map.opendrive.parser.polynomial import Polynomial @dataclass class ElevationProfile: + """ + Models elevation along s-axis of reference line. + + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-1d876c00-d69e-46d9-bbcd-709ab48f14b1 + """ + elevations: List[Elevation] def __post_init__(self): @@ -27,27 +32,21 @@ def parse(cls, elevation_profile_element: Optional[Element]) -> ElevationProfile @dataclass -class Elevation: - """TODO: Refactor and merge with other elements, e.g. LaneOffset""" - - s: float - a: float - b: float - c: float - d: float - - @classmethod - def parse(cls, elevation_element: Element) -> Elevation: - args = {key: float(elevation_element.get(key)) for key in ["s", "a", "b", "c", "d"]} - return Elevation(**args) +class Elevation(Polynomial): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + Represents an elevation profile in OpenDRIVE. + """ @dataclass class LateralProfile: + """ + Models elevation along t-axis of reference line. + + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 + """ super_elevations: List[SuperElevation] shapes: List[Shape] @@ -76,28 +75,21 @@ def parse(cls, lateral_profile_element: Optional[Element]) -> LateralProfile: @dataclass -class SuperElevation: - """TODO: Refactor and merge with other elements, e.g. Elevation, LaneOffset""" - - s: float - a: float - b: float - c: float - d: float +class SuperElevation(Polynomial): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_05_elevation.html#sec-4abf7baf-fb2f-4263-8133-ad0f64f0feac - @classmethod - def parse(cls, super_elevation_element: Element) -> SuperElevation: - args = {key: float(super_elevation_element.get(key)) for key in ["s", "a", "b", "c", "d"]} - return SuperElevation(**args) - - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + superelevation (ds) = a + b*ds + c*ds² + d*ds³ + """ @dataclass class Shape: - """TODO: Refactor and merge with other elements, e.g. Elevation, LaneOffset""" + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/1.8.0/specification/10_roads/10_05_elevation.html#sec-66ac2b58-dc5e-4538-884d-204406ea53f2 + + hShape (dt)= a + b*dt + c*dt² + d*dt³ + """ s: float t: float @@ -111,6 +103,8 @@ def parse(cls, shape_element: Element) -> Shape: args = {key: float(shape_element.get(key)) for key in ["s", "t", "a", "b", "c", "d"]} return Shape(**args) - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + def get_value(self, dt: float) -> float: + """ + Evaluate the polynomial at a given t value. + """ + return self.a + self.b * dt + self.c * dt**2 + self.d * dt**3 diff --git a/d123/dataset/conversion/map/opendrive/elements/geometry.py b/d123/dataset/conversion/map/opendrive/parser/geometry.py similarity index 76% rename from d123/dataset/conversion/map/opendrive/elements/geometry.py rename to d123/dataset/conversion/map/opendrive/parser/geometry.py index 2dd33c4e..c1f337c0 100644 --- a/d123/dataset/conversion/map/opendrive/elements/geometry.py +++ b/d123/dataset/conversion/map/opendrive/parser/geometry.py @@ -13,6 +13,10 @@ @dataclass class Geometry: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_02_road_reference_line.html + """ + s: float x: float y: float @@ -33,6 +37,10 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: @dataclass class Line(Geometry): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_03_straight_line.html + """ + @classmethod def parse(cls, geometry_element: Element) -> Geometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} @@ -53,9 +61,16 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: @dataclass class Arc(Geometry): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_05_arc.html + """ curvature: float + def __post_init__(self): + if self.curvature == 0.0: + raise ValueError("Curvature cannot be zero for Arc geometry.") + @classmethod def parse(cls, geometry_element: Element) -> Geometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} @@ -64,19 +79,20 @@ def parse(cls, geometry_element: Element) -> Geometry: def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - c = self.curvature - hdg = self.hdg - np.pi / 2 + kappa = self.curvature + radius = 1.0 / kappa if kappa != 0 else float("inf") - a = 2 / c * np.sin(s * c / 2) - alpha = (np.pi - s * c) / 2 - hdg + initial_heading = self.hdg + final_heading = initial_heading + s * kappa - dx = -1 * a * np.cos(alpha) - dy = a * np.sin(alpha) + # Parametric circle equations + dx = radius * (np.sin(final_heading) - np.sin(initial_heading)) + dy = -radius * (np.cos(final_heading) - np.cos(initial_heading)) interpolated_se2 = self.start_se2.copy() interpolated_se2[StateSE2Index.X] += dx interpolated_se2[StateSE2Index.Y] += dy - interpolated_se2[StateSE2Index.YAW] += s * self.curvature + interpolated_se2[StateSE2Index.YAW] = final_heading if t != 0.0: interpolated_se2[StateSE2Index.X] += t * np.cos(interpolated_se2[StateSE2Index.YAW] + np.pi / 2) @@ -85,32 +101,21 @@ def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: return interpolated_se2 -@dataclass -class Geometry: - s: float - x: float - y: float - hdg: float - length: float - - @property - def start_se2(self) -> npt.NDArray[np.float64]: - start_se2 = np.zeros(len(StateSE2Index), dtype=np.float64) - start_se2[StateSE2Index.X] = self.x - start_se2[StateSE2Index.Y] = self.y - start_se2[StateSE2Index.YAW] = self.hdg - return start_se2 - - def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - raise NotImplementedError - - @dataclass class Spiral(Geometry): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/09_geometries/09_04_spiral.html + https://en.wikipedia.org/wiki/Euler_spiral + """ curvature_start: float curvature_end: float + def __post_init__(self): + gamma = (self.curvature_end - self.curvature_start) / self.length + if abs(gamma) < 1e-10: + raise ValueError("Curvature change is too small, cannot define a valid spiral.") + @classmethod def parse(cls, geometry_element: Element) -> Geometry: args = {key: float(geometry_element.get(key)) for key in ["s", "x", "y", "hdg", "length"]} @@ -120,18 +125,9 @@ def parse(cls, geometry_element: Element) -> Geometry: return cls(**args) def interpolate_se2(self, s: float, t: float = 0.0) -> npt.NDArray[np.float64]: - """ - https://en.wikipedia.org/wiki/Euler_spiral - :param s: _description_ - :param t: _description_, defaults to 0.0 - :return: _description_ - """ interpolated_se2 = self.start_se2.copy() gamma = (self.curvature_end - self.curvature_start) / self.length - if abs(gamma) < 1e-10: - print(gamma) - # NOTE: doesn't consider case where gamma == 0 dx, dy = self._compute_spiral_position(s, gamma) diff --git a/d123/dataset/conversion/map/opendrive/elements/lane.py b/d123/dataset/conversion/map/opendrive/parser/lane.py similarity index 77% rename from d123/dataset/conversion/map/opendrive/elements/lane.py rename to d123/dataset/conversion/map/opendrive/parser/lane.py index 9afb0f25..ccbfc9ff 100644 --- a/d123/dataset/conversion/map/opendrive/elements/lane.py +++ b/d123/dataset/conversion/map/opendrive/parser/lane.py @@ -4,12 +4,15 @@ from typing import List, Optional from xml.etree.ElementTree import Element -import numpy as np -import numpy.typing as npt +from d123.dataset.conversion.map.opendrive.parser.polynomial import Polynomial @dataclass class Lanes: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_01_introduction.html + """ + lane_offsets: List[LaneOffset] lane_sections: List[LaneSection] @@ -32,40 +35,29 @@ def parse(cls, lanes_element: Optional[Element]) -> Lanes: return Lanes(**args) @property - def num_lane_sections(self): + def num_lane_sections(self) -> int: return len(self.lane_sections) @property - def last_lane_section_idx(self): + def last_lane_section_idx(self) -> int: return self.num_lane_sections - 1 @dataclass -class LaneOffset: - """Section 11.4""" - - s: float - a: float - b: float - c: float - d: float - - def __post_init__(self): - # NOTE: added assertion/filtering to check for element type or consistency - pass +class LaneOffset(Polynomial): + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_04_lane_offset.html - @classmethod - def parse(cls, lane_offset_element: Element) -> LaneOffset: - args = {key: float(lane_offset_element.get(key)) for key in ["s", "a", "b", "c", "d"]} - return LaneOffset(**args) - - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + offset (ds) = a + b*ds + c*ds² + d*ds³ + """ @dataclass class LaneSection: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_03_lane_sections.html + """ + s: float left_lanes: List[Lane] center_lanes: List[Lane] @@ -104,6 +96,9 @@ def parse(cls, lane_section_element: Optional[Element]) -> LaneSection: @dataclass class Lane: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_05_lane_link.html + """ id: int type: str @@ -147,6 +142,10 @@ def parse(cls, lane_element: Optional[Element]) -> Lane: @dataclass class Width: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_06_lane_geometry.html#sec-8d8ac2e0-b3d6-4048-a9ed-d5191af5c74b + """ + s_offset: Optional[float] = None a: Optional[float] = None b: Optional[float] = None @@ -163,13 +162,25 @@ def parse(cls, width_element: Optional[Element]) -> Width: args["d"] = float(width_element.get("d")) return Width(**args) - @property - def polynomial_coefficients(self) -> npt.NDArray[np.float64]: - return np.array([self.a, self.b, self.c, self.d], dtype=np.float64) + def get_polynomial(self, t_sign: float = 1.0) -> Polynomial: + """ + Returns the polynomial representation of the width. + """ + return Polynomial( + s=self.s_offset, + a=self.a * t_sign, + b=self.b * t_sign, + c=self.c * t_sign, + d=self.d * t_sign, + ) @dataclass class RoadMark: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/11_lanes/11_08_road_markings.html + """ + s_offset: float = None type: str = None material: Optional[str] = None diff --git a/d123/dataset/conversion/map/opendrive/elements/objects.py b/d123/dataset/conversion/map/opendrive/parser/objects.py similarity index 85% rename from d123/dataset/conversion/map/opendrive/elements/objects.py rename to d123/dataset/conversion/map/opendrive/parser/objects.py index 4ca2f730..4c409c54 100644 --- a/d123/dataset/conversion/map/opendrive/elements/objects.py +++ b/d123/dataset/conversion/map/opendrive/parser/objects.py @@ -7,6 +7,9 @@ @dataclass class Object: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/13_objects/13_01_introduction.html + """ id: int name: str @@ -55,6 +58,10 @@ def parse(cls, object_element: Optional[Element]) -> Object: @dataclass class CornerLocal: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/13_objects/13_03_object_outline.html#sec-cc00c8a6-eea6-49e6-b90c-37b21524c748 + """ + u: float v: float z: float diff --git a/d123/dataset/conversion/map/opendrive/elements/opendrive.py b/d123/dataset/conversion/map/opendrive/parser/opendrive.py similarity index 79% rename from d123/dataset/conversion/map/opendrive/elements/opendrive.py rename to d123/dataset/conversion/map/opendrive/parser/opendrive.py index de73ce87..2a07bb5f 100644 --- a/d123/dataset/conversion/map/opendrive/elements/opendrive.py +++ b/d123/dataset/conversion/map/opendrive/parser/opendrive.py @@ -3,14 +3,17 @@ import traceback from dataclasses import dataclass from pathlib import Path -from typing import List, Optional +from typing import List, Literal, Optional from xml.etree.ElementTree import Element, parse -from d123.dataset.conversion.map.opendrive.elements.road import Road +from d123.dataset.conversion.map.opendrive.parser.road import Road @dataclass class OpenDrive: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/06_general_architecture/06_03_root_element.html + """ header: Header @@ -55,7 +58,9 @@ def parse_from_file(cls, file_path: Path) -> OpenDrive: @dataclass class Header: - """Section 4.4.2""" + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/06_general_architecture/06_04_header.html + """ rev_major: Optional[int] = None rev_minor: Optional[int] = None @@ -95,6 +100,10 @@ def parse(cls, header_element: Optional[Element]) -> Header: @dataclass class Controller: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/14_signals/14_06_controllers.html + """ + name: str id: int sequence: int @@ -118,6 +127,9 @@ def parse(cls, controller_element: Optional[Element]) -> Junction: @dataclass class Control: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/14_signals/14_06_controllers.html + """ signal_id: str type: str @@ -132,6 +144,10 @@ def parse(cls, control_element: Optional[Element]) -> Control: @dataclass class Junction: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/12_junctions/12_02_common_junctions.html + """ + id: int name: str connections: List[Connection] @@ -153,12 +169,19 @@ def parse(cls, junction_element: Optional[Element]) -> Junction: @dataclass class Connection: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/12_junctions/12_02_common_junctions.html + """ + id: int incoming_road: int connecting_road: int - contact_point: str + contact_point: Literal["start", "end"] lane_links: List[LaneLink] + def __post_init__(self): + assert self.contact_point in ["start", "end"] + @classmethod def parse(cls, connection_element: Optional[Element]) -> Connection: args = {} @@ -178,6 +201,10 @@ def parse(cls, connection_element: Optional[Element]) -> Connection: @dataclass class LaneLink: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/1.8.0/specification/12_junctions/12_04_connecting_roads.html#top-3e9bb97e-f2ab-4751-906a-c25e9fb7ac4e + """ + start: int # NOTE: named "from" in xml end: int # NOTE: named "to" in xml diff --git a/d123/dataset/conversion/map/opendrive/parser/polynomial.py b/d123/dataset/conversion/map/opendrive/parser/polynomial.py new file mode 100644 index 00000000..497ffed5 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/parser/polynomial.py @@ -0,0 +1,32 @@ +from dataclasses import dataclass +from xml.etree.ElementTree import Element + + +@dataclass +class Polynomial: + """ + Multiple OpenDRIVE elements use polynomial coefficients, e.g. Elevation, LaneOffset, etc. + This class provides a common interface to parse and access polynomial coefficients. + + Stores a polynomial function of the third order: + value(ds) = a + b*ds + c*ds² + d*ds³ + """ + + s: float + a: float + b: float + c: float + d: float + + @classmethod + def parse(cls: type["Polynomial"], element: Element) -> "Polynomial": + args = {key: float(element.get(key)) for key in ["s", "a", "b", "c", "d"]} + return cls(**args) + + def get_value(self, ds: float) -> float: + """ + Evaluate the polynomial at a given ds value. + :param ds: The distance along the road. + :return: The evaluated polynomial value. + """ + return self.a + self.b * ds + self.c * ds**2 + self.d * ds**3 diff --git a/d123/dataset/conversion/map/opendrive/parser/reference.py b/d123/dataset/conversion/map/opendrive/parser/reference.py new file mode 100644 index 00000000..b4c88dc1 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/parser/reference.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from functools import cached_property +from typing import Final, List, Optional, Union +from xml.etree.ElementTree import Element + +import numpy as np +import numpy.typing as npt + +from d123.common.geometry.base import Point3DIndex, StateSE2Index +from d123.dataset.conversion.map.opendrive.parser.elevation import Elevation +from d123.dataset.conversion.map.opendrive.parser.geometry import Arc, Geometry, Line, Spiral +from d123.dataset.conversion.map.opendrive.parser.lane import LaneOffset, Width +from d123.dataset.conversion.map.opendrive.parser.polynomial import Polynomial + +TOLERANCE: Final[float] = 1e-3 + + +@dataclass +class PlanView: + + geometries: List[Geometry] + + def __post_init__(self): + # Ensure geometries are sorted by their starting position 's' + self.geometries.sort(key=lambda x: x.s) + + @classmethod + def parse(cls, plan_view_element: Optional[Element]) -> PlanView: + geometries: List[Geometry] = [] + for geometry_element in plan_view_element.findall("geometry"): + if geometry_element.find("line") is not None: + geometry = Line.parse(geometry_element) + elif geometry_element.find("arc") is not None: + geometry = Arc.parse(geometry_element) + elif geometry_element.find("spiral") is not None: + geometry = Spiral.parse(geometry_element) + else: + geometry_str = ET.tostring(geometry_element, encoding="unicode") + raise NotImplementedError(f"Geometry not implemented: {geometry_str}") + geometries.append(geometry) + return PlanView(geometries=geometries) + + @cached_property + def geometry_lengths(self) -> npt.NDArray[np.float64]: + return np.cumsum([0.0] + [geo.length for geo in self.geometries], dtype=np.float64) + + @property + def length(self) -> float: + return float(self.geometry_lengths[-1]) + + def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: + """ + Interpolates the SE2 state at a given longitudinal position s along the plan view. + """ + if s > self.length: + if np.isclose(s, self.length, atol=TOLERANCE): + s = self.length + else: + raise ValueError( + f"s={s} is beyond the end of the plan view (length={self.length}) with tolerance={TOLERANCE}." + ) + + # Find the geometry segment containing s + geo_idx = np.searchsorted(self.geometry_lengths, s, side="right") - 1 + geo_idx = int(np.clip(geo_idx, 0, len(self.geometries) - 1)) + + return self.geometries[geo_idx].interpolate_se2(s - self.geometry_lengths[geo_idx], t) + + +@dataclass +class ReferenceLine: + + reference_line: Union[ReferenceLine, PlanView] + width_polynomials: List[Polynomial] + elevations: List[Elevation] + s_offset: float + + @property + def length(self) -> float: + return float(self.reference_line.length) + + @classmethod + def from_plan_view( + cls, + plan_view: PlanView, + lane_offsets: List[LaneOffset], + elevations: List[Elevation], + ) -> ReferenceLine: + args = {} + args["reference_line"] = plan_view + args["width_polynomials"] = lane_offsets + args["elevations"] = elevations + args["s_offset"] = 0.0 + return ReferenceLine(**args) + + @classmethod + def from_reference_line( + cls, + reference_line: ReferenceLine, + widths: List[Width], + s_offset: float = 0.0, + t_sign: float = 1.0, + ) -> ReferenceLine: + assert t_sign in [1.0, -1.0], "t_sign must be either 1.0 or -1.0" + + args = {} + args["reference_line"] = reference_line + width_polynomials: List[Polynomial] = [] + for width in widths: + width_polynomials.append(width.get_polynomial(t_sign=t_sign)) + args["width_polynomials"] = width_polynomials + args["s_offset"] = s_offset + args["elevations"] = reference_line.elevations + + return ReferenceLine(**args) + + @staticmethod + def _find_polynomial(s: float, polynomials: List[Polynomial], lane_section_end: bool = False) -> Polynomial: + + out_polynomial = polynomials[-1] + for polynomial in polynomials[::-1]: + if lane_section_end: + if polynomial.s < s: + out_polynomial = polynomial + break + else: + if polynomial.s <= s: + out_polynomial = polynomial + break + + return out_polynomial + + def interpolate_se2(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: + + width_polynomial = self._find_polynomial(s, self.width_polynomials, lane_section_end=lane_section_end) + t_offset = width_polynomial.get_value(s - width_polynomial.s) + se2 = self.reference_line.interpolate_se2(self.s_offset + s, t=t_offset + t, lane_section_end=lane_section_end) + + return se2 + + def interpolate_3d(self, s: float, t: float = 0.0, lane_section_end: bool = False) -> npt.NDArray[np.float64]: + + se2 = self.interpolate_se2(s, t, lane_section_end=lane_section_end) + + elevation_polynomial = self._find_polynomial(s, self.elevations, lane_section_end=lane_section_end) + point_3d = np.zeros(len(Point3DIndex), dtype=np.float64) + point_3d[Point3DIndex.XY] = se2[StateSE2Index.XY] + point_3d[Point3DIndex.Z] = elevation_polynomial.get_value(s - elevation_polynomial.s) + + return point_3d diff --git a/d123/dataset/conversion/map/opendrive/elements/road.py b/d123/dataset/conversion/map/opendrive/parser/road.py similarity index 77% rename from d123/dataset/conversion/map/opendrive/elements/road.py rename to d123/dataset/conversion/map/opendrive/parser/road.py index cdef3472..e0171429 100644 --- a/d123/dataset/conversion/map/opendrive/elements/road.py +++ b/d123/dataset/conversion/map/opendrive/parser/road.py @@ -4,17 +4,21 @@ from typing import List, Optional from xml.etree.ElementTree import Element -from d123.dataset.conversion.map.opendrive.elements.elevation import ElevationProfile, LateralProfile -from d123.dataset.conversion.map.opendrive.elements.lane import Lanes -from d123.dataset.conversion.map.opendrive.elements.objects import Object -from d123.dataset.conversion.map.opendrive.elements.reference import PlanView +from d123.dataset.conversion.map.opendrive.parser.elevation import ElevationProfile, LateralProfile +from d123.dataset.conversion.map.opendrive.parser.lane import Lanes +from d123.dataset.conversion.map.opendrive.parser.objects import Object +from d123.dataset.conversion.map.opendrive.parser.reference import PlanView @dataclass class Road: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_01_introduction.html + """ + id: int junction: Optional[str] - length: float # [m] + length: float name: Optional[str] link: Link @@ -66,7 +70,9 @@ def parse(cls, road_element: Element) -> Road: @dataclass class Link: - """Section 8.2""" + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_03_road_linkage.html + """ predecessor: Optional[PredecessorSuccessor] = None successor: Optional[PredecessorSuccessor] = None @@ -84,6 +90,10 @@ def parse(cls, link_element: Optional[Element]) -> PlanView: @dataclass class PredecessorSuccessor: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_03_road_linkage.html + """ + element_type: Optional[str] = None element_id: Optional[int] = None contact_point: Optional[str] = None @@ -103,6 +113,9 @@ def parse(cls, element: Element) -> PredecessorSuccessor: @dataclass class RoadType: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_04_road_type.html + """ s: Optional[float] = None type: Optional[str] = None @@ -120,6 +133,10 @@ def parse(cls, road_type_element: Optional[Element]) -> RoadType: @dataclass class Speed: + """ + https://publications.pages.asam.net/standards/ASAM_OpenDRIVE/ASAM_OpenDRIVE_Specification/latest/specification/10_roads/10_04_road_type.html#sec-33dc6899-854e-4533-a3d9-76e9e1518ee7 + """ + max: Optional[float] = None unit: Optional[str] = None diff --git a/d123/dataset/conversion/map/opendrive/utils/collection.py b/d123/dataset/conversion/map/opendrive/utils/collection.py new file mode 100644 index 00000000..be631445 --- /dev/null +++ b/d123/dataset/conversion/map/opendrive/utils/collection.py @@ -0,0 +1,325 @@ +import logging +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from d123.dataset.conversion.map.opendrive.parser.opendrive import Junction, OpenDrive +from d123.dataset.conversion.map.opendrive.parser.reference import ReferenceLine +from d123.dataset.conversion.map.opendrive.parser.road import Road +from d123.dataset.conversion.map.opendrive.utils.id_system import ( + build_lane_id, + derive_lane_section_id, + lane_group_id_from_lane_id, + road_id_from_lane_group_id, +) +from d123.dataset.conversion.map.opendrive.utils.lane_helper import ( + OpenDriveLaneGroupHelper, + OpenDriveLaneHelper, + lane_section_to_lane_helpers, +) +from d123.dataset.conversion.map.opendrive.utils.objects_helper import OpenDriveObjectHelper, get_object_helper + +logger = logging.getLogger(__name__) + + +def collect_element_helpers( + opendrive: OpenDrive, + interpolation_step_size: float, + connection_distance_threshold: float, +) -> Tuple[ + Dict[int, Road], + Dict[int, Junction], + Dict[str, OpenDriveLaneHelper], + Dict[str, OpenDriveLaneGroupHelper], + Dict[int, OpenDriveObjectHelper], +]: + + # 1. Fill the road and junction dictionaries + road_dict: Dict[int, Road] = {road.id: road for road in opendrive.roads} + junction_dict: Dict[int, Junction] = {junction.id: junction for junction in opendrive.junctions} + + # 2. Create lane helpers from the roads + lane_helper_dict: Dict[str, OpenDriveLaneHelper] = {} + for road in opendrive.roads: + reference_line = ReferenceLine.from_plan_view( + road.plan_view, + road.lanes.lane_offsets, + road.elevation_profile.elevations, + ) + lane_section_lengths: List[float] = [ls.s for ls in road.lanes.lane_sections] + [road.length] + for idx, lane_section in enumerate(road.lanes.lane_sections): + lane_section_id = derive_lane_section_id(road.id, idx) + lane_helpers_ = lane_section_to_lane_helpers( + lane_section_id, + lane_section, + reference_line, + lane_section_lengths[idx], + lane_section_lengths[idx + 1], + road.road_types, + interpolation_step_size, + ) + lane_helper_dict.update(lane_helpers_) + + # 3. Update the connections and fill the lane helpers: + # 3.1. From links of the roads + _update_connection_from_links(lane_helper_dict, road_dict) + # 3.2. From junctions + _update_connection_from_junctions(lane_helper_dict, junction_dict, road_dict) + # 3.3. Flip the connections to align to the lane direction + _flip_and_set_connections(lane_helper_dict) + # 3.4. Remove invalid connections based on centerline distances + _post_process_connections(lane_helper_dict, connection_distance_threshold) + + # 4. Collect lane groups from lane helpers + lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = _collect_lane_groups( + lane_helper_dict, junction_dict, road_dict + ) + + # 5. Collect objects, i.e. crosswalks + crosswalk_dict = _collect_crosswalks(opendrive) + + return (road_dict, junction_dict, lane_helper_dict, lane_group_helper_dict, crosswalk_dict) + + +def _update_connection_from_links(lane_helper_dict: Dict[str, OpenDriveLaneHelper], road_dict: Dict[int, Road]) -> None: + """ + Uses the links of the roads to update the connections between lane helpers. + :param lane_helper_dict: Dictionary of lane helpers indexed by lane id. + :param road_dict: Dictionary of roads indexed by road id. + """ + + for lane_id in lane_helper_dict.keys(): + road_idx, lane_section_idx, _, lane_idx = lane_id.split("_") + road_idx, lane_section_idx, lane_idx = int(road_idx), int(lane_section_idx), int(lane_idx) + + road = road_dict[road_idx] + + successor_lane_idx = lane_helper_dict[lane_id].open_drive_lane.successor + if successor_lane_idx is not None: + successor_lane_id: Optional[str] = None + + # Not the last lane section -> Next lane section in same road + if lane_section_idx < road.lanes.last_lane_section_idx: + successor_lane_id = build_lane_id( + road_idx, + lane_section_idx + 1, + successor_lane_idx, + ) + + # Last lane section -> Next road in first lane section + # Try to get next road + elif road.link.successor is not None and road.link.successor.element_type != "junction": + successor_road = road_dict[road.link.successor.element_id] + successor_lane_section_idx = ( + 0 if road.link.successor.contact_point == "start" else successor_road.lanes.last_lane_section_idx + ) + + successor_lane_id = build_lane_id( + successor_road.id, + successor_lane_section_idx, + successor_lane_idx, + ) + + # assert successor_lane_id in lane_helper_dict.keys() + if successor_lane_id is None or successor_lane_id not in lane_helper_dict.keys(): + continue + lane_helper_dict[lane_id].successor_lane_ids.append(successor_lane_id) + lane_helper_dict[successor_lane_id].predecessor_lane_ids.append(lane_id) + + predecessor_lane_idx = lane_helper_dict[lane_id].open_drive_lane.predecessor + if predecessor_lane_idx is not None: + predecessor_lane_id: Optional[str] = None + + # Not the first lane section -> Previous lane section in same road + if lane_section_idx > 0: + predecessor_lane_id = build_lane_id( + road_idx, + lane_section_idx - 1, + predecessor_lane_idx, + ) + + # First lane section -> Previous road + # Try to get previous road + elif road.link.predecessor is not None and road.link.predecessor.element_type != "junction": + predecessor_road = road_dict[road.link.predecessor.element_id] + predecessor_lane_section_idx = ( + 0 + if road.link.predecessor.contact_point == "start" + else predecessor_road.lanes.last_lane_section_idx + ) + + predecessor_lane_id = build_lane_id( + predecessor_road.id, + predecessor_lane_section_idx, + predecessor_lane_idx, + ) + + # assert predecessor_lane_id in lane_helper_dict.keys() + if predecessor_lane_id is None or predecessor_lane_id not in lane_helper_dict.keys(): + continue + lane_helper_dict[lane_id].predecessor_lane_ids.append(predecessor_lane_id) + lane_helper_dict[predecessor_lane_id].successor_lane_ids.append(lane_id) + + +def _update_connection_from_junctions( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], + junction_dict: Dict[int, Junction], + road_dict: Dict[int, Road], +) -> None: + """ + Helper function to update the lane connections based on junctions. + :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. + :param junction_dict: Dictionary mapping junction ids to their objects. + :param road_dict: Dictionary mapping road ids to their objects. + :raises ValueError: If a connection is invalid. + """ + + for junction_idx, junction in junction_dict.items(): + for connection in junction.connections: + incoming_road = road_dict[connection.incoming_road] + connecting_road = road_dict[connection.connecting_road] + + for lane_link in connection.lane_links: + + incoming_lane_id: Optional[str] = None + connecting_lane_id: Optional[str] = None + + if connection.contact_point == "start": + incoming_lane_section_idx = incoming_road.lanes.last_lane_section_idx if lane_link.start < 0 else 0 + incoming_lane_id = build_lane_id(incoming_road.id, incoming_lane_section_idx, lane_link.start) + connecting_lane_id = build_lane_id(connecting_road.id, 0, lane_link.end) + elif connection.contact_point == "end": + incoming_lane_id = build_lane_id(incoming_road.id, 0, lane_link.start) + connecting_lane_id = build_lane_id( + connecting_road.id, + connecting_road.lanes.last_lane_section_idx, + lane_link.end, + ) + else: + raise ValueError(f"Unknown contact point: {connection.contact_point} in junction {junction_idx}") + + if incoming_lane_id is None or connecting_lane_id is None: + logger.debug(f"OpenDRIVE: Lane connection {incoming_lane_id} -> {connecting_lane_id} not valid") + continue + if incoming_lane_id not in lane_helper_dict.keys() or connecting_lane_id not in lane_helper_dict.keys(): + logger.debug( + f"OpenDRIVE: Lane connection {incoming_lane_id} -> {connecting_lane_id} not found in lane_helper_dict" + ) + continue + lane_helper_dict[incoming_lane_id].successor_lane_ids.append(connecting_lane_id) + lane_helper_dict[connecting_lane_id].predecessor_lane_ids.append(incoming_lane_id) + + +def _flip_and_set_connections(lane_helper_dict: Dict[str, OpenDriveLaneHelper]) -> None: + """ + Helper function to flip the connections of the lane helpers, to align them with the lane direction + :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. + """ + + for lane_id in lane_helper_dict.keys(): + if lane_helper_dict[lane_id].id > 0: + successors_temp = lane_helper_dict[lane_id].successor_lane_ids + lane_helper_dict[lane_id].successor_lane_ids = lane_helper_dict[lane_id].predecessor_lane_ids + lane_helper_dict[lane_id].predecessor_lane_ids = successors_temp + lane_helper_dict[lane_id].successor_lane_ids = list(set(lane_helper_dict[lane_id].successor_lane_ids)) + lane_helper_dict[lane_id].predecessor_lane_ids = list(set(lane_helper_dict[lane_id].predecessor_lane_ids)) + + +def _post_process_connections( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], + connection_distance_threshold: float, +) -> None: + """ + Helper function to post-process the connections of the lane helpers, removing invalid connections based on centerline distances. + :param lane_helper_dict: Dictionary mapping lane ids to their helper objects. + :param connection_distance_threshold: Threshold distance for valid connections. + """ + + for lane_id in lane_helper_dict.keys(): + lane_helper_dict[lane_id] + + centerline = lane_helper_dict[lane_id].center_polyline_se2 + + valid_successor_lane_ids: List[str] = [] + for successor_lane_id in lane_helper_dict[lane_id].successor_lane_ids: + successor_centerline = lane_helper_dict[successor_lane_id].center_polyline_se2 + distance = np.linalg.norm(centerline[-1, :2] - successor_centerline[0, :2]) + if distance > connection_distance_threshold: + + logger.debug( + f"OpenDRIVE: Removing connection {lane_id} -> {successor_lane_id} with distance {distance}" + ) + else: + valid_successor_lane_ids.append(successor_lane_id) + lane_helper_dict[lane_id].successor_lane_ids = valid_successor_lane_ids + + valid_predecessor_lane_ids: List[str] = [] + for predecessor_lane_id in lane_helper_dict[lane_id].predecessor_lane_ids: + predecessor_centerline = lane_helper_dict[predecessor_lane_id].center_polyline_se2 + distance = np.linalg.norm(centerline[0, :2] - predecessor_centerline[-1, :2]) + if distance > connection_distance_threshold: + logger.debug( + f"OpenDRIVE: Removing connection {predecessor_lane_id} -> {successor_lane_id} with distance {distance}" + ) + else: + valid_predecessor_lane_ids.append(predecessor_lane_id) + lane_helper_dict[lane_id].predecessor_lane_ids = valid_predecessor_lane_ids + + +def _collect_lane_groups( + lane_helper_dict: Dict[str, OpenDriveLaneHelper], + junction_dict: Dict[int, Junction], + road_dict: Dict[int, Road], +) -> None: + + lane_group_helper_dict: Dict[str, OpenDriveLaneGroupHelper] = {} + + def _collect_lane_helper_of_id(lane_group_id: str) -> List[OpenDriveLaneHelper]: + lane_helpers: List[OpenDriveLaneHelper] = [] + for lane_id, lane_helper in lane_helper_dict.items(): + if (lane_helper.type in ["driving"]) and (lane_group_id_from_lane_id(lane_id) == lane_group_id): + lane_helpers.append(lane_helper) + return lane_helpers + + def _collect_lane_group_ids_of_road(road_id: int) -> List[str]: + lane_group_ids: List[str] = [] + for lane_group_id in lane_group_helper_dict.keys(): + if int(road_id_from_lane_group_id(lane_group_id)) == road_id: + lane_group_ids.append(lane_group_id) + return lane_group_ids + + all_lane_group_ids = list(set([lane_group_id_from_lane_id(lane_id) for lane_id in lane_helper_dict.keys()])) + + for lane_group_id in all_lane_group_ids: + lane_group_lane_helper = _collect_lane_helper_of_id(lane_group_id) + if len(lane_group_lane_helper) >= 1: + lane_group_helper_dict[lane_group_id] = OpenDriveLaneGroupHelper(lane_group_id, lane_group_lane_helper) + + for junction in junction_dict.values(): + for connection in junction.connections: + connecting_road = road_dict[connection.connecting_road] + connecting_lane_group_ids = _collect_lane_group_ids_of_road(connecting_road.id) + for connecting_lane_group_id in connecting_lane_group_ids: + if connecting_lane_group_id in lane_group_helper_dict.keys(): + lane_group_helper_dict[connecting_lane_group_id].junction_id = junction.id + + return lane_group_helper_dict + + +def _collect_crosswalks(opendrive: OpenDrive) -> Dict[int, OpenDriveObjectHelper]: + + object_helper_dict: Dict[int, OpenDriveObjectHelper] = {} + for road in opendrive.roads: + if len(road.objects) == 0: + continue + reference_line = ReferenceLine.from_plan_view( + road.plan_view, + road.lanes.lane_offsets, + road.elevation_profile.elevations, + ) + for object in road.objects: + if object.type in ["crosswalk"]: + object_helper = get_object_helper(object, reference_line) + object_helper_dict[object_helper.object_id] = object_helper + + return object_helper_dict diff --git a/d123/dataset/conversion/map/opendrive/id_mapping.py b/d123/dataset/conversion/map/opendrive/utils/id_mapping.py similarity index 100% rename from d123/dataset/conversion/map/opendrive/id_mapping.py rename to d123/dataset/conversion/map/opendrive/utils/id_mapping.py diff --git a/d123/dataset/conversion/map/opendrive/conversion/id_system.py b/d123/dataset/conversion/map/opendrive/utils/id_system.py similarity index 100% rename from d123/dataset/conversion/map/opendrive/conversion/id_system.py rename to d123/dataset/conversion/map/opendrive/utils/id_system.py diff --git a/d123/dataset/conversion/map/opendrive/conversion/group_collections.py b/d123/dataset/conversion/map/opendrive/utils/lane_helper.py similarity index 70% rename from d123/dataset/conversion/map/opendrive/conversion/group_collections.py rename to d123/dataset/conversion/map/opendrive/utils/lane_helper.py index 6e577305..79e0f0c1 100644 --- a/d123/dataset/conversion/map/opendrive/conversion/group_collections.py +++ b/d123/dataset/conversion/map/opendrive/utils/lane_helper.py @@ -8,19 +8,14 @@ from d123.common.geometry.base import StateSE2Index from d123.common.geometry.units import kmph_to_mps, mph_to_mps -from d123.dataset.conversion.map.opendrive.conversion.id_system import ( +from d123.dataset.conversion.map.opendrive.parser.lane import Lane, LaneSection +from d123.dataset.conversion.map.opendrive.parser.reference import ReferenceLine +from d123.dataset.conversion.map.opendrive.parser.road import RoadType +from d123.dataset.conversion.map.opendrive.utils.id_system import ( derive_lane_group_id, derive_lane_id, lane_group_id_from_lane_id, ) -from d123.dataset.conversion.map.opendrive.elements.lane import Lane, LaneSection -from d123.dataset.conversion.map.opendrive.elements.reference import Border -from d123.dataset.conversion.map.opendrive.elements.road import RoadType - -# TODO: Add to config -STEP_SIZE = 1.0 - -# TODO: make naming consistent with objects_collections.py @dataclass @@ -30,9 +25,10 @@ class OpenDriveLaneHelper: open_drive_lane: Lane s_inner_offset: float s_range: Tuple[float, float] - inner_border: Border - outer_border: Border + inner_boundary: ReferenceLine + outer_boundary: ReferenceLine speed_limit_mps: Optional[float] + interpolation_step_size: float # lazy loaded predecessor_lane_ids: Optional[List[str]] = None @@ -54,23 +50,27 @@ def type(self) -> str: def _s_positions(self) -> npt.NDArray[np.float64]: length = self.s_range[1] - self.s_range[0] _s_positions = np.linspace( - self.s_range[0], self.s_range[1], int(np.ceil(length / STEP_SIZE)) + 1, endpoint=True, dtype=np.float64 + self.s_range[0], + self.s_range[1], + int(np.ceil(length / self.interpolation_step_size)) + 1, + endpoint=True, + dtype=np.float64, ) _s_positions[..., -1] = np.clip(_s_positions[..., -1], 0.0, self.s_range[-1]) return _s_positions @cached_property - def _is_last_mask(self) -> npt.NDArray[np.float64]: - is_last_mask = np.zeros(len(self._s_positions), dtype=bool) - is_last_mask[-1] = True - return is_last_mask + def _lane_section_end_mask(self) -> npt.NDArray[np.float64]: + lane_section_end_mask = np.zeros(len(self._s_positions), dtype=bool) + lane_section_end_mask[-1] = True + return lane_section_end_mask @cached_property def inner_polyline_se2(self) -> npt.NDArray[np.float64]: inner_polyline = np.array( [ - self.inner_border.interpolate_se2(self.s_inner_offset + s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.inner_boundary.interpolate_se2(self.s_inner_offset + s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -80,8 +80,8 @@ def inner_polyline_se2(self) -> npt.NDArray[np.float64]: def inner_polyline_3d(self) -> npt.NDArray[np.float64]: inner_polyline = np.array( [ - self.inner_border.interpolate_3d(self.s_inner_offset + s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.inner_boundary.interpolate_3d(self.s_inner_offset + s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -91,8 +91,8 @@ def inner_polyline_3d(self) -> npt.NDArray[np.float64]: def outer_polyline_se2(self) -> npt.NDArray[np.float64]: outer_polyline = np.array( [ - self.outer_border.interpolate_se2(s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.outer_boundary.interpolate_se2(s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -102,8 +102,8 @@ def outer_polyline_se2(self) -> npt.NDArray[np.float64]: def outer_polyline_3d(self) -> npt.NDArray[np.float64]: outer_polyline = np.array( [ - self.outer_border.interpolate_3d(s - self.s_range[0], is_last_pos=is_last) - for s, is_last in zip(self._s_positions, self._is_last_mask) + self.outer_boundary.interpolate_3d(s - self.s_range[0], lane_section_end=end) + for s, end in zip(self._s_positions, self._lane_section_end_mask) ], dtype=np.float64, ) @@ -198,75 +198,53 @@ def shapely_polygon(self) -> shapely.Polygon: def lane_section_to_lane_helpers( lane_section_id: str, lane_section: LaneSection, - reference_border: Border, + reference_line: ReferenceLine, s_min: float, s_max: float, road_types: List[RoadType], + interpolation_step_size: float, ) -> Dict[str, OpenDriveLaneHelper]: lane_helpers: Dict[str, OpenDriveLaneHelper] = {} - for side in ["right", "left"]: + for lanes, t_sign, side in zip([lane_section.left_lanes, lane_section.right_lanes], [1.0, -1.0], ["left", "right"]): lane_group_id = derive_lane_group_id(lane_section_id, side) - lanes = lane_section.right_lanes if side == "right" else lane_section.left_lanes - coeff_factor = -1.0 if side == "right" else 1.0 - - lane_borders = [reference_border] + lane_boundaries = [reference_line] for lane in lanes: lane_id = derive_lane_id(lane_group_id, lane.id) - s_inner_offset = lane_section.s if len(lane_borders) == 1 else 0.0 - lane_borders.append(_create_outer_lane_border(lane_borders, lane_section, lane, coeff_factor)) + s_inner_offset = lane_section.s if len(lane_boundaries) == 1 else 0.0 + lane_boundaries.append( + ReferenceLine.from_reference_line( + reference_line=lane_boundaries[-1], + widths=lane.widths, + s_offset=s_inner_offset, + t_sign=t_sign, + ) + ) lane_helper = OpenDriveLaneHelper( lane_id=lane_id, open_drive_lane=lane, s_inner_offset=s_inner_offset, s_range=(s_min, s_max), - inner_border=lane_borders[-2], - outer_border=lane_borders[-1], - speed_limit_mps=_get_speed_limit_mps(s_min, s_max, road_types), + inner_boundary=lane_boundaries[-2], + outer_boundary=lane_boundaries[-1], + speed_limit_mps=_get_speed_limit_mps(s_min, road_types), + interpolation_step_size=interpolation_step_size, ) lane_helpers[lane_id] = lane_helper return lane_helpers -def _create_outer_lane_border( - lane_borders: List[Border], - lane_section: LaneSection, - lane: Lane, - coeff_factor: float, -) -> Border: - # NOTE: This function needs to be re-written because of licensing issues. - - args = {} - if len(lane_borders) == 1: - args["s_offset"] = lane_section.s - - args["reference"] = lane_borders[-1] - args["elevation_profile"] = lane_borders[-1].elevation_profile - - width_coefficient_offsets = [] - width_coefficients = [] - - for width in lane.widths: - width_coefficient_offsets.append(width.s_offset) - width_coefficients.append([x * coeff_factor for x in width.polynomial_coefficients]) - - args["width_coefficient_offsets"] = width_coefficient_offsets - args["width_coefficients"] = width_coefficients - return Border(**args) - - -def _get_speed_limit_mps(s_min: float, s_max: float, road_types: List[RoadType]) -> Optional[float]: +def _get_speed_limit_mps(s: float, road_types: List[RoadType]) -> Optional[float]: # NOTE: Likely not correct way to extract speed limit from CARLA maps, but serves as a placeholder speed_limit_mps: Optional[float] = None s_road_types = [road_type.s for road_type in road_types] + [float("inf")] if len(road_types) > 0: - # 1. Find current road type for road_type_idx, road_type in enumerate(road_types): - if s_min >= road_type.s and s_min < s_road_types[road_type_idx + 1]: + if s >= road_type.s and s < s_road_types[road_type_idx + 1]: if road_type.speed is not None: if road_type.speed.unit == "mps": speed_limit_mps = road_type.speed.max diff --git a/d123/dataset/conversion/map/opendrive/conversion/objects_collections.py b/d123/dataset/conversion/map/opendrive/utils/objects_helper.py similarity index 85% rename from d123/dataset/conversion/map/opendrive/conversion/objects_collections.py rename to d123/dataset/conversion/map/opendrive/utils/objects_helper.py index 1eb52f77..3cbb569e 100644 --- a/d123/dataset/conversion/map/opendrive/conversion/objects_collections.py +++ b/d123/dataset/conversion/map/opendrive/utils/objects_helper.py @@ -8,8 +8,8 @@ from d123.common.geometry.base import Point2D, Point3D, Point3DIndex, StateSE2 from d123.common.geometry.transform.tranform_2d import translate_along_yaw from d123.common.geometry.utils import normalize_angle -from d123.dataset.conversion.map.opendrive.elements.objects import Object -from d123.dataset.conversion.map.opendrive.elements.reference import Border +from d123.dataset.conversion.map.opendrive.parser.objects import Object +from d123.dataset.conversion.map.opendrive.parser.reference import ReferenceLine # TODO: make naming consistent with group_collections.py @@ -29,14 +29,14 @@ def shapely_polygon(self) -> shapely.Polygon: return shapely.geometry.Polygon(self.outline_3d[:, Point3DIndex.XY]) -def get_object_helper(object: Object, reference_border: Border) -> OpenDriveObjectHelper: +def get_object_helper(object: Object, reference_line: ReferenceLine) -> OpenDriveObjectHelper: object_helper: Optional[OpenDriveObjectHelper] = None # 1. Extract object position in frenet frame of the reference line - object_se2: StateSE2 = StateSE2.from_array(reference_border.interpolate_se2(s=object.s, t=object.t)) - object_3d: Point3D = Point3D.from_array(reference_border.interpolate_3d(s=object.s, t=object.t)) + object_se2: StateSE2 = StateSE2.from_array(reference_line.interpolate_se2(s=object.s, t=object.t)) + object_3d: Point3D = Point3D.from_array(reference_line.interpolate_3d(s=object.s, t=object.t)) # Adjust yaw angle from object data object_se2.yaw = normalize_angle(object_se2.yaw + object.hdg) diff --git a/d123/dataset/dataset_specific/carla/carla_data_converter.py b/d123/dataset/dataset_specific/carla/carla_data_converter.py index 81989008..f598bb1e 100644 --- a/d123/dataset/dataset_specific/carla/carla_data_converter.py +++ b/d123/dataset/dataset_specific/carla/carla_data_converter.py @@ -21,8 +21,7 @@ from d123.common.geometry.vector import Vector3DIndex from d123.common.multithreading.worker_utils import WorkerPool, worker_map from d123.dataset.arrow.helper import open_arrow_table, write_arrow_table -from d123.dataset.conversion.map.opendrive.elements.opendrive import OpenDrive -from d123.dataset.conversion.map.opendrive.opendrive_converter import OpenDriveConverter +from d123.dataset.conversion.map.opendrive.opendrive_map_conversion import convert_from_xodr from d123.dataset.dataset_specific.raw_data_converter import DataConverterConfig, RawDataConverter from d123.dataset.logs.log_metadata import LogMetadata from d123.dataset.maps.abstract_map import AbstractMap, MapLayer @@ -130,6 +129,10 @@ def convert_logs(self, worker: WorkerPool) -> None: def convert_carla_map_to_gpkg(map_names: List[str], data_converter_config: DataConverterConfig) -> List[Any]: + + # TODO: add to config + _interpolation_step_size = 0.5 # [m] + _connection_distance_threshold = 0.1 # [m] for map_name in map_names: map_path = data_converter_config.output_path / "maps" / f"carla_{map_name.lower()}.gpkg" if data_converter_config.force_map_conversion or not map_path.exists(): @@ -145,8 +148,12 @@ def convert_carla_map_to_gpkg(map_names: List[str], data_converter_config: DataC CARLA_ROOT / "CarlaUE4" / "Content" / "Carla" / "Maps" / map_name / "OpenDrive" / f"{map_name}.xodr" ) - OpenDriveConverter(OpenDrive.parse_from_file(carla_map_path)).run(f"carla_{map_name.lower()}") - + convert_from_xodr( + carla_map_path, + f"carla_{map_name.lower()}", + _interpolation_step_size, + _connection_distance_threshold, + ) return [] diff --git a/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb index 6eb9ef23..7687c8e2 100644 --- a/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb +++ b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb @@ -6,14 +6,14 @@ "metadata": {}, "outputs": [], "source": [ - "import matplotlib.pyplot as plt\n", - "from shapely.geometry import LineString, Polygon, Point\n", - "import numpy as np\n", + "# import matplotlib.pyplot as plt\n", + "# from shapely.geometry import LineString, Polygon, Point\n", + "# import numpy as np\n", "\n", - "from typing import List\n", + "# # from typing import List\n", "\n", - "from d123.dataset.conversion.map.opendrive.elements.lane import Lane, LaneSection\n", - "from d123.dataset.conversion.map.opendrive.elements.reference import Border" + "# # from d123.dataset.conversion.map.opendrive.elements.lane import Lane, LaneSection\n", + "# # from d123.dataset.conversion.map.opendrive.elements.reference import Border" ] }, { @@ -85,7 +85,6 @@ "\n", "MAP_ROOT = Path(\"/home/daniel/d123_workspace/data\") / \"maps\"\n", "\n", - "\n", "for town_name in AVAILABLE_CARLA_MAP_LOCATIONS:\n", " map_name = town_name.lower()\n", " print(f\"Start {map_name} map...\")\n", @@ -110,6 +109,7 @@ " converter = OpenDriveConverter(opendrive)\n", " converter.run(f\"carla_{town_name.lower()}\")\n", " print(f\"Converting {map_name} map... done\")\n", + " # break\n", "\n", "\n" ] @@ -129,24 +129,24 @@ "metadata": {}, "outputs": [], "source": [ - "# from d123.dataset.maps.map_datatypes import MapLayer\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", "\n", "\n", - "# str(MapLayer.GENERIC_DRIVABLE).split(\".\")[-1].lower()\n", + "str(MapLayer.GENERIC_DRIVABLE).split(\".\")[-1].lower()\n", "\n", "\n", - "# MapLayer.GENERIC_DRIVABLE.name\n", + "MapLayer.GENERIC_DRIVABLE.name\n", "\n", - "# MapLayer.deserialize(MapLayer.GENERIC_DRIVABLE.name)\n", + "MapLayer.deserialize(MapLayer.GENERIC_DRIVABLE.name)\n", "\n", "\n", - "# MapLayer.GENERIC_DRIVABLE.name.lower().islower()\n", + "MapLayer.GENERIC_DRIVABLE.name.lower().islower()\n", "\n", "\n", - "# AVAILABLE_MAP_LAYERS = list(MapLayer)\n", - "# AVAILABLE_MAP_LAYERS\n", + "AVAILABLE_MAP_LAYERS = list(MapLayer)\n", + "AVAILABLE_MAP_LAYERS\n", "\n", - "# pyogrio.read_dataframe()." + "pyogrio.read_dataframe()." ] }, { diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index c9e58b07..7155776a 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -26,9 +26,9 @@ "\n", "\n", "# splits = [\"wopd_train\"]\n", - "# splits = [\"carla\"]\n", + "splits = [\"carla\"]\n", "# splits = [\"nuplan_private_test\"]\n", - "splits = [\"av2-sensor-mini_train\"]\n", + "# splits = [\"av2-sensor-mini_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -145,10 +145,10 @@ " MapLayer.LANE,\n", " MapLayer.LANE_GROUP,\n", " MapLayer.GENERIC_DRIVABLE,\n", - " # MapLayer.CARPARK,\n", + " MapLayer.CARPARK,\n", " MapLayer.CROSSWALK,\n", " # MapLayer.INTERSECTION,\n", - " # MapLayer.WALKWAY,\n", + " MapLayer.WALKWAY,\n", " MapLayer.ROAD_EDGE,\n", " # MapLayer.ROAD_LINE,\n", " ]\n", @@ -161,6 +161,9 @@ " for layer, map_objects in map_objects_dict.items():\n", " for map_object in map_objects:\n", " try:\n", + " if layer in [MapLayer.WALKWAY]:\n", + " print(map_object.trimesh_mesh.vertices)\n", + "\n", " if layer in [\n", " MapLayer.GENERIC_DRIVABLE,\n", " MapLayer.CARPARK,\n", @@ -247,8 +250,8 @@ " return fig, ax\n", "\n", "\n", - "scene_index = 3\n", - "iteration = 1\n", + "scene_index = 10\n", + "iteration = 2\n", "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=iteration, radius=100)\n", "plt.show()\n", "\n", diff --git a/notebooks/viz/bev_matplotlib_prediction.ipynb b/notebooks/viz/bev_matplotlib_prediction.ipynb new file mode 100644 index 00000000..041daa72 --- /dev/null +++ b/notebooks/viz/bev_matplotlib_prediction.ipynb @@ -0,0 +1,246 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.scene.scene_builder import ArrowSceneBuilder\n", + "from d123.dataset.scene.scene_filter import SceneFilter\n", + "\n", + "from d123.common.multithreading.worker_sequential import Sequential\n", + "from d123.common.datatypes.sensor.camera import CameraType" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"nuplan_private_test\"\n", + "# log_names = [\"2021.09.29.17.35.58_veh-44_00066_00432\"]\n", + "\n", + "\n", + "# splits = [\"wopd_train\"]\n", + "# splits = [\"carla\"]\n", + "# splits = [\"nuplan_private_test\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", + "# log_names = None\n", + "\n", + "\n", + "log_names = None\n", + "scene_tokens = None\n", + "\n", + "scene_filter = SceneFilter(\n", + " split_names=splits,\n", + " log_names=log_names,\n", + " scene_tokens=scene_tokens,\n", + " duration_s=8.0,\n", + " history_s=0.0,\n", + " timestamp_threshold_s=4.0,\n", + " shuffle=True,\n", + " # camera_types=[CameraType.CAM_F0],\n", + ")\n", + "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", + "worker = Sequential()\n", + "# worker = RayDistributed()\n", + "scenes = scene_builder.get_scenes(scene_filter, worker)\n", + "\n", + "print(f\"Found {len(scenes)} scenes\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from pathlib import Path\n", + "from typing import List, Optional, Tuple\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from d123.common.geometry.base import Point2D\n", + "from d123.common.visualization.color.color import BLACK, DARK_GREY, DARKER_GREY, LIGHT_GREY, NEW_TAB_10, TAB_10\n", + "from d123.common.visualization.color.config import PlotConfig\n", + "from d123.common.visualization.color.default import CENTERLINE_CONFIG, MAP_SURFACE_CONFIG, ROUTE_CONFIG\n", + "from d123.common.visualization.matplotlib.observation import (\n", + " add_box_detections_to_ax,\n", + " add_box_future_detections_to_ax,\n", + " add_default_map_on_ax,\n", + " add_ego_vehicle_to_ax,\n", + " add_traffic_lights_to_ax,\n", + ")\n", + "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", + "from d123.dataset.maps.abstract_map import AbstractMap\n", + "from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractLaneGroup\n", + "from d123.dataset.maps.map_datatypes import MapLayer\n", + "from d123.dataset.scene.abstract_scene import AbstractScene\n", + "\n", + "\n", + "\n", + "def _plot_scene_on_ax(ax: plt.Axes, scene: AbstractScene, iteration: int = 0, radius: float = 80) -> plt.Axes:\n", + "\n", + " ego_vehicle_state = scene.get_ego_state_at_iteration(iteration)\n", + " box_detections = scene.get_box_detections_at_iteration(iteration)\n", + "\n", + " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", + " # add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radiuss, route_lane_group_ids=None)\n", + " add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", + "\n", + " add_box_future_detections_to_ax(ax, scene, iteration=iteration)\n", + " add_box_detections_to_ax(ax, box_detections)\n", + " add_ego_vehicle_to_ax(ax, ego_vehicle_state)\n", + "\n", + " zoom = 1.0\n", + " ax.set_xlim(point_2d.x - radius * zoom, point_2d.x + radius * zoom)\n", + " ax.set_ylim(point_2d.y - radius * zoom, point_2d.y + radius * zoom)\n", + "\n", + " ax.set_aspect(\"equal\", adjustable=\"box\")\n", + " return ax\n", + "\n", + "\n", + "def plot_scene_at_iteration(\n", + " scene: AbstractScene, iteration: int = 0, radius: float = 80\n", + ") -> Tuple[plt.Figure, plt.Axes]:\n", + "\n", + " size = 10\n", + "\n", + " fig, ax = plt.subplots(figsize=(size, size))\n", + " _plot_scene_on_ax(ax, scene, iteration, radius)\n", + " return fig, ax\n", + "\n", + "\n", + "scene_index = 0\n", + "\n", + "for i, scene in enumerate(scenes):\n", + "\n", + " iteration = 0\n", + " fig, ax = plot_scene_at_iteration(scenes[i], iteration=iteration, radius=60)\n", + " Path(f\"/home/daniel/examples/{splits[0]}/\").mkdir(exist_ok=True, parents=True)\n", + " ax.set_xticks([])\n", + " ax.set_yticks([])\n", + " fig.tight_layout()\n", + " fig.savefig(\n", + " f\"/home/daniel/examples/{splits[0]}/{scene.log_name}_{i}.pdf\",\n", + " dpi=300,\n", + " bbox_inches=\"tight\",\n", + " )\n", + "\n", + "# camera = scenes[scene_index].get_camera_at_iteration(\n", + "# iteration=iteration, camera_type=CameraType.CAM_F0\n", + "# )\n", + "\n", + "# plt.imshow(camera.image, cmap=\"gray\", vmin=0, vmax=255)\n", + "# # fig.savefig(f\"/home/daniel/scene_{scene_index}_iteration_1.pdf\", dpi=300, bbox_inches=\"tight\")\n", + "\n", + "# scenes[scene_index].log_name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "scene = scenes[scene_index]\n", + "\n", + "\n", + "scene.get_camera_at_iteration(camera_type=CameraType.CAM_F0, iteration=0)\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import get_road_edge_linear_rings\n", + "from d123.dataset.maps.gpkg.gpkg_map import GPKGMap\n", + "\n", + "\n", + "map_api: GPKGMap = scenes[scene_index].map_api\n", + "\n", + "drivable_polygons = map_api._gpd_dataframes[MapLayer.LANE]\n", + "\n", + "\n", + "\n", + "linear_rings = get_road_edge_linear_rings(drivable_polygons.geometry.tolist())\n", + "rings_lengths = [ring.length for ring in linear_rings]\n", + "max_length_idx = np.argmax(rings_lengths)\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "size = 16\n", + "fig, ax = plt.subplots(figsize=(size, size))\n", + "\n", + "for idx, ring in enumerate(linear_rings):\n", + " if idx == max_length_idx:\n", + " ax.plot(*ring.xy, color=\"black\", linewidth=2, label=\"Longest Road Edge\")\n", + " else:\n", + " ax.plot(*ring.xy)\n", + "\n", + "\n", + "ax.set_aspect(\"equal\", adjustable=\"box\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/viz/viser_testing_v2_scene.ipynb b/notebooks/viz/viser_testing_v2_scene.ipynb index 567edd06..07d91d1d 100644 --- a/notebooks/viz/viser_testing_v2_scene.ipynb +++ b/notebooks/viz/viser_testing_v2_scene.ipynb @@ -28,9 +28,9 @@ "\n", "\n", "# splits = [\"nuplan_private_test\"]\n", - "# splits = [\"carla\"]\n", + "splits = [\"carla\"]\n", "# splits = [\"wopd_train\"]\n", - "splits = [\"av2-sensor-mini_train\"]\n", + "# splits = [\"av2-sensor-mini_train\"]\n", "log_names = None\n", "\n", "scene_tokens = None\n", @@ -39,10 +39,10 @@ " split_names=splits,\n", " log_names=log_names,\n", " scene_tokens=scene_tokens,\n", - " duration_s=15,\n", + " duration_s=10,\n", " history_s=0.0,\n", - " timestamp_threshold_s=20,\n", - " shuffle=True,\n", + " timestamp_threshold_s=10,\n", + " shuffle=False,\n", " camera_types=[CameraType.CAM_F0],\n", ")\n", "scene_builder = ArrowSceneBuilder(\"/home/daniel/d123_workspace/data\")\n", @@ -71,7 +71,7 @@ "from d123.common.visualization.viser.server import ViserVisualizationServer\n", "\n", "\n", - "visualization_server = ViserVisualizationServer(scenes, scene_index=1)" + "visualization_server = ViserVisualizationServer(scenes, scene_index=0)" ] }, { @@ -99,7 +99,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123_dev", + "display_name": "d123", "language": "python", "name": "python3" }, From 2ee7bcf4189dd3ec49b2af7554003dc0f848106b Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 22 Aug 2025 11:39:55 +0200 Subject: [PATCH 40/42] Add intersection to argoverse 2 sensor map --- d123/common/geometry/line/polylines.py | 12 +- .../av2/av2_map_conversion.py | 219 ++++++++++--- notebooks/av2/delete_me.ipynb | 4 +- notebooks/av2/delete_me_map.ipynb | 310 ++++++++++++++++++ .../test_opendrive_conversion.ipynb | 2 +- notebooks/viz/bev_matplotlib.ipynb | 34 +- 6 files changed, 504 insertions(+), 77 deletions(-) create mode 100644 notebooks/av2/delete_me_map.ipynb diff --git a/d123/common/geometry/line/polylines.py b/d123/common/geometry/line/polylines.py index 6adb11fc..37b799b5 100644 --- a/d123/common/geometry/line/polylines.py +++ b/d123/common/geometry/line/polylines.py @@ -9,7 +9,7 @@ import shapely.geometry as geom from scipy.interpolate import interp1d -from d123.common.geometry.base import Point2D, Point2DIndex, Point3D, StateSE2, StateSE2Index +from d123.common.geometry.base import Point2D, Point2DIndex, Point3D, Point3DIndex, StateSE2, StateSE2Index from d123.common.geometry.constants import DEFAULT_Z from d123.common.geometry.line.helper import get_linestring_yaws, get_path_progress from d123.common.geometry.utils import normalize_angle @@ -33,7 +33,15 @@ def from_linestring(cls, linestring: geom.LineString) -> Polyline2D: @classmethod def from_array(cls, polyline_array: npt.NDArray[np.float32]) -> Polyline2D: - raise NotImplementedError + assert polyline_array.ndim == 2 + linestring: Optional[geom.LineString] = None + if polyline_array.shape[-1] == len(Point2DIndex): + linestring = geom_creation.linestrings(polyline_array) + elif polyline_array.shape[-1] == len(Point3DIndex): + linestring = geom_creation.linestrings(polyline_array[:, Point3DIndex.XY]) + else: + raise ValueError("Array must have shape (N, 2) or (N, 3) for Point2D or Point3D respectively.") + return Polyline2D(linestring) @property def array(self) -> npt.NDArray[np.float64]: diff --git a/d123/dataset/dataset_specific/av2/av2_map_conversion.py b/d123/dataset/dataset_specific/av2/av2_map_conversion.py index b294573c..98e5fc78 100644 --- a/d123/dataset/dataset_specific/av2/av2_map_conversion.py +++ b/d123/dataset/dataset_specific/av2/av2_map_conversion.py @@ -5,11 +5,13 @@ import numpy as np import numpy.typing as npt import pandas as pd +import shapely import shapely.geometry as geom from flask import json from d123.common.geometry.base import Point3DIndex -from d123.common.geometry.line.polylines import Polyline3D +from d123.common.geometry.line.polylines import Polyline2D, Polyline3D +from d123.common.geometry.occupancy_map import OccupancyMap2D from d123.dataset.conversion.map.road_edge.road_edge_2d_utils import split_line_geometry_by_max_length from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import ( get_road_edges_3d_from_generic_drivable_area_df, @@ -32,10 +34,7 @@ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Poly return Polyline3D.from_array(polyline) map_folder = source_log_path / "map" - - next(map_folder.glob("*.npy")) log_map_archive_path = next(map_folder.glob("log_map_archive_*.json")) - next(map_folder.glob("*img_Sim2_city.json*")) with open(log_map_archive_path, "r") as f: log_map_archive = json.load(f) @@ -71,15 +70,18 @@ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Poly p3, p4 = np.array([[p["x"], p["y"], p["z"]] for p in crosswalk_dict["edge2"]], dtype=np.float64) crosswalk_dict["outline"] = Polyline3D.from_array(np.array([p1, p2, p4, p3, p1], dtype=np.float64)) + lane_group_dict = _extract_lane_group_dict(log_map_archive["lane_segments"]) + intersection_dict = _extract_intersection_dict(log_map_archive["lane_segments"], lane_group_dict) + lane_df = get_lane_df(log_map_archive["lane_segments"]) - lane_group_df = get_lane_group_df(log_map_archive["lane_segments"]) - get_empty_gdf() + lane_group_df = get_lane_group_df(lane_group_dict) + intersection_df = get_intersections_df(intersection_dict) crosswalk_df = get_crosswalk_df(log_map_archive["pedestrian_crossings"]) - get_empty_gdf() - get_empty_gdf() + walkway_df = get_empty_gdf() # NOTE: AV2 does not provide walkways, so we create an empty DataFrame. + carpark_df = get_empty_gdf() # NOTE: AV2 does not provide carparks, so we create an empty DataFrame. generic_drivable_df = get_generic_drivable_df(drivable_areas) road_edge_df = get_road_edge_df(generic_drivable_df) - get_empty_gdf() + # road_line_df = get_empty_gdf() map_file_path.unlink(missing_ok=True) if not map_file_path.parent.exists(): @@ -87,10 +89,10 @@ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Poly lane_df.to_file(map_file_path, layer=MapLayer.LANE.serialize(), driver="GPKG") lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") - # intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") + intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") - # walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") - # carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") + walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") + carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") generic_drivable_df.to_file(map_file_path, layer=MapLayer.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a") road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") # road_line_df.to_file(map_file_path, layer=MapLayer.ROAD_LINE.serialize(), driver="GPKG", mode="a") @@ -193,12 +195,9 @@ def _get_centerline_from_boundaries( return gdf -def get_lane_group_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: - - lane_group_sets = find_lane_groups(lanes) - lane_group_set_dict = {i: lane_group for i, lane_group in enumerate(lane_group_sets)} +def get_lane_group_df(lane_group_dict: Dict[int, Any]) -> gpd.GeoDataFrame: - ids = list(lane_group_set_dict.keys()) + ids = list(lane_group_dict.keys()) lane_ids = [] intersection_ids = [] predecessor_lane_group_ids = [] @@ -207,38 +206,21 @@ def get_lane_group_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: right_boundaries = [] geometries = [] - def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: - """Helper to find lane group ids that contain any of the given lane ids.""" - lane_group_ids_ = [] - for lane_group_id_, lane_group_set_ in lane_group_set_dict.items(): - if any(str(lane_id) in lane_group_set_ for lane_id in lane_ids): - lane_group_ids_.append(lane_group_id_) - return list(set(lane_group_ids_)) - - for lane_group_id, lane_group_set in lane_group_set_dict.items(): - - lane_ids.append([int(lane_id) for lane_id in lane_group_set]) - intersection_ids.append(None) # NOTE: AV2 doesn't have explicit intersection objects. - - successor_lanes = [] - predecessor_lanes = [] - for lane_id in lane_group_set: - lane_dict = lanes[str(lane_id)] - successor_lanes.extend(lane_dict["successors"]) - predecessor_lanes.extend(lane_dict["predecessors"]) + for lane_group_id, lane_group_values in lane_group_dict.items(): - left_boundary = lanes[lane_group_set[0]]["left_lane_boundary"] - right_boundary = lanes[lane_group_set[-1]]["right_lane_boundary"] + lane_ids.append(lane_group_values["lane_ids"]) + intersection_ids.append(lane_group_values["intersection_id"]) - predecessor_lane_group_ids.append(_get_lane_group_ids_of_lanes_ids(predecessor_lanes)) - successor_lane_group_ids.append(_get_lane_group_ids_of_lanes_ids(successor_lanes)) - left_boundaries.append(left_boundary.linestring) - right_boundaries.append(right_boundary.linestring) + predecessor_lane_group_ids.append(lane_group_values["predecessor_ids"]) + successor_lane_group_ids.append(lane_group_values["successor_ids"]) + left_boundaries.append(lane_group_values["left_boundary"].linestring) + right_boundaries.append(lane_group_values["right_boundary"].linestring) geometry = geom.Polygon( np.vstack( [ - left_boundary.array[:, :2], - right_boundary.array[:, :2][::-1], + lane_group_values["left_boundary"].array[:, :2], + lane_group_values["right_boundary"].array[:, :2][::-1], + lane_group_values["left_boundary"].array[0, :2][None, ...], ] ) ) @@ -259,13 +241,19 @@ def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: return gdf -def get_intersections_df() -> gpd.GeoDataFrame: +def get_intersections_df(intersection_dict: Dict[int, Any]) -> gpd.GeoDataFrame: ids = [] lane_group_ids = [] + outlines = [] geometries = [] - # NOTE: WOPD does not provide intersections, so we create an empty DataFrame. - data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids}) + for intersection_id, intersection_values in intersection_dict.items(): + ids.append(intersection_id) + lane_group_ids.append(intersection_values["lane_group_ids"]) + outlines.append(intersection_values["outline_3d"].linestring) + geometries.append(geom.Polygon(intersection_values["outline_3d"].array[:, Point3DIndex.XY])) + + data = pd.DataFrame({"id": ids, "lane_group_ids": lane_group_ids, "outline": outlines}) gdf = gpd.GeoDataFrame(data, geometry=geometries) return gdf @@ -341,7 +329,56 @@ def get_road_line_df( return gdf -def find_lane_groups(lanes) -> List[List[str]]: +def _extract_lane_group_dict(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: + + lane_group_sets = _extract_lane_group(lanes) + lane_group_set_dict = {i: lane_group for i, lane_group in enumerate(lane_group_sets)} + + lane_group_dict: Dict[int, Dict[str, Any]] = {} + + def _get_lane_group_ids_of_lanes_ids(lane_ids: List[str]) -> List[int]: + """Helper to find lane group ids that contain any of the given lane ids.""" + lane_group_ids_ = [] + for lane_group_id_, lane_group_set_ in lane_group_set_dict.items(): + if any(str(lane_id) in lane_group_set_ for lane_id in lane_ids): + lane_group_ids_.append(lane_group_id_) + return list(set(lane_group_ids_)) + + for lane_group_id, lane_group_set in lane_group_set_dict.items(): + + lane_group_dict[lane_group_id] = {} + lane_group_dict[lane_group_id]["id"] = lane_group_id + lane_group_dict[lane_group_id]["lane_ids"] = [int(lane_id) for lane_id in lane_group_set] + + successor_lanes = [] + predecessor_lanes = [] + for lane_id in lane_group_set: + lane_dict = lanes[str(lane_id)] + successor_lanes.extend(lane_dict["successors"]) + predecessor_lanes.extend(lane_dict["predecessors"]) + + left_boundary = lanes[lane_group_set[0]]["left_lane_boundary"] + right_boundary = lanes[lane_group_set[-1]]["right_lane_boundary"] + + lane_group_dict[lane_group_id]["intersection_id"] = None + lane_group_dict[lane_group_id]["predecessor_ids"] = _get_lane_group_ids_of_lanes_ids(predecessor_lanes) + lane_group_dict[lane_group_id]["successor_ids"] = _get_lane_group_ids_of_lanes_ids(successor_lanes) + lane_group_dict[lane_group_id]["left_boundary"] = left_boundary + lane_group_dict[lane_group_id]["right_boundary"] = right_boundary + outline_array = np.vstack( + [ + left_boundary.array[:, :3], + right_boundary.array[:, :3][::-1], + left_boundary.array[0, :3][None, ...], + ] + ) + + lane_group_dict[lane_group_id]["outline"] = Polyline3D.from_array(outline_array) + + return lane_group_dict + + +def _extract_lane_group(lanes) -> List[List[str]]: visited = set() lane_groups = [] @@ -391,3 +428,89 @@ def _traverse_group(start_lane_id): lane_groups.append(group) return lane_groups + + +def _extract_intersection_dict( + lanes: Dict[int, Any], lane_group_dict: Dict[int, Any], max_distance: float = 0.01 +) -> Dict[str, Any]: + + def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float: + """Interpolate Z coordinate along a 3D line segment.""" + p1, p2 = segment_coords[0], segment_coords[1] + + # Project point onto segment + segment_vec = p2[:2] - p1[:2] + point_vec = np.array([point.x, point.y]) - p1[:2] + + # Handle degenerate case + segment_length_sq = np.dot(segment_vec, segment_vec) + if segment_length_sq == 0: + return p1[2] + + # Calculate projection parameter + t = np.dot(point_vec, segment_vec) / segment_length_sq + t = np.clip(t, 0, 1) # Clamp to segment bounds + + # Interpolate Z + return p1[2] + t * (p2[2] - p1[2]) + + # 1. Collect all lane groups where at least one lane is marked as an intersection. + lane_group_intersection_dict = {} + for lane_group_id, lane_group in lane_group_dict.items(): + is_intersection_lanes = [lanes[str(lane_id)]["is_intersection"] for lane_id in lane_group["lane_ids"]] + if any(is_intersection_lanes): + lane_group_intersection_dict[lane_group_id] = lane_group + + # 2. Merge polygons of lane groups that are marked as intersections. + lane_group_intersection_geometry = { + lane_group_id: shapely.Polygon(lane_group["outline"].array[:, Point3DIndex.XY]) + for lane_group_id, lane_group in lane_group_intersection_dict.items() + } + intersection_polygons = gpd.GeoSeries(lane_group_intersection_geometry).union_all() + + # 3. Collect all intersection polygons and their lane group IDs. + intersection_dict = {} + for intersection_idx, intersection_polygon in enumerate(intersection_polygons.geoms): + if intersection_polygon.is_empty: + continue + lane_group_ids = [ + lane_group_id + for lane_group_id, lane_group_polygon in lane_group_intersection_geometry.items() + if intersection_polygon.intersects(lane_group_polygon) + ] + for lane_group_id in lane_group_ids: + lane_group_dict[lane_group_id]["intersection_id"] = intersection_idx + + intersection_dict[intersection_idx] = { + "id": intersection_idx, + "outline_2d": Polyline2D.from_array(np.array(list(intersection_polygon.exterior.coords), dtype=np.float64)), + "lane_group_ids": lane_group_ids, + } + + # 4. Lift intersection outlines to 3D. + boundary_segments = [] + for lane_group in lane_group_intersection_dict.values(): + coords = np.array(lane_group["outline"].linestring.coords, dtype=np.float64).reshape(-1, 1, 3) + segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1) + boundary_segments.append(segment_coords_boundary) + + boundary_segments = np.concatenate(boundary_segments, axis=0) + boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments) + occupancy_map = OccupancyMap2D(boundary_segment_linestrings) + + for intersection_id, intersection_data in intersection_dict.items(): + points_2d = intersection_data["outline_2d"].array + points_3d = np.zeros((len(points_2d), 3), dtype=np.float64) + points_3d[:, :2] = points_2d + + query_points = shapely.creation.points(points_2d) + results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True) + for query_idx, geometry_idx in zip(*results): + query_point = query_points[query_idx] + segment_coords = boundary_segments[geometry_idx] + best_z = _interpolate_z_on_segment(query_point, segment_coords) + points_3d[query_idx, 2] = best_z + + intersection_dict[intersection_id]["outline_3d"] = Polyline3D.from_array(points_3d) + + return intersection_dict diff --git a/notebooks/av2/delete_me.ipynb b/notebooks/av2/delete_me.ipynb index d66c98c8..e28ef5ef 100644 --- a/notebooks/av2/delete_me.ipynb +++ b/notebooks/av2/delete_me.ipynb @@ -33,7 +33,7 @@ "split = \"train\"\n", "# split = \"val\"\n", "\n", - "split_folder = Path(f\"/mnt/elements_0/argoverse/sensor_mini/{split}\")\n", + "split_folder = Path(f\"/media/nvme1/argoverse/sensor_mini/{split}\")\n", "log_folder = sorted(list(split_folder.iterdir()))[0]\n", "\n", "log_folder.name" @@ -458,7 +458,7 @@ ], "metadata": { "kernelspec": { - "display_name": "d123_dev", + "display_name": "d123", "language": "python", "name": "python3" }, diff --git a/notebooks/av2/delete_me_map.ipynb b/notebooks/av2/delete_me_map.ipynb new file mode 100644 index 00000000..79c774f7 --- /dev/null +++ b/notebooks/av2/delete_me_map.ipynb @@ -0,0 +1,310 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "0", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import io\n", + "\n", + "from PIL import Image\n", + "\n", + "import pandas as pd" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "# split = \"test\"\n", + "\n", + "\n", + "split = \"train\"\n", + "# split = \"val\"\n", + "\n", + "split_folder = Path(f\"/media/nvme1/argoverse/sensor_mini/{split}\")\n", + "log_folder = sorted(list(split_folder.iterdir()))[2]\n", + "\n", + "log_folder.name" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "from typing import Dict, List\n", + "\n", + "from d123.common.geometry.line.polylines import Polyline3D\n", + "from d123.dataset.dataset_specific.av2.av2_map_conversion import _extract_lane_group_dict\n", + "\n", + "\n", + "def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Polyline3D:\n", + " polyline = np.array([[p[\"x\"], p[\"y\"], p[\"z\"]] for p in data], dtype=np.float64)\n", + " if close:\n", + " polyline = np.vstack([polyline, polyline[0]])\n", + "\n", + " return Polyline3D.from_array(polyline)\n", + "\n", + "\n", + "map_folder = log_folder / \"map\"\n", + "log_map_archive_path = next(map_folder.glob(\"log_map_archive_*.json\"))\n", + "\n", + "with open(log_map_archive_path, \"r\") as f:\n", + " log_map_archive = json.load(f)\n", + "\n", + "drivable_areas: Dict[int, Polyline3D] = {}\n", + "\n", + "for drivable_area_id, drivable_area_dict in log_map_archive[\"drivable_areas\"].items():\n", + " # keys: [\"area_boundary\", \"id\"]\n", + " drivable_areas[int(drivable_area_id)] = _extract_polyline(drivable_area_dict[\"area_boundary\"], close=True)\n", + "\n", + "for lane_segment_id, lane_segment_dict in log_map_archive[\"lane_segments\"].items():\n", + " # keys = [\n", + " # \"id\",\n", + " # \"is_intersection\",\n", + " # \"lane_type\",\n", + " # \"left_lane_boundary\",\n", + " # \"left_lane_mark_type\",\n", + " # \"right_lane_boundary\",\n", + " # \"right_lane_mark_type\",\n", + " # \"successors\",\n", + " # \"predecessors\",\n", + " # \"right_neighbor_id\",\n", + " # \"left_neighbor_id\",\n", + " # ]\n", + " lane_segment_dict[\"left_lane_boundary\"] = _extract_polyline(lane_segment_dict[\"left_lane_boundary\"])\n", + " lane_segment_dict[\"right_lane_boundary\"] = _extract_polyline(lane_segment_dict[\"right_lane_boundary\"])\n", + "\n", + "for crosswalk_id, crosswalk_dict in log_map_archive[\"pedestrian_crossings\"].items():\n", + " # keys = [\"id\", \"outline\"]\n", + " # https://github.com/argoverse/av2-api/blob/6b22766247eda941cb1953d6a58e8d5631c561da/src/av2/map/pedestrian_crossing.py\n", + "\n", + " p1, p2 = np.array([[p[\"x\"], p[\"y\"], p[\"z\"]] for p in crosswalk_dict[\"edge1\"]], dtype=np.float64)\n", + " p3, p4 = np.array([[p[\"x\"], p[\"y\"], p[\"z\"]] for p in crosswalk_dict[\"edge2\"]], dtype=np.float64)\n", + " crosswalk_dict[\"outline\"] = Polyline3D.from_array(np.array([p1, p2, p4, p3, p1], dtype=np.float64))\n", + "\n", + "\n", + "lane_group_dict = _extract_lane_group_dict(log_map_archive[\"lane_segments\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from typing import Any\n", + "\n", + "import shapely\n", + "\n", + "from d123.common.geometry.base import Point3DIndex\n", + "import geopandas as gpd\n", + "\n", + "from d123.common.geometry.line.polylines import Polyline2D\n", + "from d123.common.geometry.occupancy_map import OccupancyMap2D\n", + "\n", + "import numpy.typing as npt\n", + "\n", + "\n", + "def _interpolate_z_on_segment(point: shapely.Point, segment_coords: npt.NDArray[np.float64]) -> float:\n", + " \"\"\"Interpolate Z coordinate along a 3D line segment.\"\"\"\n", + " p1, p2 = segment_coords[0], segment_coords[1]\n", + "\n", + " # Project point onto segment\n", + " segment_vec = p2[:2] - p1[:2]\n", + " point_vec = np.array([point.x, point.y]) - p1[:2]\n", + "\n", + " # Handle degenerate case\n", + " segment_length_sq = np.dot(segment_vec, segment_vec)\n", + " if segment_length_sq == 0:\n", + " return p1[2]\n", + "\n", + " # Calculate projection parameter\n", + " t = np.dot(point_vec, segment_vec) / segment_length_sq\n", + " t = np.clip(t, 0, 1) # Clamp to segment bounds\n", + "\n", + " # Interpolate Z\n", + " return p1[2] + t * (p2[2] - p1[2])\n", + "\n", + "\n", + "def _extract_intersection_dict(\n", + " lanes: Dict[int, Any], lane_group_dict: Dict[int, Any], max_distance: float = 0.01\n", + ") -> Dict[str, Any]:\n", + "\n", + " # 1. Collect all lane groups where at least one lane is marked as an intersection.\n", + " lane_group_intersection_dict = {}\n", + " for lane_group_id, lane_group in lane_group_dict.items():\n", + " is_intersection_lanes = [lanes[str(lane_id)][\"is_intersection\"] for lane_id in lane_group[\"lane_ids\"]]\n", + " if any(is_intersection_lanes):\n", + " lane_group_intersection_dict[lane_group_id] = lane_group\n", + "\n", + " # 2. Merge polygons of lane groups that are marked as intersections.\n", + " lane_group_intersection_geometry = {\n", + " lane_group_id: shapely.Polygon(lane_group[\"outline\"].array[:, Point3DIndex.XY])\n", + " for lane_group_id, lane_group in lane_group_intersection_dict.items()\n", + " }\n", + " intersection_polygons = gpd.GeoSeries(lane_group_intersection_geometry).union_all()\n", + "\n", + " # 3. Collect all intersection polygons and their lane group IDs.\n", + " intersection_dict = {}\n", + " for intersection_idx, intersection_polygon in enumerate(intersection_polygons.geoms):\n", + " if intersection_polygon.is_empty:\n", + " continue\n", + " lane_group_ids = [\n", + " lane_group_id\n", + " for lane_group_id, lane_group_polygon in lane_group_intersection_geometry.items()\n", + " if intersection_polygon.intersects(lane_group_polygon)\n", + " ]\n", + " intersection_dict[f\"intersection_{intersection_idx}\"] = {\n", + " \"id\": intersection_idx,\n", + " \"outline_2d\": Polyline2D.from_array(np.array(list(intersection_polygon.exterior.coords), dtype=np.float64)),\n", + " \"lane_group_ids\": lane_group_ids,\n", + " }\n", + "\n", + " # 4. Lift intersection outlines to 3D.\n", + " boundary_segments = []\n", + " for lane_group in lane_group_intersection_dict.values():\n", + " coords = np.array(lane_group[\"outline\"].linestring.coords, dtype=np.float64).reshape(-1, 1, 3)\n", + " segment_coords_boundary = np.concatenate([coords[:-1], coords[1:]], axis=1)\n", + " boundary_segments.append(segment_coords_boundary)\n", + "\n", + " boundary_segments = np.concatenate(boundary_segments, axis=0)\n", + " boundary_segment_linestrings = shapely.creation.linestrings(boundary_segments)\n", + " occupancy_map = OccupancyMap2D(boundary_segment_linestrings)\n", + "\n", + " for intersection_id, intersection_data in intersection_dict.items():\n", + " points_2d = intersection_data[\"outline_2d\"].array\n", + " points_3d = np.zeros((len(points_2d), 3), dtype=np.float64)\n", + " points_3d[:, :2] = points_2d\n", + "\n", + " query_points = shapely.creation.points(points_2d)\n", + " results = occupancy_map.query_nearest(query_points, max_distance=max_distance, exclusive=True)\n", + " for query_idx, geometry_idx in zip(*results):\n", + " query_point = query_points[query_idx]\n", + " segment_coords = boundary_segments[geometry_idx]\n", + " best_z = _interpolate_z_on_segment(query_point, segment_coords)\n", + " points_3d[query_idx, 2] = best_z\n", + "\n", + " intersection_dict[intersection_id][\"outline_3d\"] = Polyline3D.from_array(points_3d)\n", + "\n", + " return intersection_dict\n", + "\n", + "\n", + "intersection_dict = _extract_intersection_dict(log_map_archive[\"lane_segments\"], lane_group_dict)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "from matplotlib import cm\n", + "from matplotlib.colors import Normalize\n", + "\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "for intersection_id, values in intersection_dict.items():\n", + "\n", + " outline = values[\"outline_3d\"].array\n", + " print(outline[:, 2].min(), outline[:, 2].max())\n", + "\n", + " # Normalize z values to [0, 1] for colormap mapping\n", + " norm = Normalize(vmin=outline[:, 2].min(), vmax=outline[:, 2].max())\n", + " colors = cm.viridis(norm(outline[:, 2]))\n", + "\n", + " # Plot each segment with its corresponding color\n", + " for i in range(len(outline) - 1):\n", + " ax.plot(outline[i : i + 2, 0], outline[i : i + 2, 1], color=colors[i], linewidth=2)\n", + "\n", + "ax.set_aspect(\"equal\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "fig, ax = plt.subplots(figsize=(10, 10))\n", + "\n", + "\n", + "for lane in log_map_archive[\"lane_segments\"].values():\n", + " left_boundary = lane[\"left_lane_boundary\"]\n", + " right_boundary = lane[\"right_lane_boundary\"]\n", + "\n", + " polygon = np.concatenate([left_boundary.array, right_boundary.array[::-1]]).reshape(-1, 3)[:, :2]\n", + " print(polygon)\n", + " ax.fill(\n", + " polygon[:, 0], polygon[:, 1], alpha=0.5, edgecolor=\"black\", color=\"red\" if lane[\"is_intersection\"] else \"blue\"\n", + " )\n", + "\n", + " # if left_boundary and right_boundary:\n", + " # ax.plot(left_boundary.array[:, 0], left_boundary.array[:, 1], color=\"blue\", linewidth=1)\n", + " # ax.plot(right_boundary.array[:, 0], right_boundary.array[:, 1], color=\"red\", linewidth=1)\n", + "\n", + "ax.set_title(\"Lane Segments\")\n", + "ax.set_xlabel(\"X\")\n", + "ax.set_ylabel(\"Y\")\n", + "ax.set_aspect(\"equal\")\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "d123", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb index 7687c8e2..ac7aaded 100644 --- a/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb +++ b/notebooks/deprecated/map_conversion/test_opendrive_conversion.ipynb @@ -23,7 +23,7 @@ "outputs": [], "source": [ "from pathlib import Path\n", - "from d123.dataset.conversion.map.opendrive.elements.opendrive import OpenDrive\n", + "from d123.dataset.conversion.map.opendrive.parser.opendrive import OpenDrive\n", "\n", "CARLA_MAP_LOCATIONS = [\n", " \"Town01\", # A small, simple town with a river and several bridges.\n", diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 7155776a..5bbc2f08 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -26,9 +26,9 @@ "\n", "\n", "# splits = [\"wopd_train\"]\n", - "splits = [\"carla\"]\n", + "# splits = [\"carla\"]\n", "# splits = [\"nuplan_private_test\"]\n", - "# splits = [\"av2-sensor-mini_train\"]\n", + "splits = [\"av2-sensor-mini_train\"]\n", "# log_names = None\n", "\n", "\n", @@ -76,6 +76,7 @@ "from d123.common.visualization.matplotlib.utils import add_shapely_linestring_to_ax, add_shapely_polygon_to_ax\n", "from d123.dataset.maps.abstract_map import AbstractMap\n", "from d123.dataset.maps.abstract_map_objects import AbstractLane, AbstractLaneGroup\n", + "from d123.dataset.maps.gpkg.gpkg_map_objects import GPKGIntersection\n", "from d123.dataset.maps.map_datatypes import MapLayer\n", "from d123.dataset.scene.abstract_scene import AbstractScene\n", "\n", @@ -99,7 +100,7 @@ " line_color_alpha=0.5,\n", " line_width=1.0,\n", " line_style=\"-\",\n", - " zorder=3,\n", + " zorder=22,\n", ")\n", "\n", "\n", @@ -147,7 +148,7 @@ " 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", @@ -161,9 +162,6 @@ " for layer, map_objects in map_objects_dict.items():\n", " for map_object in map_objects:\n", " try:\n", - " if layer in [MapLayer.WALKWAY]:\n", - " print(map_object.trimesh_mesh.vertices)\n", - "\n", " if layer in [\n", " MapLayer.GENERIC_DRIVABLE,\n", " MapLayer.CARPARK,\n", @@ -175,26 +173,14 @@ "\n", " if layer in [MapLayer.LANE_GROUP]:\n", " map_object: AbstractLaneGroup\n", - "\n", - " successor_ids = [successor.id for successor in map_object.successors]\n", - "\n", " add_shapely_polygon_to_ax(ax, map_object.shapely_polygon, MAP_SURFACE_CONFIG[layer])\n", - " # centroid = map_object.shapely_polygon.centroid\n", - " # ax.text(\n", - " # centroid.x,\n", - " # centroid.y,\n", - " # f\"{map_object.id} -> {successor_ids}\",\n", - " # horizontalalignment=\"center\",\n", - " # verticalalignment=\"center\",\n", - " # fontsize=8,\n", - " # bbox=dict(facecolor=\"white\", alpha=0.7, boxstyle=\"round,pad=0.2\"),\n", - " # )\n", + "\n", "\n", " if layer in [MapLayer.LANE]:\n", " add_shapely_linestring_to_ax(ax, map_object.centerline.linestring, CENTERLINE_CONFIG)\n", "\n", - " if layer in [MapLayer.ROAD_EDGE]:\n", - " add_shapely_linestring_to_ax(ax, map_object.polyline_3d.linestring, ROAD_EDGE_CONFIG)\n", + " # 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", @@ -250,8 +236,8 @@ " return fig, ax\n", "\n", "\n", - "scene_index = 10\n", - "iteration = 2\n", + "scene_index = 9\n", + "iteration = 99\n", "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=iteration, radius=100)\n", "plt.show()\n", "\n", From 328448a1b7dc677cc860526a5833af743b0f9419 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 22 Aug 2025 18:11:05 +0200 Subject: [PATCH 41/42] Add road lines to argoverse map. Change road line types to argoverse convention --- .../map/opendrive/opendrive_map_conversion.py | 6 +- .../dataset_specific/av2/av2_constants.py | 20 +++++ .../av2/av2_map_conversion.py | 87 +++++++++++-------- .../nuplan/nuplan_map_conversion.py | 8 +- .../wopd/waymo_map_utils/wopd_map_utils.py | 19 ++-- d123/dataset/maps/map_datatypes.py | 25 +++--- notebooks/av2/delete_me_map.ipynb | 5 +- notebooks/viz/bev_matplotlib.ipynb | 6 +- 8 files changed, 106 insertions(+), 70 deletions(-) diff --git a/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py b/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py index c66c8613..56dfa65c 100644 --- a/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py +++ b/d123/dataset/conversion/map/opendrive/opendrive_map_conversion.py @@ -358,19 +358,19 @@ def _extract_road_line_df( if str(lane_row.right_lane_id) in ["nan", "None"]: # This is a boundary lane, e.g. a border or sidewalk ids.append(running_id) - road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) + road_line_types.append(int(RoadLineType.SOLID_WHITE)) geometries.append(lane_row.right_boundary) running_id += 1 else: # This is a regular lane ids.append(running_id) - road_line_types.append(int(RoadLineType.BROKEN_SINGLE_WHITE)) + road_line_types.append(int(RoadLineType.DASHED_WHITE)) geometries.append(lane_row.right_boundary) running_id += 1 if str(lane_row.left_lane_id) in ["nan", "None"]: # This is a boundary lane, e.g. a border or sidewalk ids.append(running_id) - road_line_types.append(int(RoadLineType.SOLID_SINGLE_WHITE)) + road_line_types.append(int(RoadLineType.DASHED_WHITE)) geometries.append(lane_row.left_boundary) running_id += 1 diff --git a/d123/dataset/dataset_specific/av2/av2_constants.py b/d123/dataset/dataset_specific/av2/av2_constants.py index 65799aae..836bb12c 100644 --- a/d123/dataset/dataset_specific/av2/av2_constants.py +++ b/d123/dataset/dataset_specific/av2/av2_constants.py @@ -1,6 +1,7 @@ from d123.common.datatypes.detection.detection_types import DetectionType from d123.common.datatypes.sensor.camera import CameraType from d123.common.utils.enums import SerialIntEnum +from d123.dataset.maps.map_datatypes import RoadLineType class AV2SensorBoxDetectionType(SerialIntEnum): @@ -85,3 +86,22 @@ class AV2SensorBoxDetectionType(SerialIntEnum): "stereo_front_left": CameraType.CAM_STEREO_L, "stereo_front_right": CameraType.CAM_STEREO_R, } + + +AV2_ROAD_LINE_TYPE_MAPPING = { + "NONE": RoadLineType.NONE, + "UNKNOWN": RoadLineType.UNKNOWN, + "DASH_SOLID_YELLOW": RoadLineType.DASH_SOLID_YELLOW, + "DASH_SOLID_WHITE": RoadLineType.DASH_SOLID_WHITE, + "DASHED_WHITE": RoadLineType.DASHED_WHITE, + "DASHED_YELLOW": RoadLineType.DASHED_YELLOW, + "DOUBLE_SOLID_YELLOW": RoadLineType.DOUBLE_SOLID_YELLOW, + "DOUBLE_SOLID_WHITE": RoadLineType.DOUBLE_SOLID_WHITE, + "DOUBLE_DASH_YELLOW": RoadLineType.DOUBLE_DASH_YELLOW, + "DOUBLE_DASH_WHITE": RoadLineType.DOUBLE_DASH_WHITE, + "SOLID_YELLOW": RoadLineType.SOLID_YELLOW, + "SOLID_WHITE": RoadLineType.SOLID_WHITE, + "SOLID_DASH_WHITE": RoadLineType.SOLID_DASH_WHITE, + "SOLID_DASH_YELLOW": RoadLineType.SOLID_DASH_YELLOW, + "SOLID_BLUE": RoadLineType.SOLID_BLUE, +} diff --git a/d123/dataset/dataset_specific/av2/av2_map_conversion.py b/d123/dataset/dataset_specific/av2/av2_map_conversion.py index 98e5fc78..390cbdb2 100644 --- a/d123/dataset/dataset_specific/av2/av2_map_conversion.py +++ b/d123/dataset/dataset_specific/av2/av2_map_conversion.py @@ -1,3 +1,4 @@ +import warnings from pathlib import Path from typing import Any, Dict, Final, List @@ -16,11 +17,16 @@ from d123.dataset.conversion.map.road_edge.road_edge_3d_utils import ( get_road_edges_3d_from_generic_drivable_area_df, ) -from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType, RoadLineType - -# TODO: -# - TODO -LANE_GROUP_MARK_TYPES = ["DASHED_WHITE", "DOUBLE_DASH_WHITE", "DASH_SOLID_WHITE", "SOLID_DASH_WHITE", "SOLID_WHITE"] +from d123.dataset.dataset_specific.av2.av2_constants import AV2_ROAD_LINE_TYPE_MAPPING +from d123.dataset.maps.map_datatypes import MapLayer, RoadEdgeType + +LANE_GROUP_MARK_TYPES: List[str] = [ + "DASHED_WHITE", + "DOUBLE_DASH_WHITE", + "DASH_SOLID_WHITE", + "SOLID_DASH_WHITE", + "SOLID_WHITE", +] MAX_ROAD_EDGE_LENGTH: Final[float] = 100.0 # TODO: Add to config @@ -73,29 +79,27 @@ def _extract_polyline(data: List[Dict[str, float]], close: bool = False) -> Poly lane_group_dict = _extract_lane_group_dict(log_map_archive["lane_segments"]) intersection_dict = _extract_intersection_dict(log_map_archive["lane_segments"], lane_group_dict) - lane_df = get_lane_df(log_map_archive["lane_segments"]) - lane_group_df = get_lane_group_df(lane_group_dict) - intersection_df = get_intersections_df(intersection_dict) - crosswalk_df = get_crosswalk_df(log_map_archive["pedestrian_crossings"]) - walkway_df = get_empty_gdf() # NOTE: AV2 does not provide walkways, so we create an empty DataFrame. - carpark_df = get_empty_gdf() # NOTE: AV2 does not provide carparks, so we create an empty DataFrame. - generic_drivable_df = get_generic_drivable_df(drivable_areas) - road_edge_df = get_road_edge_df(generic_drivable_df) - # road_line_df = get_empty_gdf() + dataframes: Dict[MapLayer, gpd.GeoDataFrame] = {} + + dataframes[MapLayer.LANE] = get_lane_df(log_map_archive["lane_segments"]) + dataframes[MapLayer.LANE_GROUP] = get_lane_group_df(lane_group_dict) + dataframes[MapLayer.INTERSECTION] = get_intersections_df(intersection_dict) + dataframes[MapLayer.CROSSWALK] = get_crosswalk_df(log_map_archive["pedestrian_crossings"]) + dataframes[MapLayer.GENERIC_DRIVABLE] = get_generic_drivable_df(drivable_areas) + dataframes[MapLayer.ROAD_EDGE] = get_road_edge_df(dataframes[MapLayer.GENERIC_DRIVABLE]) + dataframes[MapLayer.ROAD_LINE] = get_road_line_df(log_map_archive["lane_segments"]) + # NOTE: AV2 does not provide walkways or carparks, so we create an empty DataFrame. + dataframes[MapLayer.WALKWAY] = get_empty_gdf() + dataframes[MapLayer.CARPARK] = get_empty_gdf() map_file_path.unlink(missing_ok=True) if not map_file_path.parent.exists(): map_file_path.parent.mkdir(parents=True, exist_ok=True) - lane_df.to_file(map_file_path, layer=MapLayer.LANE.serialize(), driver="GPKG") - lane_group_df.to_file(map_file_path, layer=MapLayer.LANE_GROUP.serialize(), driver="GPKG", mode="a") - intersection_df.to_file(map_file_path, layer=MapLayer.INTERSECTION.serialize(), driver="GPKG", mode="a") - crosswalk_df.to_file(map_file_path, layer=MapLayer.CROSSWALK.serialize(), driver="GPKG", mode="a") - walkway_df.to_file(map_file_path, layer=MapLayer.WALKWAY.serialize(), driver="GPKG", mode="a") - carpark_df.to_file(map_file_path, layer=MapLayer.CARPARK.serialize(), driver="GPKG", mode="a") - generic_drivable_df.to_file(map_file_path, layer=MapLayer.GENERIC_DRIVABLE.serialize(), driver="GPKG", mode="a") - road_edge_df.to_file(map_file_path, layer=MapLayer.ROAD_EDGE.serialize(), driver="GPKG", mode="a") - # road_line_df.to_file(map_file_path, layer=MapLayer.ROAD_LINE.serialize(), driver="GPKG", mode="a") + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", message="'crs' was not provided") + for layer, gdf in dataframes.items(): + gdf.to_file(map_file_path, layer=layer.serialize(), driver="GPKG", mode="a") def get_empty_gdf() -> gpd.GeoDataFrame: @@ -151,10 +155,7 @@ def _get_centerline_from_boundaries( left_boundary=lane_dict["left_lane_boundary"], right_boundary=lane_dict["right_lane_boundary"], ) - lane_speed_limit_mps = None - - # ids.append(lane_id) - # lane_types.append(lanes_type[lane_id]) + lane_speed_limit_mps = None # TODO: Consider using geo reference to retrieve speed limits. lane_group_ids.append(lane_id) speed_limits_mps.append(lane_speed_limit_mps) predecessor_ids.append(lane_dict["predecessors"]) @@ -313,18 +314,28 @@ def get_road_edge_df(generic_drivable_area_df: gpd.GeoDataFrame) -> gpd.GeoDataF return gpd.GeoDataFrame(pd.DataFrame({"id": ids, "road_edge_type": road_edge_types}), geometry=geometries) -def get_road_line_df( - road_lines: Dict[int, npt.NDArray[np.float64]], road_lines_type: Dict[int, RoadLineType] -) -> gpd.GeoDataFrame: - ids = list(road_lines.keys()) - geometries = [Polyline3D.from_array(road_edge).linestring for road_edge in road_lines.values()] +def get_road_line_df(lanes: Dict[int, Any]) -> gpd.GeoDataFrame: - data = pd.DataFrame( - { - "id": ids, - "road_line_type": [int(road_line_type) for road_line_type in road_lines_type.values()], - } - ) + # TODO @DanielDauner: Allow lanes to reference road line dataframe. + + ids = [] + road_lines_type = [] + geometries = [] + + running_id = 0 + for lane in lanes.values(): + for side in ["left", "right"]: + # NOTE: We currently ignore lane markings that are NONE in the AV2 dataset. + # TODO: Review if the road line system should be changed in the future. + if lane[f"{side}_lane_mark_type"] == "NONE": + continue + + ids.append(running_id) + road_lines_type.append(AV2_ROAD_LINE_TYPE_MAPPING[lane[f"{side}_lane_mark_type"]]) + geometries.append(lane[f"{side}_lane_boundary"].linestring) + running_id += 1 + + data = pd.DataFrame({"id": ids, "road_line_type": road_lines_type}) gdf = gpd.GeoDataFrame(data, geometry=geometries) return gdf diff --git a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py index 2f0050fd..78f5d3f6 100644 --- a/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py +++ b/d123/dataset/dataset_specific/nuplan/nuplan_map_conversion.py @@ -53,8 +53,9 @@ MAX_ROAD_EDGE_LENGTH = 100.0 # meters, used to filter out very long road edges NUPLAN_ROAD_LINE_CONVERSION = { - 0: RoadLineType.BROKEN_SINGLE_WHITE, - 2: RoadLineType.SOLID_SINGLE_WHITE, + 0: RoadLineType.DASHED_WHITE, + 2: RoadLineType.SOLID_WHITE, + 3: RoadLineType.UNKNOWN, } @@ -443,9 +444,6 @@ def _extract_road_line_dataframe(self) -> gpd.GeoDataFrame: geometries = [] for idx in range(len(boundary_types)): - # NOTE @DanielDauner: We ignore boundaries of nuPlan-type, which are on intersections. - if boundary_types[idx] not in NUPLAN_ROAD_LINE_CONVERSION.keys(): - continue ids.append(fids[idx]) road_line_types.append(int(NUPLAN_ROAD_LINE_CONVERSION[boundary_types[idx]])) geometries.append(boundaries[idx]) diff --git a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py index a2de1233..6d3ef7ab 100644 --- a/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py +++ b/d123/dataset/dataset_specific/wopd/waymo_map_utils/wopd_map_utils.py @@ -20,18 +20,17 @@ # - Implement speed bumps # - Implement driveways with a different semantic type if needed # - Implement intersections and lane group logic -# - Handle lane type, e.g. TYPE_UNDEFINED = 0; TYPE_FREEWAY = 1; TYPE_SURFACE_STREET = 2; TYPE_BIKE_LANE = 3; WAYMO_ROAD_LINE_CONVERSION = { - 0: RoadLineType.UNKNOWN, - 1: RoadLineType.BROKEN_SINGLE_WHITE, - 2: RoadLineType.SOLID_SINGLE_WHITE, - 3: RoadLineType.SOLID_DOUBLE_WHITE, - 4: RoadLineType.BROKEN_SINGLE_YELLOW, - 5: RoadLineType.BROKEN_DOUBLE_YELLOW, - 6: RoadLineType.SOLID_SINGLE_YELLOW, - 7: RoadLineType.SOLID_DOUBLE_YELLOW, - 8: RoadLineType.PASSING_DOUBLE_YELLOW, + 0: RoadLineType.UNKNOWN, # aka. UNKNOWN + 1: RoadLineType.DASHED_WHITE, # aka. BROKEN_SINGLE_WHITE + 2: RoadLineType.SOLID_WHITE, # aka. SOLID_SINGLE_WHITE + 3: RoadLineType.DOUBLE_SOLID_WHITE, # aka. SOLID_DOUBLE_WHITE + 4: RoadLineType.DASHED_YELLOW, # aka. BROKEN_SINGLE_YELLOW + 5: RoadLineType.DOUBLE_DASH_YELLOW, # aka. BROKEN_DOUBLE_YELLOW + 6: RoadLineType.SOLID_YELLOW, # aka. SOLID_SINGLE_YELLOW + 7: RoadLineType.DOUBLE_SOLID_YELLOW, # aka. SOLID_DOUBLE_YELLOW + 8: RoadLineType.DOUBLE_DASH_YELLOW, # aka. PASSING_DOUBLE_YELLOW } WAYMO_ROAD_EDGE_CONVERSION = { diff --git a/d123/dataset/maps/map_datatypes.py b/d123/dataset/maps/map_datatypes.py index ab29cbf5..f948422f 100644 --- a/d123/dataset/maps/map_datatypes.py +++ b/d123/dataset/maps/map_datatypes.py @@ -25,7 +25,6 @@ class MapLayer(SerialIntEnum): class LaneType(SerialIntEnum): """ Enum for LaneType. - NOTE: We use the lane types from Waymo. https://github.com/waymo-research/waymo-open-dataset/blob/99a4cb3ff07e2fe06c2ce73da001f850f628e45a/src/waymo_open_dataset/protos/map.proto#L147 """ @@ -55,12 +54,18 @@ class RoadLineType(SerialIntEnum): https://github.com/waymo-research/waymo-open-dataset/blob/master/src/waymo_open_dataset/protos/map.proto#L208 """ - UNKNOWN = 0 - BROKEN_SINGLE_WHITE = 1 - SOLID_SINGLE_WHITE = 2 - SOLID_DOUBLE_WHITE = 3 - BROKEN_SINGLE_YELLOW = 4 - BROKEN_DOUBLE_YELLOW = 5 - SOLID_SINGLE_YELLOW = 6 - SOLID_DOUBLE_YELLOW = 7 - PASSING_DOUBLE_YELLOW = 8 + NONE = 0 + UNKNOWN = 1 + DASH_SOLID_YELLOW = 2 + DASH_SOLID_WHITE = 3 + DASHED_WHITE = 4 + DASHED_YELLOW = 5 + DOUBLE_SOLID_YELLOW = 6 + DOUBLE_SOLID_WHITE = 7 + DOUBLE_DASH_YELLOW = 8 + DOUBLE_DASH_WHITE = 9 + SOLID_YELLOW = 10 + SOLID_WHITE = 11 + SOLID_DASH_WHITE = 12 + SOLID_DASH_YELLOW = 13 + SOLID_BLUE = 14 diff --git a/notebooks/av2/delete_me_map.ipynb b/notebooks/av2/delete_me_map.ipynb index 79c774f7..03e62838 100644 --- a/notebooks/av2/delete_me_map.ipynb +++ b/notebooks/av2/delete_me_map.ipynb @@ -43,7 +43,10 @@ "id": "2", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "for lane_segment_id, lane_segment_dict in log_map_archive[\"lane_segments\"].items():\n", + " print(f\"Lane Segment ID: {lane_segment_id}\")" + ] }, { "cell_type": "code", diff --git a/notebooks/viz/bev_matplotlib.ipynb b/notebooks/viz/bev_matplotlib.ipynb index 5bbc2f08..1feef159 100644 --- a/notebooks/viz/bev_matplotlib.ipynb +++ b/notebooks/viz/bev_matplotlib.ipynb @@ -210,8 +210,8 @@ " box_detections = scene.get_box_detections_at_iteration(iteration)\n", "\n", " point_2d = ego_vehicle_state.bounding_box.center.state_se2.point_2d\n", - " add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", - " # add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " # add_debug_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", + " add_default_map_on_ax(ax, scene.map_api, point_2d, radius=radius, route_lane_group_ids=None)\n", " # add_traffic_lights_to_ax(ax, traffic_light_detections, scene.map_api)\n", "\n", " add_box_detections_to_ax(ax, box_detections)\n", @@ -238,7 +238,7 @@ "\n", "scene_index = 9\n", "iteration = 99\n", - "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=iteration, radius=100)\n", + "fig, ax = plot_scene_at_iteration(scenes[scene_index], iteration=iteration, radius=30)\n", "plt.show()\n", "\n", "# camera = scenes[scene_index].get_camera_at_iteration(\n", From a68475e4fe7d74ddae38f39fd8e9e09555d2d618 Mon Sep 17 00:00:00 2001 From: Daniel Dauner Date: Fri, 22 Aug 2025 18:13:56 +0200 Subject: [PATCH 42/42] Remove debugging print statement --- d123/dataset/dataset_specific/wopd/wopd_data_converter.py | 1 - 1 file changed, 1 deletion(-) diff --git a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py index 1295b2b0..f1145c8e 100644 --- a/d123/dataset/dataset_specific/wopd/wopd_data_converter.py +++ b/d123/dataset/dataset_specific/wopd/wopd_data_converter.py @@ -537,7 +537,6 @@ def _extract_lidar( range_image_top_pose=range_image_top_pose, keep_polar_features=True, ) - print(f"Extracted {len(points)} points from {len(frame.lasers)} lasers.") lidar_data: Dict[LiDARType, npt.NDArray[np.float32]] = {} for lidar_idx, frame_lidar in enumerate(frame.lasers):