Add nuScenes (lidar only) support (#6)

* Add nuScenes (lidar only) support

* fix existing tests

* bugfixes for loader

* Basic docs

---------

Co-authored-by: Kyle Vedder <kyle.c.vedder@gmail.com>

README.md

@@ -11,6 +11,7 @@ Currently supported datasets:
 
  - Argoverse 2 (Human Labeled and [NSFP Pseudolabeled](https://github.com/kylevedder/BucketedSceneFlowEval/blob/master/docs/GETTING_STARTED.md#argoverse-2-nsfp-pseudolabels-new))
  - Waymo Open (LiDAR only)
+ - NuScenes (LiDAR only, beta)
 
 
 If you use this repository as part of a publication, please cite:

bucketed_scene_flow_eval/datasets/__init__.py

@@ -6,11 +6,18 @@
     WaymoOpenCausalSceneFlow,
     WaymoOpenNonCausalSceneFlow,
 )
+
+from bucketed_scene_flow_eval.datasets.nuscenes import (
+    NuScenesCausalSceneFlow,
+    NuScenesNonCausalSceneFlow,
+)
 from bucketed_scene_flow_eval.interfaces import AbstractDataset
 
 importable_classes = [
     Argoverse2CausalSceneFlow,
     Argoverse2NonCausalSceneFlow,
+    NuScenesCausalSceneFlow,
+    NuScenesNonCausalSceneFlow,
     WaymoOpenCausalSceneFlow,
     WaymoOpenNonCausalSceneFlow,
 ]

bucketed_scene_flow_eval/datasets/argoverse2/__init__.py

@@ -1,8 +1,6 @@
 from .argoverse_raw_data import (
-    DEFAULT_POINT_CLOUD_RANGE,
     ArgoverseRawSequence,
     ArgoverseRawSequenceLoader,
-    PointCloudRange,
 )
 from .argoverse_scene_flow import (
     ArgoverseNoFlowSequence,

bucketed_scene_flow_eval/datasets/argoverse2/argoverse_raw_data.py

@@ -2,7 +2,7 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Optional
+from typing import Optional
 
 import cv2
 import numpy as np
@@ -196,7 +196,7 @@ def __init__(
         expected_camera_shape=(1550, 2048, 3),
     ):
         self.log_id = log_id
-        self.POINT_CLOUD_RANGE = point_cloud_range
+        self.point_cloud_range = point_cloud_range
 
         self.dataset_dir = Path(dataset_dir)
         assert self.dataset_dir.is_dir(), f"dataset_dir {dataset_dir} does not exist"
@@ -418,14 +418,14 @@ def is_ground_points(self, global_point_cloud: PointCloud) -> np.ndarray:
         return is_ground_boolean_arr
 
     def is_in_range(self, global_point_cloud: PointCloud) -> MaskArray:
-        if self.POINT_CLOUD_RANGE is None:
+        if self.point_cloud_range is None:
             return np.ones(len(global_point_cloud), dtype=bool)
-        xmin = self.POINT_CLOUD_RANGE[0]
-        ymin = self.POINT_CLOUD_RANGE[1]
-        zmin = self.POINT_CLOUD_RANGE[2]
-        xmax = self.POINT_CLOUD_RANGE[3]
-        ymax = self.POINT_CLOUD_RANGE[4]
-        zmax = self.POINT_CLOUD_RANGE[5]
+        xmin = self.point_cloud_range[0]
+        ymin = self.point_cloud_range[1]
+        zmin = self.point_cloud_range[2]
+        xmax = self.point_cloud_range[3]
+        ymax = self.point_cloud_range[4]
+        zmax = self.point_cloud_range[5]
         return global_point_cloud.within_region_mask(xmin, xmax, ymin, ymax, zmin, zmax)
 
     def __repr__(self) -> str:

bucketed_scene_flow_eval/datasets/argoverse2/argoverse_scene_flow.py

@@ -1,5 +1,3 @@
-import copy
-import dataclasses
 from pathlib import Path
 from typing import Optional, Union
 
@@ -9,7 +7,6 @@
     EgoLidarFlow,
     MaskArray,
     PointCloud,
-    PointCloudFrame,
     SemanticClassId,
     SemanticClassIdArray,
     SupervisedPointCloudFrame,
@@ -24,7 +21,7 @@
 )
 from bucketed_scene_flow_eval.utils.loaders import load_feather
 
-from . import DEFAULT_POINT_CLOUD_RANGE, ArgoverseRawSequence, PointCloudRange
+from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import DEFAULT_POINT_CLOUD_RANGE, ArgoverseRawSequence, PointCloudRange
 
 CATEGORY_MAP = {
     -1: "BACKGROUND",

bucketed_scene_flow_eval/datasets/argoverse2/dataset.py

@@ -1,6 +1,6 @@
 import copy
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Optional, Union
 
 from bucketed_scene_flow_eval.datastructures import *
 from bucketed_scene_flow_eval.eval import (
@@ -9,7 +9,6 @@
     ThreeWayEPEEvaluator,
 )
 from bucketed_scene_flow_eval.interfaces import (
-    AbstractAVLidarSequence,
     CausalSeqLoaderDataset,
     EvalType,
     NonCausalSeqLoaderDataset,

bucketed_scene_flow_eval/datasets/nuscenes/__init__.py

@@ -1 +1,22 @@
-from .nuscenes_loader import NuScenesLoader, NuScenesSequence
+from .nuscenes_raw_data import (
+    NuScenesRawSequence,
+    NuScenesRawSequenceLoader,
+)
+from .nuscenes_scene_flow import (
+    NuScenesNoFlowSequence,
+    NuScenesNoFlowSequenceLoader,
+    NuScenesSceneFlowSequence,
+    NuScenesSceneFlowSequenceLoader
+)
+from .dataset import NuScenesCausalSceneFlow, NuScenesNonCausalSceneFlow
+
+__all__ = [
+    "NuScenesCausalSceneFlow",
+    "NuScenesNonCausalSceneFlow",
+    "NuScenesNoFlowSequence",
+    "NuScenesNoFlowSequenceLoader",
+    "NuScenesRawSequence",
+    "NuScenesRawSequenceLoader",
+    "NuScenesSceneFlowSequence",
+    "NuScenesSceneFlowSequenceLoader",
+]

bucketed_scene_flow_eval/datasets/nuscenes/dataset.py

@@ -0,0 +1,134 @@
+import copy
+from pathlib import Path
+from typing import Optional, Union
+
+from bucketed_scene_flow_eval.datastructures import *
+from bucketed_scene_flow_eval.eval import (
+    BucketedEPEEvaluator,
+    Evaluator,
+    ThreeWayEPEEvaluator,
+)
+from bucketed_scene_flow_eval.interfaces import (
+    CausalSeqLoaderDataset,
+    EvalType,
+    NonCausalSeqLoaderDataset,
+)
+
+from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import DEFAULT_POINT_CLOUD_RANGE, PointCloudRange
+from .nuscenes_scene_flow import (
+    CATEGORY_MAP,
+    NuScenesNoFlowSequenceLoader,
+    NuScenesSceneFlowSequenceLoader,
+)
+from .nuscenes_metacategories import BUCKETED_METACATAGORIES, THREEWAY_EPE_METACATAGORIES
+
+
+def _make_evaluator(eval_type: EvalType, eval_args: dict) -> Evaluator:
+    eval_args_copy = copy.deepcopy(eval_args)
+    # Builds the evaluator object for this dataset.
+    if eval_type == EvalType.BUCKETED_EPE:
+        if "meta_class_lookup" not in eval_args_copy:
+            eval_args_copy["meta_class_lookup"] = BUCKETED_METACATAGORIES
+        if "class_id_to_name" not in eval_args_copy:
+            eval_args_copy["class_id_to_name"] = CATEGORY_MAP
+        return BucketedEPEEvaluator(**eval_args_copy)
+    elif eval_type == EvalType.THREEWAY_EPE:
+        if "meta_class_lookup" not in eval_args_copy:
+            eval_args_copy["meta_class_lookup"] = THREEWAY_EPE_METACATAGORIES
+        if "class_id_to_name" not in eval_args_copy:
+            eval_args_copy["class_id_to_name"] = CATEGORY_MAP
+        return ThreeWayEPEEvaluator(**eval_args_copy)
+    else:
+        raise ValueError(f"Unknown eval type {eval_type}")
+
+
+class NuScenesCausalSceneFlow(CausalSeqLoaderDataset):
+    def __init__(
+        self,
+        root_dir: Union[Path, list[Path]],
+        nuscenes_version: str,
+        subsequence_length: int = 2,
+        with_ground: bool = True,
+        with_rgb: bool = False,
+        cache_root: Path = Path("/tmp/"),
+        use_gt_flow: bool = True,
+        flow_data_path: Optional[Union[Path, list[Path]]] = None,
+        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,
+        use_cache=True,
+        load_flow: bool = True,
+    ) -> None:
+        if load_flow:
+            self.sequence_loader = NuScenesSceneFlowSequenceLoader(
+                raw_data_path=root_dir,
+                nuscenes_version=nuscenes_version,
+                with_rgb=with_rgb,
+                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 = NuScenesNoFlowSequenceLoader(
+                raw_data_path=root_dir,
+                nuscenes_version=nuscenes_version,
+                with_rgb=with_rgb,
+                expected_camera_shape=expected_camera_shape,
+                point_cloud_range=point_cloud_range,
+            )
+        super().__init__(
+            sequence_loader=self.sequence_loader,
+            subsequence_length=subsequence_length,
+            with_ground=with_ground,
+            idx_lookup_cache_root=cache_root,
+            eval_type=eval_type,
+            eval_args=eval_args,
+            use_cache=use_cache,
+        )
+
+    def evaluator(self) -> Evaluator:
+        return _make_evaluator(self.eval_type, self.eval_args)
+
+
+class NuScenesNonCausalSceneFlow(NonCausalSeqLoaderDataset):
+    def __init__(
+        self,
+        root_dir: Union[Path, list[Path]],
+        subsequence_length: int = 2,
+        with_ground: bool = True,
+        with_rgb: bool = False,
+        cache_root: Path = Path("/tmp/"),
+        use_gt_flow: bool = True,
+        flow_data_path: Optional[Union[Path, list[Path]]] = None,
+        eval_type: str = "bucketed_epe",
+        eval_args=dict(),
+        expected_camera_shape: tuple[int, int, int] = (1550, 2048, 3),
+        use_cache=True,
+        load_flow: bool = True,
+    ) -> None:
+        if load_flow:
+            self.sequence_loader = NuScenesSceneFlowSequenceLoader(
+                root_dir,
+                with_rgb=with_rgb,
+                use_gt_flow=use_gt_flow,
+                flow_data_path=flow_data_path,
+                expected_camera_shape=expected_camera_shape,
+            )
+        else:
+            self.sequence_loader = NuScenesNoFlowSequenceLoader(
+                root_dir, with_rgb=with_rgb, expected_camera_shape=expected_camera_shape
+            )
+        super().__init__(
+            sequence_loader=self.sequence_loader,
+            subsequence_length=subsequence_length,
+            with_ground=with_ground,
+            idx_lookup_cache_root=cache_root,
+            eval_type=eval_type,
+            eval_args=eval_args,
+            use_cache=use_cache,
+        )
+
+    def evaluator(self) -> Evaluator:
+        return _make_evaluator(self.eval_type, self.eval_args)

bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_metacategories.py

@@ -0,0 +1,48 @@
+BACKGROUND_CATEGORIES = ["background"]
+
+# These catagories are ignored because of labeling oddities
+STATIC_OBJECTS = [
+    "movable_object.barrier",
+    "movable_object.debris",
+    "movable_object.pushable_pullable",
+    "movable_object.trafficcone",
+    "static_object.bicycle_rack",
+]
+
+PEDESTRIAN_CATEGORIES = [
+    "animal",
+    "human.pedestrian.adult",
+    "human.pedestrian.child",
+    "human.pedestrian.construction_worker",
+    "human.pedestrian.personal_mobility",
+    "human.pedestrian.police_officer",
+    "human.pedestrian.stroller",
+    "human.pedestrian.wheelchair",
+]
+
+WHEELED_VRU = ["vehicle.bicycle", "vehicle.motorcycle"]
+
+CAR = ["vehicle.car"]
+
+OTHER_VEHICLES = [
+    "vehicle.bus.bendy",
+    "vehicle.bus.rigid",
+    "vehicle.construction",
+    "vehicle.emergency.ambulance",
+    "vehicle.emergency.police",
+    "vehicle.trailer",
+    "vehicle.truck",
+]
+
+BUCKETED_METACATAGORIES = {
+    "BACKGROUND": BACKGROUND_CATEGORIES,
+    "CAR": CAR,
+    "PEDESTRIAN": PEDESTRIAN_CATEGORIES,
+    "WHEELED_VRU": WHEELED_VRU,
+    "OTHER_VEHICLES": OTHER_VEHICLES,
+}
+
+THREEWAY_EPE_METACATAGORIES = {
+    "BACKGROUND": BACKGROUND_CATEGORIES,
+    "FOREGROUND": PEDESTRIAN_CATEGORIES + WHEELED_VRU + CAR + OTHER_VEHICLES,
+}