diff --git a/bucketed_scene_flow_eval/datasets/nuscenes/dataset.py b/bucketed_scene_flow_eval/datasets/nuscenes/dataset.py index 31313c6..a675b17 100644 --- a/bucketed_scene_flow_eval/datasets/nuscenes/dataset.py +++ b/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,6 +122,8 @@ 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, @@ -118,7 +131,7 @@ def __init__( ) 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, diff --git a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_raw_data.py b/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_raw_data.py index a52d7cb..e53b7f0 100644 --- a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_raw_data.py +++ b/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 diff --git a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_scene_flow.py b/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_scene_flow.py index cdb8adf..7036a99 100644 --- a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_scene_flow.py +++ b/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) diff --git a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py b/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py index 7955fd0..7b38a12 100644 --- a/bucketed_scene_flow_eval/datasets/nuscenes/nuscenes_utils.py +++ b/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, diff --git a/tests/datasets/nuscenes/nuscenes_tests.py b/tests/datasets/nuscenes/nuscenes_tests.py index bb6988c..84ae33f 100644 --- a/tests/datasets/nuscenes/nuscenes_tests.py +++ b/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)}"