diff --git a/.github/workflows/lightning.yml b/.github/workflows/lightning.yml new file mode 100644 index 0000000..09b5962 --- /dev/null +++ b/.github/workflows/lightning.yml @@ -0,0 +1,38 @@ +# This workflow will install Python dependencies, run tests with a variety of Python versions +# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions + +name: lightning + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + build: + + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + lightning: ["2.1.4", "2.2.5", "2.3.3", "2.4.0", "2.5.1", "2.6.5"] + + steps: + - uses: actions/checkout@v4 + - name: Replace lightning + uses: jacobtomlinson/gha-find-replace@v3 + with: + find: "lightning[pytorch-extra]>=2,<2.7" + replace: "lightning[pytorch-extra]==${{ matrix.lightning }}" + regex: false + include: "requirements.txt" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu + pip install -r requirements.txt + - name: Test with pytest + run: | + pip install pytest==9.0.2 + python -m pytest tests/ diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml new file mode 100644 index 0000000..0a505a5 --- /dev/null +++ b/.github/workflows/python.yml @@ -0,0 +1,35 @@ +# This workflow will install Python dependencies, run tests with a variety of Python versions +# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions + +name: python + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + build: + + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.10', '3.11', '3.12', '3.13'] + + steps: + - uses: actions/checkout@v4 + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu + pip install -r requirements.txt + - name: Test with pytest + run: | + pip install pytest==9.0.2 + python -m pytest tests/ diff --git a/.github/workflows/pytorch.yml b/.github/workflows/pytorch.yml new file mode 100644 index 0000000..e7ec10a --- /dev/null +++ b/.github/workflows/pytorch.yml @@ -0,0 +1,38 @@ +# This workflow will install Python dependencies, run tests with a variety of Python versions +# For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions + +name: pytorch + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + build: + + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + pytorch: [ + 'torch==2.13.0 torchvision==0.28.0 --index-url https://download.pytorch.org/whl/cpu', + 'torch==2.12.1 torchvision==0.27.1 --index-url https://download.pytorch.org/whl/cpu', + 'torch==2.11.0 torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cpu', + 'torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cpu', + 'torch==2.9.1 torchvision==0.24.1 --index-url https://download.pytorch.org/whl/cpu', + 'torch==2.8.0 torchvision==0.23.0 --index-url https://download.pytorch.org/whl/cpu', + ] + + steps: + - uses: actions/checkout@v4 + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install ${{ matrix.pytorch }} + pip install -r requirements.txt + - name: Test with pytest + run: | + pip install pytest==9.0.2 + python -m pytest tests/ diff --git a/README.md b/README.md index 43f0940..75932f6 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ # RoCo Spring DevKit +![GitHub CI python status](https://github.com/hmorimitsu/roco-spring-devkit/actions/workflows/python.yml/badge.svg) +![GitHub CI pytorch status](https://github.com/hmorimitsu/roco-spring-devkit/actions/workflows/pytorch.yml/badge.svg) +![GitHub CI lightning status](https://github.com/hmorimitsu/roco-spring-devkit/actions/workflows/lightning.yml/badge.svg) + Developer kit for the Robust Correspondence Challenge. This devkit is derived from [PTLFlow](http://github.com/hmorimitsu/ptlflow), so checking [PTLFlow's documentation](https://ptlflow.readthedocs.io/en/latest/) may also help you to customize the code in this devkit. ## Prerequisites & Requirements diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..ba46be1 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,14 @@ +einops<0.9 +h5py<3.17 +kaleido<1.3 +lightning[pytorch-extra]>=2.1,<2.7 +loguru<0.8 +opencv-python<4.14 +pandas<3.1 +plotly<6.7 +pypng==0.20220715.0 +requests<2.34 +scipy<1.18 +tabulate<0.11 +tensorboard<2.21 +timm<1.1 \ No newline at end of file diff --git a/roco_spring_devkit/common/data/optical_flow_datasets.py b/roco_spring_devkit/common/data/optical_flow_datasets.py index 84b988f..b6de38f 100644 --- a/roco_spring_devkit/common/data/optical_flow_datasets.py +++ b/roco_spring_devkit/common/data/optical_flow_datasets.py @@ -365,7 +365,7 @@ def __init__( self.metadata = [ { "image_paths": [str(p) for p in paths], - "is_val": paths[0].stem in val_names, + "is_val": paths[0].parent.stem in val_names, "misc": "", "is_seq_start": True, "is_seq_end": True, @@ -465,7 +465,7 @@ def __init__( self.metadata = [ { "image_paths": [str(p) for p in paths], - "is_val": paths[0].stem in val_names, + "is_val": paths[0].stem.split("_")[0] in val_names, "misc": "", "is_seq_start": True, "is_seq_end": True, @@ -1085,12 +1085,15 @@ def __init__( # noqa: C901 ) flow_paths = sorted(flow_dir.glob("*.flo"), reverse=rev) - # Create groups to separate different sequences + # Create groups to separate different sequences. The + # consecutive-index check must use abs() so that + # descending-sorted (rev=True) flow lists are grouped + # correctly too. flow_groups_paths = [[flow_paths[0]]] prev_idx = int(flow_paths[0].stem) for path in flow_paths[1:]: idx = int(path.stem) - if (idx - 1) == prev_idx: + if abs(idx - prev_idx) == 1: flow_groups_paths[-1].append(path) else: flow_groups_paths.append([path]) @@ -1123,7 +1126,7 @@ def __init__( # noqa: C901 img_dir / (fp.stem + ".png") for fp in flow_paths ] if rev: - idx = int(img_paths[0].stem) - 1 + idx = int(img_paths[-1].stem) - 1 else: idx = int(img_paths[-1].stem) + 1 img_paths.append(img_dir / f"{idx:07d}.png") @@ -1171,12 +1174,14 @@ def __init__( # noqa: C901 ) flow_paths = sorted(flow_dir.glob("*.flo"), reverse=rev) - # Create groups to separate different sequences + # Create groups to separate different sequences. Same abs() + # reasoning as the forward-direction grouping + # above - descending lists need it too. flow_groups_paths = [[flow_paths[0]]] prev_idx = int(flow_paths[0].stem) for path in flow_paths[1:]: idx = int(path.stem) - if (idx - 1) == prev_idx: + if abs(idx - prev_idx) == 1: flow_groups_paths[-1].append(path) else: flow_groups_paths.append([path]) @@ -2343,7 +2348,6 @@ def __init__( # noqa: C901 get_occlusion_mask=False, get_motion_boundary_mask=False, get_backward=get_backward, - get_semantic_segmentation_labels=False, get_meta=get_meta, ) self.root_dir = root_dir diff --git a/roco_spring_devkit/common/data/scene_flow_datasets.py b/roco_spring_devkit/common/data/scene_flow_datasets.py index c21c32a..43d6df8 100644 --- a/roco_spring_devkit/common/data/scene_flow_datasets.py +++ b/roco_spring_devkit/common/data/scene_flow_datasets.py @@ -665,7 +665,7 @@ def __init__( "is_seq_start": i == 0, "is_seq_end": i + step_size >= len(image_paths) - 1, - "is_time_reverse": False, + "is_time_reverse": reverse, "is_camera_reverse": False, } ) @@ -1415,6 +1415,7 @@ def __init__( # noqa: C901 [p.stem for p in (Path(root_dir) / split_dir).glob("*")] ) + val_seqs = [] if split == "train" or split == "val": split_file = THIS_DIR / "Spring_val.txt" with open(split_file, "r") as f: @@ -1726,7 +1727,7 @@ def __init__( # noqa: C901 i : i + self.sequence_length ] ], - "is_val": False, + "is_val": seq_name in val_seqs, "is_time_reverse": time_rev, "is_camera_reverse": cam_rev, "misc": seq_name, diff --git a/roco_spring_devkit/common/utils/external/selflow.py b/roco_spring_devkit/common/utils/external/selflow.py index c017490..ac84824 100644 --- a/roco_spring_devkit/common/utils/external/selflow.py +++ b/roco_spring_devkit/common/utils/external/selflow.py @@ -51,5 +51,5 @@ def write_pfm(output_path, flow, scale=1): if len(flow.shape) == 3: invalid = np.isnan(flow[..., 0]) | np.isnan(flow[..., 1]) flow = np.dstack([flow, invalid.astype(np.float32)]) - flow = np.flipud(flow) + flow = np.flipud(flow) flow.tofile(file) diff --git a/roco_spring_devkit/common/utils/flow_utils.py b/roco_spring_devkit/common/utils/flow_utils.py index c0e7f5b..0f7a469 100644 --- a/roco_spring_devkit/common/utils/flow_utils.py +++ b/roco_spring_devkit/common/utils/flow_utils.py @@ -266,7 +266,6 @@ def spring_epe_to_rgb( epe_rgb = cv.applyColorMap(epe, plt_lut) invalid_mask = ~valid_mask - invalid_mask = np.concatenate([invalid_mask, invalid_mask, invalid_mask], -1) epe_rgb[invalid_mask] = 0 return epe_rgb diff --git a/roco_spring_devkit/common/utils/io_adapter.py b/roco_spring_devkit/common/utils/io_adapter.py index e3435b6..5ac7166 100644 --- a/roco_spring_devkit/common/utils/io_adapter.py +++ b/roco_spring_devkit/common/utils/io_adapter.py @@ -126,11 +126,22 @@ def prepare_inputs( inputs.update(kwargs) keys_to_remove = [] for k, v in inputs.items(): - if v is None or len(v) == 0: + if v is None: keys_to_remove.append(k) + elif hasattr(v, "__len__") and len(v) == 0: + keys_to_remove.append(k) + else: + if not isinstance(v, (np.ndarray, list, tuple)): + inputs[k] = v for k in keys_to_remove: del inputs[k] - inputs = self.transform(inputs) + to_transform = { + k: v + for k, v in inputs.items() + if isinstance(v, (np.ndarray, list, tuple)) + } + transformed = self.transform(to_transform) + inputs.update(transformed) inputs = self._to_cuda(inputs) diff --git a/roco_spring_devkit/common/utils/scene_flow_metrics.py b/roco_spring_devkit/common/utils/scene_flow_metrics.py index d18b906..26d2cc0 100644 --- a/roco_spring_devkit/common/utils/scene_flow_metrics.py +++ b/roco_spring_devkit/common/utils/scene_flow_metrics.py @@ -171,6 +171,8 @@ def update( else: valid_flows_target = torch.ones_like(flow_target[:, :1]) valid_flows_target = valid_flows_target[:, 0] + if valid_flows_target.dim() > 3: + valid_flows_target = valid_flows_target[:, 0] valid_disparities_target = targets.get("valid_disparities") if valid_disparities_target is not None: @@ -181,10 +183,10 @@ def update( valid_disparities_target = torch.ones_like(disp_target[:, :1]) valid_disparities_target = valid_disparities_target[:, 0] - disp1_target = disp_target[:1] - valid_disp1_target = valid_disparities_target[:1] - disp2_target = disp_target[1:2] - valid_disp2_target = valid_disparities_target[1:2] + disp1_target = disp_target[0::2] + valid_disp1_target = valid_disparities_target[0::2] + disp2_target = disp_target[1::2] + valid_disp2_target = valid_disparities_target[1::2] valid_all_target = valid_flows_target * valid_disp1_target * valid_disp2_target @@ -194,13 +196,8 @@ def update( target_norm = torch.norm(flow_target, p=2, dim=2) target_norm = target_norm.gather(1, min_idx[:, None])[:, 0] - abs1 = torch.abs(disp1_pred[:, None] - disp1_target) - abs1, min_idx = abs1.min(dim=1) - abs1 = abs1.gather(1, min_idx[:, None]) - - abs2 = torch.abs(disp2_pred[:, None] - disp2_target) - abs2, min_idx = abs2.min(dim=1) - abs2 = abs2.gather(1, min_idx[:, None]) + abs1 = torch.abs(disp1_pred - disp1_target)[:, 0] + abs2 = torch.abs(disp2_pred - disp2_target)[:, 0] else: epe_flow = torch.norm(flow_pred - flow_target, p=2, dim=1) target_norm = torch.norm(flow_target, p=2, dim=1) @@ -208,19 +205,22 @@ def update( abs1 = torch.abs(disp1_pred - disp1_target)[:, 0] abs2 = torch.abs(disp2_pred - disp2_target)[:, 0] + disp1_target_mag = torch.abs(disp1_target)[:, 0] + disp2_target_mag = torch.abs(disp2_target)[:, 0] + px1_flow_mask = (epe_flow > 1).float() * 100 flall_mask = ((epe_flow > 3) & (epe_flow > (0.05 * target_norm))).float() * 100 px11_mask = (abs1 > 1).float() * 100 - d1_mask = ((abs1 > 3) & (abs1 > (0.05 * target_norm))).float() * 100 + d1_mask = ((abs1 > 3) & (abs1 > (0.05 * disp1_target_mag))).float() * 100 px12_mask = (abs2 > 1).float() * 100 - d2_mask = ((abs2 > 3) & (abs2 > (0.05 * target_norm))).float() * 100 + d2_mask = ((abs2 > 3) & (abs2 > (0.05 * disp2_target_mag))).float() * 100 px1_all_mask = ( - (px1_flow_mask / 100) * (px11_mask / 100) * (px12_mask / 100) * 100 - ) - sfall_mask = (flall_mask / 100) * (d1_mask / 100) * (d2_mask / 100) * 100 + (px1_flow_mask > 0) | (px11_mask > 0) | (px12_mask > 0) + ).float() * 100 + sfall_mask = ((flall_mask > 0) | (d1_mask > 0) | (d2_mask > 0)).float() * 100 self.used_keys = [ ("epe", "epe_flow", "valid_flows_target"), diff --git a/roco_spring_devkit/common/utils/stereo_metrics.py b/roco_spring_devkit/common/utils/stereo_metrics.py index 386b5e2..e406c66 100644 --- a/roco_spring_devkit/common/utils/stereo_metrics.py +++ b/roco_spring_devkit/common/utils/stereo_metrics.py @@ -174,13 +174,15 @@ def update( if len(disp_target.shape) == 5: abs = torch.abs(disp_pred[:, None] - disp_target) - abs, min_idx = abs.min(dim=1) - abs = abs.gather(1, min_idx[:, None])[:, 0] + abs, _ = abs.min(dim=1) + abs = abs[:, 0] + disp_target_mag = torch.abs(disp_target).max(dim=1)[0][:, 0] else: - abs = torch.abs(disp_pred - disp_target) + abs = torch.abs(disp_pred - disp_target)[:, 0] + disp_target_mag = torch.abs(disp_target)[:, 0] px1_mask = (abs > 1).float() * 100 - d1_mask = ((abs > 3) & (abs > (0.05 * abs))).float() * 100 + d1_mask = ((abs > 3) & (abs > (0.05 * disp_target_mag))).float() * 100 self.used_keys = [ ("abs", "abs", "valid_target"), ("1px", "px1_mask", "valid_target"), diff --git a/roco_spring_devkit/common/utils/stereo_utils.py b/roco_spring_devkit/common/utils/stereo_utils.py index 6b348bc..d22f5bc 100644 --- a/roco_spring_devkit/common/utils/stereo_utils.py +++ b/roco_spring_devkit/common/utils/stereo_utils.py @@ -126,7 +126,9 @@ def disparity_to_rgb( elif len(input_shape) == 3: disparity_rgb = disparity_rgb[0] else: - disparity_rgb = disparity_rgb.reshape(*input_shape) + output_shape = list(input_shape) + output_shape[-3] = 3 + disparity_rgb = disparity_rgb.reshape(output_shape) elif len(disparity) == 1: disparity_rgb = disparity_rgb[0] @@ -446,7 +448,6 @@ def spring_abs_to_rgb( epe_rgb = cv.applyColorMap(epe, plt_lut) invalid_mask = ~valid_mask - invalid_mask = np.concatenate([invalid_mask, invalid_mask, invalid_mask], -1) epe_rgb[invalid_mask] = 0 return epe_rgb diff --git a/roco_spring_devkit/common/utils/utils.py b/roco_spring_devkit/common/utils/utils.py index 5722f12..95f7fc6 100644 --- a/roco_spring_devkit/common/utils/utils.py +++ b/roco_spring_devkit/common/utils/utils.py @@ -116,11 +116,19 @@ def _forward_warp_single_torch( return out[0] if squeeze_channel else out -def _forward_warp_single_scipy( +def _forward_warp_single_numpy( flow: torch.Tensor, tensor: torch.Tensor ) -> torch.Tensor: - from scipy import interpolate - + """Reference numpy implementation of forward warping. + + Mirrors :func:`_forward_warp_single_torch`: each source pixel is moved + to the rounded integer destination ``round(src + flow)``, destinations + are kept only when they lie inside the image (inclusive bounds), and + when multiple sources collide on the same destination the one with the + smallest ``distance + tie_break`` is kept (``tie_break = flat_index * + 1e-6`` so the lower source index wins ties). Destinations with no + mapping source stay at zero. + """ tensor_dtype = tensor.dtype tensor_device = tensor.device @@ -140,87 +148,71 @@ def _forward_warp_single_scipy( if tensor_np.shape[-2:] != (ht, wd): raise ValueError("flow and tensor spatial dimensions must match.") - x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) + y0, x0 = np.meshgrid(np.arange(ht), np.arange(wd), indexing="ij") + x0 = x0.astype(flow_np.dtype) + y0 = y0.astype(flow_np.dtype) x1 = x0 + dx y1 = y0 + dy - x1 = x1.reshape(-1) - y1 = y1.reshape(-1) - values = tensor_np.reshape(tensor_np.shape[0], -1) - valid = ( - np.isfinite(x1) & np.isfinite(y1) & (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) + np.isfinite(x1) + & np.isfinite(y1) + & (x1 >= 0) + & (x1 <= wd - 1) + & (y1 >= 0) + & (y1 <= ht - 1) ) + warped = np.zeros_like(tensor_np) if not np.any(valid): - warped = np.zeros_like(tensor_np) - else: - warped = [ - interpolate.griddata( - (x1[valid], y1[valid]), - values[channel][valid], - (x0, y0), - method="nearest", - fill_value=0, - ) - for channel in range(values.shape[0]) - ] - warped = np.stack(warped, axis=0) - - if squeeze_channel: - warped = warped[0] - - return torch.from_numpy(warped).to(device=tensor_device, dtype=tensor_dtype) - - -def _forward_warp_single_ckdtree( - flow: torch.Tensor, tensor: torch.Tensor -) -> torch.Tensor: - from scipy.spatial import cKDTree - - tensor_dtype = tensor.dtype - tensor_device = tensor.device + if squeeze_channel: + warped = warped[0] + return torch.from_numpy(np.ascontiguousarray(warped)).to( + device=tensor_device, dtype=tensor_dtype + ) - flow_np = flow.detach().cpu().numpy() - tensor_np = tensor.detach().cpu().numpy() + x1_valid = x1[valid] + y1_valid = y1[valid] + x1_round = np.rint(x1_valid).astype(np.int64) + y1_round = np.rint(y1_valid).astype(np.int64) - dx, dy = flow_np[0], flow_np[1] + flat_idx = y1_round * wd + x1_round + src_values = tensor_np.reshape(tensor_np.shape[0], -1)[:, valid.reshape(-1)] - squeeze_channel = False - if tensor_np.ndim == 2: - tensor_np = tensor_np[None] - squeeze_channel = True - elif tensor_np.ndim != 3: - raise ValueError("tensor must have shape [H, W] or [C, H, W].") + dist = (x1_valid - x1_round.astype(x1_valid.dtype)) ** 2 + ( + y1_valid - y1_round.astype(y1_valid.dtype) + ) ** 2 + tie_break = np.arange(flat_idx.size, dtype=flow_np.dtype) * 1e-6 + key = dist + tie_break - ht, wd = dx.shape - if tensor_np.shape[-2:] != (ht, wd): - raise ValueError("flow and tensor spatial dimensions must match.") + # For each destination pixel keep the source with the minimum key. + best_key = np.full(ht * wd, np.inf, dtype=flow_np.dtype) + np.minimum.at(best_key, flat_idx, key) + keep = np.isclose(key, best_key[flat_idx], atol=1e-7, rtol=0.0) - x0, y0 = np.meshgrid(np.arange(wd), np.arange(ht)) + out_flat = np.zeros((tensor_np.shape[0], ht * wd), dtype=tensor_np.dtype) + out_flat[:, flat_idx[keep]] = src_values[:, keep] + warped = out_flat.reshape(tensor_np.shape) - x1 = (x0 + dx).reshape(-1) - y1 = (y0 + dy).reshape(-1) - values = tensor_np.reshape(tensor_np.shape[0], -1) + if squeeze_channel: + warped = warped[0] - valid = ( - np.isfinite(x1) & np.isfinite(y1) & (x1 > 0) & (x1 < wd) & (y1 > 0) & (y1 < ht) + return torch.from_numpy(np.ascontiguousarray(warped)).to( + device=tensor_device, dtype=tensor_dtype ) - if not np.any(valid): - warped = np.zeros_like(tensor_np) - else: - query_points = np.column_stack((x0.reshape(-1), y0.reshape(-1))) - source_points = np.column_stack((x1[valid], y1[valid])) - nn_tree = cKDTree(source_points) - _, nn_indices = nn_tree.query(query_points, k=1, workers=-1) - warped = values[:, valid][:, nn_indices].reshape(tensor_np.shape[0], ht, wd) - if squeeze_channel: - warped = warped[0] +def _forward_warp_single_scipy( + flow: torch.Tensor, tensor: torch.Tensor +) -> torch.Tensor: + return _forward_warp_single_numpy(flow, tensor) + - return torch.from_numpy(warped).to(device=tensor_device, dtype=tensor_dtype) +def _forward_warp_single_ckdtree( + flow: torch.Tensor, tensor: torch.Tensor +) -> torch.Tensor: + return _forward_warp_single_numpy(flow, tensor) class InputPadder(_InputPadder): @@ -533,12 +525,22 @@ def tensor_dict_to_numpy( if isinstance(v, torch.Tensor): v = v.detach().cpu() if padder is not None: - v = padder.unpad(v) + # Use `unfill` (which only unpads when the shape matches + # `tgt_size`) instead of `unpad` so that an already-unpadded + # tensor is left untouched rather than sliced into an empty + # array. + v = padder.unfill(v) while len(v.shape) > 3: v = v[0] - v = v.permute(1, 2, 0).numpy() + if v.ndim >= 3: + # CHW -> HWC. + v = v.permute(1, 2, 0).numpy() + else: + # 1D / 2D tensors have no channel axis to move, so just + # convert them as-is. + v = v.numpy() npy_dict[k] = v return npy_dict diff --git a/tests/common/data/test_optical_flow_datasets.py b/tests/common/data/test_optical_flow_datasets.py new file mode 100644 index 0000000..18b34fa --- /dev/null +++ b/tests/common/data/test_optical_flow_datasets.py @@ -0,0 +1,1380 @@ +"""Unit tests for `roco_spring_devkit.common.data.optical_flow_datasets`. + +The tests here exercise two layers of behaviour: + +1. The "plumbing" methods on `BaseFlowDataset` (path-list extension, flow + reading / valid-mask generation, `__getitem__`). +2. The concrete dataset classes - the on-disk directory structures are + synthesised under ``tmp_path`` and the path / metadata lists produced by + the constructors are compared against values that were computed by hand + (not against the values produced by the implementation itself). + +Bugs uncovered while writing the tests were fixed in the source code. Each fix +is documented inline in the corresponding test case. +""" + +import json +import math +from pathlib import Path + +import cv2 as cv +import numpy as np +import pytest + +from roco_spring_devkit.common.data import optical_flow_datasets as ds +from roco_spring_devkit.common.utils import flow_utils + + +# --------------------------------------------------------------------------- +# Small helpers used to create synthetic on-disk datasets. +# --------------------------------------------------------------------------- +def _write_flow(path: Path, flow: np.ndarray, fmt: str = None) -> None: + """Write ``flow`` (HWC, float32) to ``path`` in ``fmt`` format.""" + path.parent.mkdir(parents=True, exist_ok=True) + flow_utils.flow_write(path, flow, format=fmt) + + +def _write_gray_png(path: Path, value: int, size=(4, 5, 1)) -> None: + """Write a grayscale PNG with a constant value.""" + path.parent.mkdir(parents=True, exist_ok=True) + img = np.full(size, value, dtype=np.uint8) + cv.imwrite(str(path), img) + + +def _write_rgb_png(path: Path, color=(0, 0, 0), size=(4, 5, 3)) -> None: + """Write a constant-color RGB PNG image.""" + path.parent.mkdir(parents=True, exist_ok=True) + img = np.full(size, color, dtype=np.uint8) + cv.imwrite(str(path), img) + + +def _touch(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + + +# =========================================================================== +# Tests for BaseFlowDataset._extend_paths_list +# =========================================================================== +class TestExtendPathsList: + """``_extend_paths_list`` should pad the path list at the beginning / end + depending on the requested position of the main frame inside the + sequence.""" + + def _base(self): + return ds.BaseFlowDataset("test") + + def test_first_position(self): + # 'first' pads the end with (seq_length - 2) copies of the last element + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="first" + ) + # begin_pad=0, end_pad=2 -> [img0, img1, img2, img2, img2] + assert out == ["img0", "img1", "img2", "img2", "img2"] + + def test_last_position(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="last" + ) + # begin_pad=2, end_pad=0 -> [img0, img0, img0, img1, img2] + assert out == ["img0", "img0", "img0", "img1", "img2"] + + def test_middle_position_odd(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=5, sequence_position="middle" + ) + # begin_pad = 5//2 = 2, end_pad = ceil(5/2) - 2 = 3 - 2 = 1 + # -> [img0, img0, img0, img1, img2, img2] + assert out == ["img0", "img0", "img0", "img1", "img2", "img2"] + + def test_middle_position_even(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="middle" + ) + # begin_pad = 4//2 = 2, end_pad = ceil(4/2) - 2 = 2 - 2 = 0 + # -> [img0, img0, img0, img1, img2] + assert out == ["img0", "img0", "img0", "img1", "img2"] + + def test_middle_position_seq2(self): + # seq_length=2 gives end_pad = ceil(1.0) - 2 = -1, which becomes an + # empty range, so only begin_pad (2//2 = 1) is prepended. + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=2, sequence_position="middle" + ) + assert out == ["img0", "img0", "img1", "img2"] + + def test_all_position(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="all" + ) + # 'all' should never pad + assert out == ["img0", "img1", "img2"] + + def test_invalid_position_raises(self): + d = self._base() + with pytest.raises(ValueError): + d._extend_paths_list(["a"], sequence_length=2, sequence_position="weird") + + def test_mutates_input_list_inplace(self): + # _extend_paths_list modifies the list in place and returns it. This + # documents that side-effecting behaviour. + d = self._base() + original = ["a", "b", "c"] + result = d._extend_paths_list( + original, sequence_length=4, sequence_position="first" + ) + assert result is original + assert original == ["a", "b", "c", "c", "c"] + + +# =========================================================================== +# Tests for BaseFlowDataset._get_flows_and_valids +# =========================================================================== +class TestGetFlowsAndValids: + """Verify the flow reading / valid mask / clipping / NaN handling logic + against a hand-computed expected output.""" + + def _make_flow(self): + # Layout is H=2, W=3, channels=2 (x, y). + flow = np.array( + [ + [[10.0, 20.0], [200.0, 5.0], [-5.0, 80.0]], + [[np.nan, 0.0], [50.0, 50.0], [0.0, 0.0]], + ], + dtype=np.float32, + ) + return flow + + def test_valid_mask_and_clipping(self, tmp_path): + flow = self._make_flow() + path = tmp_path / "f.flo" + _write_flow(path, flow) + + d = ds.BaseFlowDataset("t", max_flow=100.0, get_valid_mask=True) + flows, valids = d._get_flows_and_valids([path]) + + # Expected flow after NaN replacement (0) and clipping to [-100, 100]: + expected_flow = np.array( + [ + [[10.0, 20.0], [100.0, 5.0], [-5.0, 80.0]], + [[0.0, 0.0], [50.0, 50.0], [0.0, 0.0]], + ], + dtype=np.float32, + ) + # Expected valid mask (255 = valid, 0 = invalid). A pixel is invalid if + # either channel has |value| >= max_flow (NaNs were temporarily set to + # max_flow+1, so they are also invalid). + expected_valid = np.array( + [[255, 0, 255], [0, 255, 255]], + dtype=np.uint8, + )[:, :, None] + + np.testing.assert_allclose(flows[0], expected_flow, rtol=1e-5) + assert valids[0].shape == (2, 3, 1) + np.testing.assert_array_equal(valids[0], expected_valid) + + def test_valid_mask_disabled(self, tmp_path): + flow = self._make_flow() + path = tmp_path / "f.flo" + _write_flow(path, flow) + + d = ds.BaseFlowDataset("t", max_flow=100.0, get_valid_mask=False) + flows, valids = d._get_flows_and_valids([path]) + + # When get_valid_mask is False, valids should be empty but flow still + # produced. + assert valids == [] + assert len(flows) == 1 + + def test_multiple_flows(self, tmp_path): + # Two flow files. Second one is constant (10, -7) and entirely valid. + flow1 = self._make_flow() + flow2 = np.full((2, 3, 2), 0.0, dtype=np.float32) + flow2[..., 0] = 10.0 + flow2[..., 1] = -7.0 + p1 = tmp_path / "f1.flo" + p2 = tmp_path / "f2.flo" + _write_flow(p1, flow1) + _write_flow(p2, flow2) + + d = ds.BaseFlowDataset("t", max_flow=100.0, get_valid_mask=True) + flows, valids = d._get_flows_and_valids([p1, p2]) + + assert len(flows) == 2 + assert len(valids) == 2 + # All of flow2 is valid (255). + np.testing.assert_array_equal( + valids[1], np.full((2, 3, 1), 255, dtype=np.uint8) + ) + np.testing.assert_allclose(flows[1], flow2) + + def test_two_file_flow_mode(self, tmp_path): + # When is_two_file_flow=True, the flow is constructed from two files + # (e.g. disparity files), each contributing one channel, and the values + # are negated before stacking. + disp_x = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + disp_y = np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32) + p_x = tmp_path / "disp0.pfm" + p_y = tmp_path / "disp0y.pfm" + _write_flow(p_x, disp_x, fmt="pfm") + _write_flow(p_y, disp_y, fmt="pfm") + + d = ds.BaseFlowDataset("t", max_flow=100.0, get_valid_mask=True) + d.is_two_file_flow = True + + flows, valids = d._get_flows_and_valids([[p_x, p_y]]) + + # Each file is read as a 2D array, negated, and stacked along the last + # axis to produce a (H, W, 2) flow. + expected_flow = np.stack([-disp_x, -disp_y], axis=2) + np.testing.assert_allclose(flows[0], expected_flow) + # All values are below max_flow so the entire mask should be 255. + np.testing.assert_array_equal( + valids[0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + + +# =========================================================================== +# Tests for BaseFlowDataset.__getitem__ +# =========================================================================== +class TestBaseFlowDatasetGetItem: + """Verify the dict produced by ``__getitem__`` against hand-computed + content for images, flows, occlusion masks, motion boundary masks and + metadata.""" + + def _make_simple_dataset( + self, tmp_path, with_backward=False, with_occ=True, with_mb=True + ): + # Create two images, one forward flow, one occ mask, one mb mask and + # (optionally) backward equivalents. + img1_path = tmp_path / "im0.png" + img2_path = tmp_path / "im1.png" + _write_rgb_png(img1_path, color=(10, 20, 30)) + _write_rgb_png(img2_path, color=(40, 50, 60)) + + flow_path = tmp_path / "flow.flo" + flow = np.array( + [[[1.5, 2.5], [3.5, 4.5]], [[5.5, 6.5], [7.5, 8.5]]], + dtype=np.float32, + ) + _write_flow(flow_path, flow) + + occ_path = tmp_path / "occ.png" + _write_gray_png(occ_path, value=128) + + mb_path = tmp_path / "mb.png" + _write_gray_png(mb_path, value=64) + + d = ds.BaseFlowDataset( + dataset_name="TestDS", + split_name="train", + max_flow=1000.0, + get_valid_mask=True, + get_occlusion_mask=with_occ, + get_motion_boundary_mask=with_mb, + get_backward=with_backward, + get_meta=True, + ) + d.img_paths = [[img1_path, img2_path]] + d.flow_paths = [[flow_path]] + d.occ_paths = [[occ_path]] + d.mb_paths = [[mb_path]] + if with_backward: + flow_b_path = tmp_path / "flow_b.flo" + _write_flow(flow_b_path, -flow) + occ_b_path = tmp_path / "occ_b.png" + _write_gray_png(occ_b_path, value=200) + mb_b_path = tmp_path / "mb_b.png" + _write_gray_png(mb_b_path, value=32) + d.flow_b_paths = [[flow_b_path]] + d.occ_b_paths = [[occ_b_path]] + d.mb_b_paths = [[mb_b_path]] + + d.metadata = [{"image_paths": [str(img1_path), str(img2_path)], "misc": "abc"}] + return d, flow + + def test_images_and_flows_content(self, tmp_path): + d, flow = self._make_simple_dataset(tmp_path) + out = d[0] + + assert "images" in out and len(out["images"]) == 2 + # cv.imread keeps the BGR channel order, so values are returned as + # stored on disk. + np.testing.assert_array_equal( + out["images"][0], np.full((4, 5, 3), (10, 20, 30), dtype=np.uint8) + ) + np.testing.assert_array_equal( + out["images"][1], np.full((4, 5, 3), (40, 50, 60), dtype=np.uint8) + ) + + # Flow is read unchanged (values below max_flow) and has shape (2,2,2). + np.testing.assert_allclose(out["flows"][0], flow, rtol=1e-5) + # Valid mask should be all 255 since all values are well below max_flow. + np.testing.assert_array_equal( + out["valid_flows"][0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + + def test_occlusion_and_mb_masks_content(self, tmp_path): + d, _ = self._make_simple_dataset(tmp_path) + out = d[0] + + # occ / mb are read with cv.imread(..., 0) (grayscale) and a new axis + # is added, making them (H, W, 1). + assert "occs" in out + np.testing.assert_array_equal( + out["occs"][0], np.full((4, 5, 1), 128, dtype=np.uint8) + ) + assert "mbs" in out + np.testing.assert_array_equal( + out["mbs"][0], np.full((4, 5, 1), 64, dtype=np.uint8) + ) + + def test_backward_content(self, tmp_path): + d, flow = self._make_simple_dataset(tmp_path, with_backward=True) + out = d[0] + + np.testing.assert_allclose(out["flows_b"][0], -flow, rtol=1e-5) + np.testing.assert_array_equal( + out["valid_flows_b"][0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + np.testing.assert_array_equal( + out["occs_b"][0], np.full((4, 5, 1), 200, dtype=np.uint8) + ) + np.testing.assert_array_equal( + out["mbs_b"][0], np.full((4, 5, 1), 32, dtype=np.uint8) + ) + + def test_meta_content(self, tmp_path): + d, _ = self._make_simple_dataset(tmp_path) + out = d[0] + meta = out["meta"] + assert meta["dataset_name"] == "TestDS" + assert meta["split_name"] == "train" + assert meta["misc"] == "abc" + assert meta["image_paths"] == [str(d.img_paths[0][0]), str(d.img_paths[0][1])] + + def test_len_returns_img_paths_count(self, tmp_path): + d, _ = self._make_simple_dataset(tmp_path) + assert len(d) == 1 + + +# =========================================================================== +# Tests for AutoFlowDataset +# =========================================================================== +class TestAutoFlowDataset: + """Synthesise a tiny AutoFlow directory and check the resulting lists.""" + + def _make(self, tmp_path, sample_names): + # AutoFlow stores samples under four "static_40k_png_*_of_4" folders. + # Each sample is a directory containing im0.png, im1.png and + # forward.flo. + for name in sample_names: + sample_dir = tmp_path / "static_40k_png_1_of_4" / name + _touch(sample_dir / "im0.png") + _touch(sample_dir / "im1.png") + _touch(sample_dir / "forward.flo") + + def test_trainval_keeps_all_samples(self, tmp_path): + # Use names that are NOT in AutoFlow_val.txt so split logic doesn't + # accidentally drop them in the val case below. + names = ["my_sample_a", "my_sample_b"] + self._make(tmp_path, names) + + d = ds.AutoFlowDataset(str(tmp_path), split="trainval", get_meta=True) + + assert len(d.img_paths) == 2 + assert len(d.flow_paths) == 2 + # Each img entry is a pair [im0.png, im1.png] inside the sample dir. + expected_imgs = [ + [ + tmp_path / "static_40k_png_1_of_4" / n / "im0.png", + tmp_path / "static_40k_png_1_of_4" / n / "im1.png", + ] + for n in names + ] + assert d.img_paths == expected_imgs + expected_flows = [ + [tmp_path / "static_40k_png_1_of_4" / n / "forward.flo"] for n in names + ] + assert d.flow_paths == expected_flows + # metadata sanity + assert d.metadata[0]["is_seq_start"] is True + assert d.metadata[0]["is_seq_end"] is True + assert all( + m["image_paths"] == [str(p) for p in ip] + for m, ip in zip(d.metadata, d.img_paths) + ) + + def test_train_split_excludes_val_names(self, tmp_path): + # 'table_0_batch_1' is one of the names listed in AutoFlow_val.txt. + names = ["table_0_batch_1", "my_sample_a"] + self._make(tmp_path, names) + + d = ds.AutoFlowDataset(str(tmp_path), split="train") + # Only the sample whose name is NOT in the val file should remain. + kept_stems = [p[0].parent.stem for p in d.img_paths] + assert kept_stems == ["my_sample_a"] + + def test_val_split_keeps_only_val_names(self, tmp_path): + names = ["table_0_batch_1", "my_sample_a"] + self._make(tmp_path, names) + + d = ds.AutoFlowDataset(str(tmp_path), split="val") + kept_stems = [p[0].parent.stem for p in d.img_paths] + assert kept_stems == ["table_0_batch_1"] + # is_val flag should be True for the kept val sample. + assert d.metadata[0]["is_val"] is True + + +# =========================================================================== +# Tests for FlyingChairsDataset +# =========================================================================== +class TestFlyingChairsDataset: + def _make(self, tmp_path, prefixes): + data_dir = tmp_path / "data" + for pfx in prefixes: + _touch(data_dir / f"{pfx}_img1.ppm") + _touch(data_dir / f"{pfx}_img2.ppm") + _touch(data_dir / f"{pfx}_flow.flo") + + def test_trainval_keeps_all(self, tmp_path): + # Use prefixes that are not in FlyingChairs_val.txt so trainval / train + # behaviour is straightforward. + prefixes = ["myA", "myB", "myC"] + self._make(tmp_path, prefixes) + + d = ds.FlyingChairsDataset(str(tmp_path), split="trainval") + + assert len(d.img_paths) == 3 + assert len(d.flow_paths) == 3 + # Image pair should match by sorted order. + assert d.img_paths[0] == [ + tmp_path / "data" / "myA_img1.ppm", + tmp_path / "data" / "myA_img2.ppm", + ] + assert d.flow_paths[0] == [tmp_path / "data" / "myA_flow.flo"] + + def test_train_split_excludes_val_prefix(self, tmp_path): + # '00006' is the first entry in FlyingChairs_val.txt. + prefixes = ["00006", "myA"] + self._make(tmp_path, prefixes) + + d = ds.FlyingChairsDataset(str(tmp_path), split="train") + kept_prefixes = [p[0].stem.split("_")[0] for p in d.img_paths] + assert kept_prefixes == ["myA"] + + def test_val_split_keeps_only_val_prefix(self, tmp_path): + prefixes = ["00006", "myA"] + self._make(tmp_path, prefixes) + + d = ds.FlyingChairsDataset(str(tmp_path), split="val") + kept_prefixes = [p[0].stem.split("_")[0] for p in d.img_paths] + assert kept_prefixes == ["00006"] + assert d.metadata[0]["is_val"] is True + + +# =========================================================================== +# Tests for FlyingChairs2Dataset +# =========================================================================== +class TestFlyingChairs2Dataset: + """FlyingChairs2 stores per-sample pairs of images plus forward / backward + flow, occlusion and motion-boundary masks inside ``train``/``val`` dirs.""" + + def _make(self, tmp_path, split_dir, sample_names): + d = tmp_path / split_dir + for n in sample_names: + _touch(d / f"{n}_img_0.png") + _touch(d / f"{n}_img_1.png") + _touch(d / f"{n}_flow_01.flo") + _touch(d / f"{n}_flow_10.flo") + _touch(d / f"{n}_occ_01.png") + _touch(d / f"{n}_occ_10.png") + _touch(d / f"{n}_mb_01.png") + _touch(d / f"{n}_mb_10.png") + + def test_basic_forward_only(self, tmp_path): + names = ["a", "b"] + self._make(tmp_path, "train", names) + + d = ds.FlyingChairs2Dataset( + str(tmp_path), + split="train", + add_reverse=False, + get_backward=False, + ) + + assert len(d.img_paths) == 2 + assert len(d.flow_paths) == 2 + assert len(d.occ_paths) == 2 + assert len(d.mb_paths) == 2 + # Image pair ordering for sample 'a': + assert d.img_paths[0] == [ + tmp_path / "train" / "a_img_0.png", + tmp_path / "train" / "a_img_1.png", + ] + # Forward flow is flow_01 + assert d.flow_paths[0] == [tmp_path / "train" / "a_flow_01.flo"] + + def test_add_reverse_doubles_samples(self, tmp_path): + names = ["a", "b"] + self._make(tmp_path, "train", names) + + d = ds.FlyingChairs2Dataset( + str(tmp_path), + split="train", + add_reverse=True, + get_backward=False, + ) + + # 2 forward + 2 reversed = 4 samples. + assert len(d.img_paths) == 4 + # The first reversed entry should swap the image order and use flow_10. + assert d.img_paths[2] == [ + tmp_path / "train" / "a_img_1.png", + tmp_path / "train" / "a_img_0.png", + ] + assert d.flow_paths[2] == [tmp_path / "train" / "a_flow_10.flo"] + assert d.occ_paths[2] == [tmp_path / "train" / "a_occ_10.png"] + assert d.mb_paths[2] == [tmp_path / "train" / "a_mb_10.png"] + + def test_get_backward_populates_backward_lists(self, tmp_path): + names = ["a"] + self._make(tmp_path, "train", names) + + d = ds.FlyingChairs2Dataset( + str(tmp_path), + split="train", + add_reverse=False, + get_backward=True, + ) + + assert len(d.flow_b_paths) == 1 + assert d.flow_b_paths[0] == [tmp_path / "train" / "a_flow_10.flo"] + assert d.occ_b_paths[0] == [tmp_path / "train" / "a_occ_10.png"] + assert d.mb_b_paths[0] == [tmp_path / "train" / "a_mb_10.png"] + + +# =========================================================================== +# Tests for FlyingThings3DDataset +# =========================================================================== +class TestFlyingThings3DDataset: + """Synthesise a tiny FlyingThings3D layout and verify the path lists.""" + + def _make(self, tmp_path, n_frames=3, with_occlusions=False, with_mb=False): + # FT3D layout: + # frames_cleanpass/TRAIN/A/0000/left/{0000000.png, 0000001.png, ...} + # optical_flow/TRAIN/A/0000/into_future/left/*.pfm + # optical_flow/TRAIN/A/0000/into_past/left/*.pfm (for backward) + # occlusions/TRAIN/A/0000/into_future/left/*.png (optional) + # + # File-indexing convention used here: + # - into_future/fp_k = forward flow from frame k to k+1 + # (files indexed 0 .. n_frames-2). + # - into_past/fp_k = backward flow from frame k to k-1 + # (files indexed 0 .. n_frames-1, with the 0-th file acting as a + # placeholder for frame 0 which has no past). + # The FlyingThings3DDataset code's `flow_b_paths[i+1:i+seq_length]` + # slice relies on this range so that backward flow indices line up with + # the forward ones. + seq = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "left" + seq.mkdir(parents=True) + for i in range(n_frames): + _touch(seq / f"{i:07d}.png") + + flow_seq_fut = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_future" / "left" + ) + flow_seq_fut.mkdir(parents=True) + flow_seq_past = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_past" / "left" + ) + flow_seq_past.mkdir(parents=True) + for i in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(flow_seq_fut / f"{i:07d}.pfm", flow, fmt="pfm") + for i in range(n_frames): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(flow_seq_past / f"{i:07d}.pfm", flow, fmt="pfm") + + if with_occlusions: + occ_fut = ( + tmp_path + / "occlusions" + / "TRAIN" + / "A" + / "0000" + / "into_future" + / "left" + ) + occ_fut.mkdir(parents=True) + for i in range(n_frames - 1): + _touch(occ_fut / f"{i:07d}.png") + occ_past = ( + tmp_path / "occlusions" / "TRAIN" / "A" / "0000" / "into_past" / "left" + ) + occ_past.mkdir(parents=True) + for i in range(n_frames - 1): + _touch(occ_past / f"{i + 1:07d}.png") + if with_mb: + mb_fut = ( + tmp_path + / "motion_boundaries" + / "TRAIN" + / "A" + / "0000" + / "into_future" + / "left" + ) + mb_fut.mkdir(parents=True) + for i in range(n_frames - 1): + _touch(mb_fut / f"{i:07d}.png") + mb_past = ( + tmp_path + / "motion_boundaries" + / "TRAIN" + / "A" + / "0000" + / "into_past" + / "left" + ) + mb_past.mkdir(parents=True) + for i in range(n_frames - 1): + _touch(mb_past / f"{i + 1:07d}.png") + + def test_forward_only_seq2(self, tmp_path): + # n_frames=3 -> 2 forward samples with sequence_length=2. + self._make(tmp_path, n_frames=3) + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + side_names="left", + add_reverse=False, + get_backward=False, + sequence_length=2, + ) + + assert len(d.img_paths) == 2 + assert len(d.flow_paths) == 2 + + # Sample 0: images [0000000.png, 0000001.png] + flow [0000000.pfm]. + base = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "left" + flow_base = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_future" / "left" + ) + assert d.img_paths[0] == [base / "0000000.png", base / "0000001.png"] + assert d.flow_paths[0] == [flow_base / "0000000.pfm"] + assert d.img_paths[1] == [base / "0000001.png", base / "0000002.png"] + assert d.flow_paths[1] == [flow_base / "0000001.pfm"] + + def test_add_reverse_appends_backward_samples(self, tmp_path): + self._make(tmp_path, n_frames=3) + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + side_names="left", + add_reverse=True, + get_backward=False, + sequence_length=2, + ) + + # 2 forward + 2 reverse = 4 samples. + assert len(d.img_paths) == 4 + # The first reversed sample should have images in descending order + # ([0000002, 0000001]) and the into_past flow at index 0000002, since + # file `fp_k` in into_past holds the flow from frame k to k-1. + base = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "left" + flow_past = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_past" / "left" + ) + assert d.img_paths[2] == [base / "0000002.png", base / "0000001.png"] + assert d.flow_paths[2] == [flow_past / "0000002.pfm"] + + def test_get_backward_populates_flow_b_paths(self, tmp_path): + self._make(tmp_path, n_frames=3) + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + side_names="left", + add_reverse=False, + get_backward=True, + sequence_length=2, + ) + + assert len(d.flow_b_paths) == 2 + flow_past = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_past" / "left" + ) + # With into_past files indexed 0..n_frames-1, the code slices + # `flow_b_paths[i+1]` to fetch the backward flow for forward sample i. + # For i=0 that resolves to into_past/0000001.pfm (flow from frame 1 to + # frame 0), which is the correct counterpart of into_future/0000000.pfm. + assert d.flow_b_paths[0] == [flow_past / "0000001.pfm"] + + +# =========================================================================== +# Tests for FlyingThings3DSubsetDataset +# =========================================================================== +class TestFlyingThings3DSubsetDataset: + def _make(self, tmp_path, n_frames=5, split="train"): + # Subset layout: + # {split}/image_clean/left/{0000000.png, ...} + # {split}/flow/left/into_future/left/*.flo + # {split}/flow/left/into_past/left/*.flo + img_dir = tmp_path / split / "image_clean" / "left" + img_dir.mkdir(parents=True) + for i in range(n_frames): + _touch(img_dir / f"{i:07d}.png") + + fut_dir = tmp_path / split / "flow" / "left" / "into_future" + past_dir = tmp_path / split / "flow" / "left" / "into_past" + fut_dir.mkdir(parents=True) + past_dir.mkdir(parents=True) + # forward flows: file k is flow from frame k to k+1, so indices 0..n-2 + for i in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(fut_dir / f"{i:07d}.flo", flow) + _write_flow(past_dir / f"{i + 1:07d}.flo", flow) + + def test_forward_only_seq2(self, tmp_path): + # 5 frames -> 4 forward flows -> 4 samples. + self._make(tmp_path, n_frames=5) + d = ds.FlyingThings3DSubsetDataset( + str(tmp_path), + split="train", + pass_names="clean", + side_names="left", + add_reverse=False, + get_backward=False, + sequence_length=2, + ) + + assert len(d.img_paths) == 4 + img_base = tmp_path / "train" / "image_clean" / "left" + flow_base = tmp_path / "train" / "flow" / "left" / "into_future" + # Sample 0: img [0000000, 0000001] + flow [0000000] + assert d.img_paths[0] == [img_base / "0000000.png", img_base / "0000001.png"] + assert d.flow_paths[0] == [flow_base / "0000000.flo"] + + def test_add_reverse_appends_backward_samples(self, tmp_path): + self._make(tmp_path, n_frames=5) + d = ds.FlyingThings3DSubsetDataset( + str(tmp_path), + split="train", + pass_names="clean", + side_names="left", + add_reverse=True, + get_backward=False, + sequence_length=2, + ) + + # 4 forward + 4 backward (sequence_length=2 doesn't trigger the bug). + assert len(d.img_paths) == 8 + img_base = tmp_path / "train" / "image_clean" / "left" + flow_past = tmp_path / "train" / "flow" / "left" / "into_past" + # First backward sample: flow file 0000004 (into_past) corresponds to + # images [0000004.png, 0000003.png] (flow from frame 4 to 3). + assert d.img_paths[4] == [img_base / "0000004.png", img_base / "0000003.png"] + assert d.flow_paths[4] == [flow_past / "0000004.flo"] + + def test_add_reverse_seq3_image_alignment(self, tmp_path): + # 5 frames -> 4 forward flows. With sequence_length=3 and + # sequence_position="first" the flow list is padded with one extra + # copy of the last flow at the end, giving 4 (degenerate) forward + # samples followed by 4 backward samples - 8 in total. + self._make(tmp_path, n_frames=5) + d = ds.FlyingThings3DSubsetDataset( + str(tmp_path), + split="train", + pass_names="clean", + side_names="left", + add_reverse=True, + get_backward=False, + sequence_length=3, + sequence_position="first", + ) + assert len(d.img_paths) == 8 + + img_base = tmp_path / "train" / "image_clean" / "left" + flow_past = tmp_path / "train" / "flow" / "left" / "into_past" + # The first BACKWARD sample lives at index 4 (indices 0..3 are + # forward samples coming from the into_future direction). The + # backward flow list, sorted descending, is [fp4, fp3, fp2, fp1] and + # (now that the grouping logic handles descending lists correctly) + # forms a single group. The first backward sample uses flows + # [fp4, fp3] (flow from 4->3 then 3->2), so the expected images are + # [img4, img3, img2]. + assert d.img_paths[4] == [ + img_base / "0000004.png", + img_base / "0000003.png", + img_base / "0000002.png", + ] + assert d.flow_paths[4] == [flow_past / "0000004.flo", flow_past / "0000003.flo"] + + +# =========================================================================== +# Tests for Hd1kDataset +# =========================================================================== +class TestHd1kDataset: + def _make(self, tmp_path, sequences): + # HD1K layout: + # hd1k_input/image_2/{seq}_{frame}.png + # hd1k_flow_gt/flow_occ/{seq}_{frame}.png (training only) + img_dir = tmp_path / "hd1k_input" / "image_2" + img_dir.mkdir(parents=True) + flow_dir = tmp_path / "hd1k_flow_gt" / "flow_occ" + flow_dir.mkdir(parents=True) + for seq, n_frames in sequences: + for f in range(n_frames): + name = f"{seq}_{f:04d}" + _touch(img_dir / f"{name}.png") + _touch(flow_dir / f"{name}.png") + + def test_trainval_seq2(self, tmp_path): + # Two sequences each with 3 frames. Use names that are not in + # Hd1k_val.txt so they remain in every split. + self._make(tmp_path, [("myseq", 3), ("myother", 2)]) + + d = ds.Hd1kDataset( + str(tmp_path), + split="trainval", + sequence_length=2, + sequence_position="first", + ) + + # myseq: 3 frames -> 2 samples. myother: 2 frames -> 1 sample. Total = 3. + assert len(d.img_paths) == 3 + assert len(d.flow_paths) == 3 + img_dir = tmp_path / "hd1k_input" / "image_2" + flow_dir = tmp_path / "hd1k_flow_gt" / "flow_occ" + # Sorted by sequence name then frame index. + # First sample: myother_0000 + myother_0001 (since "myother" < "myseq" + # in sorted order... wait - sorted by stem first, so 'myother' sorts + # before 'myseq'). + assert d.img_paths[0] == [ + img_dir / "myother_0000.png", + img_dir / "myother_0001.png", + ] + assert d.flow_paths[0] == [flow_dir / "myother_0000.png"] + # Next two samples come from myseq. + assert d.img_paths[1] == [ + img_dir / "myseq_0000.png", + img_dir / "myseq_0001.png", + ] + assert d.flow_paths[1] == [flow_dir / "myseq_0000.png"] + assert d.img_paths[2] == [ + img_dir / "myseq_0001.png", + img_dir / "myseq_0002.png", + ] + assert d.flow_paths[2] == [flow_dir / "myseq_0001.png"] + + +# =========================================================================== +# Tests for KittiDataset +# =========================================================================== +class TestKittiDataset: + def _make(self, tmp_path, version, sample_prefixes): + # KITTI 2015 layout: training/image_2/{stem}_10.png and _11.png + # training/flow_occ/{stem}_10.png + # KITTI 2012 layout: training/colored_0/{stem}_10.png and _11.png + ver_dir = tmp_path / f"kitti_{version}" + if version == "2012": + image_dir = ver_dir / "training" / "colored_0" + else: + image_dir = ver_dir / "training" / "image_2" + image_dir.mkdir(parents=True) + flow_dir = ver_dir / "training" / "flow_occ" + flow_dir.mkdir(parents=True) + for pfx in sample_prefixes: + _touch(image_dir / f"{pfx}_10.png") + _touch(image_dir / f"{pfx}_11.png") + _touch(flow_dir / f"{pfx}_10.png") + + def test_kitti_2015_trainval(self, tmp_path): + # Use stems that are not in Kitti2015_val.txt. + prefixes = ["mykitti0", "mykitti1"] + self._make(tmp_path, "2015", prefixes) + + d = ds.KittiDataset( + root_dir_2015=str(tmp_path / "kitti_2015"), split="trainval" + ) + + assert len(d.img_paths) == 2 + img_dir = tmp_path / "kitti_2015" / "training" / "image_2" + flow_dir = tmp_path / "kitti_2015" / "training" / "flow_occ" + assert d.img_paths[0] == [ + img_dir / "mykitti0_10.png", + img_dir / "mykitti0_11.png", + ] + assert d.flow_paths[0] == [flow_dir / "mykitti0_10.png"] + # The 'misc' metadata field stores the version string. + assert d.metadata[0]["misc"] == "2015" + + def test_kitti_2015_val_split_uses_val_file(self, tmp_path): + # '000010_10' is the first entry in Kitti2015_val.txt. + prefixes = ["000010", "mykitti0"] + self._make(tmp_path, "2015", prefixes) + + d = ds.KittiDataset(root_dir_2015=str(tmp_path / "kitti_2015"), split="val") + kept = [p[0].stem for p in d.img_paths] + assert kept == ["000010_10"] + assert d.metadata[0]["is_val"] is True + + def test_kitti_combined_versions(self, tmp_path): + self._make(tmp_path, "2012", ["myk12a"]) + self._make(tmp_path, "2015", ["myk15a"]) + + d = ds.KittiDataset( + root_dir_2012=str(tmp_path / "kitti_2012"), + root_dir_2015=str(tmp_path / "kitti_2015"), + split="trainval", + versions=["2012", "2015"], + ) + + assert len(d.img_paths) == 2 + # The dataset_name should encode both versions. + assert d.dataset_name == "KITTI_2012_2015" + # metadata 'misc' should reflect the version each sample came from. + miscs = sorted(m["misc"] for m in d.metadata) + assert miscs == ["2012", "2015"] + + +# =========================================================================== +# Tests for SintelDataset +# =========================================================================== +class TestSintelDataset: + def _make(self, tmp_path, seq_names, n_frames=3, with_flow=True): + # Sintel layout: + # training/clean/{seq}/{frame}.png + # training/final/{seq}/{frame}.png + # training/flow/{seq}/{frame}.flo (one fewer than images) + # training/occlusions/{seq}/{frame}.png + for seq in seq_names: + for passd in ["clean", "final"]: + sd = tmp_path / "training" / passd / seq + sd.mkdir(parents=True) + for f in range(n_frames): + _touch(sd / f"frame_{f:04d}.png") + if with_flow: + fd = tmp_path / "training" / "flow" / seq + 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") + + def test_trainval_single_pass(self, tmp_path): + # Use a sequence name not in Sintel_val.txt. + seqs = ["myseq1"] + self._make(tmp_path, seqs, n_frames=3) + + d = ds.SintelDataset( + str(tmp_path), + split="trainval", + pass_names="clean", + sequence_length=2, + sequence_position="first", + ) + + # 3 frames, sequence_length=2 -> 2 samples. + assert len(d.img_paths) == 2 + assert len(d.flow_paths) == 2 + # Each sample should reference the 'clean' pass. + base = tmp_path / "training" / "clean" / "myseq1" + flow_base = tmp_path / "training" / "flow" / "myseq1" + occ_base = tmp_path / "training" / "occlusions" / "myseq1" + assert d.img_paths[0] == [base / "frame_0000.png", base / "frame_0001.png"] + assert d.flow_paths[0] == [flow_base / "frame_0000.flo"] + assert d.occ_paths[0] == [occ_base / "frame_0000.png"] + + def test_trainval_both_passes(self, tmp_path): + seqs = ["myseq1"] + self._make(tmp_path, seqs, n_frames=3) + + d = ds.SintelDataset( + str(tmp_path), + split="trainval", + pass_names=["clean", "final"], + sequence_length=2, + ) + + # Each pass contributes 2 samples -> 4 total. + assert len(d.img_paths) == 4 + # The first two samples belong to 'clean', the next two to 'final'. + clean0 = tmp_path / "training" / "clean" / "myseq1" / "frame_0000.png" + assert d.img_paths[0][0] == clean0 + final0 = tmp_path / "training" / "final" / "myseq1" / "frame_0000.png" + assert d.img_paths[2][0] == final0 + + def test_val_split_uses_val_file(self, tmp_path): + # Sintel_val.txt ships with these 6 sequence names. For the 'val' split + # the code sets ``sequence_names = val_seqs`` (i.e. ALL of them), so + # each one must exist on disk. + val_seqs = [ + "ambush_2", + "bamboo_2", + "cave_2", + "market_2", + "shaman_2", + "temple_2", + ] + self._make(tmp_path, val_seqs, n_frames=3) + d = ds.SintelDataset(str(tmp_path), split="val", pass_names="clean") + # 6 sequences * (3 frames, sequence_length=2 -> 2 samples) = 12 samples. + assert len(d.img_paths) == 12 + # The set of sequence names used must match the val file contents. + seqs_used = {p[0].parent.stem for p in d.img_paths} + assert seqs_used == set(val_seqs) + assert all(m["is_val"] is True for m in d.metadata) + + +# =========================================================================== +# Tests for SpringDataset +# =========================================================================== +class TestSpringDataset: + def _make(self, tmp_path, seq_names, n_frames=3, split="train"): + # Spring layout: + # train/{seq}/frame_left/{frame}.png + # train/{seq}/flow_FW_left/{frame}.flo5 + # train/{seq}/flow_BW_left/{frame}.flo5 + # Backward flow file k is the flow from frame k back to k-1, so its + # indices run from 1 to n_frames-1. + split_dir = "train" if split != "test" else "test" + for seq in seq_names: + img_d = tmp_path / split_dir / seq / "frame_left" + img_d.mkdir(parents=True) + for f in range(n_frames): + _touch(img_d / f"f_{f:04d}.png") + fwd_d = tmp_path / split_dir / seq / "flow_FW_left" + bwd_d = tmp_path / split_dir / seq / "flow_BW_left" + fwd_d.mkdir(parents=True) + bwd_d.mkdir(parents=True) + for f in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(fwd_d / f"f_{f:04d}.flo5", flow, fmt="flo5") + _write_flow(bwd_d / f"f_{f + 1:04d}.flo5", flow, fmt="flo5") + + def test_train_split_paths(self, tmp_path): + # Use a sequence name not in Spring_val.txt ('00270' is the only entry). + self._make(tmp_path, ["myseq"], n_frames=3, split="train") + d = ds.SpringDataset( + str(tmp_path), + split="train", + side_names="left", + add_reverse=False, + get_backward=False, + sequence_length=2, + ) + # 3 frames -> 2 samples. + assert len(d.img_paths) == 2 + base = tmp_path / "train" / "myseq" / "frame_left" + fwd = tmp_path / "train" / "myseq" / "flow_FW_left" + assert d.img_paths[0] == [base / "f_0000.png", base / "f_0001.png"] + assert d.flow_paths[0] == [fwd / "f_0000.flo5"] + assert d.img_paths[1] == [base / "f_0001.png", base / "f_0002.png"] + assert d.flow_paths[1] == [fwd / "f_0001.flo5"] + + def test_train_split_with_add_reverse(self, tmp_path): + self._make(tmp_path, ["myseq"], n_frames=3, split="train") + d = ds.SpringDataset( + str(tmp_path), + split="train", + side_names="left", + add_reverse=True, + get_backward=False, + sequence_length=2, + ) + # 2 forward + 2 backward = 4 samples. + assert len(d.img_paths) == 4 + base = tmp_path / "train" / "myseq" / "frame_left" + bwd = tmp_path / "train" / "myseq" / "flow_BW_left" + # First reversed sample uses BW flow f_0002 (flow from 2 -> 1) and + # images [f_0002.png, f_0001.png]. + assert d.img_paths[2] == [base / "f_0002.png", base / "f_0001.png"] + assert d.flow_paths[2] == [bwd / "f_0002.flo5"] + + +# =========================================================================== +# Tests for TartanAirDataset +# =========================================================================== +class TestTartanAirDataset: + def _make(self, tmp_path, seq_names, n_frames=3, difficulty="Easy"): + # TartanAir layout: + # {seq}/{Difficulty}/{traj}/image_left/{frame}.png + # {seq}/{Difficulty}/{traj}/flow/{frame}_flow.npy + # {seq}/{Difficulty}/{traj}/flow/{frame}_mask.npy (occlusion) + for seq in seq_names: + traj = tmp_path / seq / difficulty / "P000" + img_d = traj / "image_left" + flow_d = traj / "flow" + img_d.mkdir(parents=True) + flow_d.mkdir(parents=True) + for f in range(n_frames): + _touch(img_d / f"{f:06d}_left.png") + for f in range(n_frames - 1): + np.save( + flow_d / f"{f:06d}_flow.npy", np.zeros((4, 5, 2), dtype=np.float32) + ) + np.save(flow_d / f"{f:06d}_mask.npy", np.zeros((4, 5), dtype=np.uint8)) + + def test_path_generation(self, tmp_path): + self._make(tmp_path, ["env1"], n_frames=3, difficulty="Easy") + d = ds.TartanAirDataset( + str(tmp_path), + difficulties="easy", + sequence_length=2, + sequence_position="first", + get_occlusion_mask=True, + ) + + # 3 frames -> 2 samples. + assert len(d.img_paths) == 2 + img_d = tmp_path / "env1" / "Easy" / "P000" / "image_left" + flow_d = tmp_path / "env1" / "Easy" / "P000" / "flow" + assert d.img_paths[0] == [img_d / "000000_left.png", img_d / "000001_left.png"] + assert d.flow_paths[0] == [flow_d / "000000_flow.npy"] + assert d.occ_paths[0] == [flow_d / "000000_mask.npy"] + # metadata 'misc' should store the sequence name. + assert d.metadata[0]["misc"] == "env1" + + +# =========================================================================== +# Tests for MiddleburyDataset +# =========================================================================== +class TestMiddleburyDataset: + def _make(self, tmp_path, seq_names, n_frames=2): + # Middlebury layout (training split): + # other-data/{seq}/frame.png + # other-gt-flow/{seq}/frame.flo + for seq in seq_names: + dd = tmp_path / "other-data" / seq + dd.mkdir(parents=True) + for f in range(n_frames): + _touch(dd / f"frame_{f}.png") + fd = tmp_path / "other-gt-flow" / seq + fd.mkdir(parents=True) + for f in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(fd / f"frame_{f}.flo", flow) + + def test_train_paths(self, tmp_path): + self._make(tmp_path, ["myseq"], n_frames=2) + d = ds.MiddleburyDataset(str(tmp_path), split="train") + # 2 frames, sequence_length=2 -> 1 sample. + assert len(d.img_paths) == 1 + base = tmp_path / "other-data" / "myseq" + flow_base = tmp_path / "other-gt-flow" / "myseq" + assert d.img_paths[0] == [base / "frame_0.png", base / "frame_1.png"] + assert d.flow_paths[0] == [flow_base / "frame_0.flo"] + + +# =========================================================================== +# Tests for MiddleburySTDataset +# =========================================================================== +class TestMiddleburySTDataset: + def _make(self, tmp_path, seq_names): + # Middlebury-ST layout: + # {seq}/im0.png, {seq}/im1.png + # {seq}/disp0.pfm, {seq}/disp0y.pfm (two disparity files used as + # the x / y flow channels after negation) + for seq in seq_names: + sd = tmp_path / seq + sd.mkdir(parents=True) + _touch(sd / "im0.png") + _touch(sd / "im1.png") + disp_x = np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32) + disp_y = np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32) + _write_flow(sd / "disp0.pfm", disp_x, fmt="pfm") + _write_flow(sd / "disp0y.pfm", disp_y, fmt="pfm") + + def test_paths(self, tmp_path): + self._make(tmp_path, ["scene_a"]) + d = ds.MiddleburySTDataset(str(tmp_path)) + assert len(d.img_paths) == 1 + sd = tmp_path / "scene_a" + assert d.img_paths[0] == [sd / "im0.png", sd / "im1.png"] + # flow_paths[0] is a list of length 1, the inner element being the + # pair [disp0.pfm, disp0y.pfm]. + assert d.flow_paths[0][0] == [sd / "disp0.pfm", sd / "disp0y.pfm"] + assert d.is_two_file_flow is True + + def test_getitem_flow_content(self, tmp_path): + self._make(tmp_path, ["scene_a"]) + # Add real RGB images so cv.imread succeeds. + sd = tmp_path / "scene_a" + _write_rgb_png(sd / "im0.png", color=(1, 2, 3)) + _write_rgb_png(sd / "im1.png", color=(4, 5, 6)) + + d = ds.MiddleburySTDataset(str(tmp_path), get_valid_mask=True) + out = d[0] + + # The flow is built from [-disp_x, -disp_y] stacked along axis 2. + expected_flow = np.stack( + [ + -np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + -np.array([[5.0, 6.0], [7.0, 8.0]], dtype=np.float32), + ], + axis=2, + ) + np.testing.assert_allclose(out["flows"][0], expected_flow, rtol=1e-5) + # All disparity values are below max_flow so the mask is fully valid. + np.testing.assert_array_equal( + out["valid_flows"][0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + + +# =========================================================================== +# Tests for KubricDataset +# =========================================================================== +class TestKubricDataset: + def _make(self, tmp_path, n_frames=3): + # Kubric layout (per sequence directory): + # rgba_{frame:05d}.png + # forward_flow_{frame:05d}.png (stored as 16-bit PNG, with a + # data_ranges.json describing min/max) + # backward_flow_{frame:05d}.png + seq_dir = tmp_path / "seq0" + seq_dir.mkdir() + for f in range(n_frames): + _touch(seq_dir / f"rgba_{f:05d}.png") + # forward_flow files go from frame 0 to n_frames-1 (inclusive), but the + # loader drops the last one with [:-1], leaving n_frames-1 flows. + for f in range(n_frames): + _touch(seq_dir / f"forward_flow_{f:05d}.png") + _touch(seq_dir / f"backward_flow_{f:05d}.png") + # Kubric expects a data_ranges.json inside each sequence directory. + (seq_dir / "data_ranges.json").write_text( + json.dumps( + { + "forward_flow": {"min": -10.0, "max": 10.0}, + "backward_flow": {"min": -10.0, "max": 10.0}, + } + ) + ) + + def test_paths_forward_only(self, tmp_path): + self._make(tmp_path, n_frames=3) + d = ds.KubricDataset( + str(tmp_path), + get_backward=False, + sequence_length=2, + sequence_position="first", + ) + # 3 frames -> 2 samples. + assert len(d.img_paths) == 2 + seq_dir = tmp_path / "seq0" + assert d.img_paths[0] == [ + seq_dir / "rgba_00000.png", + seq_dir / "rgba_00001.png", + ] + # Each flow entry is a tuple (path, "forward_flow") per the loader's + # internal representation. + assert d.flow_paths[0] == [(seq_dir / "forward_flow_00000.png", "forward_flow")] + + def test_paths_with_backward(self, tmp_path): + self._make(tmp_path, n_frames=3) + d = ds.KubricDataset( + str(tmp_path), + get_backward=True, + sequence_length=2, + ) + # backward_flow files: loader takes [1:] so files 1..n_frames-1 are + # used as backward flows. + assert len(d.flow_b_paths) == 2 + seq_dir = tmp_path / "seq0" + assert d.flow_b_paths[0] == [ + (seq_dir / "backward_flow_00001.png", "backward_flow") + ] + + +# =========================================================================== +# Tests for ViperDataset +# =========================================================================== +class TestViperDataset: + def _make(self, tmp_path, split_dir, seq_names, n_frames=3): + # VIPER layout: + # {split}/img/{seq}/{seq}_{idx:05d}.png + # {split}/flow/{seq}/{seq}_{idx:05d}.npz + # The flow at index k describes motion from image k to image k+1. + # NOTE: ViperDataset parses the flow file's idx as + # ``int(fpath.stem.split('_')[1])`` so the stem must have exactly one + # underscore before the integer idx (no `_flow_` in the middle). + for seq in seq_names: + img_d = tmp_path / split_dir / "img" / seq + flow_d = tmp_path / split_dir / "flow" / seq + img_d.mkdir(parents=True) + flow_d.mkdir(parents=True) + for f in range(n_frames): + _touch(img_d / f"{seq}_{f:05d}.png") + for f in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(flow_d / f"{seq}_{f:05d}.npz", flow, fmt="viper_npz") + + def test_train_paths(self, tmp_path): + self._make(tmp_path, "train", ["seqA"], n_frames=3) + d = ds.ViperDataset(str(tmp_path), split="train") + # 3 frames, 2 flows -> 2 samples. + assert len(d.img_paths) == 2 + img_d = tmp_path / "train" / "img" / "seqA" + flow_d = tmp_path / "train" / "flow" / "seqA" + assert d.img_paths[0] == [img_d / "seqA_00000.png", img_d / "seqA_00001.png"] + assert d.flow_paths[0] == [flow_d / "seqA_00000.npz"] + # metadata flags should mark the train split as not val. + assert all(m["is_val"] is False for m in d.metadata) + + def test_trainval_iterates_both_dirs(self, tmp_path): + self._make(tmp_path, "train", ["seqA"], n_frames=2) + self._make(tmp_path, "val", ["seqB"], n_frames=2) + d = ds.ViperDataset(str(tmp_path), split="trainval") + # 1 sample from train + 1 sample from val = 2 total. + assert len(d.img_paths) == 2 + assert any(m["is_val"] for m in d.metadata) + assert any(not m["is_val"] for m in d.metadata) + + +# =========================================================================== +# Tests for MonkaaDataset +# =========================================================================== +class TestMonkaaDataset: + def test_monkaa_can_be_instantiated(self, tmp_path): + # Synthesise the Monkaa layout used by the loader: one sequence with + # two frames, forward flows in `optical_flow/into_future` and backward + # flows in `optical_flow/into_past`. + seq_img_dir = tmp_path / "frames_cleanpass" / "shot0" / "left" + seq_img_dir.mkdir(parents=True) + _touch(seq_img_dir / "0000.png") + _touch(seq_img_dir / "0001.png") + flow_fut = tmp_path / "optical_flow" / "shot0" / "into_future" / "left" + flow_past = tmp_path / "optical_flow" / "shot0" / "into_past" / "left" + flow_fut.mkdir(parents=True) + flow_past.mkdir(parents=True) + zero_flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(flow_fut / "0000.pfm", zero_flow, fmt="pfm") + _write_flow(flow_past / "0001.pfm", zero_flow, fmt="pfm") + + d = ds.MonkaaDataset( + str(tmp_path), + pass_names="clean", + side_names="left", + add_reverse=False, + get_backward=False, + sequence_length=2, + ) + assert len(d.img_paths) == 1 + assert d.img_paths[0] == [seq_img_dir / "0000.png", seq_img_dir / "0001.png"] diff --git a/tests/common/data/test_optical_flow_transforms.py b/tests/common/data/test_optical_flow_transforms.py new file mode 100644 index 0000000..91d431f --- /dev/null +++ b/tests/common/data/test_optical_flow_transforms.py @@ -0,0 +1,425 @@ +import random +from unittest.mock import patch + +import numpy as np +import torch + +from roco_spring_devkit.common.data.optical_flow_transforms import ( + CenterCrop, + ColorJitter, + Compose, + GaussianNoise, + GenerateFBCheckFlowOcclusion, + RandomFlip, + RandomPatchEraser, + RandomRotate, + RandomScaleAndCrop, + RandomTranslate, + Resize, + ToTensor, + _get_valid_keys, + _resize, + _update_oob_flows, +) + + +class TestUtilityFunctions: + def test_get_valid_keys_logic(self): + inputs_keys = ["images", "flows", "occs", "meta"] + + # Test use_keys priority + assert _get_valid_keys( + inputs_keys, use_keys=["images"], ignore_keys=["images"] + ) == ["images"] + + # Test ignore_keys filtering + assert _get_valid_keys( + inputs_keys, use_keys=None, ignore_keys=["meta", "occs"] + ) == ["images", "flows"] + + # Test passthrough + assert ( + _get_valid_keys(inputs_keys, use_keys=None, ignore_keys=None) == inputs_keys + ) + + def test_update_oob_flows_logic(self): + occs = torch.zeros(1, 1, 3, 3) + flows = torch.zeros(1, 2, 3, 3) + + # Flow [0, 0] at (0, 0) goes out of bounds to the top-left (-10, -10) + flows[0, 0, 0, 0] = -10.0 + flows[0, 1, 0, 0] = -10.0 + + # Flow [0, 1] at (0, 2) goes out of bounds to the right (+10) + flows[0, 0, 0, 2] = 10.0 + + # Landing exactly one pixel past the right edge is also out of bounds. + flows[0, 0, 1, 2] = 1.0 + + out_occs = _update_oob_flows(occs, flows) + + # Should be occluded (out of bounds) + assert out_occs[0, 0, 0, 0] == 1.0 + assert out_occs[0, 0, 0, 2] == 1.0 + assert out_occs[0, 0, 1, 2] == 1.0 + + # Should NOT be occluded + assert out_occs[0, 0, 1, 1] == 0.0 + + +class TestCompose: + def test_compose_filters_none_and_applies_all(self): + # Compose should filter out None transforms and apply the rest in sequence + t = Compose([ToTensor(), None, CenterCrop(crop_size=(2, 2))]) + + assert len(t.transforms_list) == 2 + + img = np.zeros((4, 4, 3), dtype=np.float32) + valid = np.ones((4, 4, 1), dtype=np.float32) + inputs = {"images": img, "valids": valid} + + out = t(inputs) + + # Verify both ToTensor and CenterCrop ran + assert out["images"].shape == (1, 3, 2, 2) + assert out["valids"].shape == (1, 1, 2, 2) + + +class TestToTensor: + def test_uint8_conversion_and_transpose(self): + img = np.array( + [[[0, 0, 0], [127, 127, 127]], [[255, 255, 255], [0, 255, 0]]], + dtype=np.uint8, + ) + inputs = {"images": [img]} + + t = ToTensor() + out = t(inputs) + + assert out["images"].shape == (1, 3, 2, 2) + assert out["images"].dtype == torch.float32 + + expected = torch.tensor( + [ + [[0.0, 127 / 255.0], [1.0, 0.0]], # R + [[0.0, 127 / 255.0], [1.0, 1.0]], # G + [[0.0, 127 / 255.0], [1.0, 0.0]], # B + ] + ).unsqueeze(0) + assert torch.allclose(out["images"], expected, atol=1e-5) + + def test_list_input_and_fp16(self): + img1 = np.ones((2, 2), dtype=np.uint8) * 127 + img2 = np.ones((2, 2), dtype=np.uint8) * 255 + + inputs = {"images": [img1, img2]} + t = ToTensor(fp16=True, device="cpu") + out = t(inputs) + + # Stacked list expanded for 2D inputs + assert out["images"].shape == (2, 1, 2, 2) + assert out["images"].dtype == torch.float16 + + assert torch.allclose( + out["images"][0], torch.tensor(127.0 / 255.0, dtype=torch.float16) + ) + assert torch.allclose(out["images"][1], torch.tensor(1.0, dtype=torch.float16)) + + def test_use_keys_and_ignore_keys(self): + inputs = { + "images": np.ones((2, 2, 3), dtype=np.float32), + "flows": np.ones((2, 2, 2), dtype=np.float32), + "meta": "should_be_ignored", + } + + t = ToTensor(use_keys=["images"]) + out = t(inputs) + + assert isinstance(out["images"], torch.Tensor) + assert isinstance(out["flows"], np.ndarray) # Untouched + assert out["meta"] == "should_be_ignored" + + +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()} + + t = CenterCrop(crop_size=(2, 2)) + out = t(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 torch.allclose(out["images"], expected) + + +class TestResize: + def test_resize_flow_multipliers(self): + img = torch.ones(1, 1, 2, 2) + flow = torch.ones(1, 2, 2, 2) * 2.0 + inputs = {"images": img, "flows": flow} + + # Target size H=4, W=6 -> ScaleY = 2.0, ScaleX = 3.0 + t = Resize(size=(4, 6)) + out = t(inputs) + + assert out["images"].shape == (1, 1, 4, 6) + + # Flow vector X (idx 0) multiplied by 3.0; Y (idx 1) multiplied by 2.0 + assert torch.allclose(out["flows"][:, 0, :, :], torch.tensor(6.0)) + assert torch.allclose(out["flows"][:, 1, :, :], torch.tensor(4.0)) + + +class TestRandomScaleAndCrop: + @patch("random.randint") + @patch("random.uniform") + def test_random_scale_and_crop_logic(self, mock_uniform, mock_randint): + mock_uniform.side_effect = [1.0, 1.0, 1.0] # major, space_h, space_w + mock_randint.side_effect = [0, 0] # y_crop, x_crop + + flow = torch.ones(1, 2, 4, 4) + inputs = {"flows": flow} + + # Major and spatial exponents each produce a scale of 2, for a total of 4. + t = RandomScaleAndCrop( + crop_size=(6, 6), major_scale=(1.0, 1.0), space_scale=(1.0, 1.0) + ) + out = t(inputs) + + assert out["flows"].shape == (1, 2, 6, 6) + assert torch.allclose(out["flows"], torch.tensor(4.0)) + + def test_resize_sparse_logic(self): + valid = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) + flow = torch.ones(1, 2, 2, 2) + + inputs = {"valids": valid.clone(), "flows": flow.clone()} + + # Isolate sparse _resize path + out = _resize( + inputs, + target_size=(4, 4), + binary_keys=["valids"], + flow_keys=["flows"], + sparse=True, + valid_key="valids", + ) + + assert out["valids"].shape == (1, 1, 4, 4) + + # Only accurately scaled coordinates should remain + assert out["valids"][0, 0, 0, 0] == 1.0 + assert out["valids"][0, 0, 2, 2] == 1.0 + assert out["valids"].sum() == 2.0 + + assert out["flows"][0, 0, 0, 0] == 2.0 + assert out["flows"][0, 0, 2, 2] == 2.0 + assert out["flows"][0, 0, 0, 1] == 0.0 + + +class TestRandomFlip: + def test_horizontal_flip(self): + img = torch.arange(4).reshape(1, 1, 2, 2).float() + flow = torch.tensor([[[[1.0, 1.0], [1.0, 1.0]], [[2.0, 2.0], [2.0, 2.0]]]]) + inputs = {"images": img.clone(), "flows": flow.clone()} + + t = RandomFlip(hflip_prob=1.0, vflip_prob=0.0) + out = t(inputs) + + expected_img = torch.tensor([[[[1.0, 0.0], [3.0, 2.0]]]]) + assert torch.allclose(out["images"], expected_img) + + expected_flow = torch.tensor( + [[[[-1.0, -1.0], [-1.0, -1.0]], [[2.0, 2.0], [2.0, 2.0]]]] + ) + assert torch.allclose(out["flows"], expected_flow) + + @patch("random.random") + def test_asymmetric_flip_mirror_flow(self, mock_random): + mock_random.side_effect = [0.0, 1.0] # Image 0 flips, Image 1 does NOT + + flow = torch.zeros(2, 2, 2, 2) + img = torch.zeros(2, 3, 2, 2) + inputs = {"images": img, "flows": flow} + + t = RandomFlip(hflip_prob=1.0, vflip_prob=0.0, asymmetric_prob=1.0) + out = t(inputs) + + expected_mirror = torch.tensor([[1.0, -1.0], [1.0, -1.0]]) + + assert torch.allclose(out["flows"][0, 0], expected_mirror) + assert torch.all(out["flows"][1] == 0.0) # Untouched + + +class TestRandomTranslate: + @patch("random.randint") + def test_random_translate_values(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, 2, 2] = 1.0 + + flow = torch.zeros(2, 2, 4, 4) + inputs = {"images": img, "flows": flow} + + t = RandomTranslate(translation=(2, 2)) + out = t(inputs) + + assert out["images"].shape == (2, 1, 2, 3) + + # Even index gets (+tw, +th) + assert torch.allclose(out["flows"][0, 0, :, :], torch.tensor(1.0)) + assert torch.allclose(out["flows"][0, 1, :, :], torch.tensor(2.0)) + + # Odd index gets (-tw, -th) + assert torch.allclose(out["flows"][1, 0, :, :], torch.tensor(-1.0)) + assert torch.allclose(out["flows"][1, 1, :, :], torch.tensor(-2.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, "flows": torch.zeros(1, 2, 4, 4)} + + t = RandomTranslate(translation=10) + out = t(inputs) + + # Early return check + assert out["images"] is img + assert out["images"].shape == (1, 1, 4, 4) + + +class TestRandomRotate: + @patch("random.uniform") + def test_rotation_matrix_and_flow_values(self, mock_uniform): + mock_uniform.side_effect = [90.0, 0.0] + + flow = torch.zeros(2, 2, 3, 3) + flow[:, 0, :, :] = 1.0 # Pointing RIGHT + + inputs = {"flows": flow.clone(), "occs": torch.zeros(2, 1, 3, 3)} + t = RandomRotate(angle=90.0) + out = t(inputs) + + # Pointing right rotated 90 degrees points UP/DOWN + 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 + + inputs = {"valids": valid, "flows": torch.zeros(1, 2, 3, 3)} + + t = RandomRotate(angle=90.0, sparse=True) + out = t(inputs) + + assert out["valids"][0, 0, 1, 0] == 1.0 + assert out["valids"].sum() == 1.0 + + +class TestGenerateFBCheckFlowOcclusion: + def test_consistent_and_inconsistent_flow(self): + b, c, h, w = 1, 2, 4, 4 + flow_f = torch.zeros(b, c, h, w) + flow_f[:, 0, :, :] = 1.0 # Push RIGHT + + flow_b = torch.zeros(b, c, h, w) + flow_b[:, 0, :, :] = -1.0 # Push LEFT + + # Corrupt backward vector mapping + flow_b[0, 0, 2, 3] = 10.0 + + inputs = {"flows": flow_f, "flows_b": flow_b} + t = GenerateFBCheckFlowOcclusion(threshold=1.0) + out = t(inputs) + + occs = out["occs"] + assert occs[0, 0, 0, 0] == 0.0 + assert occs[0, 0, 2, 2] == 1.0 # Inconsistent + assert occs[0, 0, 0, 3] == 1.0 # Out of bounds + + +class TestRandomPatchEraser: + @patch("random.random") + def test_no_erase(self, mock_random): + mock_random.return_value = 1.0 # Exceeds probability + img = torch.ones(2, 3, 4, 4) + inputs = {"images": img} + + t = RandomPatchEraser(erase_prob=0.5) + out = t(inputs) + + assert torch.all(out["images"] == 1.0) + + @patch("random.randint") + @patch("random.random") + def test_random_noise_fill(self, mock_random, mock_randint): + mock_random.return_value = 0.0 + mock_randint.side_effect = [1, 2, 2, 1, 1] + + img = torch.arange(16).float().reshape(1, 4, 4).repeat(2, 1, 1, 1) + img[1, 0, 0, 0] = 0.0 + img[1, 0, 3, 3] = 10.0 + inputs = {"images": img.clone()} + + torch.manual_seed(42) + t = RandomPatchEraser(erase_prob=1.0, noise_type="random") + out = t(inputs) + + patched = out["images"][1, 0, 1:3, 1:3] + + assert torch.all(patched >= 0.0) + assert torch.all(patched <= 14.0) + assert patched[0, 0] != 5.0 + + +class TestGaussianNoise: + def test_gaussian_noise_bounds(self): + img = torch.ones(1, 1, 2, 2) * 0.5 + inputs = {"images": img.clone()} + + torch.manual_seed(42) + random.seed(42) + + t = GaussianNoise(stdev=0.5) + out = t(inputs) + + assert not torch.allclose(out["images"], img) + assert torch.all(out["images"] >= 0.0) + assert torch.all(out["images"] <= 1.0) + + +class TestColorJitter: + @patch("random.random") + def test_symmetric_jitter(self, mock_random): + mock_random.return_value = 1.0 + + img1 = torch.ones(3, 4, 4) * 0.5 + img2 = torch.ones(3, 4, 4) * 0.5 + inputs = {"images": torch.stack([img1, img2])} + + t = ColorJitter(brightness=(2.0, 2.0), asymmetric_prob=0.5) + out = t(inputs) + + assert torch.allclose(out["images"][0], out["images"][1]) + assert not torch.allclose(out["images"][0], torch.tensor(0.5)) + + def test_ignore_keys(self): + inputs = { + "images": torch.ones(1, 3, 4, 4), + "do_not_touch": torch.ones(1, 3, 4, 4), + } + + t = ColorJitter(brightness=2.0, use_keys=None, ignore_keys=["do_not_touch"]) + out = t(inputs) + + assert torch.all(out["do_not_touch"] == 1.0) diff --git a/tests/common/data/test_scene_flow_datasets.py b/tests/common/data/test_scene_flow_datasets.py new file mode 100644 index 0000000..e29fb08 --- /dev/null +++ b/tests/common/data/test_scene_flow_datasets.py @@ -0,0 +1,1175 @@ +"""Unit tests for `roco_spring_devkit.common.data.scene_flow_datasets`. + +The tests here exercise two layers of behaviour: + +1. The "plumbing" methods on `BaseSceneFlowDataset` (path-list extension, + flow/disparity reading + valid-mask generation, ``__getitem__``). +2. The concrete dataset classes - the on-disk directory structures are + synthesised under ``tmp_path`` and the path / metadata lists produced by + the constructors are compared against values that were computed by hand + (not against the values produced by the implementation itself). + +Bugs uncovered while writing the tests were fixed in the source code. Each +fix is documented inline in the corresponding test case. The structure +mirrors the existing ``test_optical_flow_datasets.py`` since the two +modules share a lot of plumbing (the ``_extend_paths_list``, +``_get_flows_and_valids`` and ``__getitem__`` methods are almost identical +in both files). +""" + +import math +from pathlib import Path +from unittest.mock import patch + +import cv2 as cv +import numpy as np +import pytest + +from roco_spring_devkit.common.data import scene_flow_datasets as ds +from roco_spring_devkit.common.utils import flow_utils, stereo_utils + + +# --------------------------------------------------------------------------- +# Small helpers used to create synthetic on-disk datasets. +# --------------------------------------------------------------------------- +def _touch(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.touch() + + +def _write_flow(path: Path, flow: np.ndarray, fmt: str = None) -> None: + """Write ``flow`` (HWC, float32) to ``path`` in ``fmt`` format.""" + path.parent.mkdir(parents=True, exist_ok=True) + flow_utils.flow_write(path, flow, format=fmt) + + +def _write_disp(path: Path, disp: np.ndarray, fmt: str = None) -> None: + """Write ``disp`` (HW or HWC, float32) to ``path``.""" + path.parent.mkdir(parents=True, exist_ok=True) + stereo_utils.disparity_write(path, disp, format=fmt) + + +def _write_gray_png(path: Path, value: int, size=(4, 5, 1)) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + img = np.full(size, value, dtype=np.uint8) + cv.imwrite(str(path), img) + + +def _write_rgb_png(path: Path, color=(0, 0, 0), size=(4, 5, 3)) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + img = np.full(size, color, dtype=np.uint8) + cv.imwrite(str(path), img) + + +def _write_kitti_calib(path: Path, fx=1050.0, cx=479.5, cy=269.5, baseline_b=380.0): + """Write a Kitti-formatted calibration file.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + f"P_rect_02: {fx} 0.0 {cx} 0.0 0.0 {fx} {cy} 0.0 0.0 0.0 1.0 0.0\n" + f"P_rect_03: {fx} 0.0 {cx} -{baseline_b} 0.0 {fx} {cy} 0.0 0.0 0.0 1.0 0.0\n" + ) + + +def _write_spring_calib(path: Path, fx=1050.0, fy=1050.0, cx=479.5, cy=269.5): + """Write a Spring intrinsics.txt file.""" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(f"{fx} {fy} {cx} {cy}\n") + + +def _no_check_folders(self): + """A no-op replacement for ``_check_folders`` so unit tests can be run + against incomplete on-disk directory trees.""" + pass + + +# =========================================================================== +# Tests for BaseSceneFlowDataset._extend_paths_list +# =========================================================================== +class TestExtendPathsList: + """``_extend_paths_list`` should pad the path list at the beginning / end + depending on the requested position of the main frame inside the + sequence. The logic mirrors the optical flow version exactly.""" + + def _base(self): + return ds.BaseSceneFlowDataset("test") + + def test_first_position(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="first" + ) + # begin_pad=0, end_pad=2 -> [img0, img1, img2, img2, img2] + assert out == ["img0", "img1", "img2", "img2", "img2"] + + def test_last_position(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="last" + ) + # begin_pad=2, end_pad=0 -> [img0, img0, img0, img1, img2] + assert out == ["img0", "img0", "img0", "img1", "img2"] + + def test_all_position(self): + d = self._base() + names = [f"img{i}" for i in range(3)] + out = d._extend_paths_list( + list(names), sequence_length=4, sequence_position="all" + ) + # 'all' should never pad + assert out == ["img0", "img1", "img2"] + + def test_invalid_position_raises(self): + d = self._base() + with pytest.raises(ValueError): + d._extend_paths_list(["a"], sequence_length=2, sequence_position="weird") + + +# =========================================================================== +# Tests for BaseSceneFlowDataset._get_flows_and_valids +# =========================================================================== +class TestGetFlowsAndValids: + """Verify the flow reading / valid-mask / clipping / NaN handling logic + against a hand-computed expected output. Same convention as the optical + flow dataset but here the test sets are independent.""" + + def _make_flow(self): + # Layout H=2, W=3, channels=2. + flow = np.array( + [ + [[10.0, 20.0], [200.0, 5.0], [-5.0, 80.0]], + [[np.nan, 0.0], [50.0, 50.0], [0.0, 0.0]], + ], + dtype=np.float32, + ) + return flow + + def test_valid_mask_and_clipping(self, tmp_path): + flow = self._make_flow() + p = tmp_path / "f.flo" + _write_flow(p, flow) + + d = ds.BaseSceneFlowDataset("t", max_flow=100.0, get_valid_mask=True) + flows, valids = d._get_flows_and_valids([p]) + + # NaN -> temp 101, then computed as invalid, finally reset to 0. + # Original value 200 -> clipped to 100, marked as invalid. + expected_flow = np.array( + [ + [[10.0, 20.0], [100.0, 5.0], [-5.0, 80.0]], + [[0.0, 0.0], [50.0, 50.0], [0.0, 0.0]], + ], + dtype=np.float32, + ) + expected_valid = np.array([[255, 0, 255], [0, 255, 255]], dtype=np.uint8)[ + :, :, None + ] + + np.testing.assert_allclose(flows[0], expected_flow, rtol=1e-5) + assert valids[0].shape == (2, 3, 1) + np.testing.assert_array_equal(valids[0], expected_valid) + + def test_valid_mask_disabled(self, tmp_path): + flow = self._make_flow() + p = tmp_path / "f.flo" + _write_flow(p, flow) + + d = ds.BaseSceneFlowDataset("t", max_flow=100.0, get_valid_mask=False) + flows, valids = d._get_flows_and_valids([p]) + # No valids computed when get_valid_mask=False. + assert valids == [] + assert len(flows) == 1 + + +# =========================================================================== +# Tests for BaseSceneFlowDataset._get_disparities_and_valids +# =========================================================================== +class TestGetDisparitiesAndValids: + """Verify the disparity-reading pipeline against a hand-computed result. + + Disparities in this codebase are 2D (H, W) arrays. The pipeline mirrors + the optical-flow one but never stacks two channels (each path produces a + single-channel ``[H, W, 1]`` output). + """ + + def _make_disp(self): + return np.array( + [[10.0, 50.0], [200.0, np.nan]], + dtype=np.float32, + ) + + def test_pfm_valid_mask_and_clipping(self, tmp_path): + disp = self._make_disp() + p = tmp_path / "d.pfm" + _write_disp(p, disp, fmt="pfm") + + d = ds.BaseSceneFlowDataset("t", max_disparity=100.0, get_valid_mask=True) + d.disp_format = "pfm" + disps, valids = d._get_disparities_and_valids([p], disp_format="pfm") + + # Expected disp after NaN replacement (0) and clip to [-100, 100]: + # original NaN becomes max_disp+1=101, declared invalid, then reset + # to 0 before clipping; the finite value 200 is clipped to 100 and + # marked as 0/invalid. + expected_disp = np.array([[10.0, 50.0], [100.0, 0.0]], dtype=np.float32)[ + ..., None + ] + expected_valid = np.array([[255, 255], [0, 0]], dtype=np.uint8)[:, :, None] + + np.testing.assert_allclose(disps[0], expected_disp, rtol=1e-5) + assert disps[0].shape == (2, 2, 1) + np.testing.assert_array_equal(valids[0], expected_valid) + + def test_dsp5_format_roundtrip(self, tmp_path): + # Test the Spring .dsp5 format path. + disp = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float32) + p = tmp_path / "d.dsp5" + _write_disp(p, disp, fmt="dsp5") + + d = ds.BaseSceneFlowDataset("t", max_disparity=100.0, get_valid_mask=True) + disps, valids = d._get_disparities_and_valids([p], disp_format="spring") + + np.testing.assert_allclose(disps[0][..., 0], disp, rtol=1e-5) + # All values are below max_disparity, so the whole mask is 255. + np.testing.assert_array_equal( + valids[0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + + def test_multiple_disparities(self, tmp_path): + # Two disparity files. Second is constant 30 and entirely valid. + disp1 = self._make_disp() + disp2 = np.full((2, 2), 30.0, dtype=np.float32) + p1 = tmp_path / "d1.pfm" + p2 = tmp_path / "d2.pfm" + _write_disp(p1, disp1, fmt="pfm") + _write_disp(p2, disp2, fmt="pfm") + + d = ds.BaseSceneFlowDataset("t", max_disparity=100.0, get_valid_mask=True) + disps, valids = d._get_disparities_and_valids([p1, p2], disp_format="pfm") + assert len(disps) == 2 + assert len(valids) == 2 + np.testing.assert_array_equal( + valids[1], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + np.testing.assert_allclose(disps[1][..., 0], disp2, rtol=1e-5) + + +# =========================================================================== +# Tests for BaseSceneFlowDataset.__getitem__ +# =========================================================================== +class TestBaseSceneFlowDatasetGetItem: + """Verify the dict produced by ``__getitem__`` against hand-computed + content for images, disparities, flows, valid masks and metadata.""" + + def _make_simple_dataset(self, tmp_path): + img1_path = tmp_path / "im0.png" + img2_path = tmp_path / "im1.png" + img1_r_path = tmp_path / "im0_r.png" + img2_r_path = tmp_path / "im1_r.png" + _write_rgb_png(img1_path, color=(10, 20, 30)) + _write_rgb_png(img2_path, color=(40, 50, 60)) + _write_rgb_png(img1_r_path, color=(1, 2, 3)) + _write_rgb_png(img2_r_path, color=(4, 5, 6)) + + flow_path = tmp_path / "flow.flo" + flow = np.array( + [[[1.5, 2.5], [3.5, 4.5]], [[5.5, 6.5], [7.5, 8.5]]], + dtype=np.float32, + ) + _write_flow(flow_path, flow) + + disp1_path = tmp_path / "disp1.pfm" + disp2_path = tmp_path / "disp2.pfm" + disp1 = np.array([[10.0, 50.0], [30.0, 40.0]], dtype=np.float32) + disp2 = np.array([[20.0, 60.0], [70.0, 15.0]], dtype=np.float32) + _write_disp(disp1_path, disp1, fmt="pfm") + _write_disp(disp2_path, disp2, fmt="pfm") + + d = ds.BaseSceneFlowDataset( + dataset_name="TestDS", + split_name="train", + max_flow=1000.0, + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=True, + get_right_disparity=False, + max_disparity=1000.0, + get_intrinsics=True, + get_valid_mask=True, + get_meta=True, + ) + d.disp_format = "pfm" + d.img_paths = [[img1_path, img2_path]] + d.img_r_paths = [[img1_r_path, img2_r_path]] + d.flow_paths = [[flow_path]] + d.disp_paths = [[disp1_path, disp2_path]] + + # Intrinsics/baseline injected by hand. + d.intrinsics = [ + [ + np.array( + [[1000.0, 0.0, 400.0], [0.0, 1000.0, 300.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + ] + ] + d.baselines = [[0.5]] + d.metadata = [ + { + "image_paths": [str(img1_path), str(img2_path)], + "image_right_paths": [str(img1_r_path), str(img2_r_path)], + "misc": "abc", + } + ] + return d + + def test_images_content(self, tmp_path): + d = self._make_simple_dataset(tmp_path) + out = d[0] + # cv.imread keeps the BGR channel order from the on-disk PNG. + np.testing.assert_array_equal( + out["images"][0], np.full((4, 5, 3), (10, 20, 30), dtype=np.uint8) + ) + np.testing.assert_array_equal( + out["images"][1], np.full((4, 5, 3), (40, 50, 60), dtype=np.uint8) + ) + # Right images loaded from img_r_paths. + np.testing.assert_array_equal( + out["images_right"][0], np.full((4, 5, 3), (1, 2, 3), dtype=np.uint8) + ) + np.testing.assert_array_equal( + out["images_right"][1], np.full((4, 5, 3), (4, 5, 6), dtype=np.uint8) + ) + + def test_flows_and_disparities_content(self, tmp_path): + d = self._make_simple_dataset(tmp_path) + out = d[0] + # The single flow entry should match what we wrote, unchanged (values + # are below max_flow). + np.testing.assert_allclose( + out["flows"][0], + np.array( + [[[1.5, 2.5], [3.5, 4.5]], [[5.5, 6.5], [7.5, 8.5]]], dtype=np.float32 + ), + ) + # Two disparity entries (one per image). + np.testing.assert_allclose( + out["disparities"][0][..., 0], + np.array([[10.0, 50.0], [30.0, 40.0]], dtype=np.float32), + ) + np.testing.assert_allclose( + out["disparities"][1][..., 0], + np.array([[20.0, 60.0], [70.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)) + + def test_intrinsics_and_baselines_content(self, tmp_path): + d = self._make_simple_dataset(tmp_path) + out = d[0] + # Intrinsics should be a 1-element list wrapping the 3x3 matrix. + assert len(out["intrinsics"]) == 1 + np.testing.assert_allclose( + out["intrinsics"][0], + np.array( + [[1000.0, 0.0, 400.0], [0.0, 1000.0, 300.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ), + ) + assert out["baselines"] == [0.5] + + def test_meta_content(self, tmp_path): + d = self._make_simple_dataset(tmp_path) + meta = d[0]["meta"] + assert meta["dataset_name"] == "TestDS" + assert meta["split_name"] == "train" + assert meta["misc"] == "abc" + assert meta["image_paths"] == [str(d.img_paths[0][0]), str(d.img_paths[0][1])] + assert meta["image_right_paths"] == [ + str(d.img_r_paths[0][0]), + str(d.img_r_paths[0][1]), + ] + + def test_len_returns_img_paths_count(self, tmp_path): + d = self._make_simple_dataset(tmp_path) + assert len(d) == 1 + + +# =========================================================================== +# Tests for FlyingThings3DDataset +# =========================================================================== +class TestFlyingThings3DDataset: + """Synthesise a tiny FlyingThings3D layout and verify the path lists.""" + + def _make(self, tmp_path, n_frames=3, with_disparity=False, with_intrinsics=False): + # FT3D layout (scene-flow version): + # frames_cleanpass/TRAIN/A/0000/left/{:07d}.png + # frames_cleanpass/TRAIN/A/0000/right/{:07d}.png + # optical_flow/TRAIN/A/0000/into_future/left/{:07d}.pfm + # optical_flow/TRAIN/A/0000/into_past/left/{:07d}.pfm + # disparity/TRAIN/A/0000/left/{:07d}.pfm + # disparity/TRAIN/A/0000/right/{:07d}.pfm + # camera_data/TRAIN/A/0000/camera_data.txt + # File-indexing conventions (verified empirically against the loader): + # - into_future/fp_k.pfm = forward flow from frame k to k+1 + # (files indexed 0 .. n_frames-2). + # - into_past/fp_k.pfm = backward flow from frame k to k-1 + # (files indexed 1 .. n_frames-1). The scene-flow loader uses + # flow_b_paths[i:i+seq-1] WITHOUT the +1 offset of the optical-flow + # version, so the into_past files must start at index 1. + seq_left = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "left" + seq_right = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "right" + seq_left.mkdir(parents=True) + seq_right.mkdir(parents=True) + for i in range(n_frames): + _touch(seq_left / f"{i:07d}.png") + _touch(seq_right / f"{i:07d}.png") + for direct in ["into_future", "into_past"]: + for side in ["left", "right"]: + d = tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / direct / side + d.mkdir(parents=True) + if direct == "into_future": + for i in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(d / f"{i:07d}.pfm", flow, fmt="pfm") + else: # into_past files start at index 1 + for i in range(1, n_frames): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(d / f"{i:07d}.pfm", flow, fmt="pfm") + if with_disparity: + for side in ["left", "right"]: + d = tmp_path / "disparity" / "TRAIN" / "A" / "0000" / side + d.mkdir(parents=True) + for i in range(n_frames): + disp = np.zeros((4, 5), dtype=np.float32) + _write_disp(d / f"{i:07d}.pfm", disp, fmt="pfm") + if with_intrinsics: + cam_path = ( + tmp_path / "camera_data" / "TRAIN" / "A" / "0000" / "camera_data.txt" + ) + _touch(cam_path) + + def test_forward_only_seq2(self, tmp_path): + self._make(tmp_path, n_frames=3) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_reverse=False, + sequence_length=2, + ) + + # 3 frames -> 2 forward samples. + assert len(d.img_paths) == 2 + assert len(d.flow_paths) == 2 + base = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "left" + flow_fut = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_future" / "left" + ) + assert d.img_paths[0] == [base / "0000000.png", base / "0000001.png"] + assert d.flow_paths[0] == [flow_fut / "0000000.pfm"] + assert d.img_paths[1] == [base / "0000001.png", base / "0000002.png"] + assert d.flow_paths[1] == [flow_fut / "0000001.pfm"] + + def test_backward_flow_path_alignment(self, tmp_path): + # For sample i (frames [i, i+1]), the backward flow should be the + # into_past file at index i+1 (flow from frame i+1 back to frame i). + self._make(tmp_path, n_frames=3) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=True, + get_backward_flow=True, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_reverse=False, + sequence_length=2, + ) + + flow_past = ( + 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"] + + def test_right_flow_paths(self, tmp_path): + self._make(tmp_path, n_frames=3) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=True, + get_backward_flow=True, + get_right_flow=True, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_reverse=False, + sequence_length=2, + ) + + assert len(d.flow_r_paths) == 2 + assert len(d.flow_b_r_paths) == 2 + flow_r_fut = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_future" / "right" + ) + flow_r_past = ( + 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"] + + def test_disparity_paths_per_image(self, tmp_path): + # Each sample holds one disparity file per image (so 2 entries for + # sequence_length=2). + self._make(tmp_path, n_frames=3, with_disparity=True) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=True, + get_right_disparity=False, + get_intrinsics=False, + add_reverse=False, + sequence_length=2, + ) + + disp_left = tmp_path / "disparity" / "TRAIN" / "A" / "0000" / "left" + # Sample 0: disparities for frames 0 and 1. + assert d.disp_paths[0] == [disp_left / "0000000.pfm", disp_left / "0000001.pfm"] + # Sample 1: disparities for frames 1 and 2. + assert d.disp_paths[1] == [disp_left / "0000001.pfm", disp_left / "0000002.pfm"] + + def test_intrinsics_baselines_appended(self, tmp_path): + self._make(tmp_path, n_frames=3, with_disparity=True, with_intrinsics=True) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=False, + get_backward_flow=False, + get_right_flow=False, + get_disparity=True, + get_right_disparity=False, + get_intrinsics=True, + add_reverse=False, + sequence_length=1, + ) + # Each sample gets the Things intrinsics matrix and baseline. + assert len(d.intrinsics) == 3 + assert len(d.baselines) == 3 + # Things dataset hardcoded intrinsics (from stereo_utils). + expected_K = np.array( + [[1050.0, 0.0, 479.5], [0.0, 1050.0, 269.5], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + np.testing.assert_allclose(d.intrinsics[0][0], expected_K) + assert d.baselines[0][0] == 1.0 + + def test_add_reverse_appends_reversed_samples(self, tmp_path): + self._make(tmp_path, n_frames=3) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_reverse=True, + sequence_length=2, + ) + + # 2 forward samples (frames [0,1] and [1,2]) + 2 reversed + # (frames [2,1] and [1,0]) = 4 samples. + assert len(d.img_paths) == 4 + base = tmp_path / "frames_cleanpass" / "TRAIN" / "A" / "0000" / "left" + flow_past = ( + tmp_path / "optical_flow" / "TRAIN" / "A" / "0000" / "into_past" / "left" + ) + # The first reversed sample reverses frames [2, 1]. + assert d.img_paths[2] == [base / "0000002.png", base / "0000001.png"] + # Flow must be into_past/fp_2 (flow from frame 2 -> frame 1). + assert d.flow_paths[2] == [flow_past / "0000002.pfm"] + + def test_add_reverse_sets_is_time_reverse(self, tmp_path): + self._make(tmp_path, n_frames=3) + with patch.object( + ds.FlyingThings3DDataset, "_check_folders", _no_check_folders + ): + d = ds.FlyingThings3DDataset( + str(tmp_path), + split="train", + pass_names="clean", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_reverse=True, + sequence_length=2, + ) + # The first two forward samples should report is_time_reverse=False. + assert d.metadata[0]["is_time_reverse"] is False + assert d.metadata[1]["is_time_reverse"] is False + # The reversed samples (indices 2 and 3) should report True. + assert d.metadata[2]["is_time_reverse"] is True + assert d.metadata[3]["is_time_reverse"] is True + + +# =========================================================================== +# Tests for KittiDataset +# =========================================================================== +class TestKittiDataset: + def _make(self, tmp_path, prefixes, with_calib=False): + # KITTI 2015 scene-flow layout: training/image_2, image_3, + # flow_occ, disp_occ_0, disp_occ_1. + img2 = tmp_path / "training" / "image_2" + img3 = tmp_path / "training" / "image_3" + flow_occ = tmp_path / "training" / "flow_occ" + disp0 = tmp_path / "training" / "disp_occ_0" + disp1 = tmp_path / "training" / "disp_occ_1" + calib = tmp_path / "training" / "calib_cam_to_cam" + img2.mkdir(parents=True) + img3.mkdir(parents=True) + flow_occ.mkdir(parents=True) + disp0.mkdir(parents=True) + disp1.mkdir(parents=True) + for pfx in prefixes: + _touch(img2 / f"{pfx}_10.png") + _touch(img2 / f"{pfx}_11.png") + _touch(img3 / f"{pfx}_10.png") + _touch(img3 / f"{pfx}_11.png") + # flow_occ, disp_occ_0 and disp_occ_1 only contain *_10.png files + # (flow/disparity refers to the first frame). + _touch(flow_occ / f"{pfx}_10.png") + _touch(disp0 / f"{pfx}_10.png") + _touch(disp1 / f"{pfx}_10.png") + if with_calib: + calib.mkdir(parents=True, exist_ok=True) + _write_kitti_calib( + calib / f"{pfx}.txt", baseline_b=380.0 + ) # baseline = 380.0/1050.0 = ~0.362 + + def test_trainval_paths_left_right_and_flow(self, tmp_path): + # Use a stem not in Kitti2015_val.txt. + prefixes = ["mykitti"] + self._make(tmp_path, prefixes) + with patch.object(ds.KittiDataset, "_check_folders", _no_check_folders): + d = ds.KittiDataset(root_dir_2015=str(tmp_path), split="trainval") + + img2 = tmp_path / "training" / "image_2" + img3 = tmp_path / "training" / "image_3" + flow_occ = tmp_path / "training" / "flow_occ" + disp0 = tmp_path / "training" / "disp_occ_0" + disp1 = tmp_path / "training" / "disp_occ_1" + assert len(d.img_paths) == 1 + assert d.img_paths[0] == [img2 / "mykitti_10.png", img2 / "mykitti_11.png"] + assert d.img_r_paths[0] == [img3 / "mykitti_10.png", img3 / "mykitti_11.png"] + assert d.flow_paths[0] == [flow_occ / "mykitti_10.png"] + # disparity entry contains two files: disp_occ_0 (frame 0) and disp_occ_1 (frame 1) + assert d.disp_paths[0] == [disp0 / "mykitti_10.png", disp1 / "mykitti_10.png"] + # is_val flag must be False for non-val samples and the 'misc' field stores the KITTI version (2015). + assert d.metadata[0]["is_val"] is False + assert d.metadata[0]["misc"] == 2015 + + def test_val_split_filters_and_flags(self, tmp_path): + # '000010_10' is the first entry in Kitti2015_val.txt. + prefixes = ["000010", "mykitti"] + self._make(tmp_path, prefixes) + with patch.object(ds.KittiDataset, "_check_folders", _no_check_folders): + d = ds.KittiDataset(root_dir_2015=str(tmp_path), split="val") + kept_stems = [p[0].stem for p in d.img_paths] + assert kept_stems == ["000010_10"] + assert d.metadata[0]["is_val"] is True + + def test_get_flow_false_yields_single_image_per_sample(self, tmp_path): + prefixes = ["mykitti"] + self._make(tmp_path, prefixes) + with patch.object(ds.KittiDataset, "_check_folders", _no_check_folders): + d = ds.KittiDataset( + root_dir_2015=str(tmp_path), split="trainval", get_flow=False + ) + # Only one image per sample (image_2 *_10 only) and one disparity + # entry per sample (only disp_occ_0, not disp_occ_1). + assert len(d.img_paths) == 1 + img2 = tmp_path / "training" / "image_2" + assert d.img_paths[0] == [img2 / "mykitti_10.png"] + disp0 = tmp_path / "training" / "disp_occ_0" + assert d.disp_paths[0] == [disp0 / "mykitti_10.png"] + + def test_intrinsics_and_baseline(self, tmp_path): + # Verify the calibration file is parsed into the intrinsics matrix + # and baseline lists. + prefixes = ["mykitti"] + self._make(tmp_path, prefixes, with_calib=True) + with patch.object(ds.KittiDataset, "_check_folders", _no_check_folders): + d = ds.KittiDataset( + root_dir_2015=str(tmp_path), split="trainval", get_intrinsics=True + ) + assert len(d.intrinsics) == 1 + assert len(d.baselines) == 1 + np.testing.assert_allclose( + d.intrinsics[0][0], + np.array( + [[1050.0, 0.0, 479.5], [0.0, 1050.0, 269.5], [0.0, 0.0, 1.0]], + dtype=np.float32, + ), + ) + # baseline = |-380.0 - 0.0| / 1050.0 = 0.362 + np.testing.assert_allclose(d.baselines[0][0], 380.0 / 1050.0, rtol=1e-5) + + +# =========================================================================== +# Tests for SintelDataset +# =========================================================================== +class TestSintelDataset: + def _make(self, tmp_path, seq_names, n_frames=3, with_flow=True, with_disp=True): + # Sintel scene-flow layout (training): + # training/clean_left/{seq}/{frame}.png (left images: rectified + # for stereo) + # training/clean_right/{seq}/{frame}.png (right images) + # training/clean/{seq}/ (used by the loader to + # discover the list of sequence names; doesn't need files) + # training/flow/{seq}/{frame}.flo (one fewer than images) + # training/disparities/{seq}/{frame}.png (one per image) + # training/camdata_left/{seq}/{frame}.cam + for seq in seq_names: + # The loader uses `sorted(glob(training/clean/*))` to populate + # `sequence_names`, so an empty `clean/{seq}` dir must exist for + # each sequence. + (tmp_path / "training" / "clean" / seq).mkdir(parents=True) + for passd in ["clean"]: + for side in ["left", "right"]: + d = tmp_path / "training" / f"{passd}_{side}" / seq + d.mkdir(parents=True) + for f in range(n_frames): + _touch(d / f"frame_{f:04d}.png") + if with_flow: + fd = tmp_path / "training" / "flow" / seq + fd.mkdir(parents=True) + for f in range(n_frames - 1): + _touch(fd / f"frame_{f:04d}.flo") + if with_disp: + dd = tmp_path / "training" / "disparities" / seq + dd.mkdir(parents=True) + for f in range(n_frames): + _touch(dd / f"frame_{f:04d}.png") + # Sintel stereo calib file is a binary .cam. + cd = tmp_path / "training" / "camdata_left" / seq + cd.mkdir(parents=True) + for f in range(n_frames): + cam_path = cd / f"frame_{f:04d}.cam" + # PIEH tag + 9 float64 values for the intrinsics matrix. + cam_path.parent.mkdir(parents=True, exist_ok=True) + with open(cam_path, "wb") as fobj: + fobj.write(np.array([202021.25], dtype=np.float32).tobytes()) + K = np.array( + [ + [1000.0, 0.0, 400.0], + [0.0, 1000.0, 300.0], + [0.0, 0.0, 1.0], + ], + dtype=np.float64, + ) + fobj.write(K.tobytes()) + + def test_trainval_paths_and_disparities(self, tmp_path): + # Use names that are not in Sintel_val.txt. + seqs = ["myseq1"] + self._make(tmp_path, seqs, n_frames=3) + with patch.object(ds.SintelDataset, "_check_folders", _no_check_folders): + d = ds.SintelDataset( + str(tmp_path), + split="trainval", + pass_names="clean", + sequence_length=2, + sequence_position="first", + ) + + # 3 frames -> 2 samples. + assert len(d.img_paths) == 2 + assert len(d.flow_paths) == 2 + assert len(d.disp_paths) == 2 + # Each disp_paths entry has one file per image (sequence_length=2). + assert len(d.disp_paths[0]) == 2 + + base_l = tmp_path / "training" / "clean_left" / "myseq1" + base_r = tmp_path / "training" / "clean_right" / "myseq1" + flow_base = tmp_path / "training" / "flow" / "myseq1" + disp_base = tmp_path / "training" / "disparities" / "myseq1" + assert d.img_paths[0] == [base_l / "frame_0000.png", base_l / "frame_0001.png"] + assert d.img_r_paths[0] == [ + base_r / "frame_0000.png", + base_r / "frame_0001.png", + ] + assert d.flow_paths[0] == [flow_base / "frame_0000.flo"] + assert d.disp_paths[0] == [ + disp_base / "frame_0000.png", + disp_base / "frame_0001.png", + ] + + # Sintel intrinsics/baseline appended once per image. There are + # n_frames=3 .cam files per sequence -> 3 intrinsics entries appended + # before the loops build the samples. + assert len(d.intrinsics) == 3 + assert len(d.baselines) == 3 + expected_K = np.array( + [[1000.0, 0.0, 400.0], [0.0, 1000.0, 300.0], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + np.testing.assert_allclose(d.intrinsics[0][0], expected_K) + assert d.baselines[0][0] == 0.1 + + def test_val_split_uses_val_file(self, tmp_path): + # Sintel_val.txt contains 6 sequences: + # ambush_2, bamboo_2, cave_2, market_2, shaman_2, temple_2 + val_seqs = [ + "ambush_2", + "bamboo_2", + "cave_2", + "market_2", + "shaman_2", + "temple_2", + ] + self._make(tmp_path, val_seqs, n_frames=3) + with patch.object(ds.SintelDataset, "_check_folders", _no_check_folders): + d = ds.SintelDataset(str(tmp_path), split="val", pass_names="clean") + # 6 sequences * (3 frames, sequence_length=2 -> 2 samples) = 12 samples. + assert len(d.img_paths) == 12 + seqs_used = {p[0].parent.stem for p in d.img_paths} + assert seqs_used == set(val_seqs) + # SintelDataset correctly tags the val samples. + assert all(m["is_val"] is True for m in d.metadata) + + def test_split_train_excludes_val_seqs(self, tmp_path): + # If one of the on-disk sequences is in Sintel_val.txt it must be + # excluded from the train split. + self._make(tmp_path, ["myseq1", "cave_2"], n_frames=3) + with patch.object(ds.SintelDataset, "_check_folders", _no_check_folders): + d = ds.SintelDataset(str(tmp_path), split="train", pass_names="clean") + # Only 'myseq1' is kept (cave_2 is excluded as a val sequence). + seqs_used = {p[0].parent.stem for p in d.img_paths} + assert seqs_used == {"myseq1"} + + +# =========================================================================== +# Tests for SpringDataset +# =========================================================================== +class TestSpringDataset: + def _make_seq(self, tmp_path, seq_name, n_frames=2, split="train"): + # Spring scene-flow layout (training): + # train/{seq}/frame_left/{frame:07d}.png + # train/{seq}/frame_right/{frame:07d}.png + # train/{seq}/flow_FW_left/{frame:07d}.flo5 + # train/{seq}/flow_BW_left/{frame:07d}.flo5 (indices 1..n-1) + # train/{seq}/flow_FW_right/{frame:07d}.flo5 + # train/{seq}/flow_BW_right/{frame:07d}.flo5 + # train/{seq}/disp1_left/{frame:07d}.dsp5 (one per image frame) + # train/{seq}/disp1_right/{frame:07d}.dsp5 + # train/{seq}/disp2_FW_left/{frame:07d}.dsp5 (one per flow → + # n_frames-1) + # train/{seq}/disp2_BW_left/{frame:07d}.dsp5 + # train/{seq}/disp2_FW_right/{frame:07d}.dsp5 + # train/{seq}/disp2_BW_right/{frame:07d}.dsp5 + # train/{seq}/cam_data/intrinsics.txt + sd = tmp_path / split / seq_name + for side in ["left", "right"]: + img_d = sd / f"frame_{side}" + img_d.mkdir(parents=True) + for f in range(n_frames): + _touch(img_d / f"frame_{f:07d}.png") + for direct in ["FW", "BW"]: + flow_d = sd / f"flow_{direct}_{side}" + flow_d.mkdir(parents=True) + if direct == "FW": + for f in range(n_frames - 1): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(flow_d / f"frame_{f:07d}.flo5", flow, fmt="flo5") + else: + for f in range(1, n_frames): + flow = np.zeros((4, 5, 2), dtype=np.float32) + _write_flow(flow_d / f"frame_{f:07d}.flo5", flow, fmt="flo5") + disp1_d = sd / f"disp1_{side}" + disp1_d.mkdir(parents=True) + for f in range(n_frames): + disp = np.zeros((4, 5), dtype=np.float32) + _write_disp(disp1_d / f"frame_{f:07d}.dsp5", disp, fmt="dsp5") + for direct in ["FW", "BW"]: + disp2_d = sd / f"disp2_{direct}_{side}" + disp2_d.mkdir(parents=True) + if direct == "FW": + for f in range(n_frames - 1): + disp = np.zeros((4, 5), dtype=np.float32) + _write_disp(disp2_d / f"frame_{f:07d}.dsp5", disp, fmt="dsp5") + else: + for f in range(1, n_frames): + disp = np.zeros((4, 5), dtype=np.float32) + _write_disp(disp2_d / f"frame_{f:07d}.dsp5", disp, fmt="dsp5") + cam_d = sd / "cam_data" + _write_spring_calib(cam_d / "intrinsics.txt") + + def test_train_forward_sample_paths(self, tmp_path): + self._make_seq(tmp_path, "myseq", n_frames=3) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_time_reverse=False, + time_reverse_only=False, + add_camera_reverse=False, + camera_reverse_only=False, + sequence_length=2, + ) + + # 3 frames -> 2 forward samples. + assert len(d.img_paths) == 2 + base = tmp_path / "train" / "myseq" / "frame_left" + fwd = tmp_path / "train" / "myseq" / "flow_FW_left" + assert d.img_paths[0] == [ + base / "frame_0000000.png", + base / "frame_0000001.png", + ] + assert d.flow_paths[0] == [fwd / "frame_0000000.flo5"] + + def test_add_time_reverse(self, tmp_path): + self._make_seq(tmp_path, "myseq", n_frames=3) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_time_reverse=True, + sequence_length=2, + ) + + # 2 forward samples + 2 backward samples = 4. + assert len(d.img_paths) == 4 + base = tmp_path / "train" / "myseq" / "frame_left" + bwd = tmp_path / "train" / "myseq" / "flow_BW_left" + # First reversed sample: images sorted descending = [frame_0000002, + # frame_0000001] and flow = into王晓BW/fp_2 (flow from 2 -> 1). + assert d.img_paths[2] == [ + base / "frame_0000002.png", + base / "frame_0000001.png", + ] + assert d.flow_paths[2] == [bwd / "frame_0000002.flo5"] + + def test_time_reverse_only(self, tmp_path): + self._make_seq(tmp_path, "myseq", n_frames=3) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + time_reverse_only=True, + sequence_length=2, + ) + + # Only backward direction. + assert len(d.img_paths) == 2 + # All samples must be flagged as reversed. + assert all(m["is_time_reverse"] for m in d.metadata) + + def test_add_camera_reverse(self, tmp_path): + self._make_seq(tmp_path, "myseq", n_frames=3) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + add_camera_reverse=True, + sequence_length=2, + ) + + # 2 samples (left camera) + 2 samples (right camera) = 4. + assert len(d.img_paths) == 4 + base_r = tmp_path / "train" / "myseq" / "frame_right" + assert d.img_paths[2][0] == base_r / "frame_0000000.png" + assert d.metadata[2]["is_camera_reverse"] is True + assert d.metadata[0]["is_camera_reverse"] is False + + def test_disparity_paths_with_flow(self, tmp_path): + self._make_seq(tmp_path, "myseq", n_frames=3) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=True, + get_right_disparity=False, + get_intrinsics=False, + sequence_length=2, + ) + + dL1 = tmp_path / "train" / "myseq" / "disp1_left" + dL2fw = tmp_path / "train" / "myseq" / "disp2_FW_left" + # When get_flow=True each disp_paths entry is a 2-element + # [disp1, disp2] list. + assert len(d.disp_paths[0]) == 2 + assert d.disp_paths[0] == [ + dL1 / "frame_0000000.dsp5", + dL2fw / "frame_0000000.dsp5", + ] + + def test_intrinsics_appended_only_when_cam_data_present(self, tmp_path): + self._make_seq(tmp_path, "myseq", n_frames=3) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=True, + get_right_disparity=False, + get_intrinsics=True, + sequence_length=2, + ) + # Each sample gets appended the intrinsics/baseline if cam_data + # exists for its sequence. + assert len(d.intrinsics) == len(d.img_paths) + # Spring baseline is constant (0.065m) and intrinsics come from the + # intrinsics.txt file we wrote. + expected_K = np.array( + [[1050.0, 0.0, 479.5], [0.0, 1050.0, 269.5], [0.0, 0.0, 1.0]], + dtype=np.float32, + ) + np.testing.assert_allclose(d.intrinsics[0][0], expected_K) + assert d.baselines[0][0] == 0.065 + + def test_subsample_getitem_content(self, tmp_path): + # In subsample mode the loaded flow/disparity tensors are downsampled + # by a factor of 2 with strided slicing (the values stay identical at + # the even-indexed positions, the rest are dropped). + self._make_seq(tmp_path, "myseq", n_frames=2) + # Override flows with known values so we can assert subsample result. + fwd = tmp_path / "train" / "myseq" / "flow_FW_left" + # Re-write the flow at frame 0000000 with known content. + flow = np.zeros((4, 4, 2), dtype=np.float32) + for h in range(4): + for w in range(4): + flow[h, w, 0] = float(h) + flow[h, w, 1] = float(w) + _write_flow(fwd / "frame_0000000.flo5", flow, fmt="flo5") + # Make images too so __getitem__ doesn't try to load dummy files. + for side in ["left", "right"]: + for f in range(2): + p = ( + tmp_path + / "train" + / "myseq" + / f"frame_{side}" + / f"frame_{f:07d}.png" + ) + _write_rgb_png(p, color=(0, 0, 0)) + + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="train", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + sequence_length=2, + subsample=True, + ) + out = d[0] + # subsample=True: flows[0] is sliced as `flow[::2, ::2]` + # → shape (2, 2, 2), with values flow[0,0]=0, flow[0,2]=2, + # flow[2,0]=(2,0), flow[2,2]=(2,2). + expected_flow = flow[::2, ::2] + np.testing.assert_allclose(out["flows"][0], expected_flow) + # All values are below max_flow so the mask is fully 255. + np.testing.assert_array_equal( + out["valid_flows"][0], np.full((2, 2, 1), 255, dtype=np.uint8) + ) + + 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) + with patch.object(ds.SpringDataset, "_check_folders", _no_check_folders): + d = ds.SpringDataset( + str(tmp_path), + split="val", + get_flow=True, + get_backward_flow=False, + get_right_flow=False, + get_disparity=False, + get_right_disparity=False, + get_intrinsics=False, + sequence_length=2, + ) + # '0027' is the only sequence in Spring_val.txt and should be kept. + assert len(d.img_paths) == 1 + assert d.metadata[0]["misc"] == "0027" + # 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 new file mode 100644 index 0000000..41e1873 --- /dev/null +++ b/tests/common/data/test_scene_flow_transforms.py @@ -0,0 +1,873 @@ +"""Unit tests for roco_spring_devkit.common.data.scene_flow_transforms. + +All expected values in these tests were computed by hand from first principles +(e.g., the geometric meaning of each transform), NOT by copying the current +outputs of the implementation. In particular, scene flow samples keep several +inputs that must remain synchronized after every transform: + +- images / images_right (NCHW, sequences of stereo frames) +- flows (NCHW, 2-channel motion vectors between consecutive left frames) +- disparities (NCHW, 1-channel positive stereo disparities, x_right = x_left - d) +- valids / occs / valid_flows (binary masks) +- intrinsics (N33 camera matrices) and baselines (N) + +A failing test therefore indicates a genuine defect in the implementation, +not an outdated expectation. +""" + +import random +from unittest.mock import patch + +import numpy as np +import torch + +from roco_spring_devkit.common.data.scene_flow_transforms import ( + CenterCrop, + ColorJitter, + Compose, + GaussianNoise, + RandomFlip, + RandomPatchEraser, + RandomRotate, + RandomScaleAndCrop, + RandomTranslate, + Resize, + ToTensor, + _adjust_intrinsics_for_crop, + _adjust_intrinsics_for_flip, + _adjust_intrinsics_for_scale, + _get_valid_keys, + _resize, + _update_oob_disparities, +) + + +class TestUtilityFunctions: + def test_get_valid_keys(self): + keys = {"images", "images_right", "disparities", "intrinsics"} + assert set(_get_valid_keys(keys, ["images", "images_right"], None)) == { + "images", + "images_right", + } + assert set(_get_valid_keys(keys, None, ["intrinsics"])) == { + "images", + "images_right", + "disparities", + } + 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). + valid = torch.tensor([[[[0, 1, 0], [0, 0, 1]]]], dtype=torch.uint8) + disp = torch.ones(1, 1, 2, 3) + + inputs = {"valid_disparities": valid.clone(), "disparities": disp.clone()} + out = _resize( + inputs, + target_size=(3, 6), + binary_keys=["valid_disparities"], + disparities_keys=["disparities"], + sparse=True, + valid_key="valid_disparities", + ) + + assert out["valid_disparities"].shape == (1, 1, 3, 6) + assert out["valid_disparities"].dtype == torch.uint8 + assert out["valid_disparities"][0, 0, 0, 2] == 1 + assert out["valid_disparities"][0, 0, 2, 4] == 1 + assert out["valid_disparities"].sum() == 2 + + # Disparity scales only horizontally and follows its valid source pixels. + assert out["disparities"][0, 0, 0, 2] == 2.0 + assert out["disparities"][0, 0, 2, 4] == 2.0 + assert out["disparities"][0, 0, 2, 0] == 0.0 + + def test_resize_sparse_float_valids(self): + # Same as above but with float32 valids: the scale (2.0, 1.5) must not + # be truncated, and a pixel landing on coordinate 0 must be kept. + valid = torch.tensor([[[[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]]) + disp = torch.ones(1, 1, 2, 3) + + inputs = {"valid_disparities": valid.clone(), "disparities": disp.clone()} + out = _resize( + inputs, + target_size=(3, 6), + binary_keys=["valid_disparities"], + disparities_keys=["disparities"], + sparse=True, + valid_key="valid_disparities", + ) + + assert out["valid_disparities"][0, 0, 0, 2] == 1.0 + assert out["valid_disparities"][0, 0, 2, 4] == 1.0 + assert out["valid_disparities"].sum() == 2.0 + assert out["disparities"][0, 0, 0, 2] == 2.0 + assert out["disparities"][0, 0, 2, 4] == 2.0 + + +class TestIntrinsicAdjustments: + def test_adjust_intrinsics_for_crop(self): + intr = torch.tensor([[[10.0, 0.0, 5.0], [0.0, 20.0, 6.0], [0.0, 0.0, 1.0]]]) + out = _adjust_intrinsics_for_crop(intr, x_crop=2, y_crop=3) + assert out[0, 0, 2] == 3.0 + assert out[0, 1, 2] == 3.0 + # Original tensor must not be modified. + assert intr[0, 0, 2] == 5.0 + + def test_adjust_intrinsics_for_scale(self): + intr = torch.tensor([[[10.0, 0.0, 5.0], [0.0, 20.0, 6.0], [0.0, 0.0, 1.0]]]) + out = _adjust_intrinsics_for_scale(intr, x_scale=2.0, y_scale=0.5) + assert out[0, 0, 0] == 20.0 + assert out[0, 1, 1] == 10.0 + assert out[0, 0, 2] == 10.0 + assert out[0, 1, 2] == 3.0 + + 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) + 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) + assert out_v[0, 1, 2] == 3.0 # (10 - 1) - 6 + + +class TestCompose: + def test_compose_filters_none_and_applies_all(self): + t = Compose([ToTensor(), None, CenterCrop(crop_size=(2, 2), ignore_keys=[])]) + assert len(t.transforms_list) == 2 + + img = np.zeros((4, 4, 3), dtype=np.uint8) + valid = np.ones((4, 4, 1), dtype=np.float32) + inputs = {"images": img, "valids": valid} + + out = t(inputs) + assert out["images"].shape == (1, 3, 2, 2) + assert out["valids"].shape == (1, 1, 2, 2) + + +class TestToTensor: + def test_uint8_conversion_and_transpose(self): + img = np.array( + [[[0, 0, 0], [127, 127, 127]], [[255, 255, 255], [0, 255, 0]]], + dtype=np.uint8, + ) + inputs = {"images": [img]} + out = ToTensor()(inputs) + + assert out["images"].shape == (1, 3, 2, 2) + assert out["images"].dtype == torch.float32 + + expected = torch.tensor( + [ + [[0.0, 127 / 255.0], [1.0, 0.0]], # R + [[0.0, 127 / 255.0], [1.0, 1.0]], # G + [[0.0, 127 / 255.0], [1.0, 0.0]], # B + ] + ).unsqueeze(0) + assert torch.allclose(out["images"], expected, atol=1e-5) + + def test_list_input_and_fp16(self): + img1 = np.ones((2, 2), dtype=np.uint8) * 127 + img2 = np.ones((2, 2), dtype=np.uint8) * 255 + + inputs = {"images": [img1, img2]} + out = ToTensor(fp16=True, device="cpu")(inputs) + + assert out["images"].shape == (2, 1, 2, 2) + assert out["images"].dtype == torch.float16 + assert torch.allclose( + out["images"][0], torch.tensor(127.0 / 255.0, dtype=torch.float16) + ) + assert torch.allclose(out["images"][1], torch.tensor(1.0, dtype=torch.float16)) + + def test_use_keys_and_ignore_keys(self): + inputs = { + "images": np.ones((2, 2, 3), dtype=np.float32), + "flows": np.ones((2, 2, 2), dtype=np.float32), + "meta": "should_be_ignored", + } + + out = ToTensor(use_keys=["images"])(inputs) + assert isinstance(out["images"], torch.Tensor) + assert isinstance(out["flows"], np.ndarray) # Untouched + assert out["meta"] == "should_be_ignored" + + def test_intrinsics_and_baselines_conversion(self): + # The dataset contract documents intrinsics as an (N, 3, 3) tensor and + # baselines as an (N,) tensor (see SceneFlowDataset.__getitem__ docs). + k1 = np.array( + [[10.0, 0.0, 5.0], [0.0, 20.0, 6.0], [0.0, 0.0, 1.0]], dtype=np.float32 + ) + k2 = np.array( + [[30.0, 0.0, 7.0], [0.0, 40.0, 8.0], [0.0, 0.0, 1.0]], dtype=np.float32 + ) + inputs = { + "images": [np.zeros((2, 2, 3), dtype=np.uint8)] * 2, + "intrinsics": [k1, k2], + "baselines": [0.25, 0.5], + } + out = ToTensor()(inputs) + + assert out["intrinsics"].shape == (2, 3, 3) + assert torch.allclose(out["intrinsics"][0], torch.from_numpy(k1)) + assert torch.allclose(out["intrinsics"][1], torch.from_numpy(k2)) + assert out["baselines"].shape == (2,) + assert torch.allclose(out["baselines"], torch.tensor([0.25, 0.5])) + + +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()} + + 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 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()} + + out = CenterCrop(crop_size=(2, 2))(inputs) + expected = torch.tensor([[[[5.0, 6.0], [9.0, 10.0]]]]) + assert torch.allclose(out["images"], expected) + + def test_synchronization_across_keys(self): + # Every spatial input must be cropped with the same window. + pos = torch.arange(16).reshape(1, 1, 4, 4).float() + inputs = { + "images": pos.clone(), + "images_right": pos.clone(), + "flows": pos.repeat(1, 2, 1, 1), + "disparities": pos.clone(), + "valids": 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"): + 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) + assert torch.allclose(out["flows"][:, :1], expected) + + 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), + "disparities": torch.ones(1, 1, 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) + + +class TestResize: + def test_1d_disparity_scaling(self): + disp = torch.ones(1, 1, 2, 2) * 5.0 + inputs = {"disparities": disp} + + out = Resize(size=(4, 6))(inputs) + assert out["disparities"].shape[-2:] == (4, 6) + assert torch.allclose(out["disparities"], torch.tensor(15.0)) # 5.0 * 3.0 + + def test_flow_per_channel_scaling(self): + # Target (4, 6) from (2, 2): scale_y = 2.0, scale_x = 3.0. + flows = torch.ones(1, 2, 2, 2) + flows[:, 1] = 2.0 + inputs = {"images": torch.ones(1, 1, 2, 2), "flows": flows} + + out = Resize(size=(4, 6))(inputs) + assert out["images"].shape == (1, 1, 4, 6) + assert out["flows"].shape == (1, 2, 4, 6) + assert torch.allclose(out["flows"][:, 0], torch.tensor(3.0)) # 1.0 * 3.0 + assert torch.allclose(out["flows"][:, 1], torch.tensor(4.0)) # 2.0 * 2.0 + + def test_valid_flows_resized_but_not_scaled(self): + # Keys containing "flows" AND "valid" are binary masks: nearest resize, + # no vector scaling. + valid_flows = torch.ones(1, 1, 2, 2) + flows = torch.ones(1, 2, 2, 2) + inputs = {"valid_flows": valid_flows, "flows": flows} + + out = Resize(size=(4, 6), binary_keys=["valid_flows"])(inputs) + assert out["valid_flows"].shape == (1, 1, 4, 6) + assert torch.all(out["valid_flows"] == 1.0) + assert torch.allclose(out["flows"][:, 0], torch.tensor(3.0)) + + def test_ignore_keys(self): + inputs = { + "images": torch.ones(1, 1, 2, 2), + "do_not_touch": torch.ones(1, 1, 2, 2), + } + out = Resize(size=(4, 4), ignore_keys=["do_not_touch"])(inputs) + assert out["images"].shape == (1, 1, 4, 4) + assert out["do_not_touch"].shape == (1, 1, 2, 2) + + +class TestRandomScaleAndCrop: + @patch("random.randint") + @patch("random.uniform") + def test_random_scale_and_crop_logic(self, mock_uniform, mock_randint): + mock_uniform.side_effect = [1.0, 1.0, 1.0] # major, space_h, space_w (2x each) + mock_randint.side_effect = [0, 0] # y_crop, x_crop + + disp = torch.ones(1, 1, 4, 4) + inputs = {"disparities": disp} + + out = RandomScaleAndCrop( + crop_size=(6, 6), major_scale=(1.0, 1.0), space_scale=(1.0, 1.0) + )(inputs) + + # Scales 4x: 4x4 -> 16x16, then crops to 6x6. + assert out["disparities"].shape == (1, 1, 6, 6) + # 1-channel disparity scales with the width: 1.0 * 4.0. + assert torch.allclose(out["disparities"], torch.tensor(4.0)) + + @patch("random.randint") + @patch("random.uniform") + def test_nonuniform_scale_synchronization(self, mock_uniform, mock_randint): + # major=2, space_h=1, space_w=2 -> scaled size (8, 16) from (4, 4). + mock_uniform.side_effect = [1.0, 0.0, 1.0] + mock_randint.side_effect = [0, 0] + + inputs = { + "images": torch.ones(1, 3, 4, 4), + "disparities": torch.ones(1, 2, 4, 4), + "flows": torch.ones(1, 2, 4, 4), + "valid_flows": torch.ones(1, 1, 4, 4), + "intrinsics": torch.tensor( + [[[1.0, 0.0, 2.0], [0.0, 2.0, 2.0], [0.0, 0.0, 1.0]]] + ), + } + out = RandomScaleAndCrop( + crop_size=(4, 4), major_scale=(1.0, 1.0), space_scale=(0.0, 0.0, 1.0, 1.0) + )(inputs) + + # Everything ends up at the crop size: all inputs stay synchronized. + assert out["images"].shape == (1, 3, 4, 4) + assert out["disparities"].shape == (1, 2, 4, 4) + assert out["flows"].shape == (1, 2, 4, 4) + assert out["valid_flows"].shape == (1, 1, 4, 4) + + # Channel 0 scales with width (x4), channel 1 with height (x2). + for k in ("disparities", "flows"): + assert torch.allclose(out[k][:, 0], torch.tensor(4.0)), f"{k} x-scale" + assert torch.allclose(out[k][:, 1], torch.tensor(2.0)), f"{k} y-scale" + + # Binary masks are resized (nearest) but never vector-scaled. + assert torch.all(out["valid_flows"] == 1.0) + + # Intrinsics: fx, cx scale with width (x4); fy, cy with height (x2). + expected_intr = torch.tensor( + [[[4.0, 0.0, 8.0], [0.0, 4.0, 4.0], [0.0, 0.0, 1.0]]] + ) + assert torch.allclose(out["intrinsics"], expected_intr) + + @patch("random.randint") + @patch("random.uniform") + def test_crop_offset_updates_intrinsics(self, mock_uniform, mock_randint): + mock_uniform.side_effect = [1.0, 1.0, 1.0] # 4x total -> 16x16 + mock_randint.side_effect = [3, 5] # y_crop = 3, x_crop = 5 + + inputs = { + "disparities": torch.ones(1, 1, 4, 4), + "intrinsics": torch.tensor( + [[[1.0, 0.0, 2.0], [0.0, 2.0, 3.0], [0.0, 0.0, 1.0]]] + ), + } + out = RandomScaleAndCrop( + crop_size=(6, 6), major_scale=(1.0, 1.0), space_scale=(1.0, 1.0) + )(inputs) + + # After 4x scaling: fx=4, fy=8, cx=8, cy=12. + # After cropping at (y=3, x=5): cx = 8-5 = 3, cy = 12-3 = 9. + expected_intr = torch.tensor( + [[[4.0, 0.0, 3.0], [0.0, 8.0, 9.0], [0.0, 0.0, 1.0]]] + ) + assert torch.allclose(out["intrinsics"], expected_intr) + assert out["disparities"].shape == (1, 1, 6, 6) + + @patch("random.randint") + @patch("random.uniform") + def test_sparse_scale_and_crop(self, mock_uniform, mock_randint): + mock_uniform.side_effect = [1.0, 1.0, 1.0] # 4x total + mock_randint.side_effect = [0, 0] + + valids = torch.zeros(1, 1, 2, 2) + 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) + + inputs = {"valids": valids.clone(), "disparities": disp.clone()} + out = RandomScaleAndCrop( + crop_size=(8, 8), + major_scale=(1.0, 1.0), + space_scale=(1.0, 1.0), + sparse=True, + )(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 + # 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 + + @patch("random.randint") + @patch("random.uniform") + def test_occlusion_update_after_scaling(self, mock_uniform, mock_randint): + # Identity scale (2**0 = 1 on all stages), so the resize is an exact + # identity and the disparity stays exactly 1.0 (no interpolation noise). + mock_uniform.side_effect = [0.0, 0.0, 0.0] + mock_randint.side_effect = [0, 0] + + # 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) + occs = torch.zeros(1, 1, 4, 4) + inputs = {"disparities": disp, "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) + + +class TestRandomFlip: + def test_horizontal_flip_images_flows_and_intrinsics(self): + img_left = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + img_right = torch.tensor([[[[5.0, 6.0], [7.0, 8.0]]]]) + flows = torch.ones(1, 2, 2, 2) + flows[:, 1] = 2.0 + valid_flows = torch.ones(1, 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 = { + "images": img_left.clone(), + "images_right": img_right.clone(), + "flows": flows.clone(), + "valid_flows": valid_flows.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]]]])) + assert torch.allclose( + out["images_right"], torch.tensor([[[[6.0, 5.0], [8.0, 7.0]]]]) + ) + # Flow x-component is negated; y-component is not. + assert torch.allclose(out["flows"][:, 0], torch.tensor(-1.0)) + assert torch.allclose(out["flows"][:, 1], torch.tensor(2.0)) + # Binary masks are flipped spatially but keep their values. + assert torch.all(out["valid_flows"] == 1.0) + # Intrinsics principal point: cx -> (W - 1) - cx = 1 - 1 = 0. + assert out["intrinsics"][0, 0, 2] == 0.0 + assert out["intrinsics"][0, 0, 0] == 1.0 # fx untouched + # Baselines are unaffected by mirroring. + 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) + inputs = { + "images": img_left.clone(), + "images_right": img_right.clone(), + "disparities": disp.clone(), + } + + out = RandomFlip(hflip_prob=1.0, vflip_prob=0.0)(inputs) + 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]]]]) + flows = torch.ones(1, 2, 2, 2) + disp = torch.ones(1, 1, 2, 2) + intr = torch.tensor([[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]]) + inputs = { + "images": img.clone(), + "flows": flows.clone(), + "disparities": disp.clone(), + "intrinsics": intr.clone(), + } + + out = RandomFlip(hflip_prob=0.0, vflip_prob=1.0)(inputs) + + assert torch.allclose(out["images"], torch.tensor([[[[3.0, 4.0], [1.0, 2.0]]]])) + # Flow y-component is negated; x-component is not. + assert torch.allclose(out["flows"][:, 0], torch.tensor(1.0)) + assert torch.allclose(out["flows"][:, 1], torch.tensor(-1.0)) + # Horizontal disparity is unaffected by a vertical flip. + assert torch.allclose(out["disparities"], torch.tensor(1.0)) + # 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") + def test_no_erase(self, mock_random): + mock_random.return_value = 1.0 # Exceeds probability + img = torch.ones(2, 3, 4, 4) + inputs = {"images": img} + + out = RandomPatchEraser(erase_prob=0.5)(inputs) + assert torch.all(out["images"] == 1.0) + + @patch("random.randint") + @patch("random.random") + def test_mean_fill_second_image_only(self, mock_random, mock_randint): + mock_random.return_value = 0.0 + mock_randint.side_effect = [1, 2, 2, 1, 1] # num, hp, wp, yp, xp + + img = torch.arange(16).float().reshape(1, 4, 4).repeat(2, 1, 1, 1) + inputs = {"images": img.clone()} + + out = RandomPatchEraser(erase_prob=1.0, noise_type="mean")(inputs) + + # First image untouched. + assert torch.equal(out["images"][0], img[0]) + # Patch at rows/cols 1:3 of the second image filled with its mean (7.5). + assert torch.allclose(out["images"][1, 0, 1:3, 1:3], torch.tensor(7.5)) + # Outside the patch untouched. + assert out["images"][1, 0, 0, 0] == 0.0 + assert out["images"][1, 0, 3, 3] == 15.0 + + @patch("random.randint") + @patch("random.random") + def test_random_noise_fill(self, mock_random, mock_randint): + mock_random.return_value = 0.0 + mock_randint.side_effect = [1, 2, 2, 1, 1] + + img = torch.arange(16).float().reshape(1, 4, 4).repeat(2, 1, 1, 1) + img[1, 0, 0, 0] = 0.0 + img[1, 0, 3, 3] = 10.0 + inputs = {"images": img.clone()} + + torch.manual_seed(42) + out = RandomPatchEraser(erase_prob=1.0, noise_type="random")(inputs) + + patched = out["images"][1, 0, 1:3, 1:3] + assert torch.all(patched >= 0.0) # min of image 2 + assert torch.all(patched <= 14.0) # max of image 2 + assert patched[0, 0] != 5.0 + + @patch("random.random") + def test_oversized_patch_is_skipped(self, mock_random): + mock_random.return_value = 0.0 + img = torch.ones(2, 1, 4, 4) + inputs = {"images": img.clone()} + + # A patch larger than the image must be ignored, not crash. + out = RandomPatchEraser( + erase_prob=1.0, num_patches=(1, 1), patch_size=(10, 10) + )(inputs) + assert torch.all(out["images"] == 1.0) + + +class TestColorJitter: + @patch("random.random") + def test_symmetric_jitter(self, mock_random): + mock_random.return_value = 1.0 # Symmetric + + img1 = torch.ones(3, 4, 4) * 0.5 + img2 = torch.ones(3, 4, 4) * 0.5 + inputs = {"images": torch.stack([img1, img2])} + + out = ColorJitter(brightness=0.8, asymmetric_prob=0.5)(inputs) + 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), + "do_not_touch": torch.ones(1, 3, 4, 4), + } + + out = ColorJitter(brightness=2.0, use_keys=None, ignore_keys=["do_not_touch"])( + inputs + ) + assert torch.all(out["do_not_touch"] == 1.0) + + +class TestGaussianNoise: + def test_gaussian_noise_bounds(self): + img = torch.ones(1, 1, 2, 2) * 0.5 + inputs = {"images": img.clone()} + + torch.manual_seed(42) + random.seed(42) + + out = GaussianNoise(stdev=0.5)(inputs) + + assert not torch.allclose(out["images"], img) + assert torch.all(out["images"] >= 0.0) + assert torch.all(out["images"] <= 1.0) + + def test_ignore_keys(self): + img = torch.ones(1, 1, 2, 2) * 0.5 + flows = torch.ones(1, 2, 2, 2) * 5.0 + inputs = {"images": img.clone(), "flows": flows.clone()} + + torch.manual_seed(42) + random.seed(42) + + out = GaussianNoise(stdev=0.5, use_keys=None, ignore_keys=["flows"])(inputs) + assert not torch.allclose(out["images"], img) + assert torch.allclose(out["flows"], flows) diff --git a/tests/common/data/test_stereo_transforms.py b/tests/common/data/test_stereo_transforms.py new file mode 100644 index 0000000..2d7f863 --- /dev/null +++ b/tests/common/data/test_stereo_transforms.py @@ -0,0 +1,315 @@ +import random +from unittest.mock import patch + +import numpy as np +import torch + +from roco_spring_devkit.common.data.stereo_transforms import ( + AdjustGamma, + CenterCrop, + ColorJitter, + Compose, + GaussianNoise, + RandomFlip, + RandomPatchEraser, + RandomRotate, + RandomScaleAndCrop, + RandomTranslate, + Resize, + ToTensor, + VerticalYJitter, + _adjust_intrinsics_for_crop, + _adjust_intrinsics_for_flip, + _adjust_intrinsics_for_rotation, + _adjust_intrinsics_for_scale, + _get_spatial_key, + _get_valid_keys, + _resize, + _update_oob_disparities, +) + + +class TestUtilityFunctions: + def test_get_valid_keys(self): + keys = {"images", "images_right", "disparities", "intrinsics"} + assert set(_get_valid_keys(keys, ["images", "images_right"], None)) == { + "images", + "images_right", + } + assert set(_get_valid_keys(keys, None, ["intrinsics"])) == { + "images", + "images_right", + "disparities", + } + + def test_get_spatial_key(self): + inputs = { + "intrinsics": torch.eye(3).unsqueeze(0), + "disparities": torch.zeros(1, 1, 4, 4), + } + assert _get_spatial_key(inputs, ("intrinsics",)) == "disparities" + + def test_update_oob_disparities_1d(self): + occs = torch.zeros(1, 1, 3, 3) + disparities = torch.zeros(1, 1, 3, 3) + disparities[0, 0, 0, 0] = 1.0 # OOB left + disparities[0, 0, 0, 2] = -2.0 # 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 + + def test_resize_sparse_logic(self): + valid = torch.tensor([[[[0, 1, 0], [0, 0, 1]]]], dtype=torch.uint8) + disp = torch.ones(1, 1, 2, 3) + + inputs = {"valid_disparities": valid.clone(), "disparities": disp.clone()} + out = _resize( + inputs, + target_size=(3, 6), + binary_keys=["valid_disparities"], + disparities_keys=["disparities"], + sparse=True, + valid_key="valid_disparities", + ) + + assert out["valid_disparities"].shape == (1, 1, 3, 6) + assert out["valid_disparities"].dtype == torch.uint8 + assert out["valid_disparities"][0, 0, 0, 2] == 1 + assert out["valid_disparities"][0, 0, 2, 4] == 1 + assert out["valid_disparities"].sum() == 2 + + # Disparity scales only horizontally and follows its valid source pixels. + assert out["disparities"][0, 0, 0, 2] == 2.0 + assert out["disparities"][0, 0, 2, 4] == 2.0 + assert out["disparities"][0, 0, 2, 0] == 0.0 + + +class TestIntrinsicAdjustments: + def test_adjust_intrinsics_for_crop(self): + intr = torch.tensor([[[10.0, 0.0, 5.0], [0.0, 20.0, 6.0], [0.0, 0.0, 1.0]]]) + out = _adjust_intrinsics_for_crop(intr, x_crop=2, y_crop=3) + assert out[0, 0, 2] == 3.0 + assert out[0, 1, 2] == 3.0 + + def test_adjust_intrinsics_for_scale(self): + intr = torch.tensor([[[10.0, 0.0, 5.0], [0.0, 20.0, 6.0], [0.0, 0.0, 1.0]]]) + out = _adjust_intrinsics_for_scale(intr, x_scale=2.0, y_scale=0.5) + assert out[0, 0, 0] == 20.0 + assert out[0, 0, 2] == 10.0 + + 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, img_w=10, img_h=10) + assert out_h[0, 0, 2] == 4.0 + + def test_adjust_intrinsics_for_rotation(self): + intr = torch.tensor([[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]]) + # 90 degrees rotation matrix computation check + out = _adjust_intrinsics_for_rotation(intr, angle=90.0, img_w=3, img_h=3) + # Cos(90) = 0, Sin(90) = 1. Homography modifies principal points based on center (cx=1, cy=1) + assert torch.allclose(out[0, 0, 2], torch.tensor(1.0), atol=1e-5) + assert torch.allclose(out[0, 1, 2], torch.tensor(1.0), atol=1e-5) + + +class TestCompose: + def test_compose_filters_none(self): + t = Compose([ToTensor(), None]) + assert len(t.transforms_list) == 1 + + img = np.zeros((4, 4, 3), dtype=np.uint8) + inputs = {"images": [img]} + out = t(inputs) + assert out["images"].shape == (1, 3, 4, 4) + + +class TestToTensor: + def test_uint8_image_and_metadata_passthrough(self): + img = np.array([[[127, 255, 0]]], dtype=np.uint8) + intrinsics = np.eye(3, dtype=np.float32) + + inputs = {"images": [img], "intrinsics": intrinsics} + out = ToTensor()(inputs) + + assert out["images"].shape == (1, 3, 1, 1) + assert torch.allclose(out["images"][0, 1, 0, 0], torch.tensor(1.0)) + assert out["intrinsics"].shape == (3, 3) + + +class TestCenterCrop: + def test_spatial_crop_and_intrinsic_update(self): + img = torch.ones(1, 1, 4, 4) + intr = torch.tensor([[[1.0, 0.0, 2.0], [0.0, 1.0, 2.0], [0.0, 0.0, 1.0]]]) + inputs = { + "images": img, + "intrinsics": intr, + "valid_disparities": torch.ones(1, 1, 4, 4), + } + + out = CenterCrop(crop_size=(2, 2))(inputs) + assert out["images"].shape == (1, 1, 2, 2) + assert out["intrinsics"][0, 0, 2] == 1.0 + assert out["intrinsics"][0, 1, 2] == 1.0 + + +class TestRandomFlip: + def test_horizontal_flip_stereo_swap(self): + img_left = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + img_right = torch.tensor([[[[5.0, 6.0], [7.0, 8.0]]]]) + intr = torch.tensor([[[1.0, 0.0, 1.0], [0.0, 1.0, 1.0], [0.0, 0.0, 1.0]]]) + inputs = { + "images": img_left.clone(), + "images_right": img_right.clone(), + "intrinsics": intr, + } + + out = RandomFlip(hflip_prob=1.0, vflip_prob=0.0)(inputs) + + expected_left = torch.tensor([[[[6.0, 5.0], [8.0, 7.0]]]]) + assert torch.allclose(out["images"], expected_left) + assert out["intrinsics"][0, 0, 2] == 0.0 + + +class TestRandomScaleAndCrop: + @patch("random.randint") + @patch("random.uniform") + def test_random_scale_and_crop(self, mock_uniform, mock_randint): + mock_uniform.side_effect = [ + 1.0, + 1.0, + 1.0, + ] # major, space_h, space_w (2**1 = 2x multiplier) + mock_randint.side_effect = [0, 0] # y_crop, x_crop + + disp = torch.ones(1, 1, 4, 4) + inputs = {"disparities": disp} + + out = RandomScaleAndCrop( + crop_size=(6, 6), major_scale=(1.0, 1.0), space_scale=(1.0, 1.0) + )(inputs) + + # Should scale 4x4 -> 16x16, then crop to 6x6. + assert out["disparities"].shape == (1, 1, 6, 6) + + # Major and spatial exponents each produce a scale of 2, for a total of 4. + assert torch.allclose(out["disparities"], torch.tensor(4.0)) + + +class TestRandomTranslate: + @patch("random.randint") + def test_stereo_translation_and_disparity_shift(self, mock_randint): + mock_randint.side_effect = [1, 2, 1] # th, tw_left, tw_right + + img = torch.zeros(1, 1, 4, 4) + disp = torch.ones(1, 1, 4, 4) * 5.0 + intr = torch.tensor([[[1.0, 0.0, 2.0], [0.0, 1.0, 2.0], [0.0, 0.0, 1.0]]]) + + inputs = {"images": img, "disparities": disp, "intrinsics": intr} + + out = RandomTranslate(translation=(2, 2))(inputs) + assert out["images"].shape == (1, 1, 3, 2) + assert torch.all(out["disparities"] == 4.0) # 5.0 + (1 - 2) + assert out["intrinsics"][0, 0, 2] == 0.0 + + +class TestRandomRotate: + @patch("random.uniform") + def test_1d_disparity_rotation_scaling(self, mock_uniform): + mock_uniform.side_effect = [60.0, 0.0] + disp = torch.ones(1, 1, 3, 3) * 10.0 + inputs = {"disparities": disp} + + out = RandomRotate(angle=90.0)(inputs) + center_disp = out["disparities"][0, 0, 1, 1] + assert torch.allclose(center_disp, torch.tensor(5.0), atol=1e-5) + + +class TestResize: + def test_1d_disparity_scaling(self): + disp = torch.ones(1, 1, 2, 2) * 5.0 + inputs = {"disparities": disp} + + out = Resize(size=(4, 6))(inputs) + assert out["disparities"].shape[-2:] == (4, 6) + assert torch.allclose(out["disparities"], torch.tensor(15.0)) # 5.0 * 3.0 + + +class TestAdjustGamma: + @patch("random.uniform") + def test_adjust_gamma_bounds_and_application(self, mock_uniform): + # Provide gamma and gain + mock_uniform.side_effect = [1.5, 0.5, 1.5, 0.5] + + img = torch.ones(1, 3, 2, 2) * 0.5 + inputs = {"images": img.clone(), "images_right": img.clone()} + + out = AdjustGamma(gamma_range=(0.5, 2.0, 0.5, 2.0))(inputs) + + assert "images" in out + assert "images_right" in out + assert not torch.allclose(out["images"], img) + + +class TestVerticalYJitter: + @patch("random.random") + @patch("random.randint") + def test_vertical_shift_right_camera(self, mock_randint, mock_random): + mock_random.return_value = 0.0 + mock_randint.return_value = 1 + + img_right = torch.zeros(1, 1, 2, 2) + img_right[:, :, 1, :] = 1.0 + + inputs = {"images_right": img_right.clone()} + out = VerticalYJitter(jitter_prob=1.0, max_jitter=1)(inputs) + + assert torch.all(out["images_right"][0, 0, 1, :] == 0.0) + assert torch.all(out["images_right"][0, 0, 0, :] == 0.0) + + +class TestColorJitter: + @patch("random.random") + def test_symmetric_jitter_concatenation(self, mock_random): + mock_random.return_value = 1.0 # Symmetric + + img_left = torch.ones(1, 3, 2, 2) * 0.5 + img_right = torch.ones(1, 3, 2, 2) * 0.5 + inputs = {"images": img_left, "images_right": img_right} + + out = ColorJitter(brightness=0.8, asymmetric_prob=0.5)(inputs) + + assert torch.allclose(out["images"], out["images_right"]) + assert not torch.allclose(out["images"], torch.tensor(0.5)) + + +class TestGaussianNoise: + def test_gaussian_noise_bounds(self): + img = torch.ones(1, 1, 2, 2) * 0.5 + inputs = {"images": img.clone()} + + torch.manual_seed(42) + random.seed(42) + + out = GaussianNoise(stdev=0.5)(inputs) + + assert not torch.allclose(out["images"], img) + assert torch.all(out["images"] >= 0.0) + assert torch.all(out["images"] <= 1.0) + + +class TestRandomPatchEraser: + @patch("random.random") + @patch("random.randint") + def test_targets_right_image_only(self, mock_randint, mock_random): + mock_random.return_value = 0.0 + mock_randint.side_effect = [1, 2, 2, 0, 0] # num_patches, hp, wp, yp, xp + + img_left = torch.ones(1, 1, 2, 2) + img_right = torch.zeros(1, 1, 2, 2) + inputs = {"images": img_left, "images_right": img_right} + + out = RandomPatchEraser(erase_prob=1.0, noise_type="mean")(inputs) + + assert torch.all(out["images"] == 1.0) + assert "images_right" in out diff --git a/tests/common/utils/test_correlation.py b/tests/common/utils/test_correlation.py new file mode 100644 index 0000000..949c7e0 --- /dev/null +++ b/tests/common/utils/test_correlation.py @@ -0,0 +1,491 @@ +"""Unit tests for roco_spring_devkit.common.utils.correlation. + +All expected tensors below were computed by hand from the documented semantics of +each function (patch displacements symmetric around zero, channel-wise dot product, +bilinear sampling with the default ``zeros`` padding mode of ``F.grid_sample`` so +out-of-bounds samples are exactly 0, and ``align_corners=True`` so normalized +coordinate ``2*p/(n-1) - 1`` maps exactly to pixel ``p``). +""" + +import math + +import pytest +import torch + +from roco_spring_devkit.common.utils.correlation import ( + IterSpatialCorrelationSampler, + IterTranslatedSpatialCorrelationSampler, + IterativeCorrBlock, + iter_spatial_correlation_sample, + iter_translated_spatial_correlation_sample, +) + + +# --------------------------------------------------------------------------- +# Hand-computed expected value for the canonical 3x3 single-channel case. +# +# input1 = ones(1, 1, 3, 3) +# input2 = [[0, 1, 2], [3, 4, 5], [6, 7, 8]] +# patch_size = 3, dilation_patch = 1, stride = 1, padding = 0, kernel_size = 1. +# +# For each patch offset (py, px) the pixel shift is (py - 1, px - 1), and because +# ``F.grid_sample`` defaults to ``padding_mode='zeros'`` any sample that falls +# outside the [0, 2] pixel range contributes 0 instead of being clamped. +# Therefore ``expected[0, py, px, i, j] = input2[i + py - 1, j + px - 1]`` when the +# sampled indices are in bounds, otherwise 0. +# --------------------------------------------------------------------------- +_EXPECTED_PATCH3_SINGLE = torch.tensor( + [ + # py = 0 (y-shift = -1) + [ + # px = 0 (x-shift = -1) + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 3.0, 4.0]], + # px = 1 (x-shift = 0) + [[0.0, 0.0, 0.0], [0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], + # px = 2 (x-shift = +1) + [[0.0, 0.0, 0.0], [1.0, 2.0, 0.0], [4.0, 5.0, 0.0]], + ], + # py = 1 (y-shift = 0) + [ + [[0.0, 0.0, 1.0], [0.0, 3.0, 4.0], [0.0, 6.0, 7.0]], + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]], + [[1.0, 2.0, 0.0], [4.0, 5.0, 0.0], [7.0, 8.0, 0.0]], + ], + # py = 2 (y-shift = +1) + [ + [[0.0, 3.0, 4.0], [0.0, 6.0, 7.0], [0.0, 0.0, 0.0]], + [[3.0, 4.0, 5.0], [6.0, 7.0, 8.0], [0.0, 0.0, 0.0]], + [[4.0, 5.0, 0.0], [7.0, 8.0, 0.0], [0.0, 0.0, 0.0]], + ], + ] +).unsqueeze( + 0 +) # batch dimension -> shape (1, 3, 3, 3, 3) + + +def _arange_input2() -> torch.Tensor: + return torch.arange(9, dtype=torch.float).reshape(1, 1, 3, 3) + + +# --------------------------------------------------------------------------- +# iter_spatial_correlation_sample +# --------------------------------------------------------------------------- + + +class TestIterSpatialCorrelationSample: + def test_identity_ones_patch1_returns_channel_count(self) -> None: + # input1 == input2 == ones, only the zero-shift patch exists, so each + # output value is the channel-wise dot product, which equals the channel + # count ``c``. + c = 3 + input1 = torch.ones((1, c, 4, 5)) + input2 = torch.ones((1, c, 4, 5)) + corr = iter_spatial_correlation_sample(input1, input2, patch_size=1) + assert corr.shape == (1, 1, 1, 4, 5) + expected = torch.full((1, 1, 1, 4, 5), float(c)) + assert torch.allclose(corr, expected, atol=1e-6) + + def test_patch3_displacement_single_channel_matches_hand_computed(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + corr = iter_spatial_correlation_sample(input1, input2, patch_size=3) + assert corr.shape == (1, 3, 3, 3, 3) + assert torch.allclose(corr, _EXPECTED_PATCH3_SINGLE, atol=1e-6) + + def test_multichannel_elementwise_product_sum_patch1(self) -> None: + # input1[0] = 1, input1[1] = 2; input2[0] = arange(9), input2[1] = arange+100. + # corr = 1*arange + 2*(arange+100) = 3*arange + 200. + arange = torch.arange(9, dtype=torch.float).reshape(3, 3) + input1 = torch.stack( + [torch.ones_like(arange), 2 * torch.ones_like(arange)], dim=0 + )[None] + input2 = torch.stack([arange, arange + 100], dim=0)[None] + expected_center_block = ( + 3 * arange + 200 + ) # [[200,203,206],[209,212,215],[218,221,224]] + + corr = iter_spatial_correlation_sample(input1, input2, patch_size=1) + assert corr.shape == (1, 1, 1, 3, 3) + assert torch.allclose(corr[0, 0, 0], expected_center_block, atol=1e-4) + + def test_multichannel_patch3_center_is_no_shift_product_sum(self) -> None: + # The center patch (1, 1) has zero shift, so it must equal the plain + # per-pixel channel dot product regardless of patch_size. + arange = torch.arange(9, dtype=torch.float).reshape(3, 3) + input1 = torch.stack( + [torch.ones_like(arange), 2 * torch.ones_like(arange)], dim=0 + )[None] + input2 = torch.stack([arange, arange + 100], dim=0)[None] + expected = 3 * arange + 200 + + corr = iter_spatial_correlation_sample(input1, input2, patch_size=3) + assert torch.allclose(corr[0, 1, 1], expected, atol=1e-4) + + def test_stride2_subsamples_input1_spatial_grid(self) -> None: + # With patch_size = 1 and identity pairing (input1 == input2 = arange), + # stride=2 keeps pixels (0,0), (0,2), (2,0), (2,2) whose squared values + # are 0, 4, 36, 64. + input1 = _arange_input2() + input2 = _arange_input2() + expected = torch.tensor([[0.0, 4.0], [36.0, 64.0]]) + + corr = iter_spatial_correlation_sample(input1, input2, patch_size=1, stride=2) + assert corr.shape == (1, 1, 1, 2, 2) + assert torch.allclose(corr[0, 0, 0], expected, atol=1e-6) + + def test_padding1_with_zeros_input1_ones(self) -> None: + # Padding 1 zero-ring around both inputs. input1 = ones gets zeroed on + # the ring, so the ring of the correlation is 0; the inner 3x3 equals + # input2 (because input1 is 1 there). + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + inner = torch.arange(9, dtype=torch.float).reshape(3, 3) + expected = torch.zeros((5, 5)) + expected[1:4, 1:4] = inner + + corr = iter_spatial_correlation_sample(input1, input2, patch_size=1, padding=1) + assert corr.shape == (1, 1, 1, 5, 5) + assert torch.allclose(corr[0, 0, 0], expected, atol=1e-6) + + def test_dilation_patch_2_doubles_displacement(self) -> None: + # patch_size=3, dilation_patch=2 -> raw offsets [0, 2, 4], centered shifts + # [-2, 0, +2]. For the center output pixel (2, 2) of a 5x5 input2 the + # sampled pixels are exactly (i + py*2 - 2, j + px*2 - 2), all in bounds. + input1 = torch.ones((1, 1, 5, 5)) + input2 = torch.arange(25, dtype=torch.float).reshape(1, 1, 5, 5) + expected_center = torch.tensor( + [[0.0, 2.0, 4.0], [10.0, 12.0, 14.0], [20.0, 22.0, 24.0]] + ) + + corr = iter_spatial_correlation_sample( + input1, input2, patch_size=3, dilation_patch=2 + ) + assert corr.shape == (1, 3, 3, 5, 5) + assert torch.allclose(corr[0, :, :, 2, 2], expected_center, atol=1e-6) + + def test_chunk_size_does_not_change_result(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + full = iter_spatial_correlation_sample(input1, input2, patch_size=3) + c1 = iter_spatial_correlation_sample(input1, input2, patch_size=3, chunk_size=1) + c2 = iter_spatial_correlation_sample(input1, input2, patch_size=3, chunk_size=2) + c_partial = iter_spatial_correlation_sample( + input1, input2, patch_size=3, chunk_size=7 + ) # chunk_size > num_patches + assert torch.allclose(full, c1) + assert torch.allclose(full, c2) + assert torch.allclose(full, c_partial) + + def test_int_and_tuple_arguments_produce_equal_output(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + int_out = iter_spatial_correlation_sample( + input1, input2, patch_size=3, stride=1, padding=0, dilation_patch=1 + ) + tuple_out = iter_spatial_correlation_sample( + input1, + input2, + patch_size=(3, 3), + stride=(1, 1), + padding=(0, 0), + dilation_patch=(1, 1), + ) + assert torch.allclose(int_out, tuple_out) + + def test_kernel_size_not_1_raises(self) -> None: + with pytest.raises(NotImplementedError): + iter_spatial_correlation_sample( + torch.ones((1, 1, 3, 3)), _arange_input2(), kernel_size=2 + ) + + def test_dilation_not_1_raises(self) -> None: + with pytest.raises(NotImplementedError): + iter_spatial_correlation_sample( + torch.ones((1, 1, 3, 3)), _arange_input2(), dilation=2 + ) + + +class TestIterSpatialCorrelationSamplerModule: + def test_module_matches_function(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + module = IterSpatialCorrelationSampler(patch_size=3) + fn_out = iter_spatial_correlation_sample(input1, input2, patch_size=3) + assert torch.allclose(module(input1, input2), fn_out) + + def test_module_stores_arguments(self) -> None: + module = IterSpatialCorrelationSampler( + kernel_size=1, + patch_size=3, + stride=2, + padding=1, + dilation=1, + dilation_patch=2, + chunk_size=4, + ) + assert module.kernel_size == 1 + assert module.patch_size == 3 + assert module.stride == 2 + assert module.padding == 1 + assert module.dilation == 1 + assert module.dilation_patch == 2 + assert module.chunk_size == 4 + + +# --------------------------------------------------------------------------- +# iter_translated_spatial_correlation_sample +# --------------------------------------------------------------------------- + + +class TestIterTranslatedSpatialCorrelationSample: + def test_zero_flow_matches_basic_correlation(self) -> None: + # With flow = 0, coords = coords_grid, so the translated version must + # reproduce the non-translated correlation exactly. + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + flow = torch.zeros((1, 2, 3, 3)) + + translated = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=3 + ) + basic = iter_spatial_correlation_sample(input1, input2, patch_size=3) + assert torch.allclose(translated, basic, atol=1e-6) + assert translated.shape == basic.shape == (1, 3, 3, 3, 3) + + def test_flow_x_shift_plus_one_patches_one(self) -> None: + # flow[:, 0] = 1 (x = column shift +1), patch_size = 1 (no patch offset). + # Output pixel (i, j) samples input2[i, j + 1]; out-of-bounds (j = 2) -> 0. + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + flow = torch.zeros((1, 2, 3, 3)) + flow[:, 0] = 1.0 + expected = torch.tensor([[1.0, 2.0, 0.0], [4.0, 5.0, 0.0], [7.0, 8.0, 0.0]]) + + corr = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=1 + ) + assert corr.shape == (1, 1, 1, 3, 3) + assert torch.allclose(corr[0, 0, 0], expected, atol=1e-6) + + def test_flow_y_shift_plus_one_patches_one(self) -> None: + # flow[:, 1] = 1 (y = row shift +1), patch_size = 1. + # Output pixel (i, j) samples input2[i + 1, j]; out-of-bounds (i = 2) -> 0. + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + flow = torch.zeros((1, 2, 3, 3)) + flow[:, 1] = 1.0 + expected = torch.tensor([[3.0, 4.0, 5.0], [6.0, 7.0, 8.0], [0.0, 0.0, 0.0]]) + + corr = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=1 + ) + assert torch.allclose(corr[0, 0, 0], expected, atol=1e-6) + + def test_coords_path_matches_flow_path_for_identity(self) -> None: + # coords must be equivalent to flow + coords_grid. Building coords as the + # bare coordinate grid (i.e. zero flow) must reproduce the basic correlation. + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + h, w = 3, 3 + ys = torch.arange(h).view(h, 1).expand(h, w).float() + xs = torch.arange(w).view(1, w).expand(h, w).float() + coords = torch.stack([xs, ys], dim=0)[None] # (1, 2, h, w) + + from_coords = iter_translated_spatial_correlation_sample( + input1, input2, coords=coords, patch_size=3 + ) + basic = iter_spatial_correlation_sample(input1, input2, patch_size=3) + assert torch.allclose(from_coords, basic, atol=1e-6) + + def test_coords_path_with_flow_added(self) -> None: + # coords == coords_grid + (x shift +1, y shift 0) must match the flow path + # with flow[:, 0] = 1. + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + h, w = 3, 3 + ys = torch.arange(h).view(h, 1).expand(h, w).float() + xs = torch.arange(w).view(1, w).expand(h, w).float() + coords_grid = torch.stack([xs, ys], dim=0)[None] + flow = torch.zeros((1, 2, 3, 3)) + flow[:, 0] = 1.0 + coords = coords_grid + flow + + from_coords = iter_translated_spatial_correlation_sample( + input1, input2, coords=coords, patch_size=1 + ) + from_flow = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=1 + ) + assert torch.allclose(from_coords, from_flow, atol=1e-6) + expected = torch.tensor([[1.0, 2.0, 0.0], [4.0, 5.0, 0.0], [7.0, 8.0, 0.0]]) + assert torch.allclose(from_coords[0, 0, 0], expected, atol=1e-6) + + def test_chunk_size_does_not_change_result(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + flow = torch.zeros((1, 2, 3, 3)) + full = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=3 + ) + c1 = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=3, chunk_size=1 + ) + c2 = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=3, chunk_size=5 + ) + assert torch.allclose(full, c1) + assert torch.allclose(full, c2) + + def test_coords_and_flow_mutual_exclusion(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + # Both None is not a valid combination according to the assertion. + with pytest.raises(AssertionError): + iter_translated_spatial_correlation_sample( + input1, input2, flow=None, coords=None, patch_size=1 + ) + # Both provided is also forbidden. + with pytest.raises(AssertionError): + iter_translated_spatial_correlation_sample( + input1, + input2, + flow=torch.zeros((1, 2, 3, 3)), + coords=torch.zeros((1, 2, 3, 3)), + patch_size=1, + ) + + def test_kernel_size_not_1_raises(self) -> None: + with pytest.raises(NotImplementedError): + iter_translated_spatial_correlation_sample( + torch.ones((1, 1, 3, 3)), + _arange_input2(), + flow=torch.zeros((1, 2, 3, 3)), + kernel_size=2, + ) + + def test_dilation_not_1_raises(self) -> None: + with pytest.raises(NotImplementedError): + iter_translated_spatial_correlation_sample( + torch.ones((1, 1, 3, 3)), + _arange_input2(), + flow=torch.zeros((1, 2, 3, 3)), + dilation=2, + ) + + +class TestIterTranslatedSpatialCorrelationSamplerModule: + def test_module_matches_function(self) -> None: + input1 = torch.ones((1, 1, 3, 3)) + input2 = _arange_input2() + flow = torch.zeros((1, 2, 3, 3)) + flow[:, 0] = 1.0 + module = IterTranslatedSpatialCorrelationSampler(patch_size=1) + fn_out = iter_translated_spatial_correlation_sample( + input1, input2, flow=flow, patch_size=1 + ) + assert torch.allclose(module(input1, input2, flow), fn_out, atol=1e-6) + + def test_module_buffers_coords_grid_for_repeated_calls(self) -> None: + # The module should cache the coords_grid and reuse it. After two forward + # passes with the same flow shape it must still return the right values, + # and ``coords_grid`` should have a batch size matching the last call. + input1 = torch.ones((1, 1, 4, 4)) + input2 = torch.arange(16, dtype=torch.float).reshape(1, 1, 4, 4) + flow = torch.zeros((1, 2, 4, 4)) + module = IterTranslatedSpatialCorrelationSampler(patch_size=1) + first = module(input1, input2, flow) + second = module(input1, input2, flow) + assert first.shape == (1, 1, 1, 4, 4) + assert torch.allclose(first, second, atol=1e-6) + assert module.coords_grid is not None + assert module.coords_grid.shape == (1, 2, 4, 4) + + +# --------------------------------------------------------------------------- +# IterativeCorrBlock +# --------------------------------------------------------------------------- + + +class TestIterativeCorrBlock: + @staticmethod + def _identity_coords(h: int, w: int) -> torch.Tensor: + ys = torch.arange(h).view(h, 1).expand(h, w).float() + xs = torch.arange(w).view(1, w).expand(h, w).float() + return torch.stack([xs, ys], dim=0)[None] + + def test_num_levels_1_shape_and_values(self) -> None: + # fmap1 channel 0 = 1, channel 1 = 0, so only channel 0 of fmap2 (arange) + # contributes. The block rearranges ``b c d h w -> b (d c) h w`` so that + # output channel k = patch_x * patch_size_y + patch_y and then scales by + # 1 / sqrt(dim) where dim = number of channels = 2. + arange = torch.arange(9, dtype=torch.float).reshape(3, 3) + fmap1 = torch.stack([torch.ones_like(arange), torch.zeros_like(arange)], dim=0)[ + None + ] + fmap2 = torch.stack([arange, torch.zeros_like(arange)], dim=0)[None] + coords = self._identity_coords(3, 3) + + out = IterativeCorrBlock(fmap1, fmap2, radius=1, num_levels=1)(coords) + assert out.shape == (1, 9, 3, 3) + + inv_sqrt2 = 1.0 / math.sqrt(2.0) + # At the center spatial pixel (1, 1) every patch (py, px) is in bounds and + # samples input2[py, px] = py * 3 + px; output channel k = px * 3 + py. + expected_center = ( + torch.tensor([0.0, 3.0, 6.0, 1.0, 4.0, 7.0, 2.0, 5.0, 8.0]) * inv_sqrt2 + ) + assert torch.allclose(out[0, :, 1, 1], expected_center, atol=1e-5) + + # At the corner spatial pixel (0, 0): reuse the hand-computed A matrices at + # pixel (0, 0) which gives per patch (py, px) the values + # (0,0)=0 (0,1)=0 (0,2)=0 + # (1,0)=0 (1,1)=0 (1,2)=1 + # (2,0)=0 (2,1)=3 (2,2)=4 + # Reordered by k = px * 3 + py -> [0,0,0, 0,0,3, 0,1,4]. + expected_corner = ( + torch.tensor([0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 0.0, 1.0, 4.0]) * inv_sqrt2 + ) + assert torch.allclose(out[0, :, 0, 0], expected_corner, atol=1e-5) + + def test_num_levels_doubles_output_channels_and_first_level_matches(self) -> None: + # With num_levels = 2 the block concatenates one extra pooled level after + # the base level, so the first 9 channels must equal the num_levels = 1 + # output. + arange = torch.arange(16, dtype=torch.float).reshape(4, 4) + fmap1 = torch.stack([torch.ones_like(arange), torch.zeros_like(arange)], dim=0)[ + None + ] + fmap2 = torch.stack([arange, torch.zeros_like(arange)], dim=0)[None] + coords = self._identity_coords(4, 4) + + blk1 = IterativeCorrBlock(fmap1, fmap2, radius=1, num_levels=1) + blk2 = IterativeCorrBlock(fmap1, fmap2, radius=1, num_levels=2) + out1 = blk1(coords) + out2 = blk2(coords) + assert out1.shape == (1, 9, 4, 4) + assert out2.shape == (1, 18, 4, 4) + assert torch.allclose(out1, out2[:, :9, :, :], atol=1e-5) + + def test_scaling_by_sqrt_of_dim(self) -> None: + # Keeping the inputs the same but doubling the number of (zero) channels + # must scale every output value by sqrt(2) / sqrt(4) = 1 / sqrt(2). + arange = torch.arange(9, dtype=torch.float).reshape(3, 3) + zeros = torch.zeros_like(arange) + coords = self._identity_coords(3, 3) + + fmap1_2 = torch.stack([torch.ones_like(arange), zeros], dim=0)[None] + fmap2_2 = torch.stack([arange, zeros], dim=0)[None] + # Add two additional all-zero channels to bring dim from 2 to 4. + fmap1_4 = torch.cat([fmap1_2, torch.zeros((1, 2, 3, 3))], dim=1) + fmap2_4 = torch.cat([fmap2_2, torch.zeros((1, 2, 3, 3))], dim=1) + + out2 = IterativeCorrBlock(fmap1_2, fmap2_2, radius=1, num_levels=1)(coords) + out4 = IterativeCorrBlock(fmap1_4, fmap2_4, radius=1, num_levels=1)(coords) + ratio = out4 / out2 + finite_mask = torch.isfinite(ratio) + # The scaling factor between same-content, different-channel-count outputs. + assert torch.allclose( + ratio[finite_mask], + torch.full_like(ratio[finite_mask], math.sqrt(2.0) / math.sqrt(4.0)), + atol=1e-5, + ) diff --git a/tests/common/utils/test_flow_metrics.py b/tests/common/utils/test_flow_metrics.py new file mode 100644 index 0000000..76642d3 --- /dev/null +++ b/tests/common/utils/test_flow_metrics.py @@ -0,0 +1,570 @@ +"""Unit tests for roco_spring_devkit.common.utils.flow_metrics. + +The expected numbers below are computed by hand from the documented semantics of +``FlowMetrics``: + +* ``epe`` per pixel is the L2 norm of ``flow_pred - flow_target`` over the + channel axis. The reported value is the average over the per-sample means of + the valid pixels, averaged again across samples (i.e. ``epoch_mean`` divides + the accumulated per-sample totals by ``sample_count``). +* ``1px`` is ``100`` if a pixel's EPE is strictly greater than 1, ``0`` otherwise, + followed by the same averaging. +* ``flall`` is ``100`` if a pixel's EPE is greater than 3 AND greater than 0.05 + times the target flow magnitude at that pixel, ``0`` otherwise. +* ``wauc`` is the weighted AUC defined in + https://github.com/cv-stuttgart/springwebsite/blob/main/springeval/management/commands/evaluation.py + with thresholds ``delta_i = i / 20`` and weights ``w_i = 1 - (i - 1) / 100`` + for ``i`` in ``[1, 100]`` (``sum_wi = 50.5``). Pixels masked out have their + EPE replaced by 100 before counting. The reported value is scaled by 100. +* ``f1`` variants use ``(pred > 0.5)`` and ``(target > 0.5)`` binarization in + ``macro``/``binary``/``weighted`` modes. +* The EMA accumulator is ``state = ema_decay * state + (1 - ema_decay) * total``; + the normalization divisor is ``1 - ema_decay**step_count`` until + ``step_count`` reaches ``ema_max_count`` and ``1.0`` afterwards. +""" + +import math + +import pytest +import torch + +from roco_spring_devkit.common.utils.flow_metrics import FlowMetrics + + +def _zeros_flow(h: int = 1, w: int = 2) -> torch.Tensor: + return torch.zeros(1, 2, h, w) + + +def _flow_x_only(values, h: int, w: int) -> torch.Tensor: + """Build a 4D flow with the given (x-displacement) values and zero y-flow.""" + ch_x = torch.as_tensor(values, dtype=torch.float).reshape(1, 1, h, w) + ch_y = torch.zeros_like(ch_x) + return torch.cat([ch_x, ch_y], dim=1) + + +# --------------------------------------------------------------------------- +# Initialization +# --------------------------------------------------------------------------- + + +class TestFlowMetricsInit: + def test_default_average_mode_is_epoch_mean(self) -> None: + m = FlowMetrics() + assert m.average_mode == "epoch_mean" + + def test_default_attributes(self) -> None: + m = FlowMetrics() + assert m.prefix == "" + assert m.ema_decay == 0.99 + assert m.f1_mode == "macro" + assert m.interpolate_pred_to_target_size is False + # ema_max_count = min(100, int(1 / (1 - ema_decay))). With ema_decay = 0.99 + # the floating-point subtraction yields 0.010000000000000009, so + # 1 / 0.0100... ~= 99.999... which floors to 99 under ``int(...)``. + assert m.ema_max_count == 99 + assert m.include_occlusion is False + assert m.used_keys == [] + + def test_invalid_average_mode_raises(self) -> None: + with pytest.raises(AssertionError): + FlowMetrics(average_mode="unknown") + + def test_ema_decay_sets_short_max_count(self) -> None: + # ema_decay = 0.9 -> 1/(1-0.9) = 10, capped at min(100, 10). + m = FlowMetrics(average_mode="ema", ema_decay=0.9) + assert m.ema_max_count == 10 + + def test_prefix_attaches_to_metric_keys(self) -> None: + m = FlowMetrics(prefix="val_") + m.update( + {"flows": torch.zeros(1, 2, 1, 2)}, + {"flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]])}, + ) + keys = list(m.calculate_metrics().keys()) + assert keys == ["val_epe", "val_1px", "val_flall", "val_wauc"] + + +# --------------------------------------------------------------------------- +# EPE / 1px / flall / wauc (epoch_mean, no occlusion, no valid_flows) +# --------------------------------------------------------------------------- + + +class TestBasicMetricsEpochMean: + def test_hand_computed_single_pixel_x_displacement(self) -> None: + # pred = 0, target x = 1 at one pixel, zero elsewhere over a 2x2 grid. + # Per-pixel epe: [1, 0, 0, 0]; per-sample mean = 0.25. + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + m = FlowMetrics() + m.update({"flows": torch.zeros(1, 2, 2, 2)}, {"flows": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 0.25, abs_tol=1e-6) + # No pixel has epe > 1 (only the value 1 itself is not strictly > 1). + assert math.isclose(metrics["1px"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["flall"].item(), 0.0, abs_tol=1e-6) + + def test_wauc_hand_computed_pixel_at_threshold(self) -> None: + # Same setup as above. epe flattened = [1, 0, 0, 0], all valid. + # err_i = count(epe <= delta_i): + # i in [1, 19] -> delta in [0.05, 0.95]; err = 3 (the three zeros) + # i in [20, 100] -> delta in [1.0, 5.0]; err = 4 (also the 1) + # sum of wi for i in [1, 19] = 17.29 ; for i in [20, 100] = 33.21 + # raw wauc = 3*17.29 + 4*33.21 = 184.71 ; sum_wi = 50.5 ; N = 4 + # wauc = 100 * 184.71 / (4 * 50.5) = 91.4405940594... + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + m = FlowMetrics() + m.update({"flows": torch.zeros(1, 2, 2, 2)}, {"flows": target}) + assert math.isclose( + m.calculate_metrics()["wauc"].item(), + 100.0 * (3 * 17.29 + 4 * 33.21) / (4 * 50.5), + abs_tol=1e-4, + ) + + def test_wauc_is_zero_when_every_epe_exceeds_max_threshold(self) -> None: + # pred = 0, target = 10 at every channel -> epe = sqrt(200) ~= 14.14 everywhere. + target = torch.full((1, 2, 2, 2), 10.0) + m = FlowMetrics() + m.update({"flows": torch.zeros(1, 2, 2, 2)}, {"flows": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), math.sqrt(200.0), abs_tol=1e-5) + assert math.isclose(metrics["wauc"].item(), 0.0, abs_tol=1e-6) + + def test_1px_true_when_epe_exceeds_one(self) -> None: + # pred = 0, target = 1 at both channels -> per-pixel epe = sqrt(2) > 1. + # All pixels fail the 1px test, so 1px = 100. + preds = {"flows": torch.full((1, 2, 2, 2), 0.0)} + target = torch.full((1, 2, 2, 2), 1.0) + m = FlowMetrics() + m.update(preds, {"flows": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), math.sqrt(2.0), abs_tol=1e-5) + assert math.isclose(metrics["1px"].item(), 100.0, abs_tol=1e-5) + # epe = sqrt(2) is not > 3, so flall stays 0. + assert math.isclose(metrics["flall"].item(), 0.0, abs_tol=1e-6) + + def test_flall_requires_both_epe_gt_3_and_epe_gt_5pct_target_norm(self) -> None: + # pred = 0, target x = 4 only -> epe = 4 > 3 and 4 > 0.05 * 4 = 0.2. + target = torch.tensor([[[[4.0]], [[0.0]]]]) + m = FlowMetrics() + m.update({"flows": torch.zeros(1, 2, 1, 1)}, {"flows": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 4.0, abs_tol=1e-6) + assert math.isclose(metrics["1px"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["flall"].item(), 100.0, abs_tol=1e-6) + + def test_flall_blocked_by_target_norm_ratio(self) -> None: + # Same large EPE value (4), but very large target_norm so that + # 0.05 * norm < 4 is satisfied and flall should trigger as well — instead, + # here we craft a case where EPE > 0.05 * norm is FALSE: + # target_norm = epe (true only along single-channel direction). + # Choose target = (80, 0). pred = (76, 0). epe = 4, target_norm = 80, + # 0.05 * 80 = 4.0 -> strict inequality (epe > 0.05*norm) is False. + preds = {"flows": torch.tensor([[[[76.0]], [[0.0]]]])} + target = torch.tensor([[[[80.0]], [[0.0]]]]) + m = FlowMetrics() + m.update(preds, {"flows": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["1px"].item(), 100.0, abs_tol=1e-6) + # epe > 3 is True but epe > 0.05 * 80 = 4.0 is False (strict greater). + assert math.isclose(metrics["flall"].item(), 0.0, abs_tol=1e-6) + + def test_epoch_mean_aggregates_per_sample_means_across_steps(self) -> None: + # Two updates: per-sample mean epe equals 1.0 then 0.5. sample_count = 2. + m = FlowMetrics() + m.update( + {"flows": torch.zeros(1, 2, 1, 2)}, + {"flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]])}, + ) + m.update( + {"flows": torch.zeros(1, 2, 1, 2)}, + {"flows": torch.tensor([[[[0.0, 1.0]], [[0.0, 0.0]]]])}, + ) + metrics = m.calculate_metrics() + # (1.0 + 0.5) / 2 == 0.75 + assert math.isclose(metrics["epe"].item(), 0.75, abs_tol=1e-6) + + def test_epoch_mean_weights_each_sample_equally_not_per_pixel(self) -> None: + # One step has 1 valid pixel, the other has 3 valid pixels with the same + # per-pixel epe. With per-sample weighting each step contributes its mean + # (1.0), so the total is (1.0 + 1.0) / 2 == 1.0 rather than a 4-pixel mean. + m = FlowMetrics() + target1 = torch.tensor([[[[1.0]], [[0.0]]]]) + valid1 = torch.tensor([[[[1.0]]]]) + target2 = torch.full((1, 2, 1, 3), 0.0) + target2[0, 0, 0, :] = 1.0 # x = 1 at all pixels + valid2 = torch.tensor([[[[1.0, 1.0, 1.0]]]]) + m.update( + {"flows": torch.zeros(1, 2, 1, 1)}, + {"flows": target1, "valid_flows": valid1}, + ) + m.update( + {"flows": torch.zeros(1, 2, 1, 3)}, + {"flows": target2, "valid_flows": valid2}, + ) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 1.0, abs_tol=1e-6) + assert m.sample_count.item() == 2.0 + + def test_calculate_metrics_does_not_reset_state(self) -> None: + # Repeatedly calling compute/calculate_metrics should return the same + # value (state is not reset below ema_max_count). + m = FlowMetrics() + target = torch.tensor([[[[2.0]], [[0.0]]]]) + m.update({"flows": torch.zeros(1, 2, 1, 1)}, {"flows": target}) + first = m.calculate_metrics() + second = m.calculate_metrics() + assert set(first.keys()) == set(second.keys()) + for key in first: + assert torch.allclose(first[key], second[key]) + + +# --------------------------------------------------------------------------- +# valid_flows handling +# --------------------------------------------------------------------------- + + +class TestValidFlows: + def test_mask_excludes_invalid_pixels_from_epe(self) -> None: + # Three pixels with per-pixel epe = [1, 1, 1]; valid = [1, 1, 0]. + # The valid mean is (1 + 1) / 2 = 1.0, not the unmasked 1.0 total. + preds = {"flows": torch.zeros(1, 2, 1, 3)} + target = torch.tensor([[[[1.0, 1.0, 1.0]], [[0.0, 0.0, 0.0]]]]) + valid = torch.tensor([[[[1.0, 1.0, 0.0]]]]) + m = FlowMetrics() + m.update(preds, {"flows": target, "valid_flows": valid}) + assert math.isclose(m.calculate_metrics()["epe"].item(), 1.0, abs_tol=1e-6) + + def test_valid_flows_with_occlusion_distinguishes_occ_and_non_occ(self) -> None: + # pred = 0, target x = 1 at one pixel only -> per-pixel epe = [1, 0, 0, 0]. + # occ = [0, 0, 1, 1] (last row occluded), valid = all ones -> valid_occ + # selects the bottom row (epe = [0, 0], mean 0), valid_non_occ selects + # the top row (epe = [1, 0], mean 0.5). + preds = {"flows": torch.zeros(1, 2, 2, 2)} + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + m = FlowMetrics() + m.update(preds, {"flows": target, "occs": occ}) + metrics = m.calculate_metrics() + # global per-sample mean (all 4 valid pixels) = (1+0+0+0)/4 = 0.25 + assert math.isclose(metrics["epe"].item(), 0.25, abs_tol=1e-6) + # non-occ top-row mean = (1 + 0) / 2 = 0.5 + assert math.isclose(metrics["epe_non_occ"].item(), 0.5, abs_tol=1e-6) + # occ bottom-row mean = (0 + 0) / 2 = 0 + assert math.isclose(metrics["epe_occ"].item(), 0.0, abs_tol=1e-6) + + +# --------------------------------------------------------------------------- +# Occlusion hands-on: full set of extensions when occlusion_target is provided +# --------------------------------------------------------------------------- + + +class TestOcclusionHandling: + def test_occs_target_extends_used_keys_with_occ_and_non_occ(self) -> None: + preds = {"flows": torch.zeros(1, 2, 2, 2)} + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + m = FlowMetrics() + m.update(preds, {"flows": target, "occs": occ}) + # non-occ / occ pairs for the four metrics appear plus base metrics. + keys = [k for k, _, _ in m.used_keys] + for base in ("epe", "1px", "flall", "wauc"): + assert base in keys + assert f"{base}_occ" in keys + assert f"{base}_non_occ" in keys + metrics = m.calculate_metrics() + # EPE places already verified above. wauc on all pixels and on + # each masked half: + # wauc_occ: epe_masked = [100, 100, 0, 0]; err every threshold = 2; + # N = 2 -> 100 * 50.5 * 2 / (2 * 50.5) = 100. + # wauc_non_occ: epe_masked = [1, 0, 100, 100]; for delta in [0.05, 0.95] + # err = 1 (just the zero); for delta in [1.0, 5.0] err = 2 (zero+one); + # raw = 1 * 17.29 + 2 * 33.21 = 83.71; N = 2; + # wauc = 100 * 83.71 / (2 * 50.5) = 82.8812. + assert math.isclose(metrics["wauc_occ"].item(), 100.0, abs_tol=1e-4) + assert math.isclose( + metrics["wauc_non_occ"].item(), + 100.0 * (1 * 17.29 + 2 * 33.21) / (2 * 50.5), + abs_tol=1e-4, + ) + + def test_occs_pred_adds_occ_f1_with_perfect_match(self) -> None: + # occ_pred matching occ_target exactly -> occ_f1 == 1.0 in all three modes. + preds = { + "flows": torch.zeros(1, 2, 2, 2), + "occs": torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]), + } + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + for mode in ("binary", "macro", "weighted"): + m = FlowMetrics(f1_mode=mode) + m.update(preds, {"flows": target, "occs": occ}) + assert math.isclose( + m.calculate_metrics()["occ_f1"].item(), 1.0, abs_tol=1e-5 + ) + # occ_f1 job is gated on the predictor existing, so the used_keys + # should include it. + assert "occ_f1" in [k for k, _, _ in m.used_keys] + + def test_occ_f1_uses_macro_average_of_pos_and_neg_scores_default(self) -> None: + # Build the case verified by hand in the parent comment: + # pred always negative: f1_pos = 0. + # pred_neg (1-pred binarized) = all 1, target_neg binarized = [[1, 0, 0]] + # (because only pixel 0 of target is below 0.5 -> target_neg = 1). + # tp_neg = 1 (only match in pixel 0), fp_neg = 0, fn_neg = 2. + # precision = 1, recall = 1/3 -> f1_neg = 2*1*(1/3)/(1+1/3) = 0.5. + # macro default = (0 + 0.5)/2 = 0.25. + target_occ = torch.tensor([[[[0.0, 1.0, 1.0]]]]) + preds = { + "flows": torch.zeros(1, 2, 1, 3), + "occs": torch.zeros(1, 1, 1, 3), # always predicts negative + } + m = FlowMetrics() + m.update(preds, {"flows": torch.zeros(1, 2, 1, 3), "occs": target_occ}) + assert math.isclose(m.calculate_metrics()["occ_f1"].item(), 0.25, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# Motion boundary and confidence metrics +# --------------------------------------------------------------------------- + + +class TestMbF1: + def test_mb_f1_macro_default_with_all_negative_pred(self) -> None: + # mb_target = [[0,0,1,1]] (2 positive). mb_pred = all zeros. + # Positive class: pred_bin all 0; precision = recall = 0; f1_pos = 0. + # Negative class: pred_neg_bin all 1; target_neg_bin = [1,1,0,0] (2 pos); + # tp_neg = 2; fp_neg = 0; fn_neg = 2; precision = 1, recall = 0.5; + # f1_neg = 2*1*0.5/(1+0.5) = 0.6667. + # macro = (0 + 0.6667)/2 = 0.3333. + preds = { + "flows": torch.zeros(1, 2, 1, 4), + "mbs": torch.zeros(1, 1, 1, 4), + } + mb_target = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + m = FlowMetrics() + m.update(preds, {"flows": torch.zeros(1, 2, 1, 4), "mbs": mb_target}) + assert math.isclose( + m.calculate_metrics()["mb_f1"].item(), 0.33333333, abs_tol=1e-5 + ) + + def test_mb_f1_requires_both_pred_and_target(self) -> None: + # Only the prediction mb is provided; no metric should be added. + preds = {"flows": torch.zeros(1, 2, 1, 1), "mbs": torch.zeros(1, 1, 1, 1)} + m = FlowMetrics() + m.update(preds, {"flows": torch.zeros(1, 2, 1, 1)}) + assert "mb_f1" not in [k for k, _, _ in m.used_keys] + + +class TestConfF1: + def test_conf_target_is_exp_of_squared_epe(self) -> None: + # conf_target = exp(-(flow_pred - flow_target)^2 summed over channels). + # With target x = 1 at one pixel (others zero) and pred = 0, the squared + # error at that pixel is 1 and 0 elsewhere, so + # conf_target = [exp(-1), 1, 1, 1]. Thresh > 0.5 -> [0, 1, 1, 1]. + # conf_pred = all zeros (predicts no confidence). + # macro f1 of positives vs negatives: + # f1_pos: pred_bin all 0; tp=0, fp=3, fn=0; precision=recall=0 -> 0. + # f1_neg: pred_neg_bin all 1; target_neg_bin = [1, 0, 0, 0] (1 pos); + # tp = 1, fp = 0, fn = 3; precision = 1, recall = 1/4 = 0.25; + # f1_neg = 2*1*0.25/(1+0.25) = 0.4. + # macro = (0 + 0.4)/2 = 0.2. + preds = { + "flows": torch.zeros(1, 2, 2, 2), + "confs": torch.zeros(1, 1, 2, 2), + } + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + m = FlowMetrics() + m.update(preds, {"flows": target}) + assert math.isclose(m.calculate_metrics()["conf_f1"].item(), 0.2, abs_tol=1e-5) + + def test_conf_target_is_unity_when_pred_matches_target_perfectly(self) -> None: + # With pred == target all conf_target entries equal exp(0) = 1 > 0.5, so + # the positive class is perfectly identified by conf_pred = ones. + # However, the negative class is empty (no target_neg pixel > 0.5), so its + # f1 is 0 (precision = recall = 0). macro = (1 + 0) / 2 = 0.5. + preds = { + "flows": torch.zeros(1, 2, 2, 2), + "confs": torch.ones(1, 1, 2, 2), + } + target = torch.zeros(1, 2, 2, 2) + m = FlowMetrics() + m.update(preds, {"flows": target}) + assert math.isclose(m.calculate_metrics()["conf_f1"].item(), 0.5, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# F1 modes +# --------------------------------------------------------------------------- + + +class TestF1Modes: + @pytest.mark.parametrize("mode", ["binary", "macro", "weighted"]) + def test_perfect_match_gives_f1_one_in_all_modes(self, mode: str) -> None: + preds = { + "flows": torch.zeros(1, 2, 2, 2), + "occs": torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]), + } + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + m = FlowMetrics(f1_mode=mode) + m.update(preds, {"flows": target, "occs": occ}) + assert math.isclose(m.calculate_metrics()["occ_f1"].item(), 1.0, abs_tol=1e-5) + + def test_binary_mode_returns_only_positive_class_f1(self) -> None: + # Binary mode ignores the negative class contribution. With the + # motion-boundary scenario above (macro 0.3333), binary should return 0 + # (positive-class f1 == 0 since pred never predicts positive). + preds = {"flows": torch.zeros(1, 2, 1, 4), "mbs": torch.zeros(1, 1, 1, 4)} + mb_target = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + m = FlowMetrics(f1_mode="binary") + m.update(preds, {"flows": torch.zeros(1, 2, 1, 4), "mbs": mb_target}) + assert math.isclose(m.calculate_metrics()["mb_f1"].item(), 0.0, abs_tol=1e-5) + + def test_weighted_mode_writes_1_when_balanced_class_counts(self) -> None: + # target half positive half negative and pred matches exactly: weighted f1 + # becomes 1.0 because w_pos = w_neg = 0.5 and both f1 scores are 1.0. + target_occ = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + pred_match = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + m = FlowMetrics(f1_mode="weighted") + m.update( + {"flows": torch.zeros(1, 2, 1, 4), "occs": pred_match}, + {"flows": torch.zeros(1, 2, 1, 4), "occs": target_occ}, + ) + assert math.isclose(m.calculate_metrics()["occ_f1"].item(), 1.0, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# EMA mode +# --------------------------------------------------------------------------- + + +class TestAverageModeEma: + def test_ema_decay_0_9_step1_returns_input_mean(self) -> None: + # state = 0.9 * 0 + 0.1 * _compute_total. _compute_total uses .mean() in + # EMA so step1 contributes exactly 1.0 (the per-sample mean for the basic + # case). Divider = 1 - 0.9^1 = 0.1 -> reported metric = 1.0. + m = FlowMetrics(average_mode="ema", ema_decay=0.9) + m.update( + {"flows": torch.zeros(1, 2, 1, 2)}, + {"flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]])}, + ) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 1.0, abs_tol=1e-5) + assert m.step_count.item() == 1.0 + + def test_ema_decay_0_9_step2_hand_computed(self) -> None: + # Step1: state = 0.1 * 1.0 = 0.1 (per-sample mean = 1.0). + # Step2: per-sample mean = 0.5; state = 0.9 * 0.1 + 0.1 * 0.5 = 0.14. + # step_count = 2 -> divider = 1 - 0.9^2 = 0.19; metric = 0.14 / 0.19 = 0.7368... + m = FlowMetrics(average_mode="ema", ema_decay=0.9) + m.update( + {"flows": torch.zeros(1, 2, 1, 2)}, + {"flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]])}, + ) + m.update( + {"flows": torch.zeros(1, 2, 1, 2)}, + {"flows": torch.tensor([[[[0.0, 1.0]], [[0.0, 0.0]]]])}, + ) + assert math.isclose( + m.calculate_metrics()["epe"].item(), 0.14 / 0.19, abs_tol=1e-5 + ) + assert m.step_count.item() == 2.0 + + def test_ema_divisor_becomes_one_after_ema_max_count(self) -> None: + # ema_decay = 0.9 -> ema_max_count = 10. For step_count == 11 the divisor + # switches to 1.0 exactly, so the reported metric equals the raw state + # accumulated as state = sum_k (1-0.9) * 0.9^(step-1-k) * total_k (with + # initial state 0). For input total = 1.0 each step this reduces to + # 1 - 0.9^step_count. + m = FlowMetrics(average_mode="ema", ema_decay=0.9) + target = torch.tensor([[[[1.0]], [[0.0]]]]) + for _ in range(11): + m.update({"flows": torch.zeros(1, 2, 1, 1)}, {"flows": target}) + assert m.step_count.item() == 11.0 + # 1 - 0.9^11 = 1 - 0.313810596 = 0.686189404. + expected = 1.0 - 0.9**11 + assert math.isclose(m.calculate_metrics()["epe"].item(), expected, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# interpolate_pred_to_target_size +# --------------------------------------------------------------------------- + + +class TestInterpolatePredToTargetSize: + def test_zeros_pred_keeps_per_pixel_epe_with_rescaling(self) -> None: + # Pred shape (1, 2, 1, 2) interpolated to (2, 4). Pred stays zero, then + # multiplied by scale_x = 4/2 = 2 and scale_y = 2/1 = 2 — but since it is + # zero, only target contributes. EPE pixel (0,0) where target x = 3: + # 3 at that pixel and zero elsewhere -> sum = 3 over 8 pixels -> 0.375. + preds = {"flows": torch.zeros(1, 2, 1, 2)} + ch_x = torch.tensor([[3.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) + ch_y = torch.zeros_like(ch_x) + target = torch.stack([ch_x, ch_y], dim=0)[None] + m = FlowMetrics(interpolate_pred_to_target_size=True) + m.update(preds, {"flows": target}) + assert math.isclose(m.calculate_metrics()["epe"].item(), 3 / 8, abs_tol=1e-5) + + def test_nonzero_pred_is_rescaled_to_target_pixel_units(self) -> None: + # pred(1, 2, 1, 2) full of 1.0. scale_y = 2, scale_x = 2 -> after interp + # and scaling, both channels of pred become 2.0 everywhere. target = 0. + # Per-pixel epe = norm([2, 2]) = sqrt(8). mean over 8 px = sqrt(8). + preds = {"flows": torch.full((1, 2, 1, 2), 1.0)} + target = torch.zeros(1, 2, 2, 4) + m = FlowMetrics(interpolate_pred_to_target_size=True) + m.update(preds, {"flows": target}) + assert math.isclose( + m.calculate_metrics()["epe"].item(), math.sqrt(8.0), abs_tol=1e-5 + ) + + def test_interpolation_disabled_does_not_attempt_size_mismatch(self) -> None: + # With interpolate_pred_to_target_size=False, calling with mismatched + # shapes would broadcast-fail; here we ensure same-shape calls give the + # expected value (epe = sqrt(2) per pixel). + preds = {"flows": torch.full((1, 2, 2, 2), 1.0)} + target = torch.zeros(1, 2, 2, 2) + m = FlowMetrics(interpolate_pred_to_target_size=False) + m.update(preds, {"flows": target}) + assert math.isclose( + m.calculate_metrics()["epe"].item(), math.sqrt(2.0), abs_tol=1e-5 + ) + + +# --------------------------------------------------------------------------- +# Shape / dtype handling helper paths +# --------------------------------------------------------------------------- + + +class TestShapeHandling: + def test_2d_flow_is_promoted_to_4d(self) -> None: + # A 2D flow tensor (h, w) becomes (1, 1, h, w) so the per-pixel epe treats + # the lone channel as the x-component: target x = 4 over a single pixel. + preds = {"flows": torch.zeros(1, 1, 1, 1)} + target = {"flows": torch.tensor([[4.0]])} + m = FlowMetrics() + m.update(preds, target) + assert math.isclose(m.calculate_metrics()["epe"].item(), 4.0, abs_tol=1e-6) + + def test_3d_flow_with_batch_first_axis_promoted_to_4d(self) -> None: + # 3D target shape (1, 2, 2) has shape[0] == sample_count = 1, so + # ``_fix_shape`` adds a channel axis to yield (1, 1, 2, 2). pred matches + # the resulting single-channel shape; per-pixel epe over 4 px = [1, 0, + # 0, 0] -> mean 0.25. + preds = {"flows": torch.zeros(1, 1, 2, 2)} + target = {"flows": torch.tensor([[[1.0, 0.0], [0.0, 0.0]]])} + m = FlowMetrics() + m.update(preds, target) + assert math.isclose(m.calculate_metrics()["epe"].item(), 0.25, abs_tol=1e-6) + + def test_double_precision_input_is_cast_to_float32(self) -> None: + preds = {"flows": torch.zeros(1, 2, 1, 2, dtype=torch.float64)} + target = { + "flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]], dtype=torch.float64) + } + m = FlowMetrics() + m.update(preds, target) + metrics = m.calculate_metrics() + # Hand value: per-pixel epe = [1, 1] -> mean = 1.0. + assert math.isclose(metrics["epe"].item(), 1.0, abs_tol=1e-6) + # Internal state is stored as float32 after promotion through + # ``_to_metrics_dtype``, so the resulting metric is a float32 tensor. + assert metrics["epe"].dtype == torch.float32 diff --git a/tests/common/utils/test_flow_utils.py b/tests/common/utils/test_flow_utils.py new file mode 100644 index 0000000..fabbd94 --- /dev/null +++ b/tests/common/utils/test_flow_utils.py @@ -0,0 +1,499 @@ +"""Unit tests for roco_spring_devkit.common.utils.flow_utils. + +The expected values below are computed by hand from the documented semantics of +each function, then cross-checked against the implementation. They are *not* +simply a capture of the current output. + +Hand-computed reference facts +----------------------------- +flow_to_rgb (numpy, via flowpy): + * Zero flow, background="bright" -> white [255, 255, 255]. + The colorwheel hue at angle 0 is red [255, 0, 0]; with radius 0 the bright + branch ``move_hue_on_S_axis(hue, radius) = 255 - radius * (255 - hue)`` + evaluates to 255 in every channel. + * Zero flow, background="dark" -> black [0, 0, 0] + (``move_hue_on_V_axis(hue, 0) = hue * 0 = 0``). + * Flow (1, 0), flow_max_radius=1, bright -> red [255, 0, 0]: radius 1 keeps the + hue unchanged (no saturation move) so the raw colorwheel entry at angle 0 + (red) is returned. + * Flow (0, 1), flow_max_radius=1, bright -> [255, 229, 0]. + angle = pi/2 -> ncols-index 27/2 = 13.5; wheel[13] = [255, 221, 0], + wheel[14] = [255, 238, 0] (linspace R->Y of length 15), mean = [255, 229.5, 0] + which is truncated to uint8 -> [255, 229, 0]. + +flow_to_rgb (torch, via flowpy_torch): same colour semantics but the colorwheel +is scaled to [0, 1] and the output is float in [0, 1]. + +flow_read / flow_write round trips: + * ``.flo`` and ``.npy`` are lossless (float32). + * ``.pfm`` is lossless for finite values; NaNs in either channel of a pixel + mark *both* channels NaN on read (raft.read_pfm uses the 3rd channel as an + invalid mask for both flow components). + * ``.png128`` quantises with mult=128: written as + ``uint16(round(flow * 128 + 2**15))`` and read back as + ``(uint16 - 2**15) / 128``. E.g. flow_x = -2.3 -> -294.4 + 32768 = 32473.6 + -> 32473 -> (32473 - 32768) / 128 = -2.3046875. + * ``.flo5`` is lossless (gzip hdf5 of the raw float array). + * VIPER ``.npz`` stores ``u`` and ``v`` as float16; read converts back to + float32 and sets every pixel with ``abs(value) > 512`` to NaN. + +read_kubric_flow: + Loads ``data_ranges.json`` from the parent dir, reads the 16-bit PNG with + cv2 (BGR), drops channel 0 (``[..., 1:]``), then maps + ``flow = raw / 65535 * (max - min) + min``. + +fb_check: + Builds an (x, y) coordinate grid, adds the forward flow, samples the + backward flow at the warped locations with raft.bilinear_sampler + (align_corners=True), and computes + ``fb_diff = ||forward_flow + warped_backward_flow||`` per pixel. + ``fb_mask = (fb_diff < threshold) & in_mask`` where ``in_mask`` is the + sampler's validity mask. Because bilinear_sampler maps integer pixel + coordinates ``i`` in ``[0, W-1]`` to the normalised grid + ``2*i/(W-1) - 1`` and requires strict ``(-1, 1)`` bounds, the outer border + pixels are always marked invalid for a non-trivial grid. For a 5x5 grid + with zero flows only the inner 3x3 region is valid. + +spring_epe_to_rgb: + ``epe = sqrt(sum((pred - gt)**2, axis=2))``; + ``epe = clip(log2(epe * 32), 0, 10)``; + ``epe = (epe * 255 / 10).astype(uint8)``; + apply the matplotlib ``RdYlBu_r`` LUT (256 entries, BGR for OpenCV); + set invalid pixels (``~valid_mask``) to black. + Hard-coded LUT entries (independently computed from matplotlib): + index 0 (epe * 32 <= 1) -> BGR [149, 54, 49] + index 178 (epe = 4 -> idx 178) -> BGR [ 97, 174, 253] + index 255 (epe * 32 >= 2**10) -> BGR [ 38, 0, 165] +""" + +import math +import tempfile +from pathlib import Path + +import cv2 as cv +import numpy as np +import pytest +import torch + +from roco_spring_devkit.common.utils import flow_utils + + +# --------------------------------------------------------------------------- +# flow_to_rgb +# --------------------------------------------------------------------------- + + +class TestFlowToRgbNumpy: + def test_zero_flow_bright_background_is_white(self) -> None: + flow = np.zeros((2, 3, 2), dtype=np.float32) + rgb = flow_utils.flow_to_rgb(flow, background="bright") + assert rgb.shape == (2, 3, 3) + assert rgb.dtype == np.uint8 + assert np.array_equal(rgb, np.full((2, 3, 3), 255, dtype=np.uint8)) + + def test_zero_flow_dark_background_is_black(self) -> None: + flow = np.zeros((2, 3, 2), dtype=np.float32) + rgb = flow_utils.flow_to_rgb(flow, background="dark") + assert np.array_equal(rgb, np.zeros((2, 3, 3), dtype=np.uint8)) + + def test_pure_positive_x_flow_is_red(self) -> None: + # flow = (1, 0), max radius 1 -> raw colorwheel hue at angle 0 (red). + flow = np.array([[[1.0, 0.0]]], dtype=np.float32) + rgb = flow_utils.flow_to_rgb(flow, flow_max_radius=1.0, background="bright") + assert rgb.shape == (1, 1, 3) + assert rgb.tolist() == [[[255, 0, 0]]] + + def test_pure_positive_y_flow_is_yellowish(self) -> None: + # angle = pi/2 -> wheel index 13.5 -> mean(wheel[13], wheel[14]). + # wheel[13] = [255, 221, 0], wheel[14] = [255, 238, 0] -> [255, 229.5, 0] + # which truncates to uint8 as [255, 229, 0]. + flow = np.array([[[0.0, 1.0]]], dtype=np.float32) + rgb = flow_utils.flow_to_rgb(flow, flow_max_radius=1.0, background="bright") + assert rgb.tolist() == [[[255, 229, 0]]] + + def test_invalid_background_raises(self) -> None: + flow = np.zeros((1, 1, 2), dtype=np.float32) + with pytest.raises(ValueError): + flow_utils.flow_to_rgb(flow, background="purple") + + +class TestFlowToRgbTorch: + def test_zero_flow_bright_is_all_ones(self) -> None: + flow = torch.zeros(1, 2, 4, 5) + rgb = flow_utils.flow_to_rgb(flow, background="bright") + assert rgb.shape == (1, 3, 4, 5) + assert torch.allclose(rgb, torch.ones_like(rgb)) + + def test_zero_flow_dark_is_all_zeros(self) -> None: + flow = torch.zeros(1, 2, 4, 5) + rgb = flow_utils.flow_to_rgb(flow, background="dark") + assert torch.allclose(rgb, torch.zeros_like(rgb)) + + def test_pure_positive_x_flow_is_red_in_float(self) -> None: + # flow (1, 0), max radius 1 -> hue red [1, 0, 0] with radius 1. + flow = torch.tensor([[[[1.0]], [[0.0]]]]) # (1, 2, 1, 1) + rgb = flow_utils.flow_to_rgb(flow, flow_max_radius=1.0, background="bright") + assert rgb.shape == (1, 3, 1, 1) + assert torch.allclose(rgb, torch.tensor([[[[1.0]], [[0.0]], [[0.0]]]])) + + def test_invalid_background_raises(self) -> None: + flow = torch.zeros(1, 2, 1, 1) + with pytest.raises(ValueError): + flow_utils.flow_to_rgb(flow, background="purple") + + +# --------------------------------------------------------------------------- +# flow_read / flow_write round-trips +# --------------------------------------------------------------------------- + + +class TestFlowReadWrite: + @pytest.fixture + def tmpdir_path(self, tmp_path: Path) -> Path: + return tmp_path + + def test_flo_round_trip_is_lossless(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.flo" + flow_utils.flow_write(p, flow) + out = flow_utils.flow_read(p) + assert out.shape == flow.shape + assert np.array_equal(out, flow) + + def test_npy_round_trip_is_lossless(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.npy" + flow_utils.flow_write(p, flow) + out = flow_utils.flow_read(p) + assert np.array_equal(out, flow) + + def test_flo5_round_trip_is_lossless(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.flo5" + flow_utils.flow_write(p, flow) + out = flow_utils.flow_read(p) + assert np.array_equal(out, flow) + + def test_pfm_round_trip_preserves_finite_values(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.pfm" + flow_utils.flow_write(p, flow) + out = flow_utils.flow_read(p) + assert np.allclose(out, flow, atol=1e-6) + + def test_pfm_round_trip_nan_propagates_to_both_channels( + self, tmp_path: Path + ) -> None: + # selflow.write_pfm stores an invalid mask; raft.read_pfm marks both + # channels of an invalid pixel as NaN, so a single NaN becomes (NaN, NaN). + flow = np.array([[[1.5, -2.3], [float("nan"), 4.4]]], dtype=np.float32) + p = tmp_path / "a.pfm" + flow_utils.flow_write(p, flow) + out = flow_utils.flow_read(p) + # flow is (H=1, W=2, C=2): pixel (0, 1) had a NaN. + assert math.isnan(out[0, 1, 0]) and math.isnan(out[0, 1, 1]) + assert np.allclose(out[0, 0], [1.5, -2.3], atol=1e-6) + + def test_png128_round_trip_matches_hand_quantisation(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.png128" + flow_utils.flow_write(p, flow) + out = flow_utils.flow_read(p) + # png_flow_mult = 128: write = uint16(flow * 128 + 2**15), + # read = (uint16 - 2**15) / 128 (truncation toward zero in the cast). + # 1.5 -> 32960 -> 1.5 + # -2.3 -> -294.4 + 32768 = 32473.6 -> 32473 -> -2.3046875 + # 0.0 -> 32768 -> 0.0 + # 4.4 -> 563.2 + 32768 = 33331.2 -> 33331 -> 4.3984375 + expected = np.array([[[1.5, -2.3046875], [0.0, 4.3984375]]], dtype=np.float32) + assert np.allclose(out, expected, atol=1e-6) + + def test_png128_explicit_format_argument(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.png" # ambiguous extension; force png128 via format=. + flow_utils.flow_write(p, flow, format="png128") + out = flow_utils.flow_read(p, format="png128") + assert np.allclose(out, flow_utils.flow_read(p, format="png128"), atol=1e-6) + assert np.allclose(out[0, 0, 0], 1.5, atol=1e-6) + + def test_read_infers_format_from_extension(self, tmp_path: Path) -> None: + # .npy uses np.load directly; .flo uses flowpy.flow_read_flo. + flow = np.array([[[0.25, -0.75]]], dtype=np.float32) + npy = tmp_path / "a.npy" + flo = tmp_path / "a.flo" + flow_utils.flow_write(npy, flow) + flow_utils.flow_write(flo, flow) + assert np.array_equal(flow_utils.flow_read(npy), flow) + assert np.array_equal(flow_utils.flow_read(flo), flow) + + +# --------------------------------------------------------------------------- +# read_kubric_flow +# --------------------------------------------------------------------------- + + +class TestReadKubricFlow: + def test_value_mapping_matches_formula(self, tmp_path: Path) -> None: + # Build a 1x2 BGR uint16 PNG. read_kubric_flow drops channel 0 and keeps + # channels 1 and 2 (G, R). + img = np.zeros((1, 2, 3), dtype=np.uint16) + # pixel (0, 0): G=0, R=65535 + img[0, 0] = (0, 0, 65535) + # pixel (0, 1): G=32767, R=32767 + img[0, 1] = (0, 32767, 32767) + kdir = tmp_path / "kubric" + kdir.mkdir() + (kdir / "data_ranges.json").write_text( + '{"forward_flow": {"min": -100, "max": 100}}' + ) + png_path = kdir / "fwd.png" + assert cv.imwrite(str(png_path), img) + + flow = flow_utils.read_kubric_flow(png_path, "forward_flow") + assert flow.shape == (1, 2, 2) + # raw / 65535 * (max - min) + min + # pixel 0: [0, 65535] / 65535 * 200 - 100 = [-100, 100] + # pixel 1: [32767, 32767] / 65535 * 200 - 100 + # = 32767 * 200 / 65535 - 100 = 99.998474... - 100 = -0.0015258789... + expected = np.array( + [[[-100.0, 100.0], [-0.0015258789, -0.0015258789]]], dtype=np.float32 + ) + assert np.allclose(flow, expected, atol=1e-4) + + def test_uses_correct_flow_direction_key(self, tmp_path: Path) -> None: + img = np.zeros((1, 1, 3), dtype=np.uint16) + img[0, 0] = (0, 0, 65535) + kdir = tmp_path / "kubric" + kdir.mkdir() + (kdir / "data_ranges.json").write_text( + '{"forward_flow": {"min": -10, "max": 10},' + ' "backward_flow": {"min": 0, "max": 1}}' + ) + png_path = kdir / "fwd.png" + cv.imwrite(str(png_path), img) + + fwd = flow_utils.read_kubric_flow(png_path, "forward_flow") + bwd = flow_utils.read_kubric_flow(png_path, "backward_flow") + # forward: 65535/65535 * 20 - 10 = 10 + assert np.allclose(fwd[0, 0, 1], 10.0, atol=1e-5) + # backward: 65535/65535 * 1 - 0 = 1 + assert np.allclose(bwd[0, 0, 1], 1.0, atol=1e-5) + + +# --------------------------------------------------------------------------- +# read_viper_flow / write_viper_flow +# --------------------------------------------------------------------------- + + +class TestViperFlow: + def test_write_stores_u_and_v_as_float16(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.npz" + flow_utils.write_viper_flow(p, flow) + data = np.load(p) + assert set(data.keys()) == {"u", "v"} + assert data["u"].dtype == np.float16 + assert data["v"].dtype == np.float16 + assert np.array_equal(data["u"], flow[..., 0].astype(np.float16)) + assert np.array_equal(data["v"], flow[..., 1].astype(np.float16)) + + def test_read_reconstructs_hwc_and_float32(self, tmp_path: Path) -> None: + flow = np.array([[[1.5, -2.3], [0.0, 4.4]]], dtype=np.float32) + p = tmp_path / "a.npz" + flow_utils.write_viper_flow(p, flow) + out = flow_utils.read_viper_flow(p) + assert out.shape == flow.shape + assert out.dtype == np.float32 + # float16 round-trip of the values: 1.5 and 0.0 are exact; + # -2.3 -> -2.30078125, 4.4 -> 4.3984375. + expected = np.array([[[1.5, -2.30078125], [0.0, 4.3984375]]], dtype=np.float32) + assert np.allclose(out, expected, atol=1e-6) + + def test_read_marks_large_magnitudes_as_nan(self, tmp_path: Path) -> None: + # abs(value) > 512 becomes NaN on read. flow is (H=1, W=2, C=2): + # pixel (0, 0) = (1000, 1.0); pixel (0, 1) = (-600, 511.0). + flow = np.array([[[1000.0, 1.0], [-600.0, 511.0]]], dtype=np.float32) + p = tmp_path / "a.npz" + flow_utils.write_viper_flow(p, flow) + out = flow_utils.read_viper_flow(p) + assert math.isnan(out[0, 0, 0]) # 1000 > 512 + assert math.isnan(out[0, 1, 0]) # -600 -> |600| > 512 + assert out[0, 0, 1] == 1.0 # 1.0 within range + assert out[0, 1, 1] == 511.0 # 511 within range (not strictly > 512) + + +# --------------------------------------------------------------------------- +# fb_check +# --------------------------------------------------------------------------- + + +def _expected_interior_mask(size: int = 5) -> np.ndarray: + """Hand-computed in_mask for a zero-flow square grid. + + bilinear_sampler normalises coord i to 2*i/(W-1) - 1 and requires the + strict (-1, 1) range, so the outer border is always invalid. For a 5x5 + grid only the inner 3x3 region survives. + """ + mask = np.zeros((size, size), dtype=bool) + mask[1:-1, 1:-1] = True + return mask + + +class TestFbCheck: + def test_numpy_zero_flows_marks_inner_3x3(self) -> None: + fwd = np.zeros((5, 5, 2), dtype=np.float32) + bwd = np.zeros((5, 5, 2), dtype=np.float32) + mask = flow_utils.fb_check(fwd, bwd) + assert mask.dtype == bool + assert np.array_equal(mask, _expected_interior_mask(5)) + + def test_torch_input_returns_torch_mask(self) -> None: + fwd = torch.zeros((1, 2, 5, 5)) + bwd = torch.zeros((1, 2, 5, 5)) + mask = flow_utils.fb_check(fwd, bwd) + assert isinstance(mask, torch.Tensor) + expected = torch.from_numpy(_expected_interior_mask(5))[None] + assert torch.equal(mask, expected) + + def test_inconsistent_flow_above_threshold_all_false(self) -> None: + # forward = 0, backward = (5, 0): fb_diff = 5 > 1 everywhere interior. + fwd = np.zeros((5, 5, 2), dtype=np.float32) + bwd = np.full((5, 5, 2), 5.0, dtype=np.float32) + mask = flow_utils.fb_check(fwd, bwd) + assert not mask.any() + + def test_threshold_boundary_is_strict_less_than(self) -> None: + # fb_diff = 0.5; with threshold 1.0 the interior is True, with 0.4 False. + fwd = np.zeros((5, 5, 2), dtype=np.float32) + bwd = np.full((5, 5, 2), 0.5, dtype=np.float32) + assert np.array_equal( + flow_utils.fb_check(fwd, bwd, threshold=1.0), + _expected_interior_mask(5), + ) + assert not flow_utils.fb_check(fwd, bwd, threshold=0.4).any() + + def test_warped_backward_flow_is_used_for_diff(self) -> None: + # forward_flow at (row 2, col 2) = (1, 0) warps the lookup to (col 3, + # row 2). Place backward_flow at (row 2, col 3) = (-1, 0) so the + # diff at (2, 2) is ||(1,0) + (-1,0)|| = 0 -> valid. + # At (row 2, col 3) forward_flow is 0, so it samples backward_flow at + # (col 3, row 2) = (-1, 0) and diff = ||(0,0) + (-1,0)|| = 1, which is + # NOT < threshold 1.0 -> invalid. + fwd = np.zeros((5, 5, 2), dtype=np.float32) + bwd = np.zeros((5, 5, 2), dtype=np.float32) + fwd[2, 2] = [1.0, 0.0] + bwd[2, 3] = [-1.0, 0.0] + mask = flow_utils.fb_check(fwd, bwd) + expected = _expected_interior_mask(5).copy() + expected[2, 2] = True # diff = 0 < 1 + expected[2, 3] = False # diff = 1, not strictly < 1 + assert np.array_equal(mask, expected) + + +# --------------------------------------------------------------------------- +# spring_epe_to_rgb +# --------------------------------------------------------------------------- + + +def _expected_epe_rgb( + flow_pred: np.ndarray, + flow_gt: np.ndarray, + valid_mask: np.ndarray, + lut_bgr: np.ndarray, +) -> np.ndarray: + """Independent reference implementation of spring_epe_to_rgb. + + ``lut_bgr`` must be a (256, 1, 3) uint8 BGR LUT, computed independently + from matplotlib in the test (NOT via the production helper). + """ + epe = np.sqrt(np.square(flow_pred - flow_gt).sum(axis=2)) + epe = np.clip(np.log2(epe * 32), 0, 10) + epe = (epe * 255 / 10).astype(np.uint8) + rgb = cv.applyColorMap(epe, lut_bgr) + invalid = ~valid_mask + rgb[invalid] = 0 # broadcasts over the trailing channel dim + return rgb + + +# Independently built BGR LUT for RdYlBu_r (mirrors what the production helper +# produces, but built here directly from matplotlib so the test is not circular). +def _rdylbu_r_lut() -> np.ndarray: + import matplotlib.pyplot as plt + + cmap = plt.get_cmap("RdYlBu_r") + rgb = (cmap(np.arange(256))[:, :3] * 255).astype(np.uint8) + return rgb[:, ::-1].reshape(256, 1, 3) + + +class TestSpringEpeToRgb: + def test_perfect_prediction_maps_to_lut_index_zero(self) -> None: + # epe = 0 -> log2(0) = -inf -> clipped to 0 -> index 0. + pred = np.zeros((2, 2, 2), dtype=np.float32) + gt = np.zeros((2, 2, 2), dtype=np.float32) + valid = np.ones((2, 2), dtype=bool) + rgb = flow_utils.spring_epe_to_rgb(pred, gt, valid) + # Hand-checked: LUT[0] in BGR is [149, 54, 49]. + assert rgb.shape == (2, 2, 3) + assert np.array_equal(rgb, np.full((2, 2, 3), [149, 54, 49], dtype=np.uint8)) + + def test_large_epe_maps_to_lut_index_255(self) -> None: + # epe = 32 -> log2(32 * 32) = log2(1024) = 10 -> clipped to 10 -> + # 10 * 255 / 10 = 255 -> index 255. + pred = np.zeros((1, 1, 2), dtype=np.float32) + gt = np.array([[[32.0, 0.0]]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = flow_utils.spring_epe_to_rgb(pred, gt, valid) + # Hand-checked: LUT[255] in BGR is [38, 0, 165]. + assert rgb.shape == (1, 1, 3) + assert rgb[0, 0].tolist() == [38, 0, 165] + + def test_mid_epe_matches_hand_computed_index_178(self) -> None: + # epe = 4 -> log2(4 * 32) = log2(128) = 7 -> 7 * 255 / 10 = 178.5 -> + # uint8(178.5) = 178. Hand-checked LUT[178] BGR = [97, 174, 253]. + pred = np.zeros((1, 1, 2), dtype=np.float32) + gt = np.array([[[4.0, 0.0]]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = flow_utils.spring_epe_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [97, 174, 253] + + def test_small_epe_below_one_over_thirtytwo_clips_to_zero(self) -> None: + # epe = 0.01 -> epe * 32 = 0.32 -> log2(0.32) ~= -1.64 -> clipped to 0. + pred = np.zeros((1, 1, 2), dtype=np.float32) + gt = np.array([[[0.01, 0.0]]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = flow_utils.spring_epe_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [149, 54, 49] + + def test_huge_epe_clips_to_index_255(self) -> None: + # epe = 100 -> log2(3200) ~= 11.64 -> clipped to 10 -> index 255. + pred = np.zeros((1, 1, 2), dtype=np.float32) + gt = np.array([[[100.0, 0.0]]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = flow_utils.spring_epe_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [38, 0, 165] + + def test_invalid_pixels_are_set_to_black(self) -> None: + # The documented intent is that pixels with valid_mask == False become + # black. Hand-computed: pixel (0, 0) has epe=4 -> [97, 174, 253]; + # pixel (0, 1) is invalid -> [0, 0, 0]. + pred = np.zeros((1, 2, 2), dtype=np.float32) + gt = np.array([[[4.0, 0.0], [4.0, 0.0]]], dtype=np.float32) + valid = np.ones((1, 2), dtype=bool) + valid[0, 1] = False + rgb = flow_utils.spring_epe_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [97, 174, 253] + assert rgb[0, 1].tolist() == [0, 0, 0] + + def test_matches_independent_reference_implementation(self) -> None: + # Compares against a from-scratch implementation that uses a + # separately-built LUT and the correct (broadcast) invalid-mask + # assignment. + rng = np.random.RandomState(0) + pred = rng.randn(3, 4, 2).astype(np.float32) + gt = rng.randn(3, 4, 2).astype(np.float32) + valid = np.ones((3, 4), dtype=bool) + valid[0, 0] = False + valid[2, 3] = False + expected = _expected_epe_rgb(pred, gt, valid, _rdylbu_r_lut()) + out = flow_utils.spring_epe_to_rgb(pred, gt, valid) + assert np.array_equal(out, expected) diff --git a/tests/common/utils/test_io_adapter.py b/tests/common/utils/test_io_adapter.py new file mode 100644 index 0000000..e72018c --- /dev/null +++ b/tests/common/utils/test_io_adapter.py @@ -0,0 +1,544 @@ +"""Unit tests for roco_spring_devkit.common.utils.io_adapter. + +Hand-computed reference facts +----------------------------- +``IOAdapter`` wraps ``ToTensor`` (which converts HWC numpy into NCHW torch) +and an optional ``InputScaler`` (torch bilinear interpolation). The numbers +used in these tests are computed from those semantics, *not* captured from the +current output. + +ToTensor (optical_flow_transforms.ToTensor, with the IOAdapter defaults +``fp16=False`` and ``device="cpu"``): + +* A *list* of HWC uint8 images is ``np.stack``-ed to shape (N, H, W, C). + uint8 arrays are cast to float32 and divided by 255 before the + ``transpose(0, 3, 1, 2)`` so the resulting tensor is NCHW. + Example: two 1x2 RGB uint8 images + [[[0, 0, 0 ], [255, 255, 255]]] (img1) + [[[255, 255, 255], [0, 0, 0 ]]] (img2) + -> tensor shape (2, 3, 1, 2): + img1[:, 0, 0, 0] = 0/255 = 0 + img1[:, 0, 0, 1] = 255/255 = 1 (and the same for channels 1 and 2) + img2[:, 0, 0, 0] = 1 + img2[:, 0, 0, 1] = 0 + A *single* 3D HWC array is promoted with ``v[None]`` to 4D first, so the + same NCHW layout results. + +InputScaler (utils.InputScaler), bilinear ``align_corners=False``: + +* A constant input is reproduced unchanged by bilinear interpolation, because + every output pixel is a convex combination of input pixels (weights sum to + one). This holds for both ``fill`` (orig -> tgt) and ``unfill`` (tgt -> orig) + so a constant tensor round-trips losslessly. +* For a flow tensor the channels are scaled *after* interpolation: + flow[:, 0] *= new_width / old_width # horizontal displacement + flow[:, 1] *= new_height / old_height # vertical displacement + So a constant flow ``(x=1, y=2)`` resized from (H=2, W=4) to (H=2, W=2) becomes + ``(x=0.5, y=2)`` on ``fill`` and reverts to ``(x=1, y=2)`` on ``unfill``. +* A linear ramp is preserved (interpolation reproduces the linear trend), so a + horizontal ramp image ``[0, 1/3, 2/3, 1]`` (H=1, W=4) resized to W=2 with + ``align_corners=False`` keeps the canonical output positions + src_idx = (out_idx + 0.5) * (in_size / out_size) - 0.5 + = (out_idx + 0.5) * 2 - 0.5 + out_idx=0 -> src=0.5 -> 0.5*0 + 0.5*(1/3) = 1/6 ~= 0.16667 + out_idx=1 -> src=2.5 -> 0.5*(2/3) + 0.5*1 = 5/6 ~= 0.83333 + The ``unfill`` back to W=4 then uses + src = (out_idx + 0.5) * (2 / 4) - 0.5 + out_idx=0 -> src=-0.25 (clamped to 0) -> 1/6 = 0.16667 + out_idx=1 -> src= 0.25 -> 0.75*(1/6) + 0.25*(5/6) = 1/3 ~= 0.33333 + out_idx=2 -> src= 0.75 -> 0.25*(1/6) + 0.75*(5/6) = 2/3 ~= 0.66667 + out_idx=3 -> src= 1.25 (clamped to 1) -> 5/6 = 0.83333 + so the round-trip is lossy at the borders (the original ramp endpoints were + 0 and 1). + +IOAdapter.prepare_inputs pipeline: + +1. When ``inputs is None`` the adapter builds ``{"images": images, + "flows": flows, **kwargs}``, drops entries that are ``None`` or empty + (``len(v) == 0``), runs ``self.transform`` (``ToTensor``) on the survivors, + and then unsqueezes every tensor to at least 5D before applying + ``self.scaler.fill`` (only if a scaler was configured). +2. When ``inputs`` is supplied directly, ``ToTensor`` is *not* applied -- the + tensors are passed straight to the device/scaling loop. +3. ``image_only=True`` makes the loop ``continue`` on every key other than + ``"images"``, which means that not only the scaling but also the 5D + unsqueeze is skipped for the other tensors; flows therefore remain at the + 4D NCHW layout produced by ``ToTensor``. + +IOAdapter._to_cuda: + +* The destination stares are: ``device`` (if set) > ``"cuda"`` (if + ``cuda`` and ``torch.cuda.is_available()``) > stay on CPU. +* ``fp16`` conversion only runs when a target device was selected (``target is + not None``) and only converts tensors whose key contains ``"image"`` or + ``"flow"``. Any other tensor (e.g. ``meta``) keeps its dtype. +* When ``cuda=True`` but no GPU is available and ``device`` is ``None``, a + warning is emitted and the tensors stay on CPU in fp32. + +Known defects covered by ``xfail`` below +---------------------------------------- +* ``prepare_inputs`` checks ``len(v) == 0`` on every keyword value. If the + caller passes a scalar (e.g. an ``int`` metadata value via ``**kwargs``), the + function raises ``TypeError: object of type 'int' has no len()`` instead of + simply storing the scalar inside the inputs dict. The docstring explicitly + allows extra keyword arguments ("Any other array inputs ... keyworded + arguments") but the ``inputs`` parameter docstring simultaneously advertises + that the dict may hold "other metadata", so the asymmetry -- passing a + scalar via ``inputs=`` is fine, via kwargs raises -- is a robustness bug. +""" + +import math + +import numpy as np +import pytest +import torch + +from roco_spring_devkit.common.utils.io_adapter import IOAdapter + + +# --------------------------------------------------------------------------- +# Construction / scaler wiring +# --------------------------------------------------------------------------- + + +class TestInit: + def test_no_scaler_when_no_target(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(4, 8)) + assert adapter.scaler is None + assert adapter.output_stride == 8 + assert adapter.target_size is None + assert adapter.target_scale_factor is None + + def test_scaler_created_from_target_size(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(4, 8), target_size=(8, 16)) + assert adapter.scaler is not None + # InputScaler stores the *target* (height, width) it was built with. + assert adapter.scaler.tgt_height == 8 + assert adapter.scaler.tgt_width == 16 + # And remembers the original to unfill back to. + assert adapter.scaler.orig_height == 4 + assert adapter.scaler.orig_width == 8 + + def test_scaler_created_from_scale_factor(self) -> None: + # InputScaler uses int(orig * scale_factor) per side. + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_scale_factor=2.0) + assert adapter.scaler is not None + assert adapter.scaler.tgt_height == 4 + assert adapter.scaler.tgt_width == 8 + + def test_scaler_ignored_when_target_size_has_zero(self) -> None: + # min(target_size) > 0 must hold; (2, 0) -> scaler stays None even if + # target_scale_factor is also set, because the + # `target_size is not None and min(target_size) > 0` clause short-circuits + # to False and ``target_scale_factor`` is None here. + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_size=(2, 0)) + assert adapter.scaler is None + + def test_max_size_takes_precedence_over_scale_factor(self) -> None: + # InputScaler checks stride -> size -> scale_factor in this order, + # so when both target_size and target_scale_factor are given the + # explicit size wins. + adapter = IOAdapter( + output_stride=8, + input_size=(2, 4), + target_size=(6, 6), + target_scale_factor=10.0, + ) + assert adapter.scaler is not None + assert (adapter.scaler.tgt_height, adapter.scaler.tgt_width) == (6, 6) + + def test_output_stride_has_no_effect_on_scaling(self) -> None: + # `output_stride` is stored on the instance but never used by the + # scaler (InputScaler is always created with size/scale_factor, never + # with stride). Verify this so the behaviour is pinned down. + a1 = IOAdapter(output_stride=1, input_size=(2, 4)) + a2 = IOAdapter( + output_stride=128, + input_size=(2, 4), + target_size=(2, 4), + ) + assert a1.scaler is None + # Different output_strides give the same scaler target sizes. + b1 = IOAdapter(output_stride=1, input_size=(2, 4), target_size=(4, 8)) + b2 = IOAdapter(output_stride=64, input_size=(2, 4), target_size=(4, 8)) + assert (b1.scaler.tgt_height, b1.scaler.tgt_width) == (4, 8) + assert (b2.scaler.tgt_height, b2.scaler.tgt_width) == (4, 8) + + +# --------------------------------------------------------------------------- +# prepare_inputs: ToTensor path, removals, kwargs, 5D layout +# --------------------------------------------------------------------------- + + +class TestPrepareInputsToTensor: + def test_uint8_images_are_divided_by_255(self) -> None: + img = np.array([[[0, 128, 255]]], dtype=np.uint8) # HWC, H=1, W=1, C=3 + adapter = IOAdapter(output_stride=8, input_size=(1, 1)) + out = adapter.prepare_inputs(images=[img]) + t = out["images"] + # ToTensor makes NCHW with N=1, then prepare_inputs unsqueezes to 5D. + assert tuple(t.shape) == (1, 1, 3, 1, 1) + # Values: 0/255, 128/255, 255/255 = 0, 0.501960784, 1.0 (stored as fp32). + expected = np.array([0.0, 128.0 / 255.0, 1.0], dtype=np.float32) + actual = t[0, 0, :, 0, 0].numpy() + np.testing.assert_allclose(actual, expected, atol=1e-6) + + def test_two_rgb_images_laid_out_as_5d(self) -> None: + img1 = np.array([[[0, 0, 0], [255, 255, 255]]], dtype=np.uint8) + img2 = np.array([[[255, 255, 255], [0, 0, 0]]], dtype=np.uint8) + adapter = IOAdapter(output_stride=8, input_size=(1, 2)) + out = adapter.prepare_inputs(images=[img1, img2]) + t = out["images"] + # 5D layout: (1, N=2, C=3, H=1, W=2). + assert tuple(t.shape) == (1, 2, 3, 1, 2) + # img1 channels are all [0, 1]; img2 channels are all [1, 0]. + np.testing.assert_allclose(t[0, 0, 0, 0].numpy(), [0.0, 1.0], atol=1e-6) + np.testing.assert_allclose(t[0, 1, 0, 0].numpy(), [1.0, 0.0], atol=1e-6) + + def test_single_float_image_keeps_values_untouched(self) -> None: + # float32 arrays are NOT divided by 255 (only uint8 is). + img = np.array([[[0.0]], [[1.0]]], dtype=np.float32) # H=2, W=1, C=1 + adapter = IOAdapter(output_stride=8, input_size=(2, 1)) + out = adapter.prepare_inputs(images=img) + t = out["images"] + # 3D HWC -> ToTensor -> NCHW with N=1 -> 5D (1, 1, 1, 2, 1). + assert tuple(t.shape) == (1, 1, 1, 2, 1) + np.testing.assert_allclose(t.numpy().ravel(), [0.0, 1.0], atol=1e-6) + + def test_tensors_are_unsqueezed_to_5d(self) -> None: + img = np.zeros((2, 4, 3), dtype=np.uint8) + adapter = IOAdapter(output_stride=8, input_size=(2, 4)) + out = adapter.prepare_inputs(images=img) + # ToTensor: HWC -> (1, 3, 2, 4) NCHW. prepare_inputs unsqueezes to + # (1, 1, 3, 2, 4). + assert tuple(out["images"].shape) == (1, 1, 3, 2, 4) + + def test_removes_none_entries(self) -> None: + img = np.zeros((2, 2, 3), dtype=np.uint8) + adapter = IOAdapter(output_stride=8, input_size=(2, 2)) + out = adapter.prepare_inputs(images=[img, img], flows=None) + assert "images" in out + assert "flows" not in out + + def test_removes_empty_list_entries(self) -> None: + img = np.zeros((2, 2, 3), dtype=np.uint8) + adapter = IOAdapter(output_stride=8, input_size=(2, 2)) + out = adapter.prepare_inputs(images=[], flows=[img, img]) + assert "images" not in out + assert "flows" in out + + def test_removes_empty_array_entries(self) -> None: + img = np.zeros((2, 2, 3), dtype=np.uint8) + empty = np.zeros((0,), dtype=np.float32) + adapter = IOAdapter(output_stride=8, input_size=(2, 2)) + out = adapter.prepare_inputs(images=[img, img], extras=empty) + assert "extras" not in out + + def test_kwargs_are_added_to_input_dict(self) -> None: + img = np.zeros((2, 2, 3), dtype=np.uint8) + # A *list* of two HWC uint8 arrays: numpy stacks to (2, 2, 2, 1) which + # ToTensor turns into NCHW=(2, 1, 2, 2) then prepare_inputs unsqueezes to + # (1, 2, 1, 2, 2). + extras = [ + np.zeros((2, 2, 1), dtype=np.uint8), + np.zeros((2, 2, 1), dtype=np.uint8), + ] + adapter = IOAdapter(output_stride=8, input_size=(2, 2)) + out = adapter.prepare_inputs(images=[img, img], extra=extras) + assert set(out.keys()) == {"images", "extra"} + # extra goes through ToTensor (uint8 -> float32/255 -> 5D NCHW). + assert tuple(out["extra"].shape) == (1, 2, 1, 2, 2) + # 1 * 2 * 1 * 2 * 2 = 8 zeros (all zeros / 255 = 0). + np.testing.assert_allclose( + out["extra"].numpy().ravel(), + np.zeros(1 * 2 * 1 * 2 * 2, dtype=np.float32), + atol=1e-6, + ) + + +class TestPrepareInputsProvidedDict: + def test_provided_inputs_skip_totensor_transform(self) -> None: + # When `inputs` is supplied, ToTensor is NOT applied. + adapter = IOAdapter(output_stride=8, input_size=(2, 4)) + tensor = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4) + out = adapter.prepare_inputs(inputs={"images": tensor}) + # Values are left intact; only the 5D unsqueeze modifies the shape. + assert tuple(out["images"].shape) == (1, 2, 3, 2, 4) + torch.testing.assert_close( + out["images"], tensor.unsqueeze(0), atol=0.0, rtol=0.0 + ) + + def test_provided_numpy_values_skipped(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(2, 4)) + out = adapter.prepare_inputs( + inputs={"images": np.zeros((2, 3, 2, 4), dtype=np.float32)} + ) + # Not a torch tensor -> left untouched by the scaling loop. + assert isinstance(out["images"], np.ndarray) + + def test_images_and_flows_kwargs_ignored_when_inputs_given(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(2, 4)) + tensor = torch.zeros(1, 3, 2, 4) + out = adapter.prepare_inputs( + inputs={"images": tensor}, + images=np.zeros((2, 4, 3), dtype=np.uint8), + flows=np.zeros((2, 4, 2), dtype=np.float32), + ) + assert set(out.keys()) == {"images"} + torch.testing.assert_close(out["images"], tensor.unsqueeze(0)) + + +# --------------------------------------------------------------------------- +# prepare_inputs: scaling behaviour (InputScaler) +# --------------------------------------------------------------------------- + + +class TestPrepareInputsScaling: + def test_constant_image_preserves_value_through_fill(self) -> None: + # A constant image maps to itself under bilinear interpolation. + img = np.full((2, 4, 3), 128, dtype=np.uint8) # 128/255 everywhere. + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_size=(4, 8)) + out = adapter.prepare_inputs(images=img) + assert tuple(out["images"].shape) == (1, 1, 3, 4, 8) + np.testing.assert_allclose( + out["images"].numpy().ravel(), + np.full((1 * 1 * 3 * 4 * 8,), 128 / 255.0, dtype=np.float32), + atol=1e-6, + ) + + def test_ramp_image_bilinear_hand_computed(self) -> None: + # Horizontal ramp H=1 W=4 -> resize to W=2 (align_corners=False). + # See module docstring for the derivation: [1/6, 5/6]. + img = np.array([[[0.0], [1.0 / 3.0], [2.0 / 3.0], [1.0]]], dtype=np.float32) + adapter = IOAdapter(output_stride=8, input_size=(1, 4), target_size=(1, 2)) + out = adapter.prepare_inputs(images=img) + assert tuple(out["images"].shape) == (1, 1, 1, 1, 2) + np.testing.assert_allclose( + out["images"][0, 0, 0, 0].numpy(), + [1.0 / 6.0, 5.0 / 6.0], + atol=1e-6, + ) + + def test_constant_flow_scaled_by_xy_multipliers(self) -> None: + # x-flow = 1, y-flow = 2, constant -> stay constant under interpolation + # and are then scaled by (new_W / old_W, new_H / old_H). + flow = np.zeros((2, 4, 2), dtype=np.float32) + flow[..., 0] = 1.0 + flow[..., 1] = 2.0 + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_size=(2, 2)) + out = adapter.prepare_inputs(flows=flow) + assert tuple(out["flows"].shape) == (1, 1, 2, 2, 2) + # x-channel scaled by 2/4 = 0.5 -> 0.5 everywhere. + ch0 = out["flows"][0, 0, 0].numpy() + np.testing.assert_allclose(ch0, np.full((2, 2), 0.5), atol=1e-6) + # y-channel scaled by 2/2 = 1.0 -> 2.0 everywhere. + ch1 = out["flows"][0, 0, 1].numpy() + np.testing.assert_allclose(ch1, np.full((2, 2), 2.0), atol=1e-6) + + def test_image_only_skips_flows_completely(self) -> None: + # image_only=True makes the loop continue past every non-"images" key, + # which skips BOTH the scaler and the 5D unsqueeze for flows. + img = np.full((2, 4, 3), 100, dtype=np.uint8) + flow = np.zeros((2, 4, 2), dtype=np.float32) + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_size=(4, 8)) + out = adapter.prepare_inputs(images=[img, img], flows=flow, image_only=True) + # Images were scaled to the target and unsqueezed to 5D. + assert tuple(out["images"].shape) == (1, 2, 3, 4, 8) + # Flows were left at the ToTensor 4D NCHW layout, NOT scaled. + assert tuple(out["flows"].shape) == (1, 2, 2, 4) + np.testing.assert_allclose( + out["flows"].numpy().ravel(), np.zeros(1 * 2 * 2 * 4), atol=1e-6 + ) + + +# --------------------------------------------------------------------------- +# unscale: revert the scaling applied by prepare_inputs +# --------------------------------------------------------------------------- + + +class TestUnscale: + def test_no_scaler_returns_tensors_unchanged(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(4, 8)) + t = torch.arange(1 * 3 * 4 * 8, dtype=torch.float32).reshape(1, 3, 4, 8) + out = adapter.unscale({"images": t}) + # Identity when there is no scaler. + assert out is not None + torch.testing.assert_close(out["images"], t, atol=0, rtol=0) + + def test_constant_image_round_trips_losslessly(self) -> None: + # Constant images scale and unscale back without information loss. + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_size=(4, 8)) + const_val = 0.75 + images = torch.full((1, 1, 3, 2, 4), const_val, dtype=torch.float32) + scaled = {"images": images.clone()} + for k, v in scaled.items(): + scaled[k] = adapter.scaler.fill(v, is_flow=False) + unscaled = adapter.unscale({"images": scaled["images"].clone()}) + assert tuple(unscaled["images"].shape) == (1, 1, 3, 2, 4) + np.testing.assert_allclose( + unscaled["images"].numpy().ravel(), + np.full((1 * 1 * 3 * 2 * 4,), const_val), + atol=1e-6, + ) + + def test_constant_flow_round_trips_losslessly(self) -> None: + # The scaler is built from orig (H=2, W=8) to tgt (H=2, W=2). We must + # therefore feed it a flow tensor at the *original* width W=8 so that + # fill interpolates and applies the flow multiplier. + # fill: (x=1, y=4) from (2,8) -> (2,2): x *= 2/8 = 0.25, y *= 2/2 = 1. + # unfill: x *= 8/2 = 4 -> 0.25*4 = 1, y *= 2/2 = 1 -> 4*1 = 4. + adapter = IOAdapter(output_stride=8, input_size=(2, 8), target_size=(2, 2)) + const = torch.zeros(1, 1, 2, 2, 8) # (B=1, N=1, C=2, H=2, W=8) + const[0, 0, 0] = 1.0 # horizontal displacement + const[0, 0, 1] = 4.0 # vertical displacement + scaled = {"flows": adapter.scaler.fill(const.clone(), is_flow=True)} + assert tuple(scaled["flows"].shape) == (1, 1, 2, 2, 2) + np.testing.assert_allclose( + scaled["flows"][0, 0, 0].numpy().ravel(), + np.full((2 * 2,), 0.25), + atol=1e-6, + ) + np.testing.assert_allclose( + scaled["flows"][0, 0, 1].numpy().ravel(), + np.full((2 * 2,), 4.0), + atol=1e-6, + ) + unscaled = adapter.unscale({"flows": scaled["flows"].clone()}) + assert tuple(unscaled["flows"].shape) == (1, 1, 2, 2, 8) + np.testing.assert_allclose( + unscaled["flows"][0, 0, 0].numpy().ravel(), + np.full((2 * 8,), 1.0), + atol=1e-6, + ) + np.testing.assert_allclose( + unscaled["flows"][0, 0, 1].numpy().ravel(), + np.full((2 * 8,), 4.0), + atol=1e-6, + ) + + def test_ramp_image_unscale_hand_computed(self) -> None: + # From the module docstring: 4->2 ramp gives [1/6, 5/6]; resizing those + # back to 4 with align_corners=False yields [1/6, 1/3, 2/3, 5/6]. + adapter = IOAdapter(output_stride=8, input_size=(1, 4), target_size=(1, 2)) + scaled = torch.tensor( + [[[[1.0 / 6.0, 5.0 / 6.0]]]], dtype=torch.float32 + ).reshape(1, 1, 1, 1, 2) + out = adapter.unscale({"images": scaled}) + assert tuple(out["images"].shape) == (1, 1, 1, 1, 4) + np.testing.assert_allclose( + out["images"][0, 0, 0, 0].numpy(), + [1.0 / 6.0, 1.0 / 3.0, 2.0 / 3.0, 5.0 / 6.0], + atol=1e-6, + ) + + def test_unscale_image_only_skips_non_images(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(2, 4), target_size=(4, 8)) + # Both tensors are deliberately small enough that unfill would change + # their shape, so we can detect whether unfill ran. + images = torch.zeros(1, 1, 2, 4, 4) + flows = torch.zeros(1, 1, 1, 2, 4, 8) # already at target shape + # Pre-scale images so they are at (4, 8) like the scaler would produce. + scaled_images = adapter.scaler.fill(images.clone(), is_flow=False) + assert tuple(scaled_images.shape[-2:]) == (4, 8) + out = adapter.unscale( + {"images": scaled_images.clone(), "flows": flows.clone()}, + image_only=True, + ) + # Images are restored to (2, 4). + assert tuple(out["images"].shape[-2:]) == (2, 4) + # Flows left untouched. + torch.testing.assert_close(out["flows"], flows, atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# _to_cuda: device selection, fp16 gating, and the cuda-unavailable warning +# --------------------------------------------------------------------------- + + +class TestToCuda: + def test_default_does_not_move_and_keeps_fp32(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(2, 4)) + img = np.zeros((2, 4, 3), dtype=np.uint8) + out = adapter.prepare_inputs(images=img) + assert out["images"].device.type == "cpu" + assert out["images"].dtype == torch.float32 + + def test_to_cpu_device(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(2, 4), device="cpu") + img = np.zeros((2, 4, 3), dtype=np.uint8) + out = adapter.prepare_inputs(images=img) + assert out["images"].device.type == "cpu" + assert out["images"].dtype == torch.float32 + + def test_fp16_converts_only_image_and_flow_keys(self) -> None: + adapter = IOAdapter(output_stride=8, input_size=(2, 4), device="cpu", fp16=True) + img = np.zeros((2, 4, 3), dtype=np.uint8) + flow = np.zeros((2, 4, 2), dtype=np.float32) + 0.5 + meta = np.zeros((2, 4, 1), dtype=np.uint8) + out = adapter.prepare_inputs(images=img, flows=flow, meta=meta) + assert out["images"].dtype == torch.float16 + assert out["flows"].dtype == torch.float16 + # "meta" contains neither "image" nor "flow" -> stays fp32. + assert out["meta"].dtype == torch.float32 + + def test_fp16_has_no_effect_without_device(self) -> None: + # Documented behaviour: fp16 only runs after a device move. With no + # device and cuda disabled, tensors remain fp32 even though fp16=True. + adapter = IOAdapter( + output_stride=8, + input_size=(2, 4), + device=None, + cuda=False, + fp16=True, + ) + img = np.zeros((2, 4, 3), dtype=np.uint8) + out = adapter.prepare_inputs(images=img) + assert out["images"].dtype == torch.float32 + assert out["images"].device.type == "cpu" + + def test_cuda_true_with_no_gpu_warns_and_keeps_cpu( + self, caplog: pytest.LogCaptureFixture + ) -> None: + # Make this test independent of whether a real GPU is available: + # simulate the unavailable case directly. + import torch as _torch + + adapter = IOAdapter(output_stride=8, input_size=(2, 4), cuda=True) + import roco_spring_devkit.common.utils.io_adapter as mod + + original = _torch.cuda.is_available + _torch.cuda.is_available = staticmethod(lambda: False) + try: + img = np.zeros((2, 4, 3), dtype=np.uint8) + with caplog.at_level("WARNING", logger="root"): + out = adapter.prepare_inputs(images=img) + finally: + _torch.cuda.is_available = original + + if not mod or True: + # When cuda is unavailable the tensors stay on CPU in fp32. + assert out["images"].device.type == "cpu" + assert out["images"].dtype == torch.float32 + # The warning is emitted by the stdlib root logger. + assert any( + "torch.cuda.is_available() == False" in r.message for r in caplog.records + ) + + +# --------------------------------------------------------------------------- +# Known defect: scalar kwargs crash prepare_inputs +# --------------------------------------------------------------------------- + + +class TestMetadataKwargs: + def test_scalar_kwarg_is_preserved_as_metadata(self) -> None: + # After the bug fix, scalar kwargs are kept verbatim in the inputs + # dict instead of crashing with TypeError on `len(v)`. + img = np.zeros((2, 4, 3), dtype=np.uint8) + adapter = IOAdapter(output_stride=8, input_size=(2, 4)) + out = adapter.prepare_inputs(images=[img, img], epoch=3) + assert "epoch" in out + assert out["epoch"] == 3 + # Arrays supplied as kwargs are still transformed by ToTensor. + assert "images" in out + assert isinstance(out["images"], torch.Tensor) diff --git a/tests/common/utils/test_scene_flow_metrics.py b/tests/common/utils/test_scene_flow_metrics.py new file mode 100644 index 0000000..2e6420c --- /dev/null +++ b/tests/common/utils/test_scene_flow_metrics.py @@ -0,0 +1,719 @@ +"""Unit tests for roco_spring_devkit.common.utils.scene_flow_metrics. + +The expected numbers below are computed by hand from the Spring benchmark +reference implementation +(https://github.com/cv-stuttgart/springwebsite/blob/main/springeval/management/commands/evaluation.py), +which the docstring of ``_compute_total_wauc`` cites. + +Input format (batch_size = 1 for all correct-path tests): +* ``preds``: ``flows`` (B, 2, H, W), ``disparities`` (B, 1, H, W) = disp1, + ``disparities2`` (B, 1, H, W) = disp2. +* ``targets``: ``flows`` (B, 2, H, W), ``disparities`` (2, 1, H, W) packing + [disp1_target, disp2_target] along dim 0. ``valid_disparities`` (if provided) + must be packed the same way. + +Correct metric definitions (per the Spring reference, ``get_errors_sceneflow``): +* ``epe`` = L2 norm of (flow_pred - flow_target) over the 2 channel axes. +* ``1px_flow`` = 100 if epe > 1, else 0. +* ``flall`` = 100 if (epe > 3) AND (epe > 0.05 * flow_target_norm). Else 0. +* ``wauc`` = weighted AUC with thresholds delta_i = i/20 and weights + w_i = 1 - (i-1)/100 for i in [1, 100] (sum_wi = 50.5), scaled by 100. +* ``abs1`` / ``abs2`` = |disp1_pred - disp1_target| / |disp2_pred - disp2_target|. +* ``d1`` = 100 if (abs1 > 3) AND (abs1 > 0.05 * |disp1_target|). Else 0. + (Uses the DISPARITY target magnitude, per the Spring reference.) +* ``d2`` = 100 if (abs2 > 3) AND (abs2 > 0.05 * |disp2_target|). Else 0. + (Uses the DISPARITY target magnitude, per the Spring reference.) +* ``1px_all`` = 100 if ANY of 1px_flow, 1px_d1, 1px_d2 is 100 (OR). +* ``sfall`` = 100 if ANY of flall, d1, d2 is 100 (OR). +""" + +import math + +import pytest +import torch + +from roco_spring_devkit.common.utils.scene_flow_metrics import SceneFlowMetrics + + +def _zeros_preds(h: int = 2, w: int = 2) -> dict: + return { + "flows": torch.zeros(1, 2, h, w), + "disparities": torch.zeros(1, 1, h, w), + "disparities2": torch.zeros(1, 1, h, w), + } + + +def _packed_disp_targets(disp1_values, disp2_values, h: int, w: int) -> torch.Tensor: + """Build a (2, 1, h, w) disparity target packing [disp1, disp2] on dim 0.""" + d1 = torch.as_tensor(disp1_values, dtype=torch.float).reshape(1, 1, h, w) + d2 = torch.as_tensor(disp2_values, dtype=torch.float).reshape(1, 1, h, w) + return torch.cat([d1, d2], dim=0) + + +# --------------------------------------------------------------------------- +# Initialization +# --------------------------------------------------------------------------- + + +class TestSceneFlowMetricsInit: + def test_default_average_mode_is_epoch_mean(self) -> None: + m = SceneFlowMetrics() + assert m.average_mode == "epoch_mean" + + def test_default_attributes(self) -> None: + m = SceneFlowMetrics() + assert m.prefix == "" + assert m.ema_decay == 0.99 + assert m.f1_mode == "macro" + assert m.interpolate_pred_to_target_size is False + # ema_max_count = min(100, int(1 / (1 - 0.99))). Float rounding makes + # 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 == [] + + def test_invalid_average_mode_raises(self) -> None: + with pytest.raises(AssertionError): + SceneFlowMetrics(average_mode="unknown") + + def test_ema_decay_sets_short_max_count(self) -> None: + m = SceneFlowMetrics(average_mode="ema", ema_decay=0.9) + assert m.ema_max_count == 10 + + def test_prefix_attaches_to_metric_keys(self) -> None: + m = SceneFlowMetrics(prefix="val_") + m.update( + _zeros_preds(1, 2), + { + "flows": torch.zeros(1, 2, 1, 2), + "disparities": _packed_disp_targets([0, 0], [0, 0], 1, 2), + }, + ) + keys = list(m.calculate_metrics().keys()) + assert keys == [ + "val_epe", + "val_1px_flow", + "val_flall", + "val_wauc", + "val_abs1", + "val_d1", + "val_abs2", + "val_d2", + "val_1px_all", + "val_sfall", + ] + + +# --------------------------------------------------------------------------- +# Flow-side metrics (epe, 1px_flow, flall, wauc) with zero disparities +# --------------------------------------------------------------------------- + + +class TestFlowMetrics: + def test_epe_single_pixel_x_displacement(self) -> None: + # pred = 0, target flow x = 1 at one pixel of a 2x2 grid. + # epe = [1, 0, 0, 0] -> mean 0.25. Disparities are zero. + target_flows = torch.tensor( + [[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]] + ) + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 2, 2), + } + m = SceneFlowMetrics() + m.update(_zeros_preds(2, 2), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 0.25, abs_tol=1e-6) + assert math.isclose(metrics["1px_flow"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["flall"].item(), 0.0, abs_tol=1e-6) + + def test_wauc_hand_computed(self) -> None: + # Same setup: epe = [1, 0, 0, 0]. WAUC = 100 * (3*17.29 + 4*33.21) / (4*50.5) + # = 91.440594... + target_flows = torch.tensor( + [[[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]]] + ) + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 2, 2), + } + m = SceneFlowMetrics() + m.update(_zeros_preds(2, 2), targets) + assert math.isclose( + m.calculate_metrics()["wauc"].item(), + 100.0 * (3 * 17.29 + 4 * 33.21) / (4 * 50.5), + abs_tol=1e-4, + ) + + def test_1px_flow_true_when_epe_exceeds_one(self) -> None: + target_flows = torch.tensor([[[[2.0]], [[0.0]]]]) + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 1, 1), + } + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 2.0, abs_tol=1e-6) + assert math.isclose(metrics["1px_flow"].item(), 100.0, abs_tol=1e-6) + # epe = 2 is not > 3 -> flall stays 0. + assert math.isclose(metrics["flall"].item(), 0.0, abs_tol=1e-6) + + def test_flall_triggers_when_epe_gt_3_and_gt_5pct_target_norm(self) -> None: + # epe = 4, target_norm = 4 -> 0.05*4 = 0.2 < 4 -> flall = 100. + target_flows = torch.tensor([[[[4.0]], [[0.0]]]]) + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 1, 1), + } + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["flall"].item(), 100.0, abs_tol=1e-6) + + def test_flall_triggers_with_large_flow_target_norm(self) -> None: + # epe = 80, target_norm = 80 -> 0.05*80 = 4.0. 80 > 4.0 -> flall = 100. + target_flows = torch.tensor([[[[80.0]], [[0.0]]]]) + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 1, 1), + } + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 80.0, abs_tol=1e-6) + assert math.isclose(metrics["1px_flow"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["flall"].item(), 100.0, abs_tol=1e-6) + + def test_flall_blocked_when_epe_equals_5pct_target_norm(self) -> None: + # target = (80, 0), pred = (76, 0). epe = 4, target_norm = 80, + # 0.05 * 80 = 4.0 -> strict > gives False -> flall = 0. + preds = { + "flows": torch.tensor([[[[76.0]], [[0.0]]]]), + "disparities": torch.zeros(1, 1, 1, 1), + "disparities2": torch.zeros(1, 1, 1, 1), + } + target_flows = torch.tensor([[[[80.0]], [[0.0]]]]) + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 1, 1), + } + m = SceneFlowMetrics() + m.update(preds, targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["1px_flow"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["flall"].item(), 0.0, abs_tol=1e-6) + + +# --------------------------------------------------------------------------- +# Disparity-side metrics (abs1, abs2, d1, d2) +# --------------------------------------------------------------------------- + + +class TestDisparityMetrics: + def test_abs1_and_abs2_single_pixel_mismatch(self) -> None: + # disp1 target = 1 at one pixel, disp2 target = 2 at one pixel. + # abs1 = [1, 0, 0, 0] / 4 = 0.25; abs2 = [2, 0, 0, 0] / 4 = 0.5. + # No abs > 1 -> d1 = d2 = 0 regardless of which target_norm is used. + target_disp = _packed_disp_targets( + [[1.0, 0.0], [0.0, 0.0]], [[2.0, 0.0], [0.0, 0.0]], 2, 2 + ) + targets = { + "flows": torch.zeros(1, 2, 2, 2), + "disparities": target_disp, + } + m = SceneFlowMetrics() + m.update(_zeros_preds(2, 2), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs1"].item(), 0.25, abs_tol=1e-6) + assert math.isclose(metrics["abs2"].item(), 0.5, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["d2"].item(), 0.0, abs_tol=1e-6) + + def test_d1_triggers_when_abs1_gt_3_and_gt_5pct_disp_target(self) -> None: + # abs1 = 4, disp1_target = 4, flow target = (4, 0). + # Both buggy and correct: 0.05 * 4 = 0.2 < 4 -> d1 = 100. + # (flow_norm = 4 = disp1_target = 4, so both formulas agree here.) + target_flows = torch.tensor([[[[4.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[4.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs1"].item(), 4.0, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 100.0, abs_tol=1e-6) + + def test_d1_triggers_independently_of_d2(self) -> None: + # abs1 = 4 (> 3, disp1_target = 4 -> 0.05*4=0.2 -> d1 = 100), + # abs2 = 2 (not > 3 -> d2 = 0). + # Both formulas agree because flow_norm = 4 = disp1_target = 4. + target_flows = torch.tensor([[[[4.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[4.0]], [[2.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["d1"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["d2"].item(), 0.0, abs_tol=1e-6) + + def test_d2_triggers_independently_of_d1(self) -> None: + # abs2 = 4 (> 3, disp2_target = 4 -> 0.05*4=0.2 -> d2 = 100), + # abs1 = 2 (not > 3 -> d1 = 0). + # Both formulas agree because flow_norm = 4 = disp2_target = 4. + target_flows = torch.tensor([[[[4.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[2.0]], [[4.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["d1"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["d2"].item(), 100.0, abs_tol=1e-6) + + +# --------------------------------------------------------------------------- +# Combined metrics (1px_all, sfall) +# --------------------------------------------------------------------------- +# Per the Spring reference: sf = d1 | d2 | fl (OR), and +# onepx_sf = onepx_d1 | onepx_d2 | onepx_fl (OR). +# A pixel fails the combined metric if ANY component fails. +# --------------------------------------------------------------------------- + + +class TestCombinedMetrics: + def test_all_three_fail_1px_and_sfall(self) -> None: + # epe = 4, abs1 = 4, abs2 = 4, flow_norm = 4, disp targets = 4. + # All three fail both 1px (>1) and the d1/d2/flall thresholds (>3, >0.2). + # Both AND and OR give 100 for 1px_all and sfall. + target_flows = torch.tensor([[[[4.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[4.0]], [[4.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["flall"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["d2"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["1px_all"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["sfall"].item(), 100.0, abs_tol=1e-6) + + def test_all_three_below_1px(self) -> None: + # epe = 0.5, abs1 = 0.5, abs2 = 0.5. None exceed 1 or 3. + # Both AND and OR give 0 for 1px_all and sfall. + target_flows = torch.tensor([[[[0.5]], [[0.0]]]]) + target_disp = _packed_disp_targets([[0.5]], [[0.5]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["1px_all"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["sfall"].item(), 0.0, abs_tol=1e-6) + + def test_1px_all_all_three_fail_gives_100_both_and_and_or(self) -> None: + # epe = 2, abs1 = 2, abs2 = 2. All > 1 so 1px_all = 100 under both AND/OR. + # None > 3 so flall = d1 = d2 = 0 -> sfall = 0 under both AND/OR. + target_flows = torch.tensor([[[[2.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[2.0]], [[2.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["1px_all"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["sfall"].item(), 0.0, abs_tol=1e-6) + + +# --------------------------------------------------------------------------- +# Valid mask handling +# --------------------------------------------------------------------------- + + +class TestValidMasks: + def test_valid_flows_excludes_invalid_pixels_from_epe(self) -> None: + # 3 pixels with epe = 1 each; valid_flows = [1, 1, 0]. + # Valid mean = (1 + 1) / 2 = 1.0. + preds = _zeros_preds(1, 3) + target_flows = torch.tensor([[[[1.0, 1.0, 1.0]], [[0.0, 0.0, 0.0]]]]) + target_disp = _packed_disp_targets([[0, 0, 0]], [[0, 0, 0]], 1, 3) + valid_flows = torch.tensor([[[[1.0, 1.0, 0.0]]]]) + targets = { + "flows": target_flows, + "disparities": target_disp, + "valid_flows": valid_flows, + } + m = SceneFlowMetrics() + m.update(preds, targets) + assert math.isclose(m.calculate_metrics()["epe"].item(), 1.0, abs_tol=1e-6) + + def test_valid_disparities_must_be_packed_for_disp1_and_disp2(self) -> None: + # valid_disparities packed as (2, 1, 1, 3): [valid_disp1, valid_disp2]. + # abs1 = 1 at all 3 px, valid_disp1 = [1, 1, 0] -> abs1 metric = 1.0. + # abs2 = 1 at all 3 px, valid_disp2 = [1, 1, 0] -> abs2 metric = 1.0. + preds = _zeros_preds(1, 3) + target_flows = torch.zeros(1, 2, 1, 3) + target_disp = _packed_disp_targets([[1, 1, 1]], [[1, 1, 1]], 1, 3) + valid_flows = torch.ones(1, 1, 1, 3) + valid_disp = torch.zeros(2, 1, 1, 3) + valid_disp[0, 0, 0, :] = torch.tensor([1.0, 1.0, 0.0]) + valid_disp[1, 0, 0, :] = torch.tensor([1.0, 1.0, 0.0]) + targets = { + "flows": target_flows, + "disparities": target_disp, + "valid_flows": valid_flows, + "valid_disparities": valid_disp, + } + m = SceneFlowMetrics() + m.update(preds, targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs1"].item(), 1.0, abs_tol=1e-6) + assert math.isclose(metrics["abs2"].item(), 1.0, abs_tol=1e-6) + + def test_valid_all_target_is_product_of_three_valid_masks(self) -> None: + # valid_flows = [1, 1, 0], valid_disp1 = [1, 1, 0], valid_disp2 = [1, 1, 0]. + # valid_all = [1, 1, 0]. With epe=2, abs1=2, abs2=2 -> 1px_all = 100 at + # the 2 valid pixels -> mean 100. (All three fail, so AND and OR agree.) + preds = _zeros_preds(1, 3) + target_flows = torch.tensor([[[[2.0, 2.0, 2.0]], [[0.0, 0.0, 0.0]]]]) + target_disp = _packed_disp_targets([[2, 2, 2]], [[2, 2, 2]], 1, 3) + valid_flows = torch.tensor([[[[1.0, 1.0, 0.0]]]]) + valid_disp = torch.zeros(2, 1, 1, 3) + valid_disp[0, 0, 0, :] = torch.tensor([1.0, 1.0, 0.0]) + valid_disp[1, 0, 0, :] = torch.tensor([1.0, 1.0, 0.0]) + targets = { + "flows": target_flows, + "disparities": target_disp, + "valid_flows": valid_flows, + "valid_disparities": valid_disp, + } + m = SceneFlowMetrics() + m.update(preds, targets) + assert math.isclose( + m.calculate_metrics()["1px_all"].item(), 100.0, abs_tol=1e-6 + ) + + +# --------------------------------------------------------------------------- +# Epoch mean accumulation +# --------------------------------------------------------------------------- + + +class TestEpochMeanAccumulation: + def test_two_steps_average_per_sample_means(self) -> None: + # Step 1: epe = 1.0 per sample mean. Step 2: epe = 0.5. + # epoch_mean = (1.0 + 0.5) / 2 = 0.75. + m = SceneFlowMetrics() + m.update( + _zeros_preds(1, 2), + { + "flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 2), + }, + ) + m.update( + _zeros_preds(1, 2), + { + "flows": torch.tensor([[[[0.0, 1.0]], [[0.0, 0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 2), + }, + ) + assert math.isclose(m.calculate_metrics()["epe"].item(), 0.75, abs_tol=1e-6) + assert m.sample_count.item() == 2.0 + + def test_calculate_metrics_does_not_reset_state(self) -> None: + m = SceneFlowMetrics() + m.update( + _zeros_preds(1, 1), + { + "flows": torch.tensor([[[[2.0]], [[0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 1), + }, + ) + first = m.calculate_metrics() + second = m.calculate_metrics() + for key in first: + assert torch.allclose(first[key], second[key]) + + +# --------------------------------------------------------------------------- +# EMA mode +# --------------------------------------------------------------------------- + + +class TestAverageModeEma: + def test_ema_step1_returns_input_mean(self) -> None: + m = SceneFlowMetrics(average_mode="ema", ema_decay=0.9) + m.update( + _zeros_preds(1, 2), + { + "flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 2), + }, + ) + # EPE per-sample mean = 1.0. state = 0.1 * 1.0 = 0.1. divider = 1 - 0.9 = 0.1. + # metric = 0.1 / 0.1 = 1.0. + assert math.isclose(m.calculate_metrics()["epe"].item(), 1.0, abs_tol=1e-5) + + def test_ema_step2_hand_computed(self) -> None: + m = SceneFlowMetrics(average_mode="ema", ema_decay=0.9) + m.update( + _zeros_preds(1, 2), + { + "flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 2), + }, + ) + m.update( + _zeros_preds(1, 2), + { + "flows": torch.tensor([[[[0.0, 1.0]], [[0.0, 0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 2), + }, + ) + # state = 0.9*0.1 + 0.1*0.5 = 0.14. divider = 1 - 0.9^2 = 0.19. + # metric = 0.14 / 0.19 = 0.7368... + assert math.isclose( + m.calculate_metrics()["epe"].item(), 0.14 / 0.19, abs_tol=1e-5 + ) + + def test_ema_divisor_becomes_one_after_ema_max_count(self) -> None: + m = SceneFlowMetrics(average_mode="ema", ema_decay=0.9) + targets = { + "flows": torch.tensor([[[[1.0]], [[0.0]]]]), + "disparities": torch.zeros(2, 1, 1, 1), + } + for _ in range(11): + m.update(_zeros_preds(1, 1), targets) + assert m.step_count.item() == 11.0 + expected = 1.0 - 0.9**11 + assert math.isclose(m.calculate_metrics()["epe"].item(), expected, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# interpolate_pred_to_target_size +# --------------------------------------------------------------------------- + + +class TestInterpolatePredToTargetSize: + def test_flow_and_disparity_rescaled_to_target_size(self) -> None: + # Pred shape (1, 2, 1, 2) full of 1.0. Target shape (2, 4). + # scale_y = 2, scale_x = 2. After interp: + # flow x-channel *= 2 -> 2.0; flow y-channel *= 2 -> 2.0. + # disp x-channel (channel 0) *= 2 -> 2.0. + # Target = zeros -> epe = norm([2, 2]) = sqrt(8) per pixel; abs1 = abs2 = 2. + preds = { + "flows": torch.full((1, 2, 1, 2), 1.0), + "disparities": torch.full((1, 1, 1, 2), 1.0), + "disparities2": torch.full((1, 1, 1, 2), 1.0), + } + targets = { + "flows": torch.zeros(1, 2, 2, 4), + "disparities": torch.zeros(2, 1, 2, 4), + } + m = SceneFlowMetrics(interpolate_pred_to_target_size=True) + m.update(preds, targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), math.sqrt(8.0), abs_tol=1e-5) + assert math.isclose(metrics["abs1"].item(), 2.0, abs_tol=1e-5) + assert math.isclose(metrics["abs2"].item(), 2.0, abs_tol=1e-5) + + def test_zeros_pred_uses_target_only(self) -> None: + # Pred = zeros interpolated to (2, 4). Target flow x = 3 at one pixel. + # epe = 3 at that pixel, 0 elsewhere -> sum 3 / 8 px = 0.375. + preds = _zeros_preds(1, 2) + ch_x = torch.tensor([[3.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]) + ch_y = torch.zeros_like(ch_x) + target_flows = torch.stack([ch_x, ch_y], dim=0)[None] + targets = { + "flows": target_flows, + "disparities": torch.zeros(2, 1, 2, 4), + } + m = SceneFlowMetrics(interpolate_pred_to_target_size=True) + m.update(preds, targets) + assert math.isclose(m.calculate_metrics()["epe"].item(), 3 / 8, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# Shape / dtype handling +# --------------------------------------------------------------------------- + + +class TestShapeHandling: + def test_double_precision_input_is_cast_to_float32(self) -> None: + preds = { + "flows": torch.zeros(1, 2, 1, 2, dtype=torch.float64), + "disparities": torch.zeros(1, 1, 1, 2, dtype=torch.float64), + "disparities2": torch.zeros(1, 1, 1, 2, dtype=torch.float64), + } + targets = { + "flows": torch.tensor([[[[1.0, 1.0]], [[0.0, 0.0]]]], dtype=torch.float64), + "disparities": torch.zeros(2, 1, 1, 2, dtype=torch.float64), + } + m = SceneFlowMetrics() + m.update(preds, targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["epe"].item(), 1.0, abs_tol=1e-6) + assert metrics["epe"].dtype == torch.float32 + + def test_missing_disparities2_raises_key_error(self) -> None: + preds = { + "flows": torch.zeros(1, 2, 1, 1), + "disparities": torch.zeros(1, 1, 1, 1), + } + targets = { + "flows": torch.zeros(1, 2, 1, 1), + "disparities": torch.zeros(2, 1, 1, 1), + } + m = SceneFlowMetrics() + with pytest.raises(KeyError): + m.update(preds, targets) + + +# --------------------------------------------------------------------------- +# Tests for d1/d2 using disparity target magnitudes (not flow target norm), +# sfall/1px_all using OR (not AND), batch > 1 disparity slicing, and +# multi-hypothesis flow targets via 6D inputs. +# +# Per the Spring reference (get_errors_sceneflow): +# gt_len_d1 = np.abs(gt_d1).max(axis=-1) # disparity1 target magnitude +# gt_len_d2 = np.abs(gt_d2).max(axis=-1) # disparity2 target magnitude +# gt_veclen = np.linalg.norm(gt_flow, ...) # flow target norm +# d1 = (d1_epe > 3) & (d1_epe > 0.05 * gt_len_d1) +# d2 = (d2_epe > 3) & (d2_epe > 0.05 * gt_len_d2) +# fl = (fl_epe > 3) & (fl_epe > 0.05 * gt_veclen_max) +# sf = d1 | d2 | fl # OR, not AND +# onepx_sf = onepx_d1 | onepx_d2 | onepx_fl # OR, not AND +# --------------------------------------------------------------------------- + + +class TestD1UsesDispTargetNorm: + def test_d1_uses_disp_target_not_flow_norm(self) -> None: + # abs1 = 4, disp1_target = 4, flow = (80, 0) -> flow_norm = 80. + # d1 = (4 > 3) & (4 > 0.05*4 = 0.2) = True -> 100. + # If d1 used flow_norm: (4 > 0.05*80 = 4.0) = False -> 0 (wrong). + target_flows = torch.tensor([[[[80.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[4.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + assert math.isclose(m.calculate_metrics()["d1"].item(), 100.0, abs_tol=1e-6) + + def test_d1_blocks_on_large_disp_target_not_flow(self) -> None: + # abs1 = 4 (pred=76, target=80), disp1_target = 80, flow = (0, 0). + # d1 = (4 > 3) & (4 > 0.05*80 = 4.0) = False -> 0. + # If d1 used flow_norm: (4 > 0.05*0 = 0) = True -> 100 (wrong). + preds = { + "flows": torch.zeros(1, 2, 1, 1), + "disparities": torch.tensor([[[[76.0]]]]), + "disparities2": torch.zeros(1, 1, 1, 1), + } + target_flows = torch.zeros(1, 2, 1, 1) + target_disp = _packed_disp_targets([[80.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(preds, targets) + assert math.isclose(m.calculate_metrics()["d1"].item(), 0.0, abs_tol=1e-6) + + +class TestD2UsesDispTargetNorm: + def test_d2_uses_disp_target_not_flow_norm(self) -> None: + # abs2 = 4, disp2_target = 4, flow = (80, 0) -> flow_norm = 80. + # d2 = (4 > 3) & (4 > 0.05*4 = 0.2) = True -> 100. + target_flows = torch.tensor([[[[80.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[0.0]], [[4.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + assert math.isclose(m.calculate_metrics()["d2"].item(), 100.0, abs_tol=1e-6) + + +class TestSfallUsesOr: + def test_sfall_100_when_only_flall_fails(self) -> None: + # epe = 4, flow = (4, 0) -> flall = 100. abs1 = 0, abs2 = 0 -> d1 = d2 = 0. + # sfall = flall | d1 | d2 = 100 (OR). + target_flows = torch.tensor([[[[4.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[0.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["flall"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["sfall"].item(), 100.0, abs_tol=1e-6) + + def test_sfall_100_when_only_d1_fails(self) -> None: + # flow = 0 -> flall = 0. abs1 = 4, disp1 = 4 -> d1 = 100. abs2 = 0 -> d2 = 0. + # sfall = flall | d1 | d2 = 100 (OR). + target_flows = torch.zeros(1, 2, 1, 1) + target_disp = _packed_disp_targets([[4.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["d1"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["sfall"].item(), 100.0, abs_tol=1e-6) + + +class Test1pxAllUsesOr: + def test_1px_all_100_when_only_flow_fails_1px(self) -> None: + # epe = 2 (> 1), abs1 = 0, abs2 = 0. + # 1px_all = 1px_flow | 1px_d1 | 1px_d2 = 100 (OR). + target_flows = torch.tensor([[[[2.0]], [[0.0]]]]) + target_disp = _packed_disp_targets([[0.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + metrics = m.calculate_metrics() + assert math.isclose(metrics["1px_flow"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["1px_all"].item(), 100.0, abs_tol=1e-6) + + def test_1px_all_100_when_only_disp1_fails_1px(self) -> None: + # epe = 0, abs1 = 2 (> 1), abs2 = 0. + # 1px_all = 1px_flow | 1px_d1 | 1px_d2 = 100 (OR). + target_flows = torch.zeros(1, 2, 1, 1) + target_disp = _packed_disp_targets([[2.0]], [[0.0]], 1, 1) + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(_zeros_preds(1, 1), targets) + assert math.isclose( + m.calculate_metrics()["1px_all"].item(), 100.0, abs_tol=1e-6 + ) + + +class TestBatchGtOneDispSlicing: + def test_batch_two_uses_correct_disparity_targets_per_batch(self) -> None: + preds = { + "flows": torch.zeros(2, 2, 1, 1), + "disparities": torch.zeros(2, 1, 1, 1), + "disparities2": torch.zeros(2, 1, 1, 1), + } + target_flows = torch.zeros(2, 2, 1, 1) + target_disp = torch.zeros(4, 1, 1, 1) + target_disp[0, 0, 0, 0] = 10.0 # disp1 batch 0 + target_disp[1, 0, 0, 0] = 20.0 # disp2 batch 0 + target_disp[2, 0, 0, 0] = 30.0 # disp1 batch 1 + target_disp[3, 0, 0, 0] = 40.0 # disp2 batch 1 + targets = {"flows": target_flows, "disparities": target_disp} + m = SceneFlowMetrics() + m.update(preds, targets) + metrics = m.calculate_metrics() + # abs1 = (10+30)/2 = 20, abs2 = (20+40)/2 = 30. + assert math.isclose(metrics["abs1"].item(), 20.0, abs_tol=1e-6) + assert math.isclose(metrics["abs2"].item(), 30.0, abs_tol=1e-6) + + +class TestFiveDimFlowTargetHypotheses: + def test_six_dim_flow_target_uses_min_over_hypotheses(self) -> None: + preds = { + "flows": torch.zeros(1, 2, 1, 2), + "disparities": torch.zeros(1, 1, 1, 2), + "disparities2": torch.zeros(1, 1, 1, 2), + } + h0 = torch.zeros(1, 2, 1, 2) + h1 = torch.tensor([[1.0, 1.0], [0.0, 0.0]]).reshape(1, 2, 1, 2) + flow6d = torch.stack([h0, h1], dim=1) + flow6d = flow6d[None] # (1, 1, 2, 2, 1, 2) + targets = { + "flows": flow6d, + "disparities": torch.zeros(2, 1, 1, 2), + } + m = SceneFlowMetrics() + m.update(preds, targets) + # The 5D path picks h0 (epe = 0) over h1 (epe > 0). + assert math.isclose(m.calculate_metrics()["epe"].item(), 0.0, abs_tol=1e-6) diff --git a/tests/common/utils/test_stereo_metrics.py b/tests/common/utils/test_stereo_metrics.py new file mode 100644 index 0000000..0ddb14e --- /dev/null +++ b/tests/common/utils/test_stereo_metrics.py @@ -0,0 +1,532 @@ +"""Unit tests for roco_spring_devkit.common.utils.stereo_metrics. + +The expected numbers below are computed by hand from the documented semantics of +``StereoMetrics`` and the Spring benchmark reference: + +* ``abs`` per pixel is ``|disp_pred - disp_target|`` (single channel). The + reported value is the average of the per-sample means of the valid pixels, + averaged again across samples (``epoch_mean`` divides the accumulated + per-sample totals by ``sample_count``). +* ``1px`` is ``100`` if a pixel's absolute error is strictly greater than 1, + ``0`` otherwise, with the same averaging. +* ``d1`` is ``100`` if ``(abs > 3) AND (abs > 0.05 * |disp_target|)``. Else ``0``. + Per the Spring reference, the threshold uses the ground truth disparity + magnitude, not the error magnitude. +* ``f1`` variants binarize ``pred > 0.5`` and ``target > 0.5`` and support + ``binary``/``macro``/``weighted`` modes. +* The EMA accumulator is ``state = ema_decay * state + (1 - ema_decay) * total``; + the normalization divisor is ``1 - ema_decay**step_count`` until ``step_count`` + reaches ``ema_max_count`` and ``1.0`` afterwards. +""" + +import math + +import pytest +import torch + +from roco_spring_devkit.common.utils.stereo_metrics import StereoMetrics + + +def _zeros_disp(h: int = 1, w: int = 2) -> torch.Tensor: + return torch.zeros(1, 1, h, w) + + +# --------------------------------------------------------------------------- +# Initialization +# --------------------------------------------------------------------------- + + +class TestStereoMetricsInit: + def test_default_average_mode_is_epoch_mean(self) -> None: + m = StereoMetrics() + assert m.average_mode == "epoch_mean" + + def test_default_attributes(self) -> None: + m = StereoMetrics() + assert m.prefix == "" + assert m.ema_decay == 0.99 + assert m.f1_mode == "macro" + assert m.interpolate_pred_to_target_size is False + # ema_max_count = min(100, int(1 / (1 - ema_decay))). With ema_decay = 0.99 + # the floating-point subtraction yields 0.010000000000000009, so + # 1 / 0.0100... ~= 99.999... which floors to 99 under ``int(...)``. + assert m.ema_max_count == 99 + assert m.include_occlusion is False + assert m.used_keys == [] + + def test_invalid_average_mode_raises(self) -> None: + with pytest.raises(AssertionError): + StereoMetrics(average_mode="unknown") + + def test_ema_decay_sets_short_max_count(self) -> None: + # ema_decay = 0.9 -> 1/(1-0.9) = 10, capped at min(100, 10). + m = StereoMetrics(average_mode="ema", ema_decay=0.9) + assert m.ema_max_count == 10 + + def test_prefix_attaches_to_metric_keys(self) -> None: + m = StereoMetrics(prefix="val_") + m.update( + {"disparities": torch.zeros(1, 1, 1, 2)}, + {"disparities": torch.tensor([[[[1.0, 0.0]]]])}, + ) + keys = list(m.calculate_metrics().keys()) + assert keys == ["val_abs", "val_1px", "val_d1"] + + +# --------------------------------------------------------------------------- +# abs / 1px / d1 (epoch_mean, no occlusion, no valid_disparities, batch=1) +# --------------------------------------------------------------------------- + + +class TestBasicMetricsEpochMean: + def test_hand_computed_single_pixel_disparity_mismatch(self) -> None: + # pred = 0, target = 1 at one pixel of a 2x2 grid; per-pixel abs = [1, 0, 0, 0]. + # per-sample mean = 0.25. + preds = {"disparities": torch.zeros(1, 1, 2, 2)} + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 0.25, abs_tol=1e-6) + # No pixel has abs > 1 (the only value at the threshold is not strictly >). + assert math.isclose(metrics["1px"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 0.0, abs_tol=1e-6) + + def test_1px_true_when_abs_exceeds_one(self) -> None: + # pred = 0, target = 2 at the single pixel -> abs = 2 > 1 -> 1px = 100. + preds = {"disparities": torch.zeros(1, 1, 1, 1)} + target = torch.tensor([[[[2.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 2.0, abs_tol=1e-6) + assert math.isclose(metrics["1px"].item(), 100.0, abs_tol=1e-6) + # abs = 2 is not > 3 -> d1 stays 0. + assert math.isclose(metrics["d1"].item(), 0.0, abs_tol=1e-6) + + def test_d1_triggers_when_abs_exceeds_three(self) -> None: + # pred = 0, target = 4 -> abs = 4. d1 = (abs > 3) & (abs > 0.05*abs=0.2) = True. + preds = {"disparities": torch.zeros(1, 1, 1, 1)} + target = torch.tensor([[[[4.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 4.0, abs_tol=1e-6) + assert math.isclose(metrics["1px"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 100.0, abs_tol=1e-6) + + def test_d1_blocked_by_large_disparity_target(self) -> None: + # abs = 4, |disp_target| = 80 -> 0.05*80 = 4.0. Strict > gives False. + # d1 = (4 > 3) & (4 > 4.0) = False -> 0. + # Per the Spring reference, d1 uses 0.05 * |gt_disp|, not 0.05 * |abs|. + preds = {"disparities": torch.tensor([[[[76.0]]]])} + target = torch.tensor([[[[80.0]]]]) # abs = 4, |disp_target| = 80 + m = StereoMetrics() + m.update(preds, {"disparities": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 4.0, abs_tol=1e-6) + assert math.isclose(metrics["1px"].item(), 100.0, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 0.0, abs_tol=1e-6) + + def test_d1_zero_when_abs_is_zero(self) -> None: + # abs = 0, |disp_target| = 0 -> 0.05*0 = 0. (0 > 3) is False -> d1 = 0. + preds = {"disparities": torch.zeros(1, 1, 1, 1)} + target = torch.zeros(1, 1, 1, 1) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 0.0, abs_tol=1e-6) + assert math.isclose(metrics["d1"].item(), 0.0, abs_tol=1e-6) + + def test_epoch_mean_aggregates_per_sample_means_across_steps(self) -> None: + # Two updates (batch=1 each): per-sample abs means 1.0 then 0.5. + # sample_count = 2 -> abs metric = (1.0 + 0.5) / 2 == 0.75. + m = StereoMetrics() + m.update( + {"disparities": torch.zeros(1, 1, 1, 2)}, + {"disparities": torch.tensor([[[[1.0, 1.0]]]])}, + ) + m.update( + {"disparities": torch.zeros(1, 1, 1, 2)}, + {"disparities": torch.tensor([[[[0.0, 1.0]]]])}, + ) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 0.75, abs_tol=1e-6) + assert m.sample_count.item() == 2.0 + + def test_epoch_mean_weights_each_sample_equally_not_per_pixel(self) -> None: + # Two batches with different per-sample pixel counts: the per-sample mean + # for each batch is 1.0, so the epoch mean is (1.0 + 1.0) / 2 == 1.0 + # rather than a 4-pixel average. + m = StereoMetrics() + m.update( + {"disparities": torch.zeros(1, 1, 1, 1)}, + { + "disparities": torch.tensor([[[[1.0]]]]), + "valid_disparities": torch.tensor([[[[1.0]]]]), + }, + ) + m.update( + {"disparities": torch.zeros(1, 1, 1, 3)}, + { + "disparities": torch.tensor([[[[1.0, 1.0, 1.0]]]]), + "valid_disparities": torch.tensor([[[[1.0, 1.0, 1.0]]]]), + }, + ) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 1.0, abs_tol=1e-6) + assert m.sample_count.item() == 2.0 + + def test_calculate_metrics_does_not_reset_state(self) -> None: + m = StereoMetrics() + m.update( + {"disparities": torch.zeros(1, 1, 1, 1)}, + {"disparities": torch.tensor([[[[2.0]]]])}, + ) + first = m.calculate_metrics() + second = m.calculate_metrics() + assert set(first.keys()) == set(second.keys()) + for key in first: + assert torch.allclose(first[key], second[key]) + + +# --------------------------------------------------------------------------- +# valid_disparities handling +# --------------------------------------------------------------------------- + + +class TestValidDisparities: + def test_mask_excludes_invalid_pixels_from_abs(self) -> None: + # Three pixels each with abs = 1; valid = [1, 1, 0]. The valid mean + # is (1 + 1) / 2 == 1.0, not 1.0 over three pixels. + preds = {"disparities": torch.zeros(1, 1, 1, 3)} + target = torch.tensor([[[[1.0, 1.0, 1.0]]]]) + valid = torch.tensor([[[[1.0, 1.0, 0.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target, "valid_disparities": valid}) + assert math.isclose(m.calculate_metrics()["abs"].item(), 1.0, abs_tol=1e-6) + + def test_occlusion_distinguishes_occ_and_non_occ(self) -> None: + # pred = 0, target = 1 at one pixel of a 2x2 grid; occ on the bottom row. + # Per-pixel abs = [1, 0, 0, 0]; occ splits -> valid_occ picks [0, 0] + # (mean 0), valid_non_occ picks [1, 0] (mean 0.5). Global per-sample mean + # across 4 valid pixels = 0.25. + preds = {"disparities": torch.zeros(1, 1, 2, 2)} + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target, "occs": occ}) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 0.25, abs_tol=1e-6) + assert math.isclose(metrics["abs_non_occ"].item(), 0.5, abs_tol=1e-6) + assert math.isclose(metrics["abs_occ"].item(), 0.0, abs_tol=1e-6) + + +# --------------------------------------------------------------------------- +# Occlusion prediction -> occ_f1 +# --------------------------------------------------------------------------- + + +class TestOcclusionHandling: + def test_occs_target_extends_used_keys_with_occ_and_non_occ(self) -> None: + preds = {"disparities": torch.zeros(1, 1, 2, 2)} + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target, "occs": occ}) + keys = [k for k, _, _ in m.used_keys] + for base in ("abs", "1px", "d1"): + assert base in keys + assert f"{base}_occ" in keys + assert f"{base}_non_occ" in keys + + def test_occs_pred_adds_occ_f1_with_perfect_match(self) -> None: + preds = { + "disparities": torch.zeros(1, 1, 2, 2), + "occs": torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]), + } + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + for mode in ("binary", "macro", "weighted"): + m = StereoMetrics(f1_mode=mode) + m.update(preds, {"disparities": target, "occs": occ}) + assert math.isclose( + m.calculate_metrics()["occ_f1"].item(), 1.0, abs_tol=1e-5 + ) + assert "occ_f1" in [k for k, _, _ in m.used_keys] + + def test_occ_f1_macro_default_with_all_negative_pred(self) -> None: + # Target occ = [0, 1, 1] (2 positives, 1 negative). Predict all zeros. + # Positive class: pred_bin all 0 -> tp=0, fp=2, fn=0; precision=recall=0 -> f1_pos=0. + # Negative class: pred_neg_bin all 1; target_neg_bin = [1, 0, 0] (1 pos); + # tp_neg = 1, fp_neg = 0, fn_neg = 2; precision = 1, recall = 1/3; + # f1_neg = 2*1*(1/3)/(1+1/3) = 0.5. + # macro default = (0 + 0.5)/2 = 0.25. + preds = { + "disparities": torch.zeros(1, 1, 1, 3), + "occs": torch.zeros(1, 1, 1, 3), # always predicts negative + } + target_occ = torch.tensor([[[[0.0, 1.0, 1.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": torch.zeros(1, 1, 1, 3), "occs": target_occ}) + assert math.isclose(m.calculate_metrics()["occ_f1"].item(), 0.25, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# Motion boundary and confidence metrics +# --------------------------------------------------------------------------- + + +class TestMbF1: + def test_mb_f1_macro_default_with_all_negative_pred(self) -> None: + # mb_target = [0,0,1,1] (2 positive). mb_pred = all zeros. + # f1_pos = 0 (no positive predictions). For the negative class: + # pred_neg_bin all 1, target_neg_bin = [1,1,0,0] (2 positives), + # tp_neg = 2, fp_neg = 0, fn_neg = 2; precision = 1, recall = 0.5; + # f1_neg = 2*1*0.5/(1+0.5) = 0.6667. + # macro = (0 + 0.6667)/2 = 0.3333. + preds = { + "disparities": torch.zeros(1, 1, 1, 4), + "mbs": torch.zeros(1, 1, 1, 4), + } + mb_target = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": torch.zeros(1, 1, 1, 4), "mbs": mb_target}) + assert math.isclose( + m.calculate_metrics()["mb_f1"].item(), 0.33333333, abs_tol=1e-5 + ) + + def test_mb_f1_requires_both_pred_and_target(self) -> None: + preds = {"disparities": torch.zeros(1, 1, 1, 1), "mbs": torch.zeros(1, 1, 1, 1)} + m = StereoMetrics() + m.update(preds, {"disparities": torch.zeros(1, 1, 1, 1)}) + assert "mb_f1" not in [k for k, _, _ in m.used_keys] + + +class TestConfF1: + def test_conf_target_is_exp_of_squared_disp_error(self) -> None: + # Single-channel disparity: conf_target = exp(-(disp_pred - disp_target)^2). + # With target = [1, 0, 0, 0] and pred = 0: + # conf_target = [exp(-1), 1, 1, 1] -> binarized > 0.5 -> [0, 1, 1, 1]. + # conf_pred = zeros -> pred_bin all 0. + # Positive class: tp=0, fp=3, fn=0 -> precision=recall=0 -> f1_pos=0. + # Negative class: pred_neg all 1, target_neg = [1, 0, 0, 0] (1 positive); + # tp=1, fp=0, fn=3; precision=1, recall=1/4=0.25; + # f1_neg = 2*1*0.25/(1.25) = 0.4. + # macro = (0 + 0.4)/2 = 0.2. + preds = { + "disparities": torch.zeros(1, 1, 1, 4), + "confs": torch.zeros(1, 1, 1, 4), + } + target = torch.tensor([[[[1.0, 0.0, 0.0, 0.0]]]]) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + assert math.isclose(m.calculate_metrics()["conf_f1"].item(), 0.2, abs_tol=1e-5) + + def test_conf_target_all_positive_yields_macro_half(self) -> None: + # When disp_pred == disp_target everywhere, conf_target = exp(0) = 1 at + # all pixels (positive class). With conf_pred = ones (all positive): + # f1_pos = 1 (perfect); negative class empty (target_neg all 0); + # f1_neg = 0 (tp=0, fn=0+precision/recall=0). + # macro = (1 + 0)/2 = 0.5. + preds = { + "disparities": torch.zeros(1, 1, 2, 2), + "confs": torch.ones(1, 1, 2, 2), + } + target = torch.zeros(1, 1, 2, 2) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + assert math.isclose(m.calculate_metrics()["conf_f1"].item(), 0.5, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# F1 modes +# --------------------------------------------------------------------------- + + +class TestF1Modes: + @pytest.mark.parametrize("mode", ["binary", "macro", "weighted"]) + def test_perfect_match_gives_f1_one_in_all_modes(self, mode: str) -> None: + preds = { + "disparities": torch.zeros(1, 1, 2, 2), + "occs": torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]), + } + target = torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]) + occ = torch.tensor([[[[0.0, 0.0], [1.0, 1.0]]]]) + m = StereoMetrics(f1_mode=mode) + m.update(preds, {"disparities": target, "occs": occ}) + assert math.isclose(m.calculate_metrics()["occ_f1"].item(), 1.0, abs_tol=1e-5) + + def test_binary_mode_returns_only_positive_class_f1(self) -> None: + # With the mb scenario (macro 0.3333), binary should return 0 (positive + # class f1 == 0 since pred never predicts positive). + preds = {"disparities": torch.zeros(1, 1, 1, 4), "mbs": torch.zeros(1, 1, 1, 4)} + mb_target = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + m = StereoMetrics(f1_mode="binary") + m.update(preds, {"disparities": torch.zeros(1, 1, 1, 4), "mbs": mb_target}) + assert math.isclose(m.calculate_metrics()["mb_f1"].item(), 0.0, abs_tol=1e-5) + + def test_weighted_mode_writes_one_when_balanced_class_counts(self) -> None: + # Half positive / half negative target and matching pred -> weighted f1 = 1. + target_occ = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + pred_match = torch.tensor([[[[0.0, 0.0, 1.0, 1.0]]]]) + m = StereoMetrics(f1_mode="weighted") + m.update( + {"disparities": torch.zeros(1, 1, 1, 4), "occs": pred_match}, + {"disparities": torch.zeros(1, 1, 1, 4), "occs": target_occ}, + ) + assert math.isclose(m.calculate_metrics()["occ_f1"].item(), 1.0, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# EMA mode +# --------------------------------------------------------------------------- + + +class TestAverageModeEma: + def test_ema_decay_0_9_step1_returns_input_mean(self) -> None: + # EMA state = 0.9 * 0 + 0.1 * _compute_total. _compute_total uses .mean() + # in EMA so step1 contributes exactly 1.0 (per-sample mean for the basic + # case). Divider = 1 - 0.9^1 = 0.1 -> reported metric = 1.0. + m = StereoMetrics(average_mode="ema", ema_decay=0.9) + m.update( + {"disparities": torch.zeros(1, 1, 1, 2)}, + {"disparities": torch.tensor([[[[1.0, 1.0]]]])}, + ) + metrics = m.calculate_metrics() + assert math.isclose(metrics["abs"].item(), 1.0, abs_tol=1e-5) + assert m.step_count.item() == 1.0 + + def test_ema_decay_0_9_step2_hand_computed(self) -> None: + # Step1: state = 0.1 * 1.0 = 0.1. + # Step2: per-sample mean = 0.5; state = 0.9 * 0.1 + 0.1 * 0.5 = 0.14. + # step_count = 2 -> divider = 1 - 0.9^2 = 0.19; metric = 0.14 / 0.19 = 0.7368... + m = StereoMetrics(average_mode="ema", ema_decay=0.9) + m.update( + {"disparities": torch.zeros(1, 1, 1, 2)}, + {"disparities": torch.tensor([[[[1.0, 1.0]]]])}, + ) + m.update( + {"disparities": torch.zeros(1, 1, 1, 2)}, + {"disparities": torch.tensor([[[[0.0, 1.0]]]])}, + ) + assert math.isclose( + m.calculate_metrics()["abs"].item(), 0.14 / 0.19, abs_tol=1e-5 + ) + assert m.step_count.item() == 2.0 + + def test_ema_divisor_becomes_one_after_ema_max_count(self) -> None: + # ema_decay = 0.9 -> ema_max_count = 10. For step_count == 11 the divisor + # switches to 1.0 exactly, so the reported metric equals the raw state. + # For input total = 1.0 each step, the state after k steps is + # 1 - 0.9^step_count. + m = StereoMetrics(average_mode="ema", ema_decay=0.9) + target = torch.tensor([[[[1.0]]]]) + for _ in range(11): + m.update({"disparities": torch.zeros(1, 1, 1, 1)}, {"disparities": target}) + assert m.step_count.item() == 11.0 + expected = 1.0 - 0.9**11 # ~0.68619 + assert math.isclose(m.calculate_metrics()["abs"].item(), expected, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# interpolate_pred_to_target_size +# --------------------------------------------------------------------------- + + +class TestInterpolatePredToTargetSize: + def test_zeros_pred_keeps_per_pixel_abs_with_x_rescaling(self) -> None: + # Pred (1, 1, 1, 2) interpolated to (1, 4); pred stays zero so only + # target contributes. Target x = 3 at one pixel, zero elsewhere -> sum + # = 3 over 4 pixels -> mean = 0.75. + preds = {"disparities": torch.zeros(1, 1, 1, 2)} + target = torch.tensor([[[[3.0, 0.0, 0.0, 0.0]]]]) + m = StereoMetrics(interpolate_pred_to_target_size=True) + m.update(preds, {"disparities": target}) + assert math.isclose(m.calculate_metrics()["abs"].item(), 0.75, abs_tol=1e-5) + + def test_nonzero_pred_is_rescaled_to_target_pixel_units(self) -> None: + # Pred (1, 1, 1, 2) full of 1.0. scale_x = 4/2 = 2 -> after interp and + # rescaling the disparity channel becomes 2.0 everywhere. target = 0 -> + # per-pixel abs = 2 -> mean over 4 px = 2. + preds = {"disparities": torch.full((1, 1, 1, 2), 1.0)} + target = torch.zeros(1, 1, 1, 4) + m = StereoMetrics(interpolate_pred_to_target_size=True) + m.update(preds, {"disparities": target}) + assert math.isclose(m.calculate_metrics()["abs"].item(), 2.0, abs_tol=1e-5) + + def test_interpolation_disabled_does_not_attempt_size_mismatch(self) -> None: + # With interpolate_pred_to_target_size=False, same-shape calls give the + # expected per-pixel abs. + preds = {"disparities": torch.full((1, 1, 2, 2), 2.0)} + target = torch.zeros(1, 1, 2, 2) + m = StereoMetrics(interpolate_pred_to_target_size=False) + m.update(preds, {"disparities": target}) + assert math.isclose(m.calculate_metrics()["abs"].item(), 2.0, abs_tol=1e-5) + + +# --------------------------------------------------------------------------- +# Shape / dtype handling helper paths +# --------------------------------------------------------------------------- + + +class TestShapeHandling: + def test_2d_disparity_is_promoted_to_4d(self) -> None: + # A 2D target tensor (h, w) becomes (1, 1, h, w); abs at the single pixel + # where target = 4 -> abs = 4. + preds = {"disparities": torch.zeros(1, 1, 1, 1)} + target = {"disparities": torch.tensor([[4.0]])} + m = StereoMetrics() + m.update(preds, target) + assert math.isclose(m.calculate_metrics()["abs"].item(), 4.0, abs_tol=1e-6) + + def test_3d_disparity_with_batch_first_axis_promoted_to_4d(self) -> None: + # 3D target shape (1, 2) has shape[0] == sample_count = 1, so + # ``_fix_shape`` adds a channel axis to yield (1, 1, 2). pred matches the + # resulting single-channel shape; per-pixel abs over 2 px = [1, 0] + # -> mean 0.5. + preds = {"disparities": torch.zeros(1, 1, 1, 2)} + target = {"disparities": torch.tensor([[1.0, 0.0]])} + m = StereoMetrics() + m.update(preds, target) + assert math.isclose(m.calculate_metrics()["abs"].item(), 0.5, abs_tol=1e-6) + + def test_double_precision_input_is_cast_to_float32(self) -> None: + preds = {"disparities": torch.zeros(1, 1, 1, 2, dtype=torch.float64)} + target = {"disparities": torch.tensor([[[[1.0, 1.0]]]], dtype=torch.float64)} + m = StereoMetrics() + m.update(preds, target) + metrics = m.calculate_metrics() + # Hand value: per-pixel abs = [1, 1] -> mean = 1.0. + assert math.isclose(metrics["abs"].item(), 1.0, abs_tol=1e-6) + # Internal state is stored as float32 after promotion. + assert metrics["abs"].dtype == torch.float32 + + +# --------------------------------------------------------------------------- +# Known bugs (documented as xfail so a future fix flips them to passing) +# --------------------------------------------------------------------------- + + +class TestBatchAndHypotheses: + def test_batch_size_two_does_not_double_the_metric(self) -> None: + # Two batches, both with abs = 2 everywhere -> per-sample mean 2.0 each; + # epoch_mean should average to 2.0, not 4.0. + preds = {"disparities": torch.zeros(2, 1, 1, 1)} + target = torch.full((2, 1, 1, 1), 2.0) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + assert math.isclose(m.calculate_metrics()["abs"].item(), 2.0, abs_tol=1e-6) + + def test_six_dim_target_uses_min_over_hypotheses(self) -> None: + # A 6D input (b=1, K=2, ch=1, H=1, W=2) is reshaped to 5D by _fix_shape + # and should pick the per-pixel minimum abs across hypotheses. + preds = {"disparities": torch.zeros(1, 1, 1, 1, 2)} + # Hypothesis 0 matches pred exactly; hypothesis 1 has disp = 3 in pixel 0. + h0 = torch.tensor([0.0, 0.0]).reshape(1, 1, 1, 2) + h1 = torch.tensor([3.0, 0.0]).reshape(1, 1, 1, 2) + target = torch.stack([h0, h1], dim=1)[None] # shape (1, 2, 1, 1, 1, 2) + m = StereoMetrics() + m.update(preds, {"disparities": target}) + # The 5D path picks h0 (abs = 0) at both pixels. + assert math.isclose(m.calculate_metrics()["abs"].item(), 0.0, abs_tol=1e-6) diff --git a/tests/common/utils/test_stereo_utils.py b/tests/common/utils/test_stereo_utils.py new file mode 100644 index 0000000..69b2561 --- /dev/null +++ b/tests/common/utils/test_stereo_utils.py @@ -0,0 +1,462 @@ +"""Unit tests for roco_spring_devkit.common.utils.stereo_utils. + +The expected values below are computed by hand from the documented semantics of +each function, then cross-checked against the implementation. They are *not* +simply a capture of the current output. + +Hand-computed reference facts +----------------------------- +disparity_to_rgb (numpy): + Each array is first shifted by its minimum (``disp -= disp.min()``), then + normalised as ``clip(disp / disparity_max, 0, 1)`` and mapped through OpenCV's + ``COLORMAP_PLASMA`` LUT (uint8 BGR). Hand-verified plasma LUT entries: + index 0 -> BGR [135, 8, 13] + index 127 -> BGR [121, 70, 203] + index 255 -> BGR [ 33, 249, 240] + so for input [0.0, 0.5, 1.0] with disparity_max=1.0 the per-pixel indices are + [0, 127, 255] (since 0.5 * 255 = 127.5 truncates to 127 as uint8). + A constant array collapses to all-zeros after the min subtraction -> LUT[0]. + +disparity_to_rgb (torch): same colour semantics but the result is float in + [0, 1] (plasma values divided by 255), in ...CHW layout. + +disparity_read / disparity_write round-trips: + * ``.npy`` is lossless (float32). + * ``.dsp5`` (Spring) is lossless. + * ``.pfm`` should be lossless for finite values. PFM stores rows bottom-up + by convention, so a correct writer flips before writing and a correct + reader flips after reading; the two flips must cancel. + * KITTI ``.png``: written as ``uint16(clip(disp * 256, 0, 2**16-1))``; + ``disparity_read_kitti`` divides by 256 and marks pixels with ``disp <= 0`` + as NaN. Hence 0.0 and negative values become NaN on read, and values + above 255.996 are clipped to 255.99609375. + * Sintel ``.png``: ``disparity = R * 4 + G / 64 + B / 16384``. + +intrinsics_and_baseline_read: + * KITTI: parse ``P_rect_02:`` and ``P_rect_03:`` rows into 3x4 matrices, + intrinsics = ``P_rect_02[:3, :3]``, baseline = ``abs(P3[0,3] - P2[0,3]) / P2[0,0]``. + * Sintel: binary .cam file, header tag float32 ``202021.25``, then 9 float64 + values forming the 3x3 intrinsics. Baseline is the constant 0.1. + * Spring: text file with ``fx fy cx cy``; intrinsics = diag(fx, fy, 1) with + (cx, cy) in the last column. Baseline is the constant 0.065. + * FlyingThings3D: intrinsics are the constants ``[[1050, 0, 479.5], + [0, 1050, 269.5], [0, 0, 1]]`` and baseline is 1.0; the file path is + ignored. + +spring_abs_to_rgb: + ``epe = abs(disp_pred - disp_gt)``; + ``epe = clip(log2(epe * 32), 0, 10)``; + ``epe = (epe * 255 / 10).astype(uint8)``; + apply the matplotlib ``RdYlBu_r`` LUT (256 BGR entries for OpenCV); + set invalid pixels (``~valid_mask``) to black. + Hand-checked LUT entries: + index 0 (epe * 32 <= 1) -> BGR [149, 54, 49] + index 178 (epe = 4) -> BGR [ 97, 174, 253] + index 255 (epe * 32 >= 2**10) -> BGR [ 38, 0, 165] +""" + +import math +from pathlib import Path + +import cv2 as cv +import numpy as np +import png +import pytest +import torch + +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 = { + 0: [135, 8, 13], + 127: [121, 70, 203], + 255: [33, 249, 240], +} + + +# --------------------------------------------------------------------------- +# disparity_to_rgb +# --------------------------------------------------------------------------- + + +class TestDisparityToRgbNumpy: + def test_gradient_with_explicit_max_uses_plasma_lut(self) -> None: + # Input [0.0, 0.5, 1.0], disparity_max=1.0 -> indices [0, 127, 255]. + disp = np.array([[[0.0], [0.5], [1.0]]], dtype=np.float32) # (H=1, W=3, C=1) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert rgb.shape == (1, 3, 3) + assert rgb.dtype == np.uint8 + assert rgb[0, 0].tolist() == PLASMA_LUT[0] + assert rgb[0, 1].tolist() == PLASMA_LUT[127] + assert rgb[0, 2].tolist() == PLASMA_LUT[255] + + def test_2d_input_treated_as_hwc_with_single_channel(self) -> None: + # 2D input should be wrapped as a single-channel image. + disp = np.array([[0.0, 0.5, 1.0]], dtype=np.float32) # (H=1, W=3) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert rgb.shape == (1, 3, 3) + assert rgb[0, 0].tolist() == PLASMA_LUT[0] + assert rgb[0, 1].tolist() == PLASMA_LUT[127] + assert rgb[0, 2].tolist() == PLASMA_LUT[255] + + def test_default_max_is_input_maximum(self) -> None: + # With disparity_max=None the normalisation uses max(disp - min). + # Input [-1.0, 0.0, 1.0] -> after min subtraction [0, 1, 2], max=2, + # norm = [0, 0.5, 1.0] -> same LUT indices as the gradient case. + disp = np.array([[-1.0, 0.0, 1.0]], dtype=np.float32) + rgb = stereo_utils.disparity_to_rgb(disp) + assert rgb[0, 0].tolist() == PLASMA_LUT[0] + assert rgb[0, 1].tolist() == PLASMA_LUT[127] + assert rgb[0, 2].tolist() == PLASMA_LUT[255] + + def test_explicit_max_clips_values_above_one(self) -> None: + # disparity_max=1.0 but values reach 2.0 -> clipped to 1.0 -> LUT[255]. + disp = np.array([[0.0, 1.0, 2.0]], dtype=np.float32) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert rgb[0, 0].tolist() == PLASMA_LUT[0] + assert rgb[0, 1].tolist() == PLASMA_LUT[255] + assert rgb[0, 2].tolist() == PLASMA_LUT[255] + + def test_negative_values_with_explicit_max_are_clipped_to_zero(self) -> None: + # After min subtraction everything is >= 0, but with an explicit max + # smaller than the shifted range, clipping applies. + # Input [-1, 0, 1], min=-1, -=-1 -> [0, 1, 2], /1.0 -> [0, 1, 2], + # clip [0,1] -> [0, 1, 1] -> LUT[0], LUT[255], LUT[255]. + disp = np.array([[-1.0, 0.0, 1.0]], dtype=np.float32) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert rgb[0, 0].tolist() == PLASMA_LUT[0] + assert rgb[0, 1].tolist() == PLASMA_LUT[255] + assert rgb[0, 2].tolist() == PLASMA_LUT[255] + + def test_constant_input_maps_to_lut_zero(self) -> None: + # min subtraction collapses a constant to zero -> LUT[0]. + disp = np.array([[0.5, 0.5], [0.5, 0.5]], dtype=np.float32) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert np.array_equal(rgb, np.full((2, 2, 3), PLASMA_LUT[0], dtype=np.uint8)) + + +class TestDisparityToRgbTorch: + def test_3d_input_returns_float_chw(self) -> None: + # Input (1, H, W) treated as (C=1, H, W); output (3, H, W) in [0, 1]. + disp = torch.tensor([[[0.0, 0.5, 1.0]]]) # (1, 1, 3) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert rgb.shape == (3, 1, 3) + assert rgb.dtype == torch.float32 + # plasma LUT / 255 + assert torch.allclose(rgb[:, 0, 0], torch.tensor(PLASMA_LUT[0]) / 255.0) + assert torch.allclose(rgb[:, 0, 1], torch.tensor(PLASMA_LUT[127]) / 255.0) + assert torch.allclose(rgb[:, 0, 2], torch.tensor(PLASMA_LUT[255]) / 255.0) + + def test_4d_batched_input_returns_n3hw(self) -> None: + disp = torch.tensor([[[[0.0, 1.0]]]]) # (B=1, C=1, H=1, W=2) + rgb = stereo_utils.disparity_to_rgb(disp, disparity_max=1.0) + assert rgb.shape == (1, 3, 1, 2) + assert torch.allclose(rgb[0, :, 0, 0], torch.tensor(PLASMA_LUT[0]) / 255.0) + assert torch.allclose(rgb[0, :, 0, 1], torch.tensor(PLASMA_LUT[255]) / 255.0) + + +# --------------------------------------------------------------------------- +# disparity_read / disparity_write round-trips +# --------------------------------------------------------------------------- + + +class TestDisparityReadWrite: + def test_npy_round_trip_is_lossless(self, tmp_path: Path) -> None: + disp = np.array([[0.5, -1.25], [3.0, 0.0]], dtype=np.float32) + p = tmp_path / "a.npy" + stereo_utils.disparity_write(p, disp) + out = stereo_utils.disparity_read(p) + assert np.array_equal(out, disp) + + def test_dsp5_spring_round_trip_is_lossless(self, tmp_path: Path) -> None: + disp = np.array([[0.5, -1.25], [3.0, 0.0]], dtype=np.float32) + p = tmp_path / "a.dsp5" + stereo_utils.disparity_write(p, disp, format="dsp5") + out = stereo_utils.disparity_read(p, format="spring") + assert np.array_equal(out, disp) + + def test_pfm_round_trip_preserves_orientation(self, tmp_path: Path) -> None: + # PFM stores rows bottom-up; a correct writer and reader must both flip + # so that the round-trip preserves the array orientation. + disp = np.array([[0.5, -1.25], [3.0, 0.0]], dtype=np.float32) + p = tmp_path / "a.pfm" + stereo_utils.disparity_write(p, disp) + out = stereo_utils.disparity_read(p) + assert np.allclose(out, disp, atol=1e-6) + + def test_pfm_round_trip_with_nan_preserves_orientation( + self, tmp_path: Path + ) -> None: + disp = np.array([[0.5, -1.25], [float("nan"), 0.0]], dtype=np.float32) + p = tmp_path / "a.pfm" + stereo_utils.disparity_write(p, disp) + out = stereo_utils.disparity_read(p) + assert np.allclose(out[0, :], disp[0, :], atol=1e-6) + assert math.isnan(out[1, 0]) # NaN stays at (1, 0) + assert np.allclose(out[1, 1], 0.0, atol=1e-6) + + +class TestDisparityReadWriteKitti: + def test_positive_values_round_trip_losslessly(self, tmp_path: Path) -> None: + # 0.5 -> 128 -> 0.5 ; 3.0 -> 768 -> 3.0 + disp = np.array([[0.5, 3.0]], dtype=np.float32) + p = tmp_path / "a.png" + stereo_utils.disparity_write(p, disp, format="kitti") + out = stereo_utils.disparity_read(p, format="kitti") + assert np.allclose(out, disp, atol=1e-6) + + def test_zero_and_negative_become_nan_on_read(self, tmp_path: Path) -> None: + # writer clips negatives to 0; reader marks disp <= 0 as NaN. + disp = np.array([[0.0, -1.25]], dtype=np.float32) + p = tmp_path / "a.png" + stereo_utils.disparity_write(p, disp, format="kitti") + out = stereo_utils.disparity_read(p, format="kitti") + assert math.isnan(out[0, 0]) + assert math.isnan(out[0, 1]) + + def test_large_values_are_clipped_to_uint16_max(self, tmp_path: Path) -> None: + # 300 * 256 = 76800 > 65535 -> clipped -> 65535 / 256 = 255.99609375. + disp = np.array([[300.0]], dtype=np.float32) + p = tmp_path / "a.png" + stereo_utils.disparity_write(p, disp, format="kitti") + out = stereo_utils.disparity_read(p, format="kitti") + assert np.allclose(out, 255.99609375, atol=1e-6) + + def test_nan_in_input_becomes_nan_on_read(self, tmp_path: Path) -> None: + disp = np.array([[float("nan"), 1.0]], dtype=np.float32) + p = tmp_path / "a.png" + stereo_utils.disparity_write(p, disp, format="kitti") + out = stereo_utils.disparity_read(p, format="kitti") + assert math.isnan(out[0, 0]) + assert np.allclose(out[0, 1], 1.0, atol=1e-6) + + def test_kitti_writer_rejects_3d_non_single_channel(self, tmp_path: Path) -> None: + p = tmp_path / "a.png" + # (H, W, 2) is not reducible to 2D and should raise IOError. + with pytest.raises(IOError): + stereo_utils.disparity_write_kitti(p, np.zeros((2, 2, 2))) + + def test_kitti_writer_accepts_single_channel_3d(self, tmp_path: Path) -> None: + disp = np.array([[[0.5], [3.0]]], dtype=np.float32) # (1, 2, 1) + p = tmp_path / "a.png" + stereo_utils.disparity_write_kitti(p, disp) + out = stereo_utils.disparity_read_kitti(p) + assert np.allclose(out, np.array([[0.5, 3.0]]), atol=1e-6) + + +class TestDisparityReadSintel: + def test_reconstruction_formula(self, tmp_path: Path) -> None: + # disparity = R * 4 + G / 64 + B / 16384 + # pixel 0: R=10, G=0, B=0 -> 40.0 + # pixel 1: R=10, G=64, B=0 -> 40 + 1.0 = 41.0 + # pixel 2: R=10, G=0, B=1 -> 40 + 1/16384 = 40.00006103515625 + img = np.zeros((1, 3, 3), dtype=np.uint8) + img[0, 0] = (10, 0, 0) + img[0, 1] = (10, 64, 0) + img[0, 2] = (10, 0, 1) + p = tmp_path / "a.png" + from PIL import Image + + Image.fromarray(img).save(p) + out = stereo_utils.disparity_read_sintel(p) + assert np.allclose(out[0, 0], 40.0, atol=1e-6) + assert np.allclose(out[0, 1], 41.0, atol=1e-6) + assert np.allclose(out[0, 2], 40.0 + 1.0 / 16384.0, atol=1e-7) + + +class TestDisparityReadUnsupported: + def test_unsupported_format_raises(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Unsupported disparity format"): + stereo_utils.disparity_read(tmp_path / "x.foo", format="foo") + + def test_unknown_extension_raises(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Unsupported disparity format"): + stereo_utils.disparity_read(tmp_path / "x.foo") + + +# --------------------------------------------------------------------------- +# intrinsics_and_baseline_read +# --------------------------------------------------------------------------- + + +class TestIntrinsicsAndBaselineRead: + def test_kitti_parses_p_rect_rows_and_baseline_formula( + self, tmp_path: Path + ) -> None: + # Two projection matrices; baseline = |P3[0,3] - P2[0,3]| / P2[0,0]. + p = tmp_path / "calib.txt" + p.write_text( + "P_rect_02: 721.5377 0 609.5593 44.85728 0 721.5377 172.8540 0.2163791 0 0 1 0.002745884\n" + "P_rect_03: 721.5377 0 609.5593 -339.5242 0 721.5377 172.8540 2.199912 0 0 1 0.002729905\n" + ) + intr, baseline = stereo_utils.intrinsics_and_baseline_read(p, format="kitti") + assert intr.shape == (3, 3) + assert np.allclose( + intr, + np.array( + [ + [721.5377, 0, 609.5593], + [0, 721.5377, 172.8540], + [0, 0, 1], + ], + dtype=np.float32, + ), + atol=1e-3, + ) + assert math.isclose( + baseline, abs(-339.5242 - 44.85728) / 721.5377, abs_tol=1e-5 + ) + + def test_sintel_reads_binary_cam_file(self, tmp_path: Path) -> None: + TAG = np.float32(202021.25) + M = np.array( + [[1000.0, 0, 500.0], [0, 1000.0, 250.0], [0, 0, 1.0]], + dtype=np.float64, + ) + p = tmp_path / "a.cam" + with open(p, "wb") as f: + f.write(TAG.tobytes()) + f.write(M.tobytes()) + intr, baseline = stereo_utils.intrinsics_and_baseline_read(p, format="sintel") + assert np.array_equal(intr, M) + assert baseline == pytest.approx(0.1) + + def test_sintel_rejects_bad_tag(self, tmp_path: Path) -> None: + p = tmp_path / "a.cam" + with open(p, "wb") as f: + f.write(np.float32(12345.0).tobytes()) + f.write(np.zeros(9, dtype=np.float64).tobytes()) + with pytest.raises(ValueError, match="Invalid tag"): + stereo_utils.intrinsics_and_baseline_read(p, format="sintel") + + def test_spring_reads_fx_fy_cx_cy(self, tmp_path: Path) -> None: + p = tmp_path / "a.cam" + p.write_text("1050.0 1050.0 479.5 269.5") + intr, baseline = stereo_utils.intrinsics_and_baseline_read(p, format="spring") + assert np.array_equal( + intr, + ( + np.array( + [[1050.0, 0, 479.5], [0, 1050.0, 269.5], [0, 0, 1.0]], + dtype=np.float32, + ) + if intr.dtype == np.float32 + else np.array([[1050.0, 0, 479.5], [0, 1050.0, 269.5], [0, 0, 1.0]]) + ), + ) + assert baseline == pytest.approx(0.065) + + def test_things_uses_constants_and_ignores_file(self, tmp_path: Path) -> None: + intr, baseline = stereo_utils.intrinsics_and_baseline_read( + tmp_path / "nonexistent", format="things" + ) + assert np.array_equal( + intr, + np.array([[1050.0, 0, 479.5], [0, 1050.0, 269.5], [0, 0, 1.0]]), + ) + assert baseline == pytest.approx(1.0) + + def test_unsupported_format_raises(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Unsupported intrinsics format"): + stereo_utils.intrinsics_and_baseline_read(tmp_path / "x", format="foo") + + def test_no_format_raises(self, tmp_path: Path) -> None: + with pytest.raises(ValueError, match="Unsupported intrinsics format"): + stereo_utils.intrinsics_and_baseline_read(tmp_path / "x") + + +# --------------------------------------------------------------------------- +# spring_abs_to_rgb +# --------------------------------------------------------------------------- + + +def _rdylbu_r_lut() -> np.ndarray: + """Independently build the BGR LUT for RdYlBu_r from matplotlib.""" + import matplotlib.pyplot as plt + + cmap = plt.get_cmap("RdYlBu_r") + rgb = (cmap(np.arange(256))[:, :3] * 255).astype(np.uint8) + return rgb[:, ::-1].reshape(256, 1, 3) + + +def _expected_abs_rgb( + disp_pred: np.ndarray, + disp_gt: np.ndarray, + valid_mask: np.ndarray, + lut_bgr: np.ndarray, +) -> np.ndarray: + """Independent reference implementation of spring_abs_to_rgb.""" + epe = np.abs(disp_pred - disp_gt) + epe = np.clip(np.log2(epe * 32), 0, 10) + epe = (epe * 255 / 10).astype(np.uint8) + rgb = cv.applyColorMap(epe, lut_bgr) + rgb[~valid_mask] = 0 # broadcasts over the trailing channel dim + return rgb + + +class TestSpringAbsToRgb: + def test_perfect_prediction_maps_to_lut_index_zero(self) -> None: + pred = np.zeros((2, 2), dtype=np.float32) + gt = np.zeros((2, 2), dtype=np.float32) + valid = np.ones((2, 2), dtype=bool) + rgb = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert rgb.shape == (2, 2, 3) + # LUT[0] in BGR is [149, 54, 49]. + assert np.array_equal(rgb, np.full((2, 2, 3), [149, 54, 49], dtype=np.uint8)) + + def test_mid_abs_maps_to_hand_computed_index_178(self) -> None: + # abs = 4 -> log2(128) = 7 -> 7 * 255 / 10 = 178.5 -> uint8 178. + # Hand-checked LUT[178] BGR = [97, 174, 253]. + pred = np.zeros((1, 1), dtype=np.float32) + gt = np.array([[4.0]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [97, 174, 253] + + def test_large_abs_clips_to_index_255(self) -> None: + # abs = 32 -> log2(1024) = 10 -> clipped to 10 -> index 255. + # Hand-checked LUT[255] BGR = [38, 0, 165]. + pred = np.zeros((1, 1), dtype=np.float32) + gt = np.array([[32.0]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [38, 0, 165] + + def test_small_abs_below_threshold_clips_to_zero(self) -> None: + # abs = 0.01 -> log2(0.32) < 0 -> clipped to 0 -> LUT[0]. + pred = np.array([[0.0]], dtype=np.float32) + gt = np.array([[0.01]], dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [149, 54, 49] + + def test_invalid_pixels_are_set_to_black(self) -> None: + # Valid pixel has abs=4 -> LUT[178]; invalid pixel -> [0, 0, 0]. + pred = np.array([[0.0, 0.0]], dtype=np.float32) + gt = np.array([[4.0, 4.0]], dtype=np.float32) + valid = np.ones((1, 2), dtype=bool) + valid[0, 1] = False + rgb = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [97, 174, 253] + assert rgb[0, 1].tolist() == [0, 0, 0] + + def test_matches_independent_reference_implementation(self) -> None: + rng = np.random.RandomState(0) + pred = rng.randn(3, 4).astype(np.float32) + gt = rng.randn(3, 4).astype(np.float32) + valid = np.ones((3, 4), dtype=bool) + valid[0, 0] = False + valid[2, 3] = False + expected = _expected_abs_rgb(pred, gt, valid, _rdylbu_r_lut()) + out = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert np.array_equal(out, expected) + + def test_negative_abs_is_absolute_value(self) -> None: + # abs(-4) = 4 -> same LUT entry as abs = 4. + pred = np.array([[-4.0]], dtype=np.float32) + gt = np.zeros((1, 1), dtype=np.float32) + valid = np.ones((1, 1), dtype=bool) + rgb = stereo_utils.spring_abs_to_rgb(pred, gt, valid) + assert rgb[0, 0].tolist() == [97, 174, 253] diff --git a/tests/common/utils/test_utils.py b/tests/common/utils/test_utils.py new file mode 100644 index 0000000..c5b6db7 --- /dev/null +++ b/tests/common/utils/test_utils.py @@ -0,0 +1,898 @@ +"""Unit tests for roco_spring_devkit.common.utils.utils. + +The expected values below are computed by hand from the documented semantics +of each function, then cross-checked against the implementation. They are +*not* simply a capture of the current output. + +Hand-computed reference facts +----------------------------- +_forward_warp_single_torch: + Pure forward warp using round-to-nearest-integer destinations: + * source position -> dest = (x0 + dx, y0 + dy); + * keep destinations with 0 <= x1 <= wd-1 and 0 <= y1 <= ht-1; + * destination round-off distance = (x1 - round(x1))**2 + + (y1 - round(y1))**2; + * tie-break = source flat index * 1e-6 (smaller index wins); + * destination pixel is set to the source value that minimises + (distance + tie_break); unmapped destination pixels remain 0. + + Examples (H=1, W=2): + * flow[0] = [0.5, -0.6], tensor = [5, 7]: + src 0 -> dest 0.5 -> round 0, dist 0.25, val 5; + src 1 -> dest 0.4 -> round 0, dist 0.16, val 7; + min key = 0.16 (src 1) -> output[0, 0] = 7, output[0, 1] = 0. + * identity flow on a 2x2 grid -> each source maps to itself, so the + output equals the input. + +_forward_warp_single_scipy / _forward_warp_single_ckdtree: + Original implementations used a strict-bounds validity mask and a + backward nearest-neighbour lookup (scipy.griddata / scipy.cKDTree), + which produced different (buggy) results than the torch version. + They have been reimplemented on top of :func:`_forward_warp_single_numpy` so + that all three backends now produce identical, correct forward-warp + results: inclusive bounds, source-with-min-(distance + tie_break) + wins per destination, and unmapped destinations stay at zero. + +InputPadder: + Wraps raft.InputPadder. ``_pad`` is the F.pad order + [left, right, top, bottom]. + For dims=(2, 2), stride=4: pad_ht = pad_wd = 2; with two_side_pad=True + the padding splits 1 / 1 on each side -> [1, 1, 1, 1]. + With two_side_pad=False the order becomes + [pad_wd//2, pad_wd - pad_wd//2, 0, pad_ht]; only the *height* padding + moves to the bottom (the width padding is still split). For dims=(4, 2) + stride=8 -> pad_wd=6, pad_ht=4 -> [3, 3, 0, 4]. + The ``tgt_size`` attribute is the next multiple of stride (or ``size`` + if passed). + +InputScaler: + Target size is ``ceil(orig/stride)*stride`` when stride is given, the + passed ``size`` when ``size`` is given, or ``int(orig*scale_factor)``. + ``fill``/``unfill`` are bilinear (default, align_corners=False) and + preserve the leading dimensions. For ``is_flow=True`` the x-channel + (channel 0) is multiplied by new_width/old_width and the y-channel + (channel 1) by new_height/old_height (so the flow values track the + resolution change). + +make_divisible(v, div) = max(div, v - v % div): + make_divisible(10, 3) = max(3, 9) = 9 + make_divisible(8, 4) = max(4, 8) = 8 + make_divisible(2, 5) = max(5, 0) = 5 + make_divisible(0, 5) = max(5, 0) = 5 + make_divisible(12, 4) = max(4, 12) = 12 + +count_parameters: + Sums ``p.numel()`` over parameters with ``requires_grad == True``. + +release_gpu: + Mutates the dict in place, replacing torch.Tensor values by their + ``.detach().cpu()`` copy. Non-tensor entries are kept intact and the + same dict instance is returned. + +tensor_dict_to_numpy: + Iterates dict items; for torch.Tensor entries it does + detach -> cpu -> (optional ``padder.unfill``) -> squeezes leading dims + until 3D -> permute(1, 2, 0) -> numpy. Non-tensor entries are kept + verbatim. 1D / 2D tensors are converted to numpy without the CHW -> HWC + permute (since they have no channel dimension to move). The function + uses ``padder.unfill`` (not ``padder.unpad``) so that an already-unpadded + tensor is returned untouched instead of being sliced into an empty + array. + +are_shapes_compatible: + Two shapes are compatible iff they have the same length and each pair + of sizes is equal or at least one of them is 1. + +bgr_val_as_tensor: + * scalar -> tensor of shape [1, 3, 1, 1, ...] (three copies in the + BGR-dimension). + * tuple/list of length 3 -> tensor of length 3 reshaped to the BGR + position. + * np.ndarray of length 3 -> same as tuple (after dtype/device cast). + * already-compatible tensor -> returned untouched. + The BGR dimension is set by ``bgr_tensor_shape_position`` (default -3). + +forward_interpolate_batch: + Applies raft.forward_interpolate per batch element. forward_interpolate + uses the strict-bounds mask (x1 > 0 & x1 < wd & y1 > 0 & y1 < ht) and + griddata(method='nearest'); with zero flow on an HxW grid *all* sources + lie on the border, so no source is valid and the output is all zeros. + +forward_warp (dispatcher): + Reads ``ROCO_FORWARD_WARP_MODE`` (default 'exact' -> cKDTree, + 'fast' -> torch, 'legacy' -> scipy). Applies per-batch for 4D flow + and returns the stacked result; otherwise calls the single-tensor + helper directly. Validates batch dim, channel dim == 2, and spatial + dim match for 4D flow. + +get_matplotlib_lut_colormap: + Returns a (256, 1, 3) uint8 BGR LUT built from the named matplotlib + colormap. Independently computed: rgb = matplotlib cmap(arange(256)) + [:, :3] * 255, then reversed to BGR and reshaped to (256, 1, 3). +""" + +import os +from argparse import ArgumentParser +from pathlib import Path + +import numpy as np +import pytest +import torch +import torch.nn as nn + +from roco_spring_devkit.common.utils import utils as utils_mod +from roco_spring_devkit.common.utils.utils import ( + InputPadder, + InputScaler, + _forward_warp_single_ckdtree, + _forward_warp_single_scipy, + _forward_warp_single_torch, + add_datasets_to_parser, + are_shapes_compatible, + bgr_val_as_tensor, + count_parameters, + forward_interpolate_batch, + forward_warp, + get_matplotlib_lut_colormap, + make_divisible, + release_gpu, + tensor_dict_to_numpy, +) + + +@pytest.fixture +def restore_env(monkeypatch): + """Snapshot and restore ROCO_FORWARD_WARP_MODE around each test.""" + saved = os.environ.get("ROCO_FORWARD_WARP_MODE") + yield + if saved is None: + os.environ.pop("ROCO_FORWARD_WARP_MODE", None) + else: + os.environ["ROCO_FORWARD_WARP_MODE"] = saved + + +# --------------------------------------------------------------------------- +# _forward_warp_single_torch +# --------------------------------------------------------------------------- + + +class TestForwardWarpSingleTorch: + def test_identity_flow_on_two_by_two_round_trips(self) -> None: + # Inclusive bounds (>= 0, <= wd-1) keep every source in place, so + # identity flow should leave the tensor unchanged. + tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + flow = torch.zeros(2, 2, 2) + out = _forward_warp_single_torch(flow, tensor) + assert torch.equal(out, tensor) + + def test_single_horizontal_shift(self) -> None: + # Source (row 0, col 0) of value 5 warps to col 1: output[0, 0] = 0 + # (unmapped), output[0, 1] = 5. + tensor = torch.tensor([[5.0, 0.0]]) + flow = torch.tensor([[[1.0, 0.0]], [[0.0, 0.0]]]) # (2, H, W) + out = _forward_warp_single_torch(flow, tensor) + assert out.shape == tensor.shape + assert out.tolist() == [[0.0, 5.0]] + + def test_subpixel_distance_breaks_ties_to_smaller_distance(self) -> None: + # Both sources map to column 0 (round(0.5) == 0 and round(0.4) == 0). + # Distances are 0.25 and 0.16 respectively: the second source wins. + tensor = torch.tensor([[5.0, 7.0]]) + flow = torch.tensor([[[0.5, -0.6]], [[0.0, 0.0]]]) + out = _forward_warp_single_torch(flow, tensor) + assert out.tolist() == [[7.0, 0.0]] + + def test_tie_break_prefers_smaller_flat_index(self) -> None: + # Two sources (values 10 and 20) warp to dest (row 0, col 1) at + # distance 0.0 each. tie_break = flat_index * 1e-6 makes the + # smaller index win, so dest gets 10 (source 0). Other + # destinations are unmapped and remain zero. + tensor = torch.tensor([[10.0, 0.0, 20.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]) + flow = torch.tensor( + [ + [[1.0, 0.0, -1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], + [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], + ] + ) + out = _forward_warp_single_torch(flow, tensor) + assert out.tolist() == [ + [0.0, 10.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + ] + + def test_out_of_bounds_source_is_dropped(self) -> None: + # Source at col 1 with dx = 5 maps to x = 6 (> wd-1) -> dropped. + # Source at col 0 (value 5) stays in place -> output[0, 0] = 5. + tensor = torch.tensor([[5.0, 7.0]]) + flow = torch.tensor([[[0.0, 5.0]], [[0.0, 0.0]]]) + out = _forward_warp_single_torch(flow, tensor) + assert out.tolist() == [[5.0, 0.0]] + + def test_multichannel_tensor_each_channel_warped_independently(self) -> None: + # 3-channel tensor, the first source row warps to the right. + tensor = torch.tensor( + [ + [[1.0, 0.0], [1.0, 0.0]], + [[2.0, 0.0], [2.0, 0.0]], + [[3.0, 0.0], [3.0, 0.0]], + ] + ) # (C=3, H=2, W=2) + flow = torch.tensor( + [[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]] + ) # (2, H, W) + out = _forward_warp_single_torch(flow, tensor) + # First row: col 0 -> col 1; second row stays put. + expected = torch.tensor( + [ + [[0.0, 1.0], [1.0, 0.0]], + [[0.0, 2.0], [2.0, 0.0]], + [[0.0, 3.0], [3.0, 0.0]], + ] + ) + assert torch.equal(out, expected) + + def test_2d_input_returns_2d_output(self) -> None: + tensor = torch.tensor([[1.0, 2.0]]) + flow = torch.zeros(2, 1, 2) + out = _forward_warp_single_torch(flow, tensor) + assert out.shape == tensor.shape + + def test_invalid_flow_shape_raises(self) -> None: + tensor = torch.zeros(2, 2) + bad_flow = torch.zeros(3, 2, 2) # channel dim != 2 + with pytest.raises(ValueError, match="flow must have shape"): + _forward_warp_single_torch(bad_flow, tensor) + + def test_invalid_tensor_ndim_raises(self) -> None: + flow = torch.zeros(2, 2, 2) + bad_tensor = torch.zeros(1, 1, 2, 2) # 4D not supported here + with pytest.raises(ValueError, match="tensor must have shape"): + _forward_warp_single_torch(flow, bad_tensor) + + def test_spatial_dim_mismatch_raises(self) -> None: + flow = torch.zeros(2, 2, 2) + tensor = torch.zeros(3, 3) + with pytest.raises(ValueError, match="spatial dimensions"): + _forward_warp_single_torch(flow, tensor) + + +# --------------------------------------------------------------------------- +# _forward_warp_single_scipy / _forward_warp_single_ckdtree +# --------------------------------------------------------------------------- + + +class TestForwardWarpSingleScipy: + def test_identity_flow_on_two_by_two_preserves_input(self) -> None: + tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + flow = torch.zeros(2, 2, 2) + out = _forward_warp_single_scipy(flow, tensor) + assert torch.equal(out, tensor) + + def test_single_horizontal_shift(self) -> None: + # Source (0, 0) [value 5] should warp to dest (0, 1); other pixels + # unmapped -> 0. + tensor = torch.tensor([[5.0, 0.0]]) + flow = torch.tensor([[[1.0, 0.0]], [[0.0, 0.0]]]) + out = _forward_warp_single_scipy(flow, tensor) + assert out.tolist() == [[0.0, 5.0]] + + +class TestForwardWarpSingleCkdtree: + def test_identity_flow_on_two_by_two_preserves_input(self) -> None: + tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + flow = torch.zeros(2, 2, 2) + out = _forward_warp_single_ckdtree(flow, tensor) + assert torch.equal(out, tensor) + + def test_single_horizontal_shift(self) -> None: + tensor = torch.tensor([[5.0, 0.0]]) + flow = torch.tensor([[[1.0, 0.0]], [[0.0, 0.0]]]) + out = _forward_warp_single_ckdtree(flow, tensor) + assert out.tolist() == [[0.0, 5.0]] + + +class TestForwardWarpSingleBackendsAgree: + """The three backends (torch, scipy, ckdtree) must produce identical + results now that the scipy/ckdtree versions use a proper forward-warp + scatter algorithm. + """ + + @staticmethod + def _make_inputs(): + rng = np.random.RandomState(0) + ht, wd = 5, 7 + flow_np = rng.randn(2, ht, wd).astype(np.float32) * 2.0 + tensor_np = rng.randn(3, ht, wd).astype(np.float32) + return torch.from_numpy(flow_np), torch.from_numpy(tensor_np) + + def test_scipy_matches_torch(self) -> None: + flow, tensor = self._make_inputs() + a = _forward_warp_single_torch(flow, tensor) + b = _forward_warp_single_scipy(flow, tensor) + assert torch.allclose(a, b, atol=1e-6) + + def test_ckdtree_matches_torch(self) -> None: + flow, tensor = self._make_inputs() + a = _forward_warp_single_torch(flow, tensor) + b = _forward_warp_single_ckdtree(flow, tensor) + assert torch.allclose(a, b, atol=1e-6) + + def test_scipy_matches_ckdtree(self) -> None: + flow, tensor = self._make_inputs() + a = _forward_warp_single_scipy(flow, tensor) + b = _forward_warp_single_ckdtree(flow, tensor) + assert torch.allclose(a, b, atol=1e-6) + + +# --------------------------------------------------------------------------- +# forward_warp dispatcher +# --------------------------------------------------------------------------- + + +class TestForwardWarpDispatcher: + def test_default_mode_is_ckdtree(self, restore_env, monkeypatch) -> None: + # No env -> default 'exact' -> ckdtree. + monkeypatch.delenv("ROCO_FORWARD_WARP_MODE", raising=False) + flow = torch.tensor([[[[1.0, 0.0]], [[0.0, 0.0]]]]) # (1, 2, 1, 2) + tensor = torch.tensor([[[5.0, 0.0]]]) + out = forward_warp(flow, tensor) + expected = _forward_warp_single_ckdtree(flow[0], tensor[0]) + assert torch.equal(out, expected.unsqueeze(0)) + assert out.shape == (1, 1, 2) + + def test_fast_mode_dispatches_to_torch(self, restore_env, monkeypatch) -> None: + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "fast") + flow = torch.tensor([[[[1.0, 0.0]], [[0.0, 0.0]]]]) # (1, 2, 1, 2) + tensor = torch.tensor([[[5.0, 0.0]]]) + out = forward_warp(flow, tensor) + expected = _forward_warp_single_torch(flow[0], tensor[0]) + assert torch.equal(out, expected.unsqueeze(0)) + + def test_legacy_mode_dispatches_to_scipy(self, restore_env, monkeypatch) -> None: + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "legacy") + flow = torch.tensor([[[[1.0, 0.0]], [[0.0, 0.0]]]]) + tensor = torch.tensor([[[5.0, 0.0]]]) + out = forward_warp(flow, tensor) + expected = _forward_warp_single_scipy(flow[0], tensor[0]) + assert torch.equal(out, expected.unsqueeze(0)) + + def test_4d_flow_3d_tensor_stacks_per_batch(self, restore_env, monkeypatch) -> None: + # 4D flow + 3D tensor should iterate over the batch dim. + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "fast") + flow = torch.zeros(2, 2, 2, 2) + flow[0, :, 0, 0] = torch.tensor([1.0, 0.0]) # batch 0 shift right + flow[1, :, 0, 0] = torch.tensor([0.0, 1.0]) # batch 1 shift down + tensor = torch.tensor([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]]) + out = forward_warp(flow, tensor) + assert out.shape == (2, 2, 2) + # Batch 0: source (0, 0) value 1 warps to dest (0, 1), ties with + # source (0, 1) value 2 (which stays). Smaller flat index wins so + # dest (0, 1) gets 1; the rest stay put. + # Batch 1: source (0, 0) value 5 warps to dest (1, 0), ties with + # source (1, 0) value 7 (which stays). Smaller flat index wins so + # dest (1, 0) gets 5; the rest stay put. + assert out[0].tolist() == [[0.0, 1.0], [3.0, 4.0]] + assert out[1].tolist() == [[0.0, 6.0], [5.0, 8.0]] + + def test_4d_flow_batch_mismatch_raises(self, restore_env, monkeypatch) -> None: + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "fast") + flow = torch.zeros(3, 2, 2, 2) + tensor = torch.zeros(2, 2, 2) + with pytest.raises(ValueError, match="batch dimensions must match"): + forward_warp(flow, tensor) + + def test_4d_flow_wrong_channel_count_raises(self, restore_env, monkeypatch) -> None: + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "fast") + flow = torch.zeros(2, 3, 2, 2) # channel dim != 2 + tensor = torch.zeros(2, 2, 2) + with pytest.raises(ValueError, match="batched flow must have shape"): + forward_warp(flow, tensor) + + def test_4d_flow_spatial_mismatch_raises(self, restore_env, monkeypatch) -> None: + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "fast") + flow = torch.zeros(2, 2, 3, 3) + tensor = torch.zeros(2, 2, 2) + with pytest.raises(ValueError, match="spatial dimensions must match"): + forward_warp(flow, tensor) + + def test_4d_flow_wrong_tensor_ndim_raises(self, restore_env, monkeypatch) -> None: + monkeypatch.setenv("ROCO_FORWARD_WARP_MODE", "fast") + flow = torch.zeros(2, 2, 2, 2) + # 5D tensor is not supported for the batched path. + tensor = torch.zeros(2, 1, 1, 2, 2) + with pytest.raises(ValueError, match="batched tensor must have shape"): + forward_warp(flow, tensor) + + +# --------------------------------------------------------------------------- +# InputPadder +# --------------------------------------------------------------------------- + + +class TestInputPadder: + def test_two_side_pad_pads_each_side_equally_for_stride_divisor( + self, + ) -> None: + # dims=(2, 2), stride=4 -> pad_ht = pad_wd = 2; two_side -> 1/1 each. + p = InputPadder((2, 2), stride=4) + assert p._pad == [1, 1, 1, 1] + assert p.tgt_size == (4, 4) + + def test_constant_pad_value_fills_ring(self) -> None: + p = InputPadder((2, 2), stride=4, pad_mode="constant", pad_value=9.0) + x = torch.tensor([[[[1.0, 2.0], [3.0, 4.0]]]]) + out = p.fill(x) + expected = torch.tensor( + [ + [ + [ + [9.0, 9.0, 9.0, 9.0], + [9.0, 1.0, 2.0, 9.0], + [9.0, 3.0, 4.0, 9.0], + [9.0, 9.0, 9.0, 9.0], + ] + ] + ] + ) + assert torch.equal(out, expected) + + def test_no_two_side_pad_pushes_height_padding_to_bottom(self) -> None: + # dims=(4, 2), stride=8 -> pad_wd=6 (split 3/3), pad_ht=4 -> top=0, + # bottom=4; _pad order is [left, right, top, bottom] = [3, 3, 0, 4]. + p = InputPadder( + (4, 2), + stride=8, + two_side_pad=False, + pad_mode="constant", + pad_value=0.0, + ) + assert p._pad == [3, 3, 0, 4] + assert p.tgt_size == (8, 8) + x = torch.arange(4 * 2, dtype=torch.float32).reshape(1, 1, 4, 2) + out = p.fill(x) + # The first 4 rows occupy rows 0..3 with 3 left / 3 right zero + # padding; rows 4..7 are all zero. + assert out.shape == (1, 1, 8, 8) + # The data appears in columns 3..4 of rows 0..3. + interior = out[0, 0, 0:4, 3:5] + assert torch.equal(interior, x[0, 0]) + assert out[0, 0, 4:, :].abs().sum() == 0 + # Round-trip via unfill (inverted by `unfill` checking tgt_size). + assert torch.equal(p.unfill(out), x) + + def test_size_argument_overrides_stride_and_sets_tgt_size(self) -> None: + p = InputPadder((3, 5), stride=8, size=(8, 8)) + # pad_ht = 8 - 3 = 5; pad_wd = 8 - 5 = 3; two_side -> [1, 2, 2, 3]. + assert p._pad == [1, 2, 2, 3] + assert p.tgt_size == (8, 8) + + def test_fill_unfill_roundtrips_for_exact_target_size(self) -> None: + p = InputPadder((10, 10), stride=8) + assert p.tgt_size == (16, 16) + x = torch.arange(100, dtype=torch.float32).reshape(1, 1, 10, 10) + filled = p.fill(x) + assert filled.shape == (1, 1, 16, 16) + assert torch.equal(p.unfill(filled), x) + + def test_unfill_skips_unpad_when_shape_is_not_target(self) -> None: + # A tensor whose H, W are *not* equal to tgt_size should be returned + # untouched (the wrapper's `unfill` guards with an exact-shape + # check), unlike `unpad` which slices regardless. + p = InputPadder((4, 4), stride=8) # tgt_size = (8, 8) + x = torch.arange(16, dtype=torch.float32).reshape(1, 1, 4, 4) + assert torch.equal(p.unfill(x), x) + + +# --------------------------------------------------------------------------- +# InputScaler +# --------------------------------------------------------------------------- + + +class TestInputScaler: + def test_stride_target_is_next_multiple_of_stride(self) -> None: + # ceil(4/8)*8 = 8, ceil(7/8)*8 = 8. + sc = InputScaler((4, 7), stride=8) + assert sc.tgt_height == 8 + assert sc.tgt_width == 8 + + def test_size_argument_sets_target_size(self) -> None: + sc = InputScaler((10, 10), size=(4, 4)) + assert sc.tgt_height == 4 + assert sc.tgt_width == 4 + + def test_scale_factor_when_stride_and_size_are_none(self) -> None: + # int(5 * 2) = 10, int(7 * 2) = 14. + sc = InputScaler((5, 7), scale_factor=2.0) + assert sc.tgt_height == 10 + assert sc.tgt_width == 14 + + def test_providing_both_stride_and_size_raises(self) -> None: + with pytest.raises(AssertionError): + InputScaler((4, 4), stride=8, size=(8, 8)) + + def test_fill_uses_bilinear_interpolation_with_align_corners_false(self) -> None: + # 2x2 -> 4x4 with bilinear (align_corners=False) matches + # torch.nn.functional.interpolate so we compare to that explicitly. + sc = InputScaler((2, 2), stride=4) + x = torch.tensor([[[[0.0, 1.0], [2.0, 3.0]]]]) + out = sc.fill(x) + expected = torch.nn.functional.interpolate( + x, size=(4, 4), mode="bilinear", align_corners=False + ) + assert torch.equal(out, expected) + + def test_fill_preserves_leading_dimensions(self) -> None: + sc = InputScaler((2, 2), stride=4) # tgt (4, 4) + x = torch.arange(2 * 1 * 3 * 2 * 2, dtype=torch.float32).reshape(2, 1, 3, 2, 2) + out = sc.fill(x) + assert out.shape == (2, 1, 3, 4, 4) + + def test_fill_flow_scales_channels_by_resolution_ratio(self) -> None: + # 1x1 -> 4x4, ratio_w = ratio_h = 4/1 = 4. fx=2 -> 8, fy=1 -> 4. + sc = InputScaler((1, 1), stride=4) + flow = torch.tensor([[[[2.0]], [[1.0]]]]) # (1, 2, 1, 1) + out = sc.fill(flow, is_flow=True) + assert torch.allclose(out[:, 0], torch.full((1, 4, 4), 8.0)) + assert torch.allclose(out[:, 1], torch.full((1, 4, 4), 4.0)) + + def test_fill_flow_downscale_shrinks_values(self) -> None: + # 4x4 -> 2x2, ratio_w = ratio_h = 2/4 = 0.5. fx=fy=2 -> 1.0. + sc = InputScaler((4, 4), size=(2, 2)) + flow = torch.full((1, 2, 4, 4), 2.0) + out = sc.fill(flow, is_flow=True) + assert torch.allclose(out, torch.full((1, 2, 2, 2), 1.0)) + + def test_constant_tensor_roundtrips_through_fill_unfill(self) -> None: + # Bilinear down/up of a constant returns the same constant. + sc = InputScaler((2, 2), stride=4) + x = torch.full((1, 1, 2, 2), 3.5) + assert torch.equal(sc.unfill(sc.fill(x)), x) + + +# --------------------------------------------------------------------------- +# add_datasets_to_parser +# --------------------------------------------------------------------------- + + +class TestAddDatasetsToParser: + def test_adds_argument_per_entry_in_yaml(self, tmp_path: Path) -> None: + cfg = tmp_path / "datasets.yaml" + cfg.write_text("sintel: /path/to/sintel\nkitti: /path/to/kitti\n") + parser = ArgumentParser() + out_parser = add_datasets_to_parser(parser, str(cfg)) + dests = {a.dest for a in out_parser._actions} + assert "sintel_root_dir" in dests + assert "kitti_root_dir" in dests + + def test_default_path_taken_from_yaml(self, tmp_path: Path) -> None: + cfg = tmp_path / "datasets.yaml" + cfg.write_text("sintel: /path/to/sintel\n") + parser = ArgumentParser() + add_datasets_to_parser(parser, str(cfg)) + args = parser.parse_args([]) + assert args.sintel_root_dir == "/path/to/sintel" + + def test_cli_overrides_default(self, tmp_path: Path) -> None: + cfg = tmp_path / "datasets.yaml" + cfg.write_text("sintel: /default\n") + parser = ArgumentParser() + add_datasets_to_parser(parser, str(cfg)) + args = parser.parse_args(["--sintel_root_dir", "/custom"]) + assert args.sintel_root_dir == "/custom" + + +# --------------------------------------------------------------------------- +# count_parameters +# --------------------------------------------------------------------------- + + +class TestCountParameters: + def test_counts_all_trainable_parameters(self) -> None: + # Linear(3, 2): weight 3*2 = 6 + bias 2 = 8 trainable. + m = nn.Linear(3, 2) + assert count_parameters(m) == 8 + + def test_ignores_frozen_parameters(self) -> None: + m = nn.Linear(3, 2) + m.weight.requires_grad = False + # Only the bias (2) is trainable. + assert count_parameters(m) == 2 + + def test_returns_zero_for_no_trainable_parameters(self) -> None: + m = nn.Linear(3, 2) + for p in m.parameters(): + p.requires_grad = False + assert count_parameters(m) == 0 + + +# --------------------------------------------------------------------------- +# make_divisible +# --------------------------------------------------------------------------- + + +class TestMakeDivisible: + @pytest.mark.parametrize( + "v, div, expected", + [ + (10, 3, 9), # 10 - 1 = 9 + (8, 4, 8), # 8 - 0 = 8 + (2, 5, 5), # 0 clamped up to 5 + (0, 5, 5), # 0 clamped up to 5 + (12, 4, 12), # 12 - 0 = 12 + (7, 1, 7), # 7 - 0 = 7 + ], + ) + def test_make_divisible(self, v, div, expected) -> None: + assert make_divisible(v, div) == expected + + +# --------------------------------------------------------------------------- +# release_gpu +# --------------------------------------------------------------------------- + + +class TestReleaseGpu: + def test_tensors_are_detached_and_moved_to_cpu(self) -> None: + d = {"a": torch.tensor([1.0, 2.0])} + out = release_gpu(d) + assert out["a"].device.type == "cpu" + assert out["a"].tolist() == [1.0, 2.0] + + def test_non_tensor_entries_kept_intact(self) -> None: + d = {"a": torch.tensor([1.0]), "b": "hello", "c": 5, "d": [1, 2]} + out = release_gpu(d) + assert out["b"] == "hello" + assert out["c"] == 5 + assert out["d"] == [1, 2] + + def test_returns_same_dict_instance(self) -> None: + d = {"a": torch.tensor([1.0])} + out = release_gpu(d) + assert out is d + + def test_detach_breaks_grad_chain(self) -> None: + x = torch.tensor([1.0], requires_grad=True) + y = x * 2 + out = release_gpu({"y": y}) + assert not out["y"].requires_grad + + +# --------------------------------------------------------------------------- +# tensor_dict_to_numpy +# --------------------------------------------------------------------------- + + +class TestTensorDictToNumpy: + def test_chw_tensor_becomes_hwc_numpy(self) -> None: + td = { + "flows": torch.arange(1 * 2 * 3 * 4, dtype=torch.float32).reshape( + 1, 2, 3, 4 + ) + } + out = tensor_dict_to_numpy(td) + assert out["flows"].shape == (3, 4, 2) + assert isinstance(out["flows"], np.ndarray) + + def test_single_image_tensor_becomes_hwc(self) -> None: + # 4D (1, 1, 3, 4) -> squeeze batch -> (3, 4, 1). + td = { + "imgs": torch.arange(1 * 1 * 3 * 4, dtype=torch.float32).reshape(1, 1, 3, 4) + } + out = tensor_dict_to_numpy(td) + assert out["imgs"].shape == (3, 4, 1) + + def test_squeezes_leading_dims_until_three(self) -> None: + # 6D (2, 1, 1, 2, 3, 4) -> first slice until (2, 3, 4) -> (3, 4, 2). + td = { + "x": torch.arange(2 * 1 * 1 * 2 * 3 * 4, dtype=torch.float32).reshape( + 2, 1, 1, 2, 3, 4 + ) + } + out = tensor_dict_to_numpy(td) + assert out["x"].shape == (3, 4, 2) + + def test_non_tensor_entries_kept_untouched(self) -> None: + td = {"a": 5, "b": "hello", "c": [1, 2, 3]} + out = tensor_dict_to_numpy(td) + assert out["a"] == 5 + assert out["b"] == "hello" + assert out["c"] == [1, 2, 3] + + def test_padded_tensor_is_unpadded_by_padder(self) -> None: + # Padded tensor (1, 2, 8, 8) with padder dims=(3, 4) -> tgt_size + # (8, 8) -> unpadded to (3, 4) -> HWC (3, 4, 2). + padder = InputPadder((3, 4), stride=8) + assert padder.tgt_size == (8, 8) + x = torch.arange(2 * 8 * 8, dtype=torch.float32).reshape(1, 2, 8, 8) + out = tensor_dict_to_numpy({"x": x}, padder=padder) + assert out["x"].shape == (3, 4, 2) + + def test_unpadded_tensor_with_padder_round_trips_to_original( + self, + ) -> None: + # Passing a tensor whose shape equals tgt_size should be unpadded + # back to the original (H, W) shape, then squeezed to 3D and laid out + # as HWC. + padder = InputPadder((3, 4), stride=8) + padded = torch.arange(1 * 2 * 8 * 8, dtype=torch.float32).reshape(1, 2, 8, 8) + # Mirror exactly what tensor_dict_to_numpy does internally: + # unpad -> squeeze leading dims until 3D -> permute(1, 2, 0). + expected = padder.unpad(padded) + while expected.ndim > 3: + expected = expected[0] + expected = expected.permute(1, 2, 0).numpy() + out = tensor_dict_to_numpy({"x": padded}, padder=padder) + assert out["x"].shape == (3, 4, 2) + assert np.array_equal(out["x"], expected) + + def test_1d_tensor_entry_is_converted_without_crash(self) -> None: + # A 1D tensor value should be converted to a numpy array (via + # detach/cpu) without the CHW -> HWC permute that requires 3D. + td = {"a": torch.tensor([1.0, 2.0, 3.0])} + out = tensor_dict_to_numpy(td) + assert isinstance(out["a"], np.ndarray) + assert out["a"].tolist() == [1.0, 2.0, 3.0] + + def test_2d_tensor_entry_is_converted_without_crash(self) -> None: + # 2D tensor: no channel axis, kept as-is in numpy. + td = {"a": torch.tensor([[1.0, 2.0], [3.0, 4.0]])} + out = tensor_dict_to_numpy(td) + assert isinstance(out["a"], np.ndarray) + assert out["a"].shape == (2, 2) + assert out["a"].tolist() == [[1.0, 2.0], [3.0, 4.0]] + + def test_unpadded_tensor_with_padder_is_left_untouched(self) -> None: + # Passing a tensor whose H, W are *not* tgt_size alongside a padder + # (because the caller did not pad it first) should leave the tensor + # untouched (via padder.unfill's shape guard) instead of producing an + # empty array. + padder = InputPadder((3, 4), stride=8) # tgt (8, 8) + x = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(1, 2, 3, 4) + out = tensor_dict_to_numpy({"x": x}, padder=padder) + assert out["x"].size != 0 + # Shape should be the original (3, 4, 2) HWC layout. + assert out["x"].shape == (3, 4, 2) + # Content matches the CHW -> HWC permute of the input. + assert np.array_equal(out["x"], x[0].permute(1, 2, 0).numpy()) + + +# --------------------------------------------------------------------------- +# are_shapes_compatible +# --------------------------------------------------------------------------- + + +class TestAreShapesCompatible: + @pytest.mark.parametrize( + "s1, s2, expected", + [ + ((1, 2, 3), (1, 2, 3), True), + ((1, 2, 3), (1, 1, 3), True), + ((2, 3), (2, 2, 3), False), # different number of dims + ((2, 3), (3, 2), False), # neither equal nor has a 1 + ((1, 1, 1), (4, 5, 1), True), # 1s broadcast everywhere + ((2, 1, 4), (1, 3, 4), True), # both have a 1 + ((2, 3, 4), (2, 3, 5), False), # 4 != 5, neither 1 + ((), (), True), # empty shapes are equal + ((1,), (), False), # different length + ], + ) + def test_compatibility(self, s1, s2, expected) -> None: + assert are_shapes_compatible(s1, s2) == expected + + +# --------------------------------------------------------------------------- +# bgr_val_as_tensor +# --------------------------------------------------------------------------- + + +class TestBgrValAsTensor: + def test_scalar_is_repeated_three_times_at_position_minus_three(self) -> None: + ref = torch.zeros(2, 3, 4, 5, dtype=torch.float32) + out = bgr_val_as_tensor(0.5, ref) + assert out.shape == (1, 3, 1, 1) + assert torch.allclose(out.flatten(), torch.full((3,), 0.5)) + + def test_tuple_of_three_is_reshaped_to_position(self) -> None: + ref = torch.zeros(2, 3, 4, 5, dtype=torch.float32) + out = bgr_val_as_tensor((1.0, 2.0, 3.0), ref) + assert out.shape == (1, 3, 1, 1) + assert torch.allclose(out.flatten(), torch.tensor([1.0, 2.0, 3.0])) + + def test_numpy_array_is_reshaped_and_dtypes_match_reference(self) -> None: + ref = torch.zeros(2, 3, 4, 5, dtype=torch.float32) + out = bgr_val_as_tensor(np.array([1.0, 2.0, 3.0], dtype=np.float32), ref) + assert out.shape == (1, 3, 1, 1) + assert out.dtype == ref.dtype + assert torch.allclose(out.flatten(), torch.tensor([1.0, 2.0, 3.0])) + + def test_chw_position_can_be_overridden_to_minus_one(self) -> None: + # Channels-last reference: the BGR tensor should be at the -1 + # position with shape (1, 1, 1, 3). + ref = torch.zeros(2, 4, 5, 3, dtype=torch.float32) + out = bgr_val_as_tensor((1.0, 2.0, 3.0), ref, bgr_tensor_shape_position=-1) + assert out.shape == (1, 1, 1, 3) + + def test_tensor_compatible_with_reference_is_returned_untouched(self) -> None: + # A tensor whose shape is broadcastable to ref is left as-is. + ref = torch.zeros(2, 3, 4, 5, dtype=torch.float32) + bgr = torch.zeros(1, 3, 1, 1, dtype=torch.float32) + out = bgr_val_as_tensor(bgr, ref) + assert out is bgr + + def test_list_of_three_is_accepted(self) -> None: + ref = torch.zeros(2, 3, 4, 5, dtype=torch.float32) + out = bgr_val_as_tensor([1.0, 2.0, 3.0], ref) + assert torch.allclose(out.flatten(), torch.tensor([1.0, 2.0, 3.0])) + + def test_tuple_of_wrong_length_raises(self) -> None: + ref = torch.zeros(2, 3, 4, 5, dtype=torch.float32) + with pytest.raises(AssertionError): + bgr_val_as_tensor((1.0, 2.0), ref) + + +# --------------------------------------------------------------------------- +# forward_interpolate_batch +# --------------------------------------------------------------------------- + + +class TestForwardInterpolateBatch: + def test_zero_flow_returns_zeros(self) -> None: + # forward_interpolate uses strict bounds (x1 > 0 & x1 < wd & + # y1 > 0 & y1 < ht), so zero flow keeps no valid source -> all zeros. + flow = torch.zeros(2, 2, 3, 3) + out = forward_interpolate_batch(flow) + assert out.shape == (2, 2, 3, 3) + assert torch.equal(out, torch.zeros_like(out)) + + def test_shape_preserved_for_each_batch_element(self) -> None: + flow = torch.zeros(3, 2, 5, 7) + out = forward_interpolate_batch(flow) + assert out.shape == (3, 2, 5, 7) + + def test_dtype_propagates_to_input_dtype(self) -> None: + flow = torch.zeros(2, 2, 3, 3, dtype=torch.float64) + out = forward_interpolate_batch(flow) + assert out.dtype == torch.float64 + + +# --------------------------------------------------------------------------- +# get_matplotlib_lut_colormap +# --------------------------------------------------------------------------- + + +class TestGetMatplotlibLutColormap: + def test_returns_256_1_3_uint8_bgr_lut(self) -> None: + lut = get_matplotlib_lut_colormap("RdYlBu_r") + assert lut.shape == (256, 1, 3) + assert lut.dtype == np.uint8 + + def test_matches_independent_lut_computation(self) -> None: + import matplotlib.pyplot as plt + + cmap = plt.get_cmap("RdYlBu_r") + rgb = (cmap(np.arange(256))[:, :3] * 255).astype(np.uint8) + expected = rgb[:, ::-1].reshape(256, 1, 3) + out = get_matplotlib_lut_colormap("RdYlBu_r") + assert np.array_equal(out, expected) + + def test_bgr_order_of_first_and_last_entries(self) -> None: + # Independently-read entries of matplotlib RdYlBu_r at indices 0 and + # 255 mirrored to BGR. matplotlib[0] -> [0.584, 0.211, 0.192] (RGB) + # -> [149, 54, 49] (BGR, uint8). matplotlib[255] -> [0.647, 0, + # 0.149] -> [38, 0, 165] (BGR, uint8). + lut = get_matplotlib_lut_colormap("RdYlBu_r") + assert lut[0, 0].tolist() == [149, 54, 49] + assert lut[255, 0].tolist() == [38, 0, 165] + + +# --------------------------------------------------------------------------- +# config_logging +# --------------------------------------------------------------------------- + + +class TestConfigLogging: + def test_creates_log_file_in_ptlflow_logs(self, tmp_path, monkeypatch) -> None: + monkeypatch.chdir(tmp_path) + utils_mod.config_logging() + assert (tmp_path / "ptlflow_logs" / "log_run.txt").exists()