diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py index 297b243ff9..b12b6c9289 100644 --- a/monai/transforms/utility/array.py +++ b/monai/transforms/utility/array.py @@ -197,6 +197,16 @@ def __init__(self, strict_check: bool = True, channel_dim: None | str | int = No def __call__(self, img: torch.Tensor, meta_dict: Mapping | None = None) -> torch.Tensor: """ Apply the transform to `img`. + + If the input has no channel dim, a channel dim is added and the ``spatial_ndim`` of the + output ``MetaTensor`` is derived again from its affine. + + Args: + img: the input image. + meta_dict: metadata dictionary of `img`, used when `img` is not a ``MetaTensor``. + + Returns: + the image with the channel dim first. """ if not isinstance(img, MetaTensor) and not isinstance(meta_dict, Mapping): if self.input_channel_dim is None: @@ -231,7 +241,12 @@ def __call__(self, img: torch.Tensor, meta_dict: Mapping | None = None) -> torch else: result = moveaxis(img, int(channel_dim), 0) # type: ignore - return convert_to_tensor(result, track_meta=get_track_meta()) # type: ignore + result = convert_to_tensor(result, track_meta=get_track_meta()) + if is_no_channel(channel_dim) and isinstance(result, MetaTensor) and result.affine.ndim == 2: + # the cached spatial_ndim may have been clamped while the input had no channel dim, + # so derive it again from the affine now that the channel dim exists + result.spatial_ndim = _normalize_spatial_ndim(result.affine.shape[-1] - 1, result.ndim) + return result # type: ignore class RepeatChannel(Transform): diff --git a/tests/transforms/test_ensure_channel_firstd.py b/tests/transforms/test_ensure_channel_firstd.py index e9effad951..ab06657797 100644 --- a/tests/transforms/test_ensure_channel_firstd.py +++ b/tests/transforms/test_ensure_channel_firstd.py @@ -22,7 +22,7 @@ from PIL import Image from monai.data.meta_tensor import MetaTensor -from monai.transforms import EnsureChannelFirstd, LoadImaged +from monai.transforms import EnsureChannelFirstd, LoadImaged, SpatialCropd TEST_CASE_1 = [{"keys": "img"}, ["test_image.nii.gz"], None] @@ -75,11 +75,31 @@ def test_exceptions(self): EnsureChannelFirstd("img", strict_check=False, channel_dim=None)({"img": im_nodim}) def test_default_channel_first(self): + """A channel-less tensor gets a leading channel dim.""" im = torch.rand(4, 4) result = EnsureChannelFirstd("img", channel_dim="no_channel")({"img": im}) self.assertEqual(result["img"].shape, (1, 4, 4)) + @parameterized.expand([((8, 8, 8), 3), ((8, 8), 2)]) + def test_no_channel_spatial_ndim(self, shape, expected): + """spatial_ndim follows the affine after a channel dim is added to channel-less input.""" + affine = np.eye(4) + affine[:3, 3] = [10.0, 20.0, 30.0] + im = MetaTensor(torch.zeros(shape), affine=affine) + result = EnsureChannelFirstd("img", channel_dim="no_channel")({"img": im}) + self.assertEqual(result["img"].shape, (1, *shape)) + self.assertEqual(result["img"].spatial_ndim, expected) + + def test_no_channel_crop_affine(self): + """A crop after adding a channel dim shifts the affine origin on all spatial axes.""" + affine = np.eye(4) + affine[:3, 3] = [10.0, 20.0, 30.0] + im = MetaTensor(torch.zeros((8, 8, 8)), affine=affine) + data = EnsureChannelFirstd("img", channel_dim="no_channel")({"img": im}) + data = SpatialCropd("img", roi_start=[1, 2, 3], roi_end=[6, 6, 6])(data) + np.testing.assert_allclose(data["img"].affine[:3, 3], [11.0, 22.0, 33.0]) + if __name__ == "__main__": unittest.main()