Fix spatial_ndim after EnsureChannelFirst adds a channel dim

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):

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()