Adds nuscenes log level train/val split (#7)

* Adds nuscenes log level train/val split

bucketed_scene_flow_eval/datasets/nuscenes/dataset.py

@@ -2,6 +2,10 @@
 from pathlib import Path
 from typing import Optional, Union
 
+from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import (
+    DEFAULT_POINT_CLOUD_RANGE,
+    PointCloudRange,
+)
 from bucketed_scene_flow_eval.datastructures import *
 from bucketed_scene_flow_eval.eval import (
     BucketedEPEEvaluator,
@@ -14,13 +18,15 @@
     NonCausalSeqLoaderDataset,
 )
 
-from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import DEFAULT_POINT_CLOUD_RANGE, PointCloudRange
+from .nuscenes_metacategories import (
+    BUCKETED_METACATAGORIES,
+    THREEWAY_EPE_METACATAGORIES,
+)
 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:
@@ -47,6 +53,7 @@ def __init__(
         self,
         root_dir: Union[Path, list[Path]],
         nuscenes_version: str,
+        split: str,
         subsequence_length: int = 2,
         with_ground: bool = True,
         with_rgb: bool = False,
@@ -64,6 +71,7 @@ def __init__(
             self.sequence_loader = NuScenesSceneFlowSequenceLoader(
                 raw_data_path=root_dir,
                 nuscenes_version=nuscenes_version,
+                split=split,
                 with_rgb=with_rgb,
                 use_gt_flow=use_gt_flow,
                 flow_data_path=flow_data_path,
@@ -74,6 +82,7 @@ def __init__(
             self.sequence_loader = NuScenesNoFlowSequenceLoader(
                 raw_data_path=root_dir,
                 nuscenes_version=nuscenes_version,
+                split=split,
                 with_rgb=with_rgb,
                 expected_camera_shape=expected_camera_shape,
                 point_cloud_range=point_cloud_range,
@@ -96,6 +105,8 @@ class NuScenesNonCausalSceneFlow(NonCausalSeqLoaderDataset):
     def __init__(
         self,
         root_dir: Union[Path, list[Path]],
+        nuscenes_version: str,
+        split: str,
         subsequence_length: int = 2,
         with_ground: bool = True,
         with_rgb: bool = False,
@@ -111,14 +122,16 @@ def __init__(
         if load_flow:
             self.sequence_loader = NuScenesSceneFlowSequenceLoader(
                 root_dir,
+                nuscenes_version=nuscenes_version,
+                split=split,
                 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
+                root_dir, nuscenes_version=nuscenes_version, split=split, with_rgb=with_rgb, expected_camera_shape=expected_camera_shape
             )
         super().__init__(
             sequence_loader=self.sequence_loader,

bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_raw_data.py

@@ -3,15 +3,15 @@
 from typing import Union
 
 import numpy as np
+import open3d as o3d
 from numpy.typing import NDArray
 from nuscenes.utils.data_classes import LidarPointCloud as NuscLidarPointCloud
 from PIL import Image
 from pyquaternion import Quaternion
-import open3d as o3d
 
 from bucketed_scene_flow_eval.datasets.argoverse2.argoverse_raw_data import (
-    PointCloudRange,
     DEFAULT_POINT_CLOUD_RANGE,
+    PointCloudRange,
 )
 from bucketed_scene_flow_eval.datastructures import (
     SE3,
@@ -28,7 +28,7 @@
 )
 from bucketed_scene_flow_eval.interfaces import AbstractSequence, CachedSequenceLoader
 
-from .nuscenes_utils import NuScenesWithInstanceBoxes
+from .nuscenes_utils import NuScenesWithInstanceBoxes, create_splits_tokens
 
 NuscDict = dict[str, Union[str, int, list]]
 NuscSample = dict[str, NuscDict]
@@ -350,6 +350,7 @@ class NuScenesRawSequenceLoader(CachedSequenceLoader):
     def __init__(
         self,
         sequence_dir: Path,
+        split: str,
         version: str = "v1.0-mini",
         verbose: bool = False,
         point_cloud_range: PointCloudRange | None = DEFAULT_POINT_CLOUD_RANGE,
@@ -362,6 +363,7 @@ def __init__(
             version=version, dataroot=sequence_dir, verbose=verbose
         )
         self.log_lookup: dict[str, NuscDict] = {e["token"]: e for e in self.nusc.scene}
+        self.log_lookup: dict[str, NuscDict] = {k: self.log_lookup[k] for k in create_splits_tokens(split, self.nusc)}
 
         self.point_cloud_range = point_cloud_range
 

bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_scene_flow.py

@@ -24,9 +24,14 @@
     TimeSyncedSceneFlowFrame,
     VectorArray,
 )
-from bucketed_scene_flow_eval.interfaces import AbstractAVLidarSequence, CachedSequenceLoader
+from bucketed_scene_flow_eval.interfaces import (
+    AbstractAVLidarSequence,
+    CachedSequenceLoader,
+)
 from bucketed_scene_flow_eval.utils.loaders import load_feather
 
+from .nuscenes_utils import create_splits_tokens
+
 CATEGORY_MAP = {
     -1: "background",
     0: "animal",
@@ -191,6 +196,7 @@ class NuScenesSceneFlowSequenceLoader(ArgoverseSceneFlowSequenceLoader, CachedSe
     def __init__(
         self,
         raw_data_path: Path | list[Path],
+        split: str,
         nuscenes_version: str = "v1.0-mini",
         flow_data_path: Path | list[Path] | None = None,
         use_gt_flow: bool = True,
@@ -209,6 +215,7 @@ def __init__(
         self.sequence_id_to_raw_data: dict[str, NuscDict] = {
             e["token"]: e for e in self.nuscenes.scene
         }
+        self.sequence_id_to_raw_data: dict[str, NuscDict] = {k: self.sequence_id_to_raw_data[k] for k in create_splits_tokens(split, self.nuscenes)}
         self.sequence_id_lst: list[str] = sorted(self.sequence_id_to_raw_data.keys())
         self._setup_flow_data(use_gt_flow, flow_data_path)
         self._subset_log(log_subset)
@@ -262,6 +269,7 @@ class NuScenesNoFlowSequenceLoader(NuScenesSceneFlowSequenceLoader):
     def __init__(
         self,
         raw_data_path: Path | list[Path],
+        split: str,
         nuscenes_version: str = "v1.0-mini",
         with_rgb: bool = False,
         log_subset: list[str] | None = None,
@@ -278,6 +286,7 @@ def __init__(
         self.sequence_id_to_raw_data: dict[str, NuscDict] = {
             e["token"]: e for e in self.nuscenes.scene
         }
+        self.sequence_id_to_raw_data: dict[str, NuscDict] = {k: self.sequence_id_to_raw_data[k] for k in create_splits_tokens(split, self.nuscenes)}
         self.sequence_id_lst: list[str] = sorted(self.sequence_id_to_raw_data.keys())
         self._subset_log(log_subset)
 

bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py

@@ -3,10 +3,49 @@
 import numpy as np
 from nuscenes.nuscenes import NuScenes
 from nuscenes.utils.data_classes import Box
+from nuscenes.utils.splits import create_splits_scenes
 from pyquaternion import Quaternion
 
 from bucketed_scene_flow_eval.datastructures.se3 import SE3
 
+
+def create_splits_tokens(split: str, nusc: 'NuScenes') -> list[str]:
+    """
+    Returns the logs in each dataset split of nuScenes.
+    Note: Previously this script included the teaser dataset splits. Since new scenes from those logs were added and
+          others removed in the full dataset, that code is incompatible and was removed.
+    :param split: NuScenes split.
+    :param nusc: NuScenes instance.
+    :return: A list of logs in that split.
+    """
+    # Load splits on a scene-level.
+    scene_splits = create_splits_scenes(verbose=False)
+
+    assert split in scene_splits.keys(), 'Requested split {} which is not a known nuScenes split.'.format(split)
+
+    # Check compatibility of split with nusc_version.
+    version = nusc.version
+    if split in {'train', 'val', 'train_detect', 'train_track'}:
+        assert version.endswith('trainval'), \
+            'Requested split {} which is not compatible with NuScenes version {}'.format(split, version)
+    elif split in {'mini_train', 'mini_val'}:
+        assert version.endswith('mini'), \
+            'Requested split {} which is not compatible with NuScenes version {}'.format(split, version)
+    elif split == 'test':
+        assert version.endswith('test'), \
+            'Requested split {} which is not compatible with NuScenes version {}'.format(split, version)
+    else:
+        raise ValueError('Requested split {} which this function cannot map to logs.'.format(split))
+
+    # Get logs for this split.
+    scene_to_log = {scene['name']: scene['token'] for scene in nusc.scene}
+    logs = set()
+    scenes = scene_splits[split]
+    for scene in scenes:
+        logs.add(scene_to_log[scene])
+
+    return list(logs)
+
 class InstanceBox(Box):
     def __init__(
         self,

tests/datasets/nuscenes/nuscenes_tests.py

@@ -10,13 +10,14 @@ def nuscenes_loader() -> NuScenesRawSequenceLoader:
     return NuScenesRawSequenceLoader(
         sequence_dir=Path("/tmp/nuscenes"),
         version="v1.0-mini",
+        split="mini_train",
         verbose=False,
     )
 
 
 def test_nuscenes_loader_basic_load_and_len_check(nuscenes_loader: NuScenesRawSequenceLoader):
     assert len(nuscenes_loader) > 0, f"no sequences found in {nuscenes_loader}"
-    expected_lens = [236, 239, 236, 236, 233, 223, 239, 231, 231, 228]
+    expected_lens= [236, 236, 236, 233, 223, 239, 231, 231]
     assert len(nuscenes_loader) == len(
         expected_lens
     ), f"expected {len(expected_lens)} sequences, got {len(nuscenes_loader)}"