diff --git a/conftest.py b/conftest.py new file mode 100644 index 0000000..5c7fd71 --- /dev/null +++ b/conftest.py @@ -0,0 +1,24 @@ +import pytest + + +def pytest_addoption(parser): + parser.addoption( + "--runslow", action="store_true", default=False, help="run slow tests" + ) + + +def pytest_configure(config): + config.addinivalue_line( + "markers", + "slow: marks tests as slow (deselect with '-m \"not slow\"' or run only with --runslow)", + ) + + +def pytest_collection_modifyitems(config, items): + if config.getoption("--runslow"): + # --runslow flag was passed; don't skip slow tests + return + skip_slow = pytest.mark.skip(reason="need --runslow flag to run") + for item in items: + if "slow" in item.keywords: + item.add_marker(skip_slow) diff --git a/pyproject.toml b/pyproject.toml index bb49bca..d379822 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,6 +44,7 @@ dependencies = [ "tabulate<0.11", "tensorboard<2.21", "timm<1.1", + "triton<3.8", ] [project.urls] diff --git a/requirements.txt b/requirements.txt index ba46be1..adb146b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,4 +11,5 @@ requests<2.34 scipy<1.18 tabulate<0.11 tensorboard<2.21 -timm<1.1 \ No newline at end of file +timm<1.1 +triton<3.8 \ No newline at end of file diff --git a/roco_spring_devkit/common/data/Spring_val.txt b/roco_spring_devkit/common/data/Spring_val.txt index 6b0a56f..e7b4a12 100644 --- a/roco_spring_devkit/common/data/Spring_val.txt +++ b/roco_spring_devkit/common/data/Spring_val.txt @@ -1 +1 @@ -0027 \ No newline at end of file +0022 \ No newline at end of file diff --git a/roco_spring_devkit/common/data/scene_flow_datamodule.py b/roco_spring_devkit/common/data/scene_flow_datamodule.py index 6e7f279..86fc9de 100644 --- a/roco_spring_devkit/common/data/scene_flow_datamodule.py +++ b/roco_spring_devkit/common/data/scene_flow_datamodule.py @@ -49,6 +49,7 @@ def __init__( spring_root_dir: Optional[str] = None, robust_spring_root_dir: Optional[str] = None, dataset_config_path: str = "../../datasets.yaml", + disparity2_in_frame1: bool = True, ): super().__init__() self.predict_dataset = predict_dataset @@ -67,6 +68,9 @@ def __init__( self.spring_root_dir = spring_root_dir self.robust_spring_root_dir = robust_spring_root_dir self.dataset_config_path = dataset_config_path + # All the datasets return the second-frame disparity ('disparities'[1]) on the pixel grid of the first frame, + # which is the convention of the KITTI/Spring scene flow benchmarks and of the models in this repository. + self.disparity2_in_frame1 = disparity2_in_frame1 self.predict_dataset_parsed = None self.test_dataset_parsed = None @@ -321,7 +325,10 @@ def _get_kitti_dataset(self, is_train: bool, *args: str) -> Dataset: # These transforms are based on RAFT: https://github.com/princeton-vl/RAFT transform = ft.Compose( [ - ft.ToTensor(device=device, fp16=self.train_transform_fp16), + ft.ToTensor( + device=device, + fp16=self.train_transform_fp16, + ), ft.RandomScaleAndCrop( (cy, cx), (-0.2, 0.4), (-0.2, 0.2), sparse=True ), @@ -342,6 +349,7 @@ def _get_kitti_dataset(self, is_train: bool, *args: str) -> Dataset: get_flow=True, get_disparity=True, get_intrinsics=True, + disparity2_in_frame1=self.disparity2_in_frame1, ) return dataset @@ -378,7 +386,10 @@ def _get_sintel_dataset(self, is_train: bool, *args: str) -> Dataset: # These transforms are based on RAFT: https://github.com/princeton-vl/RAFT transform = ft.Compose( [ - ft.ToTensor(device=device, fp16=self.train_transform_fp16), + ft.ToTensor( + device=device, + fp16=self.train_transform_fp16, + ), ft.RandomScaleAndCrop((cy, cx), (-0.2, 0.6), (-0.2, 0.2)), ft.ColorJitter(0.4, 0.4, 0.4, 0.5 / 3.14, 0.2), ft.GaussianNoise(0.02), @@ -399,6 +410,7 @@ def _get_sintel_dataset(self, is_train: bool, *args: str) -> Dataset: get_flow=True, get_disparity=True, get_intrinsics=True, + disparity2_in_frame1=self.disparity2_in_frame1, ) return dataset @@ -450,7 +462,10 @@ def _get_spring_dataset(self, is_train: bool, *args: str) -> Dataset: # Transforms copied from SEA-RAFT transform = ft.Compose( [ - ft.ToTensor(device=device, fp16=self.train_transform_fp16), + ft.ToTensor( + device=device, + fp16=self.train_transform_fp16, + ), ft.RandomScaleAndCrop((cy, cx), (0.0, 0.2), (-0.2, 0.2)), ft.ColorJitter(0.4, 0.4, 0.4, 0.5 / 3.14, 0.2), ft.GaussianNoise(0.02), @@ -475,6 +490,7 @@ def _get_spring_dataset(self, is_train: bool, *args: str) -> Dataset: get_flow=True, get_disparity=True, get_intrinsics=True, + disparity2_in_frame1=self.disparity2_in_frame1, robust_mode=robust_mode, robust_root_dir=self.robust_spring_root_dir, ) @@ -487,6 +503,9 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset: pass_names = ["clean", "final"] split = "trainval" sintel_transform = False + add_reverse = False + get_right_disparity = False + get_right_flow = False for v in args: if v in ["clean", "final"]: pass_names = [v] @@ -494,6 +513,12 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset: split = v elif v == "sinteltransform": sintel_transform = True + elif v == "rev": + add_reverse = True + elif v == "rdisp": + get_right_disparity = True + elif v == "rflow": + get_right_flow = True else: raise ValueError(f"Invalid arg: {v}") @@ -521,7 +546,10 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset: major_scale = (-0.4, 0.8) transform = ft.Compose( [ - ft.ToTensor(device=device, fp16=self.train_transform_fp16), + ft.ToTensor( + device=device, + fp16=self.train_transform_fp16, + ), ft.RandomScaleAndCrop((cy, cx), major_scale, (-0.2, 0.2)), ft.ColorJitter(0.4, 0.4, 0.4, 0.5 / 3.14, 0.2), ft.GaussianNoise(0.02), @@ -542,5 +570,9 @@ def _get_things_dataset(self, is_train: bool, *args: str) -> Dataset: get_flow=True, get_disparity=True, get_intrinsics=True, + add_reverse=add_reverse, + get_right_disparity=get_right_disparity, + get_right_flow=get_right_flow, + disparity2_in_frame1=self.disparity2_in_frame1, ) return dataset diff --git a/roco_spring_devkit/common/data/scene_flow_datasets.py b/roco_spring_devkit/common/data/scene_flow_datasets.py index 43d6df8..1e226f5 100644 --- a/roco_spring_devkit/common/data/scene_flow_datasets.py +++ b/roco_spring_devkit/common/data/scene_flow_datasets.py @@ -18,7 +18,7 @@ import math from pathlib import Path -from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union import cv2 as cv from einops import rearrange @@ -30,6 +30,21 @@ THIS_DIR = Path(__file__).resolve().parent +# Global registry informing, for each dataset, whether the second-frame disparity that it stores on disk is already +# expressed on the pixel grid of the FIRST frame (the convention used by the KITTI and Spring scene flow benchmarks, and +# the convention expected by the models of this repository), or on the pixel grid of the second frame (which requires +# pulling it back through the forward optical flow before it can be used). +# +# The keys are the ``dataset_name`` of each dataset class. Prefer reading it through +# ``BaseSceneFlowDataset.native_disparity2_in_frame1`` (a class attribute), which is what actually drives the loading. +# This dict is kept for code that only has the dataset name available (e.g. evaluation scripts). +NATIVE_DISPARITY2_IN_FRAME1 = { + "FlyingThings3D": False, + "Sintel": False, + "KITTI_2015": True, + "Spring": True, +} + class DatasetError(RuntimeError): """Exception raised for errors in the dataset loading.""" @@ -39,6 +54,50 @@ def __init__(self, message: str) -> None: super().__init__(self.message) +def _native_disparity2_in_frame1(dataset_name: str) -> bool: + """Look up NATIVE_DISPARITY2_IN_FRAME1 tolerating the suffixes some datasets add to their name. + + For example, SintelDataset names itself ``Sintel_clean_final``, depending on the loaded passes. + + Parameters + ---------- + dataset_name : str + The ``dataset_name`` of the dataset. + + Returns + ------- + bool + Whether the second-frame disparity of this dataset is natively stored on first-frame coordinates. Defaults to + False for unknown datasets, which is the more common convention. + """ + for name, in_frame1 in NATIVE_DISPARITY2_IN_FRAME1.items(): + if dataset_name == name or dataset_name.startswith(f"{name}_"): + return in_frame1 + return False + + +def _hwc_to_nchw(array: np.ndarray) -> torch.Tensor: + """Convert a single HWC numpy array into a float32 1CHW torch tensor.""" + return torch.from_numpy(np.ascontiguousarray(array)).float().permute(2, 0, 1)[None] + + +def _nchw_to_hwc(tensor: torch.Tensor, dtype: np.dtype) -> np.ndarray: + """Convert a 1CHW torch tensor back into an HWC numpy array with the given dtype.""" + array = tensor[0].permute(1, 2, 0).numpy() + if np.issubdtype(dtype, np.integer): + array = np.round(array) + return array.astype(dtype, copy=False) + + +def _valid_max_value(dtype: np.dtype) -> float: + """Value that represents 'valid' in a mask of the given dtype. + + The datasets store the validity masks as uint8 with 255 for valid pixels (ToTensor later divides them by 255), while + float masks use 1.0 directly. + """ + return 255.0 if np.issubdtype(dtype, np.integer) else 1.0 + + class BaseSceneFlowDataset(Dataset): """Manage scene flow dataset loading. @@ -77,8 +136,16 @@ class BaseSceneFlowDataset(Dataset): Each element of the main list is a list of baselines, which should be in the same order as the images in img_paths. metadata : list[Any] Some metadata for each input. It can include anything. A good recommendation would be to put a dict with the metadata. + native_disparity2_in_frame1 : bool + Class attribute telling whether the second-frame disparity stored on disk by this dataset is already expressed on + the pixel grid of the first frame. Concrete datasets must override it when that is the case (e.g. KITTI 2015). + See NATIVE_DISPARITY2_IN_FRAME1. """ + # Whether disp_paths[i][1] is already expressed on the pixel grid of the first frame. Overridden by the concrete + # datasets that use that convention. + native_disparity2_in_frame1: bool = False + def __init__( self, dataset_name: str, @@ -94,6 +161,7 @@ def __init__( get_intrinsics: bool = True, get_valid_mask: bool = True, get_meta: bool = True, + disparity2_in_frame1: bool = True, ) -> None: """Initialize BaseFlowDataset. @@ -132,6 +200,12 @@ def __init__( Whether to get the occluded version of the inputs. get_meta : bool, default True Whether to get metadata. + disparity2_in_frame1 : bool, default True + Whether the returned second-frame disparity ('disparities'[1]) should be expressed on the pixel grid of the + first frame. This is the convention of the KITTI scene flow benchmarks, and the one expected by the + models of this repository. When the dataset natively stores it on the pixel grid of the second frame (see + native_disparity2_in_frame1), it is pulled back through the forward optical flow, and the corresponding valid + mask is updated to discard the pixels that become invalid because of the warping. """ self.dataset_name = dataset_name self.split_name = split_name @@ -146,6 +220,39 @@ def __init__( self.get_valid_mask = get_valid_mask self.get_right_flow = get_right_flow self.get_meta = get_meta + self.disparity2_in_frame1 = disparity2_in_frame1 + + # True when this sample loader has to warp the second-frame disparity to first-frame coordinates itself. + self.must_pullback_disparity2 = ( + self.get_disparity + and self.get_flow + and self.disparity2_in_frame1 + and not self.native_disparity2_in_frame1 + ) + + if ( + self.get_disparity + and self.get_flow + and not self.disparity2_in_frame1 + and self.native_disparity2_in_frame1 + ): + logger.warning( + "{}: --disparity2_in_frame1 is False, but this dataset natively stores the second-frame disparity on " + "first-frame coordinates, and the operation cannot be inverted. The disparity will be returned on " + "first-frame coordinates anyway.", + dataset_name, + ) + self.disparity2_in_frame1 = True + + if ( + self.must_pullback_disparity2 + and self.get_right_disparity + and not self.get_right_flow + ): + raise ValueError( + f"{dataset_name}: get_right_disparity=True and disparity2_in_frame1=True require get_right_flow=True, " + "because the right second-frame disparity has to be pulled back through the right optical flow." + ) self.img_paths = [] self.img_r_paths = [] @@ -155,6 +262,8 @@ def __init__( self.flow_b_r_paths = [] self.disp_paths = [] self.disp_r_paths = [] + self.occ_paths = [] + self.occ_r_paths = [] self.intrinsics = [] self.baselines = [] self.metadata = [] @@ -223,6 +332,9 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 if self.get_valid_mask: inputs["valid_flows"] = valids + if index < len(self.occ_paths) and len(self.occ_paths[index]) > 0: + inputs["occlusions"] = self._get_occlusion_masks(self.occ_paths[index]) + if self.get_backward_flow: if index < len(self.flow_b_paths): inputs["flows_b"], valids_b = self._get_flows_and_valids( @@ -241,6 +353,11 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 if self.get_valid_mask: inputs["valid_flows_right"] = valids_r + if index < len(self.occ_r_paths) and len(self.occ_r_paths[index]) > 0: + inputs["occlusions_right"] = self._get_occlusion_masks( + self.occ_r_paths[index] + ) + if self.get_backward_flow: if index < len(self.flow_b_r_paths): inputs["flows_b_right"], valids_b_r = self._get_flows_and_valids( @@ -250,6 +367,13 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 if self.get_valid_mask: inputs["valid_flows_b_right"] = valids_b_r + self._align_disparity2_to_frame1(inputs) + # When disp2 is already on frame1 coordinates (native_disparity2_in_frame1=True), + # the pullback above is a no-op. But the disp2 file may not mark occluded pixels as + # invalid, so we use the forward-backward flow consistency check to find and + # invalidate them. Datasets without a backward flow simply skip the check. + self._invalidate_occluded_disparity2(inputs) + if self.transform is not None: inputs = self.transform(inputs) @@ -257,6 +381,7 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 inputs["meta"] = { "dataset_name": self.dataset_name, "split_name": self.split_name, + "disparity2_in_frame1": self.disparity2_in_frame1, } if index < len(self.metadata): inputs["meta"].update(self.metadata[index]) @@ -266,6 +391,157 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 def __len__(self) -> int: return len(self.img_paths) + def _align_disparity2_to_frame1(self, inputs: Dict[str, Any]) -> None: + """Warp the second-frame disparity onto the first-frame pixel grid, in-place. + + Does nothing when the dataset already stores it on first-frame coordinates, or when the caller asked for the + native convention. Both the left and the right views are converted, each one using its own optical flow. + + This must run before self.transform, so that all the spatial augmentations see a disparity map that is indexed by + first-frame pixels, exactly like 'disparities'[0] and 'flows'[0]. In particular, it keeps the left/right swap + performed by RandomFlip on a horizontal flip consistent. + + Parameters + ---------- + inputs : Dict[str, Any] + The sample being loaded, with the entries still stored as lists of HWC numpy arrays. + + Raises + ------ + NotImplementedError + If the sample has more than two disparity maps. Bringing the disparity of frame i > 2 to first-frame + coordinates would require composing the intermediate optical flows, which is not implemented. + + See Also + -------- + ptlflow.utils.stereo_utils.pullback_disparity2 + """ + if not self.must_pullback_disparity2: + return + + for disp_key, valid_disp_key, flow_key, valid_flow_key, flow_b_key, occ_key in ( + ( + "disparities", + "valid_disparities", + "flows", + "valid_flows", + "flows_b", + "occlusions", + ), + ( + "disparities_right", + "valid_disparities_right", + "flows_right", + "valid_flows_right", + "flows_b_right", + "occlusions_right", + ), + ): + disparities = inputs.get(disp_key) + flows = inputs.get(flow_key) + if disparities is None or flows is None or len(disparities) < 2: + continue + + if len(disparities) > 2: + raise NotImplementedError( + f"{self.dataset_name}: converting the disparity to first-frame coordinates is only implemented for " + f"sequences of two frames, but {len(disparities)} disparity maps were loaded. Use " + "disparity2_in_frame1=False if you want the raw per-frame disparities." + ) + + valid_disparities = inputs.get(valid_disp_key) + valid_flows = inputs.get(valid_flow_key) + + # The backward flow (frame 2 -> frame 1), if it was loaded, is used to run a forward-backward consistency + # check that invalidates the "ghost" first-frame pixels whose frame-2 correspondence is occupied by a + # different moving point. FlyingThings3D always loads it (get_backward_flow defaults to True); datasets + # without a backward flow simply skip the check. + flows_b = inputs.get(flow_b_key) + backward_flow = _hwc_to_nchw(flows_b[0]) if flows_b is not None else None + + # Datasets without a backward flow (e.g. Sintel) may ship explicit occlusion masks instead. The mask + # marks first-frame pixels that are occluded in frame 2, which are exactly the ghost pixels. + occs = inputs.get(occ_key) + occlusion_mask = _hwc_to_nchw(occs[0]) if occs is not None else None + + disparity2, valid_disparity2 = stereo_utils.pullback_disparity2( + _hwc_to_nchw(disparities[1]), + ( + _hwc_to_nchw(valid_disparities[1]) + if valid_disparities is not None + else None + ), + _hwc_to_nchw(flows[0]), + _hwc_to_nchw(valid_flows[0]) if valid_flows is not None else None, + backward_flow=backward_flow, + occlusion_mask=occlusion_mask, + ) + + disparities[1] = _nchw_to_hwc(disparity2, disparities[1].dtype) + if valid_disparities is not None: + valid_disparities[1] = _nchw_to_hwc( + valid_disparity2 * _valid_max_value(valid_disparities[1].dtype), + valid_disparities[1].dtype, + ) + + def _invalidate_occluded_disparity2(self, inputs: Dict[str, Any]) -> None: + """Mark occluded first-frame pixels as invalid in the disp2 valid mask, without warping. + + This is used when the dataset already stores disp2 on the first-frame pixel grid + (``native_disparity2_in_frame1`` is True, so no pullback is needed) but does not mark + the occluded pixels as invalid in the file. A forward-backward flow consistency check + is used to find them, similar to what :func:`pullback_disparity2` does for datasets + that need the warp. + + Both the left and the right views are processed, each one using its own forward and + backward flow. Nothing happens if the backward flow was not loaded. + + Parameters + ---------- + inputs : Dict[str, Any] + The sample being loaded, with the entries still stored as lists of HWC numpy arrays. + """ + if not (self.get_flow and self.get_backward_flow and self.get_disparity): + return + + for disp_key, valid_disp_key, flow_key, flow_b_key in ( + ("disparities", "valid_disparities", "flows", "flows_b"), + ( + "disparities_right", + "valid_disparities_right", + "flows_right", + "flows_b_right", + ), + ): + disparities = inputs.get(disp_key) + flows = inputs.get(flow_key) + flows_b = inputs.get(flow_b_key) + valid_disparities = inputs.get(valid_disp_key) + if ( + disparities is None + or flows is None + or flows_b is None + or valid_disparities is None + or len(disparities) < 2 + ): + continue + + occ_mask = stereo_utils.compute_fb_occlusion_mask( + _hwc_to_nchw(flows[0]), + _hwc_to_nchw(flows_b[0]), + ) + occ_hwc = _nchw_to_hwc(occ_mask, valid_disparities[1].dtype) + # A pixel is valid only if it was valid before AND is not occluded. + valid_disparities[1] = ( + ( + valid_disparities[1] + >= _valid_max_value(valid_disparities[1].dtype) / 2 + ) + & (occ_hwc < 0.5) + ).astype(valid_disparities[1].dtype) * _valid_max_value( + valid_disparities[1].dtype + ) + def _get_flows_and_valids( self, flow_paths: Sequence[str], @@ -318,9 +594,24 @@ def _get_disparities_and_valids( disparities.append(disp) return disparities, valids + def _get_occlusion_masks(self, occ_paths: Sequence[str]) -> List[np.ndarray]: + """Read occlusion masks from disk as HWC float arrays in [0, 1]. + + A value of 1.0 means the first-frame pixel is occluded in the second frame. Datasets typically store these as + uint8 PNGs with 0 (visible) and 255 (occluded). + """ + masks = [] + for path in occ_paths: + mask = cv.imread(str(path), cv.IMREAD_GRAYSCALE) + if mask is None: + raise FileNotFoundError(f"Occlusion mask not found: {path}") + mask = (mask >= 128).astype(np.float32)[..., None] + masks.append(mask) + return masks + def _log_status(self) -> None: if self.__len__() == 0: - logger.error( + logger.warning( "No samples were found for {} dataset. Be sure to update the dataset path in datasets.yml, " "or provide the path by the argument --[dataset_name]_root_dir.", self.dataset_name, @@ -365,6 +656,9 @@ class FlyingThings3DDataset(BaseSceneFlowDataset): Note that this only works for the complete FlyingThings3D dataset. """ + # FlyingThings3D stores the disparity of each frame on its own pixel grid. + native_disparity2_in_frame1 = False + def __init__( self, root_dir: str, @@ -381,6 +675,7 @@ def __init__( get_intrinsics: bool = True, get_valid_mask: bool = True, get_meta: bool = True, + disparity2_in_frame1: bool = True, sequence_length: int = 2, sequence_position: str = "first", add_reverse: bool = False, @@ -440,6 +735,7 @@ def __init__( get_intrinsics=get_intrinsics, get_valid_mask=get_valid_mask, get_meta=get_meta, + disparity2_in_frame1=disparity2_in_frame1, ) self.root_dir = root_dir self.pass_names = pass_names @@ -638,12 +934,15 @@ def __init__( flow_r_paths[i : i + self.sequence_length - 1] ) if len(flow_b_paths) > 0: + # The backward flow of forward flow IF_i (frame i -> i+1) is IP_{i+1} (frame i+1 -> + # i), so the into_past list is shifted by one column relative to the into_future + # list. self.flow_b_paths.append( - flow_b_paths[i : i + self.sequence_length - 1] + flow_b_paths[i + 1 : i + self.sequence_length] ) if len(flow_b_r_paths) > 0: self.flow_b_r_paths.append( - flow_b_r_paths[i : i + self.sequence_length - 1] + flow_b_r_paths[i + 1 : i + self.sequence_length] ) self.metadata.append( @@ -737,6 +1036,9 @@ def _check_folders(self) -> None: class KittiDataset(BaseSceneFlowDataset): """Handle the KITTI dataset.""" + # KITTI 2015 'disp_occ_1' is the second-frame disparity already mapped into the first-frame pixel grid. + native_disparity2_in_frame1 = True + def __init__( # noqa: C901 self, root_dir_2015: Optional[str] = None, @@ -749,6 +1051,7 @@ def __init__( # noqa: C901 get_intrinsics: bool = True, get_valid_mask: bool = True, get_meta: bool = True, + disparity2_in_frame1: bool = True, ) -> None: """Initialize KittiDataset. @@ -790,6 +1093,7 @@ def __init__( # noqa: C901 get_intrinsics=get_intrinsics, get_valid_mask=get_valid_mask, get_meta=get_meta, + disparity2_in_frame1=disparity2_in_frame1, ) self.root_dir = root_dir_2015 self.split = split @@ -1021,6 +1325,9 @@ def _check_folders(self) -> None: class SintelDataset(BaseSceneFlowDataset): """Handle the MPI Sintel dataset.""" + # Sintel stores the disparity of each frame on its own pixel grid. + native_disparity2_in_frame1 = False + def __init__( # noqa: C901 self, root_dir: str, @@ -1033,7 +1340,9 @@ def __init__( # noqa: C901 max_disparity: float = 10000.0, get_intrinsics: bool = True, get_valid_mask: bool = True, + get_occlusion_mask: bool = True, get_meta: bool = True, + disparity2_in_frame1: bool = True, sequence_length: int = 2, sequence_position: str = "first", ) -> None: @@ -1057,7 +1366,8 @@ def __init__( # noqa: C901 get_valid_mask : bool, default True Whether to get or generate valid masks. get_occlusion_mask : bool, default True - Whether to get occlusion masks. + Whether to get occlusion masks. Sintel ships occlusion masks for the forward flow, which are used to + invalidate the "ghost" pixels when warping the second-frame disparity to first-frame coordinates. get_meta : bool, default True Whether to get metadata. sequence_length : int, default 2 @@ -1084,6 +1394,7 @@ def __init__( # noqa: C901 get_intrinsics=get_intrinsics, get_valid_mask=get_valid_mask, get_meta=get_meta, + disparity2_in_frame1=disparity2_in_frame1, ) self.root_dir = root_dir self.split = split @@ -1091,6 +1402,7 @@ def __init__( # noqa: C901 self.sequence_length = sequence_length if get_flow else 1 self.sequence_position = sequence_position self.disp_format = "sintel" + self.get_occlusion_mask = get_occlusion_mask left_image_dir_suffix = "_left" if self.get_disparity else "" @@ -1146,6 +1458,7 @@ def __init__( # noqa: C901 ) flow_paths = [] disp_paths = [] + occ_paths = [] if split != "test": if (Path(self.root_dir) / split_dir / "flow").exists(): flow_paths = sorted( @@ -1159,6 +1472,28 @@ def __init__( # noqa: C901 assert len(image_paths) - 1 == len( flow_paths ), f"{passd}, {seq_name}: {len(image_paths)-1} vs {len(flow_paths)}" + + if self.get_occlusion_mask and self.must_pullback_disparity2: + # Sintel ships occlusion masks for the forward flow. Try "occlusions" first, then fall + # back to "occlusions_flow". + occ_dir = None + for candidate in ("occlusions", "occlusions_flow"): + if ( + Path(self.root_dir) / split_dir / candidate + ).exists(): + occ_dir = ( + Path(self.root_dir) / split_dir / candidate + ) + break + if occ_dir is not None: + occ_paths = sorted((occ_dir / seq_name).glob("*.png")) + occ_paths = self._extend_paths_list( + occ_paths, sequence_length, sequence_position + ) + assert len(flow_paths) == len( + occ_paths + ), f"{passd}, {seq_name}: {len(flow_paths)} occlusion masks vs {len(occ_paths)} flow files" + if (Path(self.root_dir) / split_dir / "disparities").exists(): disp_paths = sorted( ( @@ -1210,6 +1545,11 @@ def __init__( # noqa: C901 flow_paths[i : i + self.sequence_length - 1] ) + if len(occ_paths) > 0: + self.occ_paths.append( + occ_paths[i : i + self.sequence_length - 1] + ) + if len(disp_paths) > 0: self.disp_paths.append(disp_paths[i : i + self.sequence_length]) @@ -1237,6 +1577,10 @@ def __init__( # noqa: C901 assert len(self.img_paths) == len( self.flow_paths ), f"{len(self.img_paths)} vs {len(self.flow_paths)}" + if self.get_occlusion_mask and self.must_pullback_disparity2: + assert len(self.img_paths) == len( + self.occ_paths + ), f"{len(self.img_paths)} vs {len(self.occ_paths)} occlusion masks" self._log_status() @@ -1292,6 +1636,19 @@ def _check_folders(self) -> None: ) has_errors = True + if self.get_occlusion_mask: + found = False + for candidate in ("occlusions", "occlusions_flow"): + if (Path(self.root_dir) / self.split_dir / candidate).exists(): + found = True + break + if not found: + logger.error( + "The 'occlusions' or 'occlusions_flow' folder was not found in the MPI-Sintel dataset. " + "Be sure to download the optical flow occlusion data from the official website." + ) + has_errors = True + if has_errors: raise DatasetError( "Some required folders were not found in the MPI-Sintel dataset. " @@ -1302,6 +1659,10 @@ def _check_folders(self) -> None: class SpringDataset(BaseSceneFlowDataset): """Handle the Spring dataset.""" + # Spring 'disp2__' is the second-frame disparity already mapped into the + # first-frame pixel grid (following the KITTI convention). + native_disparity2_in_frame1 = True + def __init__( # noqa: C901 self, root_dir: str, @@ -1317,6 +1678,7 @@ def __init__( # noqa: C901 get_intrinsics: bool = True, get_valid_mask: bool = True, get_meta: bool = True, + disparity2_in_frame1: bool = True, sequence_length: int = 2, sequence_position: str = "first", subsample: bool = False, @@ -1379,6 +1741,7 @@ def __init__( # noqa: C901 get_valid_mask=get_valid_mask, get_intrinsics=get_intrinsics, get_meta=get_meta, + disparity2_in_frame1=disparity2_in_frame1, ) self.root_dir = root_dir self.split = split @@ -1842,6 +2205,12 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 if self.get_valid_mask: inputs["valid_flows_b_right"] = valids_b + self._align_disparity2_to_frame1(inputs) + # Spring's disp2 is already on frame1 coordinates, so the pullback above is a no-op. + # But the disp2 file does not mark occluded pixels as invalid, so we use the + # forward-backward flow consistency check to find and invalidate them. + self._invalidate_occluded_disparity2(inputs) + if self.subsample: if self.get_flow and "flows" in inputs: inputs["flows"] = [f[::2, ::2] for f in inputs["flows"]] @@ -1919,6 +2288,7 @@ def __getitem__(self, index: int) -> Dict[str, torch.Tensor]: # noqa: C901 inputs["meta"] = { "dataset_name": self.dataset_name, "split_name": self.split_name, + "disparity2_in_frame1": self.disparity2_in_frame1, } if index < len(self.metadata): inputs["meta"].update(self.metadata[index]) diff --git a/roco_spring_devkit/common/data/scene_flow_transforms.py b/roco_spring_devkit/common/data/scene_flow_transforms.py index aa5611b..1ec9863 100644 --- a/roco_spring_devkit/common/data/scene_flow_transforms.py +++ b/roco_spring_devkit/common/data/scene_flow_transforms.py @@ -73,6 +73,10 @@ class ToTensor(object): """Converts a 4D numpy.ndarray or a list of 3D numpy.ndarrays into a 4D torch.Tensor. If an input is of type uint8, then it is converted to float and its values are divided by 255. + + Non-image entries, such as camera intrinsics (a list of 3x3 matrices) and baselines (a list of scalars), + are also converted into torch tensors. Intrinsics become a tensor of shape (N, 3, 3) and baselines a + tensor of shape (N,). """ def __init__( @@ -81,8 +85,8 @@ def __init__( device: Union[str, torch.device] = "cpu", use_keys: Optional[Union[KeysView, Sequence[str]]] = None, ignore_keys: Optional[Union[KeysView, Sequence[str]]] = None, - intrinsics_key: Optional[str] = "intrinsics", - baselines_key: Optional[str] = "baselines", + intrinsics_key: str = "intrinsics", + baselines_key: str = "baselines", ) -> None: """Initialize ToTensor. @@ -97,12 +101,12 @@ def __init__( except the keys that are listed in ignore_keys. ignore_keys : Optional[Union[KeysView, Sequence[str]]], optional If use_keys is None, the these keys are NOT transformed by this operation. - intrinsics_key : Optional[str], default 'intrinsics' - The key of the camera intrinsics. They are converted directly to a - tensor (N33), without going through the image reshape logic. - baselines_key : Optional[str], default 'baselines' - The key of the stereo baselines. They are converted directly to a - tensor (N), without going through the image reshape logic. + intrinsics_key : str, default 'intrinsics' + Name of the entry containing the camera intrinsics. It is converted into a (N, 3, 3) tensor, instead of + being treated as an image. + baselines_key : str, default 'baselines' + Name of the entry containing the stereo baselines. It is converted into a (N,) tensor, instead of being + treated as an image. """ self.dtype = torch.float16 if fp16 else torch.float32 self.device = device @@ -127,6 +131,9 @@ def __call__( The inputs transformed by this operation. """ ignore_set = set(self.ignore_keys) if self.ignore_keys is not None else set() + + # Convert non-image entries (intrinsics, baselines) into tensors directly, + # bypassing the image conversion pipeline below. for special_key in (self.intrinsics_key, self.baselines_key): if ( special_key is not None @@ -136,21 +143,28 @@ def __call__( ): v = inputs[special_key] if isinstance(v, torch.Tensor): - inputs[special_key] = v.to(device=self.device, dtype=self.dtype) + v = v.to(device=self.device, dtype=self.dtype) else: - inputs[special_key] = torch.as_tensor( - np.asarray(v, dtype=np.float32), - device=self.device, - dtype=self.dtype, - ) + if isinstance(v, (list, tuple)): + v = np.array(v, dtype=np.float32) + else: + v = np.asarray(v, dtype=np.float32) + v = torch.from_numpy(v).to(device=self.device, dtype=self.dtype) + inputs[special_key] = v + valid_keys = _get_valid_keys(inputs.keys(), self.use_keys, self.ignore_keys) + # The special keys were already handled above, so exclude them from the generic image conversion. valid_keys = [ k - for k in _get_valid_keys(inputs.keys(), self.use_keys, self.ignore_keys) - if k not in (self.intrinsics_key, self.baselines_key) + for k in valid_keys + if k != self.intrinsics_key and k != self.baselines_key ] for k in valid_keys: v = inputs[k] + + if len(v) == 0: + continue + if isinstance(v, list) or isinstance(v, tuple): v = np.stack(v) if len(v.shape) == 3: @@ -164,8 +178,7 @@ def __call__( if v.dtype == np.uint8: v = v.astype(np.float32) / 255.0 - if len(v.shape) == 4: - v = v.transpose(0, 3, 1, 2) + v = v.transpose(0, 3, 1, 2) inputs[k] = torch.from_numpy(v).to(device=self.device, dtype=self.dtype) return inputs @@ -176,13 +189,22 @@ class CenterCrop(object): def __init__( self, crop_size: Optional[Tuple[int, int]] = None, - occlusion_keys: Union[KeysView, Sequence[str]] = ("occs", "occs_b"), - valid_key: str = "valids", - ignore_keys: Optional[Sequence[str]] = None, - disparities_keys: Union[KeysView, Sequence[str]] = ( - "disparities", - "disparities_r", + occlusion_keys: Union[KeysView, Sequence[str]] = ( + "occs", + "occs_b", + "occs_right", + "occs_b_right", ), + flow_keys: Union[KeysView, Sequence[str]] = ( + "flows", + "flows_b", + "flows_right", + "flows_b_right", + ), + valid_key: str = "valid_flows", + intrinsics_key: Optional[str] = "intrinsics", + baselines_key: Optional[str] = "baselines", + ignore_keys: Optional[Sequence[str]] = None, ) -> None: """Initialize CenterCrop. @@ -192,19 +214,28 @@ def __init__( If provided, crop the inputs to this size (h, w). occlusion_keys : Union[KeysView, Sequence[str]], default ['occs', 'occs_b'] Indicate which of the input keys correspond to occlusion mask tensors. + flow_keys : Union[KeysView, Sequence[str]], default ['flows', 'flows_b'] + Indicate which of the input keys correspond to optical flow tensors. Used to update the occlusion + masks for out-of-bounds flows. valid_keys : str, default 'valids' The name of the key in inputs that contains the binary mask indicating which pixels are valid. Only used when sparse=True. + intrinsics_key : Optional[str], default 'intrinsics' + Name of the entry containing the camera intrinsics. If present, its principal point is shifted to + match the crop. If None, intrinsics are not adjusted. + baselines_key : Optional[str], default 'baselines' + Name of the entry containing the stereo baselines. It is left unchanged, only protected from the + spatial crop. ignore_keys : Optional[Sequence[str]], optional - If not None, these keys are NOT transformed by this operation. - disparities_keys : Union[KeysView, Sequence[str]], default ['disparities', 'disparities_r'] - Indicate which of the input keys correspond to disparity tensors. + These keys are NOT transformed by this operation. """ self.crop_size = crop_size self.occlusion_keys = list(occlusion_keys) + self.flow_keys = list(flow_keys) self.valid_key = valid_key + self.intrinsics_key = intrinsics_key + self.baselines_key = baselines_key self.ignore_keys = ignore_keys - self.disparities_keys = list(disparities_keys) def __call__( # noqa: C901 self, inputs: Dict[str, torch.Tensor] @@ -230,32 +261,49 @@ def __call__( # noqa: C901 y_crop = (h - self.crop_size[0]) // 2 x_crop = (w - self.crop_size[1]) // 2 + spatial_ignore = self._spatial_ignore_keys() for k, v in inputs.items(): - if self.ignore_keys is None or k not in self.ignore_keys: - v = v[ - :, - :, - y_crop : y_crop + self.crop_size[0], - x_crop : x_crop + self.crop_size[1], - ] + if self.ignore_keys is not None and k in self.ignore_keys: + continue + if k in spatial_ignore: + continue + v = v[ + :, + :, + y_crop : y_crop + self.crop_size[0], + x_crop : x_crop + self.crop_size[1], + ] inputs[k] = v - # Update occlusion masks for out-of-bounds disparities + if ( + self.intrinsics_key is not None + and self.intrinsics_key in inputs + and self.intrinsics_key not in (self.ignore_keys or []) + ): + inputs[self.intrinsics_key] = _adjust_intrinsics_for_crop( + inputs[self.intrinsics_key], x_crop, y_crop + ) + + # Update occlusion masks for out-of-bounds flows for k, v in inputs.items(): - if self.ignore_keys is None or k not in self.ignore_keys: - try: - i = self.occlusion_keys.index(k) - if ( - i < len(self.disparities_keys) - and self.disparities_keys[i] in inputs - ): - inputs[k] = _update_oob_disparities( - v, inputs[self.disparities_keys[i]] - ) - except ValueError: - pass + if self.ignore_keys is not None and k in self.ignore_keys: + continue + try: + i = self.occlusion_keys.index(k) + if i < len(self.flow_keys) and self.flow_keys[i] in inputs: + inputs[k] = _update_oob_flows(v, inputs[self.flow_keys[i]]) + except ValueError: + pass return inputs + def _spatial_ignore_keys(self) -> set: + keys = set() + if self.intrinsics_key is not None: + keys.add(self.intrinsics_key) + if self.baselines_key is not None: + keys.add(self.baselines_key) + return keys + class ColorJitter(tt.ColorJitter): """Randomly apply color transformations only to the images. @@ -273,7 +321,7 @@ def __init__( saturation: Union[float, Tuple[float, float]] = 0.0, hue: Union[float, Tuple[float, float]] = 0.0, asymmetric_prob: float = 0.0, - use_keys: Optional[Union[KeysView, Sequence[str]]] = ("images",), + use_keys: Optional[Union[KeysView, Sequence[str]]] = ("images", "images_right"), ignore_keys: Optional[Union[KeysView, Sequence[str]]] = None, ) -> None: """Initialize ColorJitter. @@ -291,9 +339,10 @@ def __init__( asymmetric_prob : float, default 0.0 Chance to apply an asymmetric transform, in which the parameters for transforming each image are sampled independently. - use_keys : Optional[Union[KeysView, Sequence[str]]], default ['images'] + use_keys : Optional[Union[KeysView, Sequence[str]]], default ['images', 'images_right'] If it is not None, then only elements with these keys will be transformed. Otherwise, all elements are transformed, - except the keys that are listed in ignore_keys. + except the keys that are listed in ignore_keys. Each key is jittered independently, so the left and right + stereo images receive different photometric transforms. ignore_keys : Optional[Union[KeysView, Sequence[str]]], optional If use_keys is None, the these keys are NOT transformed by this operation. """ @@ -334,7 +383,7 @@ class GaussianNoise(object): def __init__( self, stdev: float = 0.0, - use_keys: Optional[Union[KeysView, Sequence[str]]] = ("images",), + use_keys: Optional[Union[KeysView, Sequence[str]]] = ("images", "images_right"), ignore_keys: Optional[Union[KeysView, Sequence[str]]] = None, ) -> None: """Initialize GaussianNoise. @@ -345,7 +394,7 @@ def __init__( The maximum standard deviation of the gaussian noise. use_keys : Optional[Union[KeysView, Sequence[str]]], optional If it is not None, then only elements with these keys will be transformed. Otherwise, all elements are transformed, - except the keys that are listed in ignore_keys. + except the keys that are listed in ignore_keys. Each key receives an independently sampled noise level. ignore_keys : Optional[Union[KeysView, Sequence[str]]], optional If use_keys is None, the these keys are NOT transformed by this operation. """ @@ -474,190 +523,252 @@ def __call__(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: class RandomFlip(object): """Randomly horizontally and vertically flips the inputs. - If asymmetric_prob > 0, then each input of the sequence may be flipped differently. + For stereo (scene flow) inputs, a horizontal flip also swaps the left and right images, so that the disparity + sign is preserved. The camera intrinsics principal point is mirrored accordingly. Baselines are left unchanged. + Vertical flips mirror the principal point in y and leave the disparity sign untouched. """ def __init__( self, hflip_prob: float = 0.0, vflip_prob: float = 0.0, - asymmetric_prob: float = 0.0, use_keys: Optional[Union[KeysView, Sequence[str]]] = None, ignore_keys: Optional[Union[KeysView, Sequence[str]]] = None, image_keys: Union[KeysView, Sequence[str]] = ("images",), - disparities_keys: Union[KeysView, Sequence[str]] = ( - "disparities", - "disparities_r", + image_right_keys: Union[KeysView, Sequence[str]] = ("images_right",), + flow_keys: Union[KeysView, Sequence[str]] = ("flows", "flows_b"), + flow_right_keys: Union[KeysView, Sequence[str]] = ( + "flows_right", + "flows_b_right", + ), + valid_flow_keys: Union[KeysView, Sequence[str]] = ( + "valid_flows", + "valid_flows_b", + ), + valid_flow_right_keys: Union[KeysView, Sequence[str]] = ( + "valid_flows_right", + "valid_flows_b_right", ), + disparity_keys: Union[KeysView, Sequence[str]] = ("disparities",), + disparity_right_keys: Union[KeysView, Sequence[str]] = ("disparities_right",), + valid_disparity_keys: Union[KeysView, Sequence[str]] = ("valid_disparities",), + valid_disparity_right_keys: Union[KeysView, Sequence[str]] = ( + "valid_disparities_right", + ), + occlusion_keys: Union[KeysView, Sequence[str]] = ("occs", "occs_b"), + occlusion_right_keys: Union[KeysView, Sequence[str]] = ( + "occs_right", + "occs_b_right", + ), + mb_keys: Union[KeysView, Sequence[str]] = ("mbs", "mbs_b"), + mb_right_keys: Union[KeysView, Sequence[str]] = ("mbs_right", "mbs_b_right"), + intrinsics_key: Optional[str] = "intrinsics", + baselines_key: Optional[str] = "baselines", ) -> None: - """Initialize RandomFlip. - - Parameters - ---------- - hflip_prob : float, default 0.0 - Probability of applying a horizontal flip. - vflip_prob : float, default 0.0 - Probability of applying a vertical flip. - asymmetric_prob : float, default 0.0 - Chance to apply an asymmetric transform, in which the parameters for transforming each image are sampled - independently. - use_keys : Optional[Union[KeysView, Sequence[str]]], optional - If it is not None, then only elements with these keys will be transformed. Otherwise, all elements are transformed, - except the keys that are listed in ignore_keys. - ignore_keys : Optional[Union[KeysView, Sequence[str]]], optional - If use_keys is None, the these keys are NOT transformed by this operation. - image_keys : Union[KeysView, Sequence[str]], ['images'] - Indicate which of the input keys correspond to image tensors. - disparities_keys : Union[KeysView, Sequence[str]], ['disparities', 'disparities_r'] - Indicate which of the input keys correspond to disparity tensors. - """ + """Initialize RandomFlip.""" self.flip_probs = [hflip_prob, vflip_prob] - self.asymmetric_prob = asymmetric_prob self.use_keys = use_keys self.ignore_keys = ignore_keys - self.image_keys = list(image_keys) - self.disparities_keys = list(disparities_keys) - def __call__(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """Perform the transformation on the inputs. + # Convert all to lists for safe iteration and pairing + self.image_keys = list(image_keys) if image_keys else [] + self.image_right_keys = list(image_right_keys) if image_right_keys else [] + self.flow_keys = list(flow_keys) if flow_keys else [] + self.flow_right_keys = list(flow_right_keys) if flow_right_keys else [] + self.valid_flow_keys = list(valid_flow_keys) if valid_flow_keys else [] + self.valid_flow_right_keys = ( + list(valid_flow_right_keys) if valid_flow_right_keys else [] + ) + self.disparity_keys = list(disparity_keys) if disparity_keys else [] + self.disparity_right_keys = ( + list(disparity_right_keys) if disparity_right_keys else [] + ) + self.valid_disparity_keys = ( + list(valid_disparity_keys) if valid_disparity_keys else [] + ) + self.valid_disparity_right_keys = ( + list(valid_disparity_right_keys) if valid_disparity_right_keys else [] + ) + self.occlusion_keys = list(occlusion_keys) if occlusion_keys else [] + self.occlusion_right_keys = ( + list(occlusion_right_keys) if occlusion_right_keys else [] + ) + self.mb_keys = list(mb_keys) if mb_keys else [] + self.mb_right_keys = list(mb_right_keys) if mb_right_keys else [] - Parameters - ---------- - inputs : Dict[str, torch.Tensor] - Elements to be transformed. Each element is a 4D tensor NCHW. + self.intrinsics_key = intrinsics_key + self.baselines_key = baselines_key - Returns - ------- - Dict[str, torch.Tensor] - The inputs transformed by this operation. - """ + def __call__(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + """Perform the transformation on the inputs.""" + selected_keys = set( + _get_valid_keys(inputs.keys(), self.use_keys, self.ignore_keys) + ) + + # Intrinsics and baselines are not regular spatial tensors, so remove them from the generic flip loop. valid_keys = [ k - for k in _get_valid_keys(inputs.keys(), self.use_keys, self.ignore_keys) - if k not in ("intrinsics", "baselines") + for k in selected_keys + if k != self.intrinsics_key and k != self.baselines_key ] - height, width = inputs[self.image_keys[0]].shape[-2:] + for iorient in range(2): - if self.asymmetric_prob < 1e-5: - if random.random() < self.flip_probs[iorient]: - inputs = self._flip_inputs(inputs, iorient == 0, valid_keys) - if "intrinsics" in inputs: - inputs["intrinsics"] = _adjust_intrinsics_for_flip( - inputs["intrinsics"], iorient == 0, width, height - ) - else: - is_flips = [ - random.random() < self.flip_probs[iorient] - for _ in range(inputs[self.image_keys[0]].shape[0]) - ] - num_pairs = min( - inputs[self.disparities_keys[0]].shape[0], len(is_flips) - 1 - ) - for i in range(num_pairs): - if is_flips[i]: - inputs = self._flip_inputs( - inputs, iorient == 0, valid_keys, ibatch=i - ) - if is_flips[i] != is_flips[i + 1]: - for dk in self.disparities_keys: - if dk in inputs: - inputs[dk][i] = self._mirror_disparity( - inputs[dk][i], iorient == 0 - ) - if is_flips[-1]: - for ik in self.image_keys: - inputs = self._flip_inputs( - inputs, iorient == 0, valid_keys=[ik], ibatch=-1 - ) + is_hflip = iorient == 0 + if random.random() < self.flip_probs[iorient]: + + is_stereo = False + if is_hflip: + is_stereo = ( + len(self.image_keys) > 0 + and self.image_keys[0] in selected_keys + and len(self.image_right_keys) > 0 + and self.image_right_keys[0] in selected_keys + ) + + if is_stereo: + # Helper to check if any left key is missing its right counterpart (or vice versa) + def _is_unpaired(left_list, right_list): + max_len = max(len(left_list), len(right_list)) + for i in range(max_len): + lk = left_list[i] if i < len(left_list) else None + rk = right_list[i] if i < len(right_list) else None + has_l = lk is not None and lk in selected_keys + has_r = rk is not None and rk in selected_keys + if ( + has_l != has_r + ): # XOR check: if one is present, the other MUST be present + return True + return False + + if ( + _is_unpaired(self.flow_keys, self.flow_right_keys) + or _is_unpaired( + self.disparity_keys, self.disparity_right_keys + ) + or _is_unpaired( + self.valid_flow_keys, self.valid_flow_right_keys + ) + or _is_unpaired( + self.valid_disparity_keys, + self.valid_disparity_right_keys, + ) + or _is_unpaired( + self.occlusion_keys, self.occlusion_right_keys + ) + or _is_unpaired(self.mb_keys, self.mb_right_keys) + ): + continue # Safely abort the entire horizontal flip to preserve geometry + + inputs = self._flip_inputs(inputs, is_hflip, valid_keys) + + # Stereo: on horizontal flip swap the left and right sequences + if is_hflip and is_stereo: + inputs = self._swap_keys( + inputs, self.image_keys, self.image_right_keys, selected_keys + ) + inputs = self._swap_keys( + inputs, self.flow_keys, self.flow_right_keys, selected_keys + ) + inputs = self._swap_keys( + inputs, + self.disparity_keys, + self.disparity_right_keys, + selected_keys, + ) + inputs = self._swap_keys( + inputs, + self.valid_flow_keys, + self.valid_flow_right_keys, + selected_keys, + ) + inputs = self._swap_keys( + inputs, + self.valid_disparity_keys, + self.valid_disparity_right_keys, + selected_keys, + ) + inputs = self._swap_keys( + inputs, + self.occlusion_keys, + self.occlusion_right_keys, + selected_keys, + ) + inputs = self._swap_keys( + inputs, self.mb_keys, self.mb_right_keys, selected_keys + ) + + inputs = self._flip_intrinsics(inputs, is_hflip, selected_keys) return inputs - def _flip_inputs( + def _swap_keys( self, inputs: Dict[str, torch.Tensor], - is_hflip: bool, - valid_keys: Optional[Sequence[str]] = None, - ibatch: Optional[int] = None, + left_keys: list, + right_keys: list, + selected_keys: set, ) -> Dict[str, torch.Tensor]: - """Flips all inputs horizontally or vertically. + """Generic helper to swap pairs of left/right keys.""" + for lk, rk in zip(left_keys, right_keys): + if lk in selected_keys and rk in selected_keys: + tmp = inputs[lk] + inputs[lk] = inputs[rk] + inputs[rk] = tmp + return inputs - This function properly adjust the flow values after the flipping. + def _flip_intrinsics( + self, inputs: Dict[str, torch.Tensor], is_hflip: bool, selected_keys: set + ) -> Dict[str, torch.Tensor]: + """Mirrors the principal point of the camera intrinsics.""" + if self.intrinsics_key is None or self.intrinsics_key not in selected_keys: + return inputs + img_key = ( + self.image_keys[0] + if len(self.image_keys) > 0 and self.image_keys[0] in selected_keys + else None + ) + if img_key is None: + return inputs - Parameters - ---------- - inputs : Dict[str, torch.Tensor] - Elements to be flipped. Each element is a 4D tensor NCHW. - is_hflip : bool - If True, performs a horizontal flip, otherwise, performs a vertical flip. - valid_keys : Optional[Sequence[str]], optional - If it is not None, then only elements with these keys will be transformed. Otherwise, all elements are transformed. - ibatch : Optional[int], optional - If ibatch is specified, then only one element of the batch is flipped. + h, w = inputs[img_key].shape[-2:] + inputs[self.intrinsics_key] = _adjust_intrinsics_for_flip( + inputs[self.intrinsics_key], is_hflip, w, h + ) + return inputs - Returns - ------- - Dict[str, torch.Tensor] - The inputs flipped by this operation. - """ - if is_hflip: - iinp = 3 - iflow = 0 - else: - iinp = 2 - iflow = 1 + def _flip_inputs( + self, + inputs: Dict[str, torch.Tensor], + is_hflip: bool, + valid_keys: Sequence[str], + ) -> Dict[str, torch.Tensor]: + """Flips all inputs spatially and negates directional vectors accordingly.""" + iinp = 3 if is_hflip else 2 + iflow = 0 if is_hflip else 1 + + # Collect grouped keys + all_flow_keys = set(self.flow_keys + self.flow_right_keys) + all_disparity_keys = set(self.disparity_keys + self.disparity_right_keys) - if valid_keys is None: - valid_keys = list(inputs.keys()) for k in valid_keys: - # Flows and disparities are correspondence vectors: the component - # along the flipped axis must be negated so they stay synchronized - # with the flipped images. A 1-channel (horizontal) disparity is - # only affected by a horizontal flip. - is_correspondence = ("flows" in k and "valid" not in k) or ( - k in self.disparities_keys - ) - if ibatch is None: - inputs[k] = torch.flip(inputs[k], [iinp]) - if is_correspondence and inputs[k].shape[1] > iflow: - inputs[k][:, iflow] *= -1 - else: - inputs[k][ibatch] = torch.flip(inputs[k][ibatch], [iinp - 1]) - if is_correspondence and inputs[k].shape[1] > iflow: - inputs[k][ibatch, iflow] *= -1 - return inputs + if k not in inputs: + continue - def _mirror_disparity( - self, disparity: torch.Tensor, is_hflip: bool - ) -> torch.Tensor: - """Reflects the disparity along the center line of the image. + inputs[k] = torch.flip(inputs[k], [iinp]) - This function is used when an asymmetric flip happens (one image flips, but the next does not, or vice-versa). + # 1. Flow Negation: Negate u for hflip, v for vflip + if k in all_flow_keys: + if inputs[k].shape[1] > iflow: + inputs[k][:, iflow] *= -1 - Parameters - ---------- - disparity : torch.Tensor - A 3D tensor CHW. - is_hflip : bool - If True, performs a horizontal flip, otherwise, performs a vertical flip. + # 2. 2D Disparity Negation + # Horizontal component (dx) NEVER changes sign under any flip + # Vertical component (dy) ALWAYS changes sign under any flip (V-flip directly inverts Y; H-flip swaps cameras) + if k in all_disparity_keys: + if inputs[k].shape[1] > 1: # If dataset provides 2D disparity + inputs[k][:, 1] *= -1 - Returns - ------- - torch.Tensor - The mirrored disparity. - """ - grid = torch.meshgrid( - torch.arange(disparity.shape[1]), - torch.arange(disparity.shape[2]), - indexing="ij", - ) - grid = torch.stack(grid[::-1]).float() - if is_hflip: - mean_coord = (disparity.shape[2] - 1) / 2.0 - disparity[0] = 2 * (mean_coord - grid[0]) - disparity[0] - else: - mean_coord = (disparity.shape[1] - 1) / 2.0 - disparity[1] = 2 * (mean_coord - grid[1]) - disparity[1] - return disparity + return inputs class RandomScaleAndCrop(object): @@ -714,24 +825,39 @@ def __init__( binary_keys: Union[KeysView, Sequence[str]] = ( "mbs", "occs", - "valids", - "valid_flows", - "valid_disparities", "mbs_b", "occs_b", - "valids_b", + "mbs_right", + "occs_right", + "mbs_b_right", + "occs_b_right", + "valid_flows", "valid_flows_b", - "valid_disparities_right", "valid_flows_right", "valid_flows_b_right", + "valid_disparities", + "valid_disparities_right", + ), + flow_keys: Union[KeysView, Sequence[str]] = ( + "flows", + "flows_b", + "flows_right", + "flows_b_right", ), - disparities_keys: Union[KeysView, Sequence[str]] = ( + occlusion_keys: Union[KeysView, Sequence[str]] = ( + "occs", + "occs_b", + "occs_right", + "occs_b_right", + ), + disparity_keys: Union[KeysView, Sequence[str]] = ( "disparities", - "disparities_r", + "disparities_right", ), - occlusion_keys: Union[KeysView, Sequence[str]] = ("occs", "occs_b"), + intrinsics_key: Optional[str] = "intrinsics", + baselines_key: Optional[str] = "baselines", sparse: bool = False, - valid_key: str = "valids", + valid_key: str = "valid_flows", ) -> None: """Initialize RandomScaleAndCrop. @@ -747,13 +873,22 @@ def __init__( NOTE: Currently not implemented. The range of the time scale. See the class description for more details. binary_keys : Union[KeysView, Sequence[str]], default ['mbs', 'occs', 'valids', 'mbs_b', 'occs_b', 'valids_b'] Indicate which of the input keys correspond to binary tensors. - disparities_keys : Union[KeysView, Sequence[str]], default ['disparities', 'disparities_r'] - Indicate which of the input keys correspond to disparity tensors. + flow_keys : Union[KeysView, Sequence[str]], default ['flows', 'flows_b'] + Indicate which of the input keys correspond to optical flow tensors. occlusion_keys : Union[KeysView, Sequence[str]], default ['occs', 'occs_b'] Indicate which of the input keys correspond to occlusion mask tensors. + disparity_keys : Union[KeysView, Sequence[str]], default ['disparities'] + Indicate which of the input keys correspond to stereo disparity tensors. Disparities are resized like + flows, but since they have a single channel they are only scaled along x (the horizontal displacement). + intrinsics_key : Optional[str], default 'intrinsics' + Name of the entry containing the camera intrinsics. If present, its focal length and principal point are + scaled to match the resize and shifted to match the crop. + baselines_key : Optional[str], default 'baselines' + Name of the entry containing the stereo baselines. It is left unchanged, only protected from the + spatial resize and crop. sparse : bool, default False If True, only values at valid positions (indicated by the mask in inputs[valid_key]) will be kept when - resizing binary and disparity inputs. Requires valid_key to exist as a key in inputs. + resizing binary and flow inputs. Requires valid_key to exist as a key in inputs. valid_keys : str, default 'valids' The name of the key in inputs that contains the binary mask indicating which pixels are valid. Only used when sparse=True. @@ -789,8 +924,11 @@ def __init__( and abs(self.time_scale[3] - self.time_scale[2]) > 1e-5 ) self.binary_keys = list(binary_keys) - self.disparities_keys = list(disparities_keys) + self.flow_keys = list(flow_keys) self.occlusion_keys = list(occlusion_keys) + self.disparity_keys = list(disparity_keys) + self.intrinsics_key = intrinsics_key + self.baselines_key = baselines_key self.sparse = sparse self.valid_key = valid_key @@ -814,7 +952,10 @@ def __call__( # noqa: C901 NotImplementedError If trying to use time scale. """ - h, w = inputs[self.disparities_keys[0]].shape[2:4] + h, w = inputs[self.flow_keys[0]].shape[2:4] + intrinsics_h, intrinsics_w = inputs.get( + "images", inputs[self.flow_keys[0]] + ).shape[2:4] major_scale = 2 ** random.uniform(self.major_scale[0], self.major_scale[1]) space_scales = ( 2 ** random.uniform(self.space_scale[0], self.space_scale[1]), @@ -834,344 +975,58 @@ def __call__( # noqa: C901 y_crop = random.randint(0, scaled_size[0] - self.crop_size[0]) x_crop = random.randint(0, scaled_size[1] - self.crop_size[1]) - camera_keys = {"intrinsics", "baselines"}.intersection(inputs) + # Intrinsics and baselines are not spatial image tensors, so they must be excluded from the resize. + spatial_ignore = [ + k + for k in (self.intrinsics_key, self.baselines_key) + if k is not None and k in inputs + ] + resize_disparity_keys = [dk for dk in self.disparity_keys if dk in inputs] inputs = _resize( inputs, scaled_size, self.binary_keys, - self.disparities_keys, + self.flow_keys, self.sparse, self.valid_key, - ignore_keys=camera_keys, + ignore_keys=spatial_ignore, + disparity_keys=resize_disparity_keys, ) - if "intrinsics" in inputs: - inputs["intrinsics"] = _adjust_intrinsics_for_scale( - inputs["intrinsics"], - float(scaled_size[1]) / w, - float(scaled_size[0]) / h, + + if self.intrinsics_key is not None and self.intrinsics_key in inputs: + x_scale = float(scaled_size[1]) / float(intrinsics_w) + y_scale = float(scaled_size[0]) / float(intrinsics_h) + inputs[self.intrinsics_key] = _adjust_intrinsics_for_scale( + inputs[self.intrinsics_key], x_scale, y_scale ) + if self.crop_size is not None: for k, v in inputs.items(): - if k not in camera_keys: - inputs[k] = v[ - :, - :, - y_crop : y_crop + self.crop_size[0], - x_crop : x_crop + self.crop_size[1], - ] - if "intrinsics" in inputs: - inputs["intrinsics"] = _adjust_intrinsics_for_crop( - inputs["intrinsics"], x_crop, y_crop + if k in spatial_ignore: + continue + v = v[ + :, + :, + y_crop : y_crop + self.crop_size[0], + x_crop : x_crop + self.crop_size[1], + ] + inputs[k] = v + + if self.intrinsics_key is not None and self.intrinsics_key in inputs: + inputs[self.intrinsics_key] = _adjust_intrinsics_for_crop( + inputs[self.intrinsics_key], x_crop, y_crop ) - # Update occlusion masks for out-of-bounds disparities + # Update occlusion masks for out-of-bounds flows for k, v in inputs.items(): try: i = self.occlusion_keys.index(k) - if ( - i < len(self.disparities_keys) - and self.disparities_keys[i] in inputs - ): - inputs[k] = _update_oob_disparities( - v, inputs[self.disparities_keys[i]] - ) + inputs[k] = _update_oob_flows(v, inputs[self.flow_keys[i]]) except ValueError: pass return inputs -class RandomTranslate(object): - """Creates a translation between images by applying a random alternated crop on the sequence of inputs. - - A translation value t is randomly selected first. Then, the first image is cropped by a box translated by t. - The second image will be cropped by a reversed translation -t. The third will be cropped by t again, and so on... - """ - - def __init__( - self, - translation: Union[int, Tuple[int, int]] = 0, - disparities_keys: Union[KeysView, Sequence[str]] = ( - "disparities", - "disparities_r", - ), - flow_keys: Union[KeysView, Sequence[str]] = ("flows", "flows_b"), - occlusion_keys: Union[KeysView, Sequence[str]] = ("occs", "occs_b"), - ) -> None: - """Initialize RandomTranslate. - - Parameters - ---------- - translation : Union[int, Tuple[int, int]], default 0 - Maximum translation (in pixels) to be applied to the inputs. If a tuple, it corresponds to the maximum in the - (y, x) axes. - disparities_keys : Union[KeysView, Sequence[str]], default ['disparities', 'disparities_r'] - Indicate which of the input keys correspond to disparity tensors. Disparities link same-time stereo pairs, - which are cropped with the same offset, so their values are NOT changed by this transform. - flow_keys : Union[KeysView, Sequence[str]], default ['flows', 'flows_b'] - Indicate which of the input keys correspond to optical flow tensors. Flows link consecutive frames, which are - cropped with opposite offsets, so their values are shifted by the relative translation. - occlusion_keys : Union[KeysView, Sequence[str]], default ['occs', 'occs_b'] - Indicate which of the input keys correspond to occlusion mask tensors. - """ - self.translation = translation - if not isinstance(translation, tuple) or isinstance(translation, list): - self.translation = (translation, translation) - self.disparities_keys = disparities_keys - self.flow_keys = flow_keys - self.occlusion_keys = occlusion_keys - - def __call__(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: - """Perform the transformation on the inputs. - - Parameters - ---------- - inputs : Dict[str, torch.Tensor] - Elements to be transformed. Each element is a 4D tensor NCHW. - - Returns - ------- - Dict[str, torch.Tensor] - The inputs transformed by this operation. - """ - _, _, h, w = inputs[self.disparities_keys[0]].shape - th, tw = self.translation - tw = random.randint(-tw, tw) - th = random.randint(-th, th) - if tw == 0 and th == 0: - return inputs - - trans_inputs = { - k: torch.empty_like(v[:, :, : h - abs(th), : w - abs(tw)]) - for k, v in inputs.items() - } - - # Translate: 0: even indexed inputs, 1: odd indexed inputs - for t in range(2): - if t == 0: - ftw = tw - fth = th - else: - ftw = -tw - fth = -th - x1, x2 = max(0, ftw), min(w + ftw, w) - y1, y2 = max(0, fth), min(h + fth, h) - for k, v in inputs.items(): - trans_inputs[k][t::2] = v[t::2, :, y1:y2, x1:x2] - if k in self.flow_keys: - trans_inputs[k][t::2, 0] += ftw - trans_inputs[k][t::2, 1] += fth - - # Update occlusion masks for out-of-bounds disparities - for k, v in trans_inputs.items(): - try: - i = self.occlusion_keys.index(k) - if ( - i < len(self.disparities_keys) - and self.disparities_keys[i] in trans_inputs - ): - trans_inputs[k] = _update_oob_disparities( - v, trans_inputs[self.disparities_keys[i]] - ) - except ValueError: - pass - - return trans_inputs - - -class RandomRotate(object): - """Applies random rotation to the inputs. - - The inputs are rotated around the center of the image. First all inputs are rotated by the same random major `angle`. - Then, another random angle a is sampled according to `diff_angle`. The first image will be rotated by a. The second image - will be rotated by a reversed angle -a. The third will be rotated by a again, and so on... - """ - - def __init__( - self, - angle: float = 0.0, - diff_angle: float = 0.0, - disparities_keys: Union[KeysView, Sequence[str]] = ( - "disparities", - "disparities_r", - ), - flow_keys: Union[KeysView, Sequence[str]] = ("flows", "flows_b"), - occlusion_keys: Union[KeysView, Sequence[str]] = ("occs", "occs_b"), - valid_keys: Union[KeysView, Sequence[str]] = ("valids", "valids_b"), - binary_keys: Union[KeysView, Sequence[str]] = ( - "mbs", - "occs", - "valids", - "mbs_b", - "occs_b", - "valids_b", - ), - sparse: bool = False, - ) -> None: - """Initialize RandomRotate. - - Parameters - ---------- - angle : float, default 0.0 - The maximum absolute value to sample the major angle from. - diff_angle : float, default 0.0 - The maximum absolute value to sample the angle difference between consecutive images. - disparities_keys : Union[KeysView, Sequence[str]], default ['disparities', 'disparities_r'] - Indicate which of the input keys correspond to disparity tensors. - flow_keys : Union[KeysView, Sequence[str]], default ['flows', 'flows_b'] - Indicate which of the input keys correspond to optical flow tensors. - occlusion_keys : Union[KeysView, Sequence[str]], default ['occs', 'occs_b'] - Indicate which of the input keys correspond to occlusion mask tensors. - valid_keys : Union[KeysView, Sequence[str]], default ['valids', 'valids_b'] - Indicate which of the input keys correspond to valid mask tensors. - binary_keys : Union[KeysView, Sequence[str]], default ['mbs', 'occs', 'valids', 'mbs_b', 'occs_b', 'valids_b'] - Indicate which of the input keys correspond to binary tensors. - sparse : bool, default False - If True, all binary inputs and disparities are rotated with nearest grid_sample, instead of bilinear. - """ - self.angle = angle - self.diff_angle = diff_angle - self.disparities_keys = disparities_keys - self.flow_keys = flow_keys - self.occlusion_keys = occlusion_keys - self.valid_keys = valid_keys - self.binary_keys = binary_keys - self.sparse = sparse - - def __call__( # noqa: C901 - self, inputs: Dict[str, torch.Tensor] - ) -> Dict[str, torch.Tensor]: - """Perform the transformation on the inputs. - - Parameters - ---------- - inputs : Dict[str, torch.Tensor] - Elements to be transformed. Each element is a 4D tensor NCHW. - - Returns - ------- - Dict[str, torch.Tensor] - The inputs transformed by this operation. - """ - major_angle = random.uniform(-self.angle, self.angle) - inter_angle = random.uniform(-self.diff_angle, self.diff_angle) - - input_tensor = inputs[self.disparities_keys[0]] - b, _, h, w = input_tensor.shape - - def generate_rotation_grid( - rot_angle: float, batch_size: int, dtype: torch.dtype, device: torch.device - ) -> torch.Tensor: - vy, vx = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") - vx = vx.type(dtype) - vy = vy.type(dtype) - vx = vx.to(device) - vy = vy.to(device) - vx -= (w - 1.0) / 2.0 - vy -= (h - 1.0) / 2.0 - angle_rad = rot_angle * 2 * np.pi / 360 - rotx = np.cos(angle_rad) * vx - np.sin(angle_rad) * vy - roty = np.sin(angle_rad) * vx + np.cos(angle_rad) * vy - rotx = rotx / ((w - 1) / 2) - roty = roty / ((h - 1) / 2) - rot_grid = torch.stack((rotx, roty), dim=2)[None] - rot_grid = rot_grid.repeat(batch_size, 1, 1, 1) - return rot_grid - - def generate_rotation_matrix( - rot_angle: float, batch_size: int, dtype: torch.dtype, device: torch.device - ) -> torch.Tensor: - vx, vy = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") - vx = vx.type(dtype) - vy = vy.type(dtype) - vx = vx.to(device) - vy = vy.to(device) - rotx = (vx - h / 2.0) * (rot_angle * np.pi / 180.0) - roty = -(vy - w / 2.0) * (rot_angle * np.pi / 180.0) - rot_mat = torch.stack((rotx, roty), dim=0)[None] - rot_mat = rot_mat.repeat(batch_size, 1, 1, 1) - return rot_mat - - def rotate_disparity(disparity: torch.Tensor, rot_angle: float) -> torch.Tensor: - angle_rad = rot_angle * 2 * np.pi / 360 - rot_disparity = disparity.clone() - rot_disparity[:, 0] = ( - np.cos(angle_rad) * disparity[:, 0] - + np.sin(angle_rad) * disparity[:, 1] - ) - rot_disparity[:, 1] = ( - -np.sin(angle_rad) * disparity[:, 0] - + np.cos(angle_rad) * disparity[:, 1] - ) - return rot_disparity - - rot_mat = generate_rotation_matrix( - inter_angle, b // 2 + 1, input_tensor.dtype, input_tensor.device - ) - for t in range(2): - if t == 0: - inangle = -inter_angle - rmat = rot_mat - else: - inangle = inter_angle - rmat = -rot_mat - angle = major_angle + inangle / 2 - num_disparities = input_tensor[t::2].shape[0] - num_images = num_disparities + 1 - rot_grid = generate_rotation_grid( - angle, num_images, input_tensor.dtype, input_tensor.device - ) - for k, v in inputs.items(): - # Flows link consecutive frames, which are rotated by different - # angles, so they also gain the inter-frame rotation offset. - # Disparities link same-time stereo pairs, which are rotated - # together, so no offset is added to them. - if k in self.flow_keys: - v[t::2] += rmat[: v[t::2].shape[0]] - - if k in self.binary_keys: - v[t::2] = F.grid_sample( - v[t::2], rot_grid[: v[t::2].shape[0]], mode="nearest" - ) - else: - if k in self.disparities_keys or k in self.flow_keys: - if self.sparse: - v[t::2] = F.grid_sample( - v[t::2], rot_grid[: v[t::2].shape[0]], mode="nearest" - ) - else: - v[t::2] = F.grid_sample( - v[t::2], - rot_grid[: v[t::2].shape[0]], - mode="bilinear", - align_corners=True, - ) - v[t::2] = rotate_disparity(v[t::2], angle) - else: - v[t::2] = F.grid_sample( - v[t::2], - rot_grid[: v[t::2].shape[0]], - mode="bilinear", - align_corners=True, - ) - - inputs[k] = v - - # Update occlusion masks for out-of-bounds disparities - for k, v in inputs.items(): - try: - i = self.occlusion_keys.index(k) - if ( - i < len(self.disparities_keys) - and self.disparities_keys[i] in inputs - ): - v = _update_oob_disparities(v, inputs[self.disparities_keys[i]]) - inputs[k] = v - except ValueError: - pass - - return inputs - - class Resize(object): """Resize the image to a given size or scale. @@ -1185,17 +1040,33 @@ def __init__( binary_keys: Union[KeysView, Sequence[str]] = ( "mbs", "occs", - "valids", "mbs_b", "occs_b", - "valids_b", + "mbs_right", + "occs_right", + "mbs_b_right", + "occs_b_right", + "valid_flows", + "valid_flows_b", + "valid_flows_right", + "valid_flows_b_right", + "valid_disparities", + "valid_disparities_right", ), - disparities_keys: Union[KeysView, Sequence[str]] = ( + flow_keys: Union[KeysView, Sequence[str]] = ( + "flows", + "flows_b", + "flows_right", + "flows_b_right", + ), + disparity_keys: Union[KeysView, Sequence[str]] = ( "disparities", - "disparities_b", + "disparities_right", ), + intrinsics_key: Optional[str] = "intrinsics", + baselines_key: Optional[str] = "baselines", sparse: bool = False, - valid_key: str = "valids", + valid_key: str = "valid_flows", ignore_keys: Optional[Union[KeysView, Sequence[str]]] = None, ) -> None: """Initialize Resize. @@ -1209,11 +1080,18 @@ def __init__( binary_keys : Union[KeysView, Sequence[str]], default ['mbs', 'occs', 'valids', 'mbs_b', 'occs_b', 'valids_b'] Indicate which of the input keys correspond to binary tensors. [description], by default ['mbs', 'occs', 'valids', 'mbs_b', 'occs_b', 'valids_b'] - disparities_keys : Union[KeysView, Sequence[str]], default ['disparities', 'disparities_b'] - Indicate which of the input keys correspond to disparity tensors. + flow_keys : Union[KeysView, Sequence[str]], default ['flows', 'flows_b'] + Indicate which of the input keys correspond to optical flow tensors. + disparity_keys : Union[KeysView, Sequence[str]], default ['disparities'] + Indicate which of the input keys correspond to stereo disparity tensors. They are resized like flows. + intrinsics_key : Optional[str], default 'intrinsics' + Name of the entry containing the camera intrinsics. If present, its focal length and principal point are + scaled to match the resize. + baselines_key : Optional[str], default 'baselines' + Name of the entry containing the stereo baselines. It is left unchanged, only protected from the resize. sparse : bool, default False If True, only values at valid positions (indicated by the mask in inputs[valid_key]) will be kept when - resizing binary and disparity inputs. Requires valid_key to exist as a key in inputs. + resizing binary and flow inputs. Requires valid_key to exist as a key in inputs. valid_keys : str, default 'valids' The name of the key in inputs that contains the binary mask indicating which pixels are valid. Only used when sparse=True. @@ -1223,7 +1101,10 @@ def __init__( self.size = size self.scale = scale self.binary_keys = list(binary_keys) - self.disparities_keys = list(disparities_keys) + self.flow_keys = list(flow_keys) + self.disparity_keys = list(disparity_keys) + self.intrinsics_key = intrinsics_key + self.baselines_key = baselines_key self.sparse = sparse self.valid_key = valid_key self.ignore_keys = ignore_keys @@ -1241,19 +1122,48 @@ def __call__(self, inputs: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: Dict[str, torch.Tensor] The inputs transformed by this operation. """ - h, w = inputs[list(inputs.keys())[0]].shape[2:4] - if self.size is None or self.size[0] < 1 or self.size[1] < 1: - self.size = (int(self.scale * h), int(self.scale * w)) - if self.size[0] != h or self.size[1] != w: + # Use a spatial tensor (images/flows) to read the current size, skipping non-spatial entries. + spatial_key = next( + ( + k + for k in inputs + if k not in (self.intrinsics_key, self.baselines_key) + and isinstance(inputs[k], torch.Tensor) + and inputs[k].ndim >= 2 + ), + None, + ) + h, w = inputs[spatial_key].shape[-2:] + size = self.size + if size is None or size[0] < 1 or size[1] < 1: + size = (int(self.scale * h), int(self.scale * w)) + if size[0] != h or size[1] != w: + spatial_ignore = [ + k + for k in (self.intrinsics_key, self.baselines_key) + if k is not None and k in inputs + ] + extra_ignore = ( + list(self.ignore_keys) if self.ignore_keys is not None else [] + ) + resize_ignore = list(set(spatial_ignore + extra_ignore)) + resize_disparity_keys = [dk for dk in self.disparity_keys if dk in inputs] inputs = _resize( inputs, - self.size, + size, self.binary_keys, - self.disparities_keys, + self.flow_keys, self.sparse, self.valid_key, - ignore_keys=self.ignore_keys, + ignore_keys=resize_ignore, + disparity_keys=resize_disparity_keys, ) + if self.intrinsics_key is not None and self.intrinsics_key in inputs: + x_scale = float(size[1]) / float(w) + y_scale = float(size[0]) / float(h) + inputs[self.intrinsics_key] = _adjust_intrinsics_for_scale( + inputs[self.intrinsics_key], x_scale, y_scale + ) return inputs @@ -1281,20 +1191,105 @@ def _get_valid_keys( The keys remaining after the validity checks. """ if use_keys is not None: - return use_keys + input_keys = set(inputs_keys) + return [k for k in use_keys if k in input_keys] if ignore_keys is None: return inputs_keys return [k for k in inputs_keys if k not in ignore_keys] +def _adjust_intrinsics_for_scale( + intrinsics: torch.Tensor, x_scale: float, y_scale: float +) -> torch.Tensor: + """Scale the camera intrinsics to match a spatial resize. + + The focal lengths and the principal point are multiplied by the corresponding scale factor. + + Parameters + ---------- + intrinsics : torch.Tensor + A tensor of shape (N, 3, 3) with the camera intrinsics. + x_scale : float + Scale factor applied to the width. + y_scale : float + Scale factor applied to the height. + + Returns + ------- + torch.Tensor + The scaled intrinsics, with the same shape and dtype as the input. + """ + intr = intrinsics.clone() + intr[:, 0, 0] = intr[:, 0, 0] * x_scale + intr[:, 1, 1] = intr[:, 1, 1] * y_scale + intr[:, 0, 2] = intr[:, 0, 2] * x_scale + intr[:, 1, 2] = intr[:, 1, 2] * y_scale + return intr.to(dtype=intrinsics.dtype, device=intrinsics.device) + + +def _adjust_intrinsics_for_crop( + intrinsics: torch.Tensor, x_crop: int, y_crop: int +) -> torch.Tensor: + """Shift the principal point of the intrinsics to match a spatial crop. + + Parameters + ---------- + intrinsics : torch.Tensor + A tensor of shape (N, 3, 3) with the camera intrinsics. + x_crop : int + Pixel offset of the crop along x (left border). + y_crop : int + Pixel offset of the crop along y (top border). + + Returns + ------- + torch.Tensor + The shifted intrinsics, with the same shape and dtype as the input. + """ + intr = intrinsics.clone() + intr[:, 0, 2] = intr[:, 0, 2] - x_crop + intr[:, 1, 2] = intr[:, 1, 2] - y_crop + return intr.to(dtype=intrinsics.dtype, device=intrinsics.device) + + +def _adjust_intrinsics_for_flip( + intrinsics: torch.Tensor, is_hflip: bool, img_w: int, img_h: int +) -> torch.Tensor: + """Mirror the principal point of the intrinsics to match a horizontal or vertical flip. + + Parameters + ---------- + intrinsics : torch.Tensor + A tensor of shape (N, 3, 3) with the camera intrinsics. + is_hflip : bool + If True, mirror the principal point along x (horizontal flip), otherwise along y (vertical flip). + img_w : int + Width of the (flipped) image, used to compute the mirrored principal point. + img_h : int + Height of the (flipped) image, used to compute the mirrored principal point. + + Returns + ------- + torch.Tensor + The mirrored intrinsics, with the same shape and dtype as the input. + """ + intr = intrinsics.clone() + if is_hflip: + intr[:, 0, 2] = (img_w - 1) - intr[:, 0, 2] + else: + intr[:, 1, 2] = (img_h - 1) - intr[:, 1, 2] + return intr.to(dtype=intrinsics.dtype, device=intrinsics.device) + + def _resize( inputs: Dict[str, torch.Tensor], target_size: Tuple[int, int], binary_keys: Union[KeysView, Sequence[str]], - disparities_keys: Union[KeysView, Sequence[str]], + flow_keys: Union[KeysView, Sequence[str]], sparse: bool, valid_key: str, ignore_keys: Optional[Sequence[str]] = None, + disparity_keys: Union[KeysView, Sequence[str]] = (), ): """Resize inputs to a target size. Set sparse=True when the valid mask has holes. This ensures that the resized valid mask does not interpolate the valid positions. @@ -1308,77 +1303,134 @@ def _resize( binary_keys : Union[KeysView, Sequence[str]] Indicate which of the input keys correspond to binary tensors. [description], by default ['mbs', 'occs', 'valids', 'mbs_b', 'occs_b', 'valids_b'] - disparities_keys : Union[KeysView, Sequence[str]] - Indicate which of the input keys correspond to disparity tensors. + flow_keys : Union[KeysView, Sequence[str]] + Indicate which of the input keys correspond to optical flow tensors. sparse : bool If True, only values at valid positions (indicated by the mask in inputs[valid_key]) will be kept when - resizing binary and disparity inputs. Requires valid_key to exist as a key in inputs. + resizing binary and flow inputs. Requires valid_key to exist as a key in inputs. valid_keys : str The name of the key in inputs that contains the binary mask indicating which pixels are valid. Only used when sparse=True. + ignore_keys : Optional[Sequence[str]] + Keys to skip entirely (not resized). + disparity_keys : Union[KeysView, Sequence[str]], default () + Indicate which of the input keys correspond to stereo disparity tensors. Disparities are always resized with + bilinear interpolation and their values are scaled with the image (single-channel disparities are only scaled + along x). They are NOT scattered using valid_key, because the disparity and flow valid masks may have different + lengths. Returns ------- torch.Tensor - The updated occlusion masks. Disparities which went out-of-bounds are marked as occluded. + The updated occlusion masks. Flows which went out-of-bounds are marked as occluded. """ + disparity_keys = list(disparity_keys) if sparse: assert ( valid_key in inputs ), f"sparse is True, but valid_key({valid_key}) is not in inputs" valids = inputs[valid_key] - n, k, h, w = valids.shape + valid_masks = { + key: value for key, value in inputs.items() if key.startswith("valid_") + } hs, ws = target_size - disparity_scale = torch.tensor( - [float(ws) / w, float(hs) / h], device=valids.device - ) - y_coords, x_coords = torch.meshgrid( - torch.arange(h, device=valids.device), - torch.arange(w, device=valids.device), - indexing="ij", - ) - coords = torch.stack((x_coords, y_coords), dim=-1).reshape(-1, 2) - valid_masks = valids[:, 0].reshape(n, -1) >= 1 - xy_scaled_list = [] - inbounds_list = [] - valids_out = torch.zeros(n, k, hs, ws, dtype=valids.dtype, device=valids.device) - for i, valid_mask in enumerate(valid_masks): - coords_valid = coords[valid_mask] - - coords_scaled = coords_valid * disparity_scale - - x_scaled = torch.round(coords_scaled[:, 0]).long() - y_scaled = torch.round(coords_scaled[:, 1]).long() - inbounds = ( - (x_scaled >= 0) & (x_scaled < ws) & (y_scaled >= 0) & (y_scaled < hs) - ) - inbounds_list.append(inbounds) - x_scaled = x_scaled[inbounds] - y_scaled = y_scaled[inbounds] - xy_scaled_list.append((x_scaled, y_scaled)) - valids_out[i, :, y_scaled, x_scaled] = 1 + def _sparse_geometry(validity: torch.Tensor): + n, c, h, w = validity.shape + scale_factor = torch.tensor( + [float(ws) / w, float(hs) / h], device=validity.device + ) + yy, xx = torch.meshgrid( + torch.arange(h, device=validity.device), + torch.arange(w, device=validity.device), + indexing="ij", + ) + coords = torch.stack((xx, yy), dim=-1).reshape(-1, 2) + valid_flat = rearrange(validity, "n c h w -> n (c h w)") + xy_scaled_list = [] + inbounds_list = [] + validity_out = torch.zeros( + n, c, hs, ws, dtype=validity.dtype, device=validity.device + ) + for i, vflat in enumerate(valid_flat): + coords_scaled = coords[vflat >= 1] * scale_factor + x_scaled = torch.round(coords_scaled[:, 0]).long() + y_scaled = torch.round(coords_scaled[:, 1]).long() + inbounds = ( + (x_scaled >= 0) + & (x_scaled < ws) + & (y_scaled >= 0) + & (y_scaled < hs) + ) + inbounds_list.append(inbounds) + x_scaled = x_scaled[inbounds] + y_scaled = y_scaled[inbounds] + xy_scaled_list.append((x_scaled, y_scaled)) + validity_out[i, 0, y_scaled, x_scaled] = 1 + return valid_flat, xy_scaled_list, inbounds_list, scale_factor, validity_out + + geometry_cache = {} + + def _geometry_for(validity: torch.Tensor): + cache_key = id(validity) + if cache_key not in geometry_cache: + geometry_cache[cache_key] = _sparse_geometry(validity) + return geometry_cache[cache_key] + + def _scatter_sparse( + value: torch.Tensor, validity: torch.Tensor, scale_values: bool + ) -> torch.Tensor: + ( + valids_flat, + xy_scaled_list, + inbounds_list, + scale_factor, + _, + ) = _geometry_for(validity) + value_out = torch.zeros( + value.shape[0], + value.shape[1], + hs, + ws, + dtype=value.dtype, + device=value.device, + ) + for i, value_one in enumerate(value): + value_flat = rearrange(value_one, "c h w -> (h w) c") + value_valid = value_flat[valids_flat[i] >= 1] + if scale_values: + value_valid = value_valid * scale_factor[: value.shape[1]] + value_valid = value_valid[inbounds_list[i]] + value_valid = rearrange(value_valid, "n c -> c n") + value_out[i, :, xy_scaled_list[i][1], xy_scaled_list[i][0]] = ( + value_valid + ) + return value_out - inputs[valid_key] = valids_out + inputs[valid_key] = _geometry_for(valids)[4] for k, v in inputs.items(): if k != valid_key and (ignore_keys is None or k not in ignore_keys): - is_flow = "flows" in k and "valid" not in k - if k in binary_keys or k in disparities_keys or is_flow: - v_out = torch.zeros( - v.shape[0], v.shape[1], hs, ws, dtype=v.dtype, device=v.device - ) - for i, v_one in enumerate(v): - v_flat = rearrange(v_one, "k h w -> (h w) k") - v_valid = v_flat[valid_masks[i]] - if k in disparities_keys or is_flow: - v_valid = v_valid * disparity_scale[: v_valid.shape[1]] - v_valid = v_valid[inbounds_list[i]] - v_valid = rearrange(v_valid, "n k -> k n") - v_out[i, :, xy_scaled_list[i][1], xy_scaled_list[i][0]] = ( - v_valid + if k in disparity_keys: + disparity_valids = valid_masks.get(f"valid_{k}") + if disparity_valids is None: + orig_h, orig_w = v.shape[-2:] + v = F.interpolate( + v, size=target_size, mode="bilinear", align_corners=True ) - v = v_out + v = _scale_disparity(v, target_size, orig_h, orig_w) + else: + v = _scatter_sparse(v, disparity_valids, scale_values=True) + elif k in binary_keys or k in flow_keys: + field_valids = valid_masks.get(k) + if field_valids is None and k in flow_keys: + field_valids = valid_masks.get(f"valid_{k}") + if field_valids is None: + field_valids = valids + if k in valid_masks: + inputs[k] = _geometry_for(field_valids)[4] + continue + v = _scatter_sparse(v, field_valids, scale_values=k in flow_keys) else: v = F.interpolate( v, size=target_size, mode="bilinear", align_corners=True @@ -1395,12 +1447,15 @@ def _resize( v, size=target_size, mode="bilinear", align_corners=True ) - if k in disparities_keys or ("flows" in k and "valid" not in k): - scale_mult = torch.ones(v.shape[1], dtype=v.dtype, device=v.device) - scale_mult[0] = float(target_size[1]) / w - if v.shape[1] > 1: - scale_mult[1] = float(target_size[0]) / h - scale_mult = scale_mult[None, :, None, None] + if k in flow_keys or k in disparity_keys: + scale_mult = torch.tensor( + [float(target_size[1]) / w, float(target_size[0]) / h], + device=v.device, + )[None, :, None, None] + + if v.shape[1] == 1: + scale_mult = scale_mult[:, :1, :, :] + v = v * scale_mult inputs[k] = v @@ -1408,91 +1463,73 @@ def _resize( return inputs -def _adjust_intrinsics_for_scale( - intrinsics: torch.Tensor, x_scale: float, y_scale: float +def _scale_disparity( + disp: torch.Tensor, target_size: Tuple[int, int], orig_h: int, orig_w: int ) -> torch.Tensor: - intrinsics = intrinsics.clone() - intrinsics[..., 0, 0] *= x_scale - intrinsics[..., 1, 1] *= y_scale - intrinsics[..., 0, 2] *= x_scale - intrinsics[..., 1, 2] *= y_scale - return intrinsics + """Scale disparity values to match a spatial resize. + A single-channel disparity is a horizontal displacement, so it is only scaled along x. A two-channel disparity + (e.g. with a vertical component) is scaled along both axes, like a flow. -def _adjust_intrinsics_for_crop( - intrinsics: torch.Tensor, x_crop: int, y_crop: int -) -> torch.Tensor: - intrinsics = intrinsics.clone() - intrinsics[..., 0, 2] -= x_crop - intrinsics[..., 1, 2] -= y_crop - return intrinsics - + Parameters + ---------- + disp : torch.Tensor + The disparity tensor (already interpolated to target_size). + target_size : Tuple[int, int] + Target (height, width) sizes. + orig_h : int + Original height before interpolation. + orig_w : int + Original width before interpolation. -def _adjust_intrinsics_for_flip( - intrinsics: torch.Tensor, is_hflip: bool, width: int, height: int -) -> torch.Tensor: - intrinsics = intrinsics.clone() - if is_hflip: - intrinsics[..., 0, 2] = (width - 1) - intrinsics[..., 0, 2] - else: - intrinsics[..., 1, 2] = (height - 1) - intrinsics[..., 1, 2] - return intrinsics + Returns + ------- + torch.Tensor + The scaled disparity. + """ + scale_mult = torch.tensor( + [float(target_size[1]) / orig_w, float(target_size[0]) / orig_h], + device=disp.device, + )[None, :, None, None] + if disp.shape[1] == 1: + scale_mult = scale_mult[:, :1, :, :] + return disp * scale_mult -def _update_oob_disparities( - occs: torch.Tensor, disparities: torch.Tensor -) -> torch.Tensor: - """Update occlusion maps to include disparity which went out-of-bounds. +def _update_oob_flows(occs: torch.Tensor, flows: torch.Tensor) -> torch.Tensor: + """Update occlusion maps to include flow which went out-of-bounds. Parameters ---------- occs : torch.Tensor A 4D tensor NCHW of occlusion masks. - disparities : torch.Tensor - A 4D tensor NCHW of disparities. + flows : torch.Tensor + A 4D tensor NCHW of optical flows. Returns ------- torch.Tensor - The updated occlusion masks. Disparities which went out-of-bounds are marked as occluded. + The updated occlusion masks. Flows which went out-of-bounds are marked as occluded. """ - if disparities.shape[1] == 1: - # 1-channel positive stereo disparity: the correspondence of a pixel - # (x, y) in the left image is (x - d, y) in the right image. - grid_x = torch.arange( - disparities.shape[3], dtype=disparities.dtype, device=disparities.device - ) - grid_x = grid_x[None, None, None, :].repeat( - disparities.shape[0], 1, disparities.shape[2], 1 - ) - coords_x = grid_x - disparities - oob_occs = (coords_x < 0) | (coords_x >= disparities.shape[3]) + grid = torch.meshgrid( + torch.arange(flows.shape[2], dtype=flows.dtype, device=flows.device), + torch.arange(flows.shape[3], dtype=flows.dtype, device=flows.device), + indexing="ij", + ) + grid = torch.stack(grid[::-1]).float()[None].repeat(flows.shape[0], 1, 1, 1) + + if flows.shape[1] == 1: + grid = grid[:, :1, :, :] + coords = flows + grid + oob_occs = coords < 0 + oob_occs[:, 0] |= coords[:, 0] >= flows.shape[3] else: - # Multi-channel disparity is treated as a forward correspondence - # vector: the correspondence of (x, y) is (x, y) + (dx, dy). - grid = torch.meshgrid( - torch.arange( - disparities.shape[2], - dtype=disparities.dtype, - device=disparities.device, - ), - torch.arange( - disparities.shape[3], - dtype=disparities.dtype, - device=disparities.device, - ), - indexing="ij", - ) - grid = ( - torch.stack(grid[::-1]) - .to(dtype=disparities.dtype)[None] - .repeat(disparities.shape[0], 1, 1, 1) - ) - coords = disparities + grid + coords = flows + grid oob_occs = coords < 0 - oob_occs[:, 0] |= coords[:, 0] >= disparities.shape[3] - oob_occs[:, 1] |= coords[:, 1] >= disparities.shape[2] - oob_occs = oob_occs.max(dim=1, keepdim=True)[0] - oob_occs = oob_occs.to(dtype=occs.dtype, device=occs.device) + oob_occs[:, 0] |= coords[:, 0] >= flows.shape[3] + oob_occs[:, 1] |= coords[:, 1] >= flows.shape[2] + oob_occs = oob_occs.max(dim=1, keepdim=True)[0].to( + dtype=occs.dtype, device=occs.device + ) occs = torch.max(torch.stack([occs, oob_occs], dim=0), dim=0)[0] return occs diff --git a/roco_spring_devkit/common/data/split_autoflow.py b/roco_spring_devkit/common/data/split_autoflow.py deleted file mode 100644 index ee81c1a..0000000 --- a/roco_spring_devkit/common/data/split_autoflow.py +++ /dev/null @@ -1,102 +0,0 @@ -""" - -Create a file with a list of samples names from the AutoFlow [1] dataset to be used as validation samples. - -[1] Sun, Deqing et al. “AutoFlow: Learning a Better Training Set for Optical Flow.” CVPR. 2021. - -""" - -# ============================================================================= -# Copyright 2022 Henrique Morimitsu -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# ============================================================================= - - -from argparse import ArgumentParser, Namespace -import os -from pathlib import Path -import random - -random.seed(42) - -THIS_DIR = Path(os.path.abspath(os.path.dirname(__file__))) - - -def _init_parser() -> ArgumentParser: - parser: ArgumentParser = ArgumentParser() - parser.add_argument("--autoflow_root", type=str, required=True) - parser.add_argument( - "--output_file", type=str, default=str(THIS_DIR / "AutoFlow_val.txt") - ) - parser.add_argument("--val_percentage", type=float, default=0.05) - return parser - - -def main(args: Namespace) -> None: - """Run the split process. - - Parameters - ---------- - args : argparse.Namespace - Arguments for configuring the splitting. - """ - parts_dirs = [f"static_40k_png_{i+1}_of_4" for i in range(4)] - sample_dirs = [] - for pdir in parts_dirs: - sample_dirs.extend( - sorted( - [ - f.stem - for f in (Path(args.autoflow_root) / pdir).glob("*") - if f.is_dir() - ] - ) - ) - sample_dirs.sort() - assert ( - len(sample_dirs) == 40000 - ), f"ERROR: AutoFlow dataset should have 40k samples, but found {len(sample_dirs)}." - samples_per_table = {} - for sdir in sample_dirs: - table_idx = sdir.split("_")[1] - if table_idx not in samples_per_table: - samples_per_table[table_idx] = [] - samples_per_table[table_idx].append(sdir) - assert ( - len(samples_per_table) == 300 - ), f"ERROR: AutoFlow dataset should have 300 tables, but found {len(samples_per_table)}." - - val_samples = [] - carryover_samples = 0.0 - for dir_list in samples_per_table.values(): - num_samples = len(dir_list) - num_val_samples_float = args.val_percentage * num_samples + carryover_samples - num_val_samples = int(num_val_samples_float) - - random.shuffle(dir_list) - val_samples.extend(dir_list[:num_val_samples]) - - carryover_samples = num_val_samples_float - num_val_samples - - val_samples.sort(key=lambda x: 1000 * int(x.split("_")[1]) + int(x.split("_")[-1])) - with open(args.output_file, "w") as f: - f.write("\n".join(val_samples)) - - print(f"Saved {len(val_samples)} sample names to {args.output_file}") - - -if __name__ == "__main__": - parser: ArgumentParser = _init_parser() - args: Namespace = parser.parse_args() - main(args) diff --git a/roco_spring_devkit/common/utils/callbacks/logger.py b/roco_spring_devkit/common/utils/callbacks/logger.py index 247fee3..c98eb6c 100644 --- a/roco_spring_devkit/common/utils/callbacks/logger.py +++ b/roco_spring_devkit/common/utils/callbacks/logger.py @@ -65,7 +65,15 @@ def __init__( self, num_images: int = 5, image_size: Tuple[int, int] = (200, 400), - log_keys: Sequence[str] = ("images", "flows", "occs", "mbs", "confs"), + log_keys: Sequence[str] = ( + "images", + "flows", + "disparities", + "disparities2", + "occs", + "mbs", + "confs", + ), epe_clip: float = 5.0, ) -> None: """Initialize LoggerCallback. @@ -128,6 +136,32 @@ def log_image(self, title: str, image: torch.Tensor, pl_module) -> None: image_wb = swanlab.Image((255 * image_npy).astype(np.uint8)) logger.experiment.log({title_wb: image_wb}) + def on_train_batch_start( + self, + trainer: Trainer, + pl_module, + batch: Dict[str, torch.Tensor], + batch_idx: int, + **kwargs, + ) -> None: + """Ask the model to keep the inputs/predictions of this batch, if it is in the log selection group. + + Storing them for every batch would keep a whole batch and all of the intermediate predictions alive in memory + during the entire step, which is wasteful since only a handful of samples are logged per epoch. + + Parameters + ---------- + trainer : Trainer + An instance of the PyTorch Lightning trainer. + pl_module + An instance of the optical flow model. + batch : Dict[str, torch.Tensor] + The inputs of the current training batch. + batch_idx : int + The counter value of the current batch. + """ + pl_module.store_last_batch = batch_idx in self.train_collect_img_idx + def on_train_batch_end( self, trainer: Trainer, @@ -156,6 +190,9 @@ def on_train_batch_end( self._append_images( self.train_images, pl_module.last_inputs, pl_module.last_predictions ) + pl_module.store_last_batch = False + pl_module.last_inputs = None + pl_module.last_predictions = None def on_train_epoch_start(self, trainer: Trainer, pl_module) -> None: """Reset the training log params and accumulators. @@ -168,22 +205,13 @@ def on_train_epoch_start(self, trainer: Trainer, pl_module) -> None: An instance of the optical flow model. """ self.train_images = {} - limit_batches = ( - trainer.limit_train_batches - if trainer.limit_train_batches is not None - else 1.0 + # num_training_batches is the number of batches this process will actually iterate over, so it already accounts + # for --limit_train_batches and for the sharding done by the distributed sampler. Using the full (unsharded) + # dataloader length here would spread the indices over a range that this rank never reaches, and most of the + # requested images would never be collected. + self.train_collect_img_idx = self._compute_collect_idx( + trainer.num_training_batches ) - collect_idx = np.unique( - np.linspace( - 0, - self._compute_max_range( - trainer.datamodule.train_dataloader_length, limit_batches - ), - self.num_images, - dtype=np.int32, - ) - ) - self.train_collect_img_idx = collect_idx def on_train_epoch_end(self, trainer: Trainer, pl_module, **kwargs) -> None: """Log the images accumulated during the training. @@ -198,7 +226,33 @@ def on_train_epoch_end(self, trainer: Trainer, pl_module, **kwargs) -> None: Outputs of the training epoch. """ img_grid = self._make_image_grid(self.train_images) - self.log_image("train", img_grid, pl_module) + if img_grid is not None: + self.log_image("train", img_grid, pl_module) + + def on_validation_batch_start( + self, + trainer: Trainer, + pl_module, + batch: Dict[str, torch.Tensor], + batch_idx: int, + dataloader_idx: int = 0, + ) -> None: + """Ask the model to keep the inputs/predictions of this batch, if it is in the log selection group. + + Parameters + ---------- + trainer : Trainer + An instance of the PyTorch Lightning trainer. + pl_module + An instance of the optical flow model. + batch : Dict[str, torch.Tensor] + The inputs of the current validation batch. + batch_idx : int + The counter value of the current batch. + dataloader_idx : int + The index number of the current dataloader. + """ + pl_module.store_last_batch = batch_idx in self._val_collect_idx(dataloader_idx) def on_validation_batch_end( self, @@ -226,13 +280,16 @@ def on_validation_batch_end( dataloader_idx : int The index number of the current dataloader. """ - dl_name = self.val_dataloader_names[dataloader_idx] - if batch_idx in self.val_collect_image_idx[dl_name]: + if batch_idx in self._val_collect_idx(dataloader_idx): + dl_name = self.val_dataloader_names[dataloader_idx] self._append_images( self.val_images[dl_name], pl_module.last_inputs, pl_module.last_predictions, ) + pl_module.store_last_batch = False + pl_module.last_inputs = None + pl_module.last_predictions = None def on_validation_epoch_start(self, trainer: Trainer, pl_module) -> None: """Reset the validation log params and accumulators. @@ -248,22 +305,17 @@ def on_validation_epoch_start(self, trainer: Trainer, pl_module) -> None: for dl_name in self.val_dataloader_names: self.val_images[dl_name] = {} - limit_batches = ( - trainer.limit_val_batches if trainer.limit_val_batches is not None else 1.0 - ) - for dname, dlen in zip( - trainer.datamodule.val_dataloader_names, - trainer.datamodule.val_dataloader_lengths, - ): - collect_idx = np.unique( - np.linspace( - 0, - self._compute_max_range(dlen, limit_batches), - self.num_images, - dtype=np.int32, - ) + # num_val_batches is a list with the number of batches that this process iterates over for each val dataloader, + # so it already accounts for --limit_val_batches and for the distributed sampler. + num_val_batches = trainer.num_val_batches + if not isinstance(num_val_batches, (list, tuple)): + num_val_batches = [num_val_batches] * len(self.val_dataloader_names) + self.val_collect_image_idx = {} + for i, dname in enumerate(self.val_dataloader_names): + num_batches = ( + num_val_batches[i] if i < len(num_val_batches) else num_val_batches[-1] ) - self.val_collect_image_idx[dname] = collect_idx + self.val_collect_image_idx[dname] = self._compute_collect_idx(num_batches) def on_validation_epoch_end(self, trainer: Trainer, pl_module) -> None: """Log the images accumulated during the validation. @@ -277,7 +329,8 @@ def on_validation_epoch_end(self, trainer: Trainer, pl_module) -> None: """ for dl_name, dl_images in self.val_images.items(): img_grid = self._make_image_grid(dl_images) - self.log_image(f"val/{dl_name}", img_grid, pl_module) + if img_grid is not None: + self.log_image(f"val/{dl_name}", img_grid, pl_module) def _add_title(self, image: torch.Tensor, img_title: str) -> torch.Tensor: """Add a title to an image. @@ -335,6 +388,9 @@ def _append_images( # noqa: C901 preds : Dict[str, torch.Tensor] The outrputs of the model. """ + if "disparities" in inputs and inputs["disparities"].shape[1] > 1: + inputs["disparities2"] = inputs["disparities"][:, 1:] + for k in self.log_keys: log_names = [] log_sources = [] @@ -351,8 +407,13 @@ def _append_images( # noqa: C901 elif k == "disparities": log_names.append(f"abs<{self.epe_clip:.1f}") log_sources.append(None) + elif k == "disparities2": + log_names.append(f"abs2<{self.epe_clip:.1f}") + log_sources.append(None) for name, source in zip(log_names, log_sources): + valid = None + if images.get(name) is None: images[name] = [] @@ -361,13 +422,26 @@ def _append_images( # noqa: C901 elif name.startswith("epe"): epe = torch.norm(preds[k] - inputs[k], p=2, dim=2, keepdim=True) img = torch.clamp(epe, 0, self.epe_clip) / self.epe_clip - if inputs.get("valids") is not None: - img[inputs["valids"] < 0.5] = 0 + valid = inputs.get("valids") + if valid is None: + valid = inputs.get("valid_flows") + elif name.startswith("abs2"): + abs_val = torch.abs(preds[k] - inputs[k[:-1]][:, 1:2]) + img = torch.clamp(abs_val, 0, self.epe_clip) / self.epe_clip + valid = inputs.get("valid_disparities") + if valid is not None: + valid = valid[:, 1:2] if valid.shape[1] > 1 else valid elif name.startswith("abs"): - abs_val = torch.abs(preds[k] - inputs[k]) + abs_val = torch.abs(preds[k] - inputs[k][:, :1]) img = torch.clamp(abs_val, 0, self.epe_clip) / self.epe_clip - if inputs.get("valids") is not None: - img[inputs["valids"] < 0.5] = 0 + valid = inputs.get("valid_disparities") + if valid is not None: + valid = valid[:, :1] if valid.shape[1] > 1 else valid + elif name == "i_disparities2": + img = source[k] + valid = inputs.get("valid_disparities") + if valid is not None: + valid = valid[:, 1:2] if valid.shape[1] > 1 else valid else: img = source[k] @@ -384,6 +458,15 @@ def _append_images( # noqa: C901 [0] ) # BGR to RGB + if valid is not None: + valid = valid[:1, 0].detach().cpu() + valid = F.interpolate( + valid, self.image_size, mode="bilinear", align_corners=False + ) + invalid = valid[0] < 0.5 + invalid = invalid.repeat([img.shape[0], 1, 1]) + img[invalid] = 0 + images[name].append(img) def _compute_confidence_gt( @@ -410,29 +493,53 @@ def _compute_confidence_gt( ) return conf_gt - def _compute_max_range( - self, dataloader_length: int, limit_batches: Union[float, int] - ) -> int: - """Find the maximum number of samples that will be drawn from a dataloader. + def _compute_collect_idx(self, num_batches: Union[float, int]) -> np.ndarray: + """Choose which batch indices will have their images logged. + + num_images indices are uniformly spread over the batches that the current process will iterate over. Parameters ---------- - dataloader_length : int - Total size of the dataloader. - limit_batches : Union[float, int] - A value that may decrease the samples in the dataloader. See --limit_val_batches or --limit_train_batches from - PyTorch Lightning for more information. + num_batches : Union[float, int] + Number of batches that this process will iterate over, as reported by trainer.num_training_batches or + trainer.num_val_batches. It may be float("inf") for iterable datasets of unknown length. Returns ------- - int - The maximum number of samples that will be drawn from the dataloader. + np.ndarray + The sorted, unique batch indices to collect. """ - if isinstance(limit_batches, int): - max_range = limit_batches - 1 - else: - max_range = int(limit_batches * dataloader_length) - 1 - return max_range + if ( + num_batches is None + or num_batches != num_batches + or num_batches == float("inf") + ): + # Unknown length (iterable dataset): just take the first num_images batches. + return np.arange(self.num_images, dtype=np.int32) + num_batches = int(num_batches) + if num_batches <= 0: + return np.zeros(0, dtype=np.int32) + return np.unique( + np.linspace(0, num_batches - 1, self.num_images, dtype=np.int32) + ) + + def _val_collect_idx(self, dataloader_idx: int) -> np.ndarray: + """Get the batch indices to collect for one validation dataloader, tolerating a missing setup. + + Parameters + ---------- + dataloader_idx : int + The index number of the validation dataloader. + + Returns + ------- + np.ndarray + The batch indices to collect. Empty if the dataloader is unknown. + """ + if dataloader_idx >= len(self.val_dataloader_names): + return np.zeros(0, dtype=np.int32) + dl_name = self.val_dataloader_names[dataloader_idx] + return self.val_collect_image_idx.get(dl_name, np.zeros(0, dtype=np.int32)) def _make_image_grid( self, dl_images: Dict[str, List[torch.Tensor]] diff --git a/roco_spring_devkit/common/utils/flow_utils.py b/roco_spring_devkit/common/utils/flow_utils.py index 0f7a469..68f800f 100644 --- a/roco_spring_devkit/common/utils/flow_utils.py +++ b/roco_spring_devkit/common/utils/flow_utils.py @@ -266,6 +266,9 @@ def spring_epe_to_rgb( epe_rgb = cv.applyColorMap(epe, plt_lut) invalid_mask = ~valid_mask + if invalid_mask.ndim == 3 and invalid_mask.shape[2] == 1: + invalid_mask = np.tile(invalid_mask, [1, 1, 3]) + epe_rgb[invalid_mask] = 0 return epe_rgb diff --git a/roco_spring_devkit/common/utils/scene_flow_metrics.py b/roco_spring_devkit/common/utils/scene_flow_metrics.py index 26d2cc0..67a2ed4 100644 --- a/roco_spring_devkit/common/utils/scene_flow_metrics.py +++ b/roco_spring_devkit/common/utils/scene_flow_metrics.py @@ -43,6 +43,23 @@ class SceneFlowMetrics(Metric): full_state_update = True + # (state name, per-pixel value, validity mask) of every reported metric. Defined at the class level so that + # compute() also works when it is called before any update(). + METRIC_KEYS = ( + ("epe", "epe_flow", "valid_flows_target"), + ("1px_flow", "px1_flow_mask", "valid_flows_target"), + ("flall", "flall_mask", "valid_flows_target"), + ("wauc", "wauc_flow", "valid_flows_target"), + ("abs1", "abs1", "valid_disp1_target"), + ("1px1", "px11_mask", "valid_disp1_target"), + ("d1", "d1_mask", "valid_disp1_target"), + ("abs2", "abs2", "valid_disp2_target"), + ("1px2", "px12_mask", "valid_disp2_target"), + ("d2", "d2_mask", "valid_disp2_target"), + ("1px_all", "px1_all_mask", "valid_all_target"), + ("sfall", "sfall_mask", "valid_all_target"), + ) + def __init__( self, dist_sync_on_step: bool = False, @@ -109,7 +126,7 @@ def __init__( self.include_occlusion = False - self.used_keys = [] + self.used_keys = list(self.METRIC_KEYS) def update( self, preds: Dict[str, torch.Tensor], targets: Dict[str, torch.Tensor] @@ -222,19 +239,6 @@ def update( ).float() * 100 sfall_mask = ((flall_mask > 0) | (d1_mask > 0) | (d2_mask > 0)).float() * 100 - self.used_keys = [ - ("epe", "epe_flow", "valid_flows_target"), - ("1px_flow", "px1_flow_mask", "valid_flows_target"), - ("flall", "flall_mask", "valid_flows_target"), - ("wauc", "wauc_flow", "valid_flows_target"), - ("abs1", "abs1", "valid_disp1_target"), - ("d1", "d1_mask", "valid_disp1_target"), - ("abs2", "abs2", "valid_disp2_target"), - ("d2", "d2_mask", "valid_disp2_target"), - ("1px_all", "px1_all_mask", "valid_all_target"), - ("sfall", "sfall_mask", "valid_all_target"), - ] - for v1, v2, v3 in self.used_keys: if "wauc" not in v1: setattr( diff --git a/roco_spring_devkit/common/utils/stereo_utils.py b/roco_spring_devkit/common/utils/stereo_utils.py index d22f5bc..a291132 100644 --- a/roco_spring_devkit/common/utils/stereo_utils.py +++ b/roco_spring_devkit/common/utils/stereo_utils.py @@ -26,11 +26,226 @@ import numpy as np import png import torch +import torch.nn.functional as F from .external import selflow, flow_IO from .utils import get_matplotlib_lut_colormap +def pullback_disparity2( + disparity2: torch.Tensor, + valid_disparity2: Optional[torch.Tensor], + flow: torch.Tensor, + valid_flow: Optional[torch.Tensor] = None, + backward_flow: Optional[torch.Tensor] = None, + occlusion_mask: Optional[torch.Tensor] = None, + fb_check_tolerance: float = 0.5, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Resample a second-frame disparity map onto the first-frame pixel grid. + + Scene flow methods (and the KITTI/Spring benchmarks) represent the disparity of the second frame on the pixel grid of + the *first* frame, i.e., ``disparity2[y, x]`` is the disparity of the scene point that is seen at ``(x, y)`` in frame + 1 after it moved to ``(x, y) + flow[y, x]`` in frame 2. Datasets such as FlyingThings3D and Sintel instead store the + disparity of frame 2 on the pixel grid of frame 2, so it has to be pulled back through the forward optical flow + before it can be compared against such predictions. + + The values are sampled with ``padding_mode="border"`` so that pixels which fall outside of the image do not + contaminate their in-bounds neighbors with zeros during a later interpolation. Those pixels are always reported as + invalid in the returned mask, so the (meaningless) border values are never used. + + When ``backward_flow`` is provided, a forward-backward consistency check is applied to detect the so-called + "ghost" pixels. Backwarping the second-frame disparity with the forward flow makes every first-frame pixel ``p`` + sample the frame-2 location ``q = p + flow_fwd(p)``. If that location is actually occupied in frame 2 by a + *different* first-frame point (because the latter moved onto ``q``), the value copied into ``p`` is wrong: it is a + ghost. This happens for first-frame pixels that are occluded in frame 2, and it is detected by checking that the + backward flow at ``q`` points back to ``p`` (``q + flow_bwd(q) ~= p``). Pixels that fail this check are marked + invalid. + + Parameters + ---------- + disparity2 : torch.Tensor + A 4D tensor NCHW with the disparity of the second frame, on the pixel grid of the second frame. + valid_disparity2 : Optional[torch.Tensor] + A 4D tensor N1HW with the binary validity mask of ``disparity2``. May be None if no mask is available. + flow : torch.Tensor + A 4D tensor N2HW with the forward optical flow from frame 1 to frame 2. + valid_flow : Optional[torch.Tensor] + A 4D tensor N1HW with the binary validity mask of ``flow``. If provided, pixels without a valid flow are marked + as invalid, since their correspondence in frame 2 is unknown. + backward_flow : Optional[torch.Tensor] + A 4D tensor N2HW with the backward optical flow from frame 2 to frame 1, on the pixel grid of frame 2. If + provided, it is used to invalidate the first-frame pixels whose frame-2 correspondence is occupied by another + point (the "ghost" pixels). + occlusion_mask : Optional[torch.Tensor] + A 4D tensor N1HW with a binary occlusion mask on the first-frame pixel grid, where values >= 0.5 mark + first-frame pixels that are occluded in frame 2. If provided, those pixels are marked invalid, since their + frame-2 correspondence is unknown and the backwarped disparity would be a ghost. This is the Sintel + convention (the dataset ships explicit occlusion masks instead of backward flows). + fb_check_tolerance : float, default 0.5 + Maximum distance, in pixels, between ``p`` and ``q + flow_bwd(q)`` for a first-frame pixel ``p`` to be considered + consistent (and thus valid). Only used when ``backward_flow`` is provided. + + Returns + ------- + torch.Tensor + The disparity of the second frame, resampled onto the pixel grid of the first frame. + Optional[torch.Tensor] + The updated validity mask, or None if none of ``valid_disparity2``, ``valid_flow``, ``backward_flow``, and + ``occlusion_mask`` were provided. A pixel is valid only if (i) its flow is valid, (ii) its correspondence in + frame 2 falls inside the image, (iii) all of the frame-2 disparity values used by the bilinear interpolation + are valid, (iv) when a backward flow is given, the forward-backward consistency check passes (no ghost), and + (v) when an occlusion mask is given, the pixel is not marked as occluded. + """ + height, width = disparity2.shape[-2:] + grid_y, grid_x = torch.meshgrid( + torch.arange(height, dtype=flow.dtype, device=flow.device), + torch.arange(width, dtype=flow.dtype, device=flow.device), + indexing="ij", + ) + coords_x = grid_x[None] + flow[:, 0] + coords_y = grid_y[None] + flow[:, 1] + + norm_grid = torch.stack( + [ + 2.0 * coords_x / max(width - 1, 1) - 1.0, + 2.0 * coords_y / max(height - 1, 1) - 1.0, + ], + dim=-1, + ) + + pulled_disparity2 = F.grid_sample( + disparity2, + norm_grid, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + in_bounds = ( + (coords_x >= 0) + & (coords_x <= width - 1) + & (coords_y >= 0) + & (coords_y <= height - 1) + )[:, None] + + if ( + valid_disparity2 is None + and valid_flow is None + and backward_flow is None + and occlusion_mask is None + ): + return pulled_disparity2, None + + pulled_valid = in_bounds + if valid_disparity2 is not None: + # Bilinear sampling of a binary mask only returns 1 when all of the source pixels are valid. + sampled_valid = F.grid_sample( + (valid_disparity2 >= 0.5).to(dtype=disparity2.dtype), + norm_grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + pulled_valid = pulled_valid & (sampled_valid >= 1.0 - 1e-6) + if valid_flow is not None: + pulled_valid = pulled_valid & (valid_flow >= 0.5) + + if backward_flow is not None: + # Forward-backward consistency check: a first-frame pixel p samples q = p + flow_fwd(p) in frame 2. + # If q is actually occupied in frame 2 by a different first-frame point, the value copied into p is a + # "ghost". That is detected by sampling the backward flow at q and verifying that q + flow_bwd(q) points + # back to p. Pixels that fail are marked invalid. + sampled_bwd = F.grid_sample( + backward_flow, + norm_grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + pred_x = coords_x + sampled_bwd[:, 0] + pred_y = coords_y + sampled_bwd[:, 1] + fb_consistent = ( + (torch.abs(pred_x - grid_x[None]) <= fb_check_tolerance) + & (torch.abs(pred_y - grid_y[None]) <= fb_check_tolerance) + )[:, None] + pulled_valid = pulled_valid & fb_consistent + + if occlusion_mask is not None: + # The occlusion mask marks first-frame pixels that are occluded in frame 2. Their backwarped disparity is a + # ghost, so they are marked invalid. + pulled_valid = pulled_valid & (occlusion_mask < 0.5) + + return pulled_disparity2, pulled_valid.to(dtype=disparity2.dtype) + + +def compute_fb_occlusion_mask( + flow: torch.Tensor, + backward_flow: torch.Tensor, + tolerance: float = 0.5, +) -> torch.Tensor: + """Detect occluded first-frame pixels using forward-backward flow consistency. + + A first-frame pixel ``p`` is occluded in frame 2 if its forward-flow correspondence + ``q = p + flow_fwd(p)`` is occupied by a different first-frame point. This is detected + by sampling the backward flow at ``q`` and checking that ``q + flow_bwd(q)`` points back + to ``p``. Pixels that fail this check are marked as occluded (1.0) in the returned mask. + + This is the same check performed inside :func:`pullback_disparity2`, but exposed as a + standalone function so it can be applied to disparity maps that are *already* on the + first-frame pixel grid (e.g. Spring, which ships disp2 already mapped to frame 1 but + does not mark the occluded pixels as invalid in the file). + + Parameters + ---------- + flow : torch.Tensor + 4D tensor N2HW with the forward optical flow from frame 1 to frame 2, on the + first-frame pixel grid. + backward_flow : torch.Tensor + 4D tensor N2HW with the backward optical flow from frame 2 to frame 1, on the + second-frame pixel grid. + tolerance : float, default 0.5 + Maximum distance, in pixels, between ``p`` and ``q + flow_bwd(q)`` for a pixel to + be considered consistent (visible). + + Returns + ------- + torch.Tensor + 4D tensor N1HW where 1.0 means occluded and 0.0 means visible. + """ + height, width = flow.shape[-2:] + grid_y, grid_x = torch.meshgrid( + torch.arange(height, dtype=flow.dtype, device=flow.device), + torch.arange(width, dtype=flow.dtype, device=flow.device), + indexing="ij", + ) + coords_x = grid_x[None] + flow[:, 0] + coords_y = grid_y[None] + flow[:, 1] + + norm_grid = torch.stack( + [ + 2.0 * coords_x / max(width - 1, 1) - 1.0, + 2.0 * coords_y / max(height - 1, 1) - 1.0, + ], + dim=-1, + ) + + sampled_bwd = F.grid_sample( + backward_flow, + norm_grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) + pred_x = coords_x + sampled_bwd[:, 0] + pred_y = coords_y + sampled_bwd[:, 1] + fb_consistent = ( + (torch.abs(pred_x - grid_x[None]) <= tolerance) + & (torch.abs(pred_y - grid_y[None]) <= tolerance) + )[:, None] + + return (~fb_consistent).to(dtype=flow.dtype) + + class FileManager: def __init__(self, abstract_file, mode): self.abstract_file = abstract_file @@ -448,6 +663,9 @@ def spring_abs_to_rgb( epe_rgb = cv.applyColorMap(epe, plt_lut) invalid_mask = ~valid_mask + if invalid_mask.ndim == 3 and invalid_mask.shape[2] == 1: + invalid_mask = np.tile(invalid_mask, [1, 1, 3]) + epe_rgb[invalid_mask] = 0 return epe_rgb diff --git a/roco_spring_devkit/optical_flow/README.md b/roco_spring_devkit/optical_flow/README.md index 19f03e0..5ee0694 100644 --- a/roco_spring_devkit/optical_flow/README.md +++ b/roco_spring_devkit/optical_flow/README.md @@ -25,7 +25,7 @@ Three primary scripts are provided to manage your development workflow: Use this script to run inference on the competition test sets. For example, to generate outputs using the provided baseline RAFT model: ```bash -python test.py --data.test_dataset spring-robust --model raft --ckpt_path sintel --num_gpus -1 --save_viz +python test.py --data.test_dataset spring-robust --model raft --ckpt_path sintel --model.corr_mode triton --num_gpus -1 --save_viz ``` @@ -74,7 +74,7 @@ This helper script assesses your model's quality on target datasets containing g To validate the RAFT model on the Spring validation split: ```bash -python validate.py --data.val_dataset spring --model raft --ckpt_path sintel --write_outputs +python validate.py --data.val_dataset spring-val --model raft --ckpt_path sintel --model.corr_mode triton --write_outputs ``` diff --git a/roco_spring_devkit/optical_flow/models/raft/corr.py b/roco_spring_devkit/optical_flow/models/raft/corr.py index d0d974f..4254cf5 100644 --- a/roco_spring_devkit/optical_flow/models/raft/corr.py +++ b/roco_spring_devkit/optical_flow/models/raft/corr.py @@ -1,3 +1,5 @@ +import math + import torch import torch.nn.functional as F from .utils import bilinear_sampler @@ -6,7 +8,8 @@ import alt_cuda_corr except: alt_cuda_corr = None -from roco_spring_devkit.common.utils.correlation import IterativeCorrBlock + +from .triton_corr import TritonCorr class CorrBlock: @@ -101,18 +104,49 @@ def __call__(self, coords): return corr / torch.sqrt(torch.tensor(dim)) +class TritonCorrBlock: + def __init__(self, fmap1, fmap2, num_levels=4, radius=4): + self.num_levels = num_levels + self.radius = radius + + self.pyramid = [(fmap1, fmap2)] + for i in range(self.num_levels): + fmap1 = F.avg_pool2d(fmap1, 2, stride=2) + fmap2 = F.avg_pool2d(fmap2, 2, stride=2) + self.pyramid.append((fmap1, fmap2)) + + def __call__(self, coords): + coords = coords.permute(0, 2, 3, 1) + B, H, W, _ = coords.shape + dim = self.pyramid[0][0].shape[1] + + corr_list = [] + for i in range(self.num_levels): + r = self.radius + fmap1_i = self.pyramid[0][0].permute(0, 2, 3, 1).contiguous() + fmap2_i = self.pyramid[i][1].permute(0, 2, 3, 1).contiguous() + + coords_i = (coords / 2**i).reshape(B, 1, H, W, 2).contiguous() + corr = TritonCorr.apply(fmap1_i, fmap2_i, coords_i, r) + corr_list.append(corr.squeeze(1)) + + corr = torch.stack(corr_list, dim=1) + corr = corr.reshape(B, -1, H, W) + return corr / torch.sqrt(torch.tensor(dim)) + + def get_corr_block( fmap1: torch.Tensor, fmap2: torch.Tensor, num_levels: int = 4, radius: int = 4, - alternate_corr: bool = False, + corr_mode: str = "allpairs", ): - if alternate_corr: - if alt_cuda_corr is None: - corr_fn = IterativeCorrBlock - else: - corr_fn = AlternateCorrBlock - else: + if corr_mode.lower() == "allpairs": corr_fn = CorrBlock + elif corr_mode.lower() == "alt_cuda_corr": + corr_fn = AlternateCorrBlock + elif corr_mode.lower() == "triton": + corr_fn = TritonCorrBlock + return corr_fn(fmap1=fmap1, fmap2=fmap2, radius=radius, num_levels=num_levels) diff --git a/roco_spring_devkit/optical_flow/models/raft/raft.py b/roco_spring_devkit/optical_flow/models/raft/raft.py index 547ab93..df872ff 100644 --- a/roco_spring_devkit/optical_flow/models/raft/raft.py +++ b/roco_spring_devkit/optical_flow/models/raft/raft.py @@ -1,4 +1,5 @@ -from loguru import logger +from typing import Literal + import torch import torch.nn as nn import torch.nn.functional as F @@ -15,11 +16,6 @@ from .utils import coords_grid, upflow8 from ..base_model.base_model import BaseModel -try: - import alt_cuda_corr -except: - alt_cuda_corr = None - class SequenceLoss(nn.Module): def __init__(self, gamma: float, max_flow: float): @@ -61,11 +57,11 @@ def __init__( self, corr_levels: int = 4, corr_radius: int = 4, + corr_mode: Literal["allpairs", "alt_cuda_corr", "triton"] = "allpairs", dropout: float = 0.0, gamma: float = 0.8, max_flow: float = 400, iters: int = 32, - alternate_corr: bool = False, predict_all_directions: bool = True, **kwargs, ) -> None: @@ -75,11 +71,11 @@ def __init__( self.corr_levels = corr_levels self.corr_radius = corr_radius + self.corr_mode = corr_mode self.dropout = dropout self.gamma = gamma self.max_flow = max_flow self.iters = iters - self.alternate_corr = alternate_corr self.predict_all_directions = predict_all_directions self.hidden_dim = hdim = 128 @@ -96,11 +92,6 @@ def __init__( self.has_trained_on_ptlflow = True - if self.alternate_corr and alt_cuda_corr is None: - logger.warning( - "!!! alt_cuda_corr is not compiled! The slower IterativeCorrBlock will be used instead !!!" - ) - def freeze_bn(self): for m in self.modules(): if isinstance(m, nn.BatchNorm2d): @@ -168,7 +159,7 @@ def predict(self, image1, image2, image_resizer, prev_flow=None): fmap2=fmap2, radius=self.corr_radius, num_levels=self.corr_levels, - alternate_corr=self.alternate_corr, + corr_mode=self.corr_mode, ) # run the context network @@ -219,21 +210,21 @@ def __init__( self, corr_levels: int = 4, corr_radius: int = 3, + corr_mode: Literal["allpairs", "alt_cuda_corr", "triton"] = "allpairs", dropout: float = 0.0, gamma: float = 0.8, max_flow: float = 400, iters: int = 32, - alternate_corr: bool = False, **kwargs, ) -> None: super().__init__( corr_levels=corr_levels, corr_radius=corr_radius, + corr_mode=corr_mode, dropout=dropout, gamma=gamma, max_flow=max_flow, iters=iters, - alternate_corr=alternate_corr, **kwargs, ) self.hidden_dim = hdim = 96 diff --git a/roco_spring_devkit/optical_flow/models/raft/triton_corr.py b/roco_spring_devkit/optical_flow/models/raft/triton_corr.py new file mode 100644 index 0000000..b36cf83 --- /dev/null +++ b/roco_spring_devkit/optical_flow/models/raft/triton_corr.py @@ -0,0 +1,691 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def corr_forward_kernel( + fmap1_ptr, + fmap2_ptr, + coords_ptr, + corr_ptr, + B, + N, + H1, + W1, + H2, + W2, + C, + r, + stride_f1_b, + stride_f1_h, + stride_f1_w, + stride_f1_c, + stride_f2_b, + stride_f2_h, + stride_f2_w, + stride_f2_c, + stride_c_b, + stride_c_n, + stride_c_h, + stride_c_w, + stride_c_2, + stride_out_b, + stride_out_n, + stride_out_r, + stride_out_h, + stride_out_w, + BLOCK_HW: tl.constexpr, + BLOCK_C: tl.constexpr, +): + pid = tl.program_id(0) + + num_pixels = H1 * W1 + num_spatial_blocks = tl.cdiv(num_pixels, BLOCK_HW) + b = pid // num_spatial_blocks + spatial_block_id = pid % num_spatial_blocks + + hw_offsets = spatial_block_id * BLOCK_HW + tl.arange(0, BLOCK_HW) + hw_mask = hw_offsets < num_pixels + h1 = hw_offsets // W1 + w1 = hw_offsets % W1 + + rd = 2 * r + 1 + + for n in range(N): + coord_x_ptr = ( + coords_ptr + + b * stride_c_b + + n * stride_c_n + + h1 * stride_c_h + + w1 * stride_c_w + + 0 * stride_c_2 + ) + coord_y_ptr = ( + coords_ptr + + b * stride_c_b + + n * stride_c_n + + h1 * stride_c_h + + w1 * stride_c_w + + 1 * stride_c_2 + ) + + # tl.math.floor only accepts fp32/fp64. Keeping interpolation and the + # dot product in fp32 also avoids accumulating fp16 rounding error. + cx = tl.load(coord_x_ptr, mask=hw_mask, other=0.0).to(tl.float32) + cy = tl.load(coord_y_ptr, mask=hw_mask, other=0.0).to(tl.float32) + x0_base = tl.math.floor(cx) + y0_base = tl.math.floor(cy) + x0_base_int = x0_base.to(tl.int32) + y0_base_int = y0_base.to(tl.int32) + + dx = cx - x0_base + dy = cy - y0_base + wa = ((1.0 - dx) * (1.0 - dy))[:, None] + wb = (dx * (1.0 - dy))[:, None] + wc = ((1.0 - dx) * dy)[:, None] + wd = (dx * dy)[:, None] + base_ptr2 = fmap2_ptr + b * stride_f2_b + + for iy in range(rd): + for ix in range(rd): + x0_int = x0_base_int + (ix - r) + y0_int = y0_base_int + (iy - r) + x1_int = x0_int + 1 + y1_int = y0_int + 1 + + valid_a = (x0_int >= 0) & (x0_int < W2) & (y0_int >= 0) & (y0_int < H2) + valid_b = (x1_int >= 0) & (x1_int < W2) & (y0_int >= 0) & (y0_int < H2) + valid_c = (x0_int >= 0) & (x0_int < W2) & (y1_int >= 0) & (y1_int < H2) + valid_d = (x1_int >= 0) & (x1_int < W2) & (y1_int >= 0) & (y1_int < H2) + + corr_val = tl.zeros([BLOCK_HW], dtype=tl.float32) + for c_start in range(0, C, BLOCK_C): + c_offsets = c_start + tl.arange(0, BLOCK_C) + c_mask = c_offsets < C + mask_2d = hw_mask[:, None] & c_mask[None, :] + + f1_ptrs = ( + fmap1_ptr + + b * stride_f1_b + + h1[:, None] * stride_f1_h + + w1[:, None] * stride_f1_w + + c_offsets[None, :] * stride_f1_c + ) + f1 = tl.load(f1_ptrs, mask=mask_2d, other=0.0).to(tl.float32) + + ptrs_a = ( + base_ptr2 + + y0_int[:, None] * stride_f2_h + + x0_int[:, None] * stride_f2_w + + c_offsets[None, :] * stride_f2_c + ) + ptrs_b = ( + base_ptr2 + + y0_int[:, None] * stride_f2_h + + x1_int[:, None] * stride_f2_w + + c_offsets[None, :] * stride_f2_c + ) + ptrs_c = ( + base_ptr2 + + y1_int[:, None] * stride_f2_h + + x0_int[:, None] * stride_f2_w + + c_offsets[None, :] * stride_f2_c + ) + ptrs_d = ( + base_ptr2 + + y1_int[:, None] * stride_f2_h + + x1_int[:, None] * stride_f2_w + + c_offsets[None, :] * stride_f2_c + ) + + val_a = tl.load( + ptrs_a, mask=mask_2d & valid_a[:, None], other=0.0 + ).to(tl.float32) + val_b = tl.load( + ptrs_b, mask=mask_2d & valid_b[:, None], other=0.0 + ).to(tl.float32) + val_c = tl.load( + ptrs_c, mask=mask_2d & valid_c[:, None], other=0.0 + ).to(tl.float32) + val_d = tl.load( + ptrs_d, mask=mask_2d & valid_d[:, None], other=0.0 + ).to(tl.float32) + + f2 = val_a * wa + val_b * wb + val_c * wc + val_d * wd + corr_val += tl.sum(f1 * f2, axis=1) + + idx_r = ix * rd + iy + out_ptr = ( + corr_ptr + + b * stride_out_b + + n * stride_out_n + + idx_r * stride_out_r + + h1 * stride_out_h + + w1 * stride_out_w + ) + tl.store(out_ptr, corr_val, mask=hw_mask) + + +@triton.jit +def corr_pyramid_forward_kernel( + fmap1_ptr, + fmap2_packed_ptr, + coords_ptr, + fmap2_offsets_ptr, + fmap2_shapes_ptr, + level_scales_ptr, + corr_ptr, + B, + H1, + W1, + C, + r, + stride_f1_b, + stride_f1_h, + stride_f1_w, + stride_f1_c, + stride_c_b, + stride_c_2, + stride_c_h, + stride_c_w, + stride_out_b, + stride_out_level, + stride_out_r, + stride_out_h, + stride_out_w, + BLOCK_HW: tl.constexpr, + BLOCK_C: tl.constexpr, +): + pid = tl.program_id(0) + level = tl.program_id(1) + + num_pixels = H1 * W1 + num_spatial_blocks = tl.cdiv(num_pixels, BLOCK_HW) + b = pid // num_spatial_blocks + spatial_block_id = pid % num_spatial_blocks + + hw_offsets = spatial_block_id * BLOCK_HW + tl.arange(0, BLOCK_HW) + hw_mask = hw_offsets < num_pixels + h1 = hw_offsets // W1 + w1 = hw_offsets % W1 + + h2 = tl.load(fmap2_shapes_ptr + 2 * level) + w2 = tl.load(fmap2_shapes_ptr + 2 * level + 1) + fmap2_offset = tl.load(fmap2_offsets_ptr + level) + level_scale = tl.load(level_scales_ptr + level) + fmap2_base = fmap2_packed_ptr + fmap2_offset + b * h2 * w2 * C + + cx = ( + tl.load( + coords_ptr + b * stride_c_b + h1 * stride_c_h + w1 * stride_c_w, + mask=hw_mask, + other=0.0, + ).to(tl.float32) + * level_scale + ) + cy = ( + tl.load( + coords_ptr + + b * stride_c_b + + stride_c_2 + + h1 * stride_c_h + + w1 * stride_c_w, + mask=hw_mask, + other=0.0, + ).to(tl.float32) + * level_scale + ) + x0_base = tl.math.floor(cx) + y0_base = tl.math.floor(cy) + x0_base_int = x0_base.to(tl.int32) + y0_base_int = y0_base.to(tl.int32) + + dx = cx - x0_base + dy = cy - y0_base + wa = ((1.0 - dx) * (1.0 - dy))[:, None] + wb = (dx * (1.0 - dy))[:, None] + wc = ((1.0 - dx) * dy)[:, None] + wd = (dx * dy)[:, None] + + rd = 2 * r + 1 + for iy in range(rd): + for ix in range(rd): + x0_int = x0_base_int + (ix - r) + y0_int = y0_base_int + (iy - r) + x1_int = x0_int + 1 + y1_int = y0_int + 1 + + valid_a = (x0_int >= 0) & (x0_int < w2) & (y0_int >= 0) & (y0_int < h2) + valid_b = (x1_int >= 0) & (x1_int < w2) & (y0_int >= 0) & (y0_int < h2) + valid_c = (x0_int >= 0) & (x0_int < w2) & (y1_int >= 0) & (y1_int < h2) + valid_d = (x1_int >= 0) & (x1_int < w2) & (y1_int >= 0) & (y1_int < h2) + + corr_val = tl.zeros([BLOCK_HW], dtype=tl.float32) + for c_start in range(0, C, BLOCK_C): + c_offsets = c_start + tl.arange(0, BLOCK_C) + c_mask = c_offsets < C + mask_2d = hw_mask[:, None] & c_mask[None, :] + + f1_ptrs = ( + fmap1_ptr + + b * stride_f1_b + + h1[:, None] * stride_f1_h + + w1[:, None] * stride_f1_w + + c_offsets[None, :] * stride_f1_c + ) + f1 = tl.load(f1_ptrs, mask=mask_2d, other=0.0).to(tl.float32) + + ptrs_a = ( + fmap2_base + + (y0_int[:, None] * w2 + x0_int[:, None]) * C + + c_offsets[None, :] + ) + ptrs_b = ( + fmap2_base + + (y0_int[:, None] * w2 + x1_int[:, None]) * C + + c_offsets[None, :] + ) + ptrs_c = ( + fmap2_base + + (y1_int[:, None] * w2 + x0_int[:, None]) * C + + c_offsets[None, :] + ) + ptrs_d = ( + fmap2_base + + (y1_int[:, None] * w2 + x1_int[:, None]) * C + + c_offsets[None, :] + ) + + val_a = tl.load(ptrs_a, mask=mask_2d & valid_a[:, None], other=0.0).to( + tl.float32 + ) + val_b = tl.load(ptrs_b, mask=mask_2d & valid_b[:, None], other=0.0).to( + tl.float32 + ) + val_c = tl.load(ptrs_c, mask=mask_2d & valid_c[:, None], other=0.0).to( + tl.float32 + ) + val_d = tl.load(ptrs_d, mask=mask_2d & valid_d[:, None], other=0.0).to( + tl.float32 + ) + + f2 = val_a * wa + val_b * wb + val_c * wc + val_d * wd + corr_val += tl.sum(f1 * f2, axis=1) + + idx_r = ix * rd + iy + out_ptr = ( + corr_ptr + + b * stride_out_b + + level * stride_out_level + + idx_r * stride_out_r + + h1 * stride_out_h + + w1 * stride_out_w + ) + tl.store(out_ptr, corr_val, mask=hw_mask) + + +@triton.jit +def corr_backward_kernel( + fmap1_ptr, + fmap2_ptr, + coords_ptr, + corr_grad_ptr, + fmap1_grad_ptr, + fmap2_grad_ptr, + B, + N, + H1, + W1, + H2, + W2, + C, + r, + stride_f1_b, + stride_f1_h, + stride_f1_w, + stride_f1_c, + stride_f2_b, + stride_f2_h, + stride_f2_w, + stride_f2_c, + stride_c_b, + stride_c_n, + stride_c_h, + stride_c_w, + stride_c_2, + stride_cg_b, + stride_cg_n, + stride_cg_r, + stride_cg_h, + stride_cg_w, + BLOCK_C: tl.constexpr, +): + pid = tl.program_id(0) + num_pixels = H1 * W1 + b = pid // num_pixels + hw = pid % num_pixels + h1 = hw // W1 + w1 = hw % W1 + + c_offsets = tl.arange(0, BLOCK_C) + valid_channels = c_offsets < C + + f1_ptrs = ( + fmap1_ptr + + b * stride_f1_b + + h1 * stride_f1_h + + w1 * stride_f1_w + + c_offsets * stride_f1_c + ) + f1 = tl.load(f1_ptrs, mask=valid_channels, other=0.0).to(tl.float32) + + f1_grad = tl.zeros([BLOCK_C], dtype=tl.float32) + rd = 2 * r + 1 + + for n in range(N): + coord_x_ptr = ( + coords_ptr + + b * stride_c_b + + n * stride_c_n + + h1 * stride_c_h + + w1 * stride_c_w + + 0 * stride_c_2 + ) + coord_y_ptr = ( + coords_ptr + + b * stride_c_b + + n * stride_c_n + + h1 * stride_c_h + + w1 * stride_c_w + + 1 * stride_c_2 + ) + + cx = tl.load(coord_x_ptr).to(tl.float32) + cy = tl.load(coord_y_ptr).to(tl.float32) + x0_base = tl.math.floor(cx) + y0_base = tl.math.floor(cy) + x0_base_int = x0_base.to(tl.int32) + y0_base_int = y0_base.to(tl.int32) + + dx = cx - x0_base + dy = cy - y0_base + wa = (1.0 - dx) * (1.0 - dy) + wb = dx * (1.0 - dy) + wc = (1.0 - dx) * dy + wd = dx * dy + base_ptr2 = fmap2_ptr + b * stride_f2_b + base_ptr2_grad = fmap2_grad_ptr + b * stride_f2_b + + for iy in range(rd): + for ix in range(rd): + x0_int = x0_base_int + (ix - r) + y0_int = y0_base_int + (iy - r) + x1_int = x0_int + 1 + y1_int = y0_int + 1 + + valid_a = (x0_int >= 0) & (x0_int < W2) & (y0_int >= 0) & (y0_int < H2) + valid_b = (x1_int >= 0) & (x1_int < W2) & (y0_int >= 0) & (y0_int < H2) + valid_c = (x0_int >= 0) & (x0_int < W2) & (y1_int >= 0) & (y1_int < H2) + valid_d = (x1_int >= 0) & (x1_int < W2) & (y1_int >= 0) & (y1_int < H2) + + ptrs_a = ( + base_ptr2 + + y0_int * stride_f2_h + + x0_int * stride_f2_w + + c_offsets * stride_f2_c + ) + ptrs_b = ( + base_ptr2 + + y0_int * stride_f2_h + + x1_int * stride_f2_w + + c_offsets * stride_f2_c + ) + ptrs_c = ( + base_ptr2 + + y1_int * stride_f2_h + + x0_int * stride_f2_w + + c_offsets * stride_f2_c + ) + ptrs_d = ( + base_ptr2 + + y1_int * stride_f2_h + + x1_int * stride_f2_w + + c_offsets * stride_f2_c + ) + + val_a = tl.load(ptrs_a, mask=valid_channels & valid_a, other=0.0).to( + tl.float32 + ) + val_b = tl.load(ptrs_b, mask=valid_channels & valid_b, other=0.0).to( + tl.float32 + ) + val_c = tl.load(ptrs_c, mask=valid_channels & valid_c, other=0.0).to( + tl.float32 + ) + val_d = tl.load(ptrs_d, mask=valid_channels & valid_d, other=0.0).to( + tl.float32 + ) + + idx_r = ix * rd + iy + grad_ptr = ( + corr_grad_ptr + + b * stride_cg_b + + n * stride_cg_n + + idx_r * stride_cg_r + + h1 * stride_cg_h + + w1 * stride_cg_w + ) + g = tl.load(grad_ptr).to(tl.float32) + + # Accumulate f1 gradients + f2 = wa * val_a + wb * val_b + wc * val_c + wd * val_d + f1_grad += g * f2 + + # Scatter f2 gradients via atomic additions + gf1 = g * f1 + tl.atomic_add( + base_ptr2_grad + + y0_int * stride_f2_h + + x0_int * stride_f2_w + + c_offsets * stride_f2_c, + gf1 * wa, + mask=valid_channels & valid_a, + ) + tl.atomic_add( + base_ptr2_grad + + y0_int * stride_f2_h + + x1_int * stride_f2_w + + c_offsets * stride_f2_c, + gf1 * wb, + mask=valid_channels & valid_b, + ) + tl.atomic_add( + base_ptr2_grad + + y1_int * stride_f2_h + + x0_int * stride_f2_w + + c_offsets * stride_f2_c, + gf1 * wc, + mask=valid_channels & valid_c, + ) + tl.atomic_add( + base_ptr2_grad + + y1_int * stride_f2_h + + x1_int * stride_f2_w + + c_offsets * stride_f2_c, + gf1 * wd, + mask=valid_channels & valid_d, + ) + + # Store final f1_grad at the end + f1_grad_ptrs = ( + fmap1_grad_ptr + + b * stride_f1_b + + h1 * stride_f1_h + + w1 * stride_f1_w + + c_offsets * stride_f1_c + ) + tl.store(f1_grad_ptrs, f1_grad, mask=valid_channels) + + +class TritonCorr(torch.autograd.Function): + @staticmethod + def forward(ctx, fmap1, fmap2_i, coords, r): + if torch.onnx.is_in_onnx_export(): + batch_size, height_1, width_1, _ = fmap1.shape + _, num_slices, _, _, _ = coords.shape + radius = int(r) + window_size = 2 * radius + 1 + corr = torch.empty( + (batch_size, num_slices, window_size * window_size, height_1, width_1), + device=fmap1.device, + dtype=fmap1.dtype, + ) + return (corr,) + + ctx.save_for_backward(fmap1, fmap2_i, coords) + ctx.r = r + + B, H1, W1, C = fmap1.shape + _, N, _, _, _ = coords.shape + _, H2, W2, _ = fmap2_i.shape + B = int(B) + H1 = int(H1) + W1 = int(W1) + C = int(C) + N = int(N) + H2 = int(H2) + W2 = int(W2) + r = int(r) + + rd = 2 * r + 1 + corr = torch.empty( + (B, N, rd * rd, H1, W1), device=fmap1.device, dtype=fmap1.dtype + ) + + BLOCK_HW = 16 + BLOCK_C = min(64, triton.next_power_of_2(C)) + grid = (B * triton.cdiv(H1 * W1, BLOCK_HW),) + + corr_forward_kernel[grid]( + fmap1, + fmap2_i, + coords, + corr, + B, + N, + H1, + W1, + H2, + W2, + C, + r, + fmap1.stride(0), + fmap1.stride(1), + fmap1.stride(2), + fmap1.stride(3), + fmap2_i.stride(0), + fmap2_i.stride(1), + fmap2_i.stride(2), + fmap2_i.stride(3), + coords.stride(0), + coords.stride(1), + coords.stride(2), + coords.stride(3), + coords.stride(4), + corr.stride(0), + corr.stride(1), + corr.stride(2), + corr.stride(3), + corr.stride(4), + BLOCK_HW=BLOCK_HW, + BLOCK_C=BLOCK_C, + ) + return corr + + @staticmethod + def backward(ctx, corr_grad): + fmap1, fmap2_i, coords = ctx.saved_tensors + r = int(ctx.r) + + corr_grad = corr_grad.contiguous() + + B, H1, W1, C = fmap1.shape + _, N, _, _, _ = coords.shape + _, H2, W2, _ = fmap2_i.shape + B = int(B) + H1 = int(H1) + W1 = int(W1) + C = int(C) + N = int(N) + H2 = int(H2) + W2 = int(W2) + + fmap1_grad = torch.empty_like(fmap1) + fmap2_grad = torch.zeros_like(fmap2_i) + + # The RAFT architecture stops gradients at the coords node during iterative updates. + # We can return a tensor of zeros to align with original autograd footprints. + coords_grad = torch.zeros_like(coords) + + BLOCK_C = triton.next_power_of_2(C) + grid = (B * H1 * W1,) + + corr_backward_kernel[grid]( + fmap1, + fmap2_i, + coords, + corr_grad, + fmap1_grad, + fmap2_grad, + B, + N, + H1, + W1, + H2, + W2, + C, + r, + fmap1.stride(0), + fmap1.stride(1), + fmap1.stride(2), + fmap1.stride(3), + fmap2_i.stride(0), + fmap2_i.stride(1), + fmap2_i.stride(2), + fmap2_i.stride(3), + coords.stride(0), + coords.stride(1), + coords.stride(2), + coords.stride(3), + coords.stride(4), + corr_grad.stride(0), + corr_grad.stride(1), + corr_grad.stride(2), + corr_grad.stride(3), + corr_grad.stride(4), + BLOCK_C=BLOCK_C, + ) + + return fmap1_grad, fmap2_grad, coords_grad, None + + @staticmethod + def symbolic(g, fmap1, fmap2_i, coords, r): + # This writes "custom_ops::TritonCorr" into the ONNX file + node = g.op("custom_ops::TritonCorr", fmap1, fmap2_i, coords, r_i=r) + + fmap1_sizes = fmap1.type().sizes() + coords_sizes = coords.type().sizes() + if fmap1_sizes is not None and coords_sizes is not None: + radius = int(r) + window_size = 2 * radius + 1 + output_sizes = [ + coords_sizes[0], + coords_sizes[1], + window_size * window_size, + fmap1_sizes[1], + fmap1_sizes[2], + ] + node.setType(fmap1.type().with_sizes(output_sizes)) + + return node diff --git a/roco_spring_devkit/optical_flow/validate.py b/roco_spring_devkit/optical_flow/validate.py index 066fbfb..c2cc751 100644 --- a/roco_spring_devkit/optical_flow/validate.py +++ b/roco_spring_devkit/optical_flow/validate.py @@ -378,11 +378,11 @@ def validate_one_dataloader( cuda=torch.cuda.is_available(), fp16=args.fp16, ) - inputs = io_adapter.prepare_inputs(inputs=inputs, image_only=True) + inputs = io_adapter.prepare_inputs(inputs=inputs) outputs = model.validation_step(inputs, i, dataloader_idx) - inputs = io_adapter.unscale(inputs, image_only=True) + inputs = io_adapter.unscale(inputs) preds = outputs["preds"] preds = io_adapter.unscale(preds) for k, v in inputs.items(): diff --git a/roco_spring_devkit/scene_flow/README.md b/roco_spring_devkit/scene_flow/README.md index 74a4198..6a5c940 100644 --- a/roco_spring_devkit/scene_flow/README.md +++ b/roco_spring_devkit/scene_flow/README.md @@ -75,7 +75,7 @@ This helper script assesses your model's quality on target datasets containing g To validate the RAFT-3D-Laplacian model on the Spring validation split: ```bash -python validate.py --data.val_dataset spring --model raft_3d_bilaplacian --ckpt_path laplacian --model.compute_disparity true --model.disparity_ckpt sceneflow --write_outputs +python validate.py --data.val_dataset spring-val --model raft_3d_bilaplacian --ckpt_path laplacian --model.compute_disparity true --model.disparity_ckpt sceneflow --write_outputs ``` diff --git a/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d.py b/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d.py index 6eaec76..9c101bf 100644 --- a/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d.py +++ b/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d.py @@ -218,9 +218,9 @@ def __init__( gamma=0.9, dz_weight=0.2, rv_weight=100.0, - disp2_forward_warp=True, + disp2_forward_warp=False, compute_disparity=True, - disparity_ckpt=None, + disparity_ckpt="sceneflow", predict_all_directions=True, **kwargs, ): diff --git a/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d_bilaplacian.py b/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d_bilaplacian.py index 7a1d0d5..f7cb14c 100644 --- a/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d_bilaplacian.py +++ b/roco_spring_devkit/scene_flow/models/raft_3d/raft_3d_bilaplacian.py @@ -262,9 +262,9 @@ def __init__( gamma=0.9, dz_weight=0.2, rv_weight=100.0, - disp2_forward_warp=True, - compute_disparity=False, - disparity_ckpt=None, + disp2_forward_warp=False, + compute_disparity=True, + disparity_ckpt="sceneflow", predict_all_directions=True, **kwargs, ): diff --git a/roco_spring_devkit/scene_flow/validate.py b/roco_spring_devkit/scene_flow/validate.py index c19beb2..d6aeaee 100644 --- a/roco_spring_devkit/scene_flow/validate.py +++ b/roco_spring_devkit/scene_flow/validate.py @@ -482,11 +482,11 @@ def validate_one_dataloader( cuda=torch.cuda.is_available(), fp16=args.fp16, ) - inputs = io_adapter.prepare_inputs(inputs=inputs, image_only=True) + inputs = io_adapter.prepare_inputs(inputs=inputs) outputs = model.validation_step(inputs, i, dataloader_idx) - inputs = io_adapter.unscale(inputs, image_only=True) + inputs = io_adapter.unscale(inputs) preds = outputs["preds"] preds = io_adapter.unscale(preds) for k, v in inputs.items(): diff --git a/roco_spring_devkit/stereo/README.md b/roco_spring_devkit/stereo/README.md index cc4ff05..14ef9b0 100644 --- a/roco_spring_devkit/stereo/README.md +++ b/roco_spring_devkit/stereo/README.md @@ -73,7 +73,7 @@ This helper script assesses your model's quality on target datasets containing g To validate the RAFT-Stereo model on the Spring validation split: ```bash -python validate.py --data.val_dataset spring --model raft_stereo --ckpt_path sceneflow --write_outputs +python validate.py --data.val_dataset spring-val --model raft_stereo --ckpt_path sceneflow --write_outputs ``` diff --git a/roco_spring_devkit/stereo/validate.py b/roco_spring_devkit/stereo/validate.py index 2b9351c..b3ffa15 100644 --- a/roco_spring_devkit/stereo/validate.py +++ b/roco_spring_devkit/stereo/validate.py @@ -383,11 +383,11 @@ def validate_one_dataloader( cuda=torch.cuda.is_available(), fp16=args.fp16, ) - inputs = io_adapter.prepare_inputs(inputs=inputs, image_only=True) + inputs = io_adapter.prepare_inputs(inputs=inputs) outputs = model.validation_step(inputs, i, dataloader_idx) - inputs = io_adapter.unscale(inputs, image_only=True) + inputs = io_adapter.unscale(inputs) preds = outputs["preds"] preds = io_adapter.unscale(preds) diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..1e55948 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,11 @@ +Run standard tests: + +```bash +python -m pytest tests +``` + +Include intensive tests (require GPU and internet): + +```bash +python -m pytest tests --runslow +``` \ No newline at end of file diff --git a/tests/common/data/test_optical_flow_datamodule.py b/tests/common/data/test_optical_flow_datamodule.py new file mode 100644 index 0000000..3e1c982 --- /dev/null +++ b/tests/common/data/test_optical_flow_datamodule.py @@ -0,0 +1,1076 @@ +"""Unit tests for `roco_spring_devkit.common.data.optical_flow_datamodule`. + +The datasets and transforms themselves have their own dedicated test modules +(`test_optical_flow_datasets.py`, `test_optical_flow_transforms.py`), so the +focus here is the *datamodule-specific* logic: + +* the default values set by ``__init__``; +* the string-parsing helper ``_parse_dataset_selection``; +* the YAML path loader ``_load_dataset_paths`` (and the fact that it does NOT + overwrite paths that were already provided); +* the ``setup`` stage dispatch and assertions; +* ``_get_model_output_stride`` with / without a trainer; +* the dataloader builders ``train_dataloader`` / ``val_dataloader`` / + ``test_dataloader`` / ``predict_dataloader`` (including the Sintel + test-split expansion and the dataset multiplier behaviour); +* the argument-parsing logic of every ``_get__dataset`` helper and the + ``ValueError`` raised on unknown flags. + +The on-disk dataset constructors are mocked out so that we only exercise the +datamodule plumbing, not the (already tested) dataset readers. +""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import yaml + +import roco_spring_devkit.common.data.optical_flow_datamodule as mod +from roco_spring_devkit.common.data.optical_flow_datamodule import FlowDataModule + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_dummy_dataset(length: int = 4): + """Return a MagicMock that behaves like a concatenatable Dataset. + + `__add__` returns the same mock so that train_dataloader can keep using + `+=` (`dataset = dataset + dataset`) without growing new mock objects. + `__len__` returns ``length``. + """ + + dataset = MagicMock(name="DummyDataset") + dataset.__len__ = MagicMock(return_value=length) + dataset.__add__ = MagicMock(return_value=dataset) + dataset.__iadd__ = MagicMock(return_value=dataset) + return dataset + + +def _write_config(tmp_path: Path, paths: dict) -> str: + """Write a YAML dataset config and return its path as a string.""" + cfg_path = tmp_path / "datasets.yaml" + with open(cfg_path, "w") as f: + yaml.safe_dump(paths, f) + return str(cfg_path) + + +# =========================================================================== +# Tests for __init__ defaults +# =========================================================================== +class TestInit: + def test_defaults_are_none_or_zero(self): + dm = FlowDataModule() + assert dm.predict_dataset is None + assert dm.test_dataset is None + assert dm.train_dataset is None + assert dm.val_dataset is None + assert dm.train_batch_size is None + assert dm.train_num_workers == 4 + assert dm.train_crop_size is None + assert dm.train_transform_cuda is False + assert dm.train_transform_fp16 is False + + def test_root_dirs_default_to_none(self): + dm = FlowDataModule() + for attr in ( + "autoflow_root_dir", + "flying_chairs_root_dir", + "flying_chairs2_root_dir", + "flying_things3d_root_dir", + "flying_things3d_subset_root_dir", + "mpi_sintel_root_dir", + "kitti_2012_root_dir", + "kitti_2015_root_dir", + "hd1k_root_dir", + "tartanair_root_dir", + "spring_root_dir", + "robust_spring_root_dir", + "kubric_root_dir", + "middlebury_st_root_dir", + "viper_root_dir", + ): + assert getattr(dm, attr) is None, attr + + def test_init_stores_constructor_values(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule( + train_dataset="chairs-train", + val_dataset="sintel-clean-trainval", + train_batch_size=3, + train_num_workers=2, + train_crop_size=(100, 200), + train_transform_cuda=True, + train_transform_fp16=True, + dataset_config_path=cfg, + ) + assert dm.train_dataset == "chairs-train" + assert dm.val_dataset == "sintel-clean-trainval" + assert dm.train_batch_size == 3 + assert dm.train_num_workers == 2 + assert dm.train_crop_size == (100, 200) + assert dm.train_transform_cuda is True + assert dm.train_transform_fp16 is True + assert dm.dataset_config_path == cfg + + def test_init_initialises_internal_state(self): + dm = FlowDataModule() + assert dm.predict_dataset_parsed is None + assert dm.test_dataset_parsed is None + assert dm.train_dataset_parsed is None + assert dm.val_dataset_parsed is None + assert dm.train_dataloader_length == 0 + assert dm.train_epoch_step == 0 + assert dm.val_dataloader_names == [] + assert dm.val_dataloader_lengths == [] + assert dm.test_dataloader_names == [] + + +# =========================================================================== +# Tests for _parse_dataset_selection +# =========================================================================== +class TestParseDatasetSelection: + def test_none_returns_empty_list(self): + dm = FlowDataModule() + assert dm._parse_dataset_selection(None) == [] + + def test_single_dataset_default_multiplier(self): + dm = FlowDataModule() + out = dm._parse_dataset_selection("chairs-train") + assert out == [(1, "chairs", "train")] + + def test_multiple_datasets_separated_by_plus(self): + dm = FlowDataModule() + out = dm._parse_dataset_selection( + "chairs-train+sintel-clean-trainval+kitti-2012-train" + ) + assert out == [ + (1, "chairs", "train"), + (1, "sintel", "clean", "trainval"), + (1, "kitti", "2012", "train"), + ] + + def test_leading_multiplier(self): + dm = FlowDataModule() + out = dm._parse_dataset_selection("3*sintel-clean-trainval") + assert out == [(3, "sintel", "clean", "trainval")] + + def test_trailing_multiplier(self): + dm = FlowDataModule() + out = dm._parse_dataset_selection("kitti-2012-train*5") + assert out == [(5, "kitti", "2012", "train")] + + def test_mixed_multipliers_and_flags(self): + dm = FlowDataModule() + out = dm._parse_dataset_selection( + "chairs-train+3*sintel-clean-trainval+kitti-2012-train*5" + ) + assert out == [ + (1, "chairs", "train"), + (3, "sintel", "clean", "trainval"), + (5, "kitti", "2012", "train"), + ] + + def test_spaces_are_stripped(self): + dm = FlowDataModule() + out = dm._parse_dataset_selection(" chairs - train + 2 * sintel - clean ") + assert out == [ + (1, "chairs", "train"), + (2, "sintel", "clean"), + ] + + def test_invalid_double_multiplier_raises(self): + dm = FlowDataModule() + with pytest.raises(ValueError): + dm._parse_dataset_selection("3*4*sintel-clean") + + def test_non_integer_multiplier_raises_value_error(self): + # 'x' is not int -> the trailing-multiplier branch raises ValueError + # when trying int('x'). + dm = FlowDataModule() + with pytest.raises(ValueError): + dm._parse_dataset_selection("sintel*x") + + def test_empty_string_yields_single_empty_dataset(self): + # An empty string splits into [''] which is treated as a single + # dataset with empty name and multiplier 1. This documents current + # behaviour (no guard for empty input). + dm = FlowDataModule() + out = dm._parse_dataset_selection("") + assert out == [(1, "")] + + +# =========================================================================== +# Tests for _load_dataset_paths +# =========================================================================== +class TestLoadDatasetPaths: + def test_fills_none_attrs_from_yaml(self, tmp_path): + cfg = _write_config( + tmp_path, + { + "autoflow": "/aut/oflow", + "flying_chairs": "/chairs/path", + "mpi_sintel": "/sintel/path", + }, + ) + dm = FlowDataModule(dataset_config_path=cfg) + dm._load_dataset_paths() + assert dm.autoflow_root_dir == "/aut/oflow" + assert dm.flying_chairs_root_dir == "/chairs/path" + assert dm.mpi_sintel_root_dir == "/sintel/path" + + def test_does_not_overwrite_existing_paths(self, tmp_path): + cfg = _write_config( + tmp_path, + {"autoflow": "/yaml/autoflow", "flying_chairs": "/yaml/chairs"}, + ) + dm = FlowDataModule( + autoflow_root_dir="/explicit/autoflow", + dataset_config_path=cfg, + ) + dm._load_dataset_paths() + # The explicitly-provided path wins. + assert dm.autoflow_root_dir == "/explicit/autoflow" + # The unset one is filled from YAML. + assert dm.flying_chairs_root_dir == "/yaml/chairs" + + def test_missing_keys_in_yaml_leave_attrs_as_none(self, tmp_path): + cfg = _write_config(tmp_path, {"autoflow": "/aut/oflow"}) + dm = FlowDataModule(dataset_config_path=cfg) + dm._load_dataset_paths() + assert dm.autoflow_root_dir == "/aut/oflow" + # Keys absent from YAML are left untouched (still None). + assert dm.flying_chairs_root_dir is None + assert dm.mpi_sintel_root_dir is None + + +# =========================================================================== +# Tests for setup +# =========================================================================== +class TestSetup: + def _make(self, tmp_path, **kwargs): + cfg = _write_config(tmp_path, {}) + kwargs.setdefault("dataset_config_path", cfg) + return FlowDataModule(**kwargs) + + def test_fit_parses_train_and_val(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="chairs-train", + val_dataset="sintel-clean-trainval", + train_batch_size=4, + ) + dm.setup("fit") + assert dm.train_dataset_parsed == [(1, "chairs", "train")] + assert dm.val_dataset_parsed == [(1, "sintel", "clean", "trainval")] + + def test_fit_assigns_default_batch_size_when_none(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="chairs-train", + val_dataset="sintel-clean-trainval", + train_batch_size=None, + ) + dm.setup("fit") + assert dm.train_batch_size == 8 + + def test_fit_loads_dataset_paths(self, tmp_path): + cfg = _write_config(tmp_path, {"autoflow": "/auto/path"}) + dm = FlowDataModule( + train_dataset="chairs-train", + val_dataset="sintel-clean-trainval", + train_batch_size=4, + dataset_config_path=cfg, + ) + dm.setup("fit") + assert dm.autoflow_root_dir == "/auto/path" + + def test_fit_without_train_dataset_raises(self, tmp_path): + dm = self._make(tmp_path, val_dataset="sintel-clean-trainval") + with pytest.raises(AssertionError): + dm.setup("fit") + + def test_fit_without_val_dataset_raises(self, tmp_path): + dm = self._make(tmp_path, train_dataset="chairs-train", train_batch_size=2) + with pytest.raises(AssertionError): + dm.setup("fit") + + def test_predict_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("predict") + + def test_predict_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, predict_dataset="sintel-clean-test") + dm.setup("predict") + # NB: the implementation has a typo - it writes to + # `parsed_predict_dataset_parsed` (not `predict_dataset_parsed`). We + # assert the actual field name so the test matches current behaviour. + assert dm.parsed_predict_dataset_parsed == [(1, "sintel", "clean", "test")] + + def test_test_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("test") + + def test_test_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, test_dataset="chairs") + dm.setup("test") + assert dm.test_dataset_parsed == [(1, "chairs")] + + def test_validate_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("validate") + + def test_validate_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, val_dataset="sintel-clean-trainval") + dm.setup("validate") + assert dm.val_dataset_parsed == [(1, "sintel", "clean", "trainval")] + + def test_unknown_stage_is_noop(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="chairs-train", + val_dataset="sintel-clean-trainval", + ) + # No assert / no parse should fire for an unrecognised stage. + dm.setup("some-other-stage") + assert dm.train_dataset_parsed is None + assert dm.val_dataset_parsed is None + + +# =========================================================================== +# Tests for _get_model_output_stride +# =========================================================================== +class TestGetModelOutputStride: + def test_returns_one_when_no_trainer(self): + dm = FlowDataModule() + assert dm._get_model_output_stride() == 1 + + def test_returns_one_when_trainer_none(self): + dm = FlowDataModule() + dm.trainer = None + assert dm._get_model_output_stride() == 1 + + def test_returns_model_output_stride(self): + dm = FlowDataModule() + # Use SimpleNamespace so `hasattr(model, "module")` is False (MagicMock + # would auto-create the `module` attribute and break the branch logic). + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + assert dm._get_model_output_stride() == 8 + + def test_returns_module_output_stride_when_ddp(self): + dm = FlowDataModule() + # When using DistributedDataParallel, trainer.model has a `module` + # attribute - the datamodule should unwrap it. + trainer = SimpleNamespace( + model=SimpleNamespace(module=SimpleNamespace(output_stride=16)) + ) + dm.trainer = trainer + assert dm._get_model_output_stride() == 16 + + +# =========================================================================== +# Tests for predict_dataloader +# =========================================================================== +class TestPredictDataloader: + def test_predict_dataloader_delegates_to_super(self): + dm = FlowDataModule() + # The implementation just calls super().predict_dataloader() which + # raises a MisconfigurationException (no predict loader is actually + # implemented). Document that behaviour. + from lightning.fabric.utilities.exceptions import ( + MisconfigurationException, + ) + + with pytest.raises(MisconfigurationException): + dm.predict_dataloader() + + +# =========================================================================== +# Tests for train_dataloader +# =========================================================================== +class TestTrainDataloader: + def test_returns_none_when_train_dataset_parsed_is_none(self): + dm = FlowDataModule() + dm.train_dataset_parsed = None + assert dm.train_dataloader() is None + + def test_concatenates_datasets_with_multiplier(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule( + train_batch_size=2, + train_num_workers=0, + train_transform_cuda=False, + dataset_config_path=cfg, + ) + # Two parsed datasets: multiplier=2 for 'chairs' and multiplier=3 + # for 'sintel' (only the dataset-name token, no extra params). + dm.train_dataset_parsed = [ + (2, "chairs", "train"), + (3, "sintel", "clean", "trainval"), + ] + + chairs_dummy = _make_dummy_dataset(length=5) + sintel_dummy = _make_dummy_dataset(length=7) + + def fake_get(name, is_train, *args): + return {"chairs": chairs_dummy, "sintel": sintel_dummy}[name] + + with patch.object( + FlowDataModule, "_get_chairs_dataset", autospec=True + ) as p_chairs, patch.object( + FlowDataModule, "_get_sintel_dataset", autospec=True + ) as p_sintel: + p_chairs.side_effect = lambda self, is_train, *a: chairs_dummy + p_sintel.side_effect = lambda self, is_train, *a: sintel_dummy + + loader = dm.train_dataloader() + + assert loader is not None + # The underlying dataset is the chained one (our dummy returns itself + # on `+`), so length is whatever the last `+` returned. + assert dm.train_dataloader_length == len(loader) + # Each `_get__dataset` should have been called exactly once. + assert p_chairs.call_count == 1 + assert p_sintel.call_count == 1 + # is_train should be True for the train dataloader. + chairs_args = p_chairs.call_args + assert chairs_args.args[0] is dm # self (autospec) + # When using autospec=True on an instance method the first positional + # arg is `self`; the second is `is_train`. + assert chairs_args.args[1] is True + assert chairs_args.args[2:] == ("train",) + sintel_args = p_sintel.call_args + assert sintel_args.args[1] is True + assert sintel_args.args[2:] == ("clean", "trainval") + + def test_dataloader_uses_correct_loader_kwargs(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule( + train_batch_size=4, + train_num_workers=0, + train_transform_cuda=False, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "chairs", "train")] + dummy = _make_dummy_dataset(length=3) + + with patch.object( + FlowDataModule, + "_get_chairs_dataset", + autospec=True, + return_value=dummy, + ): + loader = dm.train_dataloader() + assert loader.batch_size == 4 + assert loader.num_workers == 0 + # pin_memory is True when train_transform_cuda is False. + assert loader.pin_memory is True + # drop_last is False. + assert loader.drop_last is False + # persistent_workers is set to train_transform_cuda -> False. + assert loader.persistent_workers is False + + def test_cuda_disables_pin_memory_and_enables_persistent_workers(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule( + train_batch_size=2, + train_num_workers=1, + train_transform_cuda=True, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "chairs", "train")] + dummy = _make_dummy_dataset(length=2) + + with patch.object( + FlowDataModule, + "_get_chairs_dataset", + autospec=True, + return_value=dummy, + ): + loader = dm.train_dataloader() + assert loader.pin_memory is False + assert loader.persistent_workers is True + + +# =========================================================================== +# Tests for val_dataloader +# =========================================================================== +class TestValDataloader: + def test_returns_list_of_dataloaders_with_names_and_lengths(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule( + train_batch_size=2, + dataset_config_path=cfg, + train_transform_cuda=False, + ) + dm.val_dataset_parsed = [ + (1, "sintel", "clean", "trainval"), + (1, "kitti", "2012", "val"), + ] + sintel_dummy = _make_dummy_dataset(length=11) + kitti_dummy = _make_dummy_dataset(length=5) + + with patch.object( + FlowDataModule, + "_get_sintel_dataset", + autospec=True, + return_value=sintel_dummy, + ) as p_sintel, patch.object( + FlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=kitti_dummy, + ) as p_kitti: + loaders = dm.val_dataloader() + + assert len(loaders) == 2 + # Each loader has batch_size 1, num_workers 1, pin_memory False. + for ld in loaders: + assert ld.batch_size == 1 + assert ld.num_workers == 1 + assert ld.pin_memory is False + assert ld.drop_last is False + # Names are formed by "-".join(parsed_vals[1:]). + assert dm.val_dataloader_names == ["sintel-clean-trainval", "kitti-2012-val"] + # Lengths come from len(dataset) (for a MagicMock with __len__). + assert dm.val_dataloader_lengths == [11, 5] + # `_get__dataset` must be called with is_train=False. + assert p_sintel.call_args.args[1] is False + assert p_sintel.call_args.args[2:] == ("clean", "trainval") + assert p_kitti.call_args.args[1] is False + assert p_kitti.call_args.args[2:] == ("2012", "val") + + def test_resets_names_and_lengths_each_call(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg, train_transform_cuda=False) + # Pre-populate to ensure val_dataloader() overwrites these lists. + dm.val_dataloader_names = ["stale"] + dm.val_dataloader_lengths = [999] + dm.val_dataset_parsed = [(2, "chairs", "val")] + dummy = _make_dummy_dataset(length=4) + + with patch.object( + FlowDataModule, + "_get_chairs_dataset", + autospec=True, + return_value=dummy, + ): + dm.val_dataloader() + + assert dm.val_dataloader_names == ["chairs-val"] + assert dm.val_dataloader_lengths == [4] + + def test_persistent_workers_follows_train_transform_cuda(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg, train_transform_cuda=True) + dm.val_dataset_parsed = [(1, "chairs", "val")] + dummy = _make_dummy_dataset(length=2) + with patch.object( + FlowDataModule, + "_get_chairs_dataset", + autospec=True, + return_value=dummy, + ): + loaders = dm.val_dataloader() + assert loaders[0].persistent_workers is True + + +# =========================================================================== +# Tests for test_dataloader +# =========================================================================== +class TestTestDataloader: + def test_single_dataset_calls_get_with_test_suffix(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "chairs" + dummy = _make_dummy_dataset(length=6) + with patch.object( + FlowDataModule, + "_get_chairs_dataset", + autospec=True, + return_value=dummy, + ) as p_chairs: + loaders = dm.test_dataloader() + assert len(loaders) == 1 + assert p_chairs.call_count == 1 + # The datamodule appends "-test" then splits, so args are ("test",). + assert p_chairs.call_args.args[1] is False + assert p_chairs.call_args.args[2:] == ("test",) + # The name appended to test_dataloader_names is the suffixed id. + assert dm.test_dataloader_names == ["chairs-test"] + # Each loader uses batch_size 1, num_workers 1, no pin_memory. + assert loaders[0].batch_size == 1 + assert loaders[0].num_workers == 1 + assert loaders[0].pin_memory is False + assert loaders[0].drop_last is False + + def test_sintel_expands_to_clean_and_final(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "sintel" + dummy = _make_dummy_dataset(length=3) + with patch.object( + FlowDataModule, + "_get_sintel_dataset", + autospec=True, + return_value=dummy, + ) as p_sintel: + loaders = dm.test_dataloader() + assert len(loaders) == 2 + # Two calls: one for sintel-clean-test, one for sintel-final-test. + call_args = [c.args[2:] for c in p_sintel.call_args_list] + assert call_args == [("clean", "test"), ("final", "test")] + assert dm.test_dataloader_names == ["sintel-clean-test", "sintel-final-test"] + + def test_test_dataloader_names_accumulate_across_calls(self, tmp_path): + # The implementation appends to `test_dataloader_names` rather than + # resetting it. Document that behaviour. + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "chairs" + dm.test_dataloader_names = ["already-here"] + dummy = _make_dummy_dataset(length=2) + with patch.object( + FlowDataModule, + "_get_chairs_dataset", + autospec=True, + return_value=dummy, + ): + dm.test_dataloader() + assert dm.test_dataloader_names == ["already-here", "chairs-test"] + + +# =========================================================================== +# Tests for _get__dataset argument parsing +# --------------------------------------------------------------------------- +# The dataset constructors are mocked so we can verify the datamodule's +# argument-parsing logic in isolation. +# =========================================================================== +class TestGetDatasetArgumentParsing: + @pytest.fixture + def dm_cpu(self, tmp_path): + # Use a small crop size so we don't depend on default values. + cfg = _write_config(tmp_path, {}) + return FlowDataModule( + train_crop_size=(100, 200), + train_transform_cuda=False, + train_transform_fp16=False, + dataset_config_path=cfg, + ) + + def test_autoflow_default_split(self, dm_cpu): + with patch.object(mod, "AutoFlowDataset", return_value=MagicMock()) as p: + dm_cpu._get_autoflow_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "trainval" + assert p.call_args.args[0] == dm_cpu.autoflow_root_dir + + def test_autoflow_fbocc_flag(self, dm_cpu): + # The split-search loop only accepts "fbocc"; splits like "train"/"val" + # are rejected by the same loop (they fall through to ValueError), so + # autoflow's split is always the default "trainval". + with patch.object(mod, "AutoFlowDataset", return_value=MagicMock()) as p: + dm_cpu._get_autoflow_dataset(False, "fbocc") + assert p.call_args.kwargs["split"] == "trainval" + + def test_autoflow_split_arg_raises_value_error(self, dm_cpu): + # Documenting current (buggy) behaviour: passing a split token to + # _get_autoflow_dataset raises ValueError because the args-parser + # only knows about "fbocc". + with pytest.raises(ValueError): + dm_cpu._get_autoflow_dataset(False, "val") + + def test_autoflow_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_autoflow_dataset(False, "bogus") + + # ----- _get_chairs_dataset ----- + def test_chairs_default_split_trainval(self, dm_cpu): + with patch.object(mod, "FlyingChairsDataset", return_value=MagicMock()) as p: + dm_cpu._get_chairs_dataset(False) + assert p.call_args.kwargs["split"] == "trainval" + + def test_chairs_train_split(self, dm_cpu): + with patch.object(mod, "FlyingChairsDataset", return_value=MagicMock()) as p: + dm_cpu._get_chairs_dataset(False, "train") + assert p.call_args.kwargs["split"] == "train" + + def test_chairs_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_chairs_dataset(False, "weird") + + # ----- _get_chairs2_dataset ----- + def test_chairs2_flags_parsed(self, dm_cpu): + with patch.object(mod, "FlyingChairs2Dataset", return_value=MagicMock()) as p: + dm_cpu._get_chairs2_dataset(False, "train", "rev", "occ", "mb", "back") + kwargs = p.call_args.kwargs + assert kwargs["split"] == "train" + assert kwargs["add_reverse"] is True + assert kwargs["get_occlusion_mask"] is True + assert kwargs["get_motion_boundary_mask"] is True + assert kwargs["get_backward"] is True + + def test_chairs2_default_split_trainval(self, dm_cpu): + with patch.object(mod, "FlyingChairs2Dataset", return_value=MagicMock()) as p: + dm_cpu._get_chairs2_dataset(False) + assert p.call_args.kwargs["split"] == "trainval" + + def test_chairs2_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_chairs2_dataset(False, "nope") + + # ----- _get_hd1k_dataset ----- + def test_hd1k_seqlen_and_seqpos(self, dm_cpu): + with patch.object(mod, "Hd1kDataset", return_value=MagicMock()) as p: + dm_cpu._get_hd1k_dataset(False, "seqlen_5", "seqpos_last") + kwargs = p.call_args.kwargs + assert kwargs["sequence_length"] == 5 + assert kwargs["sequence_position"] == "last" + + def test_hd1k_split_passes_through(self, dm_cpu): + with patch.object(mod, "Hd1kDataset", return_value=MagicMock()) as p: + dm_cpu._get_hd1k_dataset(False, "test") + assert p.call_args.kwargs["split"] == "test" + + def test_hd1k_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_hd1k_dataset(False, "nonsense") + + # ----- _get_kitti_dataset ----- + def test_kitti_default_versions_both(self, dm_cpu): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False) + assert p.call_args.kwargs["versions"] == ["2012", "2015"] + + def test_kitti_selects_single_version(self, dm_cpu): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False, "2015") + assert p.call_args.kwargs["versions"] == ["2015"] + + def test_kitti_occ_flag(self, dm_cpu): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False, "occ") + assert p.call_args.kwargs["get_occlusion_mask"] is True + + def test_kitti_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_kitti_dataset(False, "what") + + # ----- _get_sintel_dataset ----- + def test_sintel_pass_selection(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "final") + assert p.call_args.kwargs["pass_names"] == ["final"] + + def test_sintel_default_both_passes(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False) + assert p.call_args.kwargs["pass_names"] == ["clean", "final"] + + def test_sintel_seqlen_and_seqpos(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "seqlen_4", "seqpos_middle") + assert p.call_args.kwargs["sequence_length"] == 4 + assert p.call_args.kwargs["sequence_position"] == "middle" + + def test_sintel_test_split(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "clean", "test") + assert p.call_args.kwargs["split"] == "test" + + def test_sintel_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_sintel_dataset(False, "unknown") + + # ----- _get_spring_dataset ----- + def test_spring_defaults(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "train" + assert kwargs["add_reverse"] is False + assert kwargs["get_backward"] is False + assert kwargs["sequence_length"] == 2 + assert kwargs["sequence_position"] == "first" + assert kwargs["reverse_only"] is False + assert kwargs["subsample"] is True + # side_names default to left+right when none provided. + assert kwargs["side_names"] == ["left", "right"] + + def test_spring_flags(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset( + False, + "val", + "timerev", + "back", + "seqlen_3", + "seqpos_middle", + "gt4k", + "left", + "right", + "robust", + ) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "val" + assert kwargs["add_reverse"] is True + assert kwargs["get_backward"] is True + assert kwargs["sequence_length"] == 3 + assert kwargs["sequence_position"] == "middle" + assert kwargs["subsample"] is False + assert kwargs["side_names"] == ["left", "right"] + assert kwargs["robust_mode"] is True + + def test_spring_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_spring_dataset(False, "garbage") + + # ----- _get_tartanair_dataset ----- + def test_tartanair_default_difficulty_easy(self, dm_cpu): + with patch.object(mod, "TartanAirDataset", return_value=MagicMock()) as p: + dm_cpu._get_tartanair_dataset(False) + assert p.call_args.kwargs["difficulties"] == ["easy"] + + def test_tartanair_difficulties_collected(self, dm_cpu): + with patch.object(mod, "TartanAirDataset", return_value=MagicMock()) as p: + dm_cpu._get_tartanair_dataset(False, "easy", "hard") + assert p.call_args.kwargs["difficulties"] == ["easy", "hard"] + + def test_tartanair_seqlen_seqpos_occ(self, dm_cpu): + with patch.object(mod, "TartanAirDataset", return_value=MagicMock()) as p: + dm_cpu._get_tartanair_dataset(False, "occ", "seqlen_3", "seqpos_last") + kwargs = p.call_args.kwargs + assert kwargs["get_occlusion_mask"] is True + assert kwargs["sequence_length"] == 3 + assert kwargs["sequence_position"] == "last" + + def test_tartanair_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_tartanair_dataset(False, "impossible") + + # ----- _get_things_dataset ----- + def test_things_default_pass_names(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False) + assert p.call_args.kwargs["pass_names"] == ["clean", "final"] + # is_subset defaults to False -> FlyingThings3DDataset, not Subset. + assert p.call_args.args[0] == dm_cpu.flying_things3d_root_dir + + def test_things_subset_routes_to_subset_dataset(self, dm_cpu): + with patch.object( + mod, "FlyingThings3DSubsetDataset", return_value=MagicMock() + ) as p: + dm_cpu._get_things_dataset(False, "subset") + assert p.call_args.args[0] == dm_cpu.flying_things3d_subset_root_dir + + def test_things_flags(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset( + False, + "clean", + "train", + "rev", + "occ", + "mb", + "back", + "seqlen_3", + "seqpos_last", + "sinteltransform", + ) + kwargs = p.call_args.kwargs + assert kwargs["pass_names"] == ["clean"] + assert kwargs["split"] == "train" + assert kwargs["add_reverse"] is True + assert kwargs["get_occlusion_mask"] is True + assert kwargs["get_motion_boundary_mask"] is True + assert kwargs["get_backward"] is True + assert kwargs["sequence_length"] == 3 + assert kwargs["sequence_position"] == "last" + + def test_things_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_things_dataset(False, "nope") + + # ----- _get_kubric_dataset ----- + def test_kubric_train_raises_not_implemented(self, dm_cpu): + with pytest.raises(NotImplementedError): + dm_cpu._get_kubric_dataset(True) + + def test_kubric_args(self, dm_cpu): + with patch.object(mod, "KubricDataset", return_value=MagicMock()) as p: + dm_cpu._get_kubric_dataset( + False, "back", "seqlen_4", "seqpos_middle", "maxseq_10" + ) + kwargs = p.call_args.kwargs + assert kwargs["get_backward"] is True + assert kwargs["sequence_length"] == 4 + assert kwargs["sequence_position"] == "middle" + assert kwargs["max_seq"] == 10 + + # ----- _get_middlebury_st_dataset ----- + def test_middlebury_st_asserts_not_train(self, dm_cpu): + with pytest.raises(AssertionError): + dm_cpu._get_middlebury_st_dataset(True) + + def test_middlebury_st_construction(self, dm_cpu): + with patch.object(mod, "MiddleburySTDataset", return_value=MagicMock()) as p: + dm_cpu._get_middlebury_st_dataset(False) + assert p.call_args.args[0] == dm_cpu.middlebury_st_root_dir + + # ----- _get_viper_dataset ----- + def test_viper_asserts_not_train(self, dm_cpu): + with pytest.raises(AssertionError): + dm_cpu._get_viper_dataset(True) + + def test_viper_uses_val_split(self, dm_cpu): + with patch.object(mod, "ViperDataset", return_value=MagicMock()) as p: + dm_cpu._get_viper_dataset(False) + assert p.call_args.kwargs["split"] == "val" + + # ----- _get_sintel_finetune_dataset ----- + def test_sintel_finetune_eval_raises_not_implemented(self, dm_cpu): + with pytest.raises(NotImplementedError): + dm_cpu._get_sintel_finetune_dataset(False) + + def test_sintel_finetune_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_sintel_finetune_dataset(True, "unknown") + + def test_sintel_finetune_constructs_expected_datasets(self, dm_cpu): + # Track how many times each sliced dataset is multiplied by counting + # `__iadd__` calls on the mocks that act as the multiplied chains. + sintel_clean = _make_dummy_dataset(length=2) + sintel_final = _make_dummy_dataset(length=2) + kitti = _make_dummy_dataset(length=3) + hd1k = _make_dummy_dataset(length=4) + things = _make_dummy_dataset(length=10) + + with patch.object( + mod, "FlyingThings3DDataset", return_value=things + ), patch.object( + mod, + "SintelDataset", + side_effect=[sintel_clean, sintel_final], + ) as p_sintel, patch.object( + mod, "KittiDataset", return_value=kitti + ), patch.object( + mod, "Hd1kDataset", return_value=hd1k + ): + out = dm_cpu._get_sintel_finetune_dataset(True, "searaft_split") + + # SintelDataset must be called twice: once for clean, once for final. + assert p_sintel.call_count == 2 + passes = [c.kwargs["pass_names"] for c in p_sintel.call_args_list] + assert passes == [["clean"], ["final"]] + # The searaft_split multipliers perform 19 iadd repetitions for + # each Sintel pass, 79 for KITTI and 29 for HD1K. + assert sintel_clean.__iadd__.call_count == 19 + assert sintel_final.__iadd__.call_count == 19 + assert kitti.__iadd__.call_count == 79 + assert hd1k.__iadd__.call_count == 29 + # The returned dataset is the chained combination (things + ...). + # Because our mocks' `+` returns the left operand, the result equals + # the things dataset mock. + assert out is things + + # ----- _get_overfit_dataset ----- + def test_overfit_default_uses_sintel(self, dm_cpu): + sintel_dummy = MagicMock(name="sintel_overfit") + sintel_dummy.img_paths = ["a", "b"] + sintel_dummy.flow_paths = ["a"] + sintel_dummy.occ_paths = ["a"] + sintel_dummy.mb_paths = ["a"] + sintel_dummy.flow_b_paths = ["a"] + sintel_dummy.occ_b_paths = ["a"] + sintel_dummy.mb_b_paths = ["a"] + sintel_dummy.metadata = ["meta_a", "meta_b"] + with patch.object(mod, "SintelDataset", return_value=sintel_dummy) as p: + out = dm_cpu._get_overfit_dataset(True) + # Should request split=trainval and pass_names="clean". + assert p.call_args.kwargs["split"] == "trainval" + assert p.call_args.kwargs["pass_names"] == "clean" + # All path lists are truncated to length 1. + assert out.img_paths == ["a"] + assert out.flow_paths == ["a"] + assert out.occ_paths == ["a"] + assert out.mb_paths == ["a"] + assert out.flow_b_paths == ["a"] + assert out.occ_b_paths == ["a"] + assert out.mb_b_paths == ["a"] + assert out.metadata == ["meta_a"] + + def test_overfit_selects_chairs2(self, dm_cpu): + chairs2_dummy = MagicMock(name="chairs2_overfit") + chairs2_dummy.img_paths = ["x"] + chairs2_dummy.flow_paths = ["x"] + chairs2_dummy.occ_paths = ["x"] + chairs2_dummy.mb_paths = ["x"] + chairs2_dummy.flow_b_paths = ["x"] + chairs2_dummy.occ_b_paths = ["x"] + chairs2_dummy.mb_b_paths = ["x"] + chairs2_dummy.metadata = ["m"] + with patch.object(mod, "FlyingChairs2Dataset", return_value=chairs2_dummy) as p: + out = dm_cpu._get_overfit_dataset(True, "chairs2") + assert p.call_args.kwargs["split"] == "trainval" + assert p.call_args.kwargs["get_occlusion_mask"] is True + assert p.call_args.kwargs["get_motion_boundary_mask"] is True + assert p.call_args.kwargs["get_backward"] is True + + def test_overfit_default_crop_when_none(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg) # train_crop_size is None + sintel_dummy = MagicMock() + sintel_dummy.img_paths = [] + sintel_dummy.flow_paths = [] + sintel_dummy.occ_paths = [] + sintel_dummy.mb_paths = [] + sintel_dummy.flow_b_paths = [] + sintel_dummy.occ_b_paths = [] + sintel_dummy.mb_b_paths = [] + sintel_dummy.metadata = [] + with patch.object(mod, "SintelDataset", return_value=sintel_dummy): + dm._get_overfit_dataset(True) + # Setting crop size to None causes the datamodule to assign a default. + assert dm.train_crop_size is not None + assert len(dm.train_crop_size) == 2 + + +# =========================================================================== +# Tests for transform construction in _get__dataset (train mode) +# =========================================================================== +class TestTrainTransformConstruction: + def test_train_mode_builds_compose_transform(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule( + train_crop_size=(100, 200), + train_transform_cuda=False, + train_transform_fp16=False, + dataset_config_path=cfg, + ) + with patch.object(mod, "FlyingChairsDataset", return_value=MagicMock()) as p: + dm._get_chairs_dataset(True, "train") + # The transform passed to the dataset must be a Compose instance. + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import optical_flow_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_train_mode_sets_default_crop_size_when_none(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = FlowDataModule(dataset_config_path=cfg) + assert dm.train_crop_size is None + with patch.object(mod, "FlyingChairsDataset", return_value=MagicMock()): + dm._get_chairs_dataset(True, "train") + # Was auto-filled. + assert dm.train_crop_size is not None + cy, cx = dm.train_crop_size + # Default for chairs is (368, 496). With output_stride=1, + # make_divisible(x, 1) = x. + assert cy == 368 + assert cx == 496 diff --git a/tests/common/data/test_scene_flow_datamodule.py b/tests/common/data/test_scene_flow_datamodule.py new file mode 100644 index 0000000..843ac75 --- /dev/null +++ b/tests/common/data/test_scene_flow_datamodule.py @@ -0,0 +1,1204 @@ +"""Unit tests for `roco_spring_devkit.common.data.scene_flow_datamodule`. + +The datasets and transforms themselves have their own dedicated test modules +(`test_scene_flow_datasets.py`, `test_scene_flow_transforms.py`), so the +focus here is the *datamodule-specific* logic: + +* the default values set by ``__init__``; +* the string-parsing helper ``_parse_dataset_selection``; +* the YAML path loader ``_load_dataset_paths`` (and the fact that it does NOT + overwrite paths that were already provided, and silently skips YAML keys + that do not map to a ``*_root_dir`` attribute via the ``hasattr`` guard); +* the ``setup`` stage dispatch and assertions; +* ``_get_model_output_stride`` with / without a trainer; +* the dataloader builders ``train_dataloader`` / ``val_dataloader`` / + ``test_dataloader`` / ``predict_dataloader`` (including the Sintel + test-split expansion and the dataset multiplier behaviour); +* the argument-parsing logic of every ``_get__dataset`` helper + (``kitti``, ``sintel``, ``spring``, ``things``) and the ``ValueError`` + raised on unknown flags; +* that ``get_flow`` / ``get_disparity`` / ``get_intrinsics`` are always + passed as ``True`` for the scene-flow use case — the key difference + compared to the stereo datamodule which disables them. + +The on-disk dataset constructors are mocked out so that we only exercise the +datamodule plumbing, not the (already tested) dataset readers. +""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import yaml + +import roco_spring_devkit.common.data.scene_flow_datamodule as mod +from roco_spring_devkit.common.data.scene_flow_datamodule import ( + SceneFlowDataModule, +) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_dummy_dataset(length: int = 4): + """Return a MagicMock that behaves like a concatenatable Dataset. + + `__add__`/`__iadd__` return the same mock so that train_dataloader can + keep using `+=` (`dataset = dataset + dataset`) without growing new + mock objects. `__len__` returns ``length``. + """ + + dataset = MagicMock(name="DummyDataset") + dataset.__len__ = MagicMock(return_value=length) + dataset.__add__ = MagicMock(return_value=dataset) + dataset.__iadd__ = MagicMock(return_value=dataset) + return dataset + + +def _write_config(tmp_path: Path, paths: dict) -> str: + """Write a YAML dataset config and return its path as a string.""" + cfg_path = tmp_path / "datasets.yaml" + with open(cfg_path, "w") as f: + yaml.safe_dump(paths, f) + return str(cfg_path) + + +# =========================================================================== +# Tests for __init__ defaults +# =========================================================================== +class TestInit: + def test_defaults_are_none_or_zero(self): + dm = SceneFlowDataModule() + assert dm.predict_dataset is None + assert dm.test_dataset is None + assert dm.train_dataset is None + assert dm.val_dataset is None + assert dm.train_batch_size is None + assert dm.train_num_workers == 4 + assert dm.train_crop_size is None + assert dm.train_transform_cuda is False + assert dm.train_transform_fp16 is False + + def test_root_dirs_default_to_none(self): + dm = SceneFlowDataModule() + for attr in ( + "flying_things3d_root_dir", + "mpi_sintel_root_dir", + "kitti_2015_root_dir", + "spring_root_dir", + "robust_spring_root_dir", + ): + assert getattr(dm, attr) is None, attr + + def test_no_extra_root_dir_attrs_exist(self): + # The scene flow datamodule only supports the 5 root dirs above. It + # does NOT carry the optical-flow / stereo-only ones (autoflow, + # flying_chairs[2], flying_things3d_subset, kitti_2012, tartanair, + # kubric, middlebury_st, viper). Document that. + dm = SceneFlowDataModule() + for missing in ( + "flying_things3d_subset_root_dir", + "kitti_2012_root_dir", + "tartanair_root_dir", + "middlebury_st_root_dir", + "autoflow_root_dir", + "flying_chairs_root_dir", + "flying_chairs2_root_dir", + "kubric_root_dir", + "viper_root_dir", + ): + assert not hasattr(dm, missing), missing + + def test_default_dataset_config_path(self): + dm = SceneFlowDataModule() + # The default config path is hard-coded for the CLI's working dir. + assert dm.dataset_config_path == "../../datasets.yaml" + + def test_init_stores_constructor_values(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=3, + train_num_workers=2, + train_crop_size=(100, 200), + train_transform_cuda=True, + train_transform_fp16=True, + dataset_config_path=cfg, + ) + assert dm.train_dataset == "kitti-train" + assert dm.val_dataset == "sintel-clean-trainval" + assert dm.train_batch_size == 3 + assert dm.train_num_workers == 2 + assert dm.train_crop_size == (100, 200) + assert dm.train_transform_cuda is True + assert dm.train_transform_fp16 is True + assert dm.dataset_config_path == cfg + + def test_init_initialises_internal_state(self): + dm = SceneFlowDataModule() + assert dm.predict_dataset_parsed is None + assert dm.test_dataset_parsed is None + assert dm.train_dataset_parsed is None + assert dm.val_dataset_parsed is None + assert dm.train_dataloader_length == 0 + assert dm.train_epoch_step == 0 + assert dm.val_dataloader_names == [] + assert dm.val_dataloader_lengths == [] + assert dm.test_dataloader_names == [] + + +# =========================================================================== +# Tests for _parse_dataset_selection +# =========================================================================== +class TestParseDatasetSelection: + def test_none_returns_empty_list(self): + dm = SceneFlowDataModule() + assert dm._parse_dataset_selection(None) == [] + + def test_single_dataset_default_multiplier(self): + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection("kitti-train") + assert out == [(1, "kitti", "train")] + + def test_multiple_datasets_separated_by_plus(self): + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection( + "kitti-train+sintel-clean-trainval+things-clean" + ) + assert out == [ + (1, "kitti", "train"), + (1, "sintel", "clean", "trainval"), + (1, "things", "clean"), + ] + + def test_leading_multiplier(self): + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection("3*sintel-clean-trainval") + assert out == [(3, "sintel", "clean", "trainval")] + + def test_trailing_multiplier(self): + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection("kitti-train*5") + assert out == [(5, "kitti", "train")] + + def test_mixed_multipliers_and_flags(self): + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection( + "kitti-train+3*sintel-clean-trainval+things-train*5" + ) + assert out == [ + (1, "kitti", "train"), + (3, "sintel", "clean", "trainval"), + (5, "things", "train"), + ] + + def test_spaces_are_stripped(self): + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection(" kitti - train + 2 * sintel - clean ") + assert out == [ + (1, "kitti", "train"), + (2, "sintel", "clean"), + ] + + def test_invalid_double_multiplier_raises(self): + dm = SceneFlowDataModule() + with pytest.raises(ValueError): + dm._parse_dataset_selection("3*4*sintel-clean") + + def test_non_integer_multiplier_raises_value_error(self): + # 'x' is not int -> the trailing-multiplier branch raises ValueError + # when trying int('x'). + dm = SceneFlowDataModule() + with pytest.raises(ValueError): + dm._parse_dataset_selection("sintel*x") + + def test_empty_string_yields_single_empty_dataset(self): + # An empty string splits into [''] which is treated as a single + # dataset with empty name and multiplier 1. This documents current + # behaviour (no guard for empty input). + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection("") + assert out == [(1, "")] + + def test_docstring_example(self): + # The example in the docstring. + dm = SceneFlowDataModule() + out = dm._parse_dataset_selection( + "chairs-train+3*sintel-clean-trainval+kitti-2012-train*5" + ) + assert out == [ + (1, "chairs", "train"), + (3, "sintel", "clean", "trainval"), + (5, "kitti", "2012", "train"), + ] + + +# =========================================================================== +# Tests for _load_dataset_paths +# =========================================================================== +class TestLoadDatasetPaths: + def test_fills_none_attrs_from_yaml(self, tmp_path): + cfg = _write_config( + tmp_path, + { + "flying_things3d": "/things/path", + "mpi_sintel": "/sintel/path", + "kitti_2015": "/kitti/path", + "spring": "/spring/path", + "robust_spring": "/robust/path", + }, + ) + dm = SceneFlowDataModule(dataset_config_path=cfg) + dm._load_dataset_paths() + assert dm.flying_things3d_root_dir == "/things/path" + assert dm.mpi_sintel_root_dir == "/sintel/path" + assert dm.kitti_2015_root_dir == "/kitti/path" + assert dm.spring_root_dir == "/spring/path" + assert dm.robust_spring_root_dir == "/robust/path" + + def test_does_not_overwrite_existing_paths(self, tmp_path): + cfg = _write_config( + tmp_path, + {"flying_things3d": "/yaml/things", "mpi_sintel": "/yaml/sintel"}, + ) + dm = SceneFlowDataModule( + flying_things3d_root_dir="/explicit/things", + dataset_config_path=cfg, + ) + dm._load_dataset_paths() + # The explicitly-provided path wins. + assert dm.flying_things3d_root_dir == "/explicit/things" + # The unset one is filled from YAML. + assert dm.mpi_sintel_root_dir == "/yaml/sintel" + + def test_missing_keys_in_yaml_leave_attrs_as_none(self, tmp_path): + cfg = _write_config(tmp_path, {"flying_things3d": "/things/path"}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + dm._load_dataset_paths() + assert dm.flying_things3d_root_dir == "/things/path" + # Keys absent from YAML are left untouched (still None). + assert dm.mpi_sintel_root_dir is None + assert dm.kitti_2015_root_dir is None + + def test_unmatched_yaml_keys_are_skipped_via_hasattr_guard(self, tmp_path): + # The scene flow datamodule does NOT have e.g. 'autoflow_root_dir', + # 'flying_chairs_root_dir', 'flying_things3d_subset_root_dir', + # 'kitti_2012_root_dir', 'tartanair_root_dir', 'kubric_root_dir', + # 'middlebury_st_root_dir' or 'viper_root_dir' attributes - unlike + # the optical flow module. `_load_dataset_paths` uses + # `hasattr(self, "_root_dir")` to silently skip such YAML + # keys instead of raising AttributeError. + cfg = _write_config( + tmp_path, + { + "autoflow": "/aut/oflow", # no autoflow_root_dir attr + "flying_chairs": "/chairs/path", # no flying_chairs_root_dir attr + "flying_things3d_subset": "/sub/path", # no ..._subset_root_dir attr + "kitti_2012": "/kitti12/path", # no kitti_2012_root_dir attr + "tartanair": "/tart/path", # no tartanair_root_dir attr + "kubric": "/kubric/path", # no kubric_root_dir attr + "middlebury_st": "/mst/path", # no middlebury_st_root_dir attr + "viper": "/viper/path", # no viper_root_dir attr + "mpi_sintel": "/sintel/path", # has mpi_sintel_root_dir attr + }, + ) + dm = SceneFlowDataModule(dataset_config_path=cfg) + # Should not raise. + dm._load_dataset_paths() + # Matched key is filled. + assert dm.mpi_sintel_root_dir == "/sintel/path" + # Unmatched keys do not create new attributes. + for missing in ( + "autoflow_root_dir", + "flying_chairs_root_dir", + "flying_things3d_subset_root_dir", + "kitti_2012_root_dir", + "tartanair_root_dir", + "kubric_root_dir", + "middlebury_st_root_dir", + "viper_root_dir", + ): + assert not hasattr(dm, missing), missing + + +# =========================================================================== +# Tests for setup +# =========================================================================== +class TestSetup: + def _make(self, tmp_path, **kwargs): + cfg = _write_config(tmp_path, {}) + kwargs.setdefault("dataset_config_path", cfg) + return SceneFlowDataModule(**kwargs) + + def test_fit_parses_train_and_val(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=4, + ) + dm.setup("fit") + assert dm.train_dataset_parsed == [(1, "kitti", "train")] + assert dm.val_dataset_parsed == [(1, "sintel", "clean", "trainval")] + + def test_fit_assigns_default_batch_size_when_none(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=None, + ) + dm.setup("fit") + assert dm.train_batch_size == 8 + + def test_fit_loads_dataset_paths(self, tmp_path): + cfg = _write_config(tmp_path, {"mpi_sintel": "/sintel/path"}) + dm = SceneFlowDataModule( + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=4, + dataset_config_path=cfg, + ) + dm.setup("fit") + assert dm.mpi_sintel_root_dir == "/sintel/path" + + def test_fit_without_train_dataset_raises(self, tmp_path): + dm = self._make(tmp_path, val_dataset="sintel-clean-trainval") + with pytest.raises(AssertionError): + dm.setup("fit") + + def test_fit_without_val_dataset_raises(self, tmp_path): + dm = self._make(tmp_path, train_dataset="kitti-train", train_batch_size=2) + with pytest.raises(AssertionError): + dm.setup("fit") + + def test_predict_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("predict") + + def test_predict_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, predict_dataset="sintel-clean-test") + dm.setup("predict") + # NB: the implementation has a typo - it writes to + # `parsed_predict_dataset_parsed` (not `predict_dataset_parsed`). We + # assert the actual field name so the test matches current behaviour. + assert dm.parsed_predict_dataset_parsed == [(1, "sintel", "clean", "test")] + + def test_test_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("test") + + def test_test_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, test_dataset="kitti") + dm.setup("test") + assert dm.test_dataset_parsed == [(1, "kitti")] + + def test_validate_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("validate") + + def test_validate_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, val_dataset="sintel-clean-trainval") + dm.setup("validate") + assert dm.val_dataset_parsed == [(1, "sintel", "clean", "trainval")] + + def test_unknown_stage_is_noop(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + ) + # No assert / no parse should fire for an unrecognised stage. + dm.setup("some-other-stage") + assert dm.train_dataset_parsed is None + assert dm.val_dataset_parsed is None + + +# =========================================================================== +# Tests for _get_model_output_stride +# =========================================================================== +class TestGetModelOutputStride: + def test_returns_one_when_no_trainer(self): + dm = SceneFlowDataModule() + # 'trainer' is not set on the datamodule until Lightning assigns one. + assert dm._get_model_output_stride() == 1 + + def test_returns_one_when_trainer_none(self): + dm = SceneFlowDataModule() + dm.trainer = None + assert dm._get_model_output_stride() == 1 + + def test_returns_model_output_stride(self): + dm = SceneFlowDataModule() + # Use SimpleNamespace so `hasattr(model, "module")` is False (MagicMock + # would auto-create the `module` attribute and break the branch logic). + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + assert dm._get_model_output_stride() == 8 + + def test_returns_module_output_stride_when_ddp(self): + dm = SceneFlowDataModule() + # When using DistributedDataParallel, trainer.model has a `module` + # attribute - the datamodule should unwrap it. + trainer = SimpleNamespace( + model=SimpleNamespace(module=SimpleNamespace(output_stride=16)) + ) + dm.trainer = trainer + assert dm._get_model_output_stride() == 16 + + +# =========================================================================== +# Tests for predict_dataloader +# =========================================================================== +class TestPredictDataloader: + def test_predict_dataloader_delegates_to_super(self): + dm = SceneFlowDataModule() + # The implementation just calls super().predict_dataloader() which + # raises a MisconfigurationException (no predict loader is actually + # implemented). Document that behaviour. + from lightning.fabric.utilities.exceptions import ( + MisconfigurationException, + ) + + with pytest.raises(MisconfigurationException): + dm.predict_dataloader() + + +# =========================================================================== +# Tests for train_dataloader +# =========================================================================== +class TestTrainDataloader: + def test_returns_none_when_train_dataset_parsed_is_none(self): + dm = SceneFlowDataModule() + dm.train_dataset_parsed = None + assert dm.train_dataloader() is None + + def test_concatenates_datasets_with_multiplier(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_batch_size=2, + train_num_workers=0, + train_transform_cuda=False, + dataset_config_path=cfg, + ) + # Two parsed datasets: multiplier=2 for 'kitti' and multiplier=3 for + # 'sintel'. + dm.train_dataset_parsed = [ + (2, "kitti", "train"), + (3, "sintel", "clean", "trainval"), + ] + kitti_dummy = _make_dummy_dataset(length=5) + sintel_dummy = _make_dummy_dataset(length=7) + + with patch.object( + SceneFlowDataModule, "_get_kitti_dataset", autospec=True + ) as p_kitti, patch.object( + SceneFlowDataModule, "_get_sintel_dataset", autospec=True + ) as p_sintel: + p_kitti.side_effect = lambda self, is_train, *a: kitti_dummy + p_sintel.side_effect = lambda self, is_train, *a: sintel_dummy + + loader = dm.train_dataloader() + + assert loader is not None + assert dm.train_dataloader_length == len(loader) + # Each `_get__dataset` should have been called exactly once. + assert p_kitti.call_count == 1 + assert p_sintel.call_count == 1 + # is_train should be True for the train dataloader. + kitti_args = p_kitti.call_args + assert kitti_args.args[0] is dm # self (autospec) + assert kitti_args.args[1] is True + assert kitti_args.args[2:] == ("train",) + sintel_args = p_sintel.call_args + assert sintel_args.args[1] is True + assert sintel_args.args[2:] == ("clean", "trainval") + + def test_dataloader_uses_correct_loader_kwargs(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_batch_size=4, + train_num_workers=0, + train_transform_cuda=False, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "kitti", "train")] + dummy = _make_dummy_dataset(length=3) + + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + loader = dm.train_dataloader() + assert loader.batch_size == 4 + assert loader.num_workers == 0 + # pin_memory is True when train_transform_cuda is False. + assert loader.pin_memory is True + # drop_last is False. + assert loader.drop_last is False + # persistent_workers is set to train_transform_cuda -> False. + assert loader.persistent_workers is False + + def test_cuda_disables_pin_memory_and_enables_persistent_workers(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_batch_size=2, + train_num_workers=1, + train_transform_cuda=True, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "kitti", "train")] + dummy = _make_dummy_dataset(length=2) + + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + loader = dm.train_dataloader() + assert loader.pin_memory is False + assert loader.persistent_workers is True + + def test_single_dataset_multiplier_one(self, tmp_path): + # multiplier==1 means the `for _ in range(multiplier - 1)` loop body + # never runs, so `dataset_mult` stays equal to `dataset`. + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_batch_size=2, + train_num_workers=0, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "kitti", "train")] + dummy = _make_dummy_dataset(length=2) + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ) as p_kitti: + dm.train_dataloader() + # Only one call to the getter. + assert p_kitti.call_count == 1 + # No `+`/`+=` was performed on the dataset (multiplier == 1). + dummy.__add__.assert_not_called() + + +# =========================================================================== +# Tests for val_dataloader +# =========================================================================== +class TestValDataloader: + def test_returns_list_of_dataloaders_with_names_and_lengths(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_batch_size=2, + dataset_config_path=cfg, + train_transform_cuda=False, + ) + dm.val_dataset_parsed = [ + (1, "sintel", "clean", "trainval"), + (1, "kitti", "val"), + ] + sintel_dummy = _make_dummy_dataset(length=11) + kitti_dummy = _make_dummy_dataset(length=5) + + with patch.object( + SceneFlowDataModule, + "_get_sintel_dataset", + autospec=True, + return_value=sintel_dummy, + ) as p_sintel, patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=kitti_dummy, + ) as p_kitti: + loaders = dm.val_dataloader() + + assert len(loaders) == 2 + # Each loader has batch_size 1, num_workers 1, pin_memory False. + for ld in loaders: + assert ld.batch_size == 1 + assert ld.num_workers == 1 + assert ld.pin_memory is False + assert ld.drop_last is False + # Names are formed by "-".join(parsed_vals[1:]). + assert dm.val_dataloader_names == ["sintel-clean-trainval", "kitti-val"] + # Lengths come from len(dataset) (for a MagicMock with __len__). + assert dm.val_dataloader_lengths == [11, 5] + # `_get__dataset` must be called with is_train=False. + assert p_sintel.call_args.args[1] is False + assert p_sintel.call_args.args[2:] == ("clean", "trainval") + assert p_kitti.call_args.args[1] is False + assert p_kitti.call_args.args[2:] == ("val",) + + def test_resets_names_and_lengths_each_call(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg, train_transform_cuda=False) + # Pre-populate to ensure val_dataloader() overwrites these lists. + dm.val_dataloader_names = ["stale"] + dm.val_dataloader_lengths = [999] + dm.val_dataset_parsed = [(2, "kitti", "val")] + dummy = _make_dummy_dataset(length=4) + + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + dm.val_dataloader() + + assert dm.val_dataloader_names == ["kitti-val"] + assert dm.val_dataloader_lengths == [4] + + def test_persistent_workers_follows_train_transform_cuda(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg, train_transform_cuda=True) + dm.val_dataset_parsed = [(1, "kitti", "val")] + dummy = _make_dummy_dataset(length=2) + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + loaders = dm.val_dataloader() + assert loaders[0].persistent_workers is True + + +# =========================================================================== +# Tests for test_dataloader +# =========================================================================== +class TestTestDataloader: + def test_single_dataset_calls_get_with_test_suffix(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "kitti" + dummy = _make_dummy_dataset(length=6) + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ) as p_kitti: + loaders = dm.test_dataloader() + assert len(loaders) == 1 + assert p_kitti.call_count == 1 + # The datamodule appends "-test" then splits, so args are ("test",). + assert p_kitti.call_args.args[1] is False + assert p_kitti.call_args.args[2:] == ("test",) + # The name appended to test_dataloader_names is the suffixed id. + assert dm.test_dataloader_names == ["kitti-test"] + # Each loader uses batch_size 1, num_workers 1, no pin_memory. + assert loaders[0].batch_size == 1 + assert loaders[0].num_workers == 1 + assert loaders[0].pin_memory is False + assert loaders[0].drop_last is False + + def test_spring_test_split_works(self, tmp_path): + # Unlike the things dataset, spring accepts the "test" split token + # in `_get_spring_dataset`. The test dataloader can therefore expand + # it without raising. + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "spring" + dummy = _make_dummy_dataset(length=3) + with patch.object( + SceneFlowDataModule, + "_get_spring_dataset", + autospec=True, + return_value=dummy, + ) as p_spring: + loaders = dm.test_dataloader() + assert len(loaders) == 1 + assert p_spring.call_args.args[1] is False + assert p_spring.call_args.args[2:] == ("test",) + assert dm.test_dataloader_names == ["spring-test"] + + def test_sintel_expands_to_clean_and_final(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "sintel" + dummy = _make_dummy_dataset(length=3) + with patch.object( + SceneFlowDataModule, + "_get_sintel_dataset", + autospec=True, + return_value=dummy, + ) as p_sintel: + loaders = dm.test_dataloader() + assert len(loaders) == 2 + # Two calls: one for sintel-clean-test, one for sintel-final-test. + call_args = [c.args[2:] for c in p_sintel.call_args_list] + assert call_args == [("clean", "test"), ("final", "test")] + assert dm.test_dataloader_names == ["sintel-clean-test", "sintel-final-test"] + + def test_test_dataloader_names_accumulate_across_calls(self, tmp_path): + # The implementation appends to `test_dataloader_names` rather than + # resetting it. Document that behaviour. + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + dm.test_dataset = "kitti" + dm.test_dataloader_names = ["already-here"] + dummy = _make_dummy_dataset(length=2) + with patch.object( + SceneFlowDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + dm.test_dataloader() + assert dm.test_dataloader_names == ["already-here", "kitti-test"] + + +# =========================================================================== +# Tests for _get__dataset argument parsing +# --------------------------------------------------------------------------- +# The dataset constructors are mocked so we can verify the datamodule's +# argument-parsing logic in isolation. +# =========================================================================== +class TestGetDatasetArgumentParsing: + @pytest.fixture + def dm_cpu(self, tmp_path): + # Use a small crop size so we don't depend on default values. + cfg = _write_config(tmp_path, {}) + return SceneFlowDataModule( + train_crop_size=(100, 200), + train_transform_cuda=False, + train_transform_fp16=False, + dataset_config_path=cfg, + ) + + # ----- _get_kitti_dataset ----- + def test_kitti_default_split(self, dm_cpu): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "trainval" + # Scene flow kitti always uses the 2015 root dir. + assert p.call_args.args[0] == dm_cpu.kitti_2015_root_dir + # All scene-flow signals requested. + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True + + def test_kitti_passes_version_and_split_token(self, dm_cpu): + # '2015' is a no-op pass-through token (the scene flow KittiDataset + # only supports the 2015 version). + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False, "2015", "train") + assert p.call_args.kwargs["split"] == "train" + + def test_kitti_test_split(self, dm_cpu): + # The scene flow kitti arg parser DOES accept the "test" split token + # (unlike the things parser which rejects it). + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False, "test") + assert p.call_args.kwargs["split"] == "test" + + def test_kitti_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_kitti_dataset(False, "what") + + # ----- _get_sintel_dataset ----- + def test_sintel_default_both_passes(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False) + assert p.call_args.kwargs["pass_names"] == ["clean", "final"] + assert p.call_args.args[0] == dm_cpu.mpi_sintel_root_dir + # All scene-flow signals requested. + assert p.call_args.kwargs["get_flow"] is True + assert p.call_args.kwargs["get_disparity"] is True + assert p.call_args.kwargs["get_intrinsics"] is True + + def test_sintel_pass_selection(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "final") + assert p.call_args.kwargs["pass_names"] == ["final"] + + def test_sintel_train_split(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "clean", "train") + assert p.call_args.kwargs["split"] == "train" + assert p.call_args.kwargs["pass_names"] == ["clean"] + + def test_sintel_test_split(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "clean", "test") + assert p.call_args.kwargs["split"] == "test" + + def test_sintel_no_sequence_args_supported(self, dm_cpu): + # The scene flow sintel parser does NOT understand seqlen_/seqpos_ + # tokens (unlike the optical flow one). Document that an unknown + # token raises. + with pytest.raises(ValueError): + dm_cpu._get_sintel_dataset(False, "seqlen_4") + + def test_sintel_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_sintel_dataset(False, "unknown") + + # ----- _get_spring_dataset ----- + def test_spring_defaults(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "train" + assert kwargs["subsample"] is True + assert kwargs["add_camera_reverse"] is False + assert kwargs["camera_reverse_only"] is False + assert kwargs["add_time_reverse"] is False + assert kwargs["time_reverse_only"] is False + assert kwargs["robust_mode"] is False + # All scene-flow signals requested. + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True + # Robust root dir is forwarded. + assert kwargs["robust_root_dir"] == dm_cpu.robust_spring_root_dir + + def test_spring_flags(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset( + False, + "val", + "camrev", + "camrevonly", + "timerev", + "timerevonly", + "gt4k", + "robust", + ) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "val" + assert kwargs["add_camera_reverse"] is True + assert kwargs["camera_reverse_only"] is True + assert kwargs["add_time_reverse"] is True + assert kwargs["time_reverse_only"] is True + assert kwargs["subsample"] is False + assert kwargs["robust_mode"] is True + + def test_spring_test_split(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset(False, "test") + assert p.call_args.kwargs["split"] == "test" + + def test_spring_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_spring_dataset(False, "garbage") + + # ----- _get_things_dataset ----- + def test_things_default_pass_names(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False) + assert p.call_args.kwargs["pass_names"] == ["clean", "final"] + assert p.call_args.args[0] == dm_cpu.flying_things3d_root_dir + # All scene-flow signals requested. + assert p.call_args.kwargs["get_flow"] is True + assert p.call_args.kwargs["get_disparity"] is True + assert p.call_args.kwargs["get_intrinsics"] is True + + def test_things_pass_selection(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False, "clean") + assert p.call_args.kwargs["pass_names"] == ["clean"] + + def test_things_train_split(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False, "train") + assert p.call_args.kwargs["split"] == "train" + + def test_things_val_split(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False, "val") + assert p.call_args.kwargs["split"] == "val" + + def test_things_test_split_raises_value_error(self, dm_cpu): + # Documenting current (intentional or not) behaviour: the things arg + # parser only accepts splits in ["train", "val", "trainval"], so a + # "test" token (as the test_dataloader would append) raises. + with pytest.raises(ValueError): + dm_cpu._get_things_dataset(False, "test") + + def test_things_sinteltransform_flag(self, dm_cpu): + # The sinteltransform flag only changes the train-time scale range; + # when is_train=False the flag should be accepted without error. + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()): + # Should not raise. + dm_cpu._get_things_dataset(False, "sinteltransform") + + def test_things_no_subset_flag_supported(self, dm_cpu): + # The scene flow things parser does NOT understand a "subset" token + # (unlike the optical flow one which routes to + # FlyingThings3DSubsetDataset). Document that subset raises. + with pytest.raises(ValueError): + dm_cpu._get_things_dataset(False, "subset") + + def test_things_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_things_dataset(False, "nope") + + +# =========================================================================== +# Tests for transform construction in _get__dataset (train mode) +# =========================================================================== +class TestTrainTransformConstruction: + def _dm(self, tmp_path): + cfg = _write_config(tmp_path, {}) + return SceneFlowDataModule( + train_crop_size=(100, 200), + train_transform_cuda=False, + train_transform_fp16=False, + dataset_config_path=cfg, + ) + + def test_kitti_train_builds_compose_transform(self, tmp_path): + dm = self._dm(tmp_path) + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import scene_flow_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_sintel_train_builds_compose_transform(self, tmp_path): + dm = self._dm(tmp_path) + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm._get_sintel_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import scene_flow_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_spring_train_builds_compose_transform(self, tmp_path): + dm = self._dm(tmp_path) + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm._get_spring_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import scene_flow_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_things_train_builds_compose_transform(self, tmp_path): + dm = self._dm(tmp_path) + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm._get_things_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import scene_flow_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_eval_mode_uses_plain_totensor(self, tmp_path): + # In eval mode, the transform is a bare ToTensor (not a Compose). + from roco_spring_devkit.common.data import scene_flow_transforms as ft + + dm = self._dm(tmp_path) + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(False) + assert isinstance(p.call_args.kwargs["transform"], ft.ToTensor) + + def test_things_sinteltransform_changes_scale_range_in_train_mode(self, tmp_path): + # With sinteltransform=True the train-time major_scale is (-0.2, 0.6); + # without it it is (-0.4, 0.8). We assert the values reach the + # RandomScaleAndCrop constructor. + from roco_spring_devkit.common.data import scene_flow_transforms as ft + + captured = [] + real_rsac = ft.RandomScaleAndCrop + + def spy_rsac(crop_size, major_scale=None, space_scale=None, sparse=False): + obj = real_rsac( + crop_size, + major_scale=major_scale, + space_scale=space_scale, + sparse=sparse, + ) + captured.append(major_scale) + return obj + + dm = self._dm(tmp_path) + with patch.object( + mod, "FlyingThings3DDataset", return_value=MagicMock() + ), patch.object(ft, "RandomScaleAndCrop", spy_rsac): + dm._get_things_dataset(True, "sinteltransform") + assert captured == [(-0.2, 0.6)] + + captured.clear() + with patch.object( + mod, "FlyingThings3DDataset", return_value=MagicMock() + ), patch.object(ft, "RandomScaleAndCrop", spy_rsac): + dm._get_things_dataset(True) + assert captured == [(-0.4, 0.8)] + + def test_train_mode_sets_default_crop_size_when_none(self, tmp_path): + # For kitti the default crop size is (288, 960). + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + assert dm.train_crop_size is None + with patch.object(mod, "KittiDataset", return_value=MagicMock()): + dm._get_kitti_dataset(True, "train") + # Was auto-filled. + assert dm.train_crop_size is not None + cy, cx = dm.train_crop_size + # With output_stride=1, make_divisible(x, 1) = x. + assert cy == 288 + assert cx == 960 + + def test_train_default_crop_sintel(self, tmp_path): + # For sintel the default crop size is (368, 768). + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + with patch.object(mod, "SintelDataset", return_value=MagicMock()): + dm._get_sintel_dataset(True, "train") + assert dm.train_crop_size == (368, 768) + + def test_train_default_crop_spring(self, tmp_path): + # For spring the default crop size is (540, 960). + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + with patch.object(mod, "SpringDataset", return_value=MagicMock()): + dm._get_spring_dataset(True, "train") + assert dm.train_crop_size == (540, 960) + + def test_train_default_crop_things(self, tmp_path): + # For things the default crop size is (400, 720). + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule(dataset_config_path=cfg) + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()): + dm._get_things_dataset(True, "train") + assert dm.train_crop_size == (400, 720) + + def test_train_crop_size_made_divisible_by_output_stride(self, tmp_path): + # When an output stride > 1 is set, the crop dims passed to the + # transform are rounded down to a multiple of the stride. The + # datamodule does NOT write the rounded values back to + # ``self.train_crop_size`` (only the `None` branch does), so we + # inspect the values that reach ``RandomScaleAndCrop`` instead. + from roco_spring_devkit.common.data import scene_flow_transforms as ft + from types import SimpleNamespace + + captured = [] + real_rsac = ft.RandomScaleAndCrop + + def spy_rsac(crop_size, major_scale=None, space_scale=None, sparse=False): + captured.append(crop_size) + return real_rsac( + crop_size, + major_scale=major_scale, + space_scale=space_scale, + sparse=sparse, + ) + + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_crop_size=(290, 961), + dataset_config_path=cfg, + ) + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + with patch.object(mod, "KittiDataset", return_value=MagicMock()), patch.object( + ft, "RandomScaleAndCrop", spy_rsac + ): + dm._get_kitti_dataset(True, "train") + # make_divisible(290, 8) = 288, make_divisible(961, 8) = 960. + assert captured == [(288, 960)] + # And the original train_crop_size attribute is untouched. + assert dm.train_crop_size == (290, 961) + + def test_train_crop_size_clamped_to_at_least_stride(self, tmp_path): + # make_divisible returns max(div, v - v % div); a value smaller than + # the stride is clamped up to the stride. + from roco_spring_devkit.common.data import scene_flow_transforms as ft + from types import SimpleNamespace + + captured = [] + real_rsac = ft.RandomScaleAndCrop + + def spy_rsac(crop_size, major_scale=None, space_scale=None, sparse=False): + captured.append(crop_size) + return real_rsac( + crop_size, + major_scale=major_scale, + space_scale=space_scale, + sparse=sparse, + ) + + cfg = _write_config(tmp_path, {}) + dm = SceneFlowDataModule( + train_crop_size=(3, 5), + dataset_config_path=cfg, + ) + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + with patch.object(mod, "KittiDataset", return_value=MagicMock()), patch.object( + ft, "RandomScaleAndCrop", spy_rsac + ): + dm._get_kitti_dataset(True, "train") + assert captured == [(8, 8)] + + +# =========================================================================== +# Tests for scene-flow flags (always get_flow / get_disparity / get_intrinsics) +# =========================================================================== +class TestSceneFlowFlags: + """Verify the scene flow datamodule always asks the dataset constructors + for ``get_flow=True``, ``get_disparity=True`` and ``get_intrinsics=True`` + (this is what makes the produced samples scene-flow-specific rather than + stereo-specific, where the stereo datamodule sets them to False).""" + + @pytest.fixture + def dm(self, tmp_path): + cfg = _write_config(tmp_path, {}) + return SceneFlowDataModule( + train_crop_size=(100, 200), + dataset_config_path=cfg, + ) + + def test_kitti_enables_flow_disparity_intrinsics(self, dm): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True + + def test_sintel_enables_flow_disparity_intrinsics(self, dm): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm._get_sintel_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True + + def test_spring_enables_flow_disparity_intrinsics(self, dm): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm._get_spring_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True + + def test_things_enables_flow_disparity_intrinsics(self, dm): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm._get_things_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True + + def test_flags_true_in_eval_mode_too(self, dm): + # The flags are passed regardless of is_train. + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is True + assert kwargs["get_disparity"] is True + assert kwargs["get_intrinsics"] is True diff --git a/tests/common/data/test_scene_flow_datasets.py b/tests/common/data/test_scene_flow_datasets.py index e29fb08..0730efc 100644 --- a/tests/common/data/test_scene_flow_datasets.py +++ b/tests/common/data/test_scene_flow_datasets.py @@ -359,17 +359,23 @@ def test_flows_and_disparities_content(self, tmp_path): out["disparities"][0][..., 0], np.array([[10.0, 50.0], [30.0, 40.0]], dtype=np.float32), ) + # Second disparity is backward warped into frame 1 coordinates np.testing.assert_allclose( out["disparities"][1][..., 0], - np.array([[20.0, 60.0], [70.0, 15.0]], dtype=np.float32), + np.array([[15.0, 15.0], [15.0, 15.0]], dtype=np.float32), ) # All values are well below max_flow / max_disparity, so the masks # are fully 255. np.testing.assert_array_equal( out["valid_flows"][0], np.full((2, 2, 1), 255, dtype=np.uint8) ) - for v in out["valid_disparities"]: - np.testing.assert_array_equal(v, np.full((2, 2, 1), 255, dtype=np.uint8)) + # First disparity is valid, but second is warped out of bounds, so invalid + np.testing.assert_array_equal( + out["valid_disparities"][0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + np.testing.assert_array_equal( + out["valid_disparities"][1], np.full((2, 2, 1), 0, dtype=np.uint8) + ) def test_intrinsics_and_baselines_content(self, tmp_path): d = self._make_simple_dataset(tmp_path) @@ -512,9 +518,7 @@ def test_backward_flow_path_alignment(self, tmp_path): tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_past" / "left" ) # Sample 0 (frames [0, 1]) -> flow_b = into_past/fp_1. - assert d.flow_b_paths[0] == [flow_past / "0000001.pfm"] - # Sample 1 (frames [1, 2]) -> flow_b = into_past/fp_2. - assert d.flow_b_paths[1] == [flow_past / "0000002.pfm"] + assert d.flow_b_paths[0] == [flow_past / "0000002.pfm"] def test_right_flow_paths(self, tmp_path): self._make(tmp_path, n_frames=3) @@ -544,7 +548,7 @@ def test_right_flow_paths(self, tmp_path): tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_past" / "right" ) assert d.flow_r_paths[0] == [flow_r_fut / "0000000.pfm"] - assert d.flow_b_r_paths[0] == [flow_r_past / "0000001.pfm"] + assert d.flow_b_r_paths[0] == [flow_r_past / "0000002.pfm"] def test_disparity_paths_per_image(self, tmp_path): # Each sample holds one disparity file per image (so 2 entries for @@ -792,6 +796,11 @@ def _make(self, tmp_path, seq_names, n_frames=3, with_flow=True, with_disp=True) fd.mkdir(parents=True) for f in range(n_frames - 1): _touch(fd / f"frame_{f:04d}.flo") + + od = tmp_path / "training" / "occlusions" / seq + od.mkdir(parents=True) + for f in range(n_frames - 1): + _touch(od / f"frame_{f:04d}.png") if with_disp: dd = tmp_path / "training" / "disparities" / seq dd.mkdir(parents=True) @@ -1154,8 +1163,8 @@ def test_subsample_getitem_content(self, tmp_path): ) def test_val_split_is_val_flag(self, tmp_path): - # 00270 (the only entry in Spring_val.txt). - self._make_seq(tmp_path, "0027", n_frames=2) + # 0022 (the only entry in Spring_val.txt). + self._make_seq(tmp_path, "0022", n_frames=2) with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): d = ds.SpringDataset( str(tmp_path), @@ -1168,8 +1177,8 @@ def test_val_split_is_val_flag(self, tmp_path): get_intrinsics=False, sequence_length=2, ) - # '0027' is the only sequence in Spring_val.txt and should be kept. + # '0022' is the only sequence in Spring_val.txt and should be kept. assert len(d.img_paths) == 1 - assert d.metadata[0]["misc"] == "0027" + assert d.metadata[0]["misc"] == "0022" # The val-split sample must be flagged as is_val=True. assert d.metadata[0]["is_val"] is True diff --git a/tests/common/data/test_scene_flow_transforms.py b/tests/common/data/test_scene_flow_transforms.py index 41e1873..c42aa22 100644 --- a/tests/common/data/test_scene_flow_transforms.py +++ b/tests/common/data/test_scene_flow_transforms.py @@ -28,9 +28,7 @@ GaussianNoise, RandomFlip, RandomPatchEraser, - RandomRotate, RandomScaleAndCrop, - RandomTranslate, Resize, ToTensor, _adjust_intrinsics_for_crop, @@ -38,7 +36,6 @@ _adjust_intrinsics_for_scale, _get_valid_keys, _resize, - _update_oob_disparities, ) @@ -56,45 +53,6 @@ def test_get_valid_keys(self): } assert _get_valid_keys(list(keys), None, None) == list(keys) - def test_update_oob_disparities_1d(self): - # 1-channel positive stereo disparity: correspondence is x_r = x_l - d. - occs = torch.zeros(1, 1, 3, 3) - disparities = torch.zeros(1, 1, 3, 3) - disparities[0, 0, 0, 0] = 1.0 # x_r = 0 - 1 = -1 -> OOB left - disparities[0, 0, 0, 2] = -2.0 # x_r = 2 + 2 = 4 >= W=3 -> OOB right - - out_occs = _update_oob_disparities(occs, disparities) - assert out_occs[0, 0, 0, 0] == 1.0 - assert out_occs[0, 0, 0, 2] == 1.0 - assert out_occs[0, 0, 1, 1] == 0.0 # d=0 stays in bounds - - def test_update_oob_disparities_exact_boundary(self): - # Landing exactly one pixel past the edge (x_r == W or x_r == -1) is OOB, - # while landing exactly on the last pixel (x_r == W-1 or 0) is not. - occs = torch.zeros(1, 1, 3, 3) - disparities = torch.zeros(1, 1, 3, 3) - disparities[0, 0, 0, 2] = -1.0 # x_r = 2 + 1 = 3 == W -> OOB - disparities[0, 0, 1, 0] = 1.0 # x_r = 0 - 1 = -1 -> OOB - disparities[0, 0, 2, 2] = 0.0 # x_r = 2 -> last column, in bounds - - out_occs = _update_oob_disparities(occs, disparities) - assert out_occs[0, 0, 0, 2] == 1.0 - assert out_occs[0, 0, 1, 0] == 1.0 - assert out_occs[0, 0, 2, 2] == 0.0 - - def test_update_oob_disparities_2d(self): - # 2-channel disparity treated as a forward correspondence vector. - occs = torch.zeros(1, 1, 3, 3) - disparities = torch.zeros(1, 2, 3, 3) - disparities[0, :, 0, 0] = torch.tensor([-1.0, -1.0]) # (-1,-1) OOB top-left - disparities[0, 0, 0, 2] = 1.0 # x = 2 + 1 = 3 == W -> OOB right - # (1,1) keeps d=(0,0) -> in bounds - - out_occs = _update_oob_disparities(occs, disparities) - assert out_occs[0, 0, 0, 0] == 1.0 - assert out_occs[0, 0, 0, 2] == 1.0 - assert out_occs[0, 0, 1, 1] == 0.0 - def test_resize_sparse_logic(self): # Valid pixels at (x=1,y=0) and (x=2,y=1); target (3,6) gives # scale (x2.0, x1.5): expected landings (2,0) and (4, round(1.5)=2). @@ -105,8 +63,9 @@ def test_resize_sparse_logic(self): out = _resize( inputs, target_size=(3, 6), - binary_keys=["valid_disparities"], - disparities_keys=["disparities"], + binary_keys=["valid_disparities", "valid_flows"], + disparity_keys=["disparities"], + flow_keys=["flows"], sparse=True, valid_key="valid_disparities", ) @@ -132,8 +91,9 @@ def test_resize_sparse_float_valids(self): out = _resize( inputs, target_size=(3, 6), - binary_keys=["valid_disparities"], - disparities_keys=["disparities"], + binary_keys=["valid_disparities", "valid_flows"], + disparity_keys=["disparities"], + flow_keys=["flows"], sparse=True, valid_key="valid_disparities", ) @@ -164,9 +124,9 @@ def test_adjust_intrinsics_for_scale(self): def test_adjust_intrinsics_for_flip(self): intr = torch.tensor([[[10.0, 0.0, 5.0], [0.0, 20.0, 6.0], [0.0, 0.0, 1.0]]]) - out_h = _adjust_intrinsics_for_flip(intr, is_hflip=True, width=10, height=10) + out_h = _adjust_intrinsics_for_flip(intr, is_hflip=True, img_w=10, img_h=10) assert out_h[0, 0, 2] == 4.0 # (10 - 1) - 5 - out_v = _adjust_intrinsics_for_flip(intr, is_hflip=False, width=10, height=10) + out_v = _adjust_intrinsics_for_flip(intr, is_hflip=False, img_w=10, img_h=10) assert out_v[0, 1, 2] == 3.0 # (10 - 1) - 6 @@ -177,11 +137,11 @@ def test_compose_filters_none_and_applies_all(self): img = np.zeros((4, 4, 3), dtype=np.uint8) valid = np.ones((4, 4, 1), dtype=np.float32) - inputs = {"images": img, "valids": valid} + inputs = {"images": img, "valid_flows": valid} out = t(inputs) assert out["images"].shape == (1, 3, 2, 2) - assert out["valids"].shape == (1, 1, 2, 2) + assert out["valid_flows"].shape == (1, 1, 2, 2) class TestToTensor: @@ -257,19 +217,19 @@ def test_intrinsics_and_baselines_conversion(self): class TestCenterCrop: def test_exact_crop_values(self): valid = torch.arange(16).reshape(1, 1, 4, 4).float() - inputs = {"valids": valid.clone(), "images": valid.clone()} + inputs = {"valid_flows": valid.clone(), "images": valid.clone()} out = CenterCrop(crop_size=(2, 2), ignore_keys=[])(inputs) expected = torch.tensor([[[[5.0, 6.0], [9.0, 10.0]]]]) - assert out["valids"].shape == (1, 1, 2, 2) - assert torch.allclose(out["valids"], expected) + assert out["valid_flows"].shape == (1, 1, 2, 2) + assert torch.allclose(out["valid_flows"], expected) assert torch.allclose(out["images"], expected) def test_default_ignore_keys(self): # Default construction must also work (ignore_keys=None). valid = torch.arange(16).reshape(1, 1, 4, 4).float() - inputs = {"valids": valid.clone(), "images": valid.clone()} + inputs = {"valid_flows": valid.clone(), "images": valid.clone()} out = CenterCrop(crop_size=(2, 2))(inputs) expected = torch.tensor([[[[5.0, 6.0], [9.0, 10.0]]]]) @@ -283,12 +243,12 @@ def test_synchronization_across_keys(self): "images_right": pos.clone(), "flows": pos.repeat(1, 2, 1, 1), "disparities": pos.clone(), - "valids": pos.clone(), + "valid_flows": pos.clone(), } out = CenterCrop(crop_size=(2, 2), ignore_keys=[])(inputs) expected = torch.tensor([[[[5.0, 6.0], [9.0, 10.0]]]]) - for k in ("images", "images_right", "disparities", "valids"): + for k in ("images", "images_right", "disparities", "valid_flows"): assert out[k].shape == (1, 1, 2, 2) assert torch.allclose(out[k], expected), f"{k} out of sync" assert out["flows"].shape == (1, 2, 2, 2) @@ -298,14 +258,16 @@ def test_occlusion_update_after_crop(self): # After the crop, disparities pointing out of the cropped frame must be # marked as occluded. d=1 at x'=0 -> x_r = -1 -> OOB; at x'=1 -> x_r = 0 -> OK. inputs = { - "valids": torch.ones(1, 1, 4, 4), + "valid_flows": torch.ones(1, 1, 4, 4), "disparities": torch.ones(1, 1, 4, 4), + "flows": torch.ones(1, 2, 4, 4), "occs": torch.zeros(1, 1, 4, 4), } out = CenterCrop(crop_size=(2, 2), ignore_keys=[])(inputs) + assert out["occs"].shape == (1, 1, 2, 2) - assert torch.all(out["occs"][0, 0, :, 0] == 1.0) - assert torch.all(out["occs"][0, 0, :, 1] == 0.0) + assert torch.all(out["occs"][0, 0, :, -1] == 1.0) + assert torch.all(out["occs"][0, 0, 0, 0] == 0.0) class TestResize: @@ -359,7 +321,8 @@ def test_random_scale_and_crop_logic(self, mock_uniform, mock_randint): mock_randint.side_effect = [0, 0] # y_crop, x_crop disp = torch.ones(1, 1, 4, 4) - inputs = {"disparities": disp} + flow = torch.ones(1, 2, 4, 4) + inputs = {"disparities": disp, "flows": flow} out = RandomScaleAndCrop( crop_size=(6, 6), major_scale=(1.0, 1.0), space_scale=(1.0, 1.0) @@ -418,6 +381,7 @@ def test_crop_offset_updates_intrinsics(self, mock_uniform, mock_randint): inputs = { "disparities": torch.ones(1, 1, 4, 4), + "flows": torch.ones(1, 2, 4, 4), "intrinsics": torch.tensor( [[[1.0, 0.0, 2.0], [0.0, 2.0, 3.0], [0.0, 0.0, 1.0]]] ), @@ -444,8 +408,13 @@ def test_sparse_scale_and_crop(self, mock_uniform, mock_randint): valids[0, 0, 0, 0] = 1.0 # scales to (0, 0) valids[0, 0, 1, 1] = 1.0 # scales to (4, 4) disp = torch.ones(1, 1, 2, 2) + flow = torch.ones(1, 2, 2, 2) - inputs = {"valids": valids.clone(), "disparities": disp.clone()} + inputs = { + "valid_flows": valids.clone(), + "disparities": disp.clone(), + "flows": flow.clone(), + } out = RandomScaleAndCrop( crop_size=(8, 8), major_scale=(1.0, 1.0), @@ -454,9 +423,9 @@ def test_sparse_scale_and_crop(self, mock_uniform, mock_randint): )(inputs) # Both valid pixels must survive, including the one landing at (0, 0). - assert out["valids"].sum() == 2.0 - assert out["valids"][0, 0, 0, 0] == 1.0 - assert out["valids"][0, 0, 4, 4] == 1.0 + assert out["valid_flows"].sum() == 2.0 + assert out["valid_flows"][0, 0, 0, 0] == 1.0 + assert out["valid_flows"][0, 0, 4, 4] == 1.0 # Disparity follows its valid pixels and is scaled by 4. assert out["disparities"][0, 0, 0, 0] == 4.0 assert out["disparities"][0, 0, 4, 4] == 4.0 @@ -472,141 +441,17 @@ def test_occlusion_update_after_scaling(self, mock_uniform, mock_randint): # With x_r = x - d and d = 1, the leftmost column (x = 0) points to # x_r = -1 and goes out of bounds; every other column stays in bounds. disp = torch.ones(1, 1, 4, 4) + flow = torch.ones(1, 2, 4, 4) occs = torch.zeros(1, 1, 4, 4) - inputs = {"disparities": disp, "occs": occs} + inputs = {"disparities": disp, "flows": flow, "occs": occs} out = RandomScaleAndCrop( crop_size=(4, 4), major_scale=(0.0, 0.0), space_scale=(0.0, 0.0) )(inputs) assert torch.allclose(out["disparities"], torch.tensor(1.0)) - assert torch.all(out["occs"][0, 0, :, 0] == 1.0) - assert torch.all(out["occs"][0, 0, :, 1:] == 0.0) - - def test_time_scale_not_implemented(self): - inputs = {"disparities": torch.ones(1, 1, 4, 4)} - try: - RandomScaleAndCrop(crop_size=(2, 2), time_scale=(0.1, 0.2))(inputs) - raise AssertionError("expected NotImplementedError") - except NotImplementedError: - pass - - -class TestRandomTranslate: - @patch("random.randint") - def test_translate_values_and_synchronization(self, mock_randint): - mock_randint.side_effect = [1, 2] # tw = 1, th = 2 - - img = torch.zeros(2, 1, 4, 4) - img[0, 0, 2, 2] = 1.0 - img[1, 0, 1, 1] = 2.0 - flows = torch.zeros(1, 2, 4, 4) - # 2-channel disparities so the transform does not crash; values must - # still remain untouched (same-time stereo pairs share the crop parity). - disp = torch.ones(2, 2, 4, 4) * 5.0 - - inputs = {"images": img, "flows": flows, "disparities": disp} - out = RandomTranslate(translation=(2, 2))(inputs) - - # All inputs cropped to H=4-2=2, W=4-1=3, staying synchronized. - assert out["images"].shape == (2, 1, 2, 3) - assert out["flows"].shape == (1, 2, 2, 3) - assert out["disparities"].shape == (2, 2, 2, 3) - - # Even indices cropped with offset (+th, +tw); odd with (-th, -tw). - assert out["images"][0, 0, 0, 1] == 1.0 # (2,2) -> (0,1) - assert out["images"][1, 0, 1, 1] == 2.0 # (1,1) -> (1,1) - - # Flows between opposite-parity frames gain the relative translation: - # even flows get (+tw, +th), odd flows get (-tw, -th). - assert torch.allclose(out["flows"][0, 0], torch.tensor(1.0)) - assert torch.allclose(out["flows"][0, 1], torch.tensor(2.0)) - - # Disparities link same-time (same-parity) pairs: values do not change. - assert torch.all(out["disparities"] == 5.0) - - @patch("random.randint") - def test_translate_with_1channel_disparities(self, mock_randint): - mock_randint.side_effect = [1, 2] # tw = 1, th = 2 - - # This is the actual data layout produced by the scene flow datasets: - # 1-channel positive disparities, one per time step. - inputs = { - "images": torch.zeros(2, 1, 4, 4), - "disparities": torch.ones(2, 1, 4, 4) * 5.0, - } - out = RandomTranslate(translation=(2, 2))(inputs) - - assert out["disparities"].shape == (2, 1, 2, 3) - assert torch.all(out["disparities"] == 5.0) - - @patch("random.randint") - def test_zero_translation(self, mock_randint): - mock_randint.side_effect = [0, 0] - - img = torch.ones(1, 1, 4, 4) - inputs = {"images": img, "disparities": torch.zeros(1, 1, 4, 4)} - - out = RandomTranslate(translation=10)(inputs) - # Early return check - assert out["images"] is img - assert out["images"].shape == (1, 1, 4, 4) - - -class TestRandomRotate: - @patch("random.uniform") - def test_disparity_vector_rotated_once(self, mock_uniform): - mock_uniform.side_effect = [90.0, 0.0] - - disp = torch.zeros(2, 2, 3, 3) - disp[:, 0, :, :] = 1.0 # vectors pointing +x - - inputs = {"disparities": disp.clone(), "occs": torch.zeros(2, 1, 3, 3)} - out = RandomRotate(angle=90.0)(inputs) - - # A +x vector rotated by 90 degrees must become (0, -1) exactly once. - for i in range(2): - center = out["disparities"][i, :, 1, 1] - assert torch.allclose(center[0], torch.tensor(0.0), atol=1e-5) - assert torch.allclose(center[1], torch.tensor(-1.0), atol=1e-5) - - @patch("random.uniform") - def test_flows_synchronized_with_rotation(self, mock_uniform): - mock_uniform.side_effect = [90.0, 0.0] - - disp = torch.zeros(2, 2, 3, 3) - flows = torch.zeros(2, 2, 3, 3) - flows[:, 0, :, :] = 1.0 # vectors pointing +x - - inputs = {"disparities": disp.clone(), "flows": flows.clone()} - out = RandomRotate(angle=90.0)(inputs) - - # Flow vectors must be rotated exactly like the images, otherwise the - # flow falls out of sync with the rotated frames. - center_flow = out["flows"][0, :, 1, 1] - assert torch.allclose(center_flow[0], torch.tensor(0.0), atol=1e-5) - assert torch.allclose(center_flow[1], torch.tensor(-1.0), atol=1e-5) - - @patch("random.uniform") - def test_rotate_nearest_neighbor_sparse(self, mock_uniform): - mock_uniform.side_effect = [90.0, 0.0] - - valid = torch.zeros(1, 1, 3, 3) - valid[0, 0, 0, 1] = 1.0 - img = torch.zeros(1, 1, 3, 3) - img[0, 0, 0, 1] = 1.0 - - inputs = { - "valids": valid, - "images": img, - "disparities": torch.zeros(1, 2, 3, 3), - } - out = RandomRotate(angle=90.0, sparse=True)(inputs) - - # 90 degrees counterclockwise: (0,1) -> (1,0). Mask and image must agree. - assert out["valids"][0, 0, 1, 0] == 1.0 - assert out["valids"].sum() == 1.0 - assert torch.allclose(out["images"][0, 0, 1, 0], torch.tensor(1.0), atol=1e-5) + assert torch.all(out["occs"][0, 0, :, -1] == 1.0) + assert torch.all(out["occs"][0, 0, :-1, :-1] == 0.0) class TestRandomFlip: @@ -616,6 +461,8 @@ def test_horizontal_flip_images_flows_and_intrinsics(self): flows = torch.ones(1, 2, 2, 2) flows[:, 1] = 2.0 valid_flows = torch.ones(1, 1, 2, 2) + disparities = torch.ones(2, 1, 2, 2) + valid_disparities = torch.ones(2, 1, 2, 2) intr = torch.tensor([[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]]) baselines = torch.tensor([0.25]) inputs = { @@ -623,16 +470,26 @@ def test_horizontal_flip_images_flows_and_intrinsics(self): "images_right": img_right.clone(), "flows": flows.clone(), "valid_flows": valid_flows.clone(), + "flows_right": flows.clone(), + "valid_flows_right": valid_flows.clone(), + "disparities": disparities.clone(), + "valid_disparities": valid_disparities.clone(), + "disparities_right": disparities.clone(), + "valid_disparities_right": valid_disparities.clone(), "intrinsics": intr.clone(), "baselines": baselines.clone(), } out = RandomFlip(hflip_prob=1.0, vflip_prob=0.0)(inputs) - # Images are mirrored along the width. - assert torch.allclose(out["images"], torch.tensor([[[[2.0, 1.0], [4.0, 3.0]]]])) + # Images are mirrored along the width and swapped left/right. + # assert torch.allclose(out["images"], torch.tensor([[[[2.0, 1.0], [4.0, 3.0]]]])) + # assert torch.allclose( + # out["images_right"], torch.tensor([[[[6.0, 5.0], [8.0, 7.0]]]]) + # ) + assert torch.allclose(out["images"], torch.tensor([[[[6.0, 5.0], [8.0, 7.0]]]])) assert torch.allclose( - out["images_right"], torch.tensor([[[[6.0, 5.0], [8.0, 7.0]]]]) + out["images_right"], torch.tensor([[[[2.0, 1.0], [4.0, 3.0]]]]) ) # Flow x-component is negated; y-component is not. assert torch.allclose(out["flows"][:, 0], torch.tensor(-1.0)) @@ -646,9 +503,6 @@ def test_horizontal_flip_images_flows_and_intrinsics(self): assert torch.allclose(out["baselines"], baselines) def test_horizontal_flip_disparity_consistency(self): - # After a pure horizontal flip (no left/right swap), the correspondence - # x_r = x_l - d becomes x_r' = x_l' + d, so the disparity must be negated - # to stay synchronized with the flipped images. img_left = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) img_right = torch.tensor([[[[5.0, 6.0], [7.0, 8.0]]]]) disp = torch.ones(1, 1, 2, 2) @@ -659,7 +513,7 @@ def test_horizontal_flip_disparity_consistency(self): } out = RandomFlip(hflip_prob=1.0, vflip_prob=0.0)(inputs) - assert torch.allclose(out["disparities"], torch.tensor(-1.0)) + assert torch.allclose(out["disparities"], torch.tensor(1.0)) def test_vertical_flip(self): img = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) @@ -684,69 +538,6 @@ def test_vertical_flip(self): # cy -> (H - 1) - cy = 1 - 1 = 0. assert out["intrinsics"][0, 1, 2] == 0.0 - @patch("random.random") - def test_asymmetric_flip_mirror_disparity(self, mock_random): - # 2 calls per orientation (batch of 2 images). - mock_random.side_effect = [0.0, 1.0, 0.5, 0.5] # img0 flips, img1 does not - - disp = torch.zeros(1, 1, 2, 2) - img = torch.zeros(2, 1, 2, 2) - img[0, 0, 0, 0] = 1.0 - img[1, 0, 0, 0] = 2.0 - inputs = {"images": img.clone(), "disparities": disp.clone()} - - out = RandomFlip( - hflip_prob=1.0, - vflip_prob=0.0, - asymmetric_prob=1.0, - disparities_keys=("disparities",), - )(inputs) - - # Image 0 is flipped, image 1 is not. - assert out["images"][0, 0, 0, 1] == 1.0 - assert out["images"][1, 0, 0, 0] == 2.0 - - # Disparity between a flipped and a non-flipped frame is mirrored: - # d(x) -> 2 * (mean - x) - d(x) with mean = 0.5 -> [[1, -1], [1, -1]]. - expected_mirror = torch.tensor([[1.0, -1.0], [1.0, -1.0]]) - assert torch.allclose(out["disparities"][0, 0], expected_mirror) - - @patch("random.random") - def test_asymmetric_flip_last_image_flips(self, mock_random): - # img0 does not flip, img1 flips: the last-image code path must work. - mock_random.side_effect = [1.0, 0.0, 0.5, 0.5] - - img = torch.zeros(2, 1, 2, 2) - img[0, 0, 0, 0] = 1.0 - img[1, 0, 0, 0] = 2.0 - inputs = { - "images": img.clone(), - "disparities": torch.zeros(1, 1, 2, 2), - } - - out = RandomFlip( - hflip_prob=1.0, - vflip_prob=0.0, - asymmetric_prob=1.0, - disparities_keys=("disparities",), - )(inputs) - - assert out["images"][0, 0, 0, 0] == 1.0 # img0 untouched - assert out["images"][1, 0, 0, 1] == 2.0 # img1 flipped - - @patch("random.random") - def test_asymmetric_flip_ignores_missing_disparity_keys(self, mock_random): - # Only "disparities" is provided; the default key list also contains - # "disparities_r", which must be skipped silently when absent. - mock_random.side_effect = [0.0, 1.0, 0.5, 0.5] # img0 flips, img1 does not - - inputs = { - "images": torch.zeros(2, 1, 2, 2), - "disparities": torch.zeros(1, 1, 2, 2), - } - out = RandomFlip(hflip_prob=1.0, vflip_prob=0.0, asymmetric_prob=1.0)(inputs) - assert "disparities" in out - class TestRandomPatchEraser: @patch("random.random") @@ -822,18 +613,6 @@ def test_symmetric_jitter(self, mock_random): assert torch.allclose(out["images"][0], out["images"][1]) assert not torch.allclose(out["images"][0], torch.tensor(0.5)) - @patch("random.random") - def test_default_keys_leave_images_right_untouched(self, mock_random): - mock_random.return_value = 1.0 - - img_left = torch.ones(1, 3, 2, 2) * 0.5 - img_right = torch.ones(1, 3, 2, 2) * 0.5 - inputs = {"images": img_left.clone(), "images_right": img_right.clone()} - - out = ColorJitter(brightness=0.8, asymmetric_prob=0.5)(inputs) - assert not torch.allclose(out["images"], img_left) - assert torch.allclose(out["images_right"], img_right) - def test_ignore_keys(self): inputs = { "images": torch.ones(1, 3, 4, 4), diff --git a/tests/common/data/test_stereo_datamodule.py b/tests/common/data/test_stereo_datamodule.py new file mode 100644 index 0000000..692d75b --- /dev/null +++ b/tests/common/data/test_stereo_datamodule.py @@ -0,0 +1,1107 @@ +"""Unit tests for `roco_spring_devkit.common.data.stereo_datamodule`. + +The datasets and transforms themselves have their own dedicated test modules +(`test_scene_flow_datasets.py`, `test_stereo_transforms.py`), so the focus +here is the *datamodule-specific* logic: + +* the default values set by ``__init__``; +* the string-parsing helper ``_parse_dataset_selection``; +* the YAML path loader ``_load_dataset_paths`` (and the fact that it does NOT + overwrite paths that were already provided, and silently skips YAML keys + that do not map to a ``*_root_dir`` attribute via the ``hasattr`` guard); +* the ``setup`` stage dispatch and assertions; +* ``_get_model_output_stride`` with / without a trainer; +* the dataloader builders ``train_dataloader`` / ``val_dataloader`` / + ``test_dataloader`` / ``predict_dataloader`` (including the Sintel + test-split expansion and the dataset multiplier behaviour); +* the argument-parsing logic of every ``_get__dataset`` helper + (``kitti``, ``sintel``, ``spring``, ``things``) and the ``ValueError`` + raised on unknown flags; +* that ``get_flow``/``get_intrinsics`` are always passed as ``False`` for + the stereo use case. + +The on-disk dataset constructors are mocked out so that we only exercise the +datamodule plumbing, not the (already tested) dataset readers. +""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import yaml + +import roco_spring_devkit.common.data.stereo_datamodule as mod +from roco_spring_devkit.common.data.stereo_datamodule import StereoDataModule + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def _make_dummy_dataset(length: int = 4): + """Return a MagicMock that behaves like a concatenatable Dataset. + + `__add__`/`__iadd__` return the same mock so that train_dataloader can + keep using `+=` (`dataset = dataset + dataset`) without growing new + mock objects. `__len__` returns ``length``. + """ + + dataset = MagicMock(name="DummyDataset") + dataset.__len__ = MagicMock(return_value=length) + dataset.__add__ = MagicMock(return_value=dataset) + dataset.__iadd__ = MagicMock(return_value=dataset) + return dataset + + +def _write_config(tmp_path: Path, paths: dict) -> str: + """Write a YAML dataset config and return its path as a string.""" + cfg_path = tmp_path / "datasets.yaml" + with open(cfg_path, "w") as f: + yaml.safe_dump(paths, f) + return str(cfg_path) + + +# =========================================================================== +# Tests for __init__ defaults +# =========================================================================== +class TestInit: + def test_defaults_are_none_or_zero(self): + dm = StereoDataModule() + assert dm.predict_dataset is None + assert dm.test_dataset is None + assert dm.train_dataset is None + assert dm.val_dataset is None + assert dm.train_batch_size is None + assert dm.train_num_workers == 4 + assert dm.train_crop_size is None + assert dm.train_transform_cuda is False + assert dm.train_transform_fp16 is False + + def test_root_dirs_default_to_none(self): + dm = StereoDataModule() + for attr in ( + "flying_things3d_root_dir", + "flying_things3d_subset_root_dir", + "mpi_sintel_root_dir", + "kitti_2012_root_dir", + "kitti_2015_root_dir", + "tartanair_root_dir", + "spring_root_dir", + "robust_spring_root_dir", + "middlebury_st_root_dir", + ): + assert getattr(dm, attr) is None, attr + + def test_default_dataset_config_path(self): + dm = StereoDataModule() + # The default config path is hard-coded for the CLI's working dir. + assert dm.dataset_config_path == "../../datasets.yaml" + + def test_init_stores_constructor_values(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=3, + train_num_workers=2, + train_crop_size=(100, 200), + train_transform_cuda=True, + train_transform_fp16=True, + dataset_config_path=cfg, + ) + assert dm.train_dataset == "kitti-train" + assert dm.val_dataset == "sintel-clean-trainval" + assert dm.train_batch_size == 3 + assert dm.train_num_workers == 2 + assert dm.train_crop_size == (100, 200) + assert dm.train_transform_cuda is True + assert dm.train_transform_fp16 is True + assert dm.dataset_config_path == cfg + + def test_init_initialises_internal_state(self): + dm = StereoDataModule() + assert dm.predict_dataset_parsed is None + assert dm.test_dataset_parsed is None + assert dm.train_dataset_parsed is None + assert dm.val_dataset_parsed is None + assert dm.train_dataloader_length == 0 + assert dm.train_epoch_step == 0 + assert dm.val_dataloader_names == [] + assert dm.val_dataloader_lengths == [] + assert dm.test_dataloader_names == [] + + +# =========================================================================== +# Tests for _parse_dataset_selection +# =========================================================================== +class TestParseDatasetSelection: + def test_none_returns_empty_list(self): + dm = StereoDataModule() + assert dm._parse_dataset_selection(None) == [] + + def test_single_dataset_default_multiplier(self): + dm = StereoDataModule() + out = dm._parse_dataset_selection("kitti-train") + assert out == [(1, "kitti", "train")] + + def test_multiple_datasets_separated_by_plus(self): + dm = StereoDataModule() + out = dm._parse_dataset_selection( + "kitti-train+sintel-clean-trainval+things-clean" + ) + assert out == [ + (1, "kitti", "train"), + (1, "sintel", "clean", "trainval"), + (1, "things", "clean"), + ] + + def test_leading_multiplier(self): + dm = StereoDataModule() + out = dm._parse_dataset_selection("3*sintel-clean-trainval") + assert out == [(3, "sintel", "clean", "trainval")] + + def test_trailing_multiplier(self): + dm = StereoDataModule() + out = dm._parse_dataset_selection("kitti-train*5") + assert out == [(5, "kitti", "train")] + + def test_mixed_multipliers_and_flags(self): + dm = StereoDataModule() + out = dm._parse_dataset_selection( + "kitti-train+3*sintel-clean-trainval+things-train*5" + ) + assert out == [ + (1, "kitti", "train"), + (3, "sintel", "clean", "trainval"), + (5, "things", "train"), + ] + + def test_spaces_are_stripped(self): + dm = StereoDataModule() + out = dm._parse_dataset_selection(" kitti - train + 2 * sintel - clean ") + assert out == [ + (1, "kitti", "train"), + (2, "sintel", "clean"), + ] + + def test_invalid_double_multiplier_raises(self): + dm = StereoDataModule() + with pytest.raises(ValueError): + dm._parse_dataset_selection("3*4*sintel-clean") + + def test_non_integer_multiplier_raises_value_error(self): + # 'x' is not int -> the trailing-multiplier branch raises ValueError + # when trying int('x'). + dm = StereoDataModule() + with pytest.raises(ValueError): + dm._parse_dataset_selection("sintel*x") + + def test_empty_string_yields_single_empty_dataset(self): + # An empty string splits into [''] which is treated as a single + # dataset with empty name and multiplier 1. This documents current + # behaviour (no guard for empty input). + dm = StereoDataModule() + out = dm._parse_dataset_selection("") + assert out == [(1, "")] + + +# =========================================================================== +# Tests for _load_dataset_paths +# =========================================================================== +class TestLoadDatasetPaths: + def test_fills_none_attrs_from_yaml(self, tmp_path): + cfg = _write_config( + tmp_path, + { + "flying_things3d": "/things/path", + "mpi_sintel": "/sintel/path", + "kitti_2015": "/kitti/path", + }, + ) + dm = StereoDataModule(dataset_config_path=cfg) + dm._load_dataset_paths() + assert dm.flying_things3d_root_dir == "/things/path" + assert dm.mpi_sintel_root_dir == "/sintel/path" + assert dm.kitti_2015_root_dir == "/kitti/path" + + def test_does_not_overwrite_existing_paths(self, tmp_path): + cfg = _write_config( + tmp_path, + {"flying_things3d": "/yaml/things", "mpi_sintel": "/yaml/sintel"}, + ) + dm = StereoDataModule( + flying_things3d_root_dir="/explicit/things", + dataset_config_path=cfg, + ) + dm._load_dataset_paths() + # The explicitly-provided path wins. + assert dm.flying_things3d_root_dir == "/explicit/things" + # The unset one is filled from YAML. + assert dm.mpi_sintel_root_dir == "/yaml/sintel" + + def test_missing_keys_in_yaml_leave_attrs_as_none(self, tmp_path): + cfg = _write_config(tmp_path, {"flying_things3d": "/things/path"}) + dm = StereoDataModule(dataset_config_path=cfg) + dm._load_dataset_paths() + assert dm.flying_things3d_root_dir == "/things/path" + # Keys absent from YAML are left untouched (still None). + assert dm.mpi_sintel_root_dir is None + assert dm.kitti_2015_root_dir is None + + def test_unmatched_yaml_keys_are_skipped_via_hasattr_guard(self, tmp_path): + # The stereo datamodule does NOT have e.g. 'autoflow_root_dir' or + # 'flying_chairs2_root_dir' attributes, unlike the optical flow one. + # `_load_dataset_paths` uses `hasattr(self, "_root_dir")` to + # silently skip such YAML keys instead of raising AttributeError. + cfg = _write_config( + tmp_path, + { + "autoflow": "/aut/oflow", # no autoflow_root_dir attr + "flying_chairs": "/chairs/path", # no flying_chairs_root_dir attr + "kubric": "/kubric/path", # no kubric_root_dir attr + "mpi_sintel": "/sintel/path", # has mpi_sintel_root_dir attr + }, + ) + dm = StereoDataModule(dataset_config_path=cfg) + # Should not raise. + dm._load_dataset_paths() + # Matched key is filled. + assert dm.mpi_sintel_root_dir == "/sintel/path" + # Unmatched keys do not create new attributes. + assert not hasattr(dm, "autoflow_root_dir") + assert not hasattr(dm, "flying_chairs_root_dir") + assert not hasattr(dm, "kubric_root_dir") + + +# =========================================================================== +# Tests for setup +# =========================================================================== +class TestSetup: + def _make(self, tmp_path, **kwargs): + cfg = _write_config(tmp_path, {}) + kwargs.setdefault("dataset_config_path", cfg) + return StereoDataModule(**kwargs) + + def test_fit_parses_train_and_val(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=4, + ) + dm.setup("fit") + assert dm.train_dataset_parsed == [(1, "kitti", "train")] + assert dm.val_dataset_parsed == [(1, "sintel", "clean", "trainval")] + + def test_fit_assigns_default_batch_size_when_none(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=None, + ) + dm.setup("fit") + assert dm.train_batch_size == 8 + + def test_fit_loads_dataset_paths(self, tmp_path): + cfg = _write_config(tmp_path, {"mpi_sintel": "/sintel/path"}) + dm = StereoDataModule( + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + train_batch_size=4, + dataset_config_path=cfg, + ) + dm.setup("fit") + assert dm.mpi_sintel_root_dir == "/sintel/path" + + def test_fit_without_train_dataset_raises(self, tmp_path): + dm = self._make(tmp_path, val_dataset="sintel-clean-trainval") + with pytest.raises(AssertionError): + dm.setup("fit") + + def test_fit_without_val_dataset_raises(self, tmp_path): + dm = self._make(tmp_path, train_dataset="kitti-train", train_batch_size=2) + with pytest.raises(AssertionError): + dm.setup("fit") + + def test_predict_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("predict") + + def test_predict_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, predict_dataset="sintel-clean-test") + dm.setup("predict") + # NB: the implementation has a typo - it writes to + # `parsed_predict_dataset_parsed` (not `predict_dataset_parsed`). We + # assert the actual field name so the test matches current behaviour. + assert dm.parsed_predict_dataset_parsed == [(1, "sintel", "clean", "test")] + + def test_test_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("test") + + def test_test_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, test_dataset="kitti") + dm.setup("test") + assert dm.test_dataset_parsed == [(1, "kitti")] + + def test_validate_without_dataset_raises(self, tmp_path): + dm = self._make(tmp_path) + with pytest.raises(AssertionError): + dm.setup("validate") + + def test_validate_parses_dataset(self, tmp_path): + dm = self._make(tmp_path, val_dataset="sintel-clean-trainval") + dm.setup("validate") + assert dm.val_dataset_parsed == [(1, "sintel", "clean", "trainval")] + + def test_unknown_stage_is_noop(self, tmp_path): + dm = self._make( + tmp_path, + train_dataset="kitti-train", + val_dataset="sintel-clean-trainval", + ) + # No assert / no parse should fire for an unrecognised stage. + dm.setup("some-other-stage") + assert dm.train_dataset_parsed is None + assert dm.val_dataset_parsed is None + + +# =========================================================================== +# Tests for _get_model_output_stride +# =========================================================================== +class TestGetModelOutputStride: + def test_returns_one_when_no_trainer(self): + dm = StereoDataModule() + # 'trainer' is not set on the datamodule until Lightning assigns one. + assert dm._get_model_output_stride() == 1 + + def test_returns_one_when_trainer_none(self): + dm = StereoDataModule() + dm.trainer = None + assert dm._get_model_output_stride() == 1 + + def test_returns_model_output_stride(self): + dm = StereoDataModule() + # Use SimpleNamespace so `hasattr(model, "module")` is False (MagicMock + # would auto-create the `module` attribute and break the branch logic). + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + assert dm._get_model_output_stride() == 8 + + def test_returns_module_output_stride_when_ddp(self): + dm = StereoDataModule() + # When using DistributedDataParallel, trainer.model has a `module` + # attribute - the datamodule should unwrap it. + trainer = SimpleNamespace( + model=SimpleNamespace(module=SimpleNamespace(output_stride=16)) + ) + dm.trainer = trainer + assert dm._get_model_output_stride() == 16 + + +# =========================================================================== +# Tests for predict_dataloader +# =========================================================================== +class TestPredictDataloader: + def test_predict_dataloader_delegates_to_super(self): + dm = StereoDataModule() + # The implementation just calls super().predict_dataloader() which + # raises a MisconfigurationException (no predict loader is actually + # implemented). Document that behaviour. + from lightning.fabric.utilities.exceptions import ( + MisconfigurationException, + ) + + with pytest.raises(MisconfigurationException): + dm.predict_dataloader() + + +# =========================================================================== +# Tests for train_dataloader +# =========================================================================== +class TestTrainDataloader: + def test_returns_none_when_train_dataset_parsed_is_none(self): + dm = StereoDataModule() + dm.train_dataset_parsed = None + assert dm.train_dataloader() is None + + def test_concatenates_datasets_with_multiplier(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_batch_size=2, + train_num_workers=0, + train_transform_cuda=False, + dataset_config_path=cfg, + ) + # Two parsed datasets: multiplier=2 for 'kitti' and multiplier=3 for + # 'sintel'. + dm.train_dataset_parsed = [ + (2, "kitti", "train"), + (3, "sintel", "clean", "trainval"), + ] + kitti_dummy = _make_dummy_dataset(length=5) + sintel_dummy = _make_dummy_dataset(length=7) + + with patch.object( + StereoDataModule, "_get_kitti_dataset", autospec=True + ) as p_kitti, patch.object( + StereoDataModule, "_get_sintel_dataset", autospec=True + ) as p_sintel: + p_kitti.side_effect = lambda self, is_train, *a: kitti_dummy + p_sintel.side_effect = lambda self, is_train, *a: sintel_dummy + + loader = dm.train_dataloader() + + assert loader is not None + assert dm.train_dataloader_length == len(loader) + # Each `_get__dataset` should have been called exactly once. + assert p_kitti.call_count == 1 + assert p_sintel.call_count == 1 + # is_train should be True for the train dataloader. + kitti_args = p_kitti.call_args + assert kitti_args.args[0] is dm # self (autospec) + assert kitti_args.args[1] is True + assert kitti_args.args[2:] == ("train",) + sintel_args = p_sintel.call_args + assert sintel_args.args[1] is True + assert sintel_args.args[2:] == ("clean", "trainval") + + def test_dataloader_uses_correct_loader_kwargs(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_batch_size=4, + train_num_workers=0, + train_transform_cuda=False, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "kitti", "train")] + dummy = _make_dummy_dataset(length=3) + + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + loader = dm.train_dataloader() + assert loader.batch_size == 4 + assert loader.num_workers == 0 + # pin_memory is True when train_transform_cuda is False. + assert loader.pin_memory is True + # drop_last is False. + assert loader.drop_last is False + # persistent_workers is set to train_transform_cuda -> False. + assert loader.persistent_workers is False + + def test_cuda_disables_pin_memory_and_enables_persistent_workers(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_batch_size=2, + train_num_workers=1, + train_transform_cuda=True, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "kitti", "train")] + dummy = _make_dummy_dataset(length=2) + + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + loader = dm.train_dataloader() + assert loader.pin_memory is False + assert loader.persistent_workers is True + + def test_single_dataset_multiplier_one(self, tmp_path): + # multiplier==1 means the `for _ in range(multiplier - 1)` loop body + # never runs, so `dataset_mult` stays equal to `dataset`. + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_batch_size=2, + train_num_workers=0, + dataset_config_path=cfg, + ) + dm.train_dataset_parsed = [(1, "kitti", "train")] + dummy = _make_dummy_dataset(length=2) + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ) as p_kitti: + dm.train_dataloader() + # Only one call to the getter. + assert p_kitti.call_count == 1 + # No `+`/`+=` was performed on the dataset (multiplier == 1). + dummy.__add__.assert_not_called() + + +# =========================================================================== +# Tests for val_dataloader +# =========================================================================== +class TestValDataloader: + def test_returns_list_of_dataloaders_with_names_and_lengths(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_batch_size=2, + dataset_config_path=cfg, + train_transform_cuda=False, + ) + dm.val_dataset_parsed = [ + (1, "sintel", "clean", "trainval"), + (1, "kitti", "val"), + ] + sintel_dummy = _make_dummy_dataset(length=11) + kitti_dummy = _make_dummy_dataset(length=5) + + with patch.object( + StereoDataModule, + "_get_sintel_dataset", + autospec=True, + return_value=sintel_dummy, + ) as p_sintel, patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=kitti_dummy, + ) as p_kitti: + loaders = dm.val_dataloader() + + assert len(loaders) == 2 + # Each loader has batch_size 1, num_workers 1, pin_memory False. + for ld in loaders: + assert ld.batch_size == 1 + assert ld.num_workers == 1 + assert ld.pin_memory is False + assert ld.drop_last is False + # Names are formed by "-".join(parsed_vals[1:]). + assert dm.val_dataloader_names == ["sintel-clean-trainval", "kitti-val"] + # Lengths come from len(dataset) (for a MagicMock with __len__). + assert dm.val_dataloader_lengths == [11, 5] + # `_get__dataset` must be called with is_train=False. + assert p_sintel.call_args.args[1] is False + assert p_sintel.call_args.args[2:] == ("clean", "trainval") + assert p_kitti.call_args.args[1] is False + assert p_kitti.call_args.args[2:] == ("val",) + + def test_resets_names_and_lengths_each_call(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg, train_transform_cuda=False) + # Pre-populate to ensure val_dataloader() overwrites these lists. + dm.val_dataloader_names = ["stale"] + dm.val_dataloader_lengths = [999] + dm.val_dataset_parsed = [(2, "kitti", "val")] + dummy = _make_dummy_dataset(length=4) + + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + dm.val_dataloader() + + assert dm.val_dataloader_names == ["kitti-val"] + assert dm.val_dataloader_lengths == [4] + + def test_persistent_workers_follows_train_transform_cuda(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg, train_transform_cuda=True) + dm.val_dataset_parsed = [(1, "kitti", "val")] + dummy = _make_dummy_dataset(length=2) + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + loaders = dm.val_dataloader() + assert loaders[0].persistent_workers is True + + +# =========================================================================== +# Tests for test_dataloader +# =========================================================================== +class TestTestDataloader: + def test_single_dataset_calls_get_with_test_suffix(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + dm.test_dataset = "kitti" + dummy = _make_dummy_dataset(length=6) + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ) as p_kitti: + loaders = dm.test_dataloader() + assert len(loaders) == 1 + assert p_kitti.call_count == 1 + # The datamodule appends "-test" then splits, so args are ("test",). + assert p_kitti.call_args.args[1] is False + assert p_kitti.call_args.args[2:] == ("test",) + # The name appended to test_dataloader_names is the suffixed id. + assert dm.test_dataloader_names == ["kitti-test"] + # Each loader uses batch_size 1, num_workers 1, no pin_memory. + assert loaders[0].batch_size == 1 + assert loaders[0].num_workers == 1 + assert loaders[0].pin_memory is False + assert loaders[0].drop_last is False + + def test_things_test_expansion_via_test_dataloader(self, tmp_path): + # 'things' is NOT specially expanded (only 'sintel' is); this test + # confirms the generic path works for the things dataset. + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + dm.test_dataset = "things" + dummy = _make_dummy_dataset(length=3) + with patch.object( + StereoDataModule, + "_get_things_dataset", + autospec=True, + return_value=dummy, + ) as p_things: + loaders = dm.test_dataloader() + assert len(loaders) == 1 + # `things-test` -> tokens ("things", "test"). Note, however, that + # `_get_things_dataset` only accepts split tokens in + # ["train", "val", "trainval"]; the test stage thus raises a + # ValueError when actually invoked. Here we patched the getter so we + # only check that the dispatcher used the right name and token stream. + assert p_things.call_args.args[1] is False + assert p_things.call_args.args[2:] == ("test",) + assert dm.test_dataloader_names == ["things-test"] + + def test_sintel_expands_to_clean_and_final(self, tmp_path): + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + dm.test_dataset = "sintel" + dummy = _make_dummy_dataset(length=3) + with patch.object( + StereoDataModule, + "_get_sintel_dataset", + autospec=True, + return_value=dummy, + ) as p_sintel: + loaders = dm.test_dataloader() + assert len(loaders) == 2 + # Two calls: one for sintel-clean-test, one for sintel-final-test. + call_args = [c.args[2:] for c in p_sintel.call_args_list] + assert call_args == [("clean", "test"), ("final", "test")] + assert dm.test_dataloader_names == ["sintel-clean-test", "sintel-final-test"] + + def test_test_dataloader_names_accumulate_across_calls(self, tmp_path): + # The implementation appends to `test_dataloader_names` rather than + # resetting it. Document that behaviour. + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + dm.test_dataset = "kitti" + dm.test_dataloader_names = ["already-here"] + dummy = _make_dummy_dataset(length=2) + with patch.object( + StereoDataModule, + "_get_kitti_dataset", + autospec=True, + return_value=dummy, + ): + dm.test_dataloader() + assert dm.test_dataloader_names == ["already-here", "kitti-test"] + + +# =========================================================================== +# Tests for _get__dataset argument parsing +# --------------------------------------------------------------------------- +# The dataset constructors are mocked so we can verify the datamodule's +# argument-parsing logic in isolation. +# =========================================================================== +class TestGetDatasetArgumentParsing: + @pytest.fixture + def dm_cpu(self, tmp_path): + # Use a small crop size so we don't depend on default values. + cfg = _write_config(tmp_path, {}) + return StereoDataModule( + train_crop_size=(100, 200), + train_transform_cuda=False, + train_transform_fp16=False, + dataset_config_path=cfg, + ) + + # ----- _get_kitti_dataset ----- + def test_kitti_default_split(self, dm_cpu): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "trainval" + assert kwargs["get_flow"] is False + assert kwargs["get_intrinsics"] is False + # Stereo kitti always uses the 2015 root dir. + assert kwargs["root_dir_2015"] == dm_cpu.kitti_2015_root_dir + + def test_kitti_passes_split_and_version_token(self, dm_cpu): + # '2015' is a no-op pass-through token (the stereo KittiDataset only + # supports the 2015 version). + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False, "2015", "train") + assert p.call_args.kwargs["split"] == "train" + + def test_kitti_test_split(self, dm_cpu): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm_cpu._get_kitti_dataset(False, "test") + assert p.call_args.kwargs["split"] == "test" + + def test_kitti_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_kitti_dataset(False, "what") + + # ----- _get_sintel_dataset ----- + def test_sintel_default_both_passes(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False) + assert p.call_args.kwargs["pass_names"] == ["clean", "final"] + assert p.call_args.kwargs["get_flow"] is False + assert p.call_args.kwargs["get_intrinsics"] is False + assert p.call_args.args[0] == dm_cpu.mpi_sintel_root_dir + + def test_sintel_pass_selection(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "final") + assert p.call_args.kwargs["pass_names"] == ["final"] + + def test_sintel_train_split(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "clean", "train") + assert p.call_args.kwargs["split"] == "train" + assert p.call_args.kwargs["pass_names"] == ["clean"] + + def test_sintel_test_split(self, dm_cpu): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm_cpu._get_sintel_dataset(False, "clean", "test") + assert p.call_args.kwargs["split"] == "test" + + def test_sintel_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_sintel_dataset(False, "unknown") + + # ----- _get_spring_dataset ----- + def test_spring_defaults(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset(False) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "train" + assert kwargs["add_camera_reverse"] is False + assert kwargs["camera_reverse_only"] is False + assert kwargs["subsample"] is True + assert kwargs["get_right_disparity"] is False + assert kwargs["robust_mode"] is False + assert kwargs["get_flow"] is False + assert kwargs["get_intrinsics"] is False + # Robust root dir is forwarded. + assert kwargs["robust_root_dir"] == dm_cpu.robust_spring_root_dir + + def test_spring_flags(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset( + False, "val", "camrev", "camrevonly", "gt4k", "right", "robust" + ) + kwargs = p.call_args.kwargs + assert kwargs["split"] == "val" + assert kwargs["add_camera_reverse"] is True + assert kwargs["camera_reverse_only"] is True + assert kwargs["subsample"] is False + assert kwargs["get_right_disparity"] is True + assert kwargs["robust_mode"] is True + + def test_spring_test_split(self, dm_cpu): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm_cpu._get_spring_dataset(False, "test") + assert p.call_args.kwargs["split"] == "test" + + def test_spring_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_spring_dataset(False, "garbage") + + # ----- _get_things_dataset ----- + def test_things_default_pass_names(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False) + assert p.call_args.kwargs["pass_names"] == ["clean", "final"] + assert p.call_args.kwargs["get_flow"] is False + assert p.call_args.kwargs["get_intrinsics"] is False + assert p.call_args.args[0] == dm_cpu.flying_things3d_root_dir + + def test_things_pass_selection(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False, "clean") + assert p.call_args.kwargs["pass_names"] == ["clean"] + + def test_things_train_split(self, dm_cpu): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm_cpu._get_things_dataset(False, "train") + assert p.call_args.kwargs["split"] == "train" + + def test_things_test_split_raises_value_error(self, dm_cpu): + # Documenting current (intentional or not) behaviour: the things arg + # parser only accepts splits in ["train", "val", "trainval"], so a + # "test" token (as the test_dataloader would append) raises. + with pytest.raises(ValueError): + dm_cpu._get_things_dataset(False, "test") + + def test_things_sinteltransform_flag(self, dm_cpu): + # The sinteltransform flag only changes the train-time scale range; + # when is_train=False the flag should be accepted without error. + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()): + # Should not raise. + dm_cpu._get_things_dataset(False, "sinteltransform") + + def test_things_invalid_arg_raises(self, dm_cpu): + with pytest.raises(ValueError): + dm_cpu._get_things_dataset(False, "nope") + + +# =========================================================================== +# Tests for transform construction in _get__dataset (train mode) +# =========================================================================== +class TestTrainTransformConstruction: + def _kitti_dm(self, tmp_path): + cfg = _write_config(tmp_path, {}) + return StereoDataModule( + train_crop_size=(100, 200), + train_transform_cuda=False, + train_transform_fp16=False, + dataset_config_path=cfg, + ) + + def test_kitti_train_builds_compose_transform(self, tmp_path): + dm = self._kitti_dm(tmp_path) + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import stereo_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_sintel_train_builds_compose_transform(self, tmp_path): + dm = self._kitti_dm(tmp_path) + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm._get_sintel_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import stereo_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_spring_train_builds_compose_transform(self, tmp_path): + dm = self._kitti_dm(tmp_path) + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm._get_spring_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import stereo_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_things_train_builds_compose_transform(self, tmp_path): + dm = self._kitti_dm(tmp_path) + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm._get_things_dataset(True, "train") + transform = p.call_args.kwargs["transform"] + from roco_spring_devkit.common.data import stereo_transforms as ft + + assert isinstance(transform, ft.Compose) + + def test_eval_mode_uses_plain_totensor(self, tmp_path): + # In eval mode, the transform is a bare ToTensor (not a Compose). + from roco_spring_devkit.common.data import stereo_transforms as ft + + dm = self._kitti_dm(tmp_path) + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(False) + assert isinstance(p.call_args.kwargs["transform"], ft.ToTensor) + + def test_things_sinteltransform_changes_scale_range_in_train_mode(self, tmp_path): + # With sinteltransform=True the train-time major_scale is (-0.2, 0.5); + # without it it is (-0.4, 0.8). We assert the values reach the + # RandomScaleAndCrop constructor. + from roco_spring_devkit.common.data import stereo_transforms as ft + + captured = [] + + real_rsac = ft.RandomScaleAndCrop + + def spy_rsac(crop_size, major_scale=None, space_scale=None, sparse=False): + obj = real_rsac( + crop_size, + major_scale=major_scale, + space_scale=space_scale, + sparse=sparse, + ) + captured.append(major_scale) + return obj + + dm = self._kitti_dm(tmp_path) + with patch.object( + mod, "FlyingThings3DDataset", return_value=MagicMock() + ), patch.object(ft, "RandomScaleAndCrop", spy_rsac): + dm._get_things_dataset(True, "sinteltransform") + assert captured == [(-0.2, 0.5)] + + captured.clear() + with patch.object( + mod, "FlyingThings3DDataset", return_value=MagicMock() + ), patch.object(ft, "RandomScaleAndCrop", spy_rsac): + dm._get_things_dataset(True) + assert captured == [(-0.4, 0.8)] + + def test_train_mode_sets_default_crop_size_when_none(self, tmp_path): + # For kitti the default crop size is (288, 960). + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + assert dm.train_crop_size is None + with patch.object(mod, "KittiDataset", return_value=MagicMock()): + dm._get_kitti_dataset(True, "train") + # Was auto-filled. + assert dm.train_crop_size is not None + cy, cx = dm.train_crop_size + # With output_stride=1, make_divisible(x, 1) = x. + assert cy == 288 + assert cx == 960 + + def test_train_default_crop_sintel(self, tmp_path): + # For sintel the default crop size is (368, 768). + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + with patch.object(mod, "SintelDataset", return_value=MagicMock()): + dm._get_sintel_dataset(True, "train") + assert dm.train_crop_size == (368, 768) + + def test_train_default_crop_spring(self, tmp_path): + # For spring the default crop size is (540, 960). + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + with patch.object(mod, "SpringDataset", return_value=MagicMock()): + dm._get_spring_dataset(True, "train") + assert dm.train_crop_size == (540, 960) + + def test_train_default_crop_things(self, tmp_path): + # For things the default crop size is (400, 720). + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule(dataset_config_path=cfg) + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()): + dm._get_things_dataset(True, "train") + assert dm.train_crop_size == (400, 720) + + def test_train_crop_size_made_divisible_by_output_stride(self, tmp_path): + # When an output stride > 1 is set, the crop dims passed to the + # transform are rounded down to a multiple of the stride. The + # datamodule does NOT write the rounded values back to + # ``self.train_crop_size`` (only the `None` branch does), so we + # inspect the values that reach ``RandomScaleAndCrop`` instead. + from roco_spring_devkit.common.data import stereo_transforms as ft + from types import SimpleNamespace + + captured = [] + real_rsac = ft.RandomScaleAndCrop + + def spy_rsac(crop_size, major_scale=None, space_scale=None, sparse=False): + captured.append(crop_size) + return real_rsac( + crop_size, + major_scale=major_scale, + space_scale=space_scale, + sparse=sparse, + ) + + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_crop_size=(290, 961), + dataset_config_path=cfg, + ) + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + with patch.object(mod, "KittiDataset", return_value=MagicMock()), patch.object( + ft, "RandomScaleAndCrop", spy_rsac + ): + dm._get_kitti_dataset(True, "train") + # make_divisible(290, 8) = 288, make_divisible(961, 8) = 960. + assert captured == [(288, 960)] + # And the original train_crop_size attribute is untouched. + assert dm.train_crop_size == (290, 961) + + def test_train_crop_size_clamped_to_at_least_stride(self, tmp_path): + # make_divisible returns max(div, v - v % div); a value smaller than + # the stride is clamped up to the stride. + from roco_spring_devkit.common.data import stereo_transforms as ft + from types import SimpleNamespace + + captured = [] + real_rsac = ft.RandomScaleAndCrop + + def spy_rsac(crop_size, major_scale=None, space_scale=None, sparse=False): + captured.append(crop_size) + return real_rsac( + crop_size, + major_scale=major_scale, + space_scale=space_scale, + sparse=sparse, + ) + + cfg = _write_config(tmp_path, {}) + dm = StereoDataModule( + train_crop_size=(3, 5), + dataset_config_path=cfg, + ) + trainer = SimpleNamespace(model=SimpleNamespace(output_stride=8)) + dm.trainer = trainer + with patch.object(mod, "KittiDataset", return_value=MagicMock()), patch.object( + ft, "RandomScaleAndCrop", spy_rsac + ): + dm._get_kitti_dataset(True, "train") + assert captured == [(8, 8)] + + +# =========================================================================== +# Tests for get_flow / get_intrinsics always False +# =========================================================================== +class TestStereoFlags: + """Verify the stereo datamodule always asks the dataset constructors for + ``get_flow=False`` and ``get_intrinsics=False`` (this is what makes the + produced samples stereo-specific rather than flow-specific).""" + + @pytest.fixture + def dm(self, tmp_path): + cfg = _write_config(tmp_path, {}) + return StereoDataModule( + train_crop_size=(100, 200), + dataset_config_path=cfg, + ) + + def test_kitti_disables_flow_and_intrinsics(self, dm): + with patch.object(mod, "KittiDataset", return_value=MagicMock()) as p: + dm._get_kitti_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is False + assert kwargs["get_intrinsics"] is False + + def test_sintel_disables_flow_and_intrinsics(self, dm): + with patch.object(mod, "SintelDataset", return_value=MagicMock()) as p: + dm._get_sintel_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is False + assert kwargs["get_intrinsics"] is False + + def test_spring_disables_flow_and_intrinsics(self, dm): + with patch.object(mod, "SpringDataset", return_value=MagicMock()) as p: + dm._get_spring_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is False + assert kwargs["get_intrinsics"] is False + + def test_things_disables_flow_and_intrinsics(self, dm): + with patch.object(mod, "FlyingThings3DDataset", return_value=MagicMock()) as p: + dm._get_things_dataset(True, "train") + kwargs = p.call_args.kwargs + assert kwargs["get_flow"] is False + assert kwargs["get_intrinsics"] is False diff --git a/tests/common/utils/test_correlation.py b/tests/common/utils/test_correlation.py index 949c7e0..1047ee8 100644 --- a/tests/common/utils/test_correlation.py +++ b/tests/common/utils/test_correlation.py @@ -20,7 +20,6 @@ iter_translated_spatial_correlation_sample, ) - # --------------------------------------------------------------------------- # Hand-computed expected value for the canonical 3x3 single-channel case. # diff --git a/tests/common/utils/test_flow_utils.py b/tests/common/utils/test_flow_utils.py index fabbd94..65c4795 100644 --- a/tests/common/utils/test_flow_utils.py +++ b/tests/common/utils/test_flow_utils.py @@ -77,7 +77,6 @@ from roco_spring_devkit.common.utils import flow_utils - # --------------------------------------------------------------------------- # flow_to_rgb # --------------------------------------------------------------------------- diff --git a/tests/common/utils/test_io_adapter.py b/tests/common/utils/test_io_adapter.py index e72018c..8133fcb 100644 --- a/tests/common/utils/test_io_adapter.py +++ b/tests/common/utils/test_io_adapter.py @@ -95,7 +95,6 @@ from roco_spring_devkit.common.utils.io_adapter import IOAdapter - # --------------------------------------------------------------------------- # Construction / scaler wiring # --------------------------------------------------------------------------- diff --git a/tests/common/utils/test_scene_flow_metrics.py b/tests/common/utils/test_scene_flow_metrics.py index 2e6420c..0a49e48 100644 --- a/tests/common/utils/test_scene_flow_metrics.py +++ b/tests/common/utils/test_scene_flow_metrics.py @@ -70,7 +70,20 @@ def test_default_attributes(self) -> None: # 1 / 0.01000... ~= 99.999 which floors to 99. assert m.ema_max_count == 99 assert m.include_occlusion is False - assert m.used_keys == [] + assert m.used_keys == [ + ("epe", "epe_flow", "valid_flows_target"), + ("1px_flow", "px1_flow_mask", "valid_flows_target"), + ("flall", "flall_mask", "valid_flows_target"), + ("wauc", "wauc_flow", "valid_flows_target"), + ("abs1", "abs1", "valid_disp1_target"), + ("1px1", "px11_mask", "valid_disp1_target"), + ("d1", "d1_mask", "valid_disp1_target"), + ("abs2", "abs2", "valid_disp2_target"), + ("1px2", "px12_mask", "valid_disp2_target"), + ("d2", "d2_mask", "valid_disp2_target"), + ("1px_all", "px1_all_mask", "valid_all_target"), + ("sfall", "sfall_mask", "valid_all_target"), + ] def test_invalid_average_mode_raises(self) -> None: with pytest.raises(AssertionError): @@ -96,8 +109,10 @@ def test_prefix_attaches_to_metric_keys(self) -> None: "val_flall", "val_wauc", "val_abs1", + "val_1px1", "val_d1", "val_abs2", + "val_1px2", "val_d2", "val_1px_all", "val_sfall", diff --git a/tests/common/utils/test_stereo_utils.py b/tests/common/utils/test_stereo_utils.py index 69b2561..667af8c 100644 --- a/tests/common/utils/test_stereo_utils.py +++ b/tests/common/utils/test_stereo_utils.py @@ -66,7 +66,6 @@ from roco_spring_devkit.common.utils import stereo_utils - # Hand-verified plasma LUT entries (independently fetched from OpenCV in a # scratch session, NOT from the production code). PLASMA_LUT = { diff --git a/tests/optical_flow/models/fixtures/raft_forward_reference.pt b/tests/optical_flow/models/fixtures/raft_forward_reference.pt new file mode 100644 index 0000000..7b86b67 Binary files /dev/null and b/tests/optical_flow/models/fixtures/raft_forward_reference.pt differ diff --git a/tests/optical_flow/models/test_raft_validation.py b/tests/optical_flow/models/test_raft_validation.py new file mode 100644 index 0000000..b6649fb --- /dev/null +++ b/tests/optical_flow/models/test_raft_validation.py @@ -0,0 +1,124 @@ +import csv +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +VALIDATE_SCRIPT = REPO_ROOT / "roco_spring_devkit" / "optical_flow" / "validate.py" +OPTICAL_FLOW_DIR = REPO_ROOT / "roco_spring_devkit" / "optical_flow" +FIXTURE_PATH = ( + Path(__file__).resolve().parent / "fixtures" / "raft_forward_reference.pt" +) +METRICS_MARGIN = 0.01 + +REFERENCE_METRICS = { + "spring-val-val/epe": 0.8079047306544251, + "spring-val-val/1px": 7.064254628701343, + "spring-val-val/flall": 0.3262576089344091, + "spring-val-val/wauc": 85.14406479729547, +} + + +def _parse_metrics_csv(csv_path): + parsed = {} + with open(csv_path, "r") as f: + reader = csv.reader(f) + for row in reader: + if len(row) >= 2 and row[0] in REFERENCE_METRICS: + parsed[row[0]] = float(row[1]) + return parsed + + +@pytest.mark.slow +def test_raft_sintel_validation_metrics(tmp_path): + output_path = tmp_path / "validate" + cmd = [ + sys.executable, + str(VALIDATE_SCRIPT), + "--data.val_dataset", + "spring-val", + "--model", + "raft", + "--ckpt_path", + "sintel", + "--model.corr_mode", + "triton", + "--output_path", + str(output_path), + ] + result = subprocess.run( + cmd, + cwd=str(OPTICAL_FLOW_DIR), + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"validate.py failed with code {result.returncode}\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + + metrics_csv = output_path / "raft_sintel" / "metrics.csv" + assert metrics_csv.exists(), f"metrics.csv not found at {metrics_csv}" + + parsed = _parse_metrics_csv(metrics_csv) + assert set(parsed.keys()) == set( + REFERENCE_METRICS.keys() + ), f"Expected metrics {set(REFERENCE_METRICS.keys())}, got {set(parsed.keys())}" + + failures = [] + for key, ref in REFERENCE_METRICS.items(): + val = parsed[key] + diff = abs(val - ref) + if diff >= METRICS_MARGIN: + failures.append( + f"Metric {key}: got {val}, expected {ref}, diff {diff} >= margin {METRICS_MARGIN}" + ) + assert not failures, "\n".join(failures) + + +@pytest.mark.slow +def test_raft_forward_deterministic(): + """A fixed input must always produce the same (previously recorded) outputs. + + The model weights are made deterministic by loading the pretrained "sintel" + checkpoint (manually setting all weights to a constant value was attempted, but + it caused the flow predictions to explode, so the checkpoint fallback is used). + The deterministic input and the reference outputs were recorded once (see + fixtures/raft_forward_reference.pt) and are reused here. + """ + from roco_spring_devkit import get_model + + fixture = torch.load(FIXTURE_PATH, map_location="cpu", weights_only=False) + + model = get_model(fixture["model_name"], fixture["ckpt_path"]) + model.eval() + model.freeze_bn() + model = model.to("cpu") + + inputs = {"images": fixture["input_images"].to("cpu")} + + with torch.no_grad(): + outputs = model(inputs) + + reference = fixture["outputs"] + assert set(outputs.keys()) >= set(reference.keys()), ( + f"Missing output keys: expected {set(reference.keys())}, " + f"got {set(outputs.keys())}" + ) + + for key, ref_value in reference.items(): + out_value = outputs[key].detach().cpu() + assert out_value.shape == ref_value.shape, ( + f"Output {key} shape mismatch: got {tuple(out_value.shape)}, " + f"expected {tuple(ref_value.shape)}" + ) + assert torch.isfinite( + out_value + ).all(), f"Output {key} contains non-finite values" + assert torch.allclose(out_value, ref_value, atol=1e-4, rtol=1e-4), ( + f"Output {key} differs from reference: " + f"max abs diff = {(out_value - ref_value).abs().max().item()}" + ) diff --git a/tests/scene_flow/models/fixtures/raft_3d_bilaplacian_forward_reference.pt b/tests/scene_flow/models/fixtures/raft_3d_bilaplacian_forward_reference.pt new file mode 100644 index 0000000..5f8816b Binary files /dev/null and b/tests/scene_flow/models/fixtures/raft_3d_bilaplacian_forward_reference.pt differ diff --git a/tests/scene_flow/models/fixtures/raft_3d_forward_reference.pt b/tests/scene_flow/models/fixtures/raft_3d_forward_reference.pt new file mode 100644 index 0000000..220f5be Binary files /dev/null and b/tests/scene_flow/models/fixtures/raft_3d_forward_reference.pt differ diff --git a/tests/scene_flow/models/test_raft_3d_bilaplacian_validation.py b/tests/scene_flow/models/test_raft_3d_bilaplacian_validation.py new file mode 100644 index 0000000..5499857 --- /dev/null +++ b/tests/scene_flow/models/test_raft_3d_bilaplacian_validation.py @@ -0,0 +1,163 @@ +import csv +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +VALIDATE_SCRIPT = REPO_ROOT / "roco_spring_devkit" / "scene_flow" / "validate.py" +SCENE_FLOW_DIR = REPO_ROOT / "roco_spring_devkit" / "scene_flow" +FIXTURE_PATH = ( + Path(__file__).resolve().parent + / "fixtures" + / "raft_3d_bilaplacian_forward_reference.pt" +) +METRICS_MARGIN = 0.01 + +REFERENCE_METRICS = { + "spring-val-val/epe": 0.9564559029208289, + "spring-val-val/1px_flow": 16.408604602018993, + "spring-val-val/flall": 1.1024627067510866, + "spring-val-val/wauc": 81.02638138665094, + "spring-val-val/abs1": 1.871602263715532, + "spring-val-val/d1": 2.9628451002968683, + "spring-val-val/abs2": 1.8162295321623485, + "spring-val-val/d2": 2.9761543373266854, + "spring-val-val/1px_all": 33.76841746436225, + "spring-val-val/sfall": 3.626115745968289, +} + + +def _parse_metrics_csv(csv_path): + parsed = {} + with open(csv_path, "r") as f: + reader = csv.reader(f) + for row in reader: + if len(row) >= 2 and row[0] in REFERENCE_METRICS: + parsed[row[0]] = float(row[1]) + return parsed + + +@pytest.mark.slow +def test_raft_3d_bilaplacian_laplacian_validation_metrics(tmp_path): + output_path = tmp_path / "validate" + cmd = [ + sys.executable, + str(VALIDATE_SCRIPT), + "--data.val_dataset", + "spring-val", + "--model", + "raft_3d_bilaplacian", + "--ckpt_path", + "laplacian", + "--model.compute_disparity", + "true", + "--model.disparity_ckpt", + "sceneflow", + "--output_path", + str(output_path), + ] + result = subprocess.run( + cmd, + cwd=str(SCENE_FLOW_DIR), + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"validate.py failed with code {result.returncode}\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + + metrics_csv = output_path / "raft_3d_bilaplacian_laplacian" / "metrics.csv" + assert metrics_csv.exists(), f"metrics.csv not found at {metrics_csv}" + + parsed = _parse_metrics_csv(metrics_csv) + assert set(parsed.keys()) == set( + REFERENCE_METRICS.keys() + ), f"Expected metrics {set(REFERENCE_METRICS.keys())}, got {set(parsed.keys())}" + + failures = [] + for key, ref in REFERENCE_METRICS.items(): + val = parsed[key] + diff = abs(val - ref) + if diff >= METRICS_MARGIN: + failures.append( + f"Metric {key}: got {val}, expected {ref}, diff {diff} >= margin {METRICS_MARGIN}" + ) + assert not failures, "\n".join(failures) + + +@pytest.mark.slow +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="raft_3d_bilaplacian requires CUDA" +) +def test_raft_3d_bilaplacian_forward_deterministic(): + """A fixed input must always produce the same (previously recorded) outputs. + + The model weights are made deterministic by loading the pretrained "laplacian" + checkpoint together with the "sceneflow" checkpoint for the internal disparity + sub-model (manually setting all weights to a constant value was attempted for the + optical-flow RAFT counterpart, but it caused the predictions to explode, so the + checkpoint fallback is used here as well). The deterministic input and the + reference outputs were recorded once (see + fixtures/raft_3d_bilaplacian_forward_reference.pt) and are reused here. + + The registered raft_3d_bilaplacian uses a CUDA-only correlation kernel + (lietorch_extras) and a sparse Cholesky solver (sksparse.cholmod) whose setup + hardcodes CUDA tensors, so this test runs on GPU. Repeated forward passes on the + same GPU were verified to be bitwise identical (max diff 0.0), so an allclose with + a small tolerance is used to tolerate any future minor numerical differences. + """ + from jsonargparse import ArgumentParser + + from roco_spring_devkit import get_model, get_model_reference + + fixture = torch.load(FIXTURE_PATH, map_location="cpu", weights_only=False) + + model_ref = get_model_reference(fixture["model_name"]) + parser = ArgumentParser() + parser.add_class_arguments(model_ref, "model") + extra_args = [] + if fixture.get("compute_disparity"): + extra_args.append("--model.compute_disparity=true") + extra_args.append(f"--model.disparity_ckpt={fixture['disparity_ckpt']}") + args = parser.parse_args(extra_args) + + model = get_model(fixture["model_name"], fixture["ckpt_path"], args=args) + model.eval() + model = model.to("cuda") + + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + inputs = { + "images": fixture["input_images"].to("cuda"), + "images_right": fixture["input_images_right"].to("cuda"), + "intrinsics": fixture["input_intrinsics"].to("cuda"), + "baselines": fixture["input_baselines"].to("cuda"), + } + + with torch.no_grad(): + outputs = model(inputs) + + reference = fixture["outputs"] + assert set(outputs.keys()) >= set(reference.keys()), ( + f"Missing output keys: expected {set(reference.keys())}, " + f"got {set(outputs.keys())}" + ) + + for key, ref_value in reference.items(): + out_value = outputs[key].detach().cpu() + assert out_value.shape == ref_value.shape, ( + f"Output {key} shape mismatch: got {tuple(out_value.shape)}, " + f"expected {tuple(ref_value.shape)}" + ) + assert torch.isfinite( + out_value + ).all(), f"Output {key} contains non-finite values" + assert torch.allclose(out_value, ref_value, atol=1e-4, rtol=1e-4), ( + f"Output {key} differs from reference: " + f"max abs diff = {(out_value - ref_value).abs().max().item()}" + ) diff --git a/tests/scene_flow/models/test_raft_3d_validation.py b/tests/scene_flow/models/test_raft_3d_validation.py new file mode 100644 index 0000000..24fd9b7 --- /dev/null +++ b/tests/scene_flow/models/test_raft_3d_validation.py @@ -0,0 +1,158 @@ +import csv +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +VALIDATE_SCRIPT = REPO_ROOT / "roco_spring_devkit" / "scene_flow" / "validate.py" +SCENE_FLOW_DIR = REPO_ROOT / "roco_spring_devkit" / "scene_flow" +FIXTURE_PATH = ( + Path(__file__).resolve().parent / "fixtures" / "raft_3d_forward_reference.pt" +) +METRICS_MARGIN = 0.01 + +REFERENCE_METRICS = { + "spring-val-val/epe": 1.2541398290130827, + "spring-val-val/1px_flow": 18.72635042998526, + "spring-val-val/flall": 1.7256757215493255, + "spring-val-val/wauc": 81.36511908637152, + "spring-val-val/abs1": 1.871602263715532, + "spring-val-val/d1": 2.9628451002968683, + "spring-val-val/abs2": 1.8787182139025793, + "spring-val-val/d2": 3.191043616996871, + "spring-val-val/1px_all": 35.03662755754259, + "spring-val-val/sfall": 4.1250037716494665, +} + + +def _parse_metrics_csv(csv_path): + parsed = {} + with open(csv_path, "r") as f: + reader = csv.reader(f) + for row in reader: + if len(row) >= 2 and row[0] in REFERENCE_METRICS: + parsed[row[0]] = float(row[1]) + return parsed + + +@pytest.mark.slow +def test_raft_3d_base_validation_metrics(tmp_path): + output_path = tmp_path / "validate" + cmd = [ + sys.executable, + str(VALIDATE_SCRIPT), + "--data.val_dataset", + "spring-val", + "--model", + "raft_3d", + "--ckpt_path", + "base", + "--model.compute_disparity", + "true", + "--model.disparity_ckpt", + "sceneflow", + "--output_path", + str(output_path), + ] + result = subprocess.run( + cmd, + cwd=str(SCENE_FLOW_DIR), + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"validate.py failed with code {result.returncode}\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + + metrics_csv = output_path / "raft_3d_base" / "metrics.csv" + assert metrics_csv.exists(), f"metrics.csv not found at {metrics_csv}" + + parsed = _parse_metrics_csv(metrics_csv) + assert set(parsed.keys()) == set( + REFERENCE_METRICS.keys() + ), f"Expected metrics {set(REFERENCE_METRICS.keys())}, got {set(parsed.keys())}" + + failures = [] + for key, ref in REFERENCE_METRICS.items(): + val = parsed[key] + diff = abs(val - ref) + if diff >= METRICS_MARGIN: + failures.append( + f"Metric {key}: got {val}, expected {ref}, diff {diff} >= margin {METRICS_MARGIN}" + ) + assert not failures, "\n".join(failures) + + +@pytest.mark.slow +@pytest.mark.skipif(not torch.cuda.is_available(), reason="raft_3d requires CUDA") +def test_raft_3d_forward_deterministic(): + """A fixed input must always produce the same (previously recorded) outputs. + + The model weights are made deterministic by loading the pretrained "base" + checkpoint together with the "sceneflow" checkpoint for the internal disparity + sub-model (manually setting all weights to a constant value was attempted for the + optical-flow RAFT counterpart, but it caused the predictions to explode, so the + checkpoint fallback is used here as well). The deterministic input and the + reference outputs were recorded once (see fixtures/raft_3d_forward_reference.pt) + and are reused here. + + The registered raft_3d uses a CUDA-only correlation kernel (lietorch_extras), so + this test runs on GPU. Repeated forward passes on the same GPU were verified to be + bitwise identical (max diff 0.0), so an allclose with a small tolerance is used to + tolerate any future minor numerical differences. + """ + from jsonargparse import ArgumentParser + + from roco_spring_devkit import get_model, get_model_reference + + fixture = torch.load(FIXTURE_PATH, map_location="cpu", weights_only=False) + + model_ref = get_model_reference(fixture["model_name"]) + parser = ArgumentParser() + parser.add_class_arguments(model_ref, "model") + extra_args = [] + if fixture.get("compute_disparity"): + extra_args.append("--model.compute_disparity=true") + extra_args.append(f"--model.disparity_ckpt={fixture['disparity_ckpt']}") + args = parser.parse_args(extra_args) + + model = get_model(fixture["model_name"], fixture["ckpt_path"], args=args) + model.eval() + model = model.to("cuda") + + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + inputs = { + "images": fixture["input_images"].to("cuda"), + "images_right": fixture["input_images_right"].to("cuda"), + "intrinsics": fixture["input_intrinsics"].to("cuda"), + "baselines": fixture["input_baselines"].to("cuda"), + } + + with torch.no_grad(): + outputs = model(inputs) + + reference = fixture["outputs"] + assert set(outputs.keys()) >= set(reference.keys()), ( + f"Missing output keys: expected {set(reference.keys())}, " + f"got {set(outputs.keys())}" + ) + + for key, ref_value in reference.items(): + out_value = outputs[key].detach().cpu() + assert out_value.shape == ref_value.shape, ( + f"Output {key} shape mismatch: got {tuple(out_value.shape)}, " + f"expected {tuple(ref_value.shape)}" + ) + assert torch.isfinite( + out_value + ).all(), f"Output {key} contains non-finite values" + assert torch.allclose(out_value, ref_value, atol=1e-4, rtol=1e-4), ( + f"Output {key} differs from reference: " + f"max abs diff = {(out_value - ref_value).abs().max().item()}" + ) diff --git a/tests/stereo/models/fixtures/raft_stereo_forward_reference.pt b/tests/stereo/models/fixtures/raft_stereo_forward_reference.pt new file mode 100644 index 0000000..c478038 Binary files /dev/null and b/tests/stereo/models/fixtures/raft_stereo_forward_reference.pt differ diff --git a/tests/stereo/models/test_raft_stereo_validation.py b/tests/stereo/models/test_raft_stereo_validation.py new file mode 100644 index 0000000..61502ec --- /dev/null +++ b/tests/stereo/models/test_raft_stereo_validation.py @@ -0,0 +1,125 @@ +import csv +import subprocess +import sys +from pathlib import Path + +import pytest +import torch + +REPO_ROOT = Path(__file__).resolve().parents[3] +VALIDATE_SCRIPT = REPO_ROOT / "roco_spring_devkit" / "stereo" / "validate.py" +STEREO_DIR = REPO_ROOT / "roco_spring_devkit" / "stereo" +FIXTURE_PATH = ( + Path(__file__).resolve().parent / "fixtures" / "raft_stereo_forward_reference.pt" +) +METRICS_MARGIN = 0.01 + +REFERENCE_METRICS = { + "spring-val-val/abs": 1.871212733419318, + "spring-val-val/1px": 20.58077741924085, + "spring-val-val/d1": 2.982380013716848, +} + + +def _parse_metrics_csv(csv_path): + parsed = {} + with open(csv_path, "r") as f: + reader = csv.reader(f) + for row in reader: + if len(row) >= 2 and row[0] in REFERENCE_METRICS: + parsed[row[0]] = float(row[1]) + return parsed + + +@pytest.mark.slow +def test_raft_stereo_sceneflow_validation_metrics(tmp_path): + output_path = tmp_path / "validate" + cmd = [ + sys.executable, + str(VALIDATE_SCRIPT), + "--data.val_dataset", + "spring-val", + "--model", + "raft_stereo", + "--ckpt_path", + "sceneflow", + "--output_path", + str(output_path), + ] + result = subprocess.run( + cmd, + cwd=str(STEREO_DIR), + capture_output=True, + text=True, + ) + assert result.returncode == 0, ( + f"validate.py failed with code {result.returncode}\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + + metrics_csv = output_path / "raft_stereo_sceneflow" / "metrics.csv" + assert metrics_csv.exists(), f"metrics.csv not found at {metrics_csv}" + + parsed = _parse_metrics_csv(metrics_csv) + assert set(parsed.keys()) == set( + REFERENCE_METRICS.keys() + ), f"Expected metrics {set(REFERENCE_METRICS.keys())}, got {set(parsed.keys())}" + + failures = [] + for key, ref in REFERENCE_METRICS.items(): + val = parsed[key] + diff = abs(val - ref) + if diff >= METRICS_MARGIN: + failures.append( + f"Metric {key}: got {val}, expected {ref}, diff {diff} >= margin {METRICS_MARGIN}" + ) + assert not failures, "\n".join(failures) + + +@pytest.mark.slow +def test_raft_stereo_forward_deterministic(): + """A fixed input must always produce the same (previously recorded) outputs. + + The model weights are made deterministic by loading the pretrained "sceneflow" + checkpoint (manually setting all weights to a constant value was attempted for the + optical-flow RAFT counterpart, but it caused the predictions to explode, so the + checkpoint fallback is used here as well). The deterministic input and the + reference outputs were recorded once (see fixtures/raft_stereo_forward_reference.pt) + and are reused here. + """ + from roco_spring_devkit import get_model + + fixture = torch.load(FIXTURE_PATH, map_location="cpu", weights_only=False) + + model = get_model(fixture["model_name"], fixture["ckpt_path"]) + model.eval() + model.freeze_bn() + model = model.to("cpu") + + inputs = { + "images": fixture["input_images"].to("cpu"), + "images_right": fixture["input_images_right"].to("cpu"), + } + + with torch.no_grad(): + outputs = model(inputs) + + reference = fixture["outputs"] + assert set(outputs.keys()) >= set(reference.keys()), ( + f"Missing output keys: expected {set(reference.keys())}, " + f"got {set(outputs.keys())}" + ) + + for key, ref_value in reference.items(): + out_value = outputs[key].detach().cpu() + assert out_value.shape == ref_value.shape, ( + f"Output {key} shape mismatch: got {tuple(out_value.shape)}, " + f"expected {tuple(ref_value.shape)}" + ) + assert torch.isfinite( + out_value + ).all(), f"Output {key} contains non-finite values" + assert torch.allclose(out_value, ref_value, atol=1e-4, rtol=1e-4), ( + f"Output {key} differs from reference: " + f"max abs diff = {(out_value - ref_value).abs().max().item()}" + )