From 15fb6330604c2edf43c0f336b19db132c74e1681 Mon Sep 17 00:00:00 2001 From: Benjamin Wilson Date: Mon, 1 May 2023 14:11:29 -0700 Subject: [PATCH] Change line length. --- pyproject.toml | 4 - .../motion_forecasting/eval/metrics.py | 47 +++- .../motion_forecasting/eval/submission.py | 31 ++- .../scenario_serialization.py | 31 ++- .../viz/scenario_visualization.py | 97 ++++++-- .../datasets/sensor/av2_sensor_dataloader.py | 81 +++++-- src/av2/datasets/sensor/sensor_dataloader.py | 137 ++++++++--- src/av2/datasets/sensor/utils.py | 6 +- src/av2/evaluation/detection/eval.py | 85 +++++-- src/av2/evaluation/detection/utils.py | 102 ++++++-- src/av2/evaluation/forecasting/constants.py | 1 + src/av2/evaluation/forecasting/eval.py | 126 +++++++--- src/av2/evaluation/forecasting/utils.py | 43 +++- src/av2/evaluation/scene_flow/constants.py | 10 +- src/av2/evaluation/scene_flow/eval.py | 107 +++++++-- .../scene_flow/example_submission.py | 17 +- .../scene_flow/make_annotation_files.py | 15 +- .../evaluation/scene_flow/make_mask_files.py | 9 +- .../scene_flow/make_submission_archive.py | 39 ++- src/av2/evaluation/scene_flow/utils.py | 31 ++- src/av2/evaluation/tracking/eval.py | 138 ++++++++--- src/av2/evaluation/tracking/utils.py | 16 +- src/av2/evaluation/typing.py | 1 - src/av2/geometry/camera/pinhole_camera.py | 36 ++- src/av2/geometry/geometry.py | 24 +- src/av2/geometry/infinity_norm_utils.py | 8 +- src/av2/geometry/interpolate.py | 34 ++- src/av2/geometry/iou.py | 4 +- src/av2/geometry/polyline_utils.py | 52 +++- src/av2/geometry/se3.py | 8 +- src/av2/geometry/sim2.py | 12 +- src/av2/geometry/utm.py | 36 ++- src/av2/map/drivable_area.py | 4 +- src/av2/map/lane_segment.py | 19 +- src/av2/map/map_api.py | 148 +++++++++--- src/av2/map/pedestrian_crossing.py | 4 +- src/av2/rendering/color.py | 4 +- src/av2/rendering/map.py | 117 +++++++-- src/av2/rendering/ops/draw.py | 20 +- src/av2/rendering/rasterize.py | 4 +- src/av2/rendering/vector.py | 8 +- src/av2/rendering/video.py | 32 ++- src/av2/structures/cuboid.py | 63 +++-- src/av2/structures/ndgrid.py | 19 +- src/av2/torch/__init__.py | 4 +- src/av2/torch/data_loaders/detection.py | 2 +- src/av2/torch/data_loaders/scene_flow.py | 4 +- src/av2/torch/structures/cuboids.py | 11 +- src/av2/torch/structures/flow.py | 26 +- src/av2/torch/structures/sweep.py | 8 +- src/av2/torch/structures/utils.py | 4 +- src/av2/utils/dataclass.py | 10 +- src/av2/utils/dense_grid_interpolation.py | 8 +- src/av2/utils/depth_map_utils.py | 9 +- src/av2/utils/dilation_utils.py | 4 +- src/av2/utils/helpers.py | 11 +- src/av2/utils/io.py | 22 +- src/av2/utils/metric_time.py | 7 +- src/av2/utils/raster.py | 8 +- src/av2/utils/synchronization_database.py | 89 +++++-- tests/integration/verify_tbv_download.py | 62 ++++- .../motion_forecasting/eval/test_metrics.py | 120 +++++++-- .../eval/test_submission.py | 47 +++- .../test_scenario_serialization.py | 41 +++- .../sensor/test_av2_sensor_dataloader.py | 4 +- .../datasets/sensor/test_sensor_dataloader.py | 39 ++- tests/unit/evaluation/detection/test_eval.py | 227 +++++++++++++++--- .../evaluation/scene_flow/test_sf_eval.py | 120 +++++++-- .../scene_flow/test_sf_submission_pipeline.py | 40 ++- tests/unit/geometry/test_geometry.py | 162 ++++++++++--- tests/unit/geometry/test_interpolate.py | 103 ++++++-- tests/unit/geometry/test_pinhole_camera.py | 150 ++++++++++-- tests/unit/geometry/test_utm.py | 8 +- tests/unit/map/test_drivable_area.py | 4 +- tests/unit/map/test_lane_segment.py | 5 +- tests/unit/map/test_map_api.py | 147 +++++++++--- tests/unit/rendering/ops/test_draw.py | 66 +++-- tests/unit/rendering/test_color.py | 4 +- tests/unit/rendering/test_map.py | 8 +- tests/unit/rendering/test_rasterize.py | 4 +- tests/unit/rendering/test_video.py | 4 +- tests/unit/structures/test_cuboid.py | 30 ++- tests/unit/structures/test_ndgrid.py | 5 +- tests/unit/structures/test_sweep.py | 13 +- .../data_loaders/test_detection_dataloader.py | 8 +- .../test_scene_flow_dataloader.py | 16 +- tests/unit/torch/structures/test_cuboids.py | 12 +- tests/unit/torch/structures/test_lidar.py | 8 +- tests/unit/torch/structures/test_sweep.py | 8 +- tests/unit/torch/structures/test_utils.py | 21 +- .../utils/test_dense_grid_interpolation.py | 29 ++- tests/unit/utils/test_depth_map_utils.py | 12 +- tests/unit/utils/test_infinity_norm_utils.py | 16 +- tests/unit/utils/test_io.py | 15 +- tests/unit/utils/test_mesh_grid.py | 16 +- tests/unit/utils/test_polyline_utils.py | 4 +- tests/unit/utils/test_raster.py | 24 +- tests/unit/utils/test_se3.py | 30 ++- tests/unit/utils/test_sim2.py | 4 +- tutorials/3d_object_detection.py | 7 +- tutorials/generate_egoview_overlaid_lidar.py | 61 ++++- .../generate_egoview_overlaid_vector_map.py | 38 ++- ...ate_forecasting_scenario_visualizations.py | 24 +- tutorials/generate_per_camera_videos.py | 31 ++- .../generate_sensor_dataset_visualizations.py | 29 ++- tutorials/map_teaser_notebook.py | 86 +++++-- tutorials/untar_tbv.py | 24 +- 107 files changed, 3109 insertions(+), 884 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index a48b0c02..d6ae3785 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,11 +61,7 @@ repository = "https://github.com/argoverse/av2-api" features = ["pyo3/extension-module"] module-name = "av2._r" -[tool.black] -line-length = 120 - [tool.ruff] -line-length = 120 select = [ "D", ] diff --git a/src/av2/datasets/motion_forecasting/eval/metrics.py b/src/av2/datasets/motion_forecasting/eval/metrics.py index 4c923643..d8424912 100644 --- a/src/av2/datasets/motion_forecasting/eval/metrics.py +++ b/src/av2/datasets/motion_forecasting/eval/metrics.py @@ -8,7 +8,9 @@ from av2.utils.typing import NDArrayBool, NDArrayFloat -def compute_ade(forecasted_trajectories: NDArrayFloat, gt_trajectory: NDArrayFloat) -> NDArrayFloat: +def compute_ade( + forecasted_trajectories: NDArrayFloat, gt_trajectory: NDArrayFloat +) -> NDArrayFloat: """Compute the average displacement error for a set of K predicted trajectories (for the same actor). Args: @@ -18,12 +20,16 @@ def compute_ade(forecasted_trajectories: NDArrayFloat, gt_trajectory: NDArrayFlo Returns: (K,) Average displacement error for each of the predicted trajectories. """ - displacement_errors = np.linalg.norm(forecasted_trajectories - gt_trajectory, axis=2) + displacement_errors = np.linalg.norm( + forecasted_trajectories - gt_trajectory, axis=2 + ) ade: NDArrayFloat = np.mean(displacement_errors, axis=1) return ade -def compute_fde(forecasted_trajectories: NDArrayFloat, gt_trajectory: NDArrayFloat) -> NDArrayFloat: +def compute_fde( + forecasted_trajectories: NDArrayFloat, gt_trajectory: NDArrayFloat +) -> NDArrayFloat: """Compute the final displacement error for a set of K predicted trajectories (for the same actor). Args: @@ -78,7 +84,9 @@ def compute_brier_ade( (K,) Probability-weighted average displacement error for each predicted trajectory. """ # Compute ADE with Brier score component - brier_score = _compute_brier_score(forecasted_trajectories, forecast_probabilities, normalize) + brier_score = _compute_brier_score( + forecasted_trajectories, forecast_probabilities, normalize + ) ade_vector = compute_ade(forecasted_trajectories, gt_trajectory) brier_ade: NDArrayFloat = ade_vector + brier_score return brier_ade @@ -102,7 +110,9 @@ def compute_brier_fde( (K,) Probability-weighted final displacement error for each predicted trajectory. """ # Compute FDE with Brier score component - brier_score = _compute_brier_score(forecasted_trajectories, forecast_probabilities, normalize) + brier_score = _compute_brier_score( + forecasted_trajectories, forecast_probabilities, normalize + ) fde_vector = compute_fde(forecasted_trajectories, gt_trajectory) brier_fde: NDArrayFloat = fde_vector + brier_score return brier_fde @@ -136,7 +146,9 @@ def _compute_brier_score( # Validate that all forecast probabilities are in the range [0, 1] if np.logical_or(forecast_probabilities < 0.0, forecast_probabilities > 1.0).any(): - raise ValueError("At least one forecast probability falls outside the range [0, 1].") + raise ValueError( + "At least one forecast probability falls outside the range [0, 1]." + ) # If enabled, normalize forecast probabilities to sum to 1 if normalize: @@ -146,7 +158,9 @@ def _compute_brier_score( return brier_score -def compute_world_fde(forecasted_world_trajectories: NDArrayFloat, gt_world_trajectories: NDArrayFloat) -> NDArrayFloat: +def compute_world_fde( + forecasted_world_trajectories: NDArrayFloat, gt_world_trajectories: NDArrayFloat +) -> NDArrayFloat: """Compute the mean final displacement error for each of K predicted worlds. Args: @@ -167,7 +181,9 @@ def compute_world_fde(forecasted_world_trajectories: NDArrayFloat, gt_world_traj return world_fdes -def compute_world_ade(forecasted_world_trajectories: NDArrayFloat, gt_world_trajectories: NDArrayFloat) -> NDArrayFloat: +def compute_world_ade( + forecasted_world_trajectories: NDArrayFloat, gt_world_trajectories: NDArrayFloat +) -> NDArrayFloat: """Compute the mean average displacement error for each of K predicted worlds. Args: @@ -189,7 +205,9 @@ def compute_world_ade(forecasted_world_trajectories: NDArrayFloat, gt_world_traj def compute_world_misses( - forecasted_world_trajectories: NDArrayFloat, gt_world_trajectories: NDArrayFloat, miss_threshold_m: float = 2.0 + forecasted_world_trajectories: NDArrayFloat, + gt_world_trajectories: NDArrayFloat, + miss_threshold_m: float = 2.0, ) -> NDArrayBool: """For each world, compute whether predictions for each actor misssed by more than a distance threshold. @@ -230,7 +248,12 @@ def compute_world_brier_fde( (K,) Mean probability-weighted final displacement error for each of the predicted worlds. """ actor_brier_fdes = [ - compute_brier_fde(forecasted_actor_trajectories, gt_actor_trajectory, forecasted_world_probabilities, normalize) + compute_brier_fde( + forecasted_actor_trajectories, + gt_actor_trajectory, + forecasted_world_probabilities, + normalize, + ) for forecasted_actor_trajectories, gt_actor_trajectory in zip( forecasted_world_trajectories, gt_world_trajectories ) @@ -256,7 +279,9 @@ def compute_world_collisions( for actor_idx in range(len(forecasted_world_trajectories)): # Compute distance from current actor to all other predicted actors at each timestep forecasted_actor_trajectories = forecasted_world_trajectories[actor_idx] - scenario_actor_dists = np.linalg.norm(forecasted_world_trajectories - forecasted_actor_trajectories, axis=-1) + scenario_actor_dists = np.linalg.norm( + forecasted_world_trajectories - forecasted_actor_trajectories, axis=-1 + ) # For each world, find the closest distance to any other predicted actor, at any time scenario_actor_dists[actor_idx, :, :] = np.inf diff --git a/src/av2/datasets/motion_forecasting/eval/submission.py b/src/av2/datasets/motion_forecasting/eval/submission.py index 230eb4e5..f748283c 100644 --- a/src/av2/datasets/motion_forecasting/eval/submission.py +++ b/src/av2/datasets/motion_forecasting/eval/submission.py @@ -49,7 +49,10 @@ def __post_init__(self) -> None: ValueError: If for any track, number of probabilities doesn't match the number of predicted trajectories. ValueError: If prediction probabilities for at least one scenario do not sum to 1. """ - for scenario_id, (scenario_probabilities, scenario_trajectories) in self.predictions.items(): + for scenario_id, ( + scenario_probabilities, + scenario_trajectories, + ) in self.predictions.items(): for track_id, track_trajectories in scenario_trajectories.items(): # Validate that predicted trajectories are of the correct shape if track_trajectories[0].shape[-2:] != EXPECTED_PREDICTION_SHAPE: @@ -94,13 +97,24 @@ def from_parquet(cls, submission_file_path: Path) -> ChallengeSubmission: for scenario_id, scenario_df in submission_df.groupby("scenario_id"): scenario_trajectories: ScenarioTrajectories = {} for track_id, track_df in scenario_df.groupby("track_id"): - predicted_trajectories_x = np.stack(track_df.loc[:, "predicted_trajectory_x"].values.tolist()) - predicted_trajectories_y = np.stack(track_df.loc[:, "predicted_trajectory_y"].values.tolist()) - predicted_trajectories = np.stack((predicted_trajectories_x, predicted_trajectories_y), axis=-1) + predicted_trajectories_x = np.stack( + track_df.loc[:, "predicted_trajectory_x"].values.tolist() + ) + predicted_trajectories_y = np.stack( + track_df.loc[:, "predicted_trajectory_y"].values.tolist() + ) + predicted_trajectories = np.stack( + (predicted_trajectories_x, predicted_trajectories_y), axis=-1 + ) scenario_trajectories[track_id] = predicted_trajectories - scenario_probabilities = np.array(track_df.loc[:, "probability"].values.tolist()) - submission_dict[scenario_id] = (scenario_probabilities, scenario_trajectories) + scenario_probabilities = np.array( + track_df.loc[:, "probability"].values.tolist() + ) + submission_dict[scenario_id] = ( + scenario_probabilities, + scenario_trajectories, + ) return cls(predictions=submission_dict) @@ -113,7 +127,10 @@ def to_parquet(self, submission_file_path: Path) -> None: prediction_rows: List[PredictionRow] = [] # Build list of rows for the submission dataframe - for scenario_id, (scenario_probabilities, scenario_trajectories) in self.predictions.items(): + for scenario_id, ( + scenario_probabilities, + scenario_trajectories, + ) in self.predictions.items(): for track_id, track_trajectories in scenario_trajectories.items(): for world_idx in range(len(track_trajectories)): prediction_rows.append( diff --git a/src/av2/datasets/motion_forecasting/scenario_serialization.py b/src/av2/datasets/motion_forecasting/scenario_serialization.py index 3c185657..a766521e 100644 --- a/src/av2/datasets/motion_forecasting/scenario_serialization.py +++ b/src/av2/datasets/motion_forecasting/scenario_serialization.py @@ -8,10 +8,18 @@ import numpy as np import pandas as pd -from av2.datasets.motion_forecasting.data_schema import ArgoverseScenario, ObjectState, ObjectType, Track, TrackCategory - - -def serialize_argoverse_scenario_parquet(save_path: Path, scenario: ArgoverseScenario) -> None: +from av2.datasets.motion_forecasting.data_schema import ( + ArgoverseScenario, + ObjectState, + ObjectType, + Track, + TrackCategory, +) + + +def serialize_argoverse_scenario_parquet( + save_path: Path, scenario: ArgoverseScenario +) -> None: """Serialize a single Argoverse scenario in parquet format and save to disk. Args: @@ -64,7 +72,9 @@ def load_argoverse_scenario_parquet(scenario_path: Path) -> ArgoverseScenario: # Interpolate scenario timestamps based on the saved start and end timestamps timestamps_ns = np.linspace( - tracks_df["start_timestamp"][0], tracks_df["end_timestamp"][0], num=tracks_df["num_timestamps"][0] + tracks_df["start_timestamp"][0], + tracks_df["end_timestamp"][0], + num=tracks_df["num_timestamps"][0], ) return ArgoverseScenario( @@ -139,7 +149,9 @@ def _load_tracks_from_tabular_format(tracks_df: pd.DataFrame) -> List[Track]: for track_id, track_df in tracks_df.groupby("track_id"): observed_states: List[bool] = track_df.loc[:, "observed"].values.tolist() object_type: ObjectType = ObjectType(track_df["object_type"].iloc[0]) - object_category: TrackCategory = TrackCategory(track_df["object_category"].iloc[0]) + object_category: TrackCategory = TrackCategory( + track_df["object_category"].iloc[0] + ) timesteps: List[int] = track_df.loc[:, "timestep"].values.tolist() positions: List[Tuple[float, float]] = list( zip( @@ -168,7 +180,12 @@ def _load_tracks_from_tabular_format(tracks_df: pd.DataFrame) -> List[Track]: ) tracks.append( - Track(track_id=track_id, object_states=object_states, object_type=object_type, category=object_category) + Track( + track_id=track_id, + object_states=object_states, + object_type=object_type, + category=object_category, + ) ) return tracks diff --git a/src/av2/datasets/motion_forecasting/viz/scenario_visualization.py b/src/av2/datasets/motion_forecasting/viz/scenario_visualization.py index 9113e82f..22402d75 100644 --- a/src/av2/datasets/motion_forecasting/viz/scenario_visualization.py +++ b/src/av2/datasets/motion_forecasting/viz/scenario_visualization.py @@ -13,7 +13,11 @@ from PIL import Image as img from PIL.Image import Image -from av2.datasets.motion_forecasting.data_schema import ArgoverseScenario, ObjectType, TrackCategory +from av2.datasets.motion_forecasting.data_schema import ( + ArgoverseScenario, + ObjectType, + TrackCategory, +) from av2.map.map_api import ArgoverseStaticMap from av2.utils.typing import NDArrayFloat, NDArrayInt @@ -35,7 +39,9 @@ _DEFAULT_ACTOR_COLOR: Final[str] = "#D3E8EF" _FOCAL_AGENT_COLOR: Final[str] = "#ECA25B" _AV_COLOR: Final[str] = "#007672" -_BOUNDING_BOX_ZORDER: Final[int] = 100 # Ensure actor bounding boxes are plotted on top of all map elements +_BOUNDING_BOX_ZORDER: Final[ + int +] = 100 # Ensure actor bounding boxes are plotted on top of all map elements _STATIC_OBJECT_TYPES: Set[ObjectType] = { ObjectType.STATIC, @@ -45,7 +51,11 @@ } -def visualize_scenario(scenario: ArgoverseScenario, scenario_static_map: ArgoverseStaticMap, save_path: Path) -> None: +def visualize_scenario( + scenario: ArgoverseScenario, + scenario_static_map: ArgoverseStaticMap, + save_path: Path, +) -> None: """Build dynamic visualization for all tracks and the local map associated with an Argoverse scenario. Note: This function uses OpenCV to create a MP4 file using the MP4V codec. @@ -69,8 +79,14 @@ def visualize_scenario(scenario: ArgoverseScenario, scenario_static_map: Argover plot_bounds = cur_plot_bounds # Set map bounds to capture focal trajectory history (with fixed buffer in all directions) - plt.xlim(plot_bounds[0] - _PLOT_BOUNDS_BUFFER_M, plot_bounds[1] + _PLOT_BOUNDS_BUFFER_M) - plt.ylim(plot_bounds[2] - _PLOT_BOUNDS_BUFFER_M, plot_bounds[3] + _PLOT_BOUNDS_BUFFER_M) + plt.xlim( + plot_bounds[0] - _PLOT_BOUNDS_BUFFER_M, + plot_bounds[1] + _PLOT_BOUNDS_BUFFER_M, + ) + plt.ylim( + plot_bounds[2] - _PLOT_BOUNDS_BUFFER_M, + plot_bounds[3] + _PLOT_BOUNDS_BUFFER_M, + ) plt.gca().set_aspect("equal", adjustable="box") # Minimize plot margins and make axes invisible @@ -98,7 +114,9 @@ def visualize_scenario(scenario: ArgoverseScenario, scenario_static_map: Argover video.release() -def _plot_static_map_elements(static_map: ArgoverseStaticMap, show_ped_xings: bool = False) -> None: +def _plot_static_map_elements( + static_map: ArgoverseStaticMap, show_ped_xings: bool = False +) -> None: """Plot all static map elements associated with an Argoverse scenario. Args: @@ -123,10 +141,16 @@ def _plot_static_map_elements(static_map: ArgoverseStaticMap, show_ped_xings: bo # Plot pedestrian crossings if show_ped_xings: for ped_xing in static_map.vector_pedestrian_crossings.values(): - _plot_polylines([ped_xing.edge1.xyz, ped_xing.edge2.xyz], alpha=1.0, color=_LANE_SEGMENT_COLOR) + _plot_polylines( + [ped_xing.edge1.xyz, ped_xing.edge2.xyz], + alpha=1.0, + color=_LANE_SEGMENT_COLOR, + ) -def _plot_actor_tracks(ax: plt.Axes, scenario: ArgoverseScenario, timestep: int) -> Optional[_PlotBounds]: +def _plot_actor_tracks( + ax: plt.Axes, scenario: ArgoverseScenario, timestep: int +) -> Optional[_PlotBounds]: """Plot all actor tracks (up to a particular time step) associated with an Argoverse scenario. Args: @@ -141,17 +165,29 @@ def _plot_actor_tracks(ax: plt.Axes, scenario: ArgoverseScenario, timestep: int) for track in scenario.tracks: # Get timesteps for which actor data is valid actor_timesteps: NDArrayInt = np.array( - [object_state.timestep for object_state in track.object_states if object_state.timestep <= timestep] + [ + object_state.timestep + for object_state in track.object_states + if object_state.timestep <= timestep + ] ) if actor_timesteps.shape[0] < 1 or actor_timesteps[-1] != timestep: continue # Get actor trajectory and heading history actor_trajectory: NDArrayFloat = np.array( - [list(object_state.position) for object_state in track.object_states if object_state.timestep <= timestep] + [ + list(object_state.position) + for object_state in track.object_states + if object_state.timestep <= timestep + ] ) actor_headings: NDArrayFloat = np.array( - [object_state.heading for object_state in track.object_states if object_state.timestep <= timestep] + [ + object_state.heading + for object_state in track.object_states + if object_state.timestep <= timestep + ] ) # Plot polyline for focal agent location history @@ -176,7 +212,10 @@ def _plot_actor_tracks(ax: plt.Axes, scenario: ArgoverseScenario, timestep: int) track_color, (_ESTIMATED_VEHICLE_LENGTH_M, _ESTIMATED_VEHICLE_WIDTH_M), ) - elif track.object_type == ObjectType.CYCLIST or track.object_type == ObjectType.MOTORCYCLIST: + elif ( + track.object_type == ObjectType.CYCLIST + or track.object_type == ObjectType.MOTORCYCLIST + ): _plot_actor_bounding_box( ax, actor_trajectory[-1], @@ -185,7 +224,13 @@ def _plot_actor_tracks(ax: plt.Axes, scenario: ArgoverseScenario, timestep: int) (_ESTIMATED_CYCLIST_LENGTH_M, _ESTIMATED_CYCLIST_WIDTH_M), ) else: - plt.plot(actor_trajectory[-1, 0], actor_trajectory[-1, 1], "o", color=track_color, markersize=4) + plt.plot( + actor_trajectory[-1, 0], + actor_trajectory[-1, 1], + "o", + color=track_color, + markersize=4, + ) return track_bounds @@ -208,10 +253,19 @@ def _plot_polylines( color: Desired color for the plotted lines. """ for polyline in polylines: - plt.plot(polyline[:, 0], polyline[:, 1], style, linewidth=line_width, color=color, alpha=alpha) + plt.plot( + polyline[:, 0], + polyline[:, 1], + style, + linewidth=line_width, + color=color, + alpha=alpha, + ) -def _plot_polygons(polygons: Sequence[NDArrayFloat], *, alpha: float = 1.0, color: str = "r") -> None: +def _plot_polygons( + polygons: Sequence[NDArrayFloat], *, alpha: float = 1.0, color: str = "r" +) -> None: """Plot a group of filled polygons with the specified config. Args: @@ -224,7 +278,11 @@ def _plot_polygons(polygons: Sequence[NDArrayFloat], *, alpha: float = 1.0, colo def _plot_actor_bounding_box( - ax: plt.Axes, cur_location: NDArrayFloat, heading: float, color: str, bbox_size: Tuple[float, float] + ax: plt.Axes, + cur_location: NDArrayFloat, + heading: float, + color: str, + bbox_size: Tuple[float, float], ) -> None: """Plot an actor bounding box centered on the actor's current location. @@ -244,6 +302,11 @@ def _plot_actor_bounding_box( pivot_y = cur_location[1] - (d / 2) * math.sin(heading + theta_2) vehicle_bounding_box = Rectangle( - (pivot_x, pivot_y), bbox_length, bbox_width, np.degrees(heading), color=color, zorder=_BOUNDING_BOX_ZORDER + (pivot_x, pivot_y), + bbox_length, + bbox_width, + np.degrees(heading), + color=color, + zorder=_BOUNDING_BOX_ZORDER, ) ax.add_patch(vehicle_bounding_box) diff --git a/src/av2/datasets/sensor/av2_sensor_dataloader.py b/src/av2/datasets/sensor/av2_sensor_dataloader.py index 1482537d..03625b25 100644 --- a/src/av2/datasets/sensor/av2_sensor_dataloader.py +++ b/src/av2/datasets/sensor/av2_sensor_dataloader.py @@ -41,7 +41,9 @@ def __init__(self, data_dir: Path, labels_dir: Path) -> None: ValueError: if input arguments are not Path objects. """ if not isinstance(data_dir, Path) or not isinstance(labels_dir, Path): - raise ValueError("Input arguments must be Path objects, representing paths to local directories") + raise ValueError( + "Input arguments must be Path objects, representing paths to local directories" + ) self._data_dir = data_dir self._labels_dir = labels_dir @@ -67,7 +69,9 @@ def get_city_SE3_ego(self, log_id: str, timestamp_ns: int) -> SE3: Raises: RuntimeError: If no recorded pose is available for the requested timestamp. """ - log_poses_df = io_utils.read_feather(self._data_dir / log_id / "city_SE3_egovehicle.feather") + log_poses_df = io_utils.read_feather( + self._data_dir / log_id / "city_SE3_egovehicle.feather" + ) pose_df = log_poses_df.loc[log_poses_df["timestamp_ns"] == timestamp_ns] if len(pose_df) == 0: @@ -76,7 +80,9 @@ def get_city_SE3_ego(self, log_id: str, timestamp_ns: int) -> SE3: city_SE3_ego = convert_pose_dataframe_to_SE3(pose_df) return city_SE3_ego - def get_subsampled_ego_trajectory(self, log_id: str, sample_rate_hz: float = 1.0) -> NDArrayFloat: + def get_subsampled_ego_trajectory( + self, log_id: str, sample_rate_hz: float = 1.0 + ) -> NDArrayFloat: """Get the trajectory of the AV (egovehicle) at an approximate sampling rate (Hz). Note: the trajectory is NOT interpolated to give an *exact* sampling rate. @@ -101,7 +107,9 @@ def get_subsampled_ego_trajectory(self, log_id: str, sample_rate_hz: float = 1.0 MAX_MEASUREMENT_FREQUENCY_HZ, ) - log_poses_df = io_utils.read_feather(self._data_dir / log_id / "city_SE3_egovehicle.feather") + log_poses_df = io_utils.read_feather( + self._data_dir / log_id / "city_SE3_egovehicle.feather" + ) # timestamp of the pose measurement. timestamp_ns = list(log_poses_df.timestamp_ns) @@ -159,7 +167,9 @@ def get_log_ids(self) -> List[str]: """Return a list of all vehicle log IDs available at the provided dataroot.""" return sorted([d.name for d in self._data_dir.glob("*") if d.is_dir()]) - def get_closest_img_fpath(self, log_id: str, cam_name: str, lidar_timestamp_ns: int) -> Optional[Path]: + def get_closest_img_fpath( + self, log_id: str, cam_name: str, lidar_timestamp_ns: int + ) -> Optional[Path]: """Return the filepath corresponding to the closest image from the lidar timestamp. Args: @@ -170,13 +180,24 @@ def get_closest_img_fpath(self, log_id: str, cam_name: str, lidar_timestamp_ns: Returns: im_fpath: path to image if one is found within the expected time interval, or else None. """ - cam_timestamp_ns = self._sdb.get_closest_cam_channel_timestamp(lidar_timestamp_ns, cam_name, log_id) + cam_timestamp_ns = self._sdb.get_closest_cam_channel_timestamp( + lidar_timestamp_ns, cam_name, log_id + ) if cam_timestamp_ns is None: return None - img_fpath = self._data_dir / log_id / "sensors" / "cameras" / cam_name / f"{cam_timestamp_ns}.jpg" + img_fpath = ( + self._data_dir + / log_id + / "sensors" + / "cameras" + / cam_name + / f"{cam_timestamp_ns}.jpg" + ) return img_fpath - def get_closest_lidar_fpath(self, log_id: str, cam_timestamp_ns: int) -> Optional[Path]: + def get_closest_lidar_fpath( + self, log_id: str, cam_timestamp_ns: int + ) -> Optional[Path]: """Get file path for LiDAR sweep accumulated to a timestamp closest to a camera timestamp. Args: @@ -186,14 +207,18 @@ def get_closest_lidar_fpath(self, log_id: str, cam_timestamp_ns: int) -> Optiona Returns: lidar_fpath: path to sweep .feather file if one is found within the expected time interval, or else None. """ - lidar_timestamp_ns = self._sdb.get_closest_lidar_timestamp(cam_timestamp_ns, log_id) + lidar_timestamp_ns = self._sdb.get_closest_lidar_timestamp( + cam_timestamp_ns, log_id + ) if lidar_timestamp_ns is None: return None lidar_fname = f"{lidar_timestamp_ns}.feather" lidar_fpath = self._data_dir / log_id / "sensors" / "lidar" / lidar_fname return lidar_fpath - def get_lidar_fpath_at_lidar_timestamp(self, log_id: str, lidar_timestamp_ns: int) -> Optional[Path]: + def get_lidar_fpath_at_lidar_timestamp( + self, log_id: str, lidar_timestamp_ns: int + ) -> Optional[Path]: """Return the file path for the LiDAR sweep accumulated to the query timestamp, if it exists. Args: @@ -233,7 +258,9 @@ def get_ordered_log_lidar_timestamps(self, log_id: str) -> List[int]: Returns: lidar_timestamps_ns: ordered timestamps, provided in nanoseconds. """ - ordered_lidar_fpaths: List[Path] = self.get_ordered_log_lidar_fpaths(log_id=log_id) + ordered_lidar_fpaths: List[Path] = self.get_ordered_log_lidar_fpaths( + log_id=log_id + ) lidar_timestamps_ns = [int(fp.stem) for fp in ordered_lidar_fpaths] return lidar_timestamps_ns @@ -261,10 +288,14 @@ def get_ordered_log_cam_fpaths(self, log_id: str, cam_name: str) -> List[Path]: Returns: List of paths, representing paths to ordered JPEG files in this log, for a specific camera. """ - cam_img_fpaths = sorted(self._data_dir.glob(f"{log_id}/sensors/cameras/{cam_name}/*.jpg")) + cam_img_fpaths = sorted( + self._data_dir.glob(f"{log_id}/sensors/cameras/{cam_name}/*.jpg") + ) return cam_img_fpaths - def get_labels_at_lidar_timestamp(self, log_id: str, lidar_timestamp_ns: int) -> CuboidList: + def get_labels_at_lidar_timestamp( + self, log_id: str, lidar_timestamp_ns: int + ) -> CuboidList: """Load the sweep annotations at the provided timestamp. Args: @@ -280,7 +311,9 @@ def get_labels_at_lidar_timestamp(self, log_id: str, lidar_timestamp_ns: int) -> # NOTE: This file contains annotations for the ENTIRE sequence. # The sweep annotations are selected below. cuboid_list = CuboidList.from_feather(annotations_feather_path) - cuboids = list(filter(lambda x: x.timestamp_ns == lidar_timestamp_ns, cuboid_list.cuboids)) + cuboids = list( + filter(lambda x: x.timestamp_ns == lidar_timestamp_ns, cuboid_list.cuboids) + ) return CuboidList(cuboids=cuboids) def project_ego_to_img_motion_compensated( @@ -310,11 +343,15 @@ def project_ego_to_img_motion_compensated( # get transformation to bring point in egovehicle frame to city frame, # at the time when camera image was recorded. - city_SE3_ego_cam_t = self.get_city_SE3_ego(log_id=log_id, timestamp_ns=cam_timestamp_ns) + city_SE3_ego_cam_t = self.get_city_SE3_ego( + log_id=log_id, timestamp_ns=cam_timestamp_ns + ) # get transformation to bring point in egovehicle frame to city frame, # at the time when the LiDAR sweep was recorded. - city_SE3_ego_lidar_t = self.get_city_SE3_ego(log_id=log_id, timestamp_ns=lidar_timestamp_ns) + city_SE3_ego_lidar_t = self.get_city_SE3_ego( + log_id=log_id, timestamp_ns=lidar_timestamp_ns + ) return pinhole_camera.project_ego_to_img_motion_compensated( points_lidar_time=points_lidar_time, @@ -337,9 +374,13 @@ def get_colored_sweep(self, log_id: str, lidar_timestamp_ns: int) -> NDArrayByte Raises: ValueError: If requested timestamp has no corresponding LiDAR sweep. """ - lidar_fpath = self.get_lidar_fpath_at_lidar_timestamp(log_id=log_id, lidar_timestamp_ns=lidar_timestamp_ns) + lidar_fpath = self.get_lidar_fpath_at_lidar_timestamp( + log_id=log_id, lidar_timestamp_ns=lidar_timestamp_ns + ) if lidar_fpath is None: - raise ValueError("Requested colored sweep at a timestamp that has no corresponding LiDAR sweep.") + raise ValueError( + "Requested colored sweep at a timestamp that has no corresponding LiDAR sweep." + ) sweep = Sweep.from_feather(lidar_fpath) n_sweep_pts = len(sweep) @@ -422,7 +463,9 @@ def get_depth_map_from_lidar( # form depth map from LiDAR if interp_depth_map: if u.max() > pinhole_camera.width_px or v.max() > pinhole_camera.height_px: - raise RuntimeError("Regular grid interpolation will fail due to out-of-bound inputs.") + raise RuntimeError( + "Regular grid interpolation will fail due to out-of-bound inputs." + ) depth_map = dense_grid_interpolation.interp_dense_grid_from_sparse( grid_img=depth_map, diff --git a/src/av2/datasets/sensor/sensor_dataloader.py b/src/av2/datasets/sensor/sensor_dataloader.py index 0eb91065..2da81108 100644 --- a/src/av2/datasets/sensor/sensor_dataloader.py +++ b/src/av2/datasets/sensor/sensor_dataloader.py @@ -22,7 +22,12 @@ from av2.structures.sweep import Sweep from av2.structures.timestamped_image import TimestampedImage from av2.utils.constants import HOME -from av2.utils.io import TimestampedCitySE3EgoPoses, read_city_SE3_ego, read_feather, read_img +from av2.utils.io import ( + TimestampedCitySE3EgoPoses, + read_city_SE3_ego, + read_feather, + read_img, +) from av2.utils.metric_time import TimeUnit, to_metric_time logger = logging.Logger(__name__) @@ -40,13 +45,19 @@ # constants defined in milliseconds # below evaluates to 50 ms -CAM_SHUTTER_INTERVAL_MS: Final[float] = to_metric_time(ts=1 / CAM_FPS, src=Second, dst=Millisecond) +CAM_SHUTTER_INTERVAL_MS: Final[float] = to_metric_time( + ts=1 / CAM_FPS, src=Second, dst=Millisecond +) # below evaluates to 100 ms -LIDAR_SWEEP_INTERVAL_MS: Final[float] = to_metric_time(ts=1 / LIDAR_FRAME_RATE_HZ, src=Second, dst=Millisecond) +LIDAR_SWEEP_INTERVAL_MS: Final[float] = to_metric_time( + ts=1 / LIDAR_FRAME_RATE_HZ, src=Second, dst=Millisecond +) ALLOWED_TIMESTAMP_BUFFER_MS: Final[int] = 2 # allow 2 ms of buffer -LIDAR_SWEEP_INTERVAL_W_BUFFER_MS: Final[float] = LIDAR_SWEEP_INTERVAL_MS + ALLOWED_TIMESTAMP_BUFFER_MS +LIDAR_SWEEP_INTERVAL_W_BUFFER_MS: Final[float] = ( + LIDAR_SWEEP_INTERVAL_MS + ALLOWED_TIMESTAMP_BUFFER_MS +) LIDAR_SWEEP_INTERVAL_W_BUFFER_NS: Final[float] = to_metric_time( ts=LIDAR_SWEEP_INTERVAL_W_BUFFER_MS, src=Millisecond, dst=Nanosecond ) @@ -151,7 +162,9 @@ def __post_init__(self) -> None: # Populate synchronization database. if self.cam_names: - synchronization_cache_path = HOME / ".cache" / "av2" / "synchronization_cache.feather" + synchronization_cache_path = ( + HOME / ".cache" / "av2" / "synchronization_cache.feather" + ) synchronization_cache_path.parent.mkdir(parents=True, exist_ok=True) # If caching is enabled AND the path exists, then load from the cache file. @@ -165,7 +178,9 @@ def __post_init__(self) -> None: self.synchronization_cache.to_feather(str(synchronization_cache_path)) # Finally, create a MultiIndex set the sync records index and sort it. - self.synchronization_cache.set_index(keys=["split", "log_id", "sensor_name"], inplace=True) + self.synchronization_cache.set_index( + keys=["split", "log_id", "sensor_name"], inplace=True + ) self.synchronization_cache.sort_index(inplace=True) @cached_property @@ -181,7 +196,9 @@ def num_sweeps(self) -> int: @cached_property def sensor_counts(self) -> pd.Series: """Return the number of records for each sensor.""" - sensor_counts: pd.Series = self.sensor_cache.index.get_level_values("sensor_name").value_counts() + sensor_counts: pd.Series = self.sensor_cache.index.get_level_values( + "sensor_name" + ).value_counts() return sensor_counts @property @@ -231,7 +248,9 @@ def _build_sensor_cache(self) -> pd.DataFrame: # 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_cache.set_index(["split", "log_id", "sensor_name", "timestamp_ns"], inplace=True) + sensor_cache.set_index( + ["split", "log_id", "sensor_name", "timestamp_ns"], inplace=True + ) sensor_cache.sort_index(inplace=True) # Return all of the sensor records. @@ -245,9 +264,12 @@ def populate_lidar_records(self) -> pd.DataFrame: N is the number of sweeps for all logs in the dataset, and the `sensor_name` column should be populated with `lidar` in every entry. """ - lidar_paths = sorted(self.dataset_dir.glob(LIDAR_PATTERN), key=lambda x: int(x.stem)) + lidar_paths = sorted( + self.dataset_dir.glob(LIDAR_PATTERN), key=lambda x: int(x.stem) + ) lidar_record_list = [ - convert_path_to_named_record(x) for x in track(lidar_paths, description="Loading lidar records ...") + convert_path_to_named_record(x) + for x in track(lidar_paths, description="Loading lidar records ...") ] # Concatenate into single dataframe (list-of-dicts to DataFrame). @@ -264,11 +286,14 @@ def populate_image_records(self) -> pd.DataFrame: every entry. """ # Get sorted list of camera paths. - cam_paths = sorted(self.dataset_dir.glob(CAMERA_PATTERN), key=lambda x: int(x.stem)) + cam_paths = sorted( + self.dataset_dir.glob(CAMERA_PATTERN), key=lambda x: int(x.stem) + ) # Load entire set of camera records. cam_record_list = [ - convert_path_to_named_record(x) for x in track(cam_paths, description="Loading camera records ...") + convert_path_to_named_record(x) + for x in track(cam_paths, description="Loading camera records ...") ] # Concatenate into single dataframe (list-of-dicts to DataFrame). @@ -310,7 +335,9 @@ def __getitem__(self, idx: int) -> SynchronizedSensorData: """ # Grab the lidar record at the specified index. # Selects data at a particular level of a MultiIndex. - record: Tuple[str, str, int] = self.sensor_cache.xs(key="lidar", level=2).iloc[idx].name + record: Tuple[str, str, int] = ( + self.sensor_cache.xs(key="lidar", level=2).iloc[idx].name + ) # Grab the identifying record fields. split, log_id, timestamp_ns = record @@ -345,7 +372,9 @@ def __getitem__(self, idx: int) -> SynchronizedSensorData: # Load camera imagery if enabled. if self.cam_names: - datum.synchronized_imagery = self._load_synchronized_cams(split, sensor_dir, log_id, timestamp_ns) + datum.synchronized_imagery = self._load_synchronized_cams( + split, sensor_dir, log_id, timestamp_ns + ) # Return datum at the specified index. return datum @@ -368,28 +397,38 @@ def _build_synchronization_cache(self) -> pd.DataFrame: # Create list to store synchronized data frames. sync_list: List[pd.DataFrame] = [] - unique_sensor_names: List[str] = self.sensor_cache.index.unique(level=2).tolist() + unique_sensor_names: List[str] = self.sensor_cache.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 = self.sensor_cache.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) + src_records = self.sensor_cache.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 = self.sensor_cache.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 - ) + target_records = self.sensor_cache.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]) + target_records[target_sensor_name] = pd.to_timedelta( + target_records[target_sensor_name] + ) tolerance = pd.to_timedelta(CAM_SHUTTER_INTERVAL_MS / 2 * 1e6) if "ring" in src_sensor_name: tolerance = pd.to_timedelta(LIDAR_SWEEP_INTERVAL_W_BUFFER_NS / 2) @@ -430,33 +469,45 @@ def find_closest_target_fpath( RuntimeError: if the synchronization database (sync_records) has not been created. """ if self.synchronization_cache is None: - raise RuntimeError("Requested synchronized data, but the synchronization database has not been created.") + raise RuntimeError( + "Requested synchronized data, but the synchronization database has not been created." + ) src_timedelta_ns = pd.Timedelta(src_timestamp_ns) - src_to_target_records = self.synchronization_cache.loc[(split, log_id, src_sensor_name)].set_index( - src_sensor_name - ) + src_to_target_records = self.synchronization_cache.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] + 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 = self.dataset_dir / split / log_id / "sensors" - valid_cameras = [x.value for x in list(RingCameras)] + [x.value for x in list(StereoCameras)] + valid_cameras = [x.value for x in list(RingCameras)] + [ + x.value for x in list(StereoCameras) + ] 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" + 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" + target_path = ( + sensor_dir / target_sensor_name / f"{timestamp_ns_str}.feather" + ) return target_path - def get_closest_img_fpath(self, split: str, log_id: str, cam_name: str, lidar_timestamp_ns: int) -> Optional[Path]: + def get_closest_img_fpath( + self, split: str, log_id: str, cam_name: str, lidar_timestamp_ns: int + ) -> Optional[Path]: """Find the file path to the closest image from the reference camera name to the lidar timestamp. Args: @@ -476,7 +527,9 @@ def get_closest_img_fpath(self, split: str, log_id: str, cam_name: str, lidar_ti target_sensor_name=cam_name, ) - def get_closest_lidar_fpath(self, split: str, log_id: str, cam_name: str, cam_timestamp_ns: int) -> Optional[Path]: + def get_closest_lidar_fpath( + self, split: str, log_id: str, cam_name: str, cam_timestamp_ns: int + ) -> Optional[Path]: """Find the file path to the closest image from the lidar to the reference camera. Args: @@ -496,7 +549,9 @@ def get_closest_lidar_fpath(self, split: str, log_id: str, cam_name: str, cam_ti target_sensor_name="lidar", ) - def _load_annotations(self, split: str, log_id: str, sweep_timestamp_ns: int) -> CuboidList: + def _load_annotations( + self, split: str, log_id: str, sweep_timestamp_ns: int + ) -> CuboidList: """Load the sweep annotations at the provided timestamp. Args: @@ -507,13 +562,17 @@ def _load_annotations(self, split: str, log_id: str, sweep_timestamp_ns: int) -> Returns: Cuboid list of annotations. """ - annotations_feather_path = self.dataset_dir / split / log_id / "annotations.feather" + annotations_feather_path = ( + self.dataset_dir / split / log_id / "annotations.feather" + ) # Load annotations from disk. # NOTE: This contains annotations for the ENTIRE sequence. # The sweep annotations are selected below. cuboid_list = CuboidList.from_feather(annotations_feather_path) - cuboids = list(filter(lambda x: x.timestamp_ns == sweep_timestamp_ns, cuboid_list.cuboids)) + cuboids = list( + filter(lambda x: x.timestamp_ns == sweep_timestamp_ns, cuboid_list.cuboids) + ) return CuboidList(cuboids=cuboids) def _load_synchronized_cams( @@ -534,7 +593,9 @@ def _load_synchronized_cams( RuntimeError: if the synchronization database (sync_records) has not been created. """ if self.synchronization_cache is None: - raise RuntimeError("Requested synchronized data, but the synchronization database has not been created.") + raise RuntimeError( + "Requested synchronized data, but the synchronization database has not been created." + ) cam_paths = [ self.find_closest_target_fpath( @@ -553,7 +614,9 @@ def _load_synchronized_cams( if p is not None: cams[p.parent.stem] = TimestampedImage( img=read_img(p, channel_order="BGR"), - camera_model=PinholeCamera.from_feather(log_dir=log_dir, cam_name=p.parent.stem), + camera_model=PinholeCamera.from_feather( + log_dir=log_dir, cam_name=p.parent.stem + ), timestamp_ns=int(p.stem), ) return cams diff --git a/src/av2/datasets/sensor/utils.py b/src/av2/datasets/sensor/utils.py index 91159fd2..9e5c5350 100644 --- a/src/av2/datasets/sensor/utils.py +++ b/src/av2/datasets/sensor/utils.py @@ -20,7 +20,11 @@ def convert_path_to_named_record(path: Path) -> Dict[str, Union[str, int]]: """ sensor_path = path.parent sensor_name = sensor_path.stem - log_path = sensor_path.parent.parent if sensor_name == "lidar" else sensor_path.parent.parent.parent + log_path = ( + sensor_path.parent.parent + if sensor_name == "lidar" + else sensor_path.parent.parent.parent + ) # log_id is 2 directories up for the lidar filepaths, but 3 levels up for images # {log_id}/sensors/cameras/ring_*/*.jpg vs. diff --git a/src/av2/evaluation/detection/eval.py b/src/av2/evaluation/detection/eval.py index 5bd4221c..df37a9d4 100644 --- a/src/av2/evaluation/detection/eval.py +++ b/src/av2/evaluation/detection/eval.py @@ -58,7 +58,9 @@ import numpy as np import pandas as pd +import polars as pl from tqdm import tqdm + from av2.evaluation.detection.constants import ( HIERARCHY, LCA, @@ -80,7 +82,6 @@ from av2.structures.cuboid import ORDERED_CUBOID_COL_NAMES from av2.utils.io import TimestampedCitySE3EgoPoses from av2.utils.typing import NDArrayBool, NDArrayFloat, NDArrayObject -import polars as pl warnings.filterwarnings("ignore", module="google") @@ -133,11 +134,15 @@ def evaluate( uuid_to_dts = { k: v[list(DTS_COLUMNS)].to_numpy().astype(float) - for k, v in dts_pl.partition_by(DETECTION_UUID_COLUMNS, maintain_order=True, as_dict=True).items() + for k, v in dts_pl.partition_by( + DETECTION_UUID_COLUMNS, maintain_order=True, as_dict=True + ).items() } uuid_to_gts = { k: v[list(GTS_COLUMNS)].to_numpy().astype(float) - for k, v in gts_pl.partition_by(DETECTION_UUID_COLUMNS, maintain_order=True, as_dict=True).items() + for k, v in gts_pl.partition_by( + DETECTION_UUID_COLUMNS, maintain_order=True, as_dict=True + ).items() } log_id_to_avm: Optional[Dict[str, ArgoverseStaticMap]] = None @@ -147,7 +152,9 @@ def evaluate( if cfg.eval_only_roi_instances and cfg.dataset_dir is not None: logger.info("Loading maps and egoposes ...") log_ids: List[str] = gts.loc[:, "log_id"].unique().tolist() - log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes(log_ids, cfg.dataset_dir) + log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes( + log_ids, cfg.dataset_dir + ) accumulate_args_list: List[ Tuple[ @@ -185,11 +192,15 @@ def evaluate( logger.info("Starting evaluation ...") with mp.get_context("spawn").Pool(processes=n_jobs) as p: - outputs: List[Tuple[NDArrayFloat, NDArrayFloat]] = p.starmap(accumulate, accumulate_args_list) + outputs: List[Tuple[NDArrayFloat, NDArrayFloat]] = p.starmap( + accumulate, accumulate_args_list + ) dts_list, gts_list = zip(*outputs) - METRIC_COLUMN_NAMES = cfg.affinity_thresholds_m + TP_ERROR_COLUMNS + ("is_evaluated",) + METRIC_COLUMN_NAMES = ( + cfg.affinity_thresholds_m + TP_ERROR_COLUMNS + ("is_evaluated",) + ) dts_metrics: NDArrayFloat = np.concatenate(dts_list) gts_metrics: NDArrayFloat = np.concatenate(gts_list) dts.loc[:, METRIC_COLUMN_NAMES] = dts_metrics @@ -218,7 +229,9 @@ def summarize_metrics( The summary metrics. """ # Sample recall values in the [0, 1] interval. - recall_interpolated: NDArrayFloat = np.linspace(0, 1, cfg.num_recall_samples, endpoint=True) + recall_interpolated: NDArrayFloat = np.linspace( + 0, 1, cfg.num_recall_samples, endpoint=True + ) # Initialize the summary metrics. summary = pd.DataFrame( @@ -226,7 +239,9 @@ def summarize_metrics( index=cfg.categories, ) - average_precisions = pd.DataFrame({t: 0.0 for t in cfg.affinity_thresholds_m}, index=cfg.categories) + average_precisions = pd.DataFrame( + {t: 0.0 for t in cfg.affinity_thresholds_m}, index=cfg.categories + ) for category in cfg.categories: # Find detections that have the current category. is_category_dts = dts["category"] == category @@ -235,7 +250,11 @@ def summarize_metrics( is_valid_dts = np.logical_and(is_category_dts, dts["is_evaluated"]) # Get valid detections and sort them in descending order. - category_dts = dts.loc[is_valid_dts].sort_values(by="score", ascending=False).reset_index(drop=True) + category_dts = ( + dts.loc[is_valid_dts] + .sort_values(by="score", ascending=False) + .reset_index(drop=True) + ) # Find annotations that have the current category. is_category_gts = gts["category"] == category @@ -248,19 +267,27 @@ def summarize_metrics( continue for affinity_threshold_m in cfg.affinity_thresholds_m: - true_positives: NDArrayBool = category_dts[affinity_threshold_m].astype(bool).to_numpy() + true_positives: NDArrayBool = ( + category_dts[affinity_threshold_m].astype(bool).to_numpy() + ) # Continue if there aren't any true positives. if len(true_positives) == 0: continue # Compute average precision for the current threshold. - threshold_average_precision, _ = compute_average_precision(true_positives, recall_interpolated, num_gts) + threshold_average_precision, _ = compute_average_precision( + true_positives, recall_interpolated, num_gts + ) # Record the average precision. - average_precisions.loc[category, affinity_threshold_m] = threshold_average_precision + average_precisions.loc[ + category, affinity_threshold_m + ] = threshold_average_precision - mean_average_precisions: NDArrayFloat = average_precisions.loc[category].to_numpy().mean() + mean_average_precisions: NDArrayFloat = ( + average_precisions.loc[category].to_numpy().mean() + ) # Select only the true positives for each instance. middle_idx = len(cfg.affinity_thresholds_m) // 2 @@ -330,21 +357,29 @@ def evaluate_hierarchy( uuid_to_dts = { cast(Tuple[str, int], k): v[list(DTS_COLUMNS)].to_numpy().astype(np.float64) - for k, v in dts_pl.partition_by(UUID_COLUMNS, maintain_order=True, as_dict=True).items() + for k, v in dts_pl.partition_by( + UUID_COLUMNS, maintain_order=True, as_dict=True + ).items() } uuid_to_gts = { cast(Tuple[str, int], k): v[list(GTS_COLUMNS)].to_numpy().astype(np.float64) - for k, v in gts_pl.partition_by(UUID_COLUMNS, maintain_order=True, as_dict=True).items() + for k, v in gts_pl.partition_by( + UUID_COLUMNS, maintain_order=True, as_dict=True + ).items() } uuid_to_dts_cats = { cast(Tuple[str, int], k): v[list(CATEGORY_COLUMN)].to_numpy().astype(np.object_) - for k, v in dts_pl.partition_by(UUID_COLUMNS, maintain_order=True, as_dict=True).items() + for k, v in dts_pl.partition_by( + UUID_COLUMNS, maintain_order=True, as_dict=True + ).items() } uuid_to_gts_cats = { cast(Tuple[str, int], k): v[list(CATEGORY_COLUMN)].to_numpy().astype(np.object_) - for k, v in gts_pl.partition_by(UUID_COLUMNS, maintain_order=True, as_dict=True).items() + for k, v in gts_pl.partition_by( + UUID_COLUMNS, maintain_order=True, as_dict=True + ).items() } log_id_to_avm: Optional[Dict[str, ArgoverseStaticMap]] = None @@ -354,7 +389,9 @@ def evaluate_hierarchy( if cfg.eval_only_roi_instances and cfg.dataset_dir is not None: logger.info("Loading maps and egoposes ...") log_ids: List[str] = gts.loc[:, "log_id"].unique().tolist() - log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes(log_ids, cfg.dataset_dir) + log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes( + log_ids, cfg.dataset_dir + ) is_evaluated_args_list: List[ Tuple[ @@ -421,9 +458,15 @@ def evaluate_hierarchy( logger.info("Starting evaluation ...") with mp.get_context("spawn").Pool(processes=n_jobs) as p: - outputs: List[Tuple[NDArrayFloat, NDArrayFloat, NDArrayObject, NDArrayObject, Tuple[str, int]]] = p.starmap( - is_evaluated, is_evaluated_args_list - ) + outputs: List[ + Tuple[ + NDArrayFloat, + NDArrayFloat, + NDArrayObject, + NDArrayObject, + Tuple[str, int], + ] + ] = p.starmap(is_evaluated, is_evaluated_args_list) dts_list: List[NDArrayFloat] = [] gts_list: List[NDArrayFloat] = [] diff --git a/src/av2/evaluation/detection/utils.py b/src/av2/evaluation/detection/utils.py index f6e6a136..60266f2e 100644 --- a/src/av2/evaluation/detection/utils.py +++ b/src/av2/evaluation/detection/utils.py @@ -16,6 +16,10 @@ from typing import Any, Dict, List, Optional, Set, Tuple, Union import numpy as np +from joblib import Parallel, delayed +from scipy.spatial.distance import cdist +from upath import UPath + from av2.evaluation import NUM_RECALL_SAMPLES, SensorCompetitionCategories from av2.evaluation.detection.constants import ( MAX_NORMALIZED_ASE, @@ -37,9 +41,6 @@ from av2.utils.constants import EPS from av2.utils.io import TimestampedCitySE3EgoPoses, read_city_SE3_ego from av2.utils.typing import NDArrayBool, NDArrayFloat, NDArrayInt, NDArrayObject -from joblib import Parallel, delayed -from scipy.spatial.distance import cdist -from upath import UPath logger = logging.getLogger(__name__) @@ -143,7 +144,9 @@ def accumulate( is_evaluated_gts &= compute_objects_in_roi_mask(gts, city_SE3_ego, avm) is_evaluated_dts &= compute_evaluated_dts_mask(dts[..., :3], cfg) - is_evaluated_gts &= compute_evaluated_gts_mask(gts[..., :3], gts[..., -1].astype(int), cfg) + is_evaluated_gts &= compute_evaluated_gts_mask( + gts[..., :3], gts[..., -1].astype(int), cfg + ) # Initialize results array. dts_augmented: NDArrayFloat = np.zeros((N, T + E + 1)) @@ -155,7 +158,9 @@ def accumulate( if is_evaluated_dts.sum() > 0 and is_evaluated_gts.sum() > 0: # Compute true positives by assigning detections and ground truths. - dts_assignments, gts_assignments = assign(dts[is_evaluated_dts], gts[is_evaluated_gts], cfg) + dts_assignments, gts_assignments = assign( + dts[is_evaluated_dts], gts[is_evaluated_gts], cfg + ) dts_augmented[is_evaluated_dts, :-1] = dts_assignments gts_augmented[is_evaluated_gts, :-1] = gts_assignments @@ -204,7 +209,9 @@ def is_evaluated( is_evaluated_gts &= compute_objects_in_roi_mask(gts, city_SE3_ego, avm) is_evaluated_dts &= compute_evaluated_dts_mask(dts[..., :3], cfg) - is_evaluated_gts &= compute_evaluated_gts_mask(gts[..., :3], gts[..., -1].astype(int), cfg) + is_evaluated_gts &= compute_evaluated_gts_mask( + gts[..., :3], gts[..., -1].astype(int), cfg + ) dts = dts[is_evaluated_dts] gts = gts[is_evaluated_gts] @@ -293,7 +300,9 @@ def accumulate_hierarchy( match_gt_idx = len(cfg.affinity_thresholds_m) * [None] if len(gts_uuids) > 0: - keep_sweep = np.all(gts_uuids == np.array([gts.shape[0] * [pred_uuid]]).squeeze(), axis=1) + keep_sweep = np.all( + gts_uuids == np.array([gts.shape[0] * [pred_uuid]]).squeeze(), axis=1 + ) else: keep_sweep = [] @@ -307,26 +316,36 @@ def accumulate_hierarchy( # Find closest match among ground truth boxes for i in range(len(cfg.affinity_thresholds_m)): - if gt_cat == cat and not (pred_uuid[0], pred_uuid[1], gt_idx) in taken[i]: + if ( + gt_cat == cat + and not (pred_uuid[0], pred_uuid[1], gt_idx) in taken[i] + ): this_distance = dist_mat[pred_idx][gt_idx] if this_distance < min_dist[i]: min_dist[i] = this_distance match_gt_idx[i] = gt_idx - is_match = [min_dist[i] < dist_th for i, dist_th in enumerate(cfg.affinity_thresholds_m)] + is_match = [ + min_dist[i] < dist_th for i, dist_th in enumerate(cfg.affinity_thresholds_m) + ] for gt in zip(gt_ind_sweep, gts_sweep, gts_cats_sweep, gts_uuids_sweep): gt_idx, gt_box, gt_cat, gt_uuid = gt # Find closest match among ground truth boxes for i in range(len(cfg.affinity_thresholds_m)): - if not is_match[i] and not (pred_uuid[0], pred_uuid[1], gt_idx) in taken[i]: + if ( + not is_match[i] + and not (pred_uuid[0], pred_uuid[1], gt_idx) in taken[i] + ): this_distance = dist_mat[pred_idx][gt_idx] if this_distance < min_dist[i]: min_dist[i] = this_distance match_gt_idx[i] = gt_idx - is_dist = [min_dist[i] < dist_th for i, dist_th in enumerate(cfg.affinity_thresholds_m)] + is_dist = [ + min_dist[i] < dist_th for i, dist_th in enumerate(cfg.affinity_thresholds_m) + ] is_match = [ True if is_dist[i] and gts_cats[match_gt_idx[i]] == cat else False for i in range(len(cfg.affinity_thresholds_m)) @@ -367,7 +386,9 @@ def accumulate_hierarchy( prec = tp[i] / (fp[i] + tp[i]) rec = tp[i] / float(npos) - rec_interp = np.linspace(0, 1, NUM_RECALL_SAMPLES) # 101 steps, from 0% to 100% recall. + rec_interp = np.linspace( + 0, 1, NUM_RECALL_SAMPLES + ) # 101 steps, from 0% to 100% recall. ap = np.mean(np.interp(rec_interp, rec, prec, right=0)) mAP.append(round(ap, NUM_DECIMALS)) @@ -375,7 +396,9 @@ def accumulate_hierarchy( return float(np.mean(mAP)), cat, lca -def assign(dts: NDArrayFloat, gts: NDArrayFloat, cfg: DetectionCfg) -> Tuple[NDArrayFloat, NDArrayFloat]: +def assign( + dts: NDArrayFloat, gts: NDArrayFloat, cfg: DetectionCfg +) -> Tuple[NDArrayFloat, NDArrayFloat]: """Attempt assignment of each detection to a ground truth label. The detections (gts) and ground truth annotations (gts) are expected to be shape (N,10) and (M,10) @@ -400,7 +423,9 @@ def assign(dts: NDArrayFloat, gts: NDArrayFloat, cfg: DetectionCfg) -> Tuple[NDA $$T$$: cfg.affinity_thresholds_m (0.5, 1.0, 2.0, 4.0 by default). $$E$$: ATE, ASE, AOE. """ - affinity_matrix = compute_affinity_matrix(dts[..., :3], gts[..., :3], cfg.affinity_type) + affinity_matrix = compute_affinity_matrix( + dts[..., :3], gts[..., :3], cfg.affinity_type + ) # Get the GT label for each max-affinity GT label, detection pair. idx_gts: NDArrayInt = affinity_matrix.argmax(axis=1)[None] @@ -408,7 +433,9 @@ def assign(dts: NDArrayFloat, gts: NDArrayFloat, cfg: DetectionCfg) -> Tuple[NDA # The affinity matrix is an N by M matrix of the detections and ground truth labels respectively. # We want to take the corresponding affinity for each of the initial assignments using `gt_matches`. # The following line grabs the max affinity for each detection to a ground truth label. - affinities: NDArrayFloat = np.take_along_axis(affinity_matrix.transpose(), idx_gts, axis=0)[0] + affinities: NDArrayFloat = np.take_along_axis( + affinity_matrix.transpose(), idx_gts, axis=0 + )[0] # Find the indices of the _first_ detection assigned to each GT. assignments: Tuple[NDArrayInt, NDArrayInt] = np.unique(idx_gts, return_index=True) @@ -436,14 +463,22 @@ def assign(dts: NDArrayFloat, gts: NDArrayFloat, cfg: DetectionCfg) -> Tuple[NDA tps_dts = dts[idx_tps_dts] tps_gts = gts[idx_tps_gts] - translation_errors = distance(tps_dts[:, :3], tps_gts[:, :3], DistanceType.TRANSLATION) + translation_errors = distance( + tps_dts[:, :3], tps_gts[:, :3], DistanceType.TRANSLATION + ) scale_errors = distance(tps_dts[:, 3:6], tps_gts[:, 3:6], DistanceType.SCALE) - orientation_errors = distance(tps_dts[:, 6:10], tps_gts[:, 6:10], DistanceType.ORIENTATION) - dts_metrics[idx_tps_dts, 4:] = np.stack((translation_errors, scale_errors, orientation_errors), axis=-1) + orientation_errors = distance( + tps_dts[:, 6:10], tps_gts[:, 6:10], DistanceType.ORIENTATION + ) + dts_metrics[idx_tps_dts, 4:] = np.stack( + (translation_errors, scale_errors, orientation_errors), axis=-1 + ) return dts_metrics, gts_metrics -def interpolate_precision(precision: NDArrayFloat, interpolation_method: InterpType = InterpType.ALL) -> NDArrayFloat: +def interpolate_precision( + precision: NDArrayFloat, interpolation_method: InterpType = InterpType.ALL +) -> NDArrayFloat: r"""Interpolate the precision at each sampled recall. This function smooths the precision-recall curve according to the method introduced in Pascal @@ -473,7 +508,9 @@ def interpolate_precision(precision: NDArrayFloat, interpolation_method: InterpT return precision_interpolated -def compute_affinity_matrix(dts: NDArrayFloat, gts: NDArrayFloat, metric: AffinityType) -> NDArrayFloat: +def compute_affinity_matrix( + dts: NDArrayFloat, gts: NDArrayFloat, metric: AffinityType +) -> NDArrayFloat: """Calculate the affinity matrix between detections and ground truth annotations. Args: @@ -523,13 +560,17 @@ def compute_average_precision( precision = interpolate_precision(precision) # Evaluate precision at different recalls. - precision_interpolated: NDArrayFloat = np.interp(recall_interpolated, recall, precision, right=0) + precision_interpolated: NDArrayFloat = np.interp( + recall_interpolated, recall, precision, right=0 + ) average_precision = np.mean(precision_interpolated).astype(float) return average_precision, precision_interpolated -def distance(dts: NDArrayFloat, gts: NDArrayFloat, metric: DistanceType) -> NDArrayFloat: +def distance( + dts: NDArrayFloat, gts: NDArrayFloat, metric: DistanceType +) -> NDArrayFloat: """Distance functions between detections and ground truth. Args: @@ -558,7 +599,9 @@ def distance(dts: NDArrayFloat, gts: NDArrayFloat, metric: DistanceType) -> NDAr raise NotImplementedError("This distance metric is not implemented!") -def compute_objects_in_roi_mask(cuboids_ego: NDArrayFloat, city_SE3_ego: SE3, avm: ArgoverseStaticMap) -> NDArrayBool: +def compute_objects_in_roi_mask( + cuboids_ego: NDArrayFloat, city_SE3_ego: SE3, avm: ArgoverseStaticMap +) -> NDArrayBool: """Compute the evaluated cuboids mask based off whether _any_ of their vertices fall into the ROI. Args: @@ -573,7 +616,9 @@ def compute_objects_in_roi_mask(cuboids_ego: NDArrayFloat, city_SE3_ego: SE3, av if len(cuboids_ego) == 0: is_within_roi = np.zeros((0,), dtype=bool) return is_within_roi - cuboid_list_ego: CuboidList = CuboidList([Cuboid.from_numpy(params) for params in cuboids_ego]) + cuboid_list_ego: CuboidList = CuboidList( + [Cuboid.from_numpy(params) for params in cuboids_ego] + ) cuboid_list_city = cuboid_list_ego.transform(city_SE3_ego) cuboid_list_vertices_m_city = cuboid_list_city.vertices_m @@ -661,9 +706,14 @@ def load_mapped_avm_and_egoposes( Raises: RuntimeError: If the process for loading maps and timestamped egoposes fails. """ - log_id_to_timestamped_poses = {log_id: read_city_SE3_ego(dataset_dir / log_id) for log_id in log_ids} + log_id_to_timestamped_poses = { + log_id: read_city_SE3_ego(dataset_dir / log_id) for log_id in log_ids + } avms: Optional[List[ArgoverseStaticMap]] = Parallel(n_jobs=-1, backend="threading")( - delayed(ArgoverseStaticMap.from_map_dir)(dataset_dir / log_id / "map", build_raster=True) for log_id in log_ids + delayed(ArgoverseStaticMap.from_map_dir)( + dataset_dir / log_id / "map", build_raster=True + ) + for log_id in log_ids ) if avms is None: diff --git a/src/av2/evaluation/forecasting/constants.py b/src/av2/evaluation/forecasting/constants.py index 14034cd7..e5a7bb7c 100644 --- a/src/av2/evaluation/forecasting/constants.py +++ b/src/av2/evaluation/forecasting/constants.py @@ -3,6 +3,7 @@ from typing import Final import numpy as np + from av2.evaluation import SensorCompetitionCategories NUM_TIMESTEPS: Final = 6 diff --git a/src/av2/evaluation/forecasting/eval.py b/src/av2/evaluation/forecasting/eval.py index aee25928..c33cbefd 100644 --- a/src/av2/evaluation/forecasting/eval.py +++ b/src/av2/evaluation/forecasting/eval.py @@ -14,19 +14,24 @@ import click import numpy as np +from scipy.spatial.transform import Rotation +from tqdm import tqdm + from av2.evaluation import NUM_RECALL_SAMPLES -from av2.evaluation.detection.utils import compute_objects_in_roi_mask, load_mapped_avm_and_egoposes +from av2.evaluation.detection.utils import ( + compute_objects_in_roi_mask, + load_mapped_avm_and_egoposes, +) from av2.evaluation.forecasting import constants, utils from av2.evaluation.forecasting.constants import ( AV2_CATEGORIES, CATEGORY_TO_VELOCITY_M_PER_S, DISTANCE_THRESHOLDS_M, - VELOCITY_TYPES, MAX_DISPLACEMENT_M, + VELOCITY_TYPES, ) from av2.utils.typing import NDArrayFloat -from scipy.spatial.transform import Rotation -from tqdm import tqdm + from ..typing import ForecastSequences, Sequences @@ -83,7 +88,9 @@ def evaluate( agent["seq_id"] = seq_id agent["timestamp_ns"] = timestamp_ns agent["velocity_m_per_s"] = utils.agent_velocity_m_per_s(agent) - agent["trajectory_type"] = utils.trajectory_type(agent, CATEGORY_TO_VELOCITY_M_PER_S) + agent["trajectory_type"] = utils.trajectory_type( + agent, CATEGORY_TO_VELOCITY_M_PER_S + ) gt_agents.append(agent) @@ -91,7 +98,9 @@ def evaluate( agent["seq_id"] = seq_id agent["timestamp_ns"] = timestamp_ns agent["velocity_m_per_s"] = utils.agent_velocity_m_per_s(agent) - agent["trajectory_type"] = utils.trajectory_type(agent, CATEGORY_TO_VELOCITY_M_PER_S) + agent["trajectory_type"] = utils.trajectory_type( + agent, CATEGORY_TO_VELOCITY_M_PER_S + ) pred_agents.append(agent) @@ -101,7 +110,15 @@ def evaluate( ): category_velocity_m_per_s = CATEGORY_TO_VELOCITY_M_PER_S[category] outputs.append( - accumulate(pred_agents, gt_agents, top_k, category, velocity_type, category_velocity_m_per_s, th) + accumulate( + pred_agents, + gt_agents, + top_k, + category, + velocity_type, + category_velocity_m_per_s, + th, + ) ) for apf, ade, fde, category, velocity_type, threshold in outputs: @@ -114,7 +131,8 @@ def evaluate( for category in AV2_CATEGORIES: for velocity_type in VELOCITY_TYPES: results[velocity_type][category]["mAP_F"] = round( - np.mean(results[velocity_type][category]["mAP_F"]), constants.NUM_DECIMALS + np.mean(results[velocity_type][category]["mAP_F"]), + constants.NUM_DECIMALS, ) results[velocity_type][category]["ADE"] = round( np.mean(results[velocity_type][category]["ADE"]), constants.NUM_DECIMALS @@ -160,7 +178,11 @@ def match(gt: str, pred: str, profile: str) -> bool: return gt == profile or gt == "ignore" and pred == profile pred = [agent for agent in pred_agents if agent["name"] == class_name] - gt = [agent for agent in gt_agents if agent["name"] == class_name and agent["trajectory_type"] == profile] + gt = [ + agent + for agent in gt_agents + if agent["name"] == class_name and agent["trajectory_type"] == profile + ] conf = [agent["detection_score"] for agent in pred] sortind = [i for (v, i) in sorted((v, i) for (i, v) in enumerate(conf))][::-1] gt_agents_by_frame = defaultdict(list) @@ -179,12 +201,15 @@ def match(gt: str, pred: str, profile: str) -> bool: min_dist = np.inf match_gt_idx = None - gt_agents_in_frame = gt_agents_by_frame[(pred_agent["seq_id"], pred_agent["timestamp_ns"])] + gt_agents_in_frame = gt_agents_by_frame[ + (pred_agent["seq_id"], pred_agent["timestamp_ns"]) + ] for gt_idx, gt_agent in enumerate(gt_agents_in_frame): if not (pred_agent["seq_id"], pred_agent["timestamp_ns"], gt_idx) in taken: # Find closest match among ground truth boxes this_distance = utils.center_distance( - gt_agent["current_translation_m"], pred_agent["current_translation_m"] + gt_agent["current_translation_m"], + pred_agent["current_translation_m"], ) if this_distance < min_dist: min_dist = this_distance @@ -198,7 +223,10 @@ def match(gt: str, pred: str, profile: str) -> bool: gt_match_agent = gt_agents_in_frame[match_gt_idx] gt_len = gt_match_agent["future_translation_m"].shape[0] - forecast_match_th = [threshold + constants.FORECAST_SCALAR[i] * velocity for i in range(gt_len + 1)] + forecast_match_th = [ + threshold + constants.FORECAST_SCALAR[i] * velocity + for i in range(gt_len + 1) + ] if top_k == 1: ind = cast(int, np.argmax(pred_agent["score"])) @@ -209,7 +237,9 @@ def match(gt: str, pred: str, profile: str) -> bool: ) for i in range(gt_len) ] - forecast_match = [dist < th for dist, th in zip(forecast_dist, forecast_match_th[1:])] + forecast_match = [ + dist < th for dist, th in zip(forecast_dist, forecast_match_th[1:]) + ] ade = cast(float, np.mean(forecast_dist)) fde = forecast_dist[-1] @@ -226,7 +256,10 @@ def match(gt: str, pred: str, profile: str) -> bool: ) for i in range(gt_len) ] - curr_forecast_match = [dist < th for dist, th in zip(curr_forecast_dist, forecast_match_th[1:])] + curr_forecast_match = [ + dist < th + for dist, th in zip(curr_forecast_dist, forecast_match_th[1:]) + ] curr_ade = cast(float, np.mean(curr_forecast_dist)) curr_fde = curr_forecast_dist[-1] @@ -276,7 +309,9 @@ def match(gt: str, pred: str, profile: str) -> bool: prec = tp_array / (fp_array + tp_array) rec = tp_array / float(npos) - rec_interp = np.linspace(0, 1, NUM_RECALL_SAMPLES) # 101 steps, from 0% to 100% recall. + rec_interp = np.linspace( + 0, 1, NUM_RECALL_SAMPLES + ) # 101 steps, from 0% to 100% recall. apf = np.mean(np.interp(rec_interp, rec, prec, right=0)) return ( @@ -303,13 +338,19 @@ def convert_forecast_labels(labels: Any) -> Any: frame_dict = {} for frame_idx, frame in enumerate(frames): forecast_instances = [] - for instance in utils.array_dict_iterator(frame, len(frame["translation_m"])): + for instance in utils.array_dict_iterator( + frame, len(frame["translation_m"]) + ): future_translations: Any = [] - for future_frame in frames[frame_idx + 1 : frame_idx + 1 + constants.NUM_TIMESTEPS]: + for future_frame in frames[ + frame_idx + 1 : frame_idx + 1 + constants.NUM_TIMESTEPS + ]: if instance["track_id"] not in future_frame["track_id"]: break future_translations.append( - future_frame["translation_m"][future_frame["track_id"] == instance["track_id"]][0] + future_frame["translation_m"][ + future_frame["track_id"] == instance["track_id"] + ][0] ) if len(future_translations) == 0: @@ -335,7 +376,9 @@ def convert_forecast_labels(labels: Any) -> Any: return forecast_labels -def filter_max_dist(forecasts: ForecastSequences, max_range_m: int) -> ForecastSequences: +def filter_max_dist( + forecasts: ForecastSequences, max_range_m: int +) -> ForecastSequences: """Remove all tracks that are beyond `max_range_m`. Args: @@ -351,7 +394,10 @@ def filter_max_dist(forecasts: ForecastSequences, max_range_m: int) -> ForecastS agent for agent in forecasts[seq_id][timestamp_ns] if "ego_translation_m" in agent - and np.linalg.norm(agent["current_translation_m"] - agent["ego_translation_m"]) < max_range_m + and np.linalg.norm( + agent["current_translation_m"] - agent["ego_translation_m"] + ) + < max_range_m ] forecasts[seq_id][timestamp_ns] = keep_forecasts @@ -370,7 +416,9 @@ def yaw_to_quaternion3d(yaw: float) -> NDArrayFloat: return np.array([qw, qx, qy, qz]) -def filter_drivable_area(forecasts: ForecastSequences, dataset_dir: str) -> ForecastSequences: +def filter_drivable_area( + forecasts: ForecastSequences, dataset_dir: str +) -> ForecastSequences: """Convert the unified label format to a format that is easier to work with for forecasting evaluation. Args: @@ -381,7 +429,9 @@ def filter_drivable_area(forecasts: ForecastSequences, dataset_dir: str) -> Fore forecasts: Dictionary of tracks. """ log_ids = list(forecasts.keys()) - log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes(log_ids, Path(dataset_dir)) + log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes( + log_ids, Path(dataset_dir) + ) for log_id in log_ids: avm = log_id_to_avm[log_id] @@ -395,7 +445,9 @@ def filter_drivable_area(forecasts: ForecastSequences, dataset_dir: str) -> Fore continue for box in forecasts[log_id][timestamp_ns]: - translation_m.append(box["current_translation_m"] - box["ego_translation_m"]) + translation_m.append( + box["current_translation_m"] - box["ego_translation_m"] + ) size.append(box["size"]) quat.append(yaw_to_quaternion3d(box["yaw"])) @@ -411,7 +463,9 @@ def filter_drivable_area(forecasts: ForecastSequences, dataset_dir: str) -> Fore ) is_evaluated = compute_objects_in_roi_mask(boxes, city_SE3_ego, avm) - forecasts[log_id][timestamp_ns] = list(np.array(forecasts[log_id][timestamp_ns])[is_evaluated]) + forecasts[log_id][timestamp_ns] = list( + np.array(forecasts[log_id][timestamp_ns])[is_evaluated] + ) return forecasts @@ -441,9 +495,27 @@ def runner( res = evaluate(predictions2, ground_truth2, top_k, max_range_m, dataset_dir) - mAP_F = np.nanmean([metrics["mAP_F"] for traj_metrics in res.values() for metrics in traj_metrics.values()]) - ADE = np.nanmean([metrics["ADE"] for traj_metrics in res.values() for metrics in traj_metrics.values()]) - FDE = np.nanmean([metrics["FDE"] for traj_metrics in res.values() for metrics in traj_metrics.values()]) + mAP_F = np.nanmean( + [ + metrics["mAP_F"] + for traj_metrics in res.values() + for metrics in traj_metrics.values() + ] + ) + ADE = np.nanmean( + [ + metrics["ADE"] + for traj_metrics in res.values() + for metrics in traj_metrics.values() + ] + ) + FDE = np.nanmean( + [ + metrics["FDE"] + for traj_metrics in res.values() + for metrics in traj_metrics.values() + ] + ) res["mean_mAP_F"] = mAP_F res["mean_ADE"] = ADE res["mean_FDE"] = FDE diff --git a/src/av2/evaluation/forecasting/utils.py b/src/av2/evaluation/forecasting/utils.py index db4ec114..61d912ce 100644 --- a/src/av2/evaluation/forecasting/utils.py +++ b/src/av2/evaluation/forecasting/utils.py @@ -3,8 +3,10 @@ from typing import Any, Dict, Iterable, List, Union, cast import numpy as np + from av2.evaluation.forecasting import constants from av2.utils.typing import NDArrayFloat, NDArrayInt + from ..typing import ForecastSequences, Frame @@ -23,18 +25,24 @@ def agent_velocity_m_per_s(agent: Dict[str, Any]) -> NDArrayFloat: if "future_translation_m" in agent: # ground_truth return cast( NDArrayFloat, - (agent["future_translation_m"][0][:2] - agent["current_translation_m"][:2]) / constants.TIME_DELTA_S, + (agent["future_translation_m"][0][:2] - agent["current_translation_m"][:2]) + / constants.TIME_DELTA_S, ) else: # predictions res = [] for i in range(agent["prediction_m"].shape[0]): - res.append((agent["prediction_m"][i][0][:2] - agent["current_translation_m"][:2]) / constants.TIME_DELTA_S) + res.append( + (agent["prediction_m"][i][0][:2] - agent["current_translation_m"][:2]) + / constants.TIME_DELTA_S + ) return np.stack(res) -def trajectory_type(agent: Dict[str, Any], category_velocity_m_per_s: Dict[str, float]) -> Union[str, List[str]]: +def trajectory_type( + agent: Dict[str, Any], category_velocity_m_per_s: Dict[str, float] +) -> Union[str, List[str]]: """Get the trajectory type, which is either static, linear or non-linear. Trajectory is static if the prediction error of a static forecast is below threshold. @@ -53,13 +61,15 @@ def trajectory_type(agent: Dict[str, Any], category_velocity_m_per_s: Dict[str, if "future_translation_m" in agent: # ground_truth time = agent["future_translation_m"].shape[0] * constants.TIME_DELTA_S static_target = agent["current_translation_m"][:2] - linear_target = agent["current_translation_m"][:2] + time * agent["velocity_m_per_s"][:2] + linear_target = ( + agent["current_translation_m"][:2] + time * agent["velocity_m_per_s"][:2] + ) final_position = agent["future_translation_m"][-1][:2] - threshold = 1 + constants.FORECAST_SCALAR[len(agent["future_translation_m"])] * category_velocity_m_per_s.get( - agent["name"], 0 - ) + threshold = 1 + constants.FORECAST_SCALAR[ + len(agent["future_translation_m"]) + ] * category_velocity_m_per_s.get(agent["name"], 0) if np.linalg.norm(final_position - static_target) < threshold: return "static" elif np.linalg.norm(final_position - linear_target) < threshold: @@ -71,12 +81,15 @@ def trajectory_type(agent: Dict[str, Any], category_velocity_m_per_s: Dict[str, res: List[str] = [] time = agent["prediction_m"].shape[1] * constants.TIME_DELTA_S - threshold = 1 + constants.FORECAST_SCALAR[len(agent["prediction_m"])] * category_velocity_m_per_s.get( - agent["name"], 0 - ) + threshold = 1 + constants.FORECAST_SCALAR[ + len(agent["prediction_m"]) + ] * category_velocity_m_per_s.get(agent["name"], 0) for i in range(agent["prediction_m"].shape[0]): static_target = agent["current_translation_m"][:2] - linear_target = agent["current_translation_m"][:2] + time * agent["velocity_m_per_s"][i][:2] + linear_target = ( + agent["current_translation_m"][:2] + + time * agent["velocity_m_per_s"][i][:2] + ) final_position = agent["prediction_m"][i][-1][:2] @@ -126,11 +139,15 @@ def index_array_values(array_dict: Frame, index: Union[int, NDArrayInt]) -> Fram Returns: Dictionary of numpy arrays, each indexed by the provided index """ - return {k: v[index] if isinstance(v, np.ndarray) else v for k, v in array_dict.items()} + return { + k: v[index] if isinstance(v, np.ndarray) else v for k, v in array_dict.items() + } def annotate_frame_metadata( - predictions: ForecastSequences, ground_truth: ForecastSequences, metadata_keys: List[str] + predictions: ForecastSequences, + ground_truth: ForecastSequences, + metadata_keys: List[str], ) -> ForecastSequences: """Index each numpy array in dictionary. diff --git a/src/av2/evaluation/scene_flow/constants.py b/src/av2/evaluation/scene_flow/constants.py index 7195b79d..d6e57cdf 100644 --- a/src/av2/evaluation/scene_flow/constants.py +++ b/src/av2/evaluation/scene_flow/constants.py @@ -10,7 +10,10 @@ SCENE_FLOW_DYNAMIC_THRESHOLD: Final = 0.05 SWEEP_PAIR_TIME_DELTA: Final = 0.1 -CATEGORY_TO_INDEX: Final = {**{"NONE": 0}, **{k.value: i + 1 for i, k in enumerate(AnnotationCategories)}} +CATEGORY_TO_INDEX: Final = { + **{"NONE": 0}, + **{k.value: i + 1 for i, k in enumerate(AnnotationCategories)}, +} @unique @@ -102,7 +105,10 @@ class MetricBreakdownCategories(str, Enum): MetricBreakdownCategories.FOREGROUND: [ CATEGORY_TO_INDEX[k.value] for k in ( - list(InanimateCategories) + list(LeggedCategories) + list(SmallVehicleCategories) + list(VehicleCategories) + list(InanimateCategories) + + list(LeggedCategories) + + list(SmallVehicleCategories) + + list(VehicleCategories) ) ], } diff --git a/src/av2/evaluation/scene_flow/eval.py b/src/av2/evaluation/scene_flow/eval.py index 5b1be250..a9496135 100644 --- a/src/av2/evaluation/scene_flow/eval.py +++ b/src/av2/evaluation/scene_flow/eval.py @@ -5,17 +5,32 @@ import zipfile from collections import defaultdict from pathlib import Path -from typing import Any, Callable, DefaultDict, Dict, Final, List, Optional, Tuple, Union, cast +from typing import ( + Any, + Callable, + DefaultDict, + Dict, + Final, + List, + Optional, + Tuple, + Union, + cast, +) from zipfile import ZipFile -import av2.evaluation.scene_flow.constants as constants import click import numpy as np import pandas as pd -from av2.evaluation.scene_flow.constants import SceneFlowMetricType, SegmentationMetricType -from av2.utils.typing import NDArrayBool, NDArrayFloat, NDArrayInt from rich.progress import track +import av2.evaluation.scene_flow.constants as constants +from av2.evaluation.scene_flow.constants import ( + SceneFlowMetricType, + SegmentationMetricType, +) +from av2.utils.typing import NDArrayBool, NDArrayFloat, NDArrayInt + ACCURACY_RELAX_DISTANCE_THRESHOLD: Final = 0.1 ACCURACY_STRICT_DISTANCE_THRESHOLD: Final = 0.05 NO_FMT_INDICES: Final = ("Background", "Dynamic") @@ -32,11 +47,15 @@ def compute_end_point_error(dts: NDArrayFloat, gts: NDArrayFloat) -> NDArrayFloa Returns: The point-wise end-point error. """ - end_point_error: NDArrayFloat = np.linalg.norm(dts - gts, axis=-1).astype(np.float64) + end_point_error: NDArrayFloat = np.linalg.norm(dts - gts, axis=-1).astype( + np.float64 + ) return end_point_error -def compute_accuracy(dts: NDArrayFloat, gts: NDArrayFloat, distance_threshold: float) -> NDArrayFloat: +def compute_accuracy( + dts: NDArrayFloat, gts: NDArrayFloat, distance_threshold: float +) -> NDArrayFloat: """Compute the percent of inliers for a given threshold for a set of prediction and ground truth vectors. Args: @@ -52,7 +71,9 @@ def compute_accuracy(dts: NDArrayFloat, gts: NDArrayFloat, distance_threshold: f relative_error = np.divide(l2_norm, gts_norm + EPS) abs_error_inlier = np.less(l2_norm, distance_threshold).astype(bool) relative_error_inlier = np.less(relative_error, distance_threshold).astype(bool) - accuracy: NDArrayFloat = np.logical_or(abs_error_inlier, relative_error_inlier).astype(np.float64) + accuracy: NDArrayFloat = np.logical_or( + abs_error_inlier, relative_error_inlier + ).astype(np.float64) return accuracy @@ -93,8 +114,12 @@ def compute_angle_error(dts: NDArrayFloat, gts: NDArrayFloat) -> NDArrayFloat: The pointwise angle errors in space-time. """ # Convert the 3D flow vectors to 4D space-time vectors. - dts_space_time = np.pad(dts, ((0, 0), (0, 1)), constant_values=constants.SWEEP_PAIR_TIME_DELTA) - gts_space_time = np.pad(gts, ((0, 0), (0, 1)), constant_values=constants.SWEEP_PAIR_TIME_DELTA) + dts_space_time = np.pad( + dts, ((0, 0), (0, 1)), constant_values=constants.SWEEP_PAIR_TIME_DELTA + ) + gts_space_time = np.pad( + gts, ((0, 0), (0, 1)), constant_values=constants.SWEEP_PAIR_TIME_DELTA + ) dts_space_time_norm = np.linalg.norm(dts_space_time, axis=-1, keepdims=True) gts_space_time_norm = np.linalg.norm(gts_space_time, axis=-1, keepdims=True) @@ -186,7 +211,9 @@ def compute_scene_flow_metrics( elif scene_flow_metric_type == SceneFlowMetricType.EPE: return compute_end_point_error(dts, gts) else: - raise NotImplementedError(f"The scene flow metric type {scene_flow_metric_type} is not implemented!") + raise NotImplementedError( + f"The scene flow metric type {scene_flow_metric_type} is not implemented!" + ) def compute_segmentation_metrics( @@ -214,7 +241,9 @@ def compute_segmentation_metrics( elif segmentation_metric_type == SegmentationMetricType.FN: return compute_false_negatives(dts, gts) else: - raise NotImplementedError(f"The segmentation metric type {segmentation_metric_type} is not implemented!") + raise NotImplementedError( + f"The segmentation metric type {segmentation_metric_type} is not implemented!" + ) def compute_metrics( @@ -275,11 +304,15 @@ def compute_metrics( if subset_size > 0: for flow_metric_type in SceneFlowMetricType: results[flow_metric_type] += [ - compute_scene_flow_metrics(pred_sub, gts_sub, flow_metric_type).mean() + compute_scene_flow_metrics( + pred_sub, gts_sub, flow_metric_type + ).mean() ] for seg_metric_type in SegmentationMetricType: results[seg_metric_type] += [ - compute_segmentation_metrics(pred_dynamic[mask], is_dynamic[mask], seg_metric_type) + compute_segmentation_metrics( + pred_dynamic[mask], is_dynamic[mask], seg_metric_type + ) ] else: for flow_metric_type in SceneFlowMetricType: @@ -289,7 +322,9 @@ def compute_metrics( return results -def evaluate_predictions(annotations_dir: Path, get_prediction: Callable[[Path], pd.DataFrame]) -> pd.DataFrame: +def evaluate_predictions( + annotations_dir: Path, get_prediction: Callable[[Path], pd.DataFrame] +) -> pd.DataFrame: """Run the evaluation on predictions and labels. Args: @@ -331,7 +366,9 @@ def evaluate_predictions(annotations_dir: Path, get_prediction: Callable[[Path], return df -def get_prediction_from_directory(annotation_name: Path, predictions_dir: Path) -> Optional[pd.DataFrame]: +def get_prediction_from_directory( + annotation_name: Path, predictions_dir: Path +) -> Optional[pd.DataFrame]: """Get the prediction corresponding annotation from a directory of prediction files. Args: @@ -348,7 +385,9 @@ def get_prediction_from_directory(annotation_name: Path, predictions_dir: Path) return pred -def get_prediction_from_zipfile(annotation_name: Path, predictions_zip: Path) -> Optional[pd.DataFrame]: +def get_prediction_from_zipfile( + annotation_name: Path, predictions_zip: Path +) -> Optional[pd.DataFrame]: """Get the prediction corresponding annotation from a zip archive of prediction files. Args: @@ -377,7 +416,9 @@ def evaluate_directories(annotations_dir: Path, predictions_dir: Path) -> pd.Dat Returns: DataFrame containing the average metrics on each subset of each example. """ - return evaluate_predictions(annotations_dir, lambda n: get_prediction_from_directory(n, predictions_dir)) + return evaluate_predictions( + annotations_dir, lambda n: get_prediction_from_directory(n, predictions_dir) + ) def evaluate_zip(annotations_dir: Path, predictions_zip: Path) -> pd.DataFrame: @@ -390,7 +431,9 @@ def evaluate_zip(annotations_dir: Path, predictions_zip: Path) -> pd.DataFrame: Returns: DataFrame containing the average metrics on each subset of each example. """ - return evaluate_predictions(annotations_dir, lambda n: get_prediction_from_zipfile(n, predictions_zip)) + return evaluate_predictions( + annotations_dir, lambda n: get_prediction_from_zipfile(n, predictions_zip) + ) def results_to_dict(frame: pd.DataFrame) -> Dict[str, float]: @@ -405,7 +448,9 @@ def results_to_dict(frame: pd.DataFrame) -> Dict[str, float]: output = {} grouped = frame.groupby(["Class", "Motion", "Distance"]) - def weighted_average(x: pd.DataFrame, metric_type: Union[SceneFlowMetricType, SegmentationMetricType]) -> float: + def weighted_average( + x: pd.DataFrame, metric_type: Union[SceneFlowMetricType, SegmentationMetricType] + ) -> float: """Weighted average of metric m using the Count column. Args: @@ -422,34 +467,46 @@ def weighted_average(x: pd.DataFrame, metric_type: Union[SceneFlowMetricType, Se return averages for metric_type in SceneFlowMetricType: - avg: pd.Series[float] = grouped.apply(lambda x, m=metric_type: weighted_average(x, metric_type=m)) + avg: pd.Series[float] = grouped.apply( + lambda x, m=metric_type: weighted_average(x, metric_type=m) + ) segments: List[Tuple[str, str, str]] = avg.index.to_list() for segment in segments: if segment[:2] == NO_FMT_INDICES: continue metric_type_str = ( - metric_type.title().replace("_", " ") if metric_type != SceneFlowMetricType.EPE else metric_type + metric_type.title().replace("_", " ") + if metric_type != SceneFlowMetricType.EPE + else metric_type ) name = metric_type_str + "/" + "/".join([str(i) for i in segment]) output[name] = avg.loc[segment] grouped = frame.groupby(["Class", "Motion"]) for metric_type in SceneFlowMetricType: - avg_nodist: pd.Series[float] = grouped.apply(lambda x, m=metric_type: weighted_average(x, metric_type=m)) + avg_nodist: pd.Series[float] = grouped.apply( + lambda x, m=metric_type: weighted_average(x, metric_type=m) + ) segments_nodist: List[Tuple[str, str, str]] = avg_nodist.index.to_list() for segment in segments_nodist: if segment[:2] == NO_FMT_INDICES: continue metric_type_str = ( - metric_type.title().replace("_", " ") if metric_type != SceneFlowMetricType.EPE else metric_type + metric_type.title().replace("_", " ") + if metric_type != SceneFlowMetricType.EPE + else metric_type ) name = metric_type_str + "/" + "/".join([str(i) for i in segment]) output[name] = avg_nodist.loc[segment] - output["Dynamic IoU"] = frame.TP.sum() / (frame.TP.sum() + frame.FP.sum() + frame.FN.sum()) + output["Dynamic IoU"] = frame.TP.sum() / ( + frame.TP.sum() + frame.FP.sum() + frame.FN.sum() + ) output["EPE 3-Way Average"] = ( - output["EPE/Foreground/Dynamic"] + output["EPE/Foreground/Static"] + output["EPE/Background/Static"] + output["EPE/Foreground/Dynamic"] + + output["EPE/Foreground/Static"] + + output["EPE/Background/Static"] ) / 3 return output diff --git a/src/av2/evaluation/scene_flow/example_submission.py b/src/av2/evaluation/scene_flow/example_submission.py index a2a7134d..17c3ca52 100644 --- a/src/av2/evaluation/scene_flow/example_submission.py +++ b/src/av2/evaluation/scene_flow/example_submission.py @@ -4,13 +4,20 @@ import click import numpy as np -from av2.evaluation.scene_flow.utils import get_eval_point_mask, get_eval_subset, write_output_file -from av2.torch.data_loaders.scene_flow import SceneFlowDataloader from kornia.geometry.linalg import transform_points from rich.progress import track +from av2.evaluation.scene_flow.utils import ( + get_eval_point_mask, + get_eval_subset, + write_output_file, +) +from av2.torch.data_loaders.scene_flow import SceneFlowDataloader + -def example_submission(output_dir: str, mask_file: str, data_dir: str, name: str) -> None: +def example_submission( + output_dir: str, mask_file: str, data_dir: str, name: str +) -> None: """Output example submission files for the leaderboard. Predicts the ego motion for every point. Args: @@ -47,7 +54,9 @@ def example_submission(output_dir: str, mask_file: str, data_dir: str, name: str help="the data should be located in //sensor/", default="av2", ) -def _example_submission_entry(output_dir: str, mask_file: str, data_dir: str, name: str) -> None: +def _example_submission_entry( + output_dir: str, mask_file: str, data_dir: str, name: str +) -> None: """Entry point for example_submission.""" example_submission(output_dir, mask_file, data_dir, name) diff --git a/src/av2/evaluation/scene_flow/make_annotation_files.py b/src/av2/evaluation/scene_flow/make_annotation_files.py index 7a735ea6..85f2bd37 100644 --- a/src/av2/evaluation/scene_flow/make_annotation_files.py +++ b/src/av2/evaluation/scene_flow/make_annotation_files.py @@ -6,10 +6,11 @@ import click import numpy as np import pandas as pd +from rich.progress import track + from av2.evaluation.scene_flow.utils import get_eval_point_mask, get_eval_subset from av2.torch.data_loaders.scene_flow import SceneFlowDataloader from av2.utils.typing import NDArrayBool, NDArrayFloat, NDArrayInt -from rich.progress import track CLOSE_DISTANCE_THRESHOLD: Final = 35.0 @@ -54,7 +55,9 @@ def write_annotation( output.to_feather(output_file) -def make_annotation_files(output_dir: str, mask_file: str, data_dir: str, name: str, split: str) -> None: +def make_annotation_files( + output_dir: str, mask_file: str, data_dir: str, name: str, split: str +) -> None: """Create annotation files for running the evaluation. Args: @@ -86,7 +89,9 @@ def make_annotation_files(output_dir: str, mask_file: str, data_dir: str, name: is_dynamic = flow_labels.is_dynamic[mask].numpy().astype(bool) pc = sweep_0.lidar.as_tensor()[mask, :3].numpy() - is_close = np.logical_and.reduce(np.abs(pc[:, :2]) <= CLOSE_DISTANCE_THRESHOLD, axis=1).astype(bool) + is_close = np.logical_and.reduce( + np.abs(pc[:, :2]) <= CLOSE_DISTANCE_THRESHOLD, axis=1 + ).astype(bool) write_annotation( category_indices, @@ -115,7 +120,9 @@ def make_annotation_files(output_dir: str, mask_file: str, data_dir: str, name: default="val", type=click.Choice(["test", "val"]), ) -def _make_annotation_files_entry(output_dir: str, mask_file: str, data_dir: str, name: str, split: str) -> None: +def _make_annotation_files_entry( + output_dir: str, mask_file: str, data_dir: str, name: str, split: str +) -> None: """Entry point for make_annotation_files.""" make_annotation_files(output_dir, mask_file, data_dir, name, split) diff --git a/src/av2/evaluation/scene_flow/make_mask_files.py b/src/av2/evaluation/scene_flow/make_mask_files.py index a1dd3c6d..de6a096b 100644 --- a/src/av2/evaluation/scene_flow/make_mask_files.py +++ b/src/av2/evaluation/scene_flow/make_mask_files.py @@ -5,11 +5,12 @@ import click import pandas as pd +from kornia.geometry.liegroup import Se3 +from rich.progress import track + from av2.evaluation.scene_flow.utils import compute_eval_point_mask, get_eval_subset from av2.torch.data_loaders.scene_flow import SceneFlowDataloader from av2.torch.structures.sweep import Sweep -from kornia.geometry.liegroup import Se3 -from rich.progress import track def get_mask( @@ -68,7 +69,9 @@ def make_mask_files(output_file: str, data_dir: str, name: str, split: str) -> N default="val", type=click.Choice(["test", "val"]), ) -def _make_mask_files_entry(output_file: str, data_dir: str, name: str, split: str) -> None: +def _make_mask_files_entry( + output_file: str, data_dir: str, name: str, split: str +) -> None: """Entry point for make_mask_files.""" make_mask_files(output_file, data_dir, name, split) diff --git a/src/av2/evaluation/scene_flow/make_submission_archive.py b/src/av2/evaluation/scene_flow/make_submission_archive.py index 175f336b..38e8fcd5 100644 --- a/src/av2/evaluation/scene_flow/make_submission_archive.py +++ b/src/av2/evaluation/scene_flow/make_submission_archive.py @@ -23,11 +23,15 @@ def validate(submission_dir: Path, mask_file: Path) -> None: ValueError: If any supplied file is malformed """ with ZipFile(mask_file, "r") as masks: - mask_files = [f.filename for f in masks.filelist if f.filename.endswith(".feather")] + mask_files = [ + f.filename for f in masks.filelist if f.filename.endswith(".feather") + ] for filename in track(mask_files, description="Validating..."): input_file = submission_dir / filename if not input_file.exists(): - raise FileNotFoundError(f"{input_file} not found in submission directory") + raise FileNotFoundError( + f"{input_file} not found in submission directory" + ) pred = pd.read_feather(input_file) expected_num_points = pd.read_feather(masks.open(filename)).sum().item() @@ -36,16 +40,22 @@ def validate(submission_dir: Path, mask_file: Path) -> None: raise ValueError(f"{input_file} does not contain {c}") if c == "is_dynamic": if pred[c].dtype != bool: - raise ValueError(f"{input_file} column {c} should be bool but is {pred[c].dtype}") + raise ValueError( + f"{input_file} column {c} should be bool but is {pred[c].dtype}" + ) else: if pred[c].dtype != np.float16: - raise ValueError(f"{input_file} column {c} should be float16 but is {pred[c].dtype}") + raise ValueError( + f"{input_file} column {c} should be float16 but is {pred[c].dtype}" + ) if len(pred.columns) > 4: raise ValueError(f"{input_file} contains extra columns") if len(pred) != expected_num_points: - raise ValueError(f"{input_file} has {len(pred)} rows but it should have {expected_num_points}") + raise ValueError( + f"{input_file} has {len(pred)} rows but it should have {expected_num_points}" + ) def zip(submission_dir: Path, mask_file: Path, output_file: Path) -> None: @@ -57,14 +67,18 @@ def zip(submission_dir: Path, mask_file: Path, output_file: Path) -> None: output_file: File to store the zip archive in. """ with ZipFile(mask_file, "r") as masks: - mask_files = [f.filename for f in masks.filelist if f.filename.endswith(".feather")] + mask_files = [ + f.filename for f in masks.filelist if f.filename.endswith(".feather") + ] with ZipFile(output_file, "w") as myzip: for filename in track(mask_files, description="Zipping..."): input_file = submission_dir / filename myzip.write(input_file, arcname=filename) -def make_submission_archive(submission_dir: str, mask_file: str, output_filename: str) -> bool: +def make_submission_archive( + submission_dir: str, mask_file: str, output_filename: str +) -> bool: """Package prediction files into a zip archive for submission. Args: @@ -89,8 +103,15 @@ def make_submission_archive(submission_dir: str, mask_file: str, output_filename @click.command() @click.argument("submission_dir", type=str) @click.argument("mask_file", type=str) -@click.option("--output_filename", type=str, help="name of the output archive file", default="submission.zip") -def _make_submission_archive_entry(submission_dir: str, mask_file: str, output_filename: str) -> bool: +@click.option( + "--output_filename", + type=str, + help="name of the output archive file", + default="submission.zip", +) +def _make_submission_archive_entry( + submission_dir: str, mask_file: str, output_filename: str +) -> bool: return make_submission_archive(submission_dir, mask_file, output_filename) diff --git a/src/av2/evaluation/scene_flow/utils.py b/src/av2/evaluation/scene_flow/utils.py index f57ff78d..27159144 100644 --- a/src/av2/evaluation/scene_flow/utils.py +++ b/src/av2/evaluation/scene_flow/utils.py @@ -7,12 +7,13 @@ import numpy as np import pandas as pd import torch +from kornia.geometry.liegroup import Se3 +from torch import BoolTensor + from av2.torch.data_loaders.scene_flow import SceneFlowDataloader from av2.torch.structures.flow import Flow from av2.torch.structures.sweep import Sweep from av2.utils.typing import NDArrayBool, NDArrayFloat -from kornia.geometry.liegroup import Se3 -from torch import BoolTensor _EVAL_ROOT: Final = Path(__file__).resolve().parent @@ -34,12 +35,18 @@ def get_eval_point_mask(sweep_uuid: Tuple[str, int], mask_file: Path) -> BoolTen """ with ZipFile(mask_file) as masks: log_id, timestamp_ns = sweep_uuid - mask = pd.read_feather(masks.open(f"{log_id}/{timestamp_ns}.feather")).to_numpy().astype(bool) + mask = ( + pd.read_feather(masks.open(f"{log_id}/{timestamp_ns}.feather")) + .to_numpy() + .astype(bool) + ) return BoolTensor(torch.from_numpy(mask).squeeze()) -def compute_eval_point_mask(datum: Tuple[Sweep, Sweep, Se3, Optional[Flow]]) -> BoolTensor: +def compute_eval_point_mask( + datum: Tuple[Sweep, Sweep, Se3, Optional[Flow]] +) -> BoolTensor: """Compute for a given sweep, a boolean mask indicating which points are evaluated on. Note: This should NOT BE USED FOR CREATING SUBMISSIONS use get_eval_point_mask to ensure consistency. @@ -54,7 +61,9 @@ def compute_eval_point_mask(datum: Tuple[Sweep, Sweep, Se3, Optional[Flow]]) -> ValueError: if datum does not have ground annotations. """ pcl = datum[0].lidar.as_tensor()[:, :3] - is_close = torch.logical_and((pcl[:, 0].abs() <= 50), (pcl[:, 1].abs() <= 50)).bool() + is_close = torch.logical_and( + (pcl[:, 0].abs() <= 50), (pcl[:, 1].abs() <= 50) + ).bool() if datum[0].is_ground is None: raise ValueError("Must have ground annotations loaded to determine eval mask") @@ -63,7 +72,10 @@ def compute_eval_point_mask(datum: Tuple[Sweep, Sweep, Se3, Optional[Flow]]) -> def write_output_file( - flow: NDArrayFloat, is_dynamic: NDArrayBool, sweep_uuid: Tuple[str, int], output_dir: Path + flow: NDArrayFloat, + is_dynamic: NDArrayBool, + sweep_uuid: Tuple[str, int], + output_dir: Path, ) -> None: """Write an output predictions file in the correct format for submission. @@ -79,6 +91,11 @@ def write_output_file( fy_m = flow[:, 1].astype(np.float16) fz_m = flow[:, 2].astype(np.float16) output = pd.DataFrame( - {"flow_tx_m": fx_m, "flow_ty_m": fy_m, "flow_tz_m": fz_m, "is_dynamic": is_dynamic.astype(bool)} + { + "flow_tx_m": fx_m, + "flow_ty_m": fy_m, + "flow_tz_m": fz_m, + "is_dynamic": is_dynamic.astype(bool), + } ) output.to_feather(output_log_dir / f"{sweep_uuid[1]}.feather") diff --git a/src/av2/evaluation/tracking/eval.py b/src/av2/evaluation/tracking/eval.py index 2438a561..56d56b9f 100644 --- a/src/av2/evaluation/tracking/eval.py +++ b/src/av2/evaluation/tracking/eval.py @@ -19,14 +19,19 @@ import click import numpy as np import trackeval -from av2.evaluation.detection.utils import compute_objects_in_roi_mask, load_mapped_avm_and_egoposes -from av2.evaluation.tracking import constants, utils -from av2.evaluation.tracking.constants import SUBMETRIC_TO_METRIC_CLASS_NAME -from av2.utils.typing import NDArrayFloat, NDArrayInt from scipy.optimize import linear_sum_assignment from scipy.spatial.transform import Rotation from tqdm import tqdm from trackeval.datasets._base_dataset import _BaseDataset + +from av2.evaluation.detection.utils import ( + compute_objects_in_roi_mask, + load_mapped_avm_and_egoposes, +) +from av2.evaluation.tracking import constants, utils +from av2.evaluation.tracking.constants import SUBMETRIC_TO_METRIC_CLASS_NAME +from av2.utils.typing import NDArrayFloat, NDArrayInt + from ..typing import Sequences @@ -66,7 +71,9 @@ def get_default_dataset_config() -> Dict[str, Any]: } return default_config - def _load_raw_file(self, tracker: str, seq_id: Union[str, int], is_gt: bool) -> Dict[str, Any]: + def _load_raw_file( + self, tracker: str, seq_id: Union[str, int], is_gt: bool + ) -> Dict[str, Any]: """Get raw track data, from either trackers or ground truth.""" tracks = (self.gt_tracks if is_gt else self.predicted_tracks)[tracker][seq_id] source = "gt" if is_gt else "tracker" @@ -77,7 +84,10 @@ def _load_raw_file(self, tracker: str, seq_id: Union[str, int], is_gt: bool) -> raw_data = { f"{source}_ids": [frame["track_id"] for frame in tracks], f"{source}_classes": [frame["label"] for frame in tracks], - f"{source}_dets": [np.concatenate((frame["translation_m"], frame["size"]), axis=-1) for frame in tracks], + f"{source}_dets": [ + np.concatenate((frame["translation_m"], frame["size"]), axis=-1) + for frame in tracks + ], "num_timesteps": len(tracks), "seq": seq_id, } @@ -85,7 +95,9 @@ def _load_raw_file(self, tracker: str, seq_id: Union[str, int], is_gt: bool) -> raw_data[f"{source}_confidences"] = [frame["score"] for frame in tracks] return raw_data - def get_preprocessed_seq_data(self, raw_data: Dict[str, Any], cls: str) -> Dict[str, Any]: + def get_preprocessed_seq_data( + self, raw_data: Dict[str, Any], cls: str + ) -> Dict[str, Any]: """Filter data to keep only one class and map id to 0 - n. Args: @@ -117,12 +129,18 @@ def get_preprocessed_seq_data(self, raw_data: Dict[str, Any], cls: str) -> Dict[ data["gt_dets"][t] = data["gt_dets"][t][gt_to_keep_mask, :] tracker_to_keep_mask = data["tracker_classes"][t] == cls_id - data["tracker_classes"][t] = data["tracker_classes"][t][tracker_to_keep_mask] + data["tracker_classes"][t] = data["tracker_classes"][t][ + tracker_to_keep_mask + ] data["tracker_ids"][t] = data["tracker_ids"][t][tracker_to_keep_mask] data["tracker_dets"][t] = data["tracker_dets"][t][tracker_to_keep_mask, :] - data["tracker_confidences"][t] = data["tracker_confidences"][t][tracker_to_keep_mask] + data["tracker_confidences"][t] = data["tracker_confidences"][t][ + tracker_to_keep_mask + ] - data["similarity_scores"][t] = data["similarity_scores"][t][:, tracker_to_keep_mask][gt_to_keep_mask] + data["similarity_scores"][t] = data["similarity_scores"][t][ + :, tracker_to_keep_mask + ][gt_to_keep_mask] # Map ids to 0 - n. unique_gt_ids = set(chain.from_iterable(data["gt_ids"])) @@ -141,13 +159,19 @@ def get_preprocessed_seq_data(self, raw_data: Dict[str, Any], cls: str) -> Dict[ def _map_ids(self, ids: List[Any], unique_ids: Iterable[Any]) -> List[NDArrayInt]: id_map = {id: i for i, id in enumerate(unique_ids)} - return [np.array([id_map[id] for id in id_array], dtype=int) for id_array in ids] + return [ + np.array([id_map[id] for id in id_array], dtype=int) for id_array in ids + ] - def _calculate_similarities(self, gt_dets_t: NDArrayFloat, tracker_dets_t: NDArrayFloat) -> NDArrayFloat: + def _calculate_similarities( + self, gt_dets_t: NDArrayFloat, tracker_dets_t: NDArrayFloat + ) -> NDArrayFloat: """Euclidean distance of the x, y translation coordinates.""" gt_xy = gt_dets_t[:, :2] tracker_xy = tracker_dets_t[:, :2] - sim = self._calculate_euclidean_similarity(gt_xy, tracker_xy, zero_distance=self.zero_distance) + sim = self._calculate_euclidean_similarity( + gt_xy, tracker_xy, zero_distance=self.zero_distance + ) return cast(NDArrayFloat, sim) @@ -193,17 +217,25 @@ def evaluate_tracking( Returns: Dictionary of metric values. """ - labels_id_ts = set((frame["seq_id"], frame["timestamp_ns"]) for frame in utils.ungroup_frames(labels)) + labels_id_ts = set( + (frame["seq_id"], frame["timestamp_ns"]) + for frame in utils.ungroup_frames(labels) + ) predictions_id_ts = set( - (frame["seq_id"], frame["timestamp_ns"]) for frame in utils.ungroup_frames(track_predictions) + (frame["seq_id"], frame["timestamp_ns"]) + for frame in utils.ungroup_frames(track_predictions) ) - assert labels_id_ts == predictions_id_ts, "sequences ids and timestamp_ns in labels and predictions don't match" + assert ( + labels_id_ts == predictions_id_ts + ), "sequences ids and timestamp_ns in labels and predictions don't match" metrics_config = { "METRICS": ["HOTA", "CLEAR"], "THRESHOLD": iou_threshold, } metric_names = cast(List[str], metrics_config["METRICS"]) - metrics_list = [getattr(trackeval.metrics, metric)(metrics_config) for metric in metric_names] + metrics_list = [ + getattr(trackeval.metrics, metric)(metrics_config) for metric in metric_names + ] dataset_config = { **TrackEvalDataset.get_default_dataset_config(), "GT_TRACKS": {tracker_name: labels}, @@ -299,10 +331,18 @@ def _tune_score_thresholds( ) metric_results = [] - for threshold_i in tqdm(range(num_thresholds), "calculating optimal track score thresholds"): - score_threshold_by_class = {n: score_thresholds_by_class[n][threshold_i] for n in classes} - filtered_predictions = utils.filter_by_class_thresholds(track_predictions, score_threshold_by_class) - with contextlib.redirect_stdout(None): # silence print statements from TrackEval + for threshold_i in tqdm( + range(num_thresholds), "calculating optimal track score thresholds" + ): + score_threshold_by_class = { + n: score_thresholds_by_class[n][threshold_i] for n in classes + } + filtered_predictions = utils.filter_by_class_thresholds( + track_predictions, score_threshold_by_class + ) + with contextlib.redirect_stdout( + None + ): # silence print statements from TrackEval result_for_threshold, _ = evaluator.evaluate( [ TrackEvalDataset( @@ -314,18 +354,26 @@ def _tune_score_thresholds( ], metrics_list, ) - metric_results.append(result_for_threshold["TrackEvalDataset"]["tracker"]["COMBINED_SEQ"]) + metric_results.append( + result_for_threshold["TrackEvalDataset"]["tracker"]["COMBINED_SEQ"] + ) optimal_score_threshold_by_class = {} optimal_metric_values_by_class = {} mean_metric_values_by_class = {} for name in classes: - metric_values = [r[name][metric_class][objective_metric] for r in metric_results] - metric_values = [np.mean(v) if isinstance(v, np.ndarray) else v for v in metric_values] + metric_values = [ + r[name][metric_class][objective_metric] for r in metric_results + ] + metric_values = [ + np.mean(v) if isinstance(v, np.ndarray) else v for v in metric_values + ] optimal_threshold = score_thresholds_by_class[name][np.argmax(metric_values)] optimal_score_threshold_by_class[name] = optimal_threshold optimal_metric_values_by_class[name] = max(0, np.max(metric_values)) - mean_metric_values_by_class[name] = np.nanmean(np.array(metric_values).clip(min=0)) + mean_metric_values_by_class[name] = np.nanmean( + np.array(metric_values).clip(min=0) + ) return ( optimal_score_threshold_by_class, optimal_metric_values_by_class, @@ -335,7 +383,10 @@ def _tune_score_thresholds( def _filter_by_class(detections: Any, name: str) -> Any: return utils.group_frames( - [utils.index_array_values(f, f["name"] == name) for f in utils.ungroup_frames(detections)] + [ + utils.index_array_values(f, f["name"] == name) + for f in utils.ungroup_frames(detections) + ] ) @@ -350,7 +401,9 @@ def _calculate_score_thresholds( recall_thresholds = np.linspace(min_recall, 1, num_thresholds).round(12)[::-1] if len(scores) == 0: return np.zeros_like(recall_thresholds) - score_thresholds = _recall_to_scores(scores, recall_threshold=recall_thresholds, n_gt=n_gt) + score_thresholds = _recall_to_scores( + scores, recall_threshold=recall_thresholds, n_gt=n_gt + ) score_thresholds = np.nan_to_num(score_thresholds, nan=0) return score_thresholds @@ -365,7 +418,9 @@ def _calculate_matched_scores( num_tp = 0 for seq_id in labels: for label_frame, prediction_frame in zip(labels[seq_id], predictions[seq_id]): - sim = sim_func(label_frame["translation_m"], prediction_frame["translation_m"]) + sim = sim_func( + label_frame["translation_m"], prediction_frame["translation_m"] + ) match_rows, match_cols = linear_sum_assignment(-sim) scores.append(prediction_frame["score"][match_cols]) n_gt += len(label_frame["translation_m"]) @@ -375,7 +430,9 @@ def _calculate_matched_scores( return scores_array, n_gt -def _recall_to_scores(scores: NDArrayFloat, recall_threshold: NDArrayFloat, n_gt: int) -> NDArrayFloat: +def _recall_to_scores( + scores: NDArrayFloat, recall_threshold: NDArrayFloat, n_gt: int +) -> NDArrayFloat: # Sort scores. scores.sort() scores = scores[::-1] @@ -392,10 +449,14 @@ def _recall_to_scores(scores: NDArrayFloat, recall_threshold: NDArrayFloat, n_gt return score_thresholds -def _xy_center_similarity(centers1: NDArrayFloat, centers2: NDArrayFloat, zero_distance: float) -> NDArrayFloat: +def _xy_center_similarity( + centers1: NDArrayFloat, centers2: NDArrayFloat, zero_distance: float +) -> NDArrayFloat: if centers1.size == 0 or centers2.size == 0: return np.zeros((len(centers1), len(centers2))) - xy_dist = np.linalg.norm(centers1[:, np.newaxis, :2] - centers2[np.newaxis, :, :2], axis=2) + xy_dist = np.linalg.norm( + centers1[:, np.newaxis, :2] - centers2[np.newaxis, :, :2], axis=2 + ) sim = np.maximum(0, 1 - xy_dist / zero_distance) return cast(NDArrayFloat, sim) @@ -416,7 +477,8 @@ def filter_max_dist(tracks: Any, max_range_m: int) -> Any: utils.index_array_values( frame, np.linalg.norm( - frame["translation_m"][:, :2] - np.array(frame["ego_translation_m"])[:2], + frame["translation_m"][:, :2] + - np.array(frame["ego_translation_m"])[:2], axis=1, ) <= max_range_m, @@ -453,7 +515,9 @@ def filter_drivable_area(tracks: Sequences, dataset_dir: Optional[str]) -> Seque return tracks log_ids = list(tracks.keys()) - log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes(log_ids, Path(dataset_dir)) + log_id_to_avm, log_id_to_timestamped_poses = load_mapped_avm_and_egoposes( + log_ids, Path(dataset_dir) + ) for log_id in log_ids: avm = log_id_to_avm[log_id] @@ -541,7 +605,9 @@ def evaluate( num_thresholds=10, match_distance_m=2, ) - filtered_track_predictions = utils.filter_by_class_thresholds(track_predictions, score_thresholds) + filtered_track_predictions = utils.filter_by_class_thresholds( + track_predictions, score_thresholds + ) res = evaluate_tracking( labels, filtered_track_predictions, @@ -576,7 +642,9 @@ def runner( track_predictions = pickle.load(open(predictions, "rb")) labels = pickle.load(open(ground_truth, "rb")) - _, _, mean_metric_values = evaluate(track_predictions, labels, objective_metric, max_range_m, dataset_dir, out) + _, _, mean_metric_values = evaluate( + track_predictions, labels, objective_metric, max_range_m, dataset_dir, out + ) pprint(mean_metric_values) diff --git a/src/av2/evaluation/tracking/utils.py b/src/av2/evaluation/tracking/utils.py index 18086ae2..49605583 100644 --- a/src/av2/evaluation/tracking/utils.py +++ b/src/av2/evaluation/tracking/utils.py @@ -11,8 +11,10 @@ from typing import Any, Dict, Iterable, List, Union, cast import numpy as np + from av2.utils.typing import NDArrayInt -from ..typing import Sequences, Frame, Frames + +from ..typing import Frame, Frames, Sequences def save(obj: Any, path: str) -> None: # noqa @@ -44,7 +46,9 @@ def load(path: str) -> Any: # noqa return pickle.load(f) -def annotate_frame_metadata(prediction_frames: Frames, label_frames: Frames, metadata_keys: List[str]) -> None: +def annotate_frame_metadata( + prediction_frames: Frames, label_frames: Frames, metadata_keys: List[str] +) -> None: """Copy annotations with provided keys from label to prediction frames. Args: @@ -96,7 +100,9 @@ def index_array_values(array_dict: Frame, index: Union[int, NDArrayInt]) -> Fram Returns: Dictionary of numpy arrays, each indexed by the provided index """ - return {k: v[index] if isinstance(v, np.ndarray) else v for k, v in array_dict.items()} + return { + k: v[index] if isinstance(v, np.ndarray) else v for k, v in array_dict.items() + } def array_dict_iterator(array_dict: Frame, length: int) -> Iterable[Frame]: @@ -140,7 +146,9 @@ def concatenate_array_values(array_dicts: Frames) -> Frame: return concatenated -def filter_by_class_thresholds(frames_by_seq_id: Sequences, thresholds_by_class: Dict[str, float]) -> Sequences: +def filter_by_class_thresholds( + frames_by_seq_id: Sequences, thresholds_by_class: Dict[str, float] +) -> Sequences: """Filter detections, keeping only detections with score higher than the provided threshold for that class. If a class threshold is not provided, all detections in that class is filtered. diff --git a/src/av2/evaluation/typing.py b/src/av2/evaluation/typing.py index f58a4400..9dfa574c 100644 --- a/src/av2/evaluation/typing.py +++ b/src/av2/evaluation/typing.py @@ -2,7 +2,6 @@ from typing import Any, Dict, List - Frame = Dict[str, Any] Frames = List[Frame] Sequences = Dict[str, Frames] diff --git a/src/av2/geometry/camera/pinhole_camera.py b/src/av2/geometry/camera/pinhole_camera.py index f5fcc846..dd18f058 100644 --- a/src/av2/geometry/camera/pinhole_camera.py +++ b/src/av2/geometry/camera/pinhole_camera.py @@ -108,7 +108,9 @@ def from_feather(cls, log_dir: Path, cam_name: str) -> PinholeCamera: cam_name=cam_name, ) - def cull_to_view_frustum(self, uv: NDArrayFloat, points_cam: NDArrayFloat) -> NDArrayBool: + def cull_to_view_frustum( + self, uv: NDArrayFloat, points_cam: NDArrayFloat + ) -> NDArrayBool: """Cull 3d points to camera view frustum. Given a set of coordinates in the image plane and corresponding points @@ -131,7 +133,9 @@ def cull_to_view_frustum(self, uv: NDArrayFloat, points_cam: NDArrayFloat) -> ND is_valid_x = np.logical_and(0 <= uv[:, 0], uv[:, 0] < self.width_px - 1) is_valid_y = np.logical_and(0 <= uv[:, 1], uv[:, 1] < self.height_px - 1) is_valid_z = points_cam[:, 2] > 0 - is_valid_points: NDArrayBool = np.logical_and.reduce([is_valid_x, is_valid_y, is_valid_z]) + is_valid_points: NDArrayBool = np.logical_and.reduce( + [is_valid_x, is_valid_y, is_valid_z] + ) return is_valid_points def project_ego_to_img( @@ -249,8 +253,12 @@ def project_ego_to_img_motion_compensated( if city_SE3_ego_lidar_t is None: raise ValueError("city_SE3_ego_lidar_t cannot be `None`!") - ego_cam_t_SE3_ego_lidar_t = city_SE3_ego_cam_t.inverse().compose(city_SE3_ego_lidar_t) - points_cam_time = ego_cam_t_SE3_ego_lidar_t.transform_point_cloud(points_lidar_time) + ego_cam_t_SE3_ego_lidar_t = city_SE3_ego_cam_t.inverse().compose( + city_SE3_ego_lidar_t + ) + points_cam_time = ego_cam_t_SE3_ego_lidar_t.transform_point_cloud( + points_lidar_time + ) return self.project_ego_to_img(points_cam_time) @cached_property @@ -332,7 +340,9 @@ def frustum_planes(self, near_clip_dist: float = 0.5) -> NDArrayFloat: bottom_plane = self.bottom_clipping_plane near_plane = self.near_clipping_plane(near_clip_dist) - planes: NDArrayFloat = np.stack([left_plane, right_plane, near_plane, bottom_plane, top_plane]) + planes: NDArrayFloat = np.stack( + [left_plane, right_plane, near_plane, bottom_plane, top_plane] + ) return planes @cached_property @@ -363,7 +373,9 @@ def fov_theta_rad(self) -> float: fov_theta_rad = 2 * np.arctan(0.5 * self.width_px / self.intrinsics.fx_px) return float(fov_theta_rad) - def compute_pixel_ray_directions(self, uv: Union[NDArrayFloat, NDArrayInt]) -> NDArrayFloat: + def compute_pixel_ray_directions( + self, uv: Union[NDArrayFloat, NDArrayInt] + ) -> NDArrayFloat: """Given (u,v) coordinates and intrinsics, generate pixel rays in the camera coordinate frame. Assume +z points out of the camera, +y is downwards, and +x is across the imager. @@ -382,7 +394,9 @@ def compute_pixel_ray_directions(self, uv: Union[NDArrayFloat, NDArrayInt]) -> N img_h, img_w = self.height_px, self.width_px if not np.isclose(fx, fy, atol=1e-3): - raise ValueError(f"Focal lengths in the x and y directions must match: {fx} != {fy}") + raise ValueError( + f"Focal lengths in the x and y directions must match: {fx} != {fy}" + ) if uv.shape[1] != 2: raise ValueError("Input (u,v) coordinates must be (N,2) in shape.") @@ -424,10 +438,14 @@ def scale(self, scale: float) -> PinholeCamera: round(self.intrinsics.width_px * scale), round(self.intrinsics.height_px * scale), ) - return PinholeCamera(ego_SE3_cam=self.ego_SE3_cam, intrinsics=intrinsics, cam_name=self.cam_name) + return PinholeCamera( + ego_SE3_cam=self.ego_SE3_cam, intrinsics=intrinsics, cam_name=self.cam_name + ) -def remove_nan_values(uv: NDArrayFloat, points_cam: NDArrayFloat) -> Tuple[NDArrayFloat, NDArrayFloat]: +def remove_nan_values( + uv: NDArrayFloat, points_cam: NDArrayFloat +) -> Tuple[NDArrayFloat, NDArrayFloat]: """Remove NaN values from camera coordinates and image plane coordinates (accepts corrupt array). Args: diff --git a/src/av2/geometry/geometry.py b/src/av2/geometry/geometry.py index 3e782cd9..c735a249 100644 --- a/src/av2/geometry/geometry.py +++ b/src/av2/geometry/geometry.py @@ -229,11 +229,15 @@ def crop_points( # Ensure that the logical operations will broadcast. if n_dim != lb_dim or n_dim != ub_dim: - raise ValueError(f"Dimensions n_dim {n_dim} must match both lb_dim {lb_dim} and ub_dim {ub_dim}") + raise ValueError( + f"Dimensions n_dim {n_dim} must match both lb_dim {lb_dim} and ub_dim {ub_dim}" + ) # Ensure that the lower bound less than or equal to the upper bound for each dimension. if not all(lb < ub for lb, ub in zip(lower_bound_inclusive, upper_bound_exclusive)): - raise ValueError("Lower bound must be less than or equal to upper bound for each dimension") + raise ValueError( + "Lower bound must be less than or equal to upper bound for each dimension" + ) # Lower bound mask. lb_mask = np.greater_equal(points, lower_bound_inclusive) @@ -246,7 +250,9 @@ def crop_points( return points[is_valid_points], is_valid_points -def compute_interior_points_mask(points_xyz: NDArrayFloat, cuboid_vertices: NDArrayFloat) -> NDArrayBool: +def compute_interior_points_mask( + points_xyz: NDArrayFloat, cuboid_vertices: NDArrayFloat +) -> NDArrayBool: r"""Compute the interior points mask for the cuboid. Reference: https://math.stackexchange.com/questions/1472049/check-if-a-point-is-inside-a-rectangular-shaped-area-3d @@ -271,7 +277,9 @@ def compute_interior_points_mask(points_xyz: NDArrayFloat, cuboid_vertices: NDAr (N,) An array of boolean flags indicating whether the points are interior to the cuboid. """ # Get three corners of the cuboid vertices. - vertices: NDArrayFloat = np.stack((cuboid_vertices[6], cuboid_vertices[3], cuboid_vertices[1])) # (3,3) + vertices: NDArrayFloat = np.stack( + (cuboid_vertices[6], cuboid_vertices[3], cuboid_vertices[1]) + ) # (3,3) # Choose reference vertex. # vertices and choice of ref_vertex are coupled. @@ -289,7 +297,11 @@ def compute_interior_points_mask(points_xyz: NDArrayFloat, cuboid_vertices: NDAr # Check 6 conditions (2 for each of the 3 orthogonal directions). # Refer to the linked reference for additional information. - constraint_a = np.logical_and(sim_uvw_ref <= sim_uvw_points, sim_uvw_points <= sim_uvw_vertices) - constraint_b = np.logical_and(sim_uvw_ref >= sim_uvw_points, sim_uvw_points >= sim_uvw_vertices) + constraint_a = np.logical_and( + sim_uvw_ref <= sim_uvw_points, sim_uvw_points <= sim_uvw_vertices + ) + constraint_b = np.logical_and( + sim_uvw_ref >= sim_uvw_points, sim_uvw_points >= sim_uvw_vertices + ) is_interior: NDArrayBool = np.logical_or(constraint_a, constraint_b).all(axis=1) return is_interior diff --git a/src/av2/geometry/infinity_norm_utils.py b/src/av2/geometry/infinity_norm_utils.py index a6f2ddd9..2e477880 100644 --- a/src/av2/geometry/infinity_norm_utils.py +++ b/src/av2/geometry/infinity_norm_utils.py @@ -7,7 +7,9 @@ from av2.utils.typing import NDArrayFloat -def has_pts_in_infinity_norm_radius(points: NDArrayFloat, window_center: NDArrayFloat, window_sz: float) -> bool: +def has_pts_in_infinity_norm_radius( + points: NDArrayFloat, window_center: NDArrayFloat, window_sz: float +) -> bool: """Check if a map entity has points within a search radius from a single query point. Note: Distance is measured by the infinity norm. @@ -24,7 +26,9 @@ def has_pts_in_infinity_norm_radius(points: NDArrayFloat, window_center: NDArray ValueError: If `points` is not in R^2 or R^3. """ if points.ndim != 2 or points.shape[1] not in (2, 3): - raise ValueError(f"Input points array must have shape (N,2) or (N,3) - received {points.shape}.") + raise ValueError( + f"Input points array must have shape (N,2) or (N,3) - received {points.shape}." + ) if points.shape[1] == 3: # take only x,y dimensions diff --git a/src/av2/geometry/interpolate.py b/src/av2/geometry/interpolate.py index 20c1d365..ec929d73 100644 --- a/src/av2/geometry/interpolate.py +++ b/src/av2/geometry/interpolate.py @@ -14,7 +14,9 @@ NUM_CENTERLINE_INTERP_PTS: Final[int] = 10 -def compute_lane_width(left_even_pts: NDArrayFloat, right_even_pts: NDArrayFloat) -> float: +def compute_lane_width( + left_even_pts: NDArrayFloat, right_even_pts: NDArrayFloat +) -> float: """Compute the width of a lane, given an explicit left and right boundary. Requires an equal number of waypoints on each boundary. For 3d polylines, this incorporates @@ -39,7 +41,9 @@ def compute_lane_width(left_even_pts: NDArrayFloat, right_even_pts: NDArrayFloat return lane_width -def compute_mid_pivot_arc(single_pt: NDArrayFloat, arc_pts: NDArrayFloat) -> Tuple[NDArrayFloat, float]: +def compute_mid_pivot_arc( + single_pt: NDArrayFloat, arc_pts: NDArrayFloat +) -> Tuple[NDArrayFloat, float]: """Compute an arc by pivoting around a single point. Given a line of points on one boundary, and a single point on the other side, @@ -90,7 +94,9 @@ def compute_midpoint_line( ValueError: If the left and right lane boundaries aren't a list of 2d or 3d waypoints. """ if left_ln_boundary.ndim != 2 or right_ln_boundary.ndim != 2: - raise ValueError("Left and right lane boundaries must consist of a sequence of 2d or 3d waypoints.") + raise ValueError( + "Left and right lane boundaries must consist of a sequence of 2d or 3d waypoints." + ) dim = left_ln_boundary.shape[1] if dim not in [2, 3]: @@ -100,11 +106,15 @@ def compute_midpoint_line( raise ValueError("Left ") if len(left_ln_boundary) == 1: - centerline_pts, lane_width = compute_mid_pivot_arc(single_pt=left_ln_boundary, arc_pts=right_ln_boundary) + centerline_pts, lane_width = compute_mid_pivot_arc( + single_pt=left_ln_boundary, arc_pts=right_ln_boundary + ) return centerline_pts[:, :2], lane_width if len(right_ln_boundary) == 1: - centerline_pts, lane_width = compute_mid_pivot_arc(single_pt=right_ln_boundary, arc_pts=left_ln_boundary) + centerline_pts, lane_width = compute_mid_pivot_arc( + single_pt=right_ln_boundary, arc_pts=left_ln_boundary + ) return centerline_pts[:, :2], lane_width # fall back to the typical case. @@ -176,7 +186,9 @@ def interp_arc(t: int, points: NDArrayFloat) -> NDArrayFloat: def linear_interpolation( - key_timestamps: Tuple[int, int], key_translations: Tuple[NDArrayFloat, NDArrayFloat], query_timestamp: int + key_timestamps: Tuple[int, int], + key_translations: Tuple[NDArrayFloat, NDArrayFloat], + query_timestamp: int, ) -> NDArrayFloat: """Given two 3d positions at specific timestamps, interpolate an intermediate position at a given timestamp. @@ -203,7 +215,9 @@ def linear_interpolation( return translation_interp -def interpolate_pose(key_timestamps: Tuple[int, int], key_poses: Tuple[SE3, SE3], query_timestamp: int) -> SE3: +def interpolate_pose( + key_timestamps: Tuple[int, int], key_poses: Tuple[SE3, SE3], query_timestamp: int +) -> SE3: """Given two SE(3) poses at specific timestamps, interpolate an intermediate pose at a given timestamp. Note: we use a straight line interpolation for the translation, while still using interpolate (aka "slerp") @@ -236,6 +250,10 @@ def interpolate_pose(key_timestamps: Tuple[int, int], key_poses: Tuple[SE3, SE3] R_interp = slerp(query_timestamp).as_matrix() key_translations = (key_poses[0].translation, key_poses[1].translation) - t_interp = linear_interpolation(key_timestamps, key_translations=key_translations, query_timestamp=query_timestamp) + t_interp = linear_interpolation( + key_timestamps, + key_translations=key_translations, + query_timestamp=query_timestamp, + ) pose_interp = SE3(rotation=R_interp, translation=t_interp) return pose_interp diff --git a/src/av2/geometry/iou.py b/src/av2/geometry/iou.py index ca20a526..aa5b0d92 100644 --- a/src/av2/geometry/iou.py +++ b/src/av2/geometry/iou.py @@ -7,7 +7,9 @@ from av2.utils.typing import NDArrayFloat -def iou_3d_axis_aligned(src_dims_m: NDArrayFloat, target_dims_m: NDArrayFloat) -> NDArrayFloat: +def iou_3d_axis_aligned( + src_dims_m: NDArrayFloat, target_dims_m: NDArrayFloat +) -> NDArrayFloat: """Compute 3d, axis-aligned (vertical axis alignment) intersection-over-union (IoU) between two sets of cuboids. Both objects are aligned to their +x axis and their centroids are placed at the origin diff --git a/src/av2/geometry/polyline_utils.py b/src/av2/geometry/polyline_utils.py index 7e212941..fce300fe 100644 --- a/src/av2/geometry/polyline_utils.py +++ b/src/av2/geometry/polyline_utils.py @@ -34,7 +34,9 @@ def get_polyline_length(polyline: NDArrayFloat) -> float: return float(np.linalg.norm(offsets, axis=1).sum()) -def interp_polyline_by_fixed_waypt_interval(polyline: NDArrayFloat, waypt_interval: float) -> Tuple[NDArrayFloat, int]: +def interp_polyline_by_fixed_waypt_interval( + polyline: NDArrayFloat, waypt_interval: float +) -> Tuple[NDArrayFloat, int]: """Resample waypoints of a polyline so that waypoints appear roughly at fixed intervals from the start. Args: @@ -61,7 +63,9 @@ def interp_polyline_by_fixed_waypt_interval(polyline: NDArrayFloat, waypt_interv return interp_polyline, num_waypts -def get_double_polylines(polyline: NDArrayFloat, width_scaling_factor: float) -> Tuple[NDArrayFloat, NDArrayFloat]: +def get_double_polylines( + polyline: NDArrayFloat, width_scaling_factor: float +) -> Tuple[NDArrayFloat, NDArrayFloat]: """Treat any polyline as a centerline, and extend a narrow strip on both sides. Dimension is preserved (2d->2d, and 3d->3d). @@ -75,20 +79,26 @@ def get_double_polylines(polyline: NDArrayFloat, width_scaling_factor: float) -> left: array of shape (N,2) or (N,3) representing left polyline. right: array of shape (N,2) or (N,3) representing right polyline. """ - double_line_polygon = centerline_to_polygon(centerline=polyline, width_scaling_factor=width_scaling_factor) + double_line_polygon = centerline_to_polygon( + centerline=polyline, width_scaling_factor=width_scaling_factor + ) num_pts = double_line_polygon.shape[0] # split index -- polygon from right boundary, left boundary, then close it w/ 0th vertex of right # we swap left and right since our polygon is generated about a boundary, not a centerline k = num_pts // 2 left = double_line_polygon[:k] - right = double_line_polygon[k:-1] # throw away the last point, since it is just a repeat + right = double_line_polygon[ + k:-1 + ] # throw away the last point, since it is just a repeat return left, right def swap_left_and_right( - condition: NDArrayBool, left_centerline: NDArrayFloat, right_centerline: NDArrayFloat + condition: NDArrayBool, + left_centerline: NDArrayFloat, + right_centerline: NDArrayFloat, ) -> Tuple[NDArrayFloat, NDArrayFloat]: """Swap points in left and right centerline according to condition. @@ -147,7 +157,9 @@ def centerline_to_polygon( x_disp = AVG_LANE_WIDTH_M * width_scaling_factor / 2.0 * np.cos(thetas) y_disp = AVG_LANE_WIDTH_M * width_scaling_factor / 2.0 * np.sin(thetas) - displacement: NDArrayFloat = np.hstack([x_disp[:, np.newaxis], y_disp[:, np.newaxis]]) + displacement: NDArrayFloat = np.hstack( + [x_disp[:, np.newaxis], y_disp[:, np.newaxis]] + ) # preserve z coordinates. right_centerline = centerline.copy() @@ -160,16 +172,24 @@ def centerline_to_polygon( subtract_cond1 = np.logical_and(dx > 0, dy < 0) subtract_cond2 = np.logical_and(dx > 0, dy > 0) subtract_cond = np.logical_or(subtract_cond1, subtract_cond2) - left_centerline, right_centerline = swap_left_and_right(subtract_cond, left_centerline, right_centerline) + left_centerline, right_centerline = swap_left_and_right( + subtract_cond, left_centerline, right_centerline + ) # right centerline also depended on if we added or subtracted y neg_disp_cond = displacement[:, 1] > 0 - left_centerline, right_centerline = swap_left_and_right(neg_disp_cond, left_centerline, right_centerline) + left_centerline, right_centerline = swap_left_and_right( + neg_disp_cond, left_centerline, right_centerline + ) if visualize: plt.scatter(centerline[:, 0], centerline[:, 1], 20, marker=".", color="b") - plt.scatter(right_centerline[:, 0], right_centerline[:, 1], 20, marker=".", color="r") - plt.scatter(left_centerline[:, 0], left_centerline[:, 1], 20, marker=".", color="g") + plt.scatter( + right_centerline[:, 0], right_centerline[:, 1], 20, marker=".", color="r" + ) + plt.scatter( + left_centerline[:, 0], left_centerline[:, 1], 20, marker=".", color="g" + ) fname = datetime.datetime.utcnow().strftime("%Y_%m_%d_%H_%M_%S_%f") plt.savefig(f"polygon_unit_tests/{fname}.png") plt.close("all") @@ -178,7 +198,9 @@ def centerline_to_polygon( return convert_lane_boundaries_to_polygon(right_centerline, left_centerline) -def convert_lane_boundaries_to_polygon(right_lane_bounds: NDArrayFloat, left_lane_bounds: NDArrayFloat) -> NDArrayFloat: +def convert_lane_boundaries_to_polygon( + right_lane_bounds: NDArrayFloat, left_lane_bounds: NDArrayFloat +) -> NDArrayFloat: """Convert lane boundaries to a polygon. Given left and right boundaries of a lane segment, provide the exterior vertices of the @@ -203,9 +225,13 @@ def convert_lane_boundaries_to_polygon(right_lane_bounds: NDArrayFloat, left_lan RuntimeError: If the last dimension of the left and right boundary polylines do not match. """ if not right_lane_bounds.shape[-1] == left_lane_bounds.shape[-1]: - raise RuntimeError("Last dimension of left and right boundary polylines must match.") + raise RuntimeError( + "Last dimension of left and right boundary polylines must match." + ) - polygon: NDArrayFloat = np.vstack([right_lane_bounds, left_lane_bounds[::-1], right_lane_bounds[0]]) + polygon: NDArrayFloat = np.vstack( + [right_lane_bounds, left_lane_bounds[::-1], right_lane_bounds[0]] + ) if not polygon.ndim == 2 or polygon.shape[1] not in [2, 3]: raise RuntimeError("Polygons must be Nx2 or Nx3 in shape.") return polygon diff --git a/src/av2/geometry/se3.py b/src/av2/geometry/se3.py index db199b97..7febb0ce 100644 --- a/src/av2/geometry/se3.py +++ b/src/av2/geometry/se3.py @@ -78,7 +78,9 @@ def inverse(self) -> SE3: Returns: instance of SE3 class, representing inverse of SE3 transformation target_SE3_src. """ - return SE3(rotation=self.rotation.T, translation=self.rotation.T.dot(-self.translation)) + return SE3( + rotation=self.rotation.T, translation=self.rotation.T.dot(-self.translation) + ) def compose(self, right_SE3: SE3) -> SE3: """Compose (right multiply) this class' transformation matrix T with another SE(3) instance. @@ -91,7 +93,9 @@ def compose(self, right_SE3: SE3) -> SE3: Returns: New instance of SE3 class. """ - chained_transform_matrix: NDArrayFloat = self.transform_matrix @ right_SE3.transform_matrix + chained_transform_matrix: NDArrayFloat = ( + self.transform_matrix @ right_SE3.transform_matrix + ) chained_SE3 = SE3( rotation=chained_transform_matrix[:3, :3], translation=chained_transform_matrix[:3, 3], diff --git a/src/av2/geometry/sim2.py b/src/av2/geometry/sim2.py index 40c88018..c84c118b 100644 --- a/src/av2/geometry/sim2.py +++ b/src/av2/geometry/sim2.py @@ -52,7 +52,9 @@ def __post_init__(self) -> None: if not isinstance(self.s, numbers.Number): raise ValueError("Scale `s` must be a numeric type!") if math.isclose(self.s, 0.0): - raise ZeroDivisionError("3x3 matrix formation would require division by zero") + raise ZeroDivisionError( + "3x3 matrix formation would require division by zero" + ) @property def theta_deg(self) -> float: @@ -74,7 +76,9 @@ def theta_deg(self) -> float: def __repr__(self) -> str: """Return a human-readable string representation of the class.""" trans = np.round(self.t, 2) - return f"Angle (deg.): {self.theta_deg:.1f}, Trans.: {trans}, Scale: {self.s:.1f}" + return ( + f"Angle (deg.): {self.theta_deg:.1f}, Trans.: {trans}, Scale: {self.s:.1f}" + ) def __eq__(self, other: object) -> bool: """Check for equality with other Sim(2) object.""" @@ -203,7 +207,9 @@ def from_json(cls, json_fpath: Union[Path, UPath]) -> Sim2: def from_matrix(cls, T: NDArrayFloat) -> Sim2: """Generate class instance from a 3x3 Numpy matrix.""" if np.isclose(T[2, 2], 0.0): - raise ZeroDivisionError("Sim(2) scale calculation would lead to division by zero.") + raise ZeroDivisionError( + "Sim(2) scale calculation would lead to division by zero." + ) R = T[:2, :2] t = T[:2, 2] diff --git a/src/av2/geometry/utm.py b/src/av2/geometry/utm.py index 41562e12..051206c0 100644 --- a/src/av2/geometry/utm.py +++ b/src/av2/geometry/utm.py @@ -50,7 +50,9 @@ class CityName(str, Enum): } -def convert_gps_to_utm(latitude: float, longitude: float, city_name: CityName) -> Tuple[float, float]: +def convert_gps_to_utm( + latitude: float, longitude: float, city_name: CityName +) -> Tuple[float, float]: """Convert GPS coordinates to UTM coordinates. Args: @@ -62,7 +64,13 @@ def convert_gps_to_utm(latitude: float, longitude: float, city_name: CityName) - easting: corresponding UTM Easting. northing: corresponding UTM Northing. """ - projector = Proj(proj="utm", zone=UTM_ZONE_MAP[city_name], ellps="WGS84", datum="WGS84", units="m") + projector = Proj( + proj="utm", + zone=UTM_ZONE_MAP[city_name], + ellps="WGS84", + datum="WGS84", + units="m", + ) # convert to UTM. easting, northing = projector(longitude, latitude) @@ -70,7 +78,9 @@ def convert_gps_to_utm(latitude: float, longitude: float, city_name: CityName) - return easting, northing -def convert_city_coords_to_utm(points_city: Union[NDArrayFloat, NDArrayInt], city_name: CityName) -> NDArrayFloat: +def convert_city_coords_to_utm( + points_city: Union[NDArrayFloat, NDArrayInt], city_name: CityName +) -> NDArrayFloat: """Convert city coordinates to UTM coordinates. Args: @@ -82,12 +92,18 @@ def convert_city_coords_to_utm(points_city: Union[NDArrayFloat, NDArrayInt], cit """ latitude, longitude = CITY_ORIGIN_LATLONG_DICT[city_name] # get (easting, northing) of origin - origin_utm = convert_gps_to_utm(latitude=latitude, longitude=longitude, city_name=city_name) - points_utm: NDArrayFloat = points_city.astype(float) + np.array(origin_utm, dtype=float) + origin_utm = convert_gps_to_utm( + latitude=latitude, longitude=longitude, city_name=city_name + ) + points_utm: NDArrayFloat = points_city.astype(float) + np.array( + origin_utm, dtype=float + ) return points_utm -def convert_city_coords_to_wgs84(points_city: Union[NDArrayFloat, NDArrayInt], city_name: CityName) -> NDArrayFloat: +def convert_city_coords_to_wgs84( + points_city: Union[NDArrayFloat, NDArrayInt], city_name: CityName +) -> NDArrayFloat: """Convert city coordinates to WGS84 coordinates. Args: @@ -99,7 +115,13 @@ def convert_city_coords_to_wgs84(points_city: Union[NDArrayFloat, NDArrayInt], c """ points_utm = convert_city_coords_to_utm(points_city, city_name) - projector = Proj(proj="utm", zone=UTM_ZONE_MAP[city_name], ellps="WGS84", datum="WGS84", units="m") + projector = Proj( + proj="utm", + zone=UTM_ZONE_MAP[city_name], + ellps="WGS84", + datum="WGS84", + units="m", + ) points_wgs84 = [] for easting, northing in points_utm: diff --git a/src/av2/map/drivable_area.py b/src/av2/map/drivable_area.py index b8a46aae..5d27ba5c 100644 --- a/src/av2/map/drivable_area.py +++ b/src/av2/map/drivable_area.py @@ -33,7 +33,9 @@ def xyz(self) -> NDArrayFloat: @classmethod def from_dict(cls, json_data: Dict[str, Any]) -> DrivableArea: """Generate object instance from dictionary read from JSON data.""" - point_list = [Point(x=v["x"], y=v["y"], z=v["z"]) for v in json_data["area_boundary"]] + point_list = [ + Point(x=v["x"], y=v["y"], z=v["z"]) for v in json_data["area_boundary"] + ] # append the first vertex to the end of vertex list point_list.append(point_list[0]) diff --git a/src/av2/map/lane_segment.py b/src/av2/map/lane_segment.py index d2bcafc6..e6c76b07 100644 --- a/src/av2/map/lane_segment.py +++ b/src/av2/map/lane_segment.py @@ -107,7 +107,9 @@ def from_dict(cls, json_data: Dict[str, Any]) -> LaneSegment: return cls( id=json_data["id"], lane_type=LaneType(json_data["lane_type"]), - right_lane_boundary=Polyline.from_json_data(json_data["right_lane_boundary"]), + right_lane_boundary=Polyline.from_json_data( + json_data["right_lane_boundary"] + ), left_lane_boundary=Polyline.from_json_data(json_data["left_lane_boundary"]), right_mark_type=LaneMarkType(json_data["right_lane_mark_type"]), left_mark_type=LaneMarkType(json_data["left_lane_mark_type"]), @@ -122,7 +124,10 @@ def from_dict(cls, json_data: Dict[str, Any]) -> LaneSegment: def left_lane_marking(self) -> LocalLaneMarking: """Retrieve the left lane marking associated with this lane segment.""" return LocalLaneMarking( - mark_type=self.left_mark_type, src_lane_id=self.id, bound_side="left", polyline=self.left_lane_boundary.xyz + mark_type=self.left_mark_type, + src_lane_id=self.id, + bound_side="left", + polyline=self.left_lane_boundary.xyz, ) @property @@ -146,7 +151,9 @@ def polygon_boundary(self) -> NDArrayFloat: self.right_lane_boundary.xyz, self.left_lane_boundary.xyz ) - def is_within_l_infinity_norm_radius(self, query_center: NDArrayFloat, search_radius_m: float) -> bool: + def is_within_l_infinity_norm_radius( + self, query_center: NDArrayFloat, search_radius_m: float + ) -> bool: """Whether any waypoint of lane boundaries falls within search_radius_m of query center, by l-infinity norm. Could have very long segment, with endpoints and all waypoints outside of radius, therefore cannot check just @@ -161,10 +168,12 @@ def is_within_l_infinity_norm_radius(self, query_center: NDArrayFloat, search_ra """ try: right_ln_bnd_interp = interp_utils.interp_arc( - t=WPT_INFINITY_NORM_INTERP_NUM, points=self.right_lane_boundary.xyz[:, :2] + t=WPT_INFINITY_NORM_INTERP_NUM, + points=self.right_lane_boundary.xyz[:, :2], ) left_ln_bnd_interp = interp_utils.interp_arc( - t=WPT_INFINITY_NORM_INTERP_NUM, points=self.left_lane_boundary.xyz[:, :2] + t=WPT_INFINITY_NORM_INTERP_NUM, + points=self.left_lane_boundary.xyz[:, :2], ) except Exception: logger.exception("Interpolation failed for lane segment %d", self.id) diff --git a/src/av2/map/map_api.py b/src/av2/map/map_api.py index 32f2f6d7..a512623d 100644 --- a/src/av2/map/map_api.py +++ b/src/av2/map/map_api.py @@ -93,7 +93,9 @@ def get_raster_values_at_coords( * (npyimage_coords[:, 0] >= 0) * (npyimage_coords[:, 0] < self.array.shape[1]) ) - raster_values[ind_valid_pts] = self.array[npyimage_coords[ind_valid_pts, 1], npyimage_coords[ind_valid_pts, 0]] + raster_values[ind_valid_pts] = self.array[ + npyimage_coords[ind_valid_pts, 1], npyimage_coords[ind_valid_pts, 0] + ] return raster_values @@ -122,13 +124,17 @@ def from_file(cls, log_map_dirpath: Union[Path, UPath]) -> GroundHeightLayer: RuntimeError: If raster ground height layer file is missing or Sim(2) mapping from city to image coordinates is missing. """ - ground_height_npy_fpaths = sorted(log_map_dirpath.glob("*_ground_height_surface____*.npy")) + ground_height_npy_fpaths = sorted( + log_map_dirpath.glob("*_ground_height_surface____*.npy") + ) if not len(ground_height_npy_fpaths) == 1: raise RuntimeError("Raster ground height layer file is missing") Sim2_json_fpaths = sorted(log_map_dirpath.glob("*___img_Sim2_city.json")) if not len(Sim2_json_fpaths) == 1: - raise RuntimeError("Sim(2) mapping from city to image coordinates is missing") + raise RuntimeError( + "Sim(2) mapping from city to image coordinates is missing" + ) # load the file with rasterized values with ground_height_npy_fpaths[0].open("rb") as f: @@ -136,7 +142,9 @@ def from_file(cls, log_map_dirpath: Union[Path, UPath]) -> GroundHeightLayer: array_Sim2_city = Sim2.from_json(Sim2_json_fpaths[0]) - return cls(array=ground_height_array.astype(float), array_Sim2_city=array_Sim2_city) + return cls( + array=ground_height_array.astype(float), array_Sim2_city=array_Sim2_city + ) def get_ground_points_boolean(self, points_xyz: NDArrayFloat) -> NDArrayBool: """Check whether each 3d point is likely to be from the ground surface. @@ -152,11 +160,15 @@ def get_ground_points_boolean(self, points_xyz: NDArrayFloat) -> NDArrayBool: ValueError: If `points_xyz` aren't 3d. """ if points_xyz.shape[1] != 3: - raise ValueError("3-dimensional points must be provided to classify them as `ground` with the map.") + raise ValueError( + "3-dimensional points must be provided to classify them as `ground` with the map." + ) ground_height_values = self.get_ground_height_at_xy(points_xyz) z = points_xyz[:, 2] - near_ground: NDArrayBool = np.absolute(z - ground_height_values) <= GROUND_HEIGHT_THRESHOLD_M + near_ground: NDArrayBool = ( + np.absolute(z - ground_height_values) <= GROUND_HEIGHT_THRESHOLD_M + ) underground: NDArrayBool = z < ground_height_values is_ground_boolean_arr: NDArrayBool = near_ground | underground return is_ground_boolean_arr @@ -181,9 +193,9 @@ def get_ground_height_at_xy(self, points_xyz: NDArrayFloat) -> NDArrayFloat: Returns: Numpy array of shape (K,) """ - ground_height_values: NDArrayFloat = self.get_raster_values_at_coords(points_xyz, fill_value=np.nan).astype( - float - ) + ground_height_values: NDArrayFloat = self.get_raster_values_at_coords( + points_xyz, fill_value=np.nan + ).astype(float) return ground_height_values @@ -195,7 +207,9 @@ class DrivableAreaMapLayer(RasterMapLayer): """ @classmethod - def from_vector_data(cls, drivable_areas: List[DrivableArea]) -> DrivableAreaMapLayer: + def from_vector_data( + cls, drivable_areas: List[DrivableArea] + ) -> DrivableAreaMapLayer: """Return a drivable area map from vector data. NOTE: This function provides "drivable area" as a binary segmentation mask in the bird's eye view. @@ -216,7 +230,9 @@ def from_vector_data(cls, drivable_areas: List[DrivableArea]) -> DrivableAreaMap img_w = int((x_max - x_min + 1) * array_s_city) # scale determines the resolution of the raster DA layer. - array_Sim2_city = Sim2(R=np.eye(2), t=np.array([-x_min, -y_min]), s=array_s_city) + array_Sim2_city = Sim2( + R=np.eye(2), t=np.array([-x_min, -y_min]), s=array_s_city + ) # convert vertices for each polygon from a 3d array in city coordinates, to a 2d array # in image/array coordinates. @@ -239,7 +255,9 @@ class RoiMapLayer(RasterMapLayer): """ @classmethod - def from_drivable_area_layer(cls, drivable_area_layer: DrivableAreaMapLayer) -> RoiMapLayer: + def from_drivable_area_layer( + cls, drivable_area_layer: DrivableAreaMapLayer + ) -> RoiMapLayer: """Rasterize and return 3d vector drivable area as a 2d array, and dilate it by 5 meters, to return a ROI mask. Args: @@ -252,13 +270,19 @@ def from_drivable_area_layer(cls, drivable_area_layer: DrivableAreaMapLayer) -> p_array = array_Sim2_city * p_city """ # initialize ROI as zero-level isocontour of drivable area, and the dilate to 5-meter isocontour - roi_mat_init: NDArrayByte = copy.deepcopy(drivable_area_layer.array).astype(np.uint8) - roi_mask = dilation_utils.dilate_by_l2(roi_mat_init, dilation_thresh=ROI_ISOCONTOUR_GRID) + roi_mat_init: NDArrayByte = copy.deepcopy(drivable_area_layer.array).astype( + np.uint8 + ) + roi_mask = dilation_utils.dilate_by_l2( + roi_mat_init, dilation_thresh=ROI_ISOCONTOUR_GRID + ) return cls(array=roi_mask, array_Sim2_city=drivable_area_layer.array_Sim2_city) -def compute_data_bounds(drivable_areas: List[DrivableArea]) -> Tuple[int, int, int, int]: +def compute_data_bounds( + drivable_areas: List[DrivableArea], +) -> Tuple[int, int, int, int]: """Find the minimum and maximum coordinates along the x and y axes for a set of drivable areas. Args: @@ -326,15 +350,22 @@ def from_json(cls, static_map_path: Union[Path, UPath]) -> ArgoverseStaticMap: log_id = static_map_path.stem.split("log_map_archive_")[1] vector_data = io.read_json_file(static_map_path) - vector_drivable_areas = {da["id"]: DrivableArea.from_dict(da) for da in vector_data["drivable_areas"].values()} - vector_lane_segments = {ls["id"]: LaneSegment.from_dict(ls) for ls in vector_data["lane_segments"].values()} + vector_drivable_areas = { + da["id"]: DrivableArea.from_dict(da) + for da in vector_data["drivable_areas"].values() + } + vector_lane_segments = { + ls["id"]: LaneSegment.from_dict(ls) + for ls in vector_data["lane_segments"].values() + } if "pedestrian_crossings" not in vector_data: logger.error("Missing Pedestrian crossings!") vector_pedestrian_crossings = {} else: vector_pedestrian_crossings = { - pc["id"]: PedestrianCrossing.from_dict(pc) for pc in vector_data["pedestrian_crossings"].values() + pc["id"]: PedestrianCrossing.from_dict(pc) + for pc in vector_data["pedestrian_crossings"].values() } return cls( @@ -348,7 +379,9 @@ def from_json(cls, static_map_path: Union[Path, UPath]) -> ArgoverseStaticMap: ) @classmethod - def from_map_dir(cls, log_map_dirpath: Union[Path, UPath], build_raster: bool = False) -> ArgoverseStaticMap: + def from_map_dir( + cls, log_map_dirpath: Union[Path, UPath], build_raster: bool = False + ) -> ArgoverseStaticMap: """Instantiate an Argoverse map object from data stored within a map data directory. Note: The ground height surface file and associated coordinate mapping is not provided for the @@ -370,7 +403,9 @@ def from_map_dir(cls, log_map_dirpath: Union[Path, UPath], build_raster: bool = # Load vector map data from JSON file vector_data_fnames = sorted(log_map_dirpath.glob("log_map_archive_*.json")) if not len(vector_data_fnames) == 1: - raise RuntimeError(f"JSON file containing vector map data is missing (searched in {log_map_dirpath})") + raise RuntimeError( + f"JSON file containing vector map data is missing (searched in {log_map_dirpath})" + ) vector_data_fname = vector_data_fnames[0].name vector_data_json_path = log_map_dirpath / vector_data_fname @@ -379,10 +414,18 @@ def from_map_dir(cls, log_map_dirpath: Union[Path, UPath], build_raster: bool = # Avoid file I/O and polygon rasterization when not needed if build_raster: - drivable_areas: List[DrivableArea] = list(static_map.vector_drivable_areas.values()) - static_map.raster_drivable_area_layer = DrivableAreaMapLayer.from_vector_data(drivable_areas=drivable_areas) - static_map.raster_roi_layer = RoiMapLayer.from_drivable_area_layer(static_map.raster_drivable_area_layer) - static_map.raster_ground_height_layer = GroundHeightLayer.from_file(log_map_dirpath) + drivable_areas: List[DrivableArea] = list( + static_map.vector_drivable_areas.values() + ) + static_map.raster_drivable_area_layer = ( + DrivableAreaMapLayer.from_vector_data(drivable_areas=drivable_areas) + ) + static_map.raster_roi_layer = RoiMapLayer.from_drivable_area_layer( + static_map.raster_drivable_area_layer + ) + static_map.raster_ground_height_layer = GroundHeightLayer.from_file( + log_map_dirpath + ) return static_map @@ -396,7 +439,9 @@ def get_scenario_vector_drivable_areas(self) -> List[DrivableArea]: """ return list(self.vector_drivable_areas.values()) - def get_lane_segment_successor_ids(self, lane_segment_id: int) -> Optional[List[int]]: + def get_lane_segment_successor_ids( + self, lane_segment_id: int + ) -> Optional[List[int]]: """Get lane id for the lane successor of the specified lane_segment_id. Args: @@ -448,8 +493,12 @@ def get_lane_segment_centerline(self, lane_segment_id: int) -> NDArrayFloat: Returns: Numpy array of shape (N,3). """ - left_ln_bound = self.vector_lane_segments[lane_segment_id].left_lane_boundary.xyz - right_ln_bound = self.vector_lane_segments[lane_segment_id].right_lane_boundary.xyz + left_ln_bound = self.vector_lane_segments[ + lane_segment_id + ].left_lane_boundary.xyz + right_ln_bound = self.vector_lane_segments[ + lane_segment_id + ].right_lane_boundary.xyz lane_centerline, _ = interp_utils.compute_midpoint_line( left_ln_boundary=left_ln_bound, @@ -488,7 +537,9 @@ def get_scenario_ped_crossings(self) -> List[PedestrianCrossing]: """ return list(self.vector_pedestrian_crossings.values()) - def get_nearby_ped_crossings(self, query_center: NDArrayFloat, search_radius_m: float) -> List[PedestrianCrossing]: + def get_nearby_ped_crossings( + self, query_center: NDArrayFloat, search_radius_m: float + ) -> List[PedestrianCrossing]: """Return nearby pedestrian crossings. Returns pedestrian crossings for which any waypoint of their boundary falls within `search_radius_m` meters @@ -513,7 +564,9 @@ def get_scenario_lane_segments(self) -> List[LaneSegment]: """ return list(self.vector_lane_segments.values()) - def get_nearby_lane_segments(self, query_center: NDArrayFloat, search_radius_m: float) -> List[LaneSegment]: + def get_nearby_lane_segments( + self, query_center: NDArrayFloat, search_radius_m: float + ) -> List[LaneSegment]: """Return the nearby lane segments. Return lane segments for which any waypoint of their lane boundaries falls @@ -528,7 +581,9 @@ def get_nearby_lane_segments(self, query_center: NDArrayFloat, search_radius_m: """ scenario_lane_segments = self.get_scenario_lane_segments() return [ - ls for ls in scenario_lane_segments if ls.is_within_l_infinity_norm_radius(query_center, search_radius_m) + ls + for ls in scenario_lane_segments + if ls.is_within_l_infinity_norm_radius(query_center, search_radius_m) ] def remove_ground_surface(self, points_xyz: NDArrayFloat) -> NDArrayFloat: @@ -575,7 +630,9 @@ def remove_non_drivable_area_points(self, points_xyz: NDArrayFloat) -> NDArrayFl Returns: subset of original point cloud, returning only those points lying within the drivable area. """ - is_da_boolean_arr = self.get_raster_layer_points_boolean(points_xyz, layer_name=RasterLayerType.DRIVABLE_AREA) + is_da_boolean_arr = self.get_raster_layer_points_boolean( + points_xyz, layer_name=RasterLayerType.DRIVABLE_AREA + ) filtered_points_xyz: NDArrayFloat = points_xyz[is_da_boolean_arr] return filtered_points_xyz @@ -590,7 +647,9 @@ def remove_non_roi_points(self, points_xyz: NDArrayFloat) -> NDArrayFloat: Returns: subset of original point cloud, returning only those points lying within the ROI. """ - is_da_boolean_arr = self.get_raster_layer_points_boolean(points_xyz, layer_name=RasterLayerType.ROI) + is_da_boolean_arr = self.get_raster_layer_points_boolean( + points_xyz, layer_name=RasterLayerType.ROI + ) filtered_points_xyz: NDArrayFloat = points_xyz[is_da_boolean_arr] return filtered_points_xyz @@ -609,8 +668,13 @@ def get_rasterized_drivable_area(self) -> Tuple[NDArrayByte, Sim2]: if self.raster_drivable_area_layer is None: raise ValueError("Raster drivable area is not loaded!") - raster_drivable_area_layer: NDArrayByte = self.raster_drivable_area_layer.array.astype(np.uint8) - return raster_drivable_area_layer, self.raster_drivable_area_layer.array_Sim2_city + raster_drivable_area_layer: NDArrayByte = ( + self.raster_drivable_area_layer.array.astype(np.uint8) + ) + return ( + raster_drivable_area_layer, + self.raster_drivable_area_layer.array_Sim2_city, + ) def get_rasterized_roi(self) -> Tuple[NDArrayByte, Sim2]: """Get the drivable area along with Sim(2) that maps matrix coordinates to city coordinates. @@ -629,7 +693,9 @@ def get_rasterized_roi(self) -> Tuple[NDArrayByte, Sim2]: raster_roi_layer: NDArrayByte = self.raster_roi_layer.array.astype(np.uint8) return raster_roi_layer, self.raster_roi_layer.array_Sim2_city - def get_raster_layer_points_boolean(self, points_xyz: NDArrayFloat, layer_name: RasterLayerType) -> NDArrayBool: + def get_raster_layer_points_boolean( + self, points_xyz: NDArrayFloat, layer_name: RasterLayerType + ) -> NDArrayBool: """Query the binary segmentation layers (drivable area and ROI) at specific coordinates, to check values. Args: @@ -648,18 +714,24 @@ def get_raster_layer_points_boolean(self, points_xyz: NDArrayFloat, layer_name: if layer_name == RasterLayerType.ROI: if self.raster_roi_layer is None: raise ValueError("Raster ROI is not loaded!") - layer_values = self.raster_roi_layer.get_raster_values_at_coords(points_xyz, fill_value=0) + layer_values = self.raster_roi_layer.get_raster_values_at_coords( + points_xyz, fill_value=0 + ) elif layer_name == RasterLayerType.DRIVABLE_AREA: if self.raster_drivable_area_layer is None: raise ValueError("Raster drivable area is not loaded!") - layer_values = self.raster_drivable_area_layer.get_raster_values_at_coords(points_xyz, fill_value=0) + layer_values = self.raster_drivable_area_layer.get_raster_values_at_coords( + points_xyz, fill_value=0 + ) else: raise ValueError("layer_name should be either `roi` or `drivable_area`.") is_layer_boolean_arr: NDArrayBool = layer_values == 1.0 return is_layer_boolean_arr - def append_height_to_2d_city_pt_cloud(self, points_xy: NDArrayFloat) -> NDArrayFloat: + def append_height_to_2d_city_pt_cloud( + self, points_xy: NDArrayFloat + ) -> NDArrayFloat: """Accept 2d point cloud in xy plane and returns a 3d point cloud (xyz) by querying map for ground height. Args: diff --git a/src/av2/map/pedestrian_crossing.py b/src/av2/map/pedestrian_crossing.py index 9eabbada..271da307 100644 --- a/src/av2/map/pedestrian_crossing.py +++ b/src/av2/map/pedestrian_crossing.py @@ -44,7 +44,9 @@ def __eq__(self, other: object) -> bool: if not isinstance(other, PedestrianCrossing): return False - return np.allclose(self.edge1.xyz, other.edge1.xyz) and np.allclose(self.edge2.xyz, other.edge2.xyz) + return np.allclose(self.edge1.xyz, other.edge1.xyz) and np.allclose( + self.edge2.xyz, other.edge2.xyz + ) @classmethod def from_dict(cls, json_data: Dict[str, Any]) -> PedestrianCrossing: diff --git a/src/av2/rendering/color.py b/src/av2/rendering/color.py index 8dee746b..d4ecfbf4 100644 --- a/src/av2/rendering/color.py +++ b/src/av2/rendering/color.py @@ -69,6 +69,8 @@ def create_colormap(color_list: Sequence[str], n_colors: int) -> NDArrayFloat: array of shape (n_colors, 3) representing a list of RGB colors in [0,1] """ cmap = LinearSegmentedColormap.from_list(name="dummy_name", colors=color_list) - colorscale: NDArrayFloat = np.array([cmap(k * 1 / n_colors) for k in range(n_colors)]) + colorscale: NDArrayFloat = np.array( + [cmap(k * 1 / n_colors) for k in range(n_colors)] + ) # ignore the 4th alpha channel return colorscale[:, :3] diff --git a/src/av2/rendering/map.py b/src/av2/rendering/map.py index 79e6052e..1768e23d 100644 --- a/src/av2/rendering/map.py +++ b/src/av2/rendering/map.py @@ -19,7 +19,12 @@ from av2.geometry.se3 import SE3 from av2.map.lane_segment import LaneMarkType, LaneSegment from av2.map.map_api import ArgoverseStaticMap -from av2.rendering.color import HANDICAP_BLUE_BGR, RED_BGR, TRAFFIC_YELLOW1_BGR, WHITE_BGR +from av2.rendering.color import ( + HANDICAP_BLUE_BGR, + RED_BGR, + TRAFFIC_YELLOW1_BGR, + WHITE_BGR, +) from av2.utils.typing import NDArrayBool, NDArrayByte, NDArrayFloat, NDArrayInt # apply uniform dashes. in reality, there is often a roughly 2:1 ratio between empty space and dashes. @@ -93,8 +98,14 @@ def render_lane_boundary_egoview( else: bound_color = RED_BGR - if ("DOUBLE" in mark_type) or ("SOLID_DASH" in mark_type) or ("DASH_SOLID" in mark_type): - left, right = polyline_utils.get_double_polylines(polyline_city_frame, width_scaling_factor=0.1) + if ( + ("DOUBLE" in mark_type) + or ("SOLID_DASH" in mark_type) + or ("DASH_SOLID" in mark_type) + ): + left, right = polyline_utils.get_double_polylines( + polyline_city_frame, width_scaling_factor=0.1 + ) if mark_type in [ LaneMarkType.SOLID_WHITE, @@ -102,19 +113,39 @@ def render_lane_boundary_egoview( LaneMarkType.SOLID_BLUE, LaneMarkType.NONE, ]: - self.render_polyline_egoview(polyline_city_frame, img_bgr, bound_color, thickness_px=line_width_px) + self.render_polyline_egoview( + polyline_city_frame, img_bgr, bound_color, thickness_px=line_width_px + ) - elif mark_type in [LaneMarkType.DOUBLE_DASH_YELLOW, LaneMarkType.DOUBLE_DASH_WHITE]: + elif mark_type in [ + LaneMarkType.DOUBLE_DASH_YELLOW, + LaneMarkType.DOUBLE_DASH_WHITE, + ]: self.draw_dashed_polyline_egoview( - left, img_bgr, bound_color, thickness_px=line_width_px, dash_interval_m=DASH_INTERVAL_M + left, + img_bgr, + bound_color, + thickness_px=line_width_px, + dash_interval_m=DASH_INTERVAL_M, ) self.draw_dashed_polyline_egoview( - right, img_bgr, bound_color, thickness_px=line_width_px, dash_interval_m=DASH_INTERVAL_M + right, + img_bgr, + bound_color, + thickness_px=line_width_px, + dash_interval_m=DASH_INTERVAL_M, ) - elif mark_type in [LaneMarkType.DOUBLE_SOLID_YELLOW, LaneMarkType.DOUBLE_SOLID_WHITE]: - self.render_polyline_egoview(left, img_bgr, bound_color, thickness_px=line_width_px) - self.render_polyline_egoview(right, img_bgr, bound_color, thickness_px=line_width_px) + elif mark_type in [ + LaneMarkType.DOUBLE_SOLID_YELLOW, + LaneMarkType.DOUBLE_SOLID_WHITE, + ]: + self.render_polyline_egoview( + left, img_bgr, bound_color, thickness_px=line_width_px + ) + self.render_polyline_egoview( + right, img_bgr, bound_color, thickness_px=line_width_px + ) elif mark_type in [LaneMarkType.DASHED_WHITE, LaneMarkType.DASHED_YELLOW]: self.draw_dashed_polyline_egoview( @@ -125,21 +156,35 @@ def render_lane_boundary_egoview( dash_interval_m=DASH_INTERVAL_M, ) - elif (mark_type in [LaneMarkType.SOLID_DASH_WHITE, LaneMarkType.SOLID_DASH_YELLOW] and side == "right") or ( - mark_type == LaneMarkType.DASH_SOLID_YELLOW and side == "left" - ): - self.render_polyline_egoview(left, img_bgr, bound_color, thickness_px=line_width_px) + elif ( + mark_type in [LaneMarkType.SOLID_DASH_WHITE, LaneMarkType.SOLID_DASH_YELLOW] + and side == "right" + ) or (mark_type == LaneMarkType.DASH_SOLID_YELLOW and side == "left"): + self.render_polyline_egoview( + left, img_bgr, bound_color, thickness_px=line_width_px + ) self.draw_dashed_polyline_egoview( - right, img_bgr, bound_color, thickness_px=line_width_px, dash_interval_m=DASH_INTERVAL_M + right, + img_bgr, + bound_color, + thickness_px=line_width_px, + dash_interval_m=DASH_INTERVAL_M, ) - elif (mark_type in [LaneMarkType.SOLID_DASH_WHITE, LaneMarkType.SOLID_DASH_YELLOW] and side == "left") or ( - mark_type == LaneMarkType.DASH_SOLID_YELLOW and side == "right" - ): + elif ( + mark_type in [LaneMarkType.SOLID_DASH_WHITE, LaneMarkType.SOLID_DASH_YELLOW] + and side == "left" + ) or (mark_type == LaneMarkType.DASH_SOLID_YELLOW and side == "right"): self.draw_dashed_polyline_egoview( - left, img_bgr, bound_color, thickness_px=line_width_px, dash_interval_m=DASH_INTERVAL_M + left, + img_bgr, + bound_color, + thickness_px=line_width_px, + dash_interval_m=DASH_INTERVAL_M, + ) + self.render_polyline_egoview( + right, img_bgr, bound_color, thickness_px=line_width_px ) - self.render_polyline_egoview(right, img_bgr, bound_color, thickness_px=line_width_px) else: raise ValueError(f"Unknown marking type {mark_type}") @@ -169,7 +214,12 @@ def draw_dashed_polyline_egoview( dash_frequency: for each dash_interval_m, we will discretize the length into n sections. 1 of n sections will contain a marked dash, and the other (n-1) spaces will be empty (non-marked). """ - interp_polyline, num_waypts = polyline_utils.interp_polyline_by_fixed_waypt_interval(polyline, dash_interval_m) + ( + interp_polyline, + num_waypts, + ) = polyline_utils.interp_polyline_by_fixed_waypt_interval( + polyline, dash_interval_m + ) for i in range(num_waypts - 1): # every other segment is a gap # (first the next dash interval is a line, and then the dash interval is empty, ...) @@ -201,11 +251,17 @@ def render_polyline_egoview( thickness_px: thickness (in pixels) to use when rendering the polyline. """ # must use interpolated, because otherwise points may lie behind camera, etc, cannot draw - interp_polyline_city = interp_utils.interp_arc(t=N_INTERP_PTS, points=polyline_city_frame) - polyline_ego_frame = self.ego_SE3_city.transform_point_cloud(interp_polyline_city) + interp_polyline_city = interp_utils.interp_arc( + t=N_INTERP_PTS, points=polyline_city_frame + ) + polyline_ego_frame = self.ego_SE3_city.transform_point_cloud( + interp_polyline_city + ) # no need to motion compensate, since these are points originally from the city frame. - uv, points_cam, is_valid_points = self.pinhole_cam.project_ego_to_img(polyline_ego_frame) + uv, points_cam, is_valid_points = self.pinhole_cam.project_ego_to_img( + polyline_ego_frame + ) if is_valid_points.sum() == 0: return @@ -215,7 +271,9 @@ def render_polyline_egoview( lane_z = points_cam[:, 2][is_valid_points] if self.depth_map is not None: - allowed_noise = depth_map_utils.compute_allowed_noise_per_point(points_cam[is_valid_points]) + allowed_noise = depth_map_utils.compute_allowed_noise_per_point( + points_cam[is_valid_points] + ) not_occluded = lane_z <= self.depth_map[v, u] + allowed_noise else: not_occluded = np.ones(lane_z.shape, dtype=bool) @@ -262,4 +320,11 @@ def draw_visible_polyline_segments_cv2( y2 = line_segments_arr_int[i + 1][1] # Use anti-aliasing (AA) for curves - image = cv2.line(image, pt1=(x1, y1), pt2=(x2, y2), color=color, thickness=thickness_px, lineType=cv2.LINE_AA) + image = cv2.line( + image, + pt1=(x1, y1), + pt2=(x2, y2), + color=color, + thickness=thickness_px, + lineType=cv2.LINE_AA, + ) diff --git a/src/av2/rendering/ops/draw.py b/src/av2/rendering/ops/draw.py index 7eaf8a29..5457ffcc 100644 --- a/src/av2/rendering/ops/draw.py +++ b/src/av2/rendering/ops/draw.py @@ -16,7 +16,9 @@ @nb.njit -def integer_linear_interpolation(x: np.uint8, y: np.uint8, alpha: np.uint16, beta: np.uint16) -> np.uint8: +def integer_linear_interpolation( + x: np.uint8, y: np.uint8, alpha: np.uint16, beta: np.uint16 +) -> np.uint8: """Approximate floating point linear interpolation. This function approximates the following: @@ -43,7 +45,9 @@ def integer_linear_interpolation(x: np.uint8, y: np.uint8, alpha: np.uint16, bet @nb.njit -def alpha_blend_kernel(fg: NDArrayByte, bg: NDArrayByte, alpha: np.uint8) -> Tuple[np.uint8, np.uint8, np.uint8]: +def alpha_blend_kernel( + fg: NDArrayByte, bg: NDArrayByte, alpha: np.uint8 +) -> Tuple[np.uint8, np.uint8, np.uint8]: """Fast integer alpha blending. Reference: https://stackoverflow.com/a/12016968 @@ -127,10 +131,14 @@ def draw_points_kernel( if (vj >= 0 and uk >= 0) and (vj < H and uk < W): alpha_vj_uk = alpha if with_anti_alias: - alpha_vj_uk *= gaussian_kernel(vj, v, sigma) * gaussian_kernel(uk, u, sigma) + alpha_vj_uk *= gaussian_kernel(vj, v, sigma) * gaussian_kernel( + uk, u, sigma + ) alpha_vj_uk *= float(UINT8_MAX) alpha_vj_uk_uint8 = np.uint8(alpha_vj_uk) - blend = alpha_blend_kernel(colors[i], img[vj, uk], alpha=alpha_vj_uk_uint8) + blend = alpha_blend_kernel( + colors[i], img[vj, uk], alpha=alpha_vj_uk_uint8 + ) img[vj, uk, 0] = blend[0] img[vj, uk, 1] = blend[1] img[vj, uk, 2] = blend[2] @@ -138,7 +146,9 @@ def draw_points_kernel( @nb.njit(nogil=True) -def clip_line_frustum(p1: NDArrayFloat, p2: NDArrayFloat, planes: NDArrayFloat) -> NDArrayFloat: +def clip_line_frustum( + p1: NDArrayFloat, p2: NDArrayFloat, planes: NDArrayFloat +) -> NDArrayFloat: """Iterate over the frustum planes and intersect them with the segment. We exploit the fact that in a camera frustum, all plane normals point inside the frustum volume. diff --git a/src/av2/rendering/rasterize.py b/src/av2/rendering/rasterize.py index 9ecdddbd..81e358cf 100644 --- a/src/av2/rendering/rasterize.py +++ b/src/av2/rendering/rasterize.py @@ -78,7 +78,9 @@ def xyz_to_bev( # Crop point cloud to the region-of-interest. lower_bound_inclusive = (0.0, 0.0, 0.0) indices_cropped, grid_boundary_reduction = crop_points( - indices, lower_bound_inclusive=lower_bound_inclusive, upper_bound_exclusive=grid_size_m + indices, + lower_bound_inclusive=lower_bound_inclusive, + upper_bound_exclusive=grid_size_m, ) # Filter the indices and intensity values. diff --git a/src/av2/rendering/vector.py b/src/av2/rendering/vector.py index 67893089..1b374a7f 100644 --- a/src/av2/rendering/vector.py +++ b/src/av2/rendering/vector.py @@ -65,7 +65,9 @@ def plot_polygon_patch_mpl( vertices = polygon_pts[:, :2] mpath = MPath(vertices, codes) - patch = mpatches.PathPatch(mpath, facecolor=color, edgecolor=color, alpha=alpha, zorder=zorder) + patch = mpatches.PathPatch( + mpath, facecolor=color, edgecolor=color, alpha=alpha, zorder=zorder + ) ax.add_patch(patch) @@ -132,5 +134,7 @@ def draw_line_frustum( uv, _, _ = cam_model.project_cam_to_img(clipped_points) uv_int: NDArrayInt = np.round(uv).astype(int) p1, p2 = uv_int[0], uv_int[1] - img = draw_line_in_img(img, p1, p2, color=color, thickness=thickness, line_type=line_type) + img = draw_line_in_img( + img, p1, p2, color=color, thickness=thickness, line_type=line_type + ) return img diff --git a/src/av2/rendering/video.py b/src/av2/rendering/video.py index eab48e97..13df7cce 100644 --- a/src/av2/rendering/video.py +++ b/src/av2/rendering/video.py @@ -35,7 +35,9 @@ class VideoCodecs(str, Enum): HEVC_VIDEOTOOLBOX = "hevc_videotoolbox" # macOS GPU acceleration. -HIGH_EFFICIENCY_VIDEO_CODECS: Final[Set[VideoCodecs]] = set([VideoCodecs.LIBX265, VideoCodecs.HEVC_VIDEOTOOLBOX]) +HIGH_EFFICIENCY_VIDEO_CODECS: Final[Set[VideoCodecs]] = set( + [VideoCodecs.LIBX265, VideoCodecs.HEVC_VIDEOTOOLBOX] +) def tile_cameras( @@ -80,32 +82,46 @@ def tile_cameras( if "ring_front_center" in named_sensors: ring_front_center = named_sensors["ring_front_center"] - tiled_im_bgr[:landscape_width, landscape_width : landscape_width + landscape_height] = ring_front_center + tiled_im_bgr[ + :landscape_width, landscape_width : landscape_width + landscape_height + ] = ring_front_center if "ring_front_right" in named_sensors: ring_front_right = named_sensors["ring_front_right"] - tiled_im_bgr[:landscape_height, landscape_width + landscape_height :] = ring_front_right + tiled_im_bgr[ + :landscape_height, landscape_width + landscape_height : + ] = ring_front_right if "ring_side_left" in named_sensors: ring_side_left = named_sensors["ring_side_left"] - tiled_im_bgr[landscape_height : 2 * landscape_height, :landscape_width] = ring_side_left + tiled_im_bgr[ + landscape_height : 2 * landscape_height, :landscape_width + ] = ring_side_left if "ring_side_right" in named_sensors: ring_side_right = named_sensors["ring_side_right"] - tiled_im_bgr[landscape_height : 2 * landscape_height, landscape_width + landscape_height :] = ring_side_right + tiled_im_bgr[ + landscape_height : 2 * landscape_height, + landscape_width + landscape_height :, + ] = ring_side_right if bev_img is not None: tiled_im_bgr[ - landscape_width : 2 * landscape_width, landscape_width : landscape_width + landscape_height + landscape_width : 2 * landscape_width, + landscape_width : landscape_width + landscape_height, ] = bev_img if "ring_rear_left" in named_sensors: ring_rear_left = named_sensors["ring_rear_left"] - tiled_im_bgr[2 * landscape_height : 3 * landscape_height, :landscape_width] = ring_rear_left + tiled_im_bgr[ + 2 * landscape_height : 3 * landscape_height, :landscape_width + ] = ring_rear_left if "ring_rear_right" in named_sensors: ring_rear_right = named_sensors["ring_rear_right"] - tiled_im_bgr[2 * landscape_height : 3 * landscape_height, width - landscape_width :] = ring_rear_right + tiled_im_bgr[ + 2 * landscape_height : 3 * landscape_height, width - landscape_width : + ] = ring_rear_right return tiled_im_bgr diff --git a/src/av2/structures/cuboid.py b/src/av2/structures/cuboid.py index a360f044..c9fe567b 100644 --- a/src/av2/structures/cuboid.py +++ b/src/av2/structures/cuboid.py @@ -20,7 +20,13 @@ from av2.rendering.color import BLUE_BGR, TRAFFIC_YELLOW1_BGR from av2.rendering.vector import draw_line_frustum from av2.utils.io import read_feather -from av2.utils.typing import NDArrayBool, NDArrayByte, NDArrayFloat, NDArrayInt, NDArrayObject +from av2.utils.typing import ( + NDArrayBool, + NDArrayByte, + NDArrayFloat, + NDArrayInt, + NDArrayObject, +) ORDERED_CUBOID_COL_NAMES: Final[List[str]] = [ "tx_m", @@ -104,12 +110,16 @@ def vertices_m(self) -> NDArrayFloat: # Transform unit polygons. vertices_obj_xyz_m: NDArrayFloat = (dims_lwh_m / 2.0) * unit_vertices_obj_xyz_m - vertices_dst_xyz_m = self.dst_SE3_object.transform_point_cloud(vertices_obj_xyz_m) + vertices_dst_xyz_m = self.dst_SE3_object.transform_point_cloud( + vertices_obj_xyz_m + ) # Finally, return the polygons. return vertices_dst_xyz_m - def compute_interior_points(self, points_xyz_m: NDArrayFloat) -> Tuple[NDArrayFloat, NDArrayBool]: + def compute_interior_points( + self, points_xyz_m: NDArrayFloat + ) -> Tuple[NDArrayFloat, NDArrayBool]: """Given a query point cloud, filter to points interior to the cuboid, and provide mask. Note: comparison is to cuboid vertices in the destination reference frame. @@ -148,7 +158,10 @@ def transform(self, target_SE3_dst: SE3) -> Cuboid: @classmethod def from_numpy( - cls, params: NDArrayObject, category: Optional[Enum] = None, timestamp_ns: Optional[int] = None + cls, + params: NDArrayObject, + category: Optional[Enum] = None, + timestamp_ns: Optional[int] = None, ) -> Cuboid: """Convert a set of cuboid parameters to a `Cuboid` object. @@ -191,13 +204,17 @@ class CuboidList: @property def xyz_center_m(self) -> NDArrayFloat: """Cartesian coordinate centers (x,y,z) in the destination reference frame.""" - center_xyz_m: NDArrayFloat = np.stack([cuboid.dst_SE3_object.translation for cuboid in self.cuboids]) + center_xyz_m: NDArrayFloat = np.stack( + [cuboid.dst_SE3_object.translation for cuboid in self.cuboids] + ) return center_xyz_m @property def dims_lwh_m(self) -> NDArrayFloat: """Object extents (length,width,height) along the (x,y,z) axes respectively in meters.""" - dims_lwh: NDArrayFloat = np.stack([cuboid.dims_lwh_m for cuboid in self.cuboids]) + dims_lwh: NDArrayFloat = np.stack( + [cuboid.dims_lwh_m for cuboid in self.cuboids] + ) return dims_lwh @cached_property @@ -219,7 +236,9 @@ def vertices_m(self) -> NDArrayFloat: Returns: (N,8,3) array of cuboid vertices. """ - vertices_m: NDArrayFloat = np.stack([cuboid.vertices_m for cuboid in self.cuboids]) + vertices_m: NDArrayFloat = np.stack( + [cuboid.vertices_m for cuboid in self.cuboids] + ) return vertices_m @property @@ -244,7 +263,9 @@ def __getitem__(self, idx: int) -> Cuboid: IndexError: if index is invalid (i.e. out of bounds). """ if idx < 0 or idx >= len(self): - raise IndexError(f"Attempted to access cuboid {idx}, but only indices [0,{len(self)-1}] are valid") + raise IndexError( + f"Attempted to access cuboid {idx}, but only indices [0,{len(self)-1}] are valid" + ) return self.cuboids[idx] def transform(self, target_SE3_dst: SE3) -> CuboidList: @@ -304,13 +325,23 @@ def project_to_cam( # Collapse first dimension to allow for vectorization. cuboids_vertices_ego = cuboids_vertices_ego.reshape(-1, D) if city_SE3_ego_cam_t is not None and city_SE3_ego_lidar_t is not None: - _, cuboids_vertices_cam, _ = cam_model.project_ego_to_img_motion_compensated( - cuboids_vertices_ego, city_SE3_ego_cam_t=city_SE3_ego_cam_t, city_SE3_ego_lidar_t=city_SE3_ego_lidar_t + ( + _, + cuboids_vertices_cam, + _, + ) = cam_model.project_ego_to_img_motion_compensated( + cuboids_vertices_ego, + city_SE3_ego_cam_t=city_SE3_ego_cam_t, + city_SE3_ego_lidar_t=city_SE3_ego_lidar_t, ) else: - _, cuboids_vertices_cam, _ = cam_model.project_ego_to_img(cuboids_vertices_ego) + _, cuboids_vertices_cam, _ = cam_model.project_ego_to_img( + cuboids_vertices_ego + ) - cuboids_vertices_cam = cuboids_vertices_cam[:, :3].reshape(N, V, D) # Unravel collapsed dimension. + cuboids_vertices_cam = cuboids_vertices_cam[:, :3].reshape( + N, V, D + ) # Unravel collapsed dimension. # Compute depth of each cuboid center (mean of the cuboid's vertices). cuboid_centers = cuboids_vertices_cam.mean(axis=1) @@ -406,12 +437,16 @@ def from_dataframe(cls, data: pd.DataFrame) -> CuboidList: Returns: Constructed cuboids. """ - cuboids_parameters: NDArrayFloat = data.loc[:, ORDERED_CUBOID_COL_NAMES].to_numpy() + cuboids_parameters: NDArrayFloat = data.loc[ + :, ORDERED_CUBOID_COL_NAMES + ].to_numpy() categories: NDArrayObject = data.loc[:, "category"].to_numpy() timestamps_ns: NDArrayInt = data.loc[:, "timestamp_ns"].to_numpy() cuboid_list = [ Cuboid.from_numpy(params, category, timestamp_ns) - for params, category, timestamp_ns in zip(cuboids_parameters, categories, timestamps_ns) + for params, category, timestamp_ns in zip( + cuboids_parameters, categories, timestamps_ns + ) ] return cls(cuboid_list) diff --git a/src/av2/structures/ndgrid.py b/src/av2/structures/ndgrid.py index 69d616a2..aebb8912 100644 --- a/src/av2/structures/ndgrid.py +++ b/src/av2/structures/ndgrid.py @@ -38,7 +38,9 @@ def __post_init__(self) -> None: the resolution is not positive. """ if not all(x < y for x, y in zip(self.min_range_m, self.max_range_m)): - raise ValueError("All minimum ranges must be less than their corresponding max ranges!") + raise ValueError( + "All minimum ranges must be less than their corresponding max ranges!" + ) if not all(x > 0 for x in self.resolution_m_per_cell): raise ValueError("Resolution per cell must be positive!") @@ -121,7 +123,10 @@ class BEVGrid(NDGrid): """ def points_to_bev_img( - self, points: NDArrayFloat, color: Tuple[int, int, int] = GRAY_BGR, diameter: int = 2 + self, + points: NDArrayFloat, + color: Tuple[int, int, int] = GRAY_BGR, + diameter: int = 2, ) -> NDArrayByte: """Convert a set of points in Cartesian space to a bird's-eye-view image. @@ -142,7 +147,11 @@ def points_to_bev_img( points_xy = points[..., :2].copy() # Prevent modifying input. indices_int = self.transform_to_grid_coordinates(points_xy) - indices, _ = crop_points(indices_int, lower_bound_inclusive=(0.0, 0.0), upper_bound_exclusive=self.dims) + indices, _ = crop_points( + indices_int, + lower_bound_inclusive=(0.0, 0.0), + upper_bound_exclusive=self.dims, + ) # Construct uv coordinates. H, W = (self.dims[0], self.dims[1]) @@ -152,6 +161,8 @@ def points_to_bev_img( shape = (H, W, C) img: NDArrayByte = np.zeros(shape, dtype=np.uint8) - colors: NDArrayByte = np.array([color for _ in range(len(points_xy))], dtype=np.uint8) + colors: NDArrayByte = np.array( + [color for _ in range(len(points_xy))], dtype=np.uint8 + ) img = draw_points_xy_in_img(img, uv, colors, diameter=diameter) return img diff --git a/src/av2/torch/__init__.py b/src/av2/torch/__init__.py index 96826ea0..f765d47c 100644 --- a/src/av2/torch/__init__.py +++ b/src/av2/torch/__init__.py @@ -5,4 +5,6 @@ LIDAR_COLUMNS: Final = ("x", "y", "z", "intensity") QWXYZ_COLUMNS: Final = ("qw", "qx", "qy", "qz") TRANSLATION_COLUMNS: Final = ("tx_m", "ty_m", "tz_m") -XYZLWH_QWXYZ_COLUMNS: Final = TRANSLATION_COLUMNS + ("length_m", "width_m", "height_m") + QWXYZ_COLUMNS +XYZLWH_QWXYZ_COLUMNS: Final = ( + TRANSLATION_COLUMNS + ("length_m", "width_m", "height_m") + QWXYZ_COLUMNS +) diff --git a/src/av2/torch/data_loaders/detection.py b/src/av2/torch/data_loaders/detection.py index 699d2a52..290a56c2 100644 --- a/src/av2/torch/data_loaders/detection.py +++ b/src/av2/torch/data_loaders/detection.py @@ -6,11 +6,11 @@ from dataclasses import dataclass, field from typing import List +import torch from torch.utils.data import Dataset import av2._r as rust from av2.utils.typing import PathType -import torch from ..structures.sweep import Sweep diff --git a/src/av2/torch/data_loaders/scene_flow.py b/src/av2/torch/data_loaders/scene_flow.py index 6359dc6a..dcc217b6 100644 --- a/src/av2/torch/data_loaders/scene_flow.py +++ b/src/av2/torch/data_loaders/scene_flow.py @@ -53,7 +53,9 @@ def __post_init__(self) -> None: self.num_accumulated_sweeps, self.memory_mapped, ) - self.data_dir = Path(self.root_dir) / self.dataset_name / "sensor" / self.split_name + self.data_dir = ( + Path(self.root_dir) / self.dataset_name / "sensor" / self.split_name + ) @cached_property def file_index(self) -> pd.DataFrame: diff --git a/src/av2/torch/structures/cuboids.py b/src/av2/torch/structures/cuboids.py index e9ca1b66..d38118e8 100644 --- a/src/av2/torch/structures/cuboids.py +++ b/src/av2/torch/structures/cuboids.py @@ -63,10 +63,17 @@ def as_tensor(self, cuboid_mode: CuboidMode = CuboidMode.XYZLWH_T) -> torch.Tens xyzlwh_qwxyz = tensor_from_frame(self._frame, list(XYZLWH_QWXYZ_COLUMNS)) if cuboid_mode == CuboidMode.XYZLWH_T: quat_wxyz = xyzlwh_qwxyz[:, 6:10] - w, x, y, z = quat_wxyz[:, 0], quat_wxyz[:, 1], quat_wxyz[:, 2], quat_wxyz[:, 3] + w, x, y, z = ( + quat_wxyz[:, 0], + quat_wxyz[:, 1], + quat_wxyz[:, 2], + quat_wxyz[:, 3], + ) _, _, yaw = euler_from_quaternion(w, x, y, z) return torch.concat([xyzlwh_qwxyz[:, :6], yaw[:, None]], dim=-1) elif cuboid_mode == CuboidMode.XYZLWH_QWXYZ: return xyzlwh_qwxyz else: - raise NotImplementedError(f"{cuboid_mode} orientation mode is not implemented.") + raise NotImplementedError( + f"{cuboid_mode} orientation mode is not implemented." + ) diff --git a/src/av2/torch/structures/flow.py b/src/av2/torch/structures/flow.py index 59cdf345..18fe5a0c 100644 --- a/src/av2/torch/structures/flow.py +++ b/src/av2/torch/structures/flow.py @@ -9,7 +9,10 @@ from kornia.geometry.linalg import transform_points from torch import BoolTensor, ByteTensor, FloatTensor -from av2.evaluation.scene_flow.constants import CATEGORY_TO_INDEX, SCENE_FLOW_DYNAMIC_THRESHOLD +from av2.evaluation.scene_flow.constants import ( + CATEGORY_TO_INDEX, + SCENE_FLOW_DYNAMIC_THRESHOLD, +) from av2.structures.cuboid import Cuboid, CuboidList from av2.torch.structures.cuboids import Cuboids from av2.torch.structures.sweep import Sweep @@ -59,14 +62,21 @@ def from_sweep_pair(cls, sweeps: Tuple[Sweep, Sweep]) -> Flow: raise ValueError("Can only create flow from sweeps with annotations") else: current_cuboids, next_cuboids = current_sweep.cuboids, next_sweep.cuboids - city_SE3_ego0, city_SE3_ego1 = current_sweep.city_SE3_ego, next_sweep.city_SE3_ego + city_SE3_ego0, city_SE3_ego1 = ( + current_sweep.city_SE3_ego, + next_sweep.city_SE3_ego, + ) ego1_SE3_ego0 = city_SE3_ego1.inverse() * city_SE3_ego0 current_cuboid_map = cuboids_to_id_cuboid_map(current_cuboids) next_cuboid_map = cuboids_to_id_cuboid_map(next_cuboids) current_pc = current_sweep.lidar.as_tensor()[:, :3] - rigid_flow = (transform_points(ego1_SE3_ego0.matrix(), current_pc[None])[0] - current_pc).float().detach() + rigid_flow = ( + (transform_points(ego1_SE3_ego0.matrix(), current_pc[None])[0] - current_pc) + .float() + .detach() + ) flow = rigid_flow.clone() is_valid = torch.ones(len(current_pc), dtype=torch.bool) @@ -77,14 +87,20 @@ def from_sweep_pair(cls, sweeps: Tuple[Sweep, Sweep]) -> Flow: c0.length_m += BOUNDING_BOX_EXPANSION # the bounding boxes are a little too tight sometimes c0.width_m += BOUNDING_BOX_EXPANSION obj_pts_npy, obj_mask_npy = c0.compute_interior_points(current_pc.numpy()) - obj_pts, obj_mask = torch.as_tensor(obj_pts_npy, dtype=torch.float32), torch.as_tensor(obj_mask_npy) + obj_pts, obj_mask = torch.as_tensor( + obj_pts_npy, dtype=torch.float32 + ), torch.as_tensor(obj_mask_npy) category_inds[obj_mask] = CATEGORY_TO_INDEX[str(c0.category)] if id in next_cuboid_map: c1 = next_cuboid_map[id] c1_SE3_c0 = c1.dst_SE3_object.compose(c0.dst_SE3_object.inverse()) flow[obj_mask] = ( - torch.as_tensor(c1_SE3_c0.transform_point_cloud(obj_pts.numpy()), dtype=torch.float32) - obj_pts + torch.as_tensor( + c1_SE3_c0.transform_point_cloud(obj_pts.numpy()), + dtype=torch.float32, + ) + - obj_pts ) else: is_valid[obj_mask] = 0 diff --git a/src/av2/torch/structures/sweep.py b/src/av2/torch/structures/sweep.py index 9c1fc596..2f58d0f5 100644 --- a/src/av2/torch/structures/sweep.py +++ b/src/av2/torch/structures/sweep.py @@ -40,7 +40,9 @@ class Sweep: is_ground: Optional[torch.Tensor] = None @classmethod - def from_rust(cls, sweep: rust.Sweep, avm: Optional[ArgoverseStaticMap] = None) -> Sweep: + def from_rust( + cls, sweep: rust.Sweep, avm: Optional[ArgoverseStaticMap] = None + ) -> Sweep: """Build a sweep from the Rust backend. Args: @@ -60,7 +62,9 @@ def from_rust(cls, sweep: rust.Sweep, avm: Optional[ArgoverseStaticMap] = None) if avm is not None: pcl_ego = lidar.as_tensor()[:, :3] pcl_city_1 = transform_points(city_SE3_ego.matrix(), pcl_ego[None])[0] - is_ground = torch.from_numpy(avm.get_ground_points_boolean(pcl_city_1.numpy()).astype(bool)) + is_ground = torch.from_numpy( + avm.get_ground_points_boolean(pcl_city_1.numpy()).astype(bool) + ) return cls( city_SE3_ego=city_SE3_ego, diff --git a/src/av2/torch/structures/utils.py b/src/av2/torch/structures/utils.py index 3e54a8a3..4b73eb47 100644 --- a/src/av2/torch/structures/utils.py +++ b/src/av2/torch/structures/utils.py @@ -47,7 +47,9 @@ def SE3_from_frame(frame: pd.DataFrame) -> Se3: quat_wxyz = Quaternion(torch.as_tensor(quaternion_npy, dtype=torch.float32)[None]) rotation = So3(quat_wxyz) - translation_npy = frame.loc[0, list(TRANSLATION_COLUMNS)].to_numpy().astype(np.float32) + translation_npy = ( + frame.loc[0, list(TRANSLATION_COLUMNS)].to_numpy().astype(np.float32) + ) translation = torch.as_tensor(translation_npy, dtype=torch.float32)[None] dst_SE3_src = Se3(rotation, translation) dst_SE3_src.rotation._q.requires_grad_(False) diff --git a/src/av2/utils/dataclass.py b/src/av2/utils/dataclass.py index 98a04b3c..a9d4045e 100644 --- a/src/av2/utils/dataclass.py +++ b/src/av2/utils/dataclass.py @@ -36,7 +36,10 @@ def dataclass_eq(base_dataclass: object, other: object) -> bool: # Check whether the dataclasses have equal values in all members base_tuple = vars(base_dataclass).values() other_tuple = vars(other).values() - return all(_dataclass_member_eq(base_mem, other_mem) for base_mem, other_mem in zip(base_tuple, other_tuple)) + return all( + _dataclass_member_eq(base_mem, other_mem) + for base_mem, other_mem in zip(base_tuple, other_tuple) + ) def _dataclass_member_eq(base: object, other: object) -> bool: @@ -55,7 +58,10 @@ def _dataclass_member_eq(base: object, other: object) -> bool: # If both objects are lists, check equality for all members if isinstance(base, list) and isinstance(other, list): - return all(_dataclass_member_eq(base_i, other_i) for base_i, other_i in itertools.zip_longest(base, other)) + return all( + _dataclass_member_eq(base_i, other_i) + for base_i, other_i in itertools.zip_longest(base, other) + ) # If both objects are np arrays, delegate equality check to numpy's built-in operation if isinstance(base, np.ndarray) and isinstance(other, np.ndarray): diff --git a/src/av2/utils/dense_grid_interpolation.py b/src/av2/utils/dense_grid_interpolation.py index af8ca7da..6e68ff45 100644 --- a/src/av2/utils/dense_grid_interpolation.py +++ b/src/av2/utils/dense_grid_interpolation.py @@ -53,9 +53,13 @@ def interp_dense_grid_from_sparse( output_dtype = values.dtype # get (x,y) tuples back - grid_coords = mesh_grid.get_mesh_grid_as_point_cloud(min_x=0, max_x=grid_w - 1, min_y=0, max_y=grid_h - 1) + grid_coords = mesh_grid.get_mesh_grid_as_point_cloud( + min_x=0, max_x=grid_w - 1, min_y=0, max_y=grid_h - 1 + ) # make RGB a function of (dim0=x,dim1=y) - interp_vals = scipy.interpolate.griddata(points, values, grid_coords, method=interp_method) + interp_vals = scipy.interpolate.griddata( + points, values, grid_coords, method=interp_method + ) u = grid_coords[:, 0].astype(np.int32) v = grid_coords[:, 1].astype(np.int32) diff --git a/src/av2/utils/depth_map_utils.py b/src/av2/utils/depth_map_utils.py index 110f4a4d..e6bfbd30 100644 --- a/src/av2/utils/depth_map_utils.py +++ b/src/av2/utils/depth_map_utils.py @@ -10,7 +10,9 @@ from av2.utils.typing import NDArrayByte, NDArrayFloat -MIN_DISTANCE_AWAY_M: Final[float] = 30.0 # assume max noise starting at this distance (meters) +MIN_DISTANCE_AWAY_M: Final[ + float +] = 30.0 # assume max noise starting at this distance (meters) MAX_ALLOWED_NOISE_M: Final[float] = 3.0 # meters @@ -39,7 +41,10 @@ def compute_allowed_noise_per_point(points_cam: NDArrayFloat) -> NDArrayFloat: def vis_depth_map( - img_rgb: NDArrayByte, depth_map: NDArrayFloat, interp_depth_map: bool, num_dilation_iters: int = 10 + img_rgb: NDArrayByte, + depth_map: NDArrayFloat, + interp_depth_map: bool, + num_dilation_iters: int = 10, ) -> None: """Visualize a depth map using Matplotlib's `inferno` colormap side by side with an RGB image. diff --git a/src/av2/utils/dilation_utils.py b/src/av2/utils/dilation_utils.py index 0b14beb3..b63d2c80 100644 --- a/src/av2/utils/dilation_utils.py +++ b/src/av2/utils/dilation_utils.py @@ -37,5 +37,7 @@ def dilate_by_l2(img: NDArrayByte, dilation_thresh: float = 5.0) -> NDArrayByte: distance_mask: NDArrayFloat = cv2.distanceTransform( mask_diff, distanceType=cv2.DIST_L2, maskSize=cv2.DIST_MASK_PRECISE ) - dilated_img: NDArrayByte = np.less_equal(distance_mask.astype(np.float32), dilation_thresh).astype(np.uint8) + dilated_img: NDArrayByte = np.less_equal( + distance_mask.astype(np.float32), dilation_thresh + ).astype(np.uint8) return dilated_img diff --git a/src/av2/utils/helpers.py b/src/av2/utils/helpers.py index 6167d000..5965b8c3 100644 --- a/src/av2/utils/helpers.py +++ b/src/av2/utils/helpers.py @@ -8,7 +8,8 @@ def assert_np_array_shape( - array: Union[NDArrayBool, NDArrayByte, NDArrayFloat, NDArrayInt], target_shape: Sequence[Optional[int]] + array: Union[NDArrayBool, NDArrayByte, NDArrayFloat, NDArrayInt], + target_shape: Sequence[Optional[int]], ) -> None: """Check for shape correctness. @@ -19,6 +20,10 @@ def assert_np_array_shape( Raises: ValueError: if array's shape does not match target_shape for any of the specified dimensions. """ - for index_dim, (array_shape_dim, target_shape_dim) in enumerate(zip(array.shape, target_shape)): + for index_dim, (array_shape_dim, target_shape_dim) in enumerate( + zip(array.shape, target_shape) + ): if target_shape_dim and array_shape_dim != target_shape_dim: - raise ValueError(f"array.shape[{index_dim}]: {array_shape_dim} != {target_shape_dim}.") + raise ValueError( + f"array.shape[{index_dim}]: {array_shape_dim} != {target_shape_dim}." + ) diff --git a/src/av2/utils/io.py b/src/av2/utils/io.py index 4e1c3697..159db41b 100644 --- a/src/av2/utils/io.py +++ b/src/av2/utils/io.py @@ -23,7 +23,9 @@ SensorPosesMapping = Dict[str, SE3] -def read_feather(path: PathType, columns: Optional[Tuple[str, ...]] = None) -> pd.DataFrame: +def read_feather( + path: PathType, columns: Optional[Tuple[str, ...]] = None +) -> pd.DataFrame: """Read Apache Feather data from a .feather file. AV2 uses .feather to serialize much of its data. This function handles the deserialization @@ -38,7 +40,9 @@ def read_feather(path: PathType, columns: Optional[Tuple[str, ...]] = None) -> p (N,len(columns)) Apache Feather data represented as a `pandas` DataFrame. """ with path.open("rb") as file_handle: - dataframe: pd.DataFrame = feather.read_feather(file_handle, columns=columns, memory_map=True) + dataframe: pd.DataFrame = feather.read_feather( + file_handle, columns=columns, memory_map=True + ) return dataframe @@ -65,7 +69,9 @@ def read_lidar_sweep(fpath: Path, attrib_spec: str = "xyz") -> NDArrayFloat: """ possible_attributes = ["x", "y", "z"] if not all([a in possible_attributes for a in attrib_spec]): - raise ValueError("Attributes must be in (x, y, z, intensity, laser_number, offset_ns).") + raise ValueError( + "Attributes must be in (x, y, z, intensity, laser_number, offset_ns)." + ) sweep_df = read_feather(fpath) @@ -108,12 +114,15 @@ def read_ego_SE3_sensor(log_dir: Path) -> SensorPosesMapping: """ ego_SE3_sensor_path = Path(log_dir, "calibration", "egovehicle_SE3_sensor.feather") ego_SE3_sensor = read_feather(ego_SE3_sensor_path) - rotations = geometry_utils.quat_to_mat(ego_SE3_sensor.loc[:, ["qw", "qx", "qy", "qz"]].to_numpy()) + rotations = geometry_utils.quat_to_mat( + ego_SE3_sensor.loc[:, ["qw", "qx", "qy", "qz"]].to_numpy() + ) translations = ego_SE3_sensor.loc[:, ["tx_m", "ty_m", "tz_m"]].to_numpy() sensor_names = ego_SE3_sensor.loc[:, "sensor_name"].to_numpy() sensor_name_to_pose: SensorPosesMapping = { - name: SE3(rotation=rotations[i], translation=translations[i]) for i, name in enumerate(sensor_names) + name: SE3(rotation=rotations[i], translation=translations[i]) + for i, name in enumerate(sensor_names) } return sensor_name_to_pose @@ -159,7 +168,8 @@ def read_city_SE3_ego(log_dir: Union[Path, UPath]) -> TimestampedCitySE3EgoPoses rotation = geometry_utils.quat_to_mat(quat_wxyz) timestamp_city_SE3_ego_dict: TimestampedCitySE3EgoPoses = { - ts: SE3(rotation=rotation[i], translation=translation_xyz_m[i]) for i, ts in enumerate(timestamps_ns) + ts: SE3(rotation=rotation[i], translation=translation_xyz_m[i]) + for i, ts in enumerate(timestamps_ns) } return timestamp_city_SE3_ego_dict diff --git a/src/av2/utils/metric_time.py b/src/av2/utils/metric_time.py index 911761f2..36e55d35 100644 --- a/src/av2/utils/metric_time.py +++ b/src/av2/utils/metric_time.py @@ -26,6 +26,11 @@ def to_metric_time(ts: Union[int, float], src: TimeUnit, dst: TimeUnit) -> float Returns: timestamp expressed now in `dst` units of metric time """ - units_per_sec = {TimeUnit.Second: 1, TimeUnit.Millisecond: 1e3, TimeUnit.Microsecond: 1e6, TimeUnit.Nanosecond: 1e9} + units_per_sec = { + TimeUnit.Second: 1, + TimeUnit.Millisecond: 1e3, + TimeUnit.Microsecond: 1e6, + TimeUnit.Nanosecond: 1e9, + } # ts is in units of `src`, which will cancel with the denominator return ts * (units_per_sec[dst] / units_per_sec[src]) diff --git a/src/av2/utils/raster.py b/src/av2/utils/raster.py index ebcdee77..cbe9c32a 100644 --- a/src/av2/utils/raster.py +++ b/src/av2/utils/raster.py @@ -11,7 +11,9 @@ from av2.utils.typing import NDArrayByte, NDArrayFloat -def get_mask_from_polygons(polygons: List[NDArrayFloat], img_h: int, img_w: int) -> NDArrayByte: +def get_mask_from_polygons( + polygons: List[NDArrayFloat], img_h: int, img_w: int +) -> NDArrayByte: """Rasterize multiple polygons onto a single 2d grid/array. NOTE: Pillow can gracefully handle the scenario when a polygon has coordinates outside of the grid, @@ -34,7 +36,9 @@ def get_mask_from_polygons(polygons: List[NDArrayFloat], img_h: int, img_w: int) return mask -def blend_images(img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7) -> NDArrayByte: +def blend_images( + img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7 +) -> NDArrayByte: """Alpha-blend two images together. Args: diff --git a/src/av2/utils/synchronization_database.py b/src/av2/utils/synchronization_database.py index 8bcd5e63..36b777e9 100644 --- a/src/av2/utils/synchronization_database.py +++ b/src/av2/utils/synchronization_database.py @@ -25,16 +25,24 @@ # constants defined in milliseconds # below evaluates to 33.3 ms -RING_CAM_SHUTTER_INTERVAL_MS: Final[float] = to_metric_time(ts=1 / RING_CAM_FPS, src=Second, dst=Millisecond) +RING_CAM_SHUTTER_INTERVAL_MS: Final[float] = to_metric_time( + ts=1 / RING_CAM_FPS, src=Second, dst=Millisecond +) # below evaluates to 200 ms -STEREO_CAM_SHUTTER_INTERVAL_MS: Final[float] = to_metric_time(ts=1 / STEREO_CAM_FPS, src=Second, dst=Millisecond) +STEREO_CAM_SHUTTER_INTERVAL_MS: Final[float] = to_metric_time( + ts=1 / STEREO_CAM_FPS, src=Second, dst=Millisecond +) # below evaluates to 100 ms -LIDAR_SWEEP_INTERVAL_MS: Final[float] = to_metric_time(ts=1 / LIDAR_FRAME_RATE_HZ, src=Second, dst=Millisecond) +LIDAR_SWEEP_INTERVAL_MS: Final[float] = to_metric_time( + ts=1 / LIDAR_FRAME_RATE_HZ, src=Second, dst=Millisecond +) ALLOWED_TIMESTAMP_BUFFER_MS: Final[int] = 2 # allow 2 ms of buffer -LIDAR_SWEEP_INTERVAL_W_BUFFER_MS: Final[float] = LIDAR_SWEEP_INTERVAL_MS + ALLOWED_TIMESTAMP_BUFFER_MS +LIDAR_SWEEP_INTERVAL_W_BUFFER_MS: Final[float] = ( + LIDAR_SWEEP_INTERVAL_MS + ALLOWED_TIMESTAMP_BUFFER_MS +) def get_timestamps_from_sensor_folder(sensor_folder_wildcard: str) -> NDArrayInt: @@ -50,11 +58,15 @@ def get_timestamps_from_sensor_folder(sensor_folder_wildcard: str) -> NDArrayInt path_generator = glob.glob(sensor_folder_wildcard) path_generator.sort() - timestamps: NDArrayInt = np.array([int(Path(jpg_fpath).stem.split("_")[-1]) for jpg_fpath in path_generator]) + timestamps: NDArrayInt = np.array( + [int(Path(jpg_fpath).stem.split("_")[-1]) for jpg_fpath in path_generator] + ) return timestamps -def find_closest_integer_in_ref_arr(query_int: int, ref_arr: NDArrayInt) -> Tuple[int, int]: +def find_closest_integer_in_ref_arr( + query_int: int, ref_arr: NDArrayInt +) -> Tuple[int, int]: """Find the closest integer to any integer inside a reference array, and the corresponding difference. In our use case, the query integer represents a nanosecond-discretized timestamp, and the @@ -72,7 +84,9 @@ def find_closest_integer_in_ref_arr(query_int: int, ref_arr: NDArrayInt) -> Tupl integer, representing the integer difference between the match and query integers """ closest_ind = np.argmin(np.absolute(ref_arr - query_int)) - closest_int = cast(int, ref_arr[closest_ind]) # mypy does not understand numpy arrays + closest_int = cast( + int, ref_arr[closest_ind] + ) # mypy does not understand numpy arrays int_diff = np.absolute(query_int - closest_int) return closest_int, int_diff @@ -101,7 +115,9 @@ class SynchronizationDB: ts=LIDAR_SWEEP_INTERVAL_W_BUFFER_MS / 2, src=Millisecond, dst=Nanosecond ) - def __init__(self, dataset_dir: str, collect_single_log_id: Optional[str] = None) -> None: + def __init__( + self, dataset_dir: str, collect_single_log_id: Optional[str] = None + ) -> None: """Build the SynchronizationDB. Note that the timestamps for each camera channel are not identical, but they are clustered together. @@ -125,8 +141,12 @@ def __init__(self, dataset_dir: str, collect_single_log_id: Optional[str] = None self.per_log_cam_timestamps_index[log_id] = {} for cam_name in list(RingCameras) + list(StereoCameras): - sensor_folder_wildcard = f"{dataset_dir}/{log_id}/sensors/cameras/{cam_name}/*.jpg" - cam_timestamps = get_timestamps_from_sensor_folder(sensor_folder_wildcard) + sensor_folder_wildcard = ( + f"{dataset_dir}/{log_id}/sensors/cameras/{cam_name}/*.jpg" + ) + cam_timestamps = get_timestamps_from_sensor_folder( + sensor_folder_wildcard + ) self.per_log_cam_timestamps_index[log_id][cam_name] = cam_timestamps sensor_folder_wildcard = f"{dataset_dir}/{log_id}/sensors/lidar/*.feather" @@ -137,7 +157,9 @@ def get_valid_logs(self) -> Iterable[str]: """Return the log_ids for which the SynchronizationDatabase contains pose information.""" return self.per_log_cam_timestamps_index.keys() - def get_closest_lidar_timestamp(self, cam_timestamp_ns: int, log_id: str) -> Optional[int]: + def get_closest_lidar_timestamp( + self, cam_timestamp_ns: int, log_id: str + ) -> Optional[int]: """Given an image timestamp, find the synchronized corresponding LiDAR timestamp. This LiDAR timestamp should have the closest absolute timestamp to the image timestamp. @@ -157,18 +179,26 @@ def get_closest_lidar_timestamp(self, cam_timestamp_ns: int, log_id: str) -> Opt if not lidar_timestamps.tolist(): return None - closest_lidar_timestamp, timestamp_diff = find_closest_integer_in_ref_arr(cam_timestamp_ns, lidar_timestamps) + closest_lidar_timestamp, timestamp_diff = find_closest_integer_in_ref_arr( + cam_timestamp_ns, lidar_timestamps + ) if timestamp_diff > self.MAX_LIDAR_ANYCAM_TIMESTAMP_DIFF: # convert to nanoseconds->milliseconds for readability logger.warning( "No corresponding LiDAR sweep: %s > %s ms", to_metric_time(ts=timestamp_diff, src=Nanosecond, dst=Millisecond), - to_metric_time(ts=self.MAX_LIDAR_ANYCAM_TIMESTAMP_DIFF, src=Nanosecond, dst=Millisecond), + to_metric_time( + ts=self.MAX_LIDAR_ANYCAM_TIMESTAMP_DIFF, + src=Nanosecond, + dst=Millisecond, + ), ) return None return closest_lidar_timestamp - def get_closest_cam_channel_timestamp(self, lidar_timestamp: int, cam_name: str, log_id: str) -> Optional[int]: + def get_closest_cam_channel_timestamp( + self, lidar_timestamp: int, cam_name: str, log_id: str + ) -> Optional[int]: """Given a LiDAR timestamp, find the synchronized corresponding image timestamp for a particular camera. This image timestamp should have the closest absolute timestamp. @@ -181,7 +211,10 @@ def get_closest_cam_channel_timestamp(self, lidar_timestamp: int, cam_name: str, Returns: closest_cam_ch_timestamp: closest timestamp """ - if log_id not in self.per_log_cam_timestamps_index or cam_name not in self.per_log_cam_timestamps_index[log_id]: + if ( + log_id not in self.per_log_cam_timestamps_index + or cam_name not in self.per_log_cam_timestamps_index[log_id] + ): return None cam_timestamps = self.per_log_cam_timestamps_index[log_id][cam_name] @@ -189,23 +222,39 @@ def get_closest_cam_channel_timestamp(self, lidar_timestamp: int, cam_name: str, if not cam_timestamps.tolist(): return None - closest_cam_ch_timestamp, timestamp_diff = find_closest_integer_in_ref_arr(lidar_timestamp, cam_timestamps) - if timestamp_diff > self.MAX_LIDAR_RING_CAM_TIMESTAMP_DIFF and cam_name in list(RingCameras): + closest_cam_ch_timestamp, timestamp_diff = find_closest_integer_in_ref_arr( + lidar_timestamp, cam_timestamps + ) + if ( + timestamp_diff > self.MAX_LIDAR_RING_CAM_TIMESTAMP_DIFF + and cam_name in list(RingCameras) + ): # convert to nanoseconds->milliseconds for readability logger.warning( "No corresponding ring image at %s: %.1f > %s ms", lidar_timestamp, to_metric_time(ts=timestamp_diff, src=Nanosecond, dst=Millisecond), - to_metric_time(ts=self.MAX_LIDAR_RING_CAM_TIMESTAMP_DIFF, src=Nanosecond, dst=Millisecond), + to_metric_time( + ts=self.MAX_LIDAR_RING_CAM_TIMESTAMP_DIFF, + src=Nanosecond, + dst=Millisecond, + ), ) return None - elif timestamp_diff > self.MAX_LIDAR_STEREO_CAM_TIMESTAMP_DIFF and cam_name in list(StereoCameras): + elif ( + timestamp_diff > self.MAX_LIDAR_STEREO_CAM_TIMESTAMP_DIFF + and cam_name in list(StereoCameras) + ): # convert to nanoseconds->milliseconds for readability logger.warning( "No corresponding stereo image at %s: %.1f > %s ms", lidar_timestamp, to_metric_time(ts=timestamp_diff, src=Nanosecond, dst=Millisecond), - to_metric_time(ts=self.MAX_LIDAR_STEREO_CAM_TIMESTAMP_DIFF, src=Nanosecond, dst=Millisecond), + to_metric_time( + ts=self.MAX_LIDAR_STEREO_CAM_TIMESTAMP_DIFF, + src=Nanosecond, + dst=Millisecond, + ), ) return None return closest_cam_ch_timestamp diff --git a/tests/integration/verify_tbv_download.py b/tests/integration/verify_tbv_download.py index e2fa0fea..d1f2bc07 100644 --- a/tests/integration/verify_tbv_download.py +++ b/tests/integration/verify_tbv_download.py @@ -67,7 +67,9 @@ def verify_log_contents(data_root: Path, log_id: str, check_image_sizes: bool) - if not check_image_sizes: continue # every image should be (H,W) = 2048x1550 (front-center) or 775x1024 for all other cameras. - for img_fpath in track(img_fpaths, description=f"Verifying image sizes for {camera_enum}"): + for img_fpath in track( + img_fpaths, description=f"Verifying image sizes for {camera_enum}" + ): img = io_utils.read_img(img_path=img_fpath, channel_order="RGB") if camera_enum == RingCameras.RING_FRONT_CENTER: assert img.shape == (2048, 1550, 3) @@ -84,14 +86,34 @@ def verify_log_contents(data_root: Path, log_id: str, check_image_sizes: bool) - assert poses_fpath.exists() # poses file should be loadable. poses_df = io_utils.read_feather(poses_fpath) - assert list(poses_df.keys()) == ["timestamp_ns", "qw", "qx", "qy", "qz", "tx_m", "ty_m", "tz_m"] + assert list(poses_df.keys()) == [ + "timestamp_ns", + "qw", + "qx", + "qy", + "qz", + "tx_m", + "ty_m", + "tz_m", + ] # every log should have an extrinsics calibration file. - extrinsics_fpath = data_root / log_id / "calibration" / "egovehicle_SE3_sensor.feather" + extrinsics_fpath = ( + data_root / log_id / "calibration" / "egovehicle_SE3_sensor.feather" + ) assert extrinsics_fpath.exists() # extrinsics should be loadable. extrinsics_df = io_utils.read_feather(extrinsics_fpath) - assert list(extrinsics_df.keys()) == ["sensor_name", "qw", "qx", "qy", "qz", "tx_m", "ty_m", "tz_m"] + assert list(extrinsics_df.keys()) == [ + "sensor_name", + "qw", + "qx", + "qy", + "qz", + "tx_m", + "ty_m", + "tz_m", + ] # extrinsics should be provided for each camera. for camera_enum in list(RingCameras): @@ -138,7 +160,9 @@ def verify_log_map(data_root: Path, log_id: str) -> None: assert log_map_dirpath.exists() # every log should have one and only one raster height map. (Note: season is stripped from uuid here). - ground_height_raster_fpaths = list(log_map_dirpath.glob("*_ground_height_surface____*.npy")) + ground_height_raster_fpaths = list( + log_map_dirpath.glob("*_ground_height_surface____*.npy") + ) assert len(ground_height_raster_fpaths) == 1 # every log should have a Sim(2) mapping from raster grid coordinates to city coordinates. @@ -155,7 +179,11 @@ def verify_log_map(data_root: Path, log_id: str) -> None: # every vector map file should have only 3 keys -- "pedestrian_crossings", "lane_segments", "drivable_areas" vector_map_json_data = io_utils.read_json_file(vector_map_fpath) - assert list(vector_map_json_data.keys()) == ["pedestrian_crossings", "lane_segments", "drivable_areas"] + assert list(vector_map_json_data.keys()) == [ + "pedestrian_crossings", + "lane_segments", + "drivable_areas", + ] for _, lane_segment_dict in vector_map_json_data["lane_segments"].items(): assert tuple(lane_segment_dict.keys()) == EXPECTED_LANE_SEGMENT_ATTRIB_KEYS @@ -164,10 +192,14 @@ def verify_log_map(data_root: Path, log_id: str) -> None: avm = ArgoverseStaticMap.from_json(static_map_path=vector_map_fpath) # every map should be loadable w/ build_raster=False - avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath=log_map_dirpath, build_raster=False) + avm = ArgoverseStaticMap.from_map_dir( + log_map_dirpath=log_map_dirpath, build_raster=False + ) # every map should be loadable w/ build_raster=True. - avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath=log_map_dirpath, build_raster=True) + avm = ArgoverseStaticMap.from_map_dir( + log_map_dirpath=log_map_dirpath, build_raster=True + ) # load every lane segment lane_segments = avm.get_scenario_lane_segments() @@ -194,7 +226,9 @@ def verify_logs_using_dataloader(data_root: Path, log_ids: Tuple[str, ...]) -> N log_ids: unique IDs of TbV vehicle logs. """ loader = AV2SensorDataLoader(data_dir=data_root, labels_dir=data_root) - for log_id in track(log_ids, description="Verify logs using an AV2 dataloader object"): + for log_id in track( + log_ids, description="Verify logs using an AV2 dataloader object" + ): logger.info("Verifying log %s", log_id) # city abbreviation should be parsable from every vector map file name, and should fall into 1 of 6 cities city_name = loader.get_city_name(log_id=log_id) @@ -203,7 +237,9 @@ def verify_logs_using_dataloader(data_root: Path, log_ids: Tuple[str, ...]) -> N # pose should be present for every lidar sweep. lidar_timestamps_ns = loader.get_ordered_log_lidar_timestamps(log_id=log_id) for lidar_timestamp_ns in lidar_timestamps_ns: - city_SE3_egovehicle = loader.get_city_SE3_ego(log_id=log_id, timestamp_ns=lidar_timestamp_ns) + city_SE3_egovehicle = loader.get_city_SE3_ego( + log_id=log_id, timestamp_ns=lidar_timestamp_ns + ) assert isinstance(city_SE3_egovehicle, SE3) @@ -230,7 +266,11 @@ def run_verify_all_tbv_logs(data_root: str, check_image_sizes: bool) -> None: for i in range(num_logs): log_id = log_ids[i] logger.info("Verifying log %d: %s", i, log_id) - verify_log_contents(data_root=Path(data_root), log_id=log_id, check_image_sizes=check_image_sizes) + verify_log_contents( + data_root=Path(data_root), + log_id=log_id, + check_image_sizes=check_image_sizes, + ) verify_logs_using_dataloader(data_root=Path(data_root), log_ids=log_ids) diff --git a/tests/unit/datasets/motion_forecasting/eval/test_metrics.py b/tests/unit/datasets/motion_forecasting/eval/test_metrics.py index 3b2ecf38..9ba1a7f1 100644 --- a/tests/unit/datasets/motion_forecasting/eval/test_metrics.py +++ b/tests/unit/datasets/motion_forecasting/eval/test_metrics.py @@ -27,7 +27,11 @@ expected_fde_stationary_k6 = np.full((6,), np.sqrt(2)) # Case 3: K=1 forecast in straight line on X axis -forecasted_trajectories_straight_k1 = np.stack([np.arange(test_N), np.zeros(test_N)], axis=1)[np.newaxis, ...] # 1xNx2 +forecasted_trajectories_straight_k1 = np.stack( + [np.arange(test_N), np.zeros(test_N)], axis=1 +)[ + np.newaxis, ... +] # 1xNx2 expected_ade_straight_k1 = np.full((1,), np.arange(test_N).mean()) expected_fde_straight_k1 = np.full((1,), test_N - 1) @@ -39,7 +43,11 @@ expected_fde_straight_k6 = np.full((6,), test_N - 1) # Case 5: K=1 forecast in diagonal line -forecasted_trajectories_diagonal_k1 = np.stack([np.arange(test_N), np.arange(test_N)], axis=1)[np.newaxis, ...] # 1xNx2 +forecasted_trajectories_diagonal_k1 = np.stack( + [np.arange(test_N), np.arange(test_N)], axis=1 +)[ + np.newaxis, ... +] # 1xNx2 expected_ade_diagonal_k1 = np.full((1,), 6.36396103) expected_fde_diagonal_k1 = np.full((1,), np.hypot(test_N - 1, test_N - 1)) @@ -55,7 +63,9 @@ ], ids=["stationary_k1", "stationary_k6", "straight_k1", "straight_k6", "diagonal_k1"], ) -def test_compute_ade(forecasted_trajectories: NDArrayFloat, expected_ade: NDArrayFloat) -> None: +def test_compute_ade( + forecasted_trajectories: NDArrayFloat, expected_ade: NDArrayFloat +) -> None: """Test that compute_ade returns the correct output with valid inputs. Args: @@ -77,7 +87,9 @@ def test_compute_ade(forecasted_trajectories: NDArrayFloat, expected_ade: NDArra ], ids=["stationary_k1", "stationary_k6", "straight_k1", "straight_k6", "diagonal_k1"], ) -def test_compute_fde(forecasted_trajectories: NDArrayFloat, expected_fde: NDArrayFloat) -> None: +def test_compute_fde( + forecasted_trajectories: NDArrayFloat, expected_fde: NDArrayFloat +) -> None: """Test that compute_fde returns the correct output with valid inputs. Args: @@ -93,14 +105,26 @@ def test_compute_fde(forecasted_trajectories: NDArrayFloat, expected_fde: NDArra [ (forecasted_trajectories_stationary_k1, 2.0, False), (forecasted_trajectories_stationary_k6, 2.0, False), - (forecasted_trajectories_straight_k1, expected_fde_straight_k1[0] + 0.01, False), + ( + forecasted_trajectories_straight_k1, + expected_fde_straight_k1[0] + 0.01, + False, + ), (forecasted_trajectories_straight_k1, expected_fde_straight_k1[0] - 0.01, True), (forecasted_trajectories_diagonal_k1, 2.0, True), ], - ids=["stationary_k1", "stationary_k6", "straight_below_threshold", "straight_above_threshold", "diagonal"], + ids=[ + "stationary_k1", + "stationary_k6", + "straight_below_threshold", + "straight_above_threshold", + "diagonal", + ], ) def test_compute_is_missed_prediction( - forecasted_trajectories: NDArrayFloat, miss_threshold_m: float, expected_is_missed_label: bool + forecasted_trajectories: NDArrayFloat, + miss_threshold_m: float, + expected_is_missed_label: bool, ) -> None: """Test that compute_is_missed_prediction returns the correct output with valid inputs. @@ -126,21 +150,49 @@ def test_compute_is_missed_prediction( wrong_shape_probabilities_k6: NDArrayFloat = np.ones((5,)) / 5 expected_bade_uniform_k1 = expected_ade_stationary_k1 -expected_bade_uniform_k6 = expected_ade_straight_k6 + np.square((1 - uniform_probabilities_k6)) -expected_bade_confident_k6 = expected_ade_straight_k6 + np.square((1 - confident_probabilities_k6)) +expected_bade_uniform_k6 = expected_ade_straight_k6 + np.square( + (1 - uniform_probabilities_k6) +) +expected_bade_confident_k6 = expected_ade_straight_k6 + np.square( + (1 - confident_probabilities_k6) +) expected_bfde_uniform_k1 = expected_fde_stationary_k1 -expected_bfde_uniform_k6 = expected_fde_straight_k6 + np.square((1 - uniform_probabilities_k6)) -expected_bfde_confident_k6 = expected_fde_straight_k6 + np.square((1 - confident_probabilities_k6)) +expected_bfde_uniform_k6 = expected_fde_straight_k6 + np.square( + (1 - uniform_probabilities_k6) +) +expected_bfde_confident_k6 = expected_fde_straight_k6 + np.square( + (1 - confident_probabilities_k6) +) @pytest.mark.parametrize( "forecasted_trajectories, forecast_probabilities, normalize, expected_brier_ade", [ - (forecasted_trajectories_stationary_k1, uniform_probabilities_k1, False, expected_bade_uniform_k1), - (forecasted_trajectories_straight_k6, uniform_probabilities_k6, False, expected_bade_uniform_k6), - (forecasted_trajectories_straight_k6, confident_probabilities_k6, False, expected_bade_confident_k6), - (forecasted_trajectories_straight_k6, non_normalized_probabilities_k6, True, expected_bade_confident_k6), + ( + forecasted_trajectories_stationary_k1, + uniform_probabilities_k1, + False, + expected_bade_uniform_k1, + ), + ( + forecasted_trajectories_straight_k6, + uniform_probabilities_k6, + False, + expected_bade_uniform_k6, + ), + ( + forecasted_trajectories_straight_k6, + confident_probabilities_k6, + False, + expected_bade_confident_k6, + ), + ( + forecasted_trajectories_straight_k6, + non_normalized_probabilities_k6, + True, + expected_bade_confident_k6, + ), ], ids=[ "uniform_stationary_k1", @@ -194,17 +246,40 @@ def test_compute_brier_ade_data_validation( """ with expectation: metrics.compute_brier_ade( - forecasted_trajectories_straight_k6, _STATIONARY_GT_TRAJ, forecast_probabilities, normalize + forecasted_trajectories_straight_k6, + _STATIONARY_GT_TRAJ, + forecast_probabilities, + normalize, ) @pytest.mark.parametrize( "forecasted_trajectories, forecast_probabilities, normalize, expected_brier_fde", [ - (forecasted_trajectories_stationary_k1, uniform_probabilities_k1, False, expected_bfde_uniform_k1), - (forecasted_trajectories_straight_k6, uniform_probabilities_k6, False, expected_bfde_uniform_k6), - (forecasted_trajectories_straight_k6, confident_probabilities_k6, False, expected_bfde_confident_k6), - (forecasted_trajectories_straight_k6, non_normalized_probabilities_k6, True, expected_bfde_confident_k6), + ( + forecasted_trajectories_stationary_k1, + uniform_probabilities_k1, + False, + expected_bfde_uniform_k1, + ), + ( + forecasted_trajectories_straight_k6, + uniform_probabilities_k6, + False, + expected_bfde_uniform_k6, + ), + ( + forecasted_trajectories_straight_k6, + confident_probabilities_k6, + False, + expected_bfde_confident_k6, + ), + ( + forecasted_trajectories_straight_k6, + non_normalized_probabilities_k6, + True, + expected_bfde_confident_k6, + ), ], ids=[ "uniform_stationary_k1", @@ -258,5 +333,8 @@ def test_compute_brier_fde_data_validation( """ with expectation: metrics.compute_brier_fde( - forecasted_trajectories_straight_k6, _STATIONARY_GT_TRAJ, forecast_probabilities, normalize + forecasted_trajectories_straight_k6, + _STATIONARY_GT_TRAJ, + forecast_probabilities, + normalize, ) diff --git a/tests/unit/datasets/motion_forecasting/eval/test_submission.py b/tests/unit/datasets/motion_forecasting/eval/test_submission.py index 0f3d4fc2..091298c0 100644 --- a/tests/unit/datasets/motion_forecasting/eval/test_submission.py +++ b/tests/unit/datasets/motion_forecasting/eval/test_submission.py @@ -20,25 +20,41 @@ ) # Build valid submission with predictions for a single track in a single scenario -valid_track_trajectories: TrackTrajectories = np.zeros((2, AV2_SCENARIO_PRED_TIMESTEPS, 2)) +valid_track_trajectories: TrackTrajectories = np.zeros( + (2, AV2_SCENARIO_PRED_TIMESTEPS, 2) +) valid_scenario_probabilities: ScenarioProbabilities = np.array([0.6, 0.4]) -valid_scenario_trajectories: ScenarioTrajectories = {"valid_track_id": valid_track_trajectories} -valid_submission_predictions = {"valid_scenario_id": (valid_scenario_probabilities, valid_scenario_trajectories)} +valid_scenario_trajectories: ScenarioTrajectories = { + "valid_track_id": valid_track_trajectories +} +valid_submission_predictions = { + "valid_scenario_id": (valid_scenario_probabilities, valid_scenario_trajectories) +} # Build invalid track submission with incorrect prediction length -too_short_track_trajectories: TrackTrajectories = np.zeros((1, AV2_SCENARIO_PRED_TIMESTEPS - 1, 2)) +too_short_track_trajectories: TrackTrajectories = np.zeros( + (1, AV2_SCENARIO_PRED_TIMESTEPS - 1, 2) +) too_short_scenario_probabilities = np.array([1.0]) too_short_scenario_trajectories = {"invalid_track_id": too_short_track_trajectories} too_short_submission_predictions = { - "invalid_scenario_id": (too_short_scenario_probabilities, too_short_scenario_trajectories) + "invalid_scenario_id": ( + too_short_scenario_probabilities, + too_short_scenario_trajectories, + ) } # Build invalid track submission with mismatched predicted trajectories and probabilities -mismatched_track_trajectories: TrackTrajectories = np.zeros((1, AV2_SCENARIO_PRED_TIMESTEPS, 2)) +mismatched_track_trajectories: TrackTrajectories = np.zeros( + (1, AV2_SCENARIO_PRED_TIMESTEPS, 2) +) mismatched_scenario_probabilities = np.array([0.5, 0.5]) mismatched_scenario_trajectories = {"invalid_track_id": mismatched_track_trajectories} mismatched_submission_predictions = { - "invalid_scenario_id": (mismatched_scenario_probabilities, mismatched_scenario_trajectories) + "invalid_scenario_id": ( + mismatched_scenario_probabilities, + mismatched_scenario_trajectories, + ) } @@ -69,7 +85,9 @@ def test_challenge_submission_data_validation( [(valid_submission_predictions)], ids=["valid_submission"], ) -def test_challenge_submission_serialization(tmpdir: Path, test_submission_dict: Dict[str, ScenarioPredictions]) -> None: +def test_challenge_submission_serialization( + tmpdir: Path, test_submission_dict: Dict[str, ScenarioPredictions] +) -> None: """Test that challenge submissions can be serialized/deserialized without changes to internal state. Args: @@ -83,9 +101,16 @@ def test_challenge_submission_serialization(tmpdir: Path, test_submission_dict: deserialized_submission = ChallengeSubmission.from_parquet(submission_file_path) # Check that deserialized data matches original data exactly - for scenario_id, (expected_probabilities, scenario_trajectories) in submission.predictions.items(): + for scenario_id, ( + expected_probabilities, + scenario_trajectories, + ) in submission.predictions.items(): for track_id, expected_trajectories in scenario_trajectories.items(): - deserialized_probabilities = deserialized_submission.predictions[scenario_id][0] - deserialized_trajectories = deserialized_submission.predictions[scenario_id][1][track_id] + deserialized_probabilities = deserialized_submission.predictions[ + scenario_id + ][0] + deserialized_trajectories = deserialized_submission.predictions[ + scenario_id + ][1][track_id] assert np.array_equal(deserialized_trajectories, expected_trajectories) assert np.array_equal(deserialized_probabilities, expected_probabilities) diff --git a/tests/unit/datasets/motion_forecasting/test_scenario_serialization.py b/tests/unit/datasets/motion_forecasting/test_scenario_serialization.py index 6da8cd37..c4be106e 100644 --- a/tests/unit/datasets/motion_forecasting/test_scenario_serialization.py +++ b/tests/unit/datasets/motion_forecasting/test_scenario_serialization.py @@ -8,13 +8,25 @@ import numpy as np from av2.datasets.motion_forecasting import scenario_serialization -from av2.datasets.motion_forecasting.data_schema import ArgoverseScenario, ObjectState, ObjectType, Track, TrackCategory +from av2.datasets.motion_forecasting.data_schema import ( + ArgoverseScenario, + ObjectState, + ObjectType, + Track, + TrackCategory, +) # Build test ArgoverseScenario _TEST_OBJECT_STATES: List[ObjectState] = [ - ObjectState(observed=True, timestep=0, position=(0.0, 0.0), heading=0.0, velocity=(0.0, 0.0)), - ObjectState(observed=True, timestep=1, position=(1.0, 1.0), heading=0.0, velocity=(1.0, 1.0)), - ObjectState(observed=True, timestep=2, position=(2.0, 2.0), heading=0.0, velocity=(2.0, 2.0)), + ObjectState( + observed=True, timestep=0, position=(0.0, 0.0), heading=0.0, velocity=(0.0, 0.0) + ), + ObjectState( + observed=True, timestep=1, position=(1.0, 1.0), heading=0.0, velocity=(1.0, 1.0) + ), + ObjectState( + observed=True, timestep=2, position=(2.0, 2.0), heading=0.0, velocity=(2.0, 2.0) + ), ] _TEST_TRACKS: List[Track] = [ Track( @@ -49,20 +61,31 @@ def test_parquet_scenario_serialization_roundtrip(tmpdir: Path) -> None: """ # Serialize Argoverse scenario to parquet and save to disk scenario_path = tmpdir / "test.parquet" - scenario_serialization.serialize_argoverse_scenario_parquet(scenario_path, _TEST_SCENARIO) + scenario_serialization.serialize_argoverse_scenario_parquet( + scenario_path, _TEST_SCENARIO + ) assert scenario_path.exists(), "Serialized Argoverse scenario not saved to disk." # Check that loading and deserializing a parquet-formatted Argoverse scenario returns an equivalent object - loaded_test_scenario = scenario_serialization.load_argoverse_scenario_parquet(scenario_path) - assert loaded_test_scenario == _TEST_SCENARIO, "Deserialized Argoverse scenario did not match original object." + loaded_test_scenario = scenario_serialization.load_argoverse_scenario_parquet( + scenario_path + ) + assert ( + loaded_test_scenario == _TEST_SCENARIO + ), "Deserialized Argoverse scenario did not match original object." def test_load_argoverse_scenario_parquet(test_data_root_dir: Path) -> None: """Try to load a real scenario from the motion forecasting dataset.""" test_scenario_id = "0a1e6f0a-1817-4a98-b02e-db8c9327d151" test_scenario_path = ( - test_data_root_dir / "forecasting_scenarios" / test_scenario_id / f"scenario_{test_scenario_id}.parquet" + test_data_root_dir + / "forecasting_scenarios" + / test_scenario_id + / f"scenario_{test_scenario_id}.parquet" ) - test_scenario = scenario_serialization.load_argoverse_scenario_parquet(test_scenario_path) + test_scenario = scenario_serialization.load_argoverse_scenario_parquet( + test_scenario_path + ) assert test_scenario.scenario_id == test_scenario_id diff --git a/tests/unit/datasets/sensor/test_av2_sensor_dataloader.py b/tests/unit/datasets/sensor/test_av2_sensor_dataloader.py index 5085f30a..b061c017 100644 --- a/tests/unit/datasets/sensor/test_av2_sensor_dataloader.py +++ b/tests/unit/datasets/sensor/test_av2_sensor_dataloader.py @@ -47,7 +47,9 @@ def test_get_city_SE3_ego(test_data_root_dir: Path) -> None: dataroot = test_data_root_dir / "sensor_dataset_logs" loader = AV2SensorDataLoader(data_dir=dataroot, labels_dir=dataroot) - city_SE3_egovehicle = loader.get_city_SE3_ego(log_id=log_id, timestamp_ns=timestamp_ns) + city_SE3_egovehicle = loader.get_city_SE3_ego( + log_id=log_id, timestamp_ns=timestamp_ns + ) assert isinstance(city_SE3_egovehicle, SE3) expected_translation: NDArrayFloat = np.array([1468.87, 211.51, 13.14]) diff --git a/tests/unit/datasets/sensor/test_sensor_dataloader.py b/tests/unit/datasets/sensor/test_sensor_dataloader.py index da83c508..6b7582dd 100644 --- a/tests/unit/datasets/sensor/test_sensor_dataloader.py +++ b/tests/unit/datasets/sensor/test_sensor_dataloader.py @@ -29,12 +29,25 @@ def _create_dummy_sensor_dataloader(log_id: str) -> SensorDataloader: for t in timestamps_ms: if "ring" in sensor_name: fpath = Path( - sensor_dataset_dir, "dummy", log_id, "sensors", "cameras", sensor_name, f"{int(t*1e6)}.jpg" + sensor_dataset_dir, + "dummy", + log_id, + "sensors", + "cameras", + sensor_name, + f"{int(t*1e6)}.jpg", ) Path(fpath).parent.mkdir(exist_ok=True, parents=True) fpath.open("w").close() elif "lidar" in sensor_name: - fpath = Path(sensor_dataset_dir, "dummy", log_id, "sensors", sensor_name, f"{int(t*1e6)}.feather") + fpath = Path( + sensor_dataset_dir, + "dummy", + log_id, + "sensors", + sensor_name, + f"{int(t*1e6)}.feather", + ) Path(fpath).parent.mkdir(exist_ok=True, parents=True) fpath.open("w").close() return SensorDataloader(dataset_dir=sensor_dataset_dir, with_cache=False) @@ -121,7 +134,9 @@ def test_sensor_data_loader_milliseconds() -> None: # use the non-pandas implementation as a "brute-force" (BF) check. # read out the dataset root from the other dataloader's attributes. - bf_loader = AV2SensorDataLoader(data_dir=loader.dataset_dir / "dummy", labels_dir=loader.dataset_dir / "dummy") + bf_loader = AV2SensorDataLoader( + data_dir=loader.dataset_dir / "dummy", labels_dir=loader.dataset_dir / "dummy" + ) # for every image, make sure query result matches the brute-force query result. for ring_camera_enum in RingCameras: @@ -129,9 +144,14 @@ def test_sensor_data_loader_milliseconds() -> None: for cam_timestamp_ms in SENSOR_TIMESTAMPS_MS_DICT[ring_camera_name]: cam_timestamp_ns = int(cam_timestamp_ms * 1e6) result = loader.get_closest_lidar_fpath( - split="dummy", log_id=log_id, cam_name=ring_camera_name, cam_timestamp_ns=cam_timestamp_ns + split="dummy", + log_id=log_id, + cam_name=ring_camera_name, + cam_timestamp_ns=cam_timestamp_ns, + ) + bf_result = bf_loader.get_closest_lidar_fpath( + log_id=log_id, cam_timestamp_ns=cam_timestamp_ns ) - bf_result = bf_loader.get_closest_lidar_fpath(log_id=log_id, cam_timestamp_ns=cam_timestamp_ns) assert result == bf_result # for every lidar sweep, make sure query result matches the brute-force query result. @@ -140,10 +160,15 @@ def test_sensor_data_loader_milliseconds() -> None: for ring_camera_enum in list(RingCameras): ring_camera_name = ring_camera_enum.value result = loader.get_closest_img_fpath( - split="dummy", log_id=log_id, cam_name=ring_camera_name, lidar_timestamp_ns=lidar_timestamp_ns + split="dummy", + log_id=log_id, + cam_name=ring_camera_name, + lidar_timestamp_ns=lidar_timestamp_ns, ) bf_result = bf_loader.get_closest_img_fpath( - log_id=log_id, cam_name=ring_camera_name, lidar_timestamp_ns=lidar_timestamp_ns + log_id=log_id, + cam_name=ring_camera_name, + lidar_timestamp_ns=lidar_timestamp_ns, ) assert result == bf_result diff --git a/tests/unit/evaluation/detection/test_eval.py b/tests/unit/evaluation/detection/test_eval.py index f17e0a0b..bbf001c8 100644 --- a/tests/unit/evaluation/detection/test_eval.py +++ b/tests/unit/evaluation/detection/test_eval.py @@ -40,14 +40,29 @@ TRANSLATION_COLS: Final = ["tx_m", "ty_m", "tz_m"] DIMS_COLS: Final = ["length_m", "width_m", "height_m"] QUAT_COLS: Final = ["qw", "qx", "qy", "qz"] -ANNO_COLS: Final = ["timestamp_ns", "category"] + DIMS_COLS + QUAT_COLS + TRANSLATION_COLS +ANNO_COLS: Final = ( + ["timestamp_ns", "category"] + DIMS_COLS + QUAT_COLS + TRANSLATION_COLS +) -CUBOID_COLS: Final = ["tx_m", "ty_m", "tz_m", "length_m", "width_m", "height_m", "qw", "qx", "qy", "qz"] +CUBOID_COLS: Final = [ + "tx_m", + "ty_m", + "tz_m", + "length_m", + "width_m", + "height_m", + "qw", + "qx", + "qy", + "qz", +] def _get_summary_identity() -> pd.DataFrame: """Define an evaluator that compares a set of results to itself.""" - detection_cfg = DetectionCfg(categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False) + detection_cfg = DetectionCfg( + categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False + ) dts: pd.DataFrame = pd.read_feather(TEST_DATA_DIR / "detections_identity.feather") gts: pd.DataFrame = dts.copy() gts.loc[:, "num_interior_pts"] = np.array([1, 1, 1, 1, 1, 1]) @@ -57,7 +72,9 @@ def _get_summary_identity() -> pd.DataFrame: def _get_summary_assignment() -> pd.DataFrame: """Define an evaluator that compares a set of results to one with an extra detection to check assignment.""" - detection_cfg = DetectionCfg(categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False) + detection_cfg = DetectionCfg( + categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False + ) dts: pd.DataFrame = pd.read_feather(TEST_DATA_DIR / "detections_assignment.feather") gts: pd.DataFrame = pd.read_feather(TEST_DATA_DIR / "labels.feather") @@ -67,7 +84,9 @@ def _get_summary_assignment() -> pd.DataFrame: def _get_summary() -> pd.DataFrame: """Get a dummy summary.""" - detection_cfg = DetectionCfg(categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False) + detection_cfg = DetectionCfg( + categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False + ) dts: pd.DataFrame = pd.read_feather(TEST_DATA_DIR / "detections.feather") gts: pd.DataFrame = pd.read_feather(TEST_DATA_DIR / "labels.feather") _, _, summary = evaluate(dts, gts, detection_cfg) @@ -120,17 +139,29 @@ def test_orientation_quarter_angles() -> None: """ # Check all of the 90 degree angles expected_result: float = (2 * PI) / 4 - quarter_angles: List[NDArrayFloat] = [np.array([0, 0, angle]) for angle in np.arange(0, 2 * PI, expected_result)] + quarter_angles: List[NDArrayFloat] = [ + np.array([0, 0, angle]) for angle in np.arange(0, 2 * PI, expected_result) + ] for i in range(len(quarter_angles) - 1): - quat_xyzw_dts: NDArrayFloat = Rotation.from_rotvec(quarter_angles[i : i + 1]).as_quat() - quat_xyzw_gts: NDArrayFloat = Rotation.from_rotvec(quarter_angles[i + 1 : i + 2]).as_quat() + quat_xyzw_dts: NDArrayFloat = Rotation.from_rotvec( + quarter_angles[i : i + 1] + ).as_quat() + quat_xyzw_gts: NDArrayFloat = Rotation.from_rotvec( + quarter_angles[i + 1 : i + 2] + ).as_quat() quat_wxyz_dts = quat_xyzw_dts[..., [3, 0, 1, 2]] quat_wxyz_gts = quat_xyzw_gts[..., [3, 0, 1, 2]] - assert np.isclose(distance(quat_wxyz_dts, quat_wxyz_gts, DistanceType.ORIENTATION), expected_result) - assert np.isclose(distance(quat_wxyz_gts, quat_wxyz_dts, DistanceType.ORIENTATION), expected_result) + assert np.isclose( + distance(quat_wxyz_dts, quat_wxyz_gts, DistanceType.ORIENTATION), + expected_result, + ) + assert np.isclose( + distance(quat_wxyz_gts, quat_wxyz_dts, DistanceType.ORIENTATION), + expected_result, + ) def test_orientation_eighth_angles() -> None: @@ -140,7 +171,9 @@ def test_orientation_eighth_angles() -> None: between the detection and ground truth label. """ expected_result: float = (2 * PI) / 8 - eigth_angles: List[NDArrayFloat] = [np.array([0, 0, angle]) for angle in np.arange(0, 2 * PI, expected_result)] + eigth_angles: List[NDArrayFloat] = [ + np.array([0, 0, angle]) for angle in np.arange(0, 2 * PI, expected_result) + ] for i in range(len(eigth_angles) - 1): quat_xyzw_dts = Rotation.from_rotvec(eigth_angles[i : i + 1]).as_quat() @@ -149,8 +182,14 @@ def test_orientation_eighth_angles() -> None: quat_wxyz_dts = quat_xyzw_dts[..., [3, 0, 1, 2]] quat_wxyz_gts = quat_xyzw_gts[..., [3, 0, 1, 2]] - assert np.isclose(distance(quat_wxyz_dts, quat_wxyz_gts, DistanceType.ORIENTATION), expected_result) - assert np.isclose(distance(quat_wxyz_gts, quat_wxyz_dts, DistanceType.ORIENTATION), expected_result) + assert np.isclose( + distance(quat_wxyz_dts, quat_wxyz_gts, DistanceType.ORIENTATION), + expected_result, + ) + assert np.isclose( + distance(quat_wxyz_gts, quat_wxyz_dts, DistanceType.ORIENTATION), + expected_result, + ) def test_wrap_angle() -> None: @@ -166,7 +205,11 @@ def test_accumulate() -> None: gts: pd.DataFrame = pd.read_feather(TEST_DATA_DIR / "labels.feather") for _, group in gts.groupby(["log_id", "timestamp_ns"]): - job = (group.loc[:, CUBOID_COLS].to_numpy(), group.loc[:, CUBOID_COLS + ["num_interior_pts"]].to_numpy(), cfg) + job = ( + group.loc[:, CUBOID_COLS].to_numpy(), + group.loc[:, CUBOID_COLS + ["num_interior_pts"]].to_numpy(), + cfg, + ) dts, gts = accumulate(*job) # Check that there's a true positive under every threshold. @@ -253,7 +296,10 @@ def test_translation_error() -> None: """Test that ATE is 0 for the self-compared results.""" expected_result_identity: float = 0.0 expected_result_det: float = 0.017 # 0.1 / 6, one of six dets is off by 0.1 - assert _get_summary_identity().loc["AVERAGE_METRICS", "ATE"] == expected_result_identity + assert ( + _get_summary_identity().loc["AVERAGE_METRICS", "ATE"] + == expected_result_identity + ) assert _get_summary().loc["AVERAGE_METRICS", "ATE"] == expected_result_det @@ -261,7 +307,10 @@ def test_scale_error() -> None: """Test that ASE is 0 for the self-compared results.""" expected_result_identity: float = 0.0 expected_result_det: float = 0.033 # 0.2 / 6, one of six dets is off by 20% in IoU - assert _get_summary_identity().loc["AVERAGE_METRICS", "ASE"] == expected_result_identity + assert ( + _get_summary_identity().loc["AVERAGE_METRICS", "ASE"] + == expected_result_identity + ) assert _get_summary().loc["AVERAGE_METRICS", "ASE"] == expected_result_det @@ -270,7 +319,10 @@ def test_orientation_error() -> None: expected_result_identity = 0.0 expected_result_det = 0.524 # pi / 6, since one of six dets is off by pi - assert _get_summary_identity().loc["AVERAGE_METRICS", "AOE"] == expected_result_identity + assert ( + _get_summary_identity().loc["AVERAGE_METRICS", "AOE"] + == expected_result_identity + ) assert _get_summary().loc["AVERAGE_METRICS", "AOE"] == expected_result_det @@ -278,13 +330,59 @@ def test_compute_evaluated_dts_mask() -> None: """Unit test for computing valid detections cuboids.""" dts: NDArrayFloat = np.array( [ - [5.0, 5.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0], # In bounds with at least 1 point. - [175, 175.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0], # Out of bounds with at least 1 point. - [-175.0, -175.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0], # Out of bounds with at least 1 point. - [1.0, 1.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0], # In bounds with at least 1 point. + [ + 5.0, + 5.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + ], # In bounds with at least 1 point. + [ + 175, + 175.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + ], # Out of bounds with at least 1 point. + [ + -175.0, + -175.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + ], # Out of bounds with at least 1 point. + [ + 1.0, + 1.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + ], # In bounds with at least 1 point. ], ) - detection_cfg = DetectionCfg(categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False) + detection_cfg = DetectionCfg( + categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False + ) dts_mask = compute_evaluated_dts_mask(dts, detection_cfg) dts_mask_: NDArrayBool = np.array([True, False, False, True]) np.testing.assert_array_equal(dts_mask, dts_mask_) @@ -292,11 +390,13 @@ def test_compute_evaluated_dts_mask() -> None: def test_compute_evaluated_dts_mask_2() -> None: """Randomly generate detections and ensure that they never exceed the maximum detection limit.""" - detection_cfg = DetectionCfg(categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False) + detection_cfg = DetectionCfg( + categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False + ) for i in range(1000): - dts: NDArrayFloat = np.random.randint(0, 250, size=(detection_cfg.max_num_dts_per_category + i, 10)).astype( - float - ) + dts: NDArrayFloat = np.random.randint( + 0, 250, size=(detection_cfg.max_num_dts_per_category + i, 10) + ).astype(float) dts_mask = compute_evaluated_dts_mask(dts, detection_cfg) assert dts_mask.sum() <= detection_cfg.max_num_dts_per_category @@ -305,13 +405,63 @@ def test_compute_evaluated_gts_mask() -> None: """Unit test for computing valid ground truth cuboids.""" gts: NDArrayFloat = np.array( [ - [5.0, 5.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0, 5], # In bounds with at least 1 point. - [175, 175.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0, 5], # Out of bounds with at least 1 point. - [-175.0, -175.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0, 5], # Out of bounds with at least 1 point. - [1.0, 1.0, 5.0, 1.0, 0.0, 0.0, 0.0, 3.0, 4.0, 0.0, 0], # In bounds with at least 1 point. + [ + 5.0, + 5.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + 5, + ], # In bounds with at least 1 point. + [ + 175, + 175.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + 5, + ], # Out of bounds with at least 1 point. + [ + -175.0, + -175.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + 5, + ], # Out of bounds with at least 1 point. + [ + 1.0, + 1.0, + 5.0, + 1.0, + 0.0, + 0.0, + 0.0, + 3.0, + 4.0, + 0.0, + 0, + ], # In bounds with at least 1 point. ], ) - detection_cfg = DetectionCfg(categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False) + detection_cfg = DetectionCfg( + categories=("REGULAR_VEHICLE",), eval_only_roi_instances=False + ) gts_xyz_ego = gts[..., :3] num_interior_pts = gts[..., -1].astype(int) @@ -338,16 +488,23 @@ def test_compute_objects_in_roi_mask() -> None: "a7c8f6a2-26b6-4610-9eb3-294799f9846c", # Two vertices within ROI. ] avm = ArgoverseStaticMap.from_map_dir(map_dir, build_raster=True) - annotations = read_feather(TEST_DATA_DIR / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" / "annotations.feather") - timestamped_city_SE3_egoposes = read_city_SE3_ego(TEST_DATA_DIR / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76") + annotations = read_feather( + TEST_DATA_DIR / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" / "annotations.feather" + ) + timestamped_city_SE3_egoposes = read_city_SE3_ego( + TEST_DATA_DIR / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" + ) selected_cuboids_mask = np.logical_and( - annotations.timestamp_ns == timestamp_ns, annotations["track_uuid"].isin(track_uuids) + annotations.timestamp_ns == timestamp_ns, + annotations["track_uuid"].isin(track_uuids), ) sweep_annotations = annotations.loc[selected_cuboids_mask] mask = compute_objects_in_roi_mask( - sweep_annotations.loc[:, ORDERED_CUBOID_COL_NAMES].to_numpy(), timestamped_city_SE3_egoposes[timestamp_ns], avm + sweep_annotations.loc[:, ORDERED_CUBOID_COL_NAMES].to_numpy(), + timestamped_city_SE3_egoposes[timestamp_ns], + avm, ) mask_: NDArrayBool = np.array([True, False, True]) np.testing.assert_array_equal(mask, mask_) diff --git a/tests/unit/evaluation/scene_flow/test_sf_eval.py b/tests/unit/evaluation/scene_flow/test_sf_eval.py index aa27b4ea..51df13f9 100644 --- a/tests/unit/evaluation/scene_flow/test_sf_eval.py +++ b/tests/unit/evaluation/scene_flow/test_sf_eval.py @@ -32,12 +32,20 @@ dts_close: NDArrayFloat = gts + 0.05 dts_zero: NDArrayFloat = np.zeros_like(gts).astype(np.float64) -gts_dynamic: NDArrayBool = np.array([False, False, False, False, True, True, True, True, False, False]) +gts_dynamic: NDArrayBool = np.array( + [False, False, False, False, True, True, True, True, False, False] +) gts_classes: NDArrayInt = np.array([0, 0, 0, 0, 17, 17, 1, 11, 0, 23]) -gts_valid: NDArrayBool = np.array([False, True, True, True, True, True, True, True, True, True]) -gts_close: NDArrayBool = np.array([True, True, False, False, True, True, False, False, True, True]) +gts_valid: NDArrayBool = np.array( + [False, True, True, True, True, True, True, True, True, True] +) +gts_close: NDArrayBool = np.array( + [True, True, False, False, True, True, False, False, True, True] +) -dts_dynamic: NDArrayBool = np.array([False, True, False, True, False, True, False, True, False, True]) +dts_dynamic: NDArrayBool = np.array( + [False, True, False, True, False, True, False, True, False, True] +) def test_end_point_error() -> None: @@ -46,12 +54,21 @@ def test_end_point_error() -> None: Verify that calculated end-point error matches the expected value. """ epe_perfect = np.zeros(10) - epe_very_close = (np.sqrt(0.01**2 * 3, dtype=np.float64) * np.ones(10)).astype(np.float64) - epe_close = (np.sqrt(0.05**2 * 3, dtype=np.float64) * np.ones(10)).astype(np.float64) - epe_zero = np.array([0.0e00, 1.0e-3, 1.0e-3, 1.0e-3, 1.0e3, 1.0e3, 1.0e3, 1.0e00, 1.0e00, 1.0e00], dtype=np.float64) + epe_very_close = (np.sqrt(0.01**2 * 3, dtype=np.float64) * np.ones(10)).astype( + np.float64 + ) + epe_close = (np.sqrt(0.05**2 * 3, dtype=np.float64) * np.ones(10)).astype( + np.float64 + ) + epe_zero = np.array( + [0.0e00, 1.0e-3, 1.0e-3, 1.0e-3, 1.0e3, 1.0e3, 1.0e3, 1.0e00, 1.0e00, 1.0e00], + dtype=np.float64, + ) assert np.allclose(eval.compute_end_point_error(dts_perfect, gts), epe_perfect) - assert np.allclose(eval.compute_end_point_error(dts_very_close, gts), epe_very_close) + assert np.allclose( + eval.compute_end_point_error(dts_very_close, gts), epe_very_close + ) assert np.allclose(eval.compute_end_point_error(dts_close, gts), epe_close) assert np.allclose(eval.compute_end_point_error(dts_zero, gts), epe_zero) @@ -67,7 +84,9 @@ def test_accuracy_strict() -> None: accS_zero = np.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) assert np.allclose(eval.compute_accuracy_strict(dts_perfect, gts), accS_perfect) - assert np.allclose(eval.compute_accuracy_strict(dts_very_close, gts), accS_very_close) + assert np.allclose( + eval.compute_accuracy_strict(dts_very_close, gts), accS_very_close + ) assert np.allclose(eval.compute_accuracy_strict(dts_close, gts), accS_close) assert np.allclose(eval.compute_accuracy_strict(dts_zero, gts), accS_zero) @@ -83,7 +102,9 @@ def test_accuracy_relax() -> None: accS_zero = np.array([1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) assert np.allclose(eval.compute_accuracy_relax(dts_perfect, gts), accS_perfect) - assert np.allclose(eval.compute_accuracy_relax(dts_very_close, gts), accS_very_close) + assert np.allclose( + eval.compute_accuracy_relax(dts_very_close, gts), accS_very_close + ) assert np.allclose(eval.compute_accuracy_relax(dts_close, gts), accS_close) assert np.allclose(eval.compute_accuracy_relax(dts_zero, gts), accS_zero) @@ -99,7 +120,9 @@ def test_angle_error() -> None: dts_perfect_norm = np.sqrt((dts_perfect**2).sum(-1) + 0.1**2) dts_zero_norm = np.sqrt((dts_zero**2).sum(-1) + 0.1**2) - gts_nz_value = np.array([0.0, 1e-3, 1e-3, 1e-3, 1e3, 1e3, 1e3, 1.0, 1.0, 1.0], dtype=np.float64) + gts_nz_value = np.array( + [0.0, 1e-3, 1e-3, 1e-3, 1e3, 1e3, 1e3, 1.0, 1.0, 1.0], dtype=np.float64 + ) ae_perfect = np.arccos( np.clip( @@ -110,7 +133,8 @@ def test_angle_error() -> None: ) ) ae_close = np.arccos( - 0.1 / dts_close_norm * 0.1 / gts_norm + (gts_nz_value / gts_norm) * ((gts_nz_value + 0.05) / dts_close_norm) + 0.1 / dts_close_norm * 0.1 / gts_norm + + (gts_nz_value / gts_norm) * ((gts_nz_value + 0.05) / dts_close_norm) ) ae_very_close = np.arccos( 0.1 / dts_very_close_norm * 0.1 / gts_norm @@ -203,16 +227,42 @@ def test_metrics() -> None: assert nz_subsets == 5 assert np.allclose( - sum([accr for accr, count in zip(metrics["ACCURACY_RELAX"], metrics["Count"]) if count > 0]) / nz_subsets, 1.0 + sum( + [ + accr + for accr, count in zip(metrics["ACCURACY_RELAX"], metrics["Count"]) + if count > 0 + ] + ) + / nz_subsets, + 1.0, ) assert np.allclose( - sum([accs for accs, count in zip(metrics["ACCURACY_STRICT"], metrics["Count"]) if count > 0]) / nz_subsets, 1.0 + sum( + [ + accs + for accs, count in zip(metrics["ACCURACY_STRICT"], metrics["Count"]) + if count > 0 + ] + ) + / nz_subsets, + 1.0, ) assert np.allclose( - sum([ae for ae, count in zip(metrics["ANGLE_ERROR"], metrics["Count"]) if count > 0]) / nz_subsets, 0.0 + sum( + [ + ae + for ae, count in zip(metrics["ANGLE_ERROR"], metrics["Count"]) + if count > 0 + ] + ) + / nz_subsets, + 0.0, ) assert np.allclose( - sum([epe for epe, count in zip(metrics["EPE"], metrics["Count"]) if count > 0]) / nz_subsets, 0.0 + sum([epe for epe, count in zip(metrics["EPE"], metrics["Count"]) if count > 0]) + / nz_subsets, + 0.0, ) @@ -232,10 +282,34 @@ def test_average_metrics() -> None: timestamp_ns_1 = 111111111111111111 timestamp_ns_2 = 111111111111111112 - write_annotation(gts_classes, gts_close, gts_dynamic, gts_valid, gts, ("log", timestamp_ns_1), anno_dir) - write_annotation(gts_classes, gts_close, gts_dynamic, gts_valid, gts, ("log", timestamp_ns_2), anno_dir) + write_annotation( + gts_classes, + gts_close, + gts_dynamic, + gts_valid, + gts, + ("log", timestamp_ns_1), + anno_dir, + ) + write_annotation( + gts_classes, + gts_close, + gts_dynamic, + gts_valid, + gts, + ("log", timestamp_ns_2), + anno_dir, + ) - write_annotation(gts_classes, gts_close, gts_dynamic, gts_valid, gts, ("log_missing", timestamp_ns_1), anno_dir) + write_annotation( + gts_classes, + gts_close, + gts_dynamic, + gts_valid, + gts, + ("log_missing", timestamp_ns_1), + anno_dir, + ) write_output_file(dts_perfect, dts_dynamic, ("log", timestamp_ns_1), pred_dir) write_output_file(dts_perfect, dts_dynamic, ("log", timestamp_ns_2), pred_dir) @@ -255,11 +329,15 @@ def test_average_metrics() -> None: assert results_df.FP.sum() == 3 * 2 assert results_df.FN.sum() == 2 * 2 - assert results_df.groupby(["Class", "Motion"]).Count.sum().Background.Dynamic == 0 + assert ( + results_df.groupby(["Class", "Motion"]).Count.sum().Background.Dynamic == 0 + ) results_dict = eval.results_to_dict(results_df) assert len(results_dict) == 38 - assert np.isnan([results_dict[k] for k in results_dict if k.endswith("Foreground/Static/Far")]).all() + assert np.isnan( + [results_dict[k] for k in results_dict if k.endswith("Foreground/Static/Far")] + ).all() assert len([True for k in results_dict if "Background/Dyamic" in k]) == 0 assert results_dict["Dynamic IoU"] == 2 / (3 + 2 + 2) assert results_dict["EPE 3-Way Average"] == 0.0 diff --git a/tests/unit/evaluation/scene_flow/test_sf_submission_pipeline.py b/tests/unit/evaluation/scene_flow/test_sf_submission_pipeline.py index 8c863523..7dd8e23d 100644 --- a/tests/unit/evaluation/scene_flow/test_sf_submission_pipeline.py +++ b/tests/unit/evaluation/scene_flow/test_sf_submission_pipeline.py @@ -12,7 +12,10 @@ from av2.evaluation.scene_flow.example_submission import example_submission from av2.evaluation.scene_flow.make_annotation_files import make_annotation_files from av2.evaluation.scene_flow.make_mask_files import make_mask_files -from av2.evaluation.scene_flow.make_submission_archive import make_submission_archive, validate +from av2.evaluation.scene_flow.make_submission_archive import ( + make_submission_archive, + validate, +) from av2.evaluation.scene_flow.utils import compute_eval_point_mask from av2.torch.data_loaders.scene_flow import SceneFlowDataloader @@ -39,10 +42,18 @@ def test_submission() -> None: make_mask_files(str(mask_file), str(_TEST_DATA_ROOT), "test_data", "val") annotations_dir = test_dir / "annotations" - make_annotation_files(str(annotations_dir), str(mask_file), str(_TEST_DATA_ROOT), "test_data", "val") + make_annotation_files( + str(annotations_dir), + str(mask_file), + str(_TEST_DATA_ROOT), + "test_data", + "val", + ) predictions_dir = test_dir / "output" - example_submission(str(predictions_dir), str(mask_file), str(_TEST_DATA_ROOT), "test_data") + example_submission( + str(predictions_dir), str(mask_file), str(_TEST_DATA_ROOT), "test_data" + ) results = eval.evaluate(str(annotations_dir), str(predictions_dir)) for metric in results: @@ -50,7 +61,9 @@ def test_submission() -> None: if "Static" in metric: assert np.allclose(results[metric], 1.0) elif "Dynamic" in metric and "Strict" in metric: - assert np.isnan(results[metric]) or np.allclose(results[metric], 0.0) + assert np.isnan(results[metric]) or np.allclose( + results[metric], 0.0 + ) elif metric.startswith("EPE"): if "Static" in metric: if "Background" in metric: @@ -62,23 +75,34 @@ def test_submission() -> None: assert results[metric] < 1e-4 output_file = test_dir / "submission.zip" - success = make_submission_archive(str(predictions_dir), str(mask_file), str(output_file)) + success = make_submission_archive( + str(predictions_dir), str(mask_file), str(output_file) + ) assert success assert output_file.stat().st_size > 0 annotation_files = list(annotations_dir.rglob("*.feather")) - print([anno_file.relative_to(annotations_dir).as_posix() for anno_file in annotation_files]) + print( + [ + anno_file.relative_to(annotations_dir).as_posix() + for anno_file in annotation_files + ] + ) with ZipFile(output_file, "r") as zf: files = {f.filename for f in zf.filelist} print(files) - results_zip = eval.results_to_dict(eval.evaluate_zip(annotations_dir, output_file)) + results_zip = eval.results_to_dict( + eval.evaluate_zip(annotations_dir, output_file) + ) for metric in results: assert np.allclose(results[metric], results_zip[metric], equal_nan=True) empty_predictions_dir = test_dir / "bad_output_1" empty_predictions_dir.mkdir() - success = make_submission_archive(str(empty_predictions_dir), str(mask_file), str(output_file)) + success = make_submission_archive( + str(empty_predictions_dir), str(mask_file), str(output_file) + ) assert not success failed = False diff --git a/tests/unit/geometry/test_geometry.py b/tests/unit/geometry/test_geometry.py index 83fad36a..141f0c86 100644 --- a/tests/unit/geometry/test_geometry.py +++ b/tests/unit/geometry/test_geometry.py @@ -82,9 +82,14 @@ def test_cart_to_hom_3d() -> None: np.array([[1, -1], [1, -2], [0, -2], [0, -1]]).astype(np.float64), ), ], - ids=[f"Cartesian to texture coordinates conversion (Test Case: {idx + 1})" for idx in range(2)], + ids=[ + f"Cartesian to texture coordinates conversion (Test Case: {idx + 1})" + for idx in range(2) + ], ) -def test_xy_to_uv(xy: NDArrayFloat, width: int, height: int, expected_uv: NDArrayFloat) -> None: +def test_xy_to_uv( + xy: NDArrayFloat, width: int, height: int, expected_uv: NDArrayFloat +) -> None: """Test conversion of coordinates in R^2 (x,y) to texture coordinates (u,v) in R^2. Args: @@ -122,15 +127,24 @@ def test_quat_to_mat_3d(quat_wxyz: NDArrayFloat) -> None: @pytest.mark.parametrize( "cart_xyz, expected_sph_theta_phi_r", [ - (np.array([1, 1, 1]).astype(np.float64), np.array([0.78539816, 0.61547971, 1.73205081])), + ( + np.array([1, 1, 1]).astype(np.float64), + np.array([0.78539816, 0.61547971, 1.73205081]), + ), ( np.array([[1, 1, 1], [1, 2, 0]]).astype(np.float64), - np.array([[0.78539816, 0.61547971, 1.73205081], [1.10714872, 0.0, 2.23606798]]), + np.array( + [[0.78539816, 0.61547971, 1.73205081], [1.10714872, 0.0, 2.23606798]] + ), ), ], - ids=[f"Cartesian to Spherical coordinates (Test Case: {idx + 1})" for idx in range(2)], + ids=[ + f"Cartesian to Spherical coordinates (Test Case: {idx + 1})" for idx in range(2) + ], ) -def test_cart_to_sph_3d(cart_xyz: NDArrayFloat, expected_sph_theta_phi_r: NDArrayFloat) -> None: +def test_cart_to_sph_3d( + cart_xyz: NDArrayFloat, expected_sph_theta_phi_r: NDArrayFloat +) -> None: """Test conversion of cartesian coordinates to spherical coordinates. Args: @@ -144,38 +158,83 @@ def test_cart_to_sph_3d(cart_xyz: NDArrayFloat, expected_sph_theta_phi_r: NDArra "points_xyz, lower_bound_inclusive, upper_bound_exclusive, expected_crop_points, expected_mask", [ ( - np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [0, 0, 1], [1, 0, 1], [1, 1, 1]]).astype( - np.int64 - ), + np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 1, 1], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + ] + ).astype(np.int64), (0, 0, 0), (1.5, 1.5, 1.5), - np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [0, 0, 1], [1, 0, 1], [1, 1, 1]]).astype( - np.int64 - ), + np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 1, 1], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + ] + ).astype(np.int64), np.array([True] * 8), ), ( - np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [0, 0, 1], [1, 0, 1], [1, 1, 1]]).astype( - np.int64 - ), + np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 1, 1], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + ] + ).astype(np.int64), (0, 0, 0), (0.5, 0.5, 0.5), np.array([[0, 0, 0]]).astype(np.int64), np.array([True] + [False] * 7), ), ( - np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [0, 0, 1], [1, 0, 1], [1, 1, 1]]).astype( - np.int64 - ), + np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 1, 1], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + ] + ).astype(np.int64), (0, 0, 0), (1.25, 1.25, 1.0), np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0]]).astype(np.int64), np.array([True] * 4 + [False] * 4), ), ( - np.array([[0, 0, 0], [1, 0, 0], [1, 1, 0], [0, 1, 0], [0, 1, 1], [0, 0, 1], [1, 0, 1], [1, 1, 1]]).astype( - np.int64 - ), + np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 1, 1], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + ] + ).astype(np.int64), (-1.0, -1.0, -1.0), (0.0, 1.0, 1.0), np.empty((0, 3)).astype(np.int64), @@ -205,7 +264,9 @@ def test_crop_points( expected_crop_points: (...,3) Expected tuple of cropped Cartesian coordinates. expected_mask: (...,) Expected boolean mask. """ - cropped_xyz, mask = geometry_utils.crop_points(points_xyz, lower_bound_inclusive, upper_bound_exclusive) + cropped_xyz, mask = geometry_utils.crop_points( + points_xyz, lower_bound_inclusive, upper_bound_exclusive + ) np.testing.assert_array_equal(expected_crop_points, cropped_xyz) np.testing.assert_array_equal(expected_mask, mask) @@ -267,7 +328,9 @@ def test_crop_points( "No points lie inside the bounding box.", ], ) -def test_compute_interior_points_mask(points_xyz: NDArrayFloat, expected_is_interior: NDArrayBool) -> None: +def test_compute_interior_points_mask( + points_xyz: NDArrayFloat, expected_is_interior: NDArrayBool +) -> None: r"""Test finding the points interior to an axis-aligned cuboid. Reference: https://math.stackexchange.com/questions/1472049/check-if-a-point-is-inside-a-rectangular-shaped-area-3d @@ -300,11 +363,15 @@ def test_compute_interior_points_mask(points_xyz: NDArrayFloat, expected_is_inte category=AnnotationCategories.REGULAR_VEHICLE, timestamp_ns=0, ) - is_interior = geometry_utils.compute_interior_points_mask(points_xyz, cuboid.vertices_m) + is_interior = geometry_utils.compute_interior_points_mask( + points_xyz, cuboid.vertices_m + ) assert np.array_equal(is_interior, expected_is_interior) -def test_benchmark_compute_interior_points_mask_optimized(benchmark: Callable[..., Any]) -> None: +def test_benchmark_compute_interior_points_mask_optimized( + benchmark: Callable[..., Any] +) -> None: """Benchmark compute_interior_pts on 100000 random points.""" rotation: NDArrayFloat = np.eye(3) translation: NDArrayFloat = np.array([0.5, 0.5, 0.5]) @@ -320,10 +387,14 @@ def test_benchmark_compute_interior_points_mask_optimized(benchmark: Callable[.. N = 100000 points_xyz: NDArrayFloat = 100.0 * np.random.rand(N, 3) - benchmark(geometry_utils.compute_interior_points_mask, points_xyz, cuboid.vertices_m) + benchmark( + geometry_utils.compute_interior_points_mask, points_xyz, cuboid.vertices_m + ) -def test_benchmark_compute_interior_points_mask_slow(benchmark: Callable[..., Any]) -> None: +def test_benchmark_compute_interior_points_mask_slow( + benchmark: Callable[..., Any] +) -> None: """Benchmark compute_interior_points_mask on 100000 random points.""" rotation: NDArrayFloat = np.eye(3) translation: NDArrayFloat = np.array([0.5, 0.5, 0.5]) @@ -340,7 +411,9 @@ def test_benchmark_compute_interior_points_mask_slow(benchmark: Callable[..., An N = 100000 points_xyz: NDArrayFloat = 100.0 * np.random.rand(N, 3) - def compute_interior_points_mask_slow(points_xyz: NDArrayFloat, cuboid_vertices: NDArrayFloat) -> NDArrayBool: + def compute_interior_points_mask_slow( + points_xyz: NDArrayFloat, cuboid_vertices: NDArrayFloat + ) -> NDArrayBool: """Compute the interior points mask with the older slow version. Args: @@ -357,30 +430,38 @@ def compute_interior_points_mask_slow(points_xyz: NDArrayFloat, cuboid_vertices: # point x lies within the box when the following # constraints are respected valid_u1 = np.logical_and( - u.dot(cuboid_vertices[2]) <= points_xyz.dot(u), points_xyz.dot(u) <= u.dot(cuboid_vertices[6]) + u.dot(cuboid_vertices[2]) <= points_xyz.dot(u), + points_xyz.dot(u) <= u.dot(cuboid_vertices[6]), ) valid_v1 = np.logical_and( - v.dot(cuboid_vertices[2]) <= points_xyz.dot(v), points_xyz.dot(v) <= v.dot(cuboid_vertices[3]) + v.dot(cuboid_vertices[2]) <= points_xyz.dot(v), + points_xyz.dot(v) <= v.dot(cuboid_vertices[3]), ) valid_w1 = np.logical_and( - w.dot(cuboid_vertices[2]) <= points_xyz.dot(w), points_xyz.dot(w) <= w.dot(cuboid_vertices[1]) + w.dot(cuboid_vertices[2]) <= points_xyz.dot(w), + points_xyz.dot(w) <= w.dot(cuboid_vertices[1]), ) valid_u2 = np.logical_and( - u.dot(cuboid_vertices[2]) >= points_xyz.dot(u), points_xyz.dot(u) >= u.dot(cuboid_vertices[6]) + u.dot(cuboid_vertices[2]) >= points_xyz.dot(u), + points_xyz.dot(u) >= u.dot(cuboid_vertices[6]), ) valid_v2 = np.logical_and( - v.dot(cuboid_vertices[2]) >= points_xyz.dot(v), points_xyz.dot(v) >= v.dot(cuboid_vertices[3]) + v.dot(cuboid_vertices[2]) >= points_xyz.dot(v), + points_xyz.dot(v) >= v.dot(cuboid_vertices[3]), ) valid_w2 = np.logical_and( - w.dot(cuboid_vertices[2]) >= points_xyz.dot(w), points_xyz.dot(w) >= w.dot(cuboid_vertices[1]) + w.dot(cuboid_vertices[2]) >= points_xyz.dot(w), + points_xyz.dot(w) >= w.dot(cuboid_vertices[1]), ) valid_u = np.logical_or(valid_u1, valid_u2) valid_v = np.logical_or(valid_v1, valid_v2) valid_w = np.logical_or(valid_w1, valid_w2) - is_interior: NDArrayBool = np.logical_and(np.logical_and(valid_u, valid_v), valid_w) + is_interior: NDArrayBool = np.logical_and( + np.logical_and(valid_u, valid_v), valid_w + ) return is_interior benchmark(compute_interior_points_mask_slow, points_xyz, cuboid.vertices_m) @@ -433,7 +514,9 @@ def test_mat_to_xyz_constrained() -> None: xyz[0] = 0 # Set roll to zero. mat_constrained = xyz_to_mat(xyz) - xyz_expected = np.deg2rad([0, 0, 90]) # [45, 0, 90] -> constrain roll to zero -> [0, 0, 90]. + xyz_expected = np.deg2rad( + [0, 0, 90] + ) # [45, 0, 90] -> constrain roll to zero -> [0, 0, 90]. mat_constrained_expected = xyz_to_mat(xyz_expected) np.testing.assert_allclose(mat_constrained, mat_constrained_expected) @@ -510,7 +593,12 @@ def test_xyz_to_mat_vs_gtsam() -> None: def test_constrain_cuboid_pose() -> None: """Unit test to constrain cuboid pose.""" - path = Path(__file__).parent.resolve() / "data" / "b87683ae-14c5-321f-8af3-623e7bafc3a7" / "annotations.feather" + path = ( + Path(__file__).parent.resolve() + / "data" + / "b87683ae-14c5-321f-8af3-623e7bafc3a7" + / "annotations.feather" + ) cuboid_list = CuboidList.from_feather(path) for cuboid in cuboid_list.cuboids: pose = cuboid.dst_SE3_object diff --git a/tests/unit/geometry/test_interpolate.py b/tests/unit/geometry/test_interpolate.py index aad85044..45a42c33 100644 --- a/tests/unit/geometry/test_interpolate.py +++ b/tests/unit/geometry/test_interpolate.py @@ -91,8 +91,12 @@ def test_compute_lane_width_curved_width1() -> None: \\ ----- \\----- """ - left_even_pts: NDArrayFloat = np.array([[0, 2], [-2, 2], [-3, 1], [-3, 0], [-2, -1], [0, -1]]) - right_even_pts: NDArrayFloat = np.array([[0, 3], [-2, 3], [-4, 1], [-4, 0], [-2, -2], [0, -2]]) + left_even_pts: NDArrayFloat = np.array( + [[0, 2], [-2, 2], [-3, 1], [-3, 0], [-2, -1], [0, -1]] + ) + right_even_pts: NDArrayFloat = np.array( + [[0, 3], [-2, 3], [-4, 1], [-4, 0], [-2, -2], [0, -2]] + ) lane_width = interp_utils.compute_lane_width(left_even_pts, right_even_pts) gt_lane_width = 1.0 assert np.isclose(lane_width, gt_lane_width) @@ -115,9 +119,13 @@ def test_compute_lane_width_curved_not_width1() -> None: We get waypoint distances of [1,1,1,1,0.707..., 1,1] """ - left_even_pts: NDArrayFloat = np.array([[0, 2], [-2, 2], [-3, 1], [-3, 0], [-2.5, -0.5], [-2, -1], [0, -1]]) + left_even_pts: NDArrayFloat = np.array( + [[0, 2], [-2, 2], [-3, 1], [-3, 0], [-2.5, -0.5], [-2, -1], [0, -1]] + ) - right_even_pts: NDArrayFloat = np.array([[0, 3], [-2, 3], [-4, 1], [-4, 0], [-3, -1], [-2, -2], [0, -2]]) + right_even_pts: NDArrayFloat = np.array( + [[0, 3], [-2, 3], [-4, 1], [-4, 0], [-3, -1], [-2, -2], [0, -2]] + ) lane_width = interp_utils.compute_lane_width(left_even_pts, right_even_pts) gt_lane_width = 0.9581581115980783 @@ -236,7 +244,9 @@ def test_compute_mid_pivot_arc_5pt_cul_de_sac() -> None: # centerline_pts: Numpy array of shape (N,3) centerline_pts, lane_width = interp_utils.compute_mid_pivot_arc(single_pt, arc_pts) - gt_centerline_pts: NDArrayFloat = np.array([[0, 1], [0.5, 0.5], [1, 0], [0.5, -0.5], [0, -1]]) + gt_centerline_pts: NDArrayFloat = np.array( + [[0, 1], [0.5, 0.5], [1, 0], [0.5, -0.5], [0, -1]] + ) gt_lane_width = (2 + 2 + 2 + np.sqrt(2) + np.sqrt(2)) / 5 assert np.allclose(centerline_pts, gt_centerline_pts) assert np.isclose(lane_width, gt_lane_width) @@ -251,9 +261,13 @@ def test_compute_midpoint_line_cul_de_sac_right_onept() -> None: left_ln_bnds: NDArrayFloat = np.array([[0, 2], [1, 1], [2, 0], [1, -1], [0, -2]]) right_ln_bnds: NDArrayFloat = np.array([[0, 0]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=5) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=5 + ) - gt_centerline_pts: NDArrayFloat = np.array([[0, 1], [0.5, 0.5], [1, 0], [0.5, -0.5], [0, -1]]) + gt_centerline_pts: NDArrayFloat = np.array( + [[0, 1], [0.5, 0.5], [1, 0], [0.5, -0.5], [0, -1]] + ) gt_lane_width = (2 + 2 + 2 + np.sqrt(2) + np.sqrt(2)) / 5 assert np.allclose(centerline_pts, gt_centerline_pts) @@ -269,9 +283,13 @@ def test_compute_midpoint_line_cul_de_sac_left_onept() -> None: right_ln_bnds: NDArrayFloat = np.array([[0, 2], [1, 1], [2, 0], [1, -1], [0, -2]]) left_ln_bnds: NDArrayFloat = np.array([[0, 0]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=5) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=5 + ) - gt_centerline_pts: NDArrayFloat = np.array([[0, 1], [0.5, 0.5], [1, 0], [0.5, -0.5], [0, -1]]) + gt_centerline_pts: NDArrayFloat = np.array( + [[0, 1], [0.5, 0.5], [1, 0], [0.5, -0.5], [0, -1]] + ) gt_lane_width = (2 + 2 + 2 + np.sqrt(2) + np.sqrt(2)) / 5 assert np.allclose(centerline_pts, gt_centerline_pts) @@ -284,12 +302,18 @@ def test_compute_midpoint_line_straightline_maintain_5_waypts() -> None: Make sure that if we provide left and right boundary polylines in 2d, we can get the correct centerline by averaging left and right waypoints. """ - right_ln_bnds: NDArrayFloat = np.array([[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]]) + right_ln_bnds: NDArrayFloat = np.array( + [[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]] + ) left_ln_bnds: NDArrayFloat = np.array([[2, 4], [2, 2], [2, 0], [2, -2], [2, -4]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=5) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=5 + ) - gt_centerline_pts: NDArrayFloat = np.array([[0.5, 4], [0.5, 2], [0.5, 0], [0.5, -2], [0.5, -4]]) + gt_centerline_pts: NDArrayFloat = np.array( + [[0.5, 4], [0.5, 2], [0.5, 0], [0.5, -2], [0.5, -4]] + ) gt_lane_width = 3.0 assert np.allclose(centerline_pts, gt_centerline_pts) assert np.isclose(lane_width, gt_lane_width) @@ -301,12 +325,18 @@ def test_compute_midpoint_line_straightline_maintain_4_waypts() -> None: Make sure that if we provide left and right boundary polylines in 2d, we can get the correct centerline by averaging left and right waypoints. """ - right_ln_bnds: NDArrayFloat = np.array([[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]]) + right_ln_bnds: NDArrayFloat = np.array( + [[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]] + ) left_ln_bnds: NDArrayFloat = np.array([[2, 4], [2, 2], [2, 0], [2, -2], [2, -4]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=4) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=4 + ) - gt_centerline_pts: NDArrayFloat = np.array([[0.5, 4], [0.5, 4 / 3], [0.5, -4 / 3], [0.5, -4]]) + gt_centerline_pts: NDArrayFloat = np.array( + [[0.5, 4], [0.5, 4 / 3], [0.5, -4 / 3], [0.5, -4]] + ) gt_lane_width = 3.0 assert np.allclose(centerline_pts, gt_centerline_pts) assert np.isclose(lane_width, gt_lane_width) @@ -318,10 +348,14 @@ def test_compute_midpoint_line_straightline_maintain_3_waypts() -> None: Make sure that if we provide left and right boundary polylines in 2d, we can get the correct centerline by averaging left and right waypoints. """ - right_ln_bnds: NDArrayFloat = np.array([[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]]) + right_ln_bnds: NDArrayFloat = np.array( + [[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]] + ) left_ln_bnds: NDArrayFloat = np.array([[2, 4], [2, 2], [2, 0], [2, -2], [2, -4]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=3) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=3 + ) gt_centerline_pts: NDArrayFloat = np.array([[0.5, 4], [0.5, 0], [0.5, -4]]) gt_lane_width = 3.0 @@ -335,10 +369,14 @@ def test_compute_midpoint_line_straightline_maintain_2_waypts() -> None: Make sure that if we provide left and right boundary polylines in 2d, we can get the correct centerline by averaging left and right waypoints. """ - right_ln_bnds: NDArrayFloat = np.array([[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]]) + right_ln_bnds: NDArrayFloat = np.array( + [[-1, 4], [-1, 2], [-1, 0], [-1, -2], [-1, -4]] + ) left_ln_bnds: NDArrayFloat = np.array([[2, 4], [2, 2], [2, 0], [2, -2], [2, -4]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=2) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=2 + ) gt_centerline_pts: NDArrayFloat = np.array([[0.5, 4], [0.5, -4]]) gt_lane_width = 3.0 @@ -358,7 +396,9 @@ def test_compute_midpoint_line_curved_maintain_4_waypts() -> None: right_ln_bnds: NDArrayFloat = np.array([[-1, 3], [1, 3], [4, 0], [4, -2]]) left_ln_bnds: NDArrayFloat = np.array([[-1, 1], [1, 1], [2, 0], [2, -2]]) - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=4) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=4 + ) # from argoverse.utils.mpl_plotting_utils import draw_polygon_mpl @@ -403,7 +443,9 @@ def test_compute_midpoint_line_straightline_maintain_3_waypts_3dpolylines() -> N ) # fmt: on - centerline_pts, lane_width = interp_utils.compute_midpoint_line(left_ln_bnds, right_ln_bnds, num_interp_pts=3) + centerline_pts, lane_width = interp_utils.compute_midpoint_line( + left_ln_bnds, right_ln_bnds, num_interp_pts=3 + ) # fmt: off gt_centerline_pts: NDArrayFloat = np.array( [ @@ -558,14 +600,17 @@ def test_interpolate_pose() -> None: city_SE3_egot0 = SE3(rotation=np.eye(3), translation=np.array([5, 0, 0])) city_SE3_egot1 = SE3( - rotation=Rotation.from_euler("z", 90, degrees=True).as_matrix(), translation=np.array([0, 5, 0]) + rotation=Rotation.from_euler("z", 90, degrees=True).as_matrix(), + translation=np.array([0, 5, 0]), ) t0 = 0 t1 = 10 for query_timestamp in np.arange(11): pose = interp_utils.interpolate_pose( - key_timestamps=(t0, t1), key_poses=(city_SE3_egot0, city_SE3_egot1), query_timestamp=query_timestamp + key_timestamps=(t0, t1), + key_poses=(city_SE3_egot0, city_SE3_egot1), + query_timestamp=query_timestamp, ) if visualize: _plot_pose(pose) @@ -608,16 +653,22 @@ def test_linear_interpolation() -> None: X1: NDArrayFloat = np.array([-1, 2, 10], dtype=float) # at start of interval (@5 sec) - Xt_5 = interp_utils.linear_interpolation(key_timestamps=(5, 15), key_translations=(X0, X1), query_timestamp=5) + Xt_5 = interp_utils.linear_interpolation( + key_timestamps=(5, 15), key_translations=(X0, X1), query_timestamp=5 + ) expected_Xt_5: NDArrayFloat = np.array([1, 0, 0], dtype=float) assert np.array_equal(Xt_5, expected_Xt_5) # midway through interval (@10 sec) - Xt_10 = interp_utils.linear_interpolation(key_timestamps=(5, 15), key_translations=(X0, X1), query_timestamp=10) + Xt_10 = interp_utils.linear_interpolation( + key_timestamps=(5, 15), key_translations=(X0, X1), query_timestamp=10 + ) expected_Xt_10: NDArrayFloat = np.array([0, 1, 5], dtype=float) assert np.array_equal(Xt_10, expected_Xt_10) # at end of interval (@15 sec) - Xt_15 = interp_utils.linear_interpolation(key_timestamps=(5, 15), key_translations=(X0, X1), query_timestamp=15) + Xt_15 = interp_utils.linear_interpolation( + key_timestamps=(5, 15), key_translations=(X0, X1), query_timestamp=15 + ) expected_Xt_15: NDArrayFloat = np.array([-1, 2, 10], dtype=float) assert np.array_equal(Xt_15, expected_Xt_15) diff --git a/tests/unit/geometry/test_pinhole_camera.py b/tests/unit/geometry/test_pinhole_camera.py index 8a6a8c2c..c61c2ab7 100644 --- a/tests/unit/geometry/test_pinhole_camera.py +++ b/tests/unit/geometry/test_pinhole_camera.py @@ -30,12 +30,21 @@ def _create_pinhole_camera( translation: NDArrayFloat = np.zeros(3) ego_SE3_cam = SE3(rotation=rotation, translation=translation) - intrinsics = Intrinsics(fx_px=fx_px, fy_px=fy_px, cx_px=cx_px, cy_px=cy_px, width_px=width_px, height_px=height_px) + intrinsics = Intrinsics( + fx_px=fx_px, + fy_px=fy_px, + cx_px=cx_px, + cy_px=cy_px, + width_px=width_px, + height_px=height_px, + ) pinhole_camera = PinholeCamera(ego_SE3_cam, intrinsics, cam_name) return pinhole_camera -def _fit_plane_to_point_cloud(points_xyz: NDArrayFloat) -> Tuple[float, float, float, float]: +def _fit_plane_to_point_cloud( + points_xyz: NDArrayFloat, +) -> Tuple[float, float, float, float]: """Use SVD with at least 3 points to fit a plane. Args: @@ -45,7 +54,9 @@ def _fit_plane_to_point_cloud(points_xyz: NDArrayFloat) -> Tuple[float, float, f (4,) Plane coefficients. Defining ax + by + cz = d for the plane. """ center_xyz: NDArrayFloat = np.mean(points_xyz, axis=0) - out: Tuple[NDArrayFloat, NDArrayFloat, NDArrayFloat] = np.linalg.svd(points_xyz - center_xyz) + out: Tuple[NDArrayFloat, NDArrayFloat, NDArrayFloat] = np.linalg.svd( + points_xyz - center_xyz + ) vh = out[2] # Get the unitary normal vector @@ -63,8 +74,17 @@ def test_intrinsics_constructor() -> None: cx_px, cy_px = 1024, 775 - intrinsics = Intrinsics(fx_px=fx_px, fy_px=fy_px, cx_px=cx_px, cy_px=cy_px, width_px=width_px, height_px=height_px) - K_expected: NDArrayFloat = np.array(([1000, 0, 1024], [0, 1001, 775], [0, 0, 1]), dtype=np.float64) + intrinsics = Intrinsics( + fx_px=fx_px, + fy_px=fy_px, + cx_px=cx_px, + cy_px=cy_px, + width_px=width_px, + height_px=height_px, + ) + K_expected: NDArrayFloat = np.array( + ([1000, 0, 1024], [0, 1001, 775], [0, 0, 1]), dtype=np.float64 + ) assert np.array_equal(intrinsics.K, K_expected) @@ -93,7 +113,13 @@ def test_right_clipping_plane() -> None: fx_px = 10.0 width_px = 30 pinhole_camera = _create_pinhole_camera( - fx_px=fx_px, fy_px=0, cx_px=0, cy_px=0, height_px=30, width_px=width_px, cam_name="ring_front_center" + fx_px=fx_px, + fy_px=0, + cx_px=0, + cy_px=0, + height_px=30, + width_px=width_px, + cam_name="ring_front_center", ) right_plane = pinhole_camera.right_clipping_plane @@ -139,7 +165,13 @@ def test_left_clipping_plane() -> None: width_px = 30 pinhole_camera = _create_pinhole_camera( - fx_px=fx_px, fy_px=0, cx_px=0, cy_px=0, height_px=30, width_px=width_px, cam_name="ring_front_center" + fx_px=fx_px, + fy_px=0, + cx_px=0, + cy_px=0, + height_px=30, + width_px=width_px, + cam_name="ring_front_center", ) left_plane = pinhole_camera.left_clipping_plane @@ -181,7 +213,13 @@ def test_top_clipping_plane() -> None: fx_px = 10.0 height_px = 45 pinhole_camera = _create_pinhole_camera( - fx_px=fx_px, fy_px=0, cx_px=0, cy_px=0, height_px=height_px, width_px=1000, cam_name="ring_front_center" + fx_px=fx_px, + fy_px=0, + cx_px=0, + cy_px=0, + height_px=height_px, + width_px=1000, + cam_name="ring_front_center", ) top_plane = pinhole_camera.top_clipping_plane @@ -227,7 +265,13 @@ def test_bottom_clipping_plane() -> None: width_px = 10000 pinhole_camera = _create_pinhole_camera( - fx_px=fx_px, fy_px=1, cx_px=0, cy_px=0, height_px=height_px, width_px=width_px, cam_name="ring_front_center" + fx_px=fx_px, + fy_px=1, + cx_px=0, + cy_px=0, + height_px=height_px, + width_px=width_px, + cam_name="ring_front_center", ) bottom_plane = pinhole_camera.bottom_clipping_plane @@ -258,7 +302,13 @@ def test_form_near_clipping_plane() -> None: near_clip_dist = 30.0 pinhole_camera = _create_pinhole_camera( - fx_px=1, fy_px=0, cx_px=0, cy_px=0, height_px=30, width_px=width_px, cam_name="ring_front_center" + fx_px=1, + fy_px=0, + cx_px=0, + cy_px=0, + height_px=30, + width_px=width_px, + cam_name="ring_front_center", ) near_plane = pinhole_camera.near_clipping_plane(near_clip_dist) @@ -298,9 +348,21 @@ def test_frustum_planes_ring_cam() -> None: width_px = 2048 pinhole_camera = _create_pinhole_camera( - fx_px=fx_px, fy_px=fy_px, cx_px=cx_px, cy_px=cy_px, height_px=height_px, width_px=width_px, cam_name=camera_name + fx_px=fx_px, + fy_px=fy_px, + cx_px=cx_px, + cy_px=cy_px, + height_px=height_px, + width_px=width_px, + cam_name=camera_name, ) - left_plane, right_plane, near_plane, bottom_plane, top_plane = pinhole_camera.frustum_planes(near_clip_dist) + ( + left_plane, + right_plane, + near_plane, + bottom_plane, + top_plane, + ) = pinhole_camera.frustum_planes(near_clip_dist) left_plane_expected: NDArrayFloat = np.array([fx_px, 0.0, width_px / 2.0, 0.0]) right_plane_expected: NDArrayFloat = np.array([-fx_px, 0.0, width_px / 2.0, 0.0]) @@ -308,10 +370,18 @@ def test_frustum_planes_ring_cam() -> None: bottom_plane_expected: NDArrayFloat = np.array([0.0, -fx_px, height_px / 2.0, 0.0]) top_plane_expected: NDArrayFloat = np.array([0.0, fx_px, height_px / 2.0, 0.0]) - assert np.allclose(left_plane, left_plane_expected / np.linalg.norm(left_plane_expected)) - assert np.allclose(right_plane, right_plane_expected / np.linalg.norm(right_plane_expected)) - assert np.allclose(bottom_plane, bottom_plane_expected / np.linalg.norm(bottom_plane_expected)) - assert np.allclose(top_plane, top_plane_expected / np.linalg.norm(top_plane_expected)) + assert np.allclose( + left_plane, left_plane_expected / np.linalg.norm(left_plane_expected) + ) + assert np.allclose( + right_plane, right_plane_expected / np.linalg.norm(right_plane_expected) + ) + assert np.allclose( + bottom_plane, bottom_plane_expected / np.linalg.norm(bottom_plane_expected) + ) + assert np.allclose( + top_plane, top_plane_expected / np.linalg.norm(top_plane_expected) + ) assert np.allclose(near_plane, near_plane_expected) @@ -336,9 +406,21 @@ def test_generate_frustum_planes_stereo() -> None: width_px = 2048 pinhole_camera = _create_pinhole_camera( - fx_px=fx_px, fy_px=fy_px, cx_px=cx_px, cy_px=cy_px, height_px=height_px, width_px=width_px, cam_name=camera_name + fx_px=fx_px, + fy_px=fy_px, + cx_px=cx_px, + cy_px=cy_px, + height_px=height_px, + width_px=width_px, + cam_name=camera_name, ) - left_plane, right_plane, near_plane, bottom_plane, top_plane = pinhole_camera.frustum_planes(near_clip_dist) + ( + left_plane, + right_plane, + near_plane, + bottom_plane, + top_plane, + ) = pinhole_camera.frustum_planes(near_clip_dist) left_plane_expected: NDArrayFloat = np.array([fx_px, 0.0, width_px / 2.0, 0.0]) right_plane_expected: NDArrayFloat = np.array([-fx_px, 0.0, width_px / 2.0, 0.0]) @@ -346,10 +428,18 @@ def test_generate_frustum_planes_stereo() -> None: bottom_plane_expected: NDArrayFloat = np.array([0.0, -fx_px, height_px / 2.0, 0.0]) top_plane_expected: NDArrayFloat = np.array([0.0, fx_px, height_px / 2.0, 0.0]) - assert np.allclose(left_plane, left_plane_expected / np.linalg.norm(left_plane_expected)) - assert np.allclose(right_plane, right_plane_expected / np.linalg.norm(right_plane_expected)) - assert np.allclose(bottom_plane, bottom_plane_expected / np.linalg.norm(bottom_plane_expected)) - assert np.allclose(top_plane, top_plane_expected / np.linalg.norm(top_plane_expected)) + assert np.allclose( + left_plane, left_plane_expected / np.linalg.norm(left_plane_expected) + ) + assert np.allclose( + right_plane, right_plane_expected / np.linalg.norm(right_plane_expected) + ) + assert np.allclose( + bottom_plane, bottom_plane_expected / np.linalg.norm(bottom_plane_expected) + ) + assert np.allclose( + top_plane, top_plane_expected / np.linalg.norm(top_plane_expected) + ) assert np.allclose(near_plane, near_plane_expected) @@ -488,7 +578,9 @@ def test_compute_pixel_rays() -> None: assert np.allclose(gt_ray_dir, ray_dir) -def _compute_pixel_ray_direction(u: float, v: float, fx: float, fy: float, img_w: int, img_h: int) -> NDArrayFloat: +def _compute_pixel_ray_direction( + u: float, v: float, fx: float, fy: float, img_w: int, img_h: int +) -> NDArrayFloat: r"""Generate rays in the camera coordinate frame. Note: only used as a test utility. @@ -523,7 +615,9 @@ def _compute_pixel_ray_direction(u: float, v: float, fx: float, fy: float, img_w ValueError: If horizontal and vertical focal lengths are not close (within 1e-3). """ if not np.isclose(fx, fy, atol=1e-3): - raise ValueError(f"Focal lengths in the x and y directions must match: {fx} != {fy}") + raise ValueError( + f"Focal lengths in the x and y directions must match: {fx} != {fy}" + ) # approximation for principal point px = img_w / 2 @@ -591,9 +685,13 @@ def test_get_egovehicle_yaw_cam() -> None: for cam_enum in list(RingCameras): cam_name = cam_enum.value - pinhole_camera = PinholeCamera.from_feather(log_dir=sample_log_dir, cam_name=cam_name) + pinhole_camera = PinholeCamera.from_feather( + log_dir=sample_log_dir, cam_name=cam_name + ) ego_yaw_cam_deg = np.rad2deg(pinhole_camera.egovehicle_yaw_cam_rad) - assert np.isclose(ego_yaw_cam_deg, expected_ego_yaw_cam_deg_dict[cam_name], atol=0.1) + assert np.isclose( + ego_yaw_cam_deg, expected_ego_yaw_cam_deg_dict[cam_name], atol=0.1 + ) np.rad2deg(pinhole_camera.fov_theta_rad) diff --git a/tests/unit/geometry/test_utm.py b/tests/unit/geometry/test_utm.py index 45ed9960..eafbfcb3 100644 --- a/tests/unit/geometry/test_utm.py +++ b/tests/unit/geometry/test_utm.py @@ -19,7 +19,9 @@ def test_convert_city_coords_to_wgs84_atx() -> None: ] ) - wgs84_coords = geo_utils.convert_city_coords_to_wgs84(points_city, city_name=CityName.ATX) + wgs84_coords = geo_utils.convert_city_coords_to_wgs84( + points_city, city_name=CityName.ATX + ) expected_wgs84_coords: NDArrayFloat = np.array( [ @@ -40,7 +42,9 @@ def test_convert_city_coords_to_wgs84_wdc() -> None: ] ) - wgs84_coords = geo_utils.convert_city_coords_to_wgs84(points_city, city_name=CityName.WDC) + wgs84_coords = geo_utils.convert_city_coords_to_wgs84( + points_city, city_name=CityName.WDC + ) expected_wgs84_coords: NDArrayFloat = np.array( [ [38.9299801515994, -77.0168603173312], diff --git a/tests/unit/map/test_drivable_area.py b/tests/unit/map/test_drivable_area.py index 3d6b4a8c..56deb5d1 100644 --- a/tests/unit/map/test_drivable_area.py +++ b/tests/unit/map/test_drivable_area.py @@ -29,4 +29,6 @@ def test_from_dict(self) -> None: assert isinstance(drivable_area, DrivableArea) assert drivable_area.id == 4499430 - assert len(drivable_area.area_boundary) == 4 # first vertex is repeated as the last vertex + assert ( + len(drivable_area.area_boundary) == 4 + ) # first vertex is repeated as the last vertex diff --git a/tests/unit/map/test_lane_segment.py b/tests/unit/map/test_lane_segment.py index 9ee9997c..160b8de1 100644 --- a/tests/unit/map/test_lane_segment.py +++ b/tests/unit/map/test_lane_segment.py @@ -24,7 +24,10 @@ def test_from_dict(self) -> None: ], "left_lane_mark_type": "SOLID_YELLOW", "left_neighbor_id": None, - "right_lane_boundary": [{"x": 874.01, "y": -105.15, "z": -19.58}, {"x": 890.58, "y": -104.26, "z": -19.58}], + "right_lane_boundary": [ + {"x": 874.01, "y": -105.15, "z": -19.58}, + {"x": 890.58, "y": -104.26, "z": -19.58}, + ], "right_lane_mark_type": "SOLID_WHITE", "right_neighbor_id": 93269520, "predecessors": [], diff --git a/tests/unit/map/test_map_api.py b/tests/unit/map/test_map_api.py index 4935628c..aa07b9e6 100644 --- a/tests/unit/map/test_map_api.py +++ b/tests/unit/map/test_map_api.py @@ -29,7 +29,9 @@ def dummy_static_map(test_data_root_dir: Path) -> ArgoverseStaticMap: Static map instantiated from dummy test data. """ log_map_dirpath = ( - test_data_root_dir / "static_maps" / "dummy_log_map_gs1B8ZCv7DMi8cMt5aN5rSYjQidJXvGP__2020-07-21-Z1F0076" + test_data_root_dir + / "static_maps" + / "dummy_log_map_gs1B8ZCv7DMi8cMt5aN5rSYjQidJXvGP__2020-07-21-Z1F0076" ) return ArgoverseStaticMap.from_map_dir(log_map_dirpath, build_raster=True) @@ -46,7 +48,9 @@ def full_static_map(test_data_root_dir: Path) -> ArgoverseStaticMap: Static map instantiated from full test data. """ log_map_dirpath = ( - test_data_root_dir / "static_maps" / "full_log_map_gs1B8ZCv7DMi8cMt5aN5rSYjQidJXvGP__2020-07-21-Z1F0076" + test_data_root_dir + / "static_maps" + / "full_log_map_gs1B8ZCv7DMi8cMt5aN5rSYjQidJXvGP__2020-07-21-Z1F0076" ) return ArgoverseStaticMap.from_map_dir(log_map_dirpath, build_raster=True) @@ -56,7 +60,10 @@ class TestPolyline: def test_from_list(self) -> None: """Ensure object is generated correctly from a list of dictionaries.""" - points_dict_list = [{"x": 874.01, "y": -105.15, "z": -19.58}, {"x": 890.58, "y": -104.26, "z": -19.58}] + points_dict_list = [ + {"x": 874.01, "y": -105.15, "z": -19.58}, + {"x": 890.58, "y": -104.26, "z": -19.58}, + ] polyline = Polyline.from_json_data(points_dict_list) assert isinstance(polyline, Polyline) @@ -90,8 +97,14 @@ def test_from_dict(self) -> None: """Ensure object is generated correctly from a dictionary.""" json_data = { "id": 6310421, - "edge1": [{"x": 899.17, "y": -91.52, "z": -19.58}, {"x": 915.68, "y": -93.93, "z": -19.53}], - "edge2": [{"x": 899.44, "y": -95.37, "z": -19.48}, {"x": 918.25, "y": -98.05, "z": -19.4}], + "edge1": [ + {"x": 899.17, "y": -91.52, "z": -19.58}, + {"x": 915.68, "y": -93.93, "z": -19.53}, + ], + "edge2": [ + {"x": 899.44, "y": -95.37, "z": -19.48}, + {"x": 918.25, "y": -98.05, "z": -19.4}, + ], } pedestrian_crossing = PedestrianCrossing.from_dict(json_data) @@ -101,7 +114,9 @@ def test_from_dict(self) -> None: class TestArgoverseStaticMap: """Unit test for the Argoverse 2.0 per-log map.""" - def test_get_lane_segment_successor_ids(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_lane_segment_successor_ids( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure lane segment successors are fetched properly.""" lane_segment_id = 93269421 successor_ids = dummy_static_map.get_lane_segment_successor_ids(lane_segment_id) @@ -118,7 +133,9 @@ def test_get_lane_segment_successor_ids(self, dummy_static_map: ArgoverseStaticM expected_successor_ids = [93269526] assert successor_ids == expected_successor_ids - def test_lane_is_in_intersection(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_lane_is_in_intersection( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure the attribute describing if a lane segment is located with an intersection is fetched properly.""" lane_segment_id = 93269421 in_intersection = dummy_static_map.lane_is_in_intersection(lane_segment_id) @@ -135,44 +152,64 @@ def test_lane_is_in_intersection(self, dummy_static_map: ArgoverseStaticMap) -> assert isinstance(in_intersection, bool) assert not in_intersection - def test_get_lane_segment_left_neighbor_id(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_lane_segment_left_neighbor_id( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Test getting a lane segment id from the left neighbor.""" # Ensure id of lane segment (if any) that is the left neighbor to the query lane segment can be fetched properly lane_segment_id = 93269421 - l_neighbor_id = dummy_static_map.get_lane_segment_left_neighbor_id(lane_segment_id) + l_neighbor_id = dummy_static_map.get_lane_segment_left_neighbor_id( + lane_segment_id + ) assert l_neighbor_id is None lane_segment_id = 93269500 - l_neighbor_id = dummy_static_map.get_lane_segment_left_neighbor_id(lane_segment_id) + l_neighbor_id = dummy_static_map.get_lane_segment_left_neighbor_id( + lane_segment_id + ) assert l_neighbor_id is None lane_segment_id = 93269520 - l_neighbor_id = dummy_static_map.get_lane_segment_left_neighbor_id(lane_segment_id) + l_neighbor_id = dummy_static_map.get_lane_segment_left_neighbor_id( + lane_segment_id + ) assert l_neighbor_id == 93269421 - def test_get_lane_segment_right_neighbor_id(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_lane_segment_right_neighbor_id( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Test getting a lane segment id from the right neighbor.""" # Ensure id of lane segment (if any) that is the right neighbor to the query lane segment can be fetched lane_segment_id = 93269421 - r_neighbor_id = dummy_static_map.get_lane_segment_right_neighbor_id(lane_segment_id) + r_neighbor_id = dummy_static_map.get_lane_segment_right_neighbor_id( + lane_segment_id + ) assert r_neighbor_id == 93269520 lane_segment_id = 93269500 - r_neighbor_id = dummy_static_map.get_lane_segment_right_neighbor_id(lane_segment_id) + r_neighbor_id = dummy_static_map.get_lane_segment_right_neighbor_id( + lane_segment_id + ) assert r_neighbor_id == 93269526 lane_segment_id = 93269520 - r_neighbor_id = dummy_static_map.get_lane_segment_right_neighbor_id(lane_segment_id) + r_neighbor_id = dummy_static_map.get_lane_segment_right_neighbor_id( + lane_segment_id + ) assert r_neighbor_id == 93269458 - def test_get_scenario_lane_segment_ids(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_scenario_lane_segment_ids( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure ids of all lane segments in the local map can be fetched properly.""" lane_segment_ids = dummy_static_map.get_scenario_lane_segment_ids() expected_lane_segment_ids = [93269421, 93269500, 93269520] assert lane_segment_ids == expected_lane_segment_ids - def test_get_lane_segment_polygon(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_lane_segment_polygon( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure lane segment polygons are fetched properly.""" lane_segment_id = 93269421 @@ -191,7 +228,9 @@ def test_get_lane_segment_polygon(self, dummy_static_map: ArgoverseStaticMap) -> ) np.testing.assert_allclose(ls_polygon, expected_ls_polygon) - def test_get_lane_segment_centerline(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_lane_segment_centerline( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure lane segment centerlines can be inferred and fetched properly.""" lane_segment_id = 93269421 @@ -214,14 +253,18 @@ def test_get_lane_segment_centerline(self, dummy_static_map: ArgoverseStaticMap) ) np.testing.assert_allclose(centerline, expected_centerline) - def test_get_scenario_lane_segments(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_scenario_lane_segments( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure that all LaneSegment objects in the local map can be returned as a list.""" vector_lane_segments = dummy_static_map.get_scenario_lane_segments() assert isinstance(vector_lane_segments, list) assert all([isinstance(vls, LaneSegment) for vls in vector_lane_segments]) assert len(vector_lane_segments) == 3 - def test_get_scenario_ped_crossings(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_scenario_ped_crossings( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure that all PedCrossing objects in the local map can be returned as a list.""" ped_crossings = dummy_static_map.get_scenario_ped_crossings() assert isinstance(ped_crossings, list) @@ -260,9 +303,16 @@ def test_get_scenario_ped_crossings(self, dummy_static_map: ArgoverseStaticMap) ] # fmt: on assert len(ped_crossings) == len(expected_ped_crossings) - assert all([pc == expected_pc for pc, expected_pc in zip(ped_crossings, expected_ped_crossings)]) + assert all( + [ + pc == expected_pc + for pc, expected_pc in zip(ped_crossings, expected_ped_crossings) + ] + ) - def test_get_scenario_vector_drivable_areas(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_scenario_vector_drivable_areas( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure that drivable areas are loaded and formatted correctly.""" vector_das = dummy_static_map.get_scenario_vector_drivable_areas() assert isinstance(vector_das, list) @@ -286,21 +336,35 @@ def test_get_scenario_vector_drivable_areas(self, dummy_static_map: ArgoverseSta # fmt: on np.testing.assert_allclose(vector_da.xyz[:4], expected_first4_vertices) - def test_get_ground_height_at_xy(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_ground_height_at_xy( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure that ground height at (x,y) locations can be retrieved properly.""" point_cloud: NDArrayFloat = np.array( [ [770.6398, -105.8351, -19.4105], # ego-vehicle pose at one timestamp [943.5386, -49.6295, -19.3291], # ego-vehicle pose at one timestamp [918.0960, 82.5588, -20.5742], # ego-vehicle pose at one timestamp - [9999999, 999999, 0], # obviously out of bounds value for city coordinate system - [-999999, -999999, 0], # obviously out of bounds value for city coordinate system + [ + 9999999, + 999999, + 0, + ], # obviously out of bounds value for city coordinate system + [ + -999999, + -999999, + 0, + ], # obviously out of bounds value for city coordinate system ] ) assert dummy_static_map.raster_ground_height_layer is not None - ground_height_z = dummy_static_map.raster_ground_height_layer.get_ground_height_at_xy(point_cloud) + ground_height_z = ( + dummy_static_map.raster_ground_height_layer.get_ground_height_at_xy( + point_cloud + ) + ) assert ground_height_z.shape[0] == point_cloud.shape[0] assert ground_height_z.dtype == np.dtype(np.float64) @@ -310,17 +374,29 @@ def test_get_ground_height_at_xy(self, dummy_static_map: ArgoverseStaticMap) -> # based on grid resolution, ground should be within 7 centimeters of 30cm under back axle. expected_ground = point_cloud[:3, 2] - 0.30 - assert np.allclose(np.absolute(expected_ground - ground_height_z[:3]), 0, atol=0.07) + assert np.allclose( + np.absolute(expected_ground - ground_height_z[:3]), 0, atol=0.07 + ) - def test_get_ground_points_boolean(self, dummy_static_map: ArgoverseStaticMap) -> None: + def test_get_ground_points_boolean( + self, dummy_static_map: ArgoverseStaticMap + ) -> None: """Ensure that points close to the ground surface are correctly classified as `ground` category.""" point_cloud: NDArrayFloat = np.array( [ [770.6398, -105.8351, -19.4105], # ego-vehicle pose at one timestamp [943.5386, -49.6295, -19.3291], # ego-vehicle pose at one timestamp [918.0960, 82.5588, -20.5742], # ego-vehicle pose at one timestamp - [9999999, 999999, 0], # obviously out of bounds value for city coordinate system - [-999999, -999999, 0], # obviously out of bounds value for city coordinate system + [ + 9999999, + 999999, + 0, + ], # obviously out of bounds value for city coordinate system + [ + -999999, + -999999, + 0, + ], # obviously out of bounds value for city coordinate system ] ) @@ -330,7 +406,11 @@ def test_get_ground_points_boolean(self, dummy_static_map: ArgoverseStaticMap) - assert dummy_static_map.raster_ground_height_layer is not None - is_ground_pt = dummy_static_map.raster_ground_height_layer.get_ground_points_boolean(point_cloud) + is_ground_pt = ( + dummy_static_map.raster_ground_height_layer.get_ground_points_boolean( + point_cloud + ) + ) expected_is_ground_pt: NDArrayBool = np.array([True, True, True, False, False]) assert is_ground_pt.dtype == np.dtype(bool) assert np.array_equal(is_ground_pt, expected_is_ground_pt) @@ -340,7 +420,10 @@ def test_load_motion_forecasting_map(test_data_root_dir: Path) -> None: """Try to load a real map from the motion forecasting dataset.""" mf_scenario_id = "0a1e6f0a-1817-4a98-b02e-db8c9327d151" mf_scenario_map_path = ( - test_data_root_dir / "forecasting_scenarios" / mf_scenario_id / f"log_map_archive_{mf_scenario_id}.json" + test_data_root_dir + / "forecasting_scenarios" + / mf_scenario_id + / f"log_map_archive_{mf_scenario_id}.json" ) mf_map = ArgoverseStaticMap.from_json(mf_scenario_map_path) diff --git a/tests/unit/rendering/ops/test_draw.py b/tests/unit/rendering/ops/test_draw.py index 3916eef8..5b50ef91 100644 --- a/tests/unit/rendering/ops/test_draw.py +++ b/tests/unit/rendering/ops/test_draw.py @@ -7,12 +7,20 @@ import cv2 import numpy as np -from av2.rendering.ops.draw import alpha_blend_kernel, draw_points_kernel, gaussian_kernel +from av2.rendering.ops.draw import ( + alpha_blend_kernel, + draw_points_kernel, + gaussian_kernel, +) from av2.utils.typing import NDArrayByte, NDArrayInt def _draw_points_cv2( - img: NDArrayByte, points_xy: NDArrayInt, colors: NDArrayByte, radius: int, with_anti_alias: bool = True + img: NDArrayByte, + points_xy: NDArrayInt, + colors: NDArrayByte, + radius: int, + with_anti_alias: bool = True, ) -> NDArrayByte: """Draw points in an image using OpenCV functionality. @@ -81,7 +89,12 @@ def test_draw_points_kernel_3x3_antialiased() -> None: dtype=np.uint8, ) img = draw_points_kernel( - img=img, points_uv=points_xy, colors=colors, diameter=diameter, sigma=sigma, with_anti_alias=True + img=img, + points_uv=points_xy, + colors=colors, + diameter=diameter, + sigma=sigma, + with_anti_alias=True, ) assert np.array_equal(img, expected_img) @@ -105,7 +118,12 @@ def test_draw_points_kernel_9x9_aliased() -> None: dtype=np.uint8, ) img = draw_points_kernel( - img=img, points_uv=points_xy, colors=colors, diameter=diameter, sigma=sigma, with_anti_alias=False + img=img, + points_uv=points_xy, + colors=colors, + diameter=diameter, + sigma=sigma, + with_anti_alias=False, ) assert np.array_equal(img, expected_img) @@ -113,26 +131,42 @@ def test_draw_points_kernel_9x9_aliased() -> None: def test_benchmark_draw_points_kernel_aliased(benchmark: Callable[..., Any]) -> None: """Benchmark the draw points kernel _without_ anti-aliasing.""" img: NDArrayByte = np.zeros((2048, 2048, 3), dtype=np.uint8) - points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype(np.int64) - colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype(np.uint8) + points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype( + np.int64 + ) + colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype( + np.uint8 + ) diameter = 10 benchmark(draw_points_kernel, img, points_xy, colors, diameter) -def test_benchmark_draw_points_kernel_anti_aliased(benchmark: Callable[..., Any]) -> None: +def test_benchmark_draw_points_kernel_anti_aliased( + benchmark: Callable[..., Any] +) -> None: """Benchmark the draw points kernel _with_ anti-aliasing.""" img: NDArrayByte = np.zeros((2048, 2048, 3), dtype=np.uint8) - points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype(np.int64) - colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype(np.uint8) + points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype( + np.int64 + ) + colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype( + np.uint8 + ) diameter = 10 - benchmark(draw_points_kernel, img, points_xy, colors, diameter, with_anti_alias=True) + benchmark( + draw_points_kernel, img, points_xy, colors, diameter, with_anti_alias=True + ) def test_benchmark_draw_points_cv2_aliased(benchmark: Callable[..., Any]) -> None: """Benchmark the draw points method from OpenCV _without_ anti-aliasing.""" img: NDArrayByte = np.zeros((2048, 2048, 3), dtype=np.uint8) - points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype(np.int64) - colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype(np.uint8) + points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype( + np.int64 + ) + colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype( + np.uint8 + ) radius = 10 benchmark(_draw_points_cv2, img, points_xy, colors, radius, with_anti_alias=False) @@ -140,7 +174,11 @@ def test_benchmark_draw_points_cv2_aliased(benchmark: Callable[..., Any]) -> Non def test_benchmark_draw_points_cv2_anti_aliased(benchmark: Callable[..., Any]) -> None: """Benchmark the draw points method from OpenCV _with_ anti-aliasing.""" img: NDArrayByte = np.zeros((2048, 2048, 3), dtype=np.uint8) - points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype(np.int64) - colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype(np.uint8) + points_xy: NDArrayInt = np.random.randint(low=0, high=2048, size=(60000, 2)).astype( + np.int64 + ) + colors: NDArrayByte = np.random.randint(low=0, high=255, size=(60000, 3)).astype( + np.uint8 + ) radius = 10 benchmark(_draw_points_cv2, img, points_xy, colors, radius, with_anti_alias=True) diff --git a/tests/unit/rendering/test_color.py b/tests/unit/rendering/test_color.py index c30d07b5..0e22ce2e 100644 --- a/tests/unit/rendering/test_color.py +++ b/tests/unit/rendering/test_color.py @@ -10,7 +10,9 @@ def test_create_colormap() -> None: """Ensure we can create a red-to-green RGB colormap with values in [0,1].""" - colors_arr_rgb = color_utils.create_colormap(color_list=[RED_HEX, GREEN_HEX], n_colors=10) + colors_arr_rgb = color_utils.create_colormap( + color_list=[RED_HEX, GREEN_HEX], n_colors=10 + ) assert np.logical_and(0 <= colors_arr_rgb, colors_arr_rgb <= 1).all() assert colors_arr_rgb.shape == (10, 3) diff --git a/tests/unit/rendering/test_map.py b/tests/unit/rendering/test_map.py index 7d8735ae..2abbe275 100644 --- a/tests/unit/rendering/test_map.py +++ b/tests/unit/rendering/test_map.py @@ -14,7 +14,9 @@ def test_draw_visible_polyline_segments_cv2_some_visible() -> None: """Test rendering when one vertex is marked as occluded, and so two line segments are dropped out.""" visualize = False # 6 vertices in the polyline. - line_segments_arr: NDArrayInt = np.array([[50, 0], [50, 20], [50, 40], [50, 60], [50, 80], [60, 120]]) + line_segments_arr: NDArrayInt = np.array( + [[50, 0], [50, 20], [50, 40], [50, 60], [50, 80], [60, 120]] + ) valid_pts_bool: NDArrayBool = np.array([True, True, True, False, True, True]) img_bgr: NDArrayByte = np.zeros((100, 100, 3), dtype=np.uint8) @@ -34,7 +36,9 @@ def test_draw_visible_polyline_segments_cv2_all_visible() -> None: """Test rendering when all vertices are visible (and thus all line segments are visible).""" visualize = False # 6 vertices in the polyline. - line_segments_arr: NDArrayInt = np.array([[50, 0], [50, 20], [50, 40], [50, 60], [50, 80], [60, 120]]) + line_segments_arr: NDArrayInt = np.array( + [[50, 0], [50, 20], [50, 40], [50, 60], [50, 80], [60, 120]] + ) valid_pts_bool: NDArrayBool = np.array([True, True, True, True, True, True]) img_bgr: NDArrayByte = np.zeros((100, 100, 3), dtype=np.uint8) diff --git a/tests/unit/rendering/test_rasterize.py b/tests/unit/rendering/test_rasterize.py index dda574b3..1830466f 100644 --- a/tests/unit/rendering/test_rasterize.py +++ b/tests/unit/rendering/test_rasterize.py @@ -10,7 +10,9 @@ def _build_dummy_raster_inputs( n: int, d: int -) -> Tuple[NDArrayFloat, Tuple[float, float, float], Tuple[float, float, float], NDArrayFloat]: +) -> Tuple[ + NDArrayFloat, Tuple[float, float, float], Tuple[float, float, float], NDArrayFloat +]: """Build dummy inputs for the rasterize function. Args: diff --git a/tests/unit/rendering/test_video.py b/tests/unit/rendering/test_video.py index 18947f70..f76ec0f3 100644 --- a/tests/unit/rendering/test_video.py +++ b/tests/unit/rendering/test_video.py @@ -87,5 +87,7 @@ def test_crop_video_to_even_dims() -> None: save_fpath = Path(NamedTemporaryFile(suffix=".mp4").name) assert not save_fpath.exists() - video_utils.write_video(video=cropped_video, dst=save_fpath, fps=10, preset="medium") + video_utils.write_video( + video=cropped_video, dst=save_fpath, fps=10, preset="medium" + ) assert save_fpath.exists() diff --git a/tests/unit/structures/test_cuboid.py b/tests/unit/structures/test_cuboid.py index 158e8e90..f0843080 100644 --- a/tests/unit/structures/test_cuboid.py +++ b/tests/unit/structures/test_cuboid.py @@ -63,7 +63,9 @@ def test_compute_interior_points() -> None: ], dtype=float ) # fmt: on - expected_is_interior: NDArrayBool = np.array([False, True, True, True, True, True, False]) + expected_is_interior: NDArrayBool = np.array( + [False, True, True, True, True, True, False] + ) dst_SE3_object = SE3(rotation=np.eye(3), translation=np.array([2, 0, 0])) @@ -108,7 +110,14 @@ def _get_dummy_cuboid_list_params(num_cuboids: int) -> List[Cuboid]: """Create a cuboid list of length `num_cuboids`.""" cuboids: List[Cuboid] = [] for i in range(num_cuboids): - ego_SE3_object, length_m, width_m, height_m, category, timestamp_ns = _get_dummy_cuboid_params() + ( + ego_SE3_object, + length_m, + width_m, + height_m, + category, + timestamp_ns, + ) = _get_dummy_cuboid_params() cuboid = Cuboid( dst_SE3_object=ego_SE3_object, length_m=length_m + i, @@ -123,7 +132,14 @@ def _get_dummy_cuboid_list_params(num_cuboids: int) -> List[Cuboid]: def test_cuboid_constructor() -> None: """Test initializing a single cuboid.""" - ego_SE3_object, length_m, width_m, height_m, category, timestamp_ns = _get_dummy_cuboid_params() + ( + ego_SE3_object, + length_m, + width_m, + height_m, + category, + timestamp_ns, + ) = _get_dummy_cuboid_params() cuboid = Cuboid( dst_SE3_object=ego_SE3_object, length_m=length_m, @@ -229,7 +245,9 @@ def benchmark_transform(cuboids: List[Cuboid], target_SE3_ego: SE3) -> List[Cubo def test_benchmark_transform_list_comprehension(benchmark: Callable[..., Any]) -> None: """Benchmark cuboid transform with list comprehension.""" - def benchmark_transform_list_comprehension(cuboids: List[Cuboid], target_SE3_ego: SE3) -> List[Cuboid]: + def benchmark_transform_list_comprehension( + cuboids: List[Cuboid], target_SE3_ego: SE3 + ) -> List[Cuboid]: transformed_cuboids: List[Cuboid] = [ Cuboid( dst_SE3_object=target_SE3_ego.compose(cuboid.dst_SE3_object), @@ -245,7 +263,9 @@ def benchmark_transform_list_comprehension(cuboids: List[Cuboid], target_SE3_ego num_cuboids = 1000 cuboids = _get_dummy_cuboid_list_params(num_cuboids) - benchmark(benchmark_transform_list_comprehension, cuboids, cuboids[0].dst_SE3_object) + benchmark( + benchmark_transform_list_comprehension, cuboids, cuboids[0].dst_SE3_object + ) if __name__ == "__main__": diff --git a/tests/unit/structures/test_ndgrid.py b/tests/unit/structures/test_ndgrid.py index 97068142..d32aa5da 100644 --- a/tests/unit/structures/test_ndgrid.py +++ b/tests/unit/structures/test_ndgrid.py @@ -52,7 +52,10 @@ def test_bev_grid() -> None: grid_coordinates_expected: NDArrayInt = np.array([61, 72], dtype=int) values_expected = GRAY_BGR - assert np.array_equal(bev_img[grid_coordinates_expected[1], grid_coordinates_expected[0]], values_expected) + assert np.array_equal( + bev_img[grid_coordinates_expected[1], grid_coordinates_expected[0]], + values_expected, + ) def test_BEVGrid_non_integer_multiple() -> None: diff --git a/tests/unit/structures/test_sweep.py b/tests/unit/structures/test_sweep.py index ebb1c3ce..3b5fb42a 100644 --- a/tests/unit/structures/test_sweep.py +++ b/tests/unit/structures/test_sweep.py @@ -14,13 +14,22 @@ @pytest.fixture def dummy_sweep(test_data_root_dir: Path) -> Sweep: """Get a fake sweep containing two points.""" - path = test_data_root_dir / "sensor_dataset_logs" / "dummy" / "sensors" / "lidar" / "315968663259918000.feather" + path = ( + test_data_root_dir + / "sensor_dataset_logs" + / "dummy" + / "sensors" + / "lidar" + / "315968663259918000.feather" + ) return Sweep.from_feather(path) def test_sweep_from_feather(dummy_sweep: Sweep) -> None: """Test loading a sweep from a feather file.""" - xyz_expected: NDArrayFloat = np.array([[-22.1875, 20.484375, 0.55029296875], [-20.609375, 19.1875, 1.30078125]]) + xyz_expected: NDArrayFloat = np.array( + [[-22.1875, 20.484375, 0.55029296875], [-20.609375, 19.1875, 1.30078125]] + ) intensity_expected: NDArrayByte = np.array([38, 5], dtype=np.uint8) laser_number_expected: NDArrayByte = np.array([19, 3], dtype=np.uint8) offset_ns_expected: NDArrayInt = np.array([253440, 283392], dtype=np.int32) diff --git a/tests/unit/torch/data_loaders/test_detection_dataloader.py b/tests/unit/torch/data_loaders/test_detection_dataloader.py index 8ca831d9..a7cf547b 100644 --- a/tests/unit/torch/data_loaders/test_detection_dataloader.py +++ b/tests/unit/torch/data_loaders/test_detection_dataloader.py @@ -5,7 +5,9 @@ from av2.torch.data_loaders.detection import DetectionDataLoader -TEST_DATA_DIR: Final = Path(__file__).parent.parent.parent.resolve() / "test_data" / "sensor_dataset_logs" +TEST_DATA_DIR: Final = ( + Path(__file__).parent.parent.parent.resolve() / "test_data" / "sensor_dataset_logs" +) def test_build_data_loader() -> None: @@ -13,6 +15,8 @@ def test_build_data_loader() -> None: root_dir = TEST_DATA_DIR dataset_name = "av2" split_name = "val" - data_loader = DetectionDataLoader(root_dir=root_dir, dataset_name=dataset_name, split_name=split_name) + data_loader = DetectionDataLoader( + root_dir=root_dir, dataset_name=dataset_name, split_name=split_name + ) for datum in data_loader: assert datum is not None diff --git a/tests/unit/torch/data_loaders/test_scene_flow_dataloader.py b/tests/unit/torch/data_loaders/test_scene_flow_dataloader.py index 247fcd4e..ed481b25 100644 --- a/tests/unit/torch/data_loaders/test_scene_flow_dataloader.py +++ b/tests/unit/torch/data_loaders/test_scene_flow_dataloader.py @@ -19,7 +19,9 @@ def test_scene_flow_dataloader() -> None: The computed flow should check the visually confirmed labels in flow_labels.feather. """ - dl_test = av2.torch.data_loaders.scene_flow.SceneFlowDataloader(_TEST_DATA_ROOT, "test_data", "test") + dl_test = av2.torch.data_loaders.scene_flow.SceneFlowDataloader( + _TEST_DATA_ROOT, "test_data", "test" + ) sweep_0, sweep_1, ego, not_flow = dl_test[0] assert not_flow is None rust_sweep = dl_test._backend.get(0) @@ -40,7 +42,9 @@ def test_scene_flow_dataloader() -> None: failed = True assert failed - data_loader = av2.torch.data_loaders.scene_flow.SceneFlowDataloader(_TEST_DATA_ROOT, "test_data", "val") + data_loader = av2.torch.data_loaders.scene_flow.SceneFlowDataloader( + _TEST_DATA_ROOT, "test_data", "val" + ) assert len(data_loader) == 1 assert data_loader.get_log_id(0) == "7fab2350-7eaf-3b7e-a39d-6937a4c1bede" @@ -51,7 +55,9 @@ def test_scene_flow_dataloader() -> None: flow: Flow = maybe_flow assert len(flow) == len(sweep_0.lidar.as_tensor()) - log_dir = _TEST_DATA_ROOT / "test_data/sensor/val/7fab2350-7eaf-3b7e-a39d-6937a4c1bede" + log_dir = ( + _TEST_DATA_ROOT / "test_data/sensor/val/7fab2350-7eaf-3b7e-a39d-6937a4c1bede" + ) flow_labels = pd.read_feather(log_dir / "flow_labels.feather") FLOW_COLS = ["flow_tx_m", "flow_ty_m", "flow_tz_m"] @@ -65,7 +71,9 @@ def test_scene_flow_dataloader() -> None: assert np.allclose(flow.category_indices.numpy(), flow_labels.classes.to_numpy()) assert np.allclose(flow.is_dynamic.numpy(), flow_labels.dynamic.to_numpy()) assert sweep_0.is_ground is not None - ground_match: NDArrayBool = sweep_0.is_ground.numpy() == flow_labels.is_ground_0.to_numpy() + ground_match: NDArrayBool = ( + sweep_0.is_ground.numpy() == flow_labels.is_ground_0.to_numpy() + ) assert np.logical_not(ground_match).sum() < 10 gt_ego = np.load(log_dir / "ego_motion.npz") diff --git a/tests/unit/torch/structures/test_cuboids.py b/tests/unit/torch/structures/test_cuboids.py index 1d2dae05..ae2ed18c 100644 --- a/tests/unit/torch/structures/test_cuboids.py +++ b/tests/unit/torch/structures/test_cuboids.py @@ -13,14 +13,18 @@ from av2.torch.structures.cuboids import CuboidMode, Cuboids TEST_DATA_DIR: Final = Path(__file__).parent.parent.parent.resolve() / "test_data" -SAMPLE_LOG_DIR: Final = TEST_DATA_DIR / "sensor_dataset_logs" / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" +SAMPLE_LOG_DIR: Final = ( + TEST_DATA_DIR / "sensor_dataset_logs" / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" +) def test_build_cuboids() -> None: """Test building the Cuboids structure.""" annotations_path = SAMPLE_LOG_DIR / "annotations.feather" annotations_frame = pd.read_feather(annotations_path) - cuboids_npy = annotations_frame[list(XYZLWH_QWXYZ_COLUMNS)].to_numpy().astype(np.float32) + cuboids_npy = ( + annotations_frame[list(XYZLWH_QWXYZ_COLUMNS)].to_numpy().astype(np.float32) + ) cuboids = Cuboids(annotations_frame) cuboids_xyzlwht = cuboids.as_tensor() @@ -31,7 +35,9 @@ def test_build_cuboids() -> None: w, x, y, z = cuboids_xyzlwh_qwxyz[:, 6:10].t() _, _, yaw = euler_from_quaternion(w, x, y, z) assert_close(cuboids_xyzlwht[:, 6], yaw) - assert_close(cuboids.as_tensor(cuboid_mode=CuboidMode.XYZLWH_QWXYZ), cuboids_xyzlwh_qwxyz) + assert_close( + cuboids.as_tensor(cuboid_mode=CuboidMode.XYZLWH_QWXYZ), cuboids_xyzlwh_qwxyz + ) track_uuid_expected: List[str] = annotations_frame["track_uuid"].to_list() assert cuboids.track_uuid == track_uuid_expected diff --git a/tests/unit/torch/structures/test_lidar.py b/tests/unit/torch/structures/test_lidar.py index 96be6e92..2e7604a9 100644 --- a/tests/unit/torch/structures/test_lidar.py +++ b/tests/unit/torch/structures/test_lidar.py @@ -12,7 +12,9 @@ from av2.torch.structures.lidar import Lidar TEST_DATA_DIR: Final = Path(__file__).parent.parent.parent.resolve() / "test_data" -SAMPLE_LOG_DIR: Final = TEST_DATA_DIR / "sensor_dataset_logs" / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" +SAMPLE_LOG_DIR: Final = ( + TEST_DATA_DIR / "sensor_dataset_logs" / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" +) def test_build_lidar() -> None: @@ -20,7 +22,9 @@ def test_build_lidar() -> None: lidar_paths = sorted((SAMPLE_LOG_DIR / "sensors" / "lidar").glob("*.feather")) lidar_path = lidar_paths[0] frame = pd.read_feather(lidar_path) - lidar_tensor = torch.as_tensor(frame[list(LIDAR_COLUMNS)].to_numpy().astype(np.float32)) + lidar_tensor = torch.as_tensor( + frame[list(LIDAR_COLUMNS)].to_numpy().astype(np.float32) + ) lidar = Lidar(frame) assert_close(lidar.as_tensor(), lidar_tensor) diff --git a/tests/unit/torch/structures/test_sweep.py b/tests/unit/torch/structures/test_sweep.py index e44fe304..758f7c48 100644 --- a/tests/unit/torch/structures/test_sweep.py +++ b/tests/unit/torch/structures/test_sweep.py @@ -11,7 +11,9 @@ from av2.torch.structures.utils import SE3_from_frame TEST_DATA_DIR: Final = Path(__file__).parent.parent.parent.resolve() / "test_data" -SAMPLE_LOG_DIR: Final = TEST_DATA_DIR / "sensor_dataset_logs" / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" +SAMPLE_LOG_DIR: Final = ( + TEST_DATA_DIR / "sensor_dataset_logs" / "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" +) def test_build_sweep() -> None: @@ -30,5 +32,7 @@ def test_build_sweep() -> None: city_SE3_ego = SE3_from_frame(city_pose_frame) sweep_uuid = annotations_path.parent.stem, int(lidar_path.stem) - sweep = Sweep(city_SE3_ego=city_SE3_ego, lidar=lidar, sweep_uuid=sweep_uuid, cuboids=cuboids) + sweep = Sweep( + city_SE3_ego=city_SE3_ego, lidar=lidar, sweep_uuid=sweep_uuid, cuboids=cuboids + ) assert sweep is not None diff --git a/tests/unit/torch/structures/test_utils.py b/tests/unit/torch/structures/test_utils.py index baaa6a2d..43a8712c 100644 --- a/tests/unit/torch/structures/test_utils.py +++ b/tests/unit/torch/structures/test_utils.py @@ -33,7 +33,14 @@ def test_tensor_from_frame() -> None: tensor = tensor_from_frame(frame, columns=["qw", "qx", "qy", "qz"]) tensor_expected = torch.as_tensor( - [[frame.loc[0, "qw"], frame.loc[0, "qx"], frame.loc[0, "qy"], frame.loc[0, "qz"]]] + [ + [ + frame.loc[0, "qw"], + frame.loc[0, "qx"], + frame.loc[0, "qy"], + frame.loc[0, "qz"], + ] + ] ) assert_close(tensor, tensor_expected) @@ -42,12 +49,18 @@ def test_SE3_from_frame() -> None: """Test converting a data-frame into an SE(3) object.""" frame = _build_dummy_frame() - quat_wxyz_tensor = torch.as_tensor(frame[list(QWXYZ_COLUMNS)].to_numpy().astype(np.float32)) - translation = torch.as_tensor(frame[list(TRANSLATION_COLUMNS)].to_numpy().astype(np.float32)) + quat_wxyz_tensor = torch.as_tensor( + frame[list(QWXYZ_COLUMNS)].to_numpy().astype(np.float32) + ) + translation = torch.as_tensor( + frame[list(TRANSLATION_COLUMNS)].to_numpy().astype(np.float32) + ) quat_wxyz = Quaternion(quat_wxyz_tensor) rotation = So3(quat_wxyz) city_SE3_ego_expected = Se3(rotation, translation) city_SE3_ego = SE3_from_frame(frame) assert_close(city_SE3_ego.translation, city_SE3_ego_expected.translation) - assert_close(city_SE3_ego.rotation.matrix(), city_SE3_ego_expected.rotation.matrix()) + assert_close( + city_SE3_ego.rotation.matrix(), city_SE3_ego_expected.rotation.matrix() + ) diff --git a/tests/unit/utils/test_dense_grid_interpolation.py b/tests/unit/utils/test_dense_grid_interpolation.py index 3b36023a..dc2148dd 100644 --- a/tests/unit/utils/test_dense_grid_interpolation.py +++ b/tests/unit/utils/test_dense_grid_interpolation.py @@ -25,7 +25,12 @@ def test_interp_dense_grid_from_sparse_insufficient_points_simplex() -> None: grid_w = 10 bev_img_interp = dense_grid_interpolation.interp_dense_grid_from_sparse( - grid_img=bev_img, points=points, values=rgb_values, grid_h=grid_h, grid_w=grid_w, interp_method="linear" + grid_img=bev_img, + points=points, + values=rgb_values, + grid_h=grid_h, + grid_w=grid_w, + interp_method="linear", ) assert np.allclose(bev_img_interp, np.zeros((10, 10, 3), dtype=np.uint8)) @@ -44,12 +49,19 @@ def test_interp_dense_grid_from_sparse_byte() -> None: # provided as (x,y) tuples points: NDArrayInt = np.array([[0, 0], [0, 3], [3, 3], [3, 0]]) - rgb_values: NDArrayByte = np.array([RED_RGB, GREEN_RGB, BLUE_RGB, RED_RGB], dtype=np.uint8) + rgb_values: NDArrayByte = np.array( + [RED_RGB, GREEN_RGB, BLUE_RGB, RED_RGB], dtype=np.uint8 + ) grid_h = 4 grid_w = 4 bev_img_interp = dense_grid_interpolation.interp_dense_grid_from_sparse( - grid_img=bev_img, points=points, values=rgb_values, grid_h=grid_h, grid_w=grid_w, interp_method="linear" + grid_img=bev_img, + points=points, + values=rgb_values, + grid_h=grid_h, + grid_w=grid_w, + interp_method="linear", ) assert bev_img_interp.dtype == np.dtype(np.uint8) @@ -77,12 +89,19 @@ def test_interp_dense_grid_from_sparse_float() -> None: # provided as (x,y) tuples points: NDArrayInt = np.array([[0, 0], [0, 3], [3, 3], [3, 0]], dtype=int) - rgb_values: NDArrayFloat = np.array([RED_RGB, GREEN_RGB, BLUE_RGB, RED_RGB], dtype=float) + rgb_values: NDArrayFloat = np.array( + [RED_RGB, GREEN_RGB, BLUE_RGB, RED_RGB], dtype=float + ) grid_h = 4 grid_w = 4 bev_img_interp = dense_grid_interpolation.interp_dense_grid_from_sparse( - grid_img=bev_img, points=points, values=rgb_values, grid_h=grid_h, grid_w=grid_w, interp_method="linear" + grid_img=bev_img, + points=points, + values=rgb_values, + grid_h=grid_h, + grid_w=grid_w, + interp_method="linear", ) assert bev_img_interp.dtype == np.dtype(np.float64) diff --git a/tests/unit/utils/test_depth_map_utils.py b/tests/unit/utils/test_depth_map_utils.py index d0a0390a..2f2e6f91 100644 --- a/tests/unit/utils/test_depth_map_utils.py +++ b/tests/unit/utils/test_depth_map_utils.py @@ -18,14 +18,20 @@ def test_vis_depth_map() -> None: img_rgb: NDArrayByte = np.zeros((H, W, 3), dtype=np.uint8) img_rgb[:, :, 0] = 255 # channels will be (255,0,0) for red. - depth_map: NDArrayFloat = np.arange(H * W).reshape(H, W).astype(np.float32) / (H * W) * 255 + depth_map: NDArrayFloat = ( + np.arange(H * W).reshape(H, W).astype(np.float32) / (H * W) * 255 + ) - depth_map_utils.vis_depth_map(img_rgb=img_rgb, depth_map=depth_map, interp_depth_map=True) + depth_map_utils.vis_depth_map( + img_rgb=img_rgb, depth_map=depth_map, interp_depth_map=True + ) if visualize: plt.show() plt.close("all") - depth_map_utils.vis_depth_map(img_rgb=img_rgb, depth_map=depth_map, interp_depth_map=False) + depth_map_utils.vis_depth_map( + img_rgb=img_rgb, depth_map=depth_map, interp_depth_map=False + ) if visualize: plt.show() plt.close("all") diff --git a/tests/unit/utils/test_infinity_norm_utils.py b/tests/unit/utils/test_infinity_norm_utils.py index 369575f0..e2497fbd 100644 --- a/tests/unit/utils/test_infinity_norm_utils.py +++ b/tests/unit/utils/test_infinity_norm_utils.py @@ -18,7 +18,9 @@ def test_has_pts_in_infinity_norm_radius1() -> None: [5.1, 5.1] ]) # fmt: on - within = infinity_norm_utils.has_pts_in_infinity_norm_radius(pts, window_center=np.zeros(2), window_sz=5) + within = infinity_norm_utils.has_pts_in_infinity_norm_radius( + pts, window_center=np.zeros(2), window_sz=5 + ) assert not within @@ -32,7 +34,9 @@ def test_has_pts_in_infinity_norm_radius2() -> None: [5.1, 5.1] ]) # fmt: on - within = infinity_norm_utils.has_pts_in_infinity_norm_radius(pts, window_center=np.zeros(2), window_sz=5) + within = infinity_norm_utils.has_pts_in_infinity_norm_radius( + pts, window_center=np.zeros(2), window_sz=5 + ) assert within @@ -46,12 +50,16 @@ def test_has_pts_in_infinity_norm_radius3() -> None: [4.9, 4.9] ]) # fmt: on - within = infinity_norm_utils.has_pts_in_infinity_norm_radius(pts, window_center=np.zeros(2), window_sz=5) + within = infinity_norm_utils.has_pts_in_infinity_norm_radius( + pts, window_center=np.zeros(2), window_sz=5 + ) assert within def test_has_pts_in_infinity_norm_radius4() -> None: """All pts within radius.""" pts: NDArrayFloat = np.array([[4.9, 4.9]]) - within = infinity_norm_utils.has_pts_in_infinity_norm_radius(pts, window_center=np.zeros(2), window_sz=5) + within = infinity_norm_utils.has_pts_in_infinity_norm_radius( + pts, window_center=np.zeros(2), window_sz=5 + ) assert within diff --git a/tests/unit/utils/test_io.py b/tests/unit/utils/test_io.py index be66b96f..ecfd8f3d 100644 --- a/tests/unit/utils/test_io.py +++ b/tests/unit/utils/test_io.py @@ -12,7 +12,13 @@ def test_read_feather(test_data_root_dir: Path) -> None: """Read an Apache Feather file.""" - feather_path = test_data_root_dir / "sensor_dataset_logs" / "test_log" / "calibration" / "intrinsics.feather" + feather_path = ( + test_data_root_dir + / "sensor_dataset_logs" + / "test_log" + / "calibration" + / "intrinsics.feather" + ) feather_file = read_feather(feather_path) assert feather_file is not None @@ -31,7 +37,12 @@ def test_read_ego_SE3_sensor(test_data_root_dir: Path) -> None: def test_read_lidar_sweep() -> None: """Read 3d point coordinates from a LiDAR sweep file from an example log.""" log_id = "adcf7d18-0510-35b0-a2fa-b4cea13a6d76" - EXAMPLE_LOG_DATA_ROOT = Path(__file__).resolve().parent.parent / "test_data" / "sensor_dataset_logs" / log_id + EXAMPLE_LOG_DATA_ROOT = ( + Path(__file__).resolve().parent.parent + / "test_data" + / "sensor_dataset_logs" + / log_id + ) fpath = EXAMPLE_LOG_DATA_ROOT / "sensors" / "lidar" / "315973157959879000.feather" arr = io_utils.read_lidar_sweep(fpath, attrib_spec="xyz") diff --git a/tests/unit/utils/test_mesh_grid.py b/tests/unit/utils/test_mesh_grid.py index 731aa881..2d01e09d 100644 --- a/tests/unit/utils/test_mesh_grid.py +++ b/tests/unit/utils/test_mesh_grid.py @@ -18,7 +18,9 @@ def test_get_mesh_grid_as_point_cloud_3x3square() -> None: max_y = 4 # integer, maximum y-coordinate of 2D grid # return pts, a Numpy array of shape (N,2) - pts = mesh_grid_utils.get_mesh_grid_as_point_cloud(min_x, max_x, min_y, max_y, downsample_factor=1.0) + pts = mesh_grid_utils.get_mesh_grid_as_point_cloud( + min_x, max_x, min_y, max_y, downsample_factor=1.0 + ) assert pts.shape == (9, 2) gt_pts: NDArrayFloat = np.array( @@ -46,7 +48,9 @@ def test_get_mesh_grid_as_point_cloud_3x2rect() -> None: max_y = 3 # integer, maximum y-coordinate of 2D grid # return pts, a Numpy array of shape (N,2) - pts = mesh_grid_utils.get_mesh_grid_as_point_cloud(min_x, max_x, min_y, max_y, downsample_factor=1.0) + pts = mesh_grid_utils.get_mesh_grid_as_point_cloud( + min_x, max_x, min_y, max_y, downsample_factor=1.0 + ) assert pts.shape == (6, 2) # fmt: off @@ -71,7 +75,9 @@ def test_get_mesh_grid_as_point_cloud_single_pt() -> None: max_y = 2 # integer, maximum y-coordinate of 2D grid # return pts, a Numpy array of shape (N,2) - pts = mesh_grid_utils.get_mesh_grid_as_point_cloud(min_x, max_x, min_y, max_y, downsample_factor=1.0) + pts = mesh_grid_utils.get_mesh_grid_as_point_cloud( + min_x, max_x, min_y, max_y, downsample_factor=1.0 + ) assert pts.shape == (1, 2) gt_pts: NDArrayFloat = np.array([[-3.0, 2.0]]) @@ -87,7 +93,9 @@ def test_get_mesh_grid_as_point_cloud_downsample() -> None: max_y = 5 # integer, maximum y-coordinate of 2D grid # return pts, a Numpy array of shape (N,2) - pts = mesh_grid_utils.get_mesh_grid_as_point_cloud(min_x, max_x, min_y, max_y, downsample_factor=3.0) + pts = mesh_grid_utils.get_mesh_grid_as_point_cloud( + min_x, max_x, min_y, max_y, downsample_factor=3.0 + ) assert pts.shape == (4, 2) diff --git a/tests/unit/utils/test_polyline_utils.py b/tests/unit/utils/test_polyline_utils.py index e103c47f..9f82d55e 100644 --- a/tests/unit/utils/test_polyline_utils.py +++ b/tests/unit/utils/test_polyline_utils.py @@ -22,7 +22,9 @@ def test_convert_lane_boundaries_to_polygon_3d() -> None: [14, -1, 8] ]) # fmt: on - polygon = polyline_utils.convert_lane_boundaries_to_polygon(right_ln_bnd, left_ln_bnd) + polygon = polyline_utils.convert_lane_boundaries_to_polygon( + right_ln_bnd, left_ln_bnd + ) # fmt: off gt_polygon: NDArrayFloat = np.array( diff --git a/tests/unit/utils/test_raster.py b/tests/unit/utils/test_raster.py index 8a8d929b..a3d64327 100644 --- a/tests/unit/utils/test_raster.py +++ b/tests/unit/utils/test_raster.py @@ -30,7 +30,9 @@ def test_get_mask_from_polygon() -> None: ] ) # fmt: on - mask = raster_utils.get_mask_from_polygons(polygons=[triangle, rectangle], img_h=7, img_w=7) + mask = raster_utils.get_mask_from_polygons( + polygons=[triangle, rectangle], img_h=7, img_w=7 + ) # fmt: off expected_mask: NDArrayByte = np.array( @@ -73,7 +75,9 @@ def test_get_mask_from_polygon_repeated_coords() -> None: ] ) # fmt: on - mask = raster_utils.get_mask_from_polygons(polygons=[triangle, rectangle], img_h=7, img_w=7) + mask = raster_utils.get_mask_from_polygons( + polygons=[triangle, rectangle], img_h=7, img_w=7 + ) # fmt: off expected_mask: NDArrayByte = np.array( @@ -124,7 +128,9 @@ def test_get_mask_from_polygon_coords_out_of_bounds() -> None: def test_benchmark_blend_images_cv2(benchmark: Callable[..., Any]) -> None: """Benchmark opencv implementation of alpha blending.""" - def blend_images(img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7) -> NDArrayByte: + def blend_images( + img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7 + ) -> NDArrayByte: """Alpha-blend two images together using OpenCV `addWeighted`. Args: @@ -147,7 +153,9 @@ def blend_images(img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7) -> ND def test_benchmark_blend_images_npy(benchmark: Callable[..., Any]) -> None: """Benchmark numpy implementation of alpha blending.""" - def blend_images(img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7) -> NDArrayByte: + def blend_images( + img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7 + ) -> NDArrayByte: """Alpha-blend two images together using `numpy`. Args: @@ -158,7 +166,9 @@ def blend_images(img0: NDArrayByte, img1: NDArrayByte, alpha: float = 0.7) -> ND Returns: uint8 array of shape (H,W,3) """ - blended: NDArrayFloat = np.multiply(img0.astype(np.float32), alpha, dtype=float) + np.multiply( + blended: NDArrayFloat = np.multiply( + img0.astype(np.float32), alpha, dtype=float + ) + np.multiply( img1.astype(np.float32), (1 - alpha), dtype=float, @@ -187,6 +197,8 @@ def test_blend_images() -> None: img_b: NDArrayByte = np.zeros((H, W, 3), dtype=np.uint8) img_b[:, :, :3] = np.array([2, 4, 8]) # column 0 has uniform intensity - blended_img_expected: NDArrayByte = np.round(img_a * alpha + img_b * beta + gamma).astype(np.uint8) + blended_img_expected: NDArrayByte = np.round( + img_a * alpha + img_b * beta + gamma + ).astype(np.uint8) blended_img = raster_utils.blend_images(img_a, img_b, alpha=alpha) assert np.array_equal(blended_img, blended_img_expected) diff --git a/tests/unit/utils/test_se3.py b/tests/unit/utils/test_se3.py index f442209e..94737572 100644 --- a/tests/unit/utils/test_se3.py +++ b/tests/unit/utils/test_se3.py @@ -53,7 +53,9 @@ def test_SE3_transform_point_cloud_identity() -> None: def test_SE3_transform_point_cloud_by_quaternion() -> None: """Test rotating points by a given quaternion, and then adding translation vector to each point.""" - pts: NDArrayFloat = np.array([[1.0, 1.0, 1.1], [1.0, 1.0, 2.1], [1.0, 1.0, 3.1], [1.0, 1.0, 4.1]]) + pts: NDArrayFloat = np.array( + [[1.0, 1.0, 1.1], [1.0, 1.0, 2.1], [1.0, 1.0, 3.1], [1.0, 1.0, 4.1]] + ) # x, y, z of cuboid center t: NDArrayFloat = np.array([-34.7128603513203, 5.29461762417753, 0.10328996181488]) @@ -103,7 +105,9 @@ def test_SE3_inverse_transform_point_cloud_identity() -> None: Since the transformation was the identity, the points should not be affected. """ - transformed_pts: NDArrayFloat = np.array([[1.0, 1.0, 1.1], [1.0, 1.0, 2.1], [1.0, 1.0, 3.1], [1.0, 1.0, 4.1]]) + transformed_pts: NDArrayFloat = np.array( + [[1.0, 1.0, 1.1], [1.0, 1.0, 2.1], [1.0, 1.0, 3.1], [1.0, 1.0, 4.1]] + ) dst_se3_src = SE3(rotation=np.eye(3), translation=np.zeros(3)) pts = dst_se3_src.inverse().transform_point_cloud(transformed_pts.copy()) assert np.allclose(pts, transformed_pts) @@ -128,7 +132,9 @@ def test_SE3_inverse_transform_point_cloud() -> None: dst_se3_src = SE3(rotation=R, translation=t) pts = dst_se3_src.inverse().transform_point_cloud(transformed_pts) - gt_pts: NDArrayFloat = np.array([[1.0, 1.0, 1.1], [1.0, 1.0, 2.1], [1.0, 1.0, 3.1], [1.0, 1.0, 4.1]]) + gt_pts: NDArrayFloat = np.array( + [[1.0, 1.0, 1.1], [1.0, 1.0, 2.1], [1.0, 1.0, 3.1], [1.0, 1.0, 4.1]] + ) assert np.allclose(pts, gt_pts) @@ -165,7 +171,9 @@ def test_SE3_chaining_transforms() -> None: assert np.allclose(fr2_se3_fr0.translation, np.zeros(3)) -def test_benchmark_SE3_transform_point_cloud_optimized(benchmark: Callable[..., Any]) -> None: +def test_benchmark_SE3_transform_point_cloud_optimized( + benchmark: Callable[..., Any] +) -> None: """Ensure that our transform_point_cloud() implementation is faster than a naive implementation.""" num_pts = 100000 points_src = np.random.randn(num_pts, 3) @@ -182,7 +190,9 @@ def test_benchmark_SE3_transform_point_cloud_optimized(benchmark: Callable[..., benchmark(dst_SE3_src.transform_point_cloud, points_src) -def test_benchmark_SE3_transform_point_cloud_unoptimized(benchmark: Callable[..., Any]) -> None: +def test_benchmark_SE3_transform_point_cloud_unoptimized( + benchmark: Callable[..., Any] +) -> None: """Benchmark unoptimized SE(3) transformation.""" def benchmark_SE3_transform_point_cloud_unoptimized( @@ -192,7 +202,9 @@ def benchmark_SE3_transform_point_cloud_unoptimized( # convert to homogeneous num_pts = len(point_cloud) homogeneous_pts: NDArrayFloat = np.hstack([point_cloud, np.ones((num_pts, 1))]) - transformed_point_cloud: NDArrayFloat = homogeneous_pts.dot(transform_matrix.T)[:, :3] + transformed_point_cloud: NDArrayFloat = homogeneous_pts.dot(transform_matrix.T)[ + :, :3 + ] return transformed_point_cloud num_pts = 100000 @@ -207,4 +219,8 @@ def benchmark_SE3_transform_point_cloud_unoptimized( dst_SE3_src = SE3(rotation=R.copy(), translation=t.copy()) - benchmark(benchmark_SE3_transform_point_cloud_unoptimized, points_src, dst_SE3_src.transform_matrix) + benchmark( + benchmark_SE3_transform_point_cloud_unoptimized, + points_src, + dst_SE3_src.transform_matrix, + ) diff --git a/tests/unit/utils/test_sim2.py b/tests/unit/utils/test_sim2.py index 3f1c8cc0..53e97275 100644 --- a/tests/unit/utils/test_sim2.py +++ b/tests/unit/utils/test_sim2.py @@ -206,7 +206,9 @@ def test_transform_from_backwards() -> None: """Test Similarity(2) backward transform.""" img_pts: NDArrayFloat = np.array([[6, 4], [4, 6], [0, 0], [1, 7]]) - expected_world_pts: NDArrayFloat = np.array([[2, -1], [1, 0], [-1, -3], [-0.5, 0.5]]) + expected_world_pts: NDArrayFloat = np.array( + [[2, -1], [1, 0], [-1, -3], [-0.5, 0.5]] + ) scale = 0.5 wSimg = Sim2(R=np.eye(2), t=np.array([-2.0, -6.0]), s=scale) diff --git a/tutorials/3d_object_detection.py b/tutorials/3d_object_detection.py index cb34226b..db2497f2 100644 --- a/tutorials/3d_object_detection.py +++ b/tutorials/3d_object_detection.py @@ -34,7 +34,12 @@ def main( max_iterations: Maximum number of iterations for the data-loader example. """ logger.info("Starting detection data-loader example ...") - data_loader = DetectionDataLoader(root_dir, dataset_name, split_name, num_accumulated_sweeps=num_accumulated_sweeps) + data_loader = DetectionDataLoader( + root_dir, + dataset_name, + split_name, + num_accumulated_sweeps=num_accumulated_sweeps, + ) for i, sweep in enumerate(tqdm(data_loader)): # 4x4 matrix representing the SE(3) transformation to city from ego-vehicle coordinates. city_SE3_ego_mat4 = sweep.city_SE3_ego.matrix() diff --git a/tutorials/generate_egoview_overlaid_lidar.py b/tutorials/generate_egoview_overlaid_lidar.py index 37366a35..dd1101e4 100644 --- a/tutorials/generate_egoview_overlaid_lidar.py +++ b/tutorials/generate_egoview_overlaid_lidar.py @@ -31,7 +31,11 @@ def generate_egoview_overlaid_lidar( - data_root: Path, output_dir: Path, log_id: str, render_ground_pts_only: bool, dump_single_frames: bool + data_root: Path, + output_dir: Path, + log_id: str, + render_ground_pts_only: bool, + dump_single_frames: bool, ) -> None: """Render LiDAR points from a particular camera's viewpoint (color by ground surface, and apply ROI filtering). @@ -52,7 +56,9 @@ def generate_egoview_overlaid_lidar( avm = ArgoverseStaticMap.from_map_dir(log_map_dirpath, build_raster=True) # repeat red to green colormap every 50 m. - colors_arr_rgb = color_utils.create_colormap(color_list=[RED_HEX, GREEN_HEX], n_colors=NUM_RANGE_BINS) + colors_arr_rgb = color_utils.create_colormap( + color_list=[RED_HEX, GREEN_HEX], n_colors=NUM_RANGE_BINS + ) colors_arr_rgb = (colors_arr_rgb * 255).astype(np.uint8) colors_arr_bgr: NDArrayByte = np.fliplr(colors_arr_rgb) @@ -63,7 +69,9 @@ def generate_egoview_overlaid_lidar( video_list = [] for i, im_fpath in enumerate(cam_im_fpaths): if i % 50 == 0: - logging.info(f"\tOn file {i}/{num_cam_imgs} of camera {cam_name} of {log_id}") + logging.info( + f"\tOn file {i}/{num_cam_imgs} of camera {cam_name} of {log_id}" + ) cam_timestamp_ns = int(im_fpath.stem) city_SE3_ego = loader.get_city_SE3_ego(log_id, cam_timestamp_ns) @@ -74,7 +82,10 @@ def generate_egoview_overlaid_lidar( # load feather file path, e.g. '315978406032859416.feather" lidar_fpath = loader.get_closest_lidar_fpath(log_id, cam_timestamp_ns) if lidar_fpath is None: - logger.info("No LiDAR sweep found within the synchronization interval for %s, so skipping...", cam_name) + logger.info( + "No LiDAR sweep found within the synchronization interval for %s, so skipping...", + cam_name, + ) continue img_bgr = io_utils.read_img(im_fpath, channel_order="BGR") @@ -86,11 +97,19 @@ def generate_egoview_overlaid_lidar( lidar_points_city = city_SE3_ego.transform_point_cloud(lidar_points_ego) lidar_points_city = avm.remove_non_drivable_area_points(lidar_points_city) is_ground_logicals = avm.get_ground_points_boolean(lidar_points_city) - lidar_points_city = lidar_points_city[is_ground_logicals if render_ground_pts_only else ~is_ground_logicals] - lidar_points_ego = city_SE3_ego.inverse().transform_point_cloud(lidar_points_city) + lidar_points_city = lidar_points_city[ + is_ground_logicals if render_ground_pts_only else ~is_ground_logicals + ] + lidar_points_ego = city_SE3_ego.inverse().transform_point_cloud( + lidar_points_city + ) # motion compensate always - uv, points_cam, is_valid_points = loader.project_ego_to_img_motion_compensated( + ( + uv, + points_cam, + is_valid_points, + ) = loader.project_ego_to_img_motion_compensated( points_lidar_time=lidar_points_ego, cam_name=cam_name, cam_timestamp_ns=cam_timestamp_ns, @@ -113,26 +132,38 @@ def generate_egoview_overlaid_lidar( uv_colors_bgr = colors_arr_bgr[color_bins] img_empty = np.full_like(img_bgr, fill_value=255) - img_empty = raster_rendering_utils.draw_points_xy_in_img(img_empty, uv_int, uv_colors_bgr, diameter=10) + img_empty = raster_rendering_utils.draw_points_xy_in_img( + img_empty, uv_int, uv_colors_bgr, diameter=10 + ) blended_bgr = raster_utils.blend_images(img_bgr, img_empty) frame_rgb = blended_bgr[:, :, ::-1] if dump_single_frames: save_dir = output_dir / log_id / cam_name os.makedirs(save_dir, exist_ok=True) - cv2.imwrite(str(save_dir / f"{cam_name}_{lidar_timestamp_ns}.jpg"), blended_bgr) + cv2.imwrite( + str(save_dir / f"{cam_name}_{lidar_timestamp_ns}.jpg"), blended_bgr + ) video_list.append(frame_rgb) if len(video_list) == 0: - raise RuntimeError("No video frames were found; log data was not found on disk.") + raise RuntimeError( + "No video frames were found; log data was not found on disk." + ) video: NDArrayByte = np.stack(video_list).astype(np.uint8) video_output_dir = output_dir / "videos" - video_utils.write_video(video=video, dst=video_output_dir / f"{log_id}_{cam_name}.mp4", fps=RING_CAMERA_FPS) + video_utils.write_video( + video=video, + dst=video_output_dir / f"{log_id}_{cam_name}.mp4", + fps=RING_CAMERA_FPS, + ) -@click.command(help="Generate LiDAR + map visualizations from the Argoverse 2 Sensor Dataset.") +@click.command( + help="Generate LiDAR + map visualizations from the Argoverse 2 Sensor Dataset." +) @click.option( "-d", "--data-root", @@ -170,7 +201,11 @@ def generate_egoview_overlaid_lidar( type=bool, ) def run_generate_egoview_overlaid_lidar( - data_root: str, output_dir: str, log_id: str, render_ground_pts_only: bool, dump_single_frames: bool + data_root: str, + output_dir: str, + log_id: str, + render_ground_pts_only: bool, + dump_single_frames: bool, ) -> None: """Click entry point for visualizing LiDAR returns rendered on top of sensor imagery.""" logging.basicConfig(stream=sys.stdout, level=logging.INFO) diff --git a/tutorials/generate_egoview_overlaid_vector_map.py b/tutorials/generate_egoview_overlaid_vector_map.py index 584bd2de..ebbf2ce2 100644 --- a/tutorials/generate_egoview_overlaid_vector_map.py +++ b/tutorials/generate_egoview_overlaid_vector_map.py @@ -67,7 +67,9 @@ def generate_egoview_overlaid_map( video_list = [] for i, img_fpath in enumerate(cam_im_fpaths): if i % 50 == 0: - logging.info(f"\tOn file {i}/{num_cam_imgs} of camera {cam_name} of {log_id}") + logging.info( + f"\tOn file {i}/{num_cam_imgs} of camera {cam_name} of {log_id}" + ) cam_timestamp_ns = int(img_fpath.stem) city_SE3_ego = loader.get_city_SE3_ego(log_id, cam_timestamp_ns) @@ -96,7 +98,10 @@ def generate_egoview_overlaid_map( depth_map = None egoview_renderer = EgoViewMapRenderer( - depth_map=depth_map, city_SE3_ego=city_SE3_ego, pinhole_cam=pinhole_cam, avm=avm + depth_map=depth_map, + city_SE3_ego=city_SE3_ego, + pinhole_cam=pinhole_cam, + avm=avm, ) frame_rgb = render_egoview( output_dir=output_dir, @@ -148,7 +153,9 @@ def render_egoview( # we only create log-specific directories, if dumping individual frames. save_dir.mkdir(exist_ok=True, parents=True) - img_fname = f"{egoview_renderer.pinhole_cam.cam_name}_{cam_timestamp_ns}_vectormap.jpg" + img_fname = ( + f"{egoview_renderer.pinhole_cam.cam_name}_{cam_timestamp_ns}_vectormap.jpg" + ) save_fpath = save_dir / img_fname if save_fpath.exists(): @@ -183,7 +190,10 @@ def render_egoview( def render_egoview_with_occlusion_checks( - img_canvas: NDArrayByte, egoview_renderer: EgoViewMapRenderer, max_range_m: float, line_width_px: int = 10 + img_canvas: NDArrayByte, + egoview_renderer: EgoViewMapRenderer, + max_range_m: float, + line_width_px: int = 10, ) -> NDArrayByte: """Render pedestrian crossings and lane segments in the ego-view. @@ -200,8 +210,12 @@ def render_egoview_with_occlusion_checks( array of shape (H,W,3) and type uint8 representing a RGB image. """ for ls in egoview_renderer.avm.get_scenario_lane_segments(): - img_canvas = egoview_renderer.render_lane_boundary_egoview(img_canvas, ls, "right", line_width_px) - img_canvas = egoview_renderer.render_lane_boundary_egoview(img_canvas, ls, "left", line_width_px) + img_canvas = egoview_renderer.render_lane_boundary_egoview( + img_canvas, ls, "right", line_width_px + ) + img_canvas = egoview_renderer.render_lane_boundary_egoview( + img_canvas, ls, "left", line_width_px + ) for pc in egoview_renderer.avm.get_scenario_ped_crossings(): EPS = 1e-5 @@ -213,8 +227,12 @@ def render_egoview_with_occlusion_checks( N_INTERP_PTS = 100 # For pixel-perfect rendering, querying crosswalk boundary ground height at waypoints throughout # the street is much more accurate than 3d linear interpolation using only the 4 annotated corners. - polygon_city_frame = interp_utils.interp_arc(t=N_INTERP_PTS, points=xwalk_polygon[:, :2]) - polygon_city_frame = egoview_renderer.avm.append_height_to_2d_city_pt_cloud(points_xy=polygon_city_frame) + polygon_city_frame = interp_utils.interp_arc( + t=N_INTERP_PTS, points=xwalk_polygon[:, :2] + ) + polygon_city_frame = egoview_renderer.avm.append_height_to_2d_city_pt_cloud( + points_xy=polygon_city_frame + ) egoview_renderer.render_polyline_egoview( polygon_city_frame, img_canvas, @@ -250,7 +268,9 @@ def parse_camera_enum_types(cam_names: Tuple[str, ...]) -> List[RingCameras]: return cam_enums -@click.command(help="Generate map visualizations on ego-view imagery from the Argoverse 2 Sensor or TbV Datasets.") +@click.command( + help="Generate map visualizations on ego-view imagery from the Argoverse 2 Sensor or TbV Datasets." +) @click.option( "--data-root", required=True, diff --git a/tutorials/generate_forecasting_scenario_visualizations.py b/tutorials/generate_forecasting_scenario_visualizations.py index f2050954..c2273670 100644 --- a/tutorials/generate_forecasting_scenario_visualizations.py +++ b/tutorials/generate_forecasting_scenario_visualizations.py @@ -11,7 +11,9 @@ from rich.progress import track from av2.datasets.motion_forecasting import scenario_serialization -from av2.datasets.motion_forecasting.viz.scenario_visualization import visualize_scenario +from av2.datasets.motion_forecasting.viz.scenario_visualization import ( + visualize_scenario, +) from av2.map.map_api import ArgoverseStaticMap _DEFAULT_N_JOBS: Final[int] = -2 # Use all but one CPUs @@ -60,7 +62,9 @@ def generate_scenario_visualization(scenario_path: Path) -> None: scenario_path: Path to the parquet file corresponding to the Argoverse scenario to visualize. """ scenario_id = scenario_path.stem.split("_")[-1] - static_map_path = scenario_path.parents[0] / f"log_map_archive_{scenario_id}.json" + static_map_path = ( + scenario_path.parents[0] / f"log_map_archive_{scenario_id}.json" + ) viz_save_path = viz_output_dir / f"{scenario_id}.mp4" scenario = scenario_serialization.load_argoverse_scenario_parquet(scenario_path) @@ -73,7 +77,8 @@ def generate_scenario_visualization(scenario_path: Path) -> None: generate_scenario_visualization(scenario_path) else: Parallel(n_jobs=_DEFAULT_N_JOBS)( - delayed(generate_scenario_visualization)(scenario_path) for scenario_path in track(scenario_file_list) + delayed(generate_scenario_visualization)(scenario_path) + for scenario_path in track(scenario_file_list) ) @@ -104,9 +109,18 @@ def generate_scenario_visualization(scenario_path: Path) -> None: help="Controls how scenarios are selected for visualization - either the first available or at random.", type=click.Choice(["first", "random"], case_sensitive=False), ) -@click.option("--debug", is_flag=True, default=False, help="Runs preprocessing in single-threaded mode when enabled.") +@click.option( + "--debug", + is_flag=True, + default=False, + help="Runs preprocessing in single-threaded mode when enabled.", +) def run_generate_scenario_visualizations( - argoverse_scenario_dir: str, viz_output_dir: str, num_scenarios: int, selection_criteria: str, debug: bool + argoverse_scenario_dir: str, + viz_output_dir: str, + num_scenarios: int, + selection_criteria: str, + debug: bool, ) -> None: """Click entry point for generation of Argoverse scenario visualizations.""" generate_scenario_visualizations( diff --git a/tutorials/generate_per_camera_videos.py b/tutorials/generate_per_camera_videos.py index 45e45079..9a1bda2f 100644 --- a/tutorials/generate_per_camera_videos.py +++ b/tutorials/generate_per_camera_videos.py @@ -25,7 +25,9 @@ logger = logging.getLogger(__name__) -def generate_per_camera_videos(data_root: Path, output_dir: Path, num_workers: int) -> None: +def generate_per_camera_videos( + data_root: Path, output_dir: Path, num_workers: int +) -> None: """Launch jobs to render ring camera .mp4 videos for all sensor logs available on disk. Args: @@ -38,14 +40,17 @@ def generate_per_camera_videos(data_root: Path, output_dir: Path, num_workers: i if num_workers > 1: Parallel(n_jobs=num_workers)( - delayed(render_log_ring_camera_videos)(output_dir, loader, log_id) for log_id in log_ids + delayed(render_log_ring_camera_videos)(output_dir, loader, log_id) + for log_id in log_ids ) else: for log_id in log_ids: render_log_ring_camera_videos(output_dir, loader, log_id) -def render_log_ring_camera_videos(output_dir: Path, loader: AV2SensorDataLoader, log_id: str) -> None: +def render_log_ring_camera_videos( + output_dir: Path, loader: AV2SensorDataLoader, log_id: str +) -> None: """Render .mp4 videos for all ring cameras of a single log. Args: @@ -60,7 +65,9 @@ def render_log_ring_camera_videos(output_dir: Path, loader: AV2SensorDataLoader, for camera_name in list(RingCameras): video_save_fpath = Path(output_dir) / f"{log_id}_{camera_name}.mp4" if video_save_fpath.exists(): - logger.info("Video already exists for %s, %s, so skipping...", log_id, camera_name) + logger.info( + "Video already exists for %s, %s, so skipping...", log_id, camera_name + ) continue cam_im_fpaths = loader.get_ordered_log_cam_fpaths(log_id, camera_name) @@ -69,7 +76,9 @@ def render_log_ring_camera_videos(output_dir: Path, loader: AV2SensorDataLoader, video_list: List[NDArrayByte] = [] for i, im_fpath in enumerate(cam_im_fpaths): if i % PRINT_EVERY == 0: - logger.info(f"\tOn file {i}/{num_cam_imgs} of camera {camera_name} of {log_id}") + logger.info( + f"\tOn file {i}/{num_cam_imgs} of camera {camera_name} of {log_id}" + ) img_rgb = io_utils.read_img(im_fpath, channel_order="RGB") video_list.append(img_rgb) @@ -82,7 +91,9 @@ def render_log_ring_camera_videos(output_dir: Path, loader: AV2SensorDataLoader, ) -@click.command(help="Generate map visualizations on ego-view imagery from the Argoverse 2 Sensor or TbV Datasets.") +@click.command( + help="Generate map visualizations on ego-view imagery from the Argoverse 2 Sensor or TbV Datasets." +) @click.option( "-d", "--data-root", @@ -103,10 +114,14 @@ def render_log_ring_camera_videos(output_dir: Path, loader: AV2SensorDataLoader, help="Number of worker processes to use for rendering.", type=int, ) -def run_generate_per_camera_videos(data_root: str, output_dir: str, num_workers: int) -> None: +def run_generate_per_camera_videos( + data_root: str, output_dir: str, num_workers: int +) -> None: """Click entry point for ring camera .mp4 video generation.""" logging.basicConfig(stream=sys.stdout, level=logging.DEBUG) - generate_per_camera_videos(data_root=Path(data_root), output_dir=Path(output_dir), num_workers=num_workers) + generate_per_camera_videos( + data_root=Path(data_root), output_dir=Path(output_dir), num_workers=num_workers + ) if __name__ == "__main__": diff --git a/tutorials/generate_sensor_dataset_visualizations.py b/tutorials/generate_sensor_dataset_visualizations.py index f07c9597..c8a188fa 100644 --- a/tutorials/generate_sensor_dataset_visualizations.py +++ b/tutorials/generate_sensor_dataset_visualizations.py @@ -24,7 +24,9 @@ # Model an xy grid in the Bird's-eye view. BEV_GRID: Final[BEVGrid] = BEVGrid( - min_range_m=MIN_RANGE_M, max_range_m=MAX_RANGE_M, resolution_m_per_cell=RESOLUTION_M_PER_CELL + min_range_m=MIN_RANGE_M, + max_range_m=MAX_RANGE_M, + resolution_m_per_cell=RESOLUTION_M_PER_CELL, ) @@ -64,9 +66,15 @@ def generate_sensor_dataset_visualizations( and sweep.timestamp_ns in timestamp_city_SE3_ego_dict ): city_SE3_ego_cam_t = timestamp_city_SE3_ego_dict[cam.timestamp_ns] - city_SE3_ego_lidar_t = timestamp_city_SE3_ego_dict[sweep.timestamp_ns] - - uv, points_cam, is_valid_points = cam.camera_model.project_ego_to_img_motion_compensated( + city_SE3_ego_lidar_t = timestamp_city_SE3_ego_dict[ + sweep.timestamp_ns + ] + + ( + uv, + points_cam, + is_valid_points, + ) = cam.camera_model.project_ego_to_img_motion_compensated( sweep.xyz, city_SE3_ego_cam_t=city_SE3_ego_cam_t, city_SE3_ego_lidar_t=city_SE3_ego_lidar_t, @@ -75,11 +83,20 @@ def generate_sensor_dataset_visualizations( uv_int: NDArrayInt = np.round(uv[is_valid_points]).astype(int) colors = create_range_map(points_cam[is_valid_points, :3]) img = draw_points_xy_in_img( - cam.img, uv_int, colors=colors, alpha=0.85, diameter=5, sigma=1.0, with_anti_alias=True + cam.img, + uv_int, + colors=colors, + alpha=0.85, + diameter=5, + sigma=1.0, + with_anti_alias=True, ) if annotations is not None: img = annotations.project_to_cam( - img, cam.camera_model, city_SE3_ego_cam_t, city_SE3_ego_lidar_t + img, + cam.camera_model, + city_SE3_ego_cam_t, + city_SE3_ego_lidar_t, ) cam_name_to_img[cam_name] = img if len(cam_name_to_img) < len(cam_names): diff --git a/tutorials/map_teaser_notebook.py b/tutorials/map_teaser_notebook.py index f1756bae..dc041179 100644 --- a/tutorials/map_teaser_notebook.py +++ b/tutorials/map_teaser_notebook.py @@ -27,7 +27,11 @@ # scaled to [0,1] for matplotlib. PURPLE_RGB: Final[Tuple[int, int, int]] = (201, 71, 245) -PURPLE_RGB_MPL: Final[Tuple[float, float, float]] = (PURPLE_RGB[0] / 255, PURPLE_RGB[1] / 255, PURPLE_RGB[2] / 255) +PURPLE_RGB_MPL: Final[Tuple[float, float, float]] = ( + PURPLE_RGB[0] / 255, + PURPLE_RGB[1] / 255, + PURPLE_RGB[2] / 255, +) DARK_GRAY_RGB: Final[Tuple[int, int, int]] = (40, 39, 38) DARK_GRAY_RGB_MPL: Final[Tuple[float, float, float]] = ( @@ -54,13 +58,19 @@ def single_log_teaser(data_root: Path, log_id: str, save_figures: bool) -> None: ax = fig.add_subplot() for _, ls in avm.vector_lane_segments.items(): - vector_plotting_utils.draw_polygon_mpl(ax, ls.polygon_boundary, color="g", linewidth=0.5) - vector_plotting_utils.plot_polygon_patch_mpl(ls.polygon_boundary, ax, color="g", alpha=0.2) + vector_plotting_utils.draw_polygon_mpl( + ax, ls.polygon_boundary, color="g", linewidth=0.5 + ) + vector_plotting_utils.plot_polygon_patch_mpl( + ls.polygon_boundary, ax, color="g", alpha=0.2 + ) # plot all pedestrian crossings for _, pc in avm.vector_pedestrian_crossings.items(): vector_plotting_utils.draw_polygon_mpl(ax, pc.polygon, color="m", linewidth=0.5) - vector_plotting_utils.plot_polygon_patch_mpl(pc.polygon, ax, color="m", alpha=0.2) + vector_plotting_utils.plot_polygon_patch_mpl( + pc.polygon, ax, color="m", alpha=0.2 + ) plt.axis("equal") plt.tight_layout() @@ -73,7 +83,9 @@ def single_log_teaser(data_root: Path, log_id: str, save_figures: bool) -> None: ax = fig.add_subplot() for da in list(avm.vector_drivable_areas.values()): vector_plotting_utils.draw_polygon_mpl(ax, da.xyz, color="gray", linewidth=0.5) - vector_plotting_utils.plot_polygon_patch_mpl(da.xyz, ax, color="gray", alpha=0.2) + vector_plotting_utils.plot_polygon_patch_mpl( + da.xyz, ax, color="gray", alpha=0.2 + ) plt.axis("equal") plt.tight_layout() @@ -176,7 +188,8 @@ def overlaid_maps_all_logs_teaser(data_root: Path) -> None: pts_ego = city_SE3_egot0.inverse().transform_point_cloud(pts_city) for bound_type, bound_city in zip( - [ls.left_mark_type, ls.right_mark_type], [ls.left_lane_boundary, ls.right_lane_boundary] + [ls.left_mark_type, ls.right_mark_type], + [ls.left_lane_boundary, ls.right_lane_boundary], ): if "YELLOW" in bound_type: mark_color = "y" @@ -194,7 +207,9 @@ def overlaid_maps_all_logs_teaser(data_root: Path) -> None: else: linestyle = "solid" - bound_ego = city_SE3_egot0.inverse().transform_point_cloud(bound_city.xyz) + bound_ego = city_SE3_egot0.inverse().transform_point_cloud( + bound_city.xyz + ) ax.plot( bound_ego[:, 0], bound_ego[:, 1], @@ -205,7 +220,11 @@ def overlaid_maps_all_logs_teaser(data_root: Path) -> None: ) vector_plotting_utils.plot_polygon_patch_mpl( - polygon_pts=pts_ego, ax=ax, color=color, alpha=OVERLAID_MAPS_ALPHA, zorder=i + polygon_pts=pts_ego, + ax=ax, + color=color, + alpha=OVERLAID_MAPS_ALPHA, + zorder=i, ) plt.axis("equal") @@ -216,7 +235,9 @@ def overlaid_maps_all_logs_teaser(data_root: Path) -> None: def plot_lane_segments( - ax: Axes, lane_segments: Sequence[LaneSegment], lane_color: Tuple[float, float, float] = DARK_GRAY_RGB_MPL + ax: Axes, + lane_segments: Sequence[LaneSegment], + lane_color: Tuple[float, float, float] = DARK_GRAY_RGB_MPL, ) -> None: """Plot lane segments onto a Matplotlib canvas, according to their lane marking boundary type/color. @@ -237,7 +258,8 @@ def plot_lane_segments( mark_color: str = "" linestyle: Union[str, Tuple[int, Tuple[int, int]]] = "" for bound_type, bound_city in zip( - [ls.left_mark_type, ls.right_mark_type], [ls.left_lane_boundary, ls.right_lane_boundary] + [ls.left_mark_type, ls.right_mark_type], + [ls.left_lane_boundary, ls.right_lane_boundary], ): if "YELLOW" in bound_type: mark_color = "y" @@ -259,8 +281,22 @@ def plot_lane_segments( left, right = polyline_utils.get_double_polylines( polyline=bound_city.xyz[:, :2], width_scaling_factor=0.1 ) - ax.plot(left[:, 0], left[:, 1], color=mark_color, alpha=ALPHA, linestyle=linestyle, zorder=2) - ax.plot(right[:, 0], right[:, 1], color=mark_color, alpha=ALPHA, linestyle=linestyle, zorder=2) + ax.plot( + left[:, 0], + left[:, 1], + color=mark_color, + alpha=ALPHA, + linestyle=linestyle, + zorder=2, + ) + ax.plot( + right[:, 0], + right[:, 1], + color=mark_color, + alpha=ALPHA, + linestyle=linestyle, + zorder=2, + ) else: ax.plot( bound_city.xyz[:, 0], @@ -272,7 +308,9 @@ def plot_lane_segments( ) -def visualize_ego_pose_and_lane_markings(data_root: Path, log_id: str, save_figures: bool) -> None: +def visualize_ego_pose_and_lane_markings( + data_root: Path, log_id: str, save_figures: bool +) -> None: """Visualize both ego-vehicle poses and the per-log local vector map. Crosswalks are plotted in purple. Lane segments plotted in dark gray. Ego-pose in red. @@ -318,7 +356,15 @@ def visualize_ego_pose_and_lane_markings(data_root: Path, log_id: str, save_figu # Plot nearly continuous line for ego-pose, and show the AV's pose @ 1 Hz w/ red unfilled circles. ax.plot(traj_ns[:, 0], traj_ns[:, 1], color="r", zorder=4, label="Ego-vehicle pose") - ax.scatter(x=traj_1hz[:, 0], y=traj_1hz[:, 1], s=100, marker="o", facecolors="none", edgecolors="r", zorder=4) + ax.scatter( + x=traj_1hz[:, 0], + y=traj_1hz[:, 1], + s=100, + marker="o", + facecolors="none", + edgecolors="r", + zorder=4, + ) plt.axis("equal") plt.xlim(*xlims) @@ -362,9 +408,15 @@ def run_map_tutorial(data_root: str, log_id: str, save_figures: bool) -> None: SAVE_DIR.mkdir(exist_ok=True, parents=True) logger.info("data_root: %s, log_id: %s", data_root_path, log_id) - single_log_teaser(data_root=data_root_path, log_id=log_id, save_figures=save_figures) - visualize_raster_layers(data_root=data_root_path, log_id=log_id, save_figures=save_figures) - visualize_ego_pose_and_lane_markings(data_root=data_root_path, log_id=log_id, save_figures=save_figures) + single_log_teaser( + data_root=data_root_path, log_id=log_id, save_figures=save_figures + ) + visualize_raster_layers( + data_root=data_root_path, log_id=log_id, save_figures=save_figures + ) + visualize_ego_pose_and_lane_markings( + data_root=data_root_path, log_id=log_id, save_figures=save_figures + ) overlaid_maps_all_logs_teaser(data_root=data_root_path) diff --git a/tutorials/untar_tbv.py b/tutorials/untar_tbv.py index 923091ca..38666fb2 100644 --- a/tutorials/untar_tbv.py +++ b/tutorials/untar_tbv.py @@ -12,7 +12,9 @@ NUM_TBV_SHARDS: Final[int] = 21 -def run_command(cmd: str, return_output: bool = False) -> Tuple[Optional[bytes], Optional[bytes]]: +def run_command( + cmd: str, return_output: bool = False +) -> Tuple[Optional[bytes], Optional[bytes]]: """Execute a system call, and block until the system call completes. Args: @@ -23,14 +25,18 @@ def run_command(cmd: str, return_output: bool = False) -> Tuple[Optional[bytes], Tuple of (stdout, stderr) output if return_output is True, else None """ print(cmd) - (stdout_data, stderr_data) = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE).communicate() + (stdout_data, stderr_data) = subprocess.Popen( + cmd, shell=True, stdout=subprocess.PIPE + ).communicate() if return_output: return stdout_data, stderr_data return None, None -def untar_tbv_dataset(num_workers: int, shard_dirpath: Path, desired_tbv_dataroot: Path) -> None: +def untar_tbv_dataset( + num_workers: int, shard_dirpath: Path, desired_tbv_dataroot: Path +) -> None: """Untar each of the tar.gz archives. Args: @@ -58,7 +64,9 @@ def untar_tbv_dataset(num_workers: int, shard_dirpath: Path, desired_tbv_dataroo run_command(cmd) -@click.command(help="Extract TbV tar.gz archives that were previously downloaded to a local disk.") +@click.command( + help="Extract TbV tar.gz archives that were previously downloaded to a local disk." +) @click.option( "--num-workers", required=True, @@ -77,10 +85,14 @@ def untar_tbv_dataset(num_workers: int, shard_dirpath: Path, desired_tbv_dataroo help="Path to local directory, where TbV logs will be extracted.", type=str, ) -def run_untar_tbv_dataset(num_workers: int, shard_dirpath: str, desired_tbv_dataroot: str) -> None: +def run_untar_tbv_dataset( + num_workers: int, shard_dirpath: str, desired_tbv_dataroot: str +) -> None: """Click entry point for TbV tar.gz file extraction.""" untar_tbv_dataset( - num_workers=num_workers, shard_dirpath=Path(shard_dirpath), desired_tbv_dataroot=Path(desired_tbv_dataroot) + num_workers=num_workers, + shard_dirpath=Path(shard_dirpath), + desired_tbv_dataroot=Path(desired_tbv_dataroot), )