diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index 297b243ff91..03dca190549 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -30,7 +30,7 @@ from monai.config import DtypeLike from monai.config.type_definitions import NdarrayOrTensor from monai.data.meta_obj import get_track_meta -from monai.data.meta_tensor import MetaTensor, _normalize_spatial_ndim, get_spatial_ndim +from monai.data.meta_tensor import MetaTensor, _has_explicit_no_channel, _normalize_spatial_ndim, get_spatial_ndim from monai.data.utils import is_no_channel, no_collation, orientation_ras_lps from monai.networks.layers.simplelayers import ( ApplyFilter, @@ -231,6 +231,17 @@ def __call__(self, img: torch.Tensor, meta_dict: Mapping | None = None) -> torch else: result = moveaxis(img, int(channel_dim), 0) # type: ignore + # When a channel dim is added (channel_dim="no_channel"), the affine stays + # the same but the tensor gains a channel axis, so the cached `spatial_ndim` + # (derived from the input tensor's shape) would be too low. Refresh it from + # the affine, which is the source of truth for spatial dimensions, so that + # downstream lazy transforms (e.g. spatial crop) build translations with the + # correct number of axes and preserve the affine offset in world space. + if is_no_channel(channel_dim) and isinstance(result, MetaTensor): + result.spatial_ndim = _normalize_spatial_ndim( + result.affine.shape[-1] - 1, result.ndim, no_channel=_has_explicit_no_channel(result.meta) + ) + return convert_to_tensor(result, track_meta=get_track_meta()) # type: ignore diff --git a/tests/transforms/test_ensure_channel_first.py b/tests/transforms/test_ensure_channel_first.py index fe046a4cdf9..9574aad1477 100644 --- a/tests/transforms/test_ensure_channel_first.py +++ b/tests/transforms/test_ensure_channel_first.py @@ -22,7 +22,7 @@ from PIL import Image from monai.data.meta_tensor import MetaTensor -from monai.transforms import EnsureChannelFirst, LoadImage +from monai.transforms import EnsureChannelFirst, LoadImage, SpatialCrop from monai.utils import optional_import itk, has_itk = optional_import("itk", allow_namespace_pkg=True) @@ -108,6 +108,23 @@ def test_default_channel_first(self): self.assertEqual(result.shape, (1, 4, 4)) + def test_ensure_channel_first_no_channel_affine_regression(self): + # Regression test for https://github.com/Project-MONAI/MONAI/issues/9184. + # EnsureChannelFirst(channel_dim="no_channel") must refresh the cached + # `spatial_ndim` from the affine; otherwise a following spatial crop only + # shifts the first (N-1) axes and silently drops the last affine offset. + affine = np.eye(4) + affine[:3, 3] = [10, 20, 30] + img = MetaTensor(np.zeros((40, 40, 40), dtype=np.float32), affine=affine) + + result = EnsureChannelFirst(channel_dim="no_channel")(img) + self.assertEqual(result.spatial_ndim, 3) + self.assertEqual(result.ndim, 4) + + cropped = SpatialCrop(roi_start=[5, 5, 5], roi_end=[20, 20, 20])(result) + self.assertEqual(cropped.spatial_ndim, 3) + self.assertEqual(cropped.affine[:3, 3].tolist(), [15, 25, 35]) + if __name__ == "__main__": unittest.main()