From 5cadf2f5d8f0582e40d0ba5ee6da6234bebe0d27 Mon Sep 17 00:00:00 2001 From: Kyle Vedder Date: Wed, 15 May 2024 18:57:09 -0400 Subject: [PATCH] v2.0.10: Added volume based clustering support for AV2 --- .../argoverse2/argoverse_scene_flow.py | 13 +- .../datasets/argoverse2/av2_metacategories.py | 7 + .../datasets/argoverse2/dataset.py | 26 ++- .../base_dataset_abstract_seq_loader.py | 1 + data_prep_scripts/argo/count_boxes.py | 89 +++++++++++ data_prep_scripts/argo/create_gt_flow.py | 68 +++++++- data_prep_scripts/argo/plot_boxes.py | 150 ++++++++++++++++++ pyproject.toml | 2 +- 8 files changed, 341 insertions(+), 15 deletions(-) create mode 100644 data_prep_scripts/argo/count_boxes.py create mode 100644 data_prep_scripts/argo/plot_boxes.py diff --git a/bucketed_scene_flow_eval/datasets/argoverse2/argoverse_scene_flow.py b/bucketed_scene_flow_eval/datasets/argoverse2/argoverse_scene_flow.py index 96fad5e..9e73971 100644 --- a/bucketed_scene_flow_eval/datasets/argoverse2/argoverse_scene_flow.py +++ b/bucketed_scene_flow_eval/datasets/argoverse2/argoverse_scene_flow.py @@ -3,6 +3,11 @@ import numpy as np +from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import ( + DEFAULT_POINT_CLOUD_RANGE, + ArgoverseRawSequence, + PointCloudRange, +) from bucketed_scene_flow_eval.datastructures import ( EgoLidarFlow, MaskArray, @@ -21,8 +26,6 @@ ) from bucketed_scene_flow_eval.utils.loaders import load_feather -from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import DEFAULT_POINT_CLOUD_RANGE, ArgoverseRawSequence, PointCloudRange - CATEGORY_MAP = { -1: "BACKGROUND", 0: "ANIMAL", @@ -116,7 +119,11 @@ def _load_flow_feather( assert idx >= 0, f"idx {idx} is out of range" flow_data_file = self.flow_data_files[idx] flow_data = load_feather(flow_data_file, verbose=False) - is_valid_arr = flow_data["is_valid"].values + is_valid_arr = flow_data["is_valid"].values.astype(bool) + # Ensure that is_valid_arr is a boolean array. + assert ( + is_valid_arr.dtype == bool + ), f"is_valid_arr must be a boolean array, got {is_valid_arr.dtype} from {flow_data_file.absolute()}" # The flow data is stored as 3 1D arrays, one for each dimension. xs = flow_data["flow_tx_m"].values diff --git a/bucketed_scene_flow_eval/datasets/argoverse2/av2_metacategories.py b/bucketed_scene_flow_eval/datasets/argoverse2/av2_metacategories.py index 8635576..45a69ba 100644 --- a/bucketed_scene_flow_eval/datasets/argoverse2/av2_metacategories.py +++ b/bucketed_scene_flow_eval/datasets/argoverse2/av2_metacategories.py @@ -49,3 +49,10 @@ "BACKGROUND": BACKGROUND_CATEGORIES, "FOREGROUND": PEDESTRIAN_CATEGORIES + WHEELED_VRU + CAR + OTHER_VEHICLES, } + +BUCKETED_VOLUME_METACATAGORIES = { + "BACKGROUND": ["BACKGROUND"], + "SMALL": ["SMALL"], + "MEDIUM": ["MEDIUM"], + "LARGE": ["LARGE"], +} diff --git a/bucketed_scene_flow_eval/datasets/argoverse2/dataset.py b/bucketed_scene_flow_eval/datasets/argoverse2/dataset.py index 7cf0eb4..806934b 100644 --- a/bucketed_scene_flow_eval/datasets/argoverse2/dataset.py +++ b/bucketed_scene_flow_eval/datasets/argoverse2/dataset.py @@ -20,7 +20,11 @@ ArgoverseNoFlowSequenceLoader, ArgoverseSceneFlowSequenceLoader, ) -from .av2_metacategories import BUCKETED_METACATAGORIES, THREEWAY_EPE_METACATAGORIES +from .av2_metacategories import ( + BUCKETED_METACATAGORIES, + BUCKETED_VOLUME_METACATAGORIES, + THREEWAY_EPE_METACATAGORIES, +) def _make_av2_evaluator(eval_type: EvalType, eval_args: dict) -> Evaluator: @@ -38,6 +42,17 @@ def _make_av2_evaluator(eval_type: EvalType, eval_args: dict) -> Evaluator: if "class_id_to_name" not in eval_args_copy: eval_args_copy["class_id_to_name"] = CATEGORY_MAP return ThreeWayEPEEvaluator(**eval_args_copy) + elif eval_type == EvalType.BUCKETED_VOLUME_EPE: + if "meta_class_lookup" not in eval_args_copy: + eval_args_copy["meta_class_lookup"] = BUCKETED_VOLUME_METACATAGORIES + if "class_id_to_name" not in eval_args_copy: + eval_args_copy["class_id_to_name"] = { + -1: "BACKGROUND", + 0: "SMALL", + 1: "MEDIUM", + 2: "LARGE", + } + return BucketedEPEEvaluator(**eval_args_copy) else: raise ValueError(f"Unknown eval type {eval_type}") @@ -55,7 +70,7 @@ def __init__( eval_type: str = "bucketed_epe", eval_args=dict(), expected_camera_shape: tuple[int, int, int] = (1550, 2048, 3), - point_cloud_range: Optional[PointCloudRange] = DEFAULT_POINT_CLOUD_RANGE, + point_cloud_range: PointCloudRange | None = DEFAULT_POINT_CLOUD_RANGE, use_cache=True, load_flow: bool = True, ) -> None: @@ -102,6 +117,7 @@ def __init__( eval_type: str = "bucketed_epe", eval_args=dict(), expected_camera_shape: tuple[int, int, int] = (1550, 2048, 3), + point_cloud_range: PointCloudRange | None = DEFAULT_POINT_CLOUD_RANGE, use_cache=True, load_flow: bool = True, ) -> None: @@ -112,10 +128,14 @@ def __init__( use_gt_flow=use_gt_flow, flow_data_path=flow_data_path, expected_camera_shape=expected_camera_shape, + point_cloud_range=point_cloud_range, ) else: self.sequence_loader = ArgoverseNoFlowSequenceLoader( - root_dir, with_rgb=with_rgb, expected_camera_shape=expected_camera_shape + root_dir, + with_rgb=with_rgb, + expected_camera_shape=expected_camera_shape, + point_cloud_range=point_cloud_range, ) super().__init__( sequence_loader=self.sequence_loader, diff --git a/bucketed_scene_flow_eval/interfaces/base_dataset_abstract_seq_loader.py b/bucketed_scene_flow_eval/interfaces/base_dataset_abstract_seq_loader.py index 27a6f84..7291b38 100644 --- a/bucketed_scene_flow_eval/interfaces/base_dataset_abstract_seq_loader.py +++ b/bucketed_scene_flow_eval/interfaces/base_dataset_abstract_seq_loader.py @@ -20,6 +20,7 @@ class EvalType(enum.Enum): BUCKETED_EPE = 0 THREEWAY_EPE = 1 + BUCKETED_VOLUME_EPE = 2 CacheLookup = list[tuple[int, tuple[int, int]]] diff --git a/data_prep_scripts/argo/count_boxes.py b/data_prep_scripts/argo/count_boxes.py new file mode 100644 index 0000000..b860872 --- /dev/null +++ b/data_prep_scripts/argo/count_boxes.py @@ -0,0 +1,89 @@ +import json +from argparse import ArgumentParser +from collections import defaultdict +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + + +def load_cuboid_metadata(json_file: Path): + """Load cuboid metadata from a JSON file.""" + with open(json_file, "r") as f: + cuboid_metadata = json.load(f) + return cuboid_metadata + + +def plot_histogram(cuboid_metadata, output_file: Path): + """Plot a histogram of boxes by volume, colored by class name.""" + volumes_by_class = defaultdict(list) + + # Collect volumes by class + for entry in cuboid_metadata: + class_name = entry["class_name"] + volume = entry["volume"] + volumes_by_class[class_name].append(volume) + + # Prepare data for the histogram + classes = list(volumes_by_class.keys()) + volumes = [volumes_by_class[class_name] for class_name in classes] + + # Create histogram bins + all_volumes = np.concatenate(volumes) + bin_edges = np.histogram_bin_edges(all_volumes, bins="auto") + + # Compute histogram data for each class + histogram_data = [] + for class_volumes in volumes: + counts, _ = np.histogram(class_volumes, bins=bin_edges) + histogram_data.append(counts) + + histogram_data = np.array(histogram_data) + + # Plot stacked histogram + plt.figure(figsize=(10, 6), dpi=300) + bottom = np.zeros(len(bin_edges) - 1) + + for class_name, class_counts in zip(classes, histogram_data): + plt.bar( + bin_edges[:-1], + class_counts, + width=np.diff(bin_edges), + bottom=bottom, + label=class_name, + align="edge", + ) + bottom += class_counts + + plt.xlabel("Volume") + plt.ylabel("Count") + plt.title("Histogram of Boxes by Volume and Class") + plt.legend() + plt.grid(True) + + # Save plot as high-resolution PNG + plt.savefig(output_file, format="png") + plt.close() + + +if __name__ == "__main__": + parser = ArgumentParser(description="Plot histogram of cuboid volumes by class") + parser.add_argument( + "--json_file", + type=str, + required=True, + help="Path to the JSON file containing cuboid metadata", + ) + parser.add_argument( + "--output_file", + type=str, + required=True, + help="Path to save the output PNG file", + ) + + args = parser.parse_args() + json_file = Path(args.json_file) + output_file = Path(args.output_file) + + cuboid_metadata = load_cuboid_metadata(json_file) + plot_histogram(cuboid_metadata, output_file) diff --git a/data_prep_scripts/argo/create_gt_flow.py b/data_prep_scripts/argo/create_gt_flow.py index 4159ac0..7e5bec5 100644 --- a/data_prep_scripts/argo/create_gt_flow.py +++ b/data_prep_scripts/argo/create_gt_flow.py @@ -2,6 +2,7 @@ os.environ["OMP_NUM_THREADS"] = "1" +import enum import multiprocessing import os from argparse import ArgumentParser @@ -23,6 +24,32 @@ from bucketed_scene_flow_eval.utils.loaders import save_feather +def category_to_class_id(category: str) -> int: + assert isinstance(category, str), f"Expected str, got {type(category)}" + assert category in CATEGORY_MAP_INV, f"Unknown category: {category}" + return CATEGORY_MAP_INV[category] + + +def volume_to_class_id(volume: float) -> int: + # Determined by clustering. + if volume < 9.5: + return 0 # SMALL + elif volume < 40: + return 1 # MEDIUM + return 2 # LARGE + + +class ClassType(enum.Enum): + SEMANTIC = "SEMANTIC" + VOLUME = "VOLUME" + + +CLASS_TYPE_DICT = { + ClassType.SEMANTIC.value: category_to_class_id, + ClassType.VOLUME.value: volume_to_class_id, +} + + def get_ids_and_cuboids_at_lidar_timestamps( dataset: AV2SensorDataLoader, log_id: str, lidar_timestamps_ns: list[int] ) -> list[dict[str, Cuboid]]: @@ -56,13 +83,14 @@ def get_ids_and_cuboids_at_lidar_timestamps( def compute_sceneflow( - dataset: AV2SensorDataLoader, log_id: str, timestamps: tuple[int, int] + dataset: AV2SensorDataLoader, log_id: str, timestamps: tuple[int, int], class_type: ClassType ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Compute sceneflow between the sweeps at the given timestamps. Args: dataset: Sensor dataset. log_id: unique id. timestamps: the timestamps of the lidar sweeps to compute flow between + class_type: ClassType to use for the output classes Returns: dictionary with fields: pcl_0: Nx3 array containing the points at time 0 @@ -91,7 +119,15 @@ def compute_flow(sweeps, cuboids, poses): ) c0.width_m += 0.2 obj_pts, obj_mask = c0.compute_interior_points(sweeps[0].xyz) - classes_0[obj_mask] = CATEGORY_MAP_INV[c0.category] + + match class_type: + case ClassType.SEMANTIC: + classes_0[obj_mask] = CLASS_TYPE_DICT[class_type.value](c0.category) + # CATEGORY_MAP_INV[c0.category] + case ClassType.VOLUME: + classes_0[obj_mask] = CLASS_TYPE_DICT[class_type.value]( + c0.length_m * c0.width_m * c0.height_m + ) if id in cuboids[1]: c1 = cuboids[1][id] @@ -113,7 +149,11 @@ def compute_flow(sweeps, cuboids, poses): def process_log( - dataset: AV2SensorDataLoader, log_id: str, output_dir: Path, n: Optional[int] = None + dataset: AV2SensorDataLoader, + log_id: str, + output_dir: Path, + class_type: ClassType, + n: Optional[int] = None, ): """Outputs sceneflow and auxillary information for each pair of pointclouds in the dataset. Output files have the format /_.npz @@ -121,6 +161,7 @@ def process_log( dataset: Sensor dataset to process. log_id: Log unique id. output_dir: Output_directory. + class_type: ClassType to use for the output classes n: the position to use for the progress bar Returns: None @@ -138,7 +179,7 @@ def process_log( ) for ts0, ts1 in iter_bar: - flow_0_1, classes_0, valid_0 = compute_sceneflow(dataset, log_id, (ts0, ts1)) + flow_0_1, classes_0, valid_0 = compute_sceneflow(dataset, log_id, (ts0, ts1), class_type) df = pd.DataFrame( { "flow_tx_m": flow_0_1[:, 0], @@ -160,7 +201,7 @@ def process_log_wrapper(x, ignore_current_process=False): process_log(*x, n=pos) -def process_logs(data_dir: Path, output_dir: Path, nproc: int): +def process_logs(data_dir: Path, output_dir: Path, nproc: int, class_type: ClassType): """Compute sceneflow for all logs in the dataset. Logs are processed in parallel. Args: data_dir: Argoverse 2.0 directory @@ -176,7 +217,7 @@ def process_logs(data_dir: Path, output_dir: Path, nproc: int): dataset = AV2SensorDataLoader(data_dir=data_dir, labels_dir=data_dir) logs = dataset.get_log_ids() - args = sorted([(dataset, log, split_output_dir) for log in logs]) + args = sorted([(dataset, log, split_output_dir, class_type) for log in logs]) print(f"Using {nproc} processes") if nproc <= 1: @@ -196,15 +237,26 @@ def process_logs(data_dir: Path, output_dir: Path, nproc: int): parser.add_argument( "--argo_dir", type=str, + required=True, help="The top level directory contating the input dataset", ) parser.add_argument( - "--output_dir", type=str, help="The location to output the sceneflow files to" + "--output_dir", + type=str, + required=True, + help="The location to output the sceneflow files to", ) parser.add_argument("--nproc", type=int, default=(multiprocessing.cpu_count() - 1)) + parser.add_argument( + "--class_type", + type=str, + default=ClassType.SEMANTIC.value, + choices=[ClassType.SEMANTIC.value, ClassType.VOLUME.value], + ) args = parser.parse_args() data_root = Path(args.argo_dir) output_dir = Path(args.output_dir) + class_type = ClassType[args.class_type] - process_logs(data_root, output_dir, args.nproc) + process_logs(data_root, output_dir, args.nproc, class_type) diff --git a/data_prep_scripts/argo/plot_boxes.py b/data_prep_scripts/argo/plot_boxes.py new file mode 100644 index 0000000..8a982e8 --- /dev/null +++ b/data_prep_scripts/argo/plot_boxes.py @@ -0,0 +1,150 @@ +import json +from argparse import ArgumentParser +from collections import defaultdict +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np + +# Define the category mappings +BACKGROUND_CATEGORIES = ["BACKGROUND"] + +ROAD_SIGNS = [ + "BOLLARD", + "CONSTRUCTION_BARREL", + "CONSTRUCTION_CONE", + "MOBILE_PEDESTRIAN_CROSSING_SIGN", + "SIGN", + "STOP_SIGN", + "MESSAGE_BOARD_TRAILER", + "TRAFFIC_LIGHT_TRAILER", +] + +PEDESTRIAN_CATEGORIES = ["PEDESTRIAN", "STROLLER", "WHEELCHAIR", "OFFICIAL_SIGNALER"] + +WHEELED_VRU = [ + "BICYCLE", + "BICYCLIST", + "MOTORCYCLE", + "MOTORCYCLIST", + "WHEELED_DEVICE", + "WHEELED_RIDER", +] + +CAR = ["REGULAR_VEHICLE"] + +OTHER_VEHICLES = [ + "BOX_TRUCK", + "LARGE_VEHICLE", + "RAILED_VEHICLE", + "TRUCK", + "TRUCK_CAB", + "VEHICULAR_TRAILER", + "ARTICULATED_BUS", + "BUS", + "SCHOOL_BUS", +] + +BUCKETED_METACATAGORIES = { + "BACKGROUND": BACKGROUND_CATEGORIES, + "CAR": CAR, + "PEDESTRIAN": PEDESTRIAN_CATEGORIES, + "WHEELED_VRU": WHEELED_VRU, + "OTHER_VEHICLES": OTHER_VEHICLES, +} + +# Reverse mapping from specific category to meta category +CATEGORY_TO_META = { + category: meta + for meta, categories in BUCKETED_METACATAGORIES.items() + for category in categories +} + + +def load_cuboid_metadata(json_file: Path): + """Load cuboid metadata from a JSON file.""" + print("Loading cuboid metadata from", json_file) + with open(json_file, "r") as f: + cuboid_metadata = json.load(f) + print("Loaded", len(cuboid_metadata), "cuboids") + return cuboid_metadata + + +def plot_histogram(cuboid_metadata, output_file: Path): + """Plot a histogram of boxes by volume, colored by class name.""" + volumes_by_class = defaultdict(list) + + # Collect volumes by class + for entry in cuboid_metadata: + class_name = entry["class_name"] + if class_name in CATEGORY_TO_META: + meta_class_name = CATEGORY_TO_META[class_name] + volume = entry["volume"] + volumes_by_class[meta_class_name].append(volume) + + # Prepare data for the histogram + classes = list(volumes_by_class.keys()) + volumes = [volumes_by_class[class_name] for class_name in classes] + + # Create histogram bins + all_volumes = np.concatenate(volumes) + bin_edges = np.histogram_bin_edges(all_volumes, bins="auto") + + # Compute histogram data for each class + histogram_data = [] + for class_volumes in volumes: + counts, _ = np.histogram(class_volumes, bins=bin_edges) + histogram_data.append(counts) + + histogram_data = np.array(histogram_data) + + # Plot stacked histogram + plt.figure(figsize=(10, 6), dpi=300) + bottom = np.zeros(len(bin_edges) - 1) + + for class_name, class_counts in zip(classes, histogram_data): + plt.bar( + bin_edges[:-1], + class_counts, + width=np.diff(bin_edges), + bottom=bottom, + label=class_name, + align="edge", + ) + bottom += class_counts + + plt.yscale("log") + plt.xlabel("Volume") + plt.ylabel("Count") + plt.title("Histogram of Boxes by Volume and Class") + # Set x-axis to limit of 0 to 60 + plt.xlim(0, 60) + plt.legend() + plt.grid(True, which="both", ls="--") + + # Save plot as high-resolution PNG + plt.savefig(output_file, format="png") + plt.close() + + +if __name__ == "__main__": + parser = ArgumentParser(description="Plot histogram of cuboid volumes by class") + parser.add_argument( + "--json_file", + type=str, + required=True, + help="Path to the JSON file containing cuboid metadata", + ) + parser.add_argument( + "--output_file", + type=str, + required=True, + help="Path to save the output PNG file", + ) + + args = parser.parse_args() + json_file = Path(args.json_file) + output_file = Path(args.output_file) + + cuboid_metadata = load_cuboid_metadata(json_file) + plot_histogram(cuboid_metadata, output_file) diff --git a/pyproject.toml b/pyproject.toml index 5cb7b2d..ab353f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,7 +15,7 @@ exclude = [ [project] name = "bucketed_scene_flow_eval" -version = "2.0.9" +version = "2.0.10" authors = [ { name="Kyle Vedder", email="kvedder@seas.upenn.edu" }, ]