Merge branch 'master' of github.com:kylevedder/BucketedSceneFlowEval

bucketed_scene_flow_eval/datasets/argoverse2/__init__.py

@@ -1,13 +1,15 @@
-from .argoverse_raw_data import (
-    ArgoverseRawSequence,
-    ArgoverseRawSequenceLoader,
-)
+from .argoverse_raw_data import ArgoverseRawSequence, ArgoverseRawSequenceLoader
 from .argoverse_scene_flow import (
     ArgoverseNoFlowSequence,
     ArgoverseNoFlowSequenceLoader,
     ArgoverseSceneFlowSequence,
     ArgoverseSceneFlowSequenceLoader,
 )
+
+from .argoverse_box_annotations import (
+    ArgoverseBoxAnnotationSequence,
+    ArgoverseBoxAnnotationSequenceLoader,
+)
 from .dataset import Argoverse2CausalSceneFlow, Argoverse2NonCausalSceneFlow
 
 __all__ = [

bucketed_scene_flow_eval/datasets/argoverse2/argoverse_box_annotations.py

@@ -0,0 +1,81 @@
+from dataclasses import dataclass
+from pathlib import Path
+
+from bucketed_scene_flow_eval.datastructures import (
+    SE3,
+    BoundingBox,
+    TimeSyncedAVLidarData,
+    TimeSyncedSceneFlowBoxFrame,
+    TimeSyncedSceneFlowFrame,
+)
+from bucketed_scene_flow_eval.utils import load_feather
+
+from .argoverse_scene_flow import ArgoverseNoFlowSequence, ArgoverseNoFlowSequenceLoader
+
+
+class ArgoverseBoxAnnotationSequence(ArgoverseNoFlowSequence):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.timestamp_to_boxes = self._prep_bbox_annotations()
+
+    def _prep_bbox_annotations(self) -> dict[int, list[BoundingBox]]:
+        annotations_file = self.dataset_dir / "annotations.feather"
+        assert annotations_file.exists(), f"Annotations file {annotations_file} does not exist"
+        annotation_df = load_feather(annotations_file)
+        # Index(['timestamp_ns', 'track_uuid', 'category', 'length_m', 'width_m',
+        #         'height_m', 'qw', 'qx', 'qy', 'qz', 'tx_m', 'ty_m', 'tz_m',
+        #         'num_interior_pts'],
+        #         dtype='object')
+
+        # Convert to dictionary keyed by timestamp_ns int
+        timestamp_to_annotations: dict[int, list[BoundingBox]] = {}
+        for _, row in annotation_df.iterrows():
+            timestamp_ns = row["timestamp_ns"]
+            if timestamp_ns not in timestamp_to_annotations:
+                timestamp_to_annotations[timestamp_ns] = []
+            pose = SE3.from_rot_w_x_y_z_translation_x_y_z(
+                row["qw"],
+                row["qx"],
+                row["qy"],
+                row["qz"],
+                row["tx_m"],
+                row["ty_m"],
+                row["tz_m"],
+            )
+            timestamp_to_annotations[timestamp_ns].append(
+                BoundingBox(
+                    pose=pose,
+                    length=row["length_m"],
+                    width=row["width_m"],
+                    height=row["height_m"],
+                    track_uuid=row["track_uuid"],
+                    category=row["category"],
+                )
+            )
+        return timestamp_to_annotations
+
+    def load(
+        self, idx: int, relative_to_idx: int, with_flow: bool = False
+    ) -> tuple[TimeSyncedSceneFlowBoxFrame, TimeSyncedAVLidarData]:
+        scene_flow_frame, lidar_data = super().load(idx, relative_to_idx, with_flow)
+        timestamp = self.timestamp_list[idx]
+        boxes = self.timestamp_to_boxes.get(timestamp, [])
+        return TimeSyncedSceneFlowBoxFrame(**vars(scene_flow_frame), boxes=boxes), lidar_data
+
+
+class ArgoverseBoxAnnotationSequenceLoader(ArgoverseNoFlowSequenceLoader):
+
+    def _load_sequence_uncached(self, sequence_id: str) -> ArgoverseBoxAnnotationSequence:
+        assert (
+            sequence_id in self.sequence_id_to_raw_data
+        ), f"sequence_id {sequence_id} does not exist"
+        return ArgoverseBoxAnnotationSequence(
+            sequence_id,
+            self.sequence_id_to_raw_data[sequence_id],
+            self.sequence_id_to_raw_data[sequence_id],
+            with_classes=False,
+            **self.load_sequence_kwargs,
+        )
+
+    def cache_folder_name(self) -> str:
+        return f"av2_box_data_use_gt_flow_{self.use_gt_flow}_raw_data_path_{self.raw_data_path}_No_flow_data_path"

bucketed_scene_flow_eval/datasets/argoverse2/argoverse_scene_flow.py

@@ -264,6 +264,7 @@ def _subset_log(self, log_subset: Optional[list[str]]):
             self.sequence_id_lst = [
                 sequence_id for sequence_id in self.sequence_id_lst if sequence_id in log_subset
             ]
+            assert len(self.sequence_id_lst) > 0, f"No sequences found in log_subset {log_subset}"
 
     def _sanitize_raw_data_path(self, raw_data_path: Union[Path, list[Path]]) -> list[Path]:
         if isinstance(raw_data_path, str):
@@ -359,7 +360,7 @@ def _load_flow_feather(
         raise NotImplementedError("No flow data available for ArgoverseNoFlowSequence")
 
     def load(
-        self, idx: int, relative_to_idx: int, with_flow: bool = True
+        self, idx: int, relative_to_idx: int, with_flow: bool = False
     ) -> tuple[TimeSyncedSceneFlowFrame, TimeSyncedAVLidarData]:
         return super().load(idx, relative_to_idx, with_flow=False)
 

bucketed_scene_flow_eval/datasets/argoverse2/dataset.py

@@ -14,6 +14,7 @@
     NonCausalSeqLoaderDataset,
 )
 
+from .argoverse_box_annotations import ArgoverseBoxAnnotationSequenceLoader
 from .argoverse_raw_data import DEFAULT_POINT_CLOUD_RANGE, PointCloudRange
 from .argoverse_scene_flow import (
     CATEGORY_MAP,
@@ -69,11 +70,17 @@ def __init__(
         flow_data_path: Optional[Union[Path, list[Path]]] = None,
         eval_type: str = "bucketed_epe",
         eval_args=dict(),
+        load_boxes: bool = False,
         load_flow: bool = True,
         use_cache=True,
         **kwargs,
     ) -> None:
-        if load_flow:
+        if load_boxes:
+            self.sequence_loader = ArgoverseBoxAnnotationSequenceLoader(
+                root_dir,
+                **kwargs,
+            )
+        elif load_flow:
             self.sequence_loader = ArgoverseSceneFlowSequenceLoader(
                 root_dir,
                 use_gt_flow=use_gt_flow,
@@ -113,10 +120,16 @@ def __init__(
         eval_type: str = "bucketed_epe",
         eval_args=dict(),
         use_cache=True,
+        load_boxes: bool = False,
         load_flow: bool = True,
         **kwargs,
     ) -> None:
-        if load_flow:
+        if load_boxes:
+            self.sequence_loader = ArgoverseBoxAnnotationSequenceLoader(
+                raw_data_path=root_dir,
+                **kwargs,
+            )
+        elif load_flow:
             self.sequence_loader = ArgoverseSceneFlowSequenceLoader(
                 raw_data_path=root_dir,
                 use_gt_flow=use_gt_flow,

bucketed_scene_flow_eval/datastructures/__init__.py

@@ -1,5 +1,6 @@
 from .camera_projection import CameraModel, CameraProjection
 from .dataclasses import (
+    BoundingBox,
     EgoLidarDistance,
     EgoLidarFlow,
     MaskArray,
@@ -13,6 +14,7 @@
     TimeSyncedAVLidarData,
     TimeSyncedBaseAuxilaryData,
     TimeSyncedRawFrame,
+    TimeSyncedSceneFlowBoxFrame,
     TimeSyncedSceneFlowFrame,
     VectorArray,
 )

bucketed_scene_flow_eval/datastructures/dataclasses.py

@@ -272,6 +272,16 @@ def __len__(self) -> int:
         return len(self.lookup)
 
 
+@dataclass
+class BoundingBox:
+    pose: SE3
+    length: float
+    width: float
+    height: float
+    track_uuid: str
+    category: str
+
+
 @dataclass(kw_only=True)
 class TimeSyncedBaseAuxilaryData:
     pass
@@ -310,3 +320,14 @@ def __post_init__(self):
         assert len(self.flow.full_flow) == len(
             self.pc.full_pc
         ), f"flow and pc must have the same length, got {len(self.flow.full_flow)} and {len(self.pc.full_pc)}"
+
+
+@dataclass(kw_only=True)
+class TimeSyncedSceneFlowBoxFrame(TimeSyncedSceneFlowFrame):
+    boxes: list[BoundingBox]
+
+    def __post_init__(self):
+        assert isinstance(self.boxes, list), f"boxes must be a list, got {type(self.boxes)}"
+        assert all(
+            isinstance(box, BoundingBox) for box in self.boxes
+        ), f"all boxes must be BoundingBox objects, got {self.boxes}"