From 6fd1e02caec960c22b919ea4aa1a617f070675a3 Mon Sep 17 00:00:00 2001 From: tritsystem Date: Fri, 16 Oct 2026 02:27:47 -0700 Subject: [PATCH] fix(lazy): keep the image's affine rank on the lazy crop/pad path lazy.utils.resample reduced the affine to the spatial rank and, on the crop/pad/flip shortcut, assigned that reduced affine to the output. A 2D image with the default 4x4 affine therefore came back with a 3x3 affine and lost its out-of-plane spacing and origin (eager mode keeps them), e.g. a lazily cropped 2D slice saved to NIfTI is written with the wrong geometry. Compose the transform at the image's own affine rank instead. The interpolating path is unchanged. Assisted-by: Claude Opus 5.5 Signed-off-by: tritsystem --- monai/transforms/lazy/utils.py | 9 ++++++-- tests/transforms/functional/test_resample.py | 22 ++++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/monai/transforms/lazy/utils.py b/monai/transforms/lazy/utils.py index 1f8506031a9..d19f5024056 100644 --- a/monai/transforms/lazy/utils.py +++ b/monai/transforms/lazy/utils.py @@ -193,6 +193,11 @@ def resample(data: torch.Tensor, matrix: NdarrayOrTensor, kwargs: dict | None = ndim = len(matrix) - 1 img = convert_to_tensor(data=data, track_meta=monai.data.get_track_meta()) init_affine = monai.data.to_affine_nd(ndim, img.affine) + # The crop/pad shortcut below writes the affine directly, so compose at the image's own affine rank: a 2D image + # with a 4x4 affine must keep its out-of-plane spacing and origin, as the eager transforms do. + full_rank = len(img.affine) - 1 + full_xform = monai.data.to_affine_nd(full_rank, matrix) + full_dst_affine = img.affine @ monai.utils.convert_to_dst_type(full_xform, img.affine)[0] spatial_size = kwargs.get(LazyAttr.SHAPE, None) out_spatial_size = img.peek_pending_shape() if spatial_size is None else spatial_size out_spatial_size = convert_to_numpy(out_spatial_size, wrap_sequence=True) @@ -225,13 +230,13 @@ def resample(data: torch.Tensor, matrix: NdarrayOrTensor, kwargs: dict | None = and len(in_shape) == len(out_spatial_size) and allclose(convert_to_numpy(in_shape, wrap_sequence=True), out_spatial_size) ): - img.affine = call_kwargs["dst_affine"] + img.affine = full_dst_affine img = img.to(torch.float32) # consistent with monai.transforms.spatial.functional.spatial_resample return img # pyrefly: ignore [bad-argument-type, implicit-import] img = monai.transforms.crop_or_pad_nd(img, matrix_np, out_spatial_size, mode=call_kwargs["padding_mode"]) img = img.to(torch.float32) # consistent with monai.transforms.spatial.functional.spatial_resample - img.affine = call_kwargs["dst_affine"] + img.affine = full_dst_affine return img # pyrefly: ignore [bad-argument-type, implicit-import] diff --git a/tests/transforms/functional/test_resample.py b/tests/transforms/functional/test_resample.py index 40d264598d3..90505699e08 100644 --- a/tests/transforms/functional/test_resample.py +++ b/tests/transforms/functional/test_resample.py @@ -16,6 +16,8 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor +from monai.transforms import CenterSpatialCropd, Compose, Flipd, SpatialPadd from monai.transforms.lazy.functional import resample from monai.utils import convert_to_tensor from tests.test_utils import assert_allclose, get_arange_img @@ -45,6 +47,26 @@ def test_resample_function_impl(self, img, matrix, expected): out_1 = resample(img, matrix, {"lazy_resample_mode": "other value", "lazy_dtype": torch.float}) self.assertIs(out.dtype, out_1.dtype) # testing dtype in different lazy_resample_mode + @parameterized.expand( + [ + ["flip", lambda: Flipd("img", spatial_axis=0)], + ["crop", lambda: CenterSpatialCropd("img", (5, 4))], + ["pad", lambda: SpatialPadd("img", (12, 12))], + ] + ) + def test_lazy_crop_pad_keeps_affine_rank(self, _, make_transform): + """A 2D image with a 4x4 affine keeps its out-of-plane spacing and origin on the lazy crop/pad path.""" + affine = torch.diag(torch.tensor([1.5, 0.8, 2.0, 1.0], dtype=torch.float64)) + affine[:3, 3] = torch.tensor([3.0, -2.0, 1.0], dtype=torch.float64) + img = MetaTensor(torch.rand(1, 10, 8), affine=affine) + + eager = Compose([make_transform()], lazy=False)({"img": img.clone()})["img"] + lazy = Compose([make_transform()], lazy=True)({"img": img.clone()})["img"] + + self.assertEqual(lazy.affine.shape, (4, 4)) + assert_allclose(lazy.affine, eager.affine, type_test=False) + assert_allclose(lazy, eager, type_test=False) + if __name__ == "__main__": unittest.main()