From 6640b598c3e4713977aaa7fd9318746d9b4ac2d2 Mon Sep 17 00:00:00 2001 From: Shreyansh Goyal Date: Thu, 16 Jul 2026 00:18:47 +0530 Subject: [PATCH] Fix affine-elastic composition order --- src/torchio/transforms/spatial/spatial.py | 11 +++++---- tests/test_spatial.py | 28 +++++++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/torchio/transforms/spatial/spatial.py b/src/torchio/transforms/spatial/spatial.py index 117c46c3c..d19991274 100644 --- a/src/torchio/transforms/spatial/spatial.py +++ b/src/torchio/transforms/spatial/spatial.py @@ -1568,13 +1568,14 @@ def _build_sampling_grid( ) if affine_first: - # Affine first: map to input space, then add elastic offset. - input_voxels = _apply_voxel_mapping(output_coords, mapping) - input_voxels = input_voxels + displacement / input_spacing_t - else: - # Elastic first: deform in output space, then map to input. + # Resampling follows the inverse transform, so apply the operations + # in reverse order: undo the elastic field, then the affine mapping. deformed_output = output_coords + displacement / output_spacing_t input_voxels = _apply_voxel_mapping(deformed_output, mapping) + else: + # Undo the affine mapping first, then the elastic field. + input_voxels = _apply_voxel_mapping(output_coords, mapping) + input_voxels = input_voxels + displacement / input_spacing_t return input_voxels diff --git a/tests/test_spatial.py b/tests/test_spatial.py index 4c98800cc..1850de953 100644 --- a/tests/test_spatial.py +++ b/tests/test_spatial.py @@ -88,6 +88,34 @@ def test_affine_first_changes_result(self) -> None: assert not torch.allclose(first.t1.data, second.t1.data) + def test_affine_first_composes_inverse_grid_in_reverse_order(self) -> None: + identity = AffineMatrix(np.eye(4)) + affine_matrix = np.diag([2.0, 1.0, 1.0, 1.0]) + control_points = torch.zeros(4, 4, 4, 3) + control_points[..., 0] = 1.0 + kwargs = { + "input_shape": (3, 3, 3), + "input_affine": identity, + "output_shape": (3, 3, 3), + "output_affine": identity, + "affine_matrix": affine_matrix, + "control_points": control_points, + "max_displacement": (1.0, 0.0, 0.0), + "device": torch.device("cpu"), + } + + affine_then_elastic = _build_sampling_grid(affine_first=True, **kwargs) + elastic_then_affine = _build_sampling_grid(affine_first=False, **kwargs) + + torch.testing.assert_close( + affine_then_elastic[2, 0, 0], + torch.tensor([1.5, 0.0, 0.0]), + ) + torch.testing.assert_close( + elastic_then_affine[2, 0, 0], + torch.tensor([2.0, 0.0, 0.0]), + ) + def test_2d_suppresses_out_of_plane(self) -> None: data = torch.rand(1, 8, 8, 1) subject = tio.Subject(t1=tio.ScalarImage(data))