v2.0.10: Added volume based clustering support for AV2

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

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"],
+}

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,

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]]]

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)

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,14 +149,19 @@ 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 <output_dir>/<log_id>_<sweep_1_timestamp>.npz
      Args:
        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)