4731 squeeze dim update meta (#5041)

Signed-off-by: Wenqi Li <[email protected]>

Fixes #4731

(additionally fixes Project-MONAI/tutorials#569)

### Description
squeeze dim drops a spatial axis should also update the affine

### Status
**Ready**

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not applicable items -->
- [x] Non-breaking change (fix or new feature that would not break existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing functionality to change).
- [ ] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u --net --coverage`.
- [ ] Quick tests passed locally by running `./runtests.sh --quick --unittests  --disttests`.
- [ ] In-line docstrings updated.
- [ ] Documentation updated, tested `make html` command in the `docs/` folder.

Signed-off-by: Wenqi Li <[email protected]>

Author: wyli

monai/transforms/utility/array.py

@@ -603,11 +603,12 @@ class SqueezeDim(Transform):
 
     backend = [TransformBackends.TORCH, TransformBackends.NUMPY]
 
-    def __init__(self, dim: Optional[int] = 0) -> None:
+    def __init__(self, dim: Optional[int] = 0, update_meta=True) -> None:
         """
         Args:
             dim: dimension to be squeezed. Default = 0
                 "None" works when the input is numpy array.
+            update_meta: whether to update the meta info if the input is a metatensor. Default is ``True``.
 
         Raises:
             TypeError: When ``dim`` is not an ``Optional[int]``.
@@ -616,6 +617,7 @@ def __init__(self, dim: Optional[int] = 0) -> None:
         if dim is not None and not isinstance(dim, int):
             raise TypeError(f"dim must be None or a int but is {type(dim).__name__}.")
         self.dim = dim
+        self.update_meta = update_meta
 
     def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
@@ -624,11 +626,25 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
         img = convert_to_tensor(img, track_meta=get_track_meta())
         if self.dim is None:
+            if self.update_meta:
+                warnings.warn("update_meta=True is ignored when dim=None.")
             return img.squeeze()
+        dim = (self.dim + len(img.shape)) if self.dim < 0 else self.dim
         # for pytorch/numpy unification
-        if img.shape[self.dim] != 1:
-            raise ValueError(f"Can only squeeze singleton dimension, got shape {img.shape}.")
-        return img.squeeze(self.dim)
+        if img.shape[dim] != 1:
+            raise ValueError(f"Can only squeeze singleton dimension, got shape {img.shape[dim]} of {img.shape}.")
+        img = img.squeeze(dim)
+        if self.update_meta and isinstance(img, MetaTensor) and dim > 0 and len(img.affine.shape) == 2:
+            h, w = img.affine.shape
+            affine, device = img.affine, img.affine.device if isinstance(img.affine, torch.Tensor) else None
+            if h > dim:
+                affine = affine[torch.arange(0, h, device=device) != dim - 1]
+            if w > dim:
+                affine = affine[:, torch.arange(0, w, device=device) != dim - 1]
+            if (affine.shape[0] == affine.shape[1]) and not np.linalg.det(convert_to_numpy(affine, wrap_sequence=True)):
+                warnings.warn(f"After SqueezeDim, img.affine is ill-posed: \n{img.affine}.")
+            img.affine = affine
+        return img
 
 
 class DataStats(Transform):

monai/transforms/utility/dictionary.py

@@ -763,16 +763,19 @@ class SqueezeDimd(MapTransform):
 
     backend = SqueezeDim.backend
 
-    def __init__(self, keys: KeysCollection, dim: int = 0, allow_missing_keys: bool = False) -> None:
+    def __init__(
+        self, keys: KeysCollection, dim: int = 0, update_meta: bool = True, allow_missing_keys: bool = False
+    ) -> None:
         """
         Args:
             keys: keys of the corresponding items to be transformed.
                 See also: :py:class:`monai.transforms.compose.MapTransform`
             dim: dimension to be squeezed. Default: 0 (the first dimension)
+            update_meta: whether to update the meta info if the input is a metatensor. Default is ``True``.
             allow_missing_keys: don't raise exception if key is missing.
         """
         super().__init__(keys, allow_missing_keys)
-        self.converter = SqueezeDim(dim=dim)
+        self.converter = SqueezeDim(dim=dim, update_meta=update_meta)
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
         d = dict(data)

monai/utils/jupyter_utils.py

@@ -14,6 +14,7 @@
 Matplotlib produce common plots for metrics and images.
 """
 
+import copy
 from enum import Enum
 from threading import RLock, Thread
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
@@ -340,7 +341,7 @@ def status_dict(self) -> Dict[str, str]:
 
     def status(self) -> str:
         """Returns a status string for the current state of the engine."""
-        stats = self.status_dict
+        stats = copy.deepcopy(self.status_dict)
 
         msgs = [stats.pop(StatusMembers.STATUS.value), "Iters: " + str(stats.pop(StatusMembers.ITERS.value, 0))]
 

tests/test_squeezedim.py

@@ -14,8 +14,9 @@
 import numpy as np
 from parameterized import parameterized
 
+from monai.data import MetaTensor
 from monai.transforms import SqueezeDim
-from tests.utils import TEST_NDARRAYS
+from tests.utils import TEST_NDARRAYS, assert_allclose
 
 TESTS, TESTS_FAIL = [], []
 for p in TEST_NDARRAYS:
@@ -34,13 +35,28 @@ def test_shape(self, input_param, test_data, expected_shape):
 
         result = SqueezeDim(**input_param)(test_data)
         self.assertTupleEqual(result.shape, expected_shape)
+        if "dim" in input_param and input_param["dim"] == 2 and isinstance(result, MetaTensor):
+            assert_allclose(result.affine.shape, [3, 3])
 
     @parameterized.expand(TESTS_FAIL)
     def test_invalid_inputs(self, exception, input_param, test_data):
 
         with self.assertRaises(exception):
             SqueezeDim(**input_param)(test_data)
 
+    def test_affine_ill_inputs(self):
+        img = MetaTensor(
+            np.random.rand(1, 2, 1, 3),
+            affine=[
+                [-0.7422, 0.0, 0.0, 186.3210],
+                [0.0, 0.0, -3.0, 70.6580],
+                [0.0, -0.7422, 0.0, 189.4130],
+                [0.0, 0.0, 0.0, 1.0],
+            ],
+        )
+        with self.assertWarns(UserWarning):
+            SqueezeDim(dim=2)(img)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_squeezedimd.py

@@ -14,8 +14,9 @@
 import numpy as np
 from parameterized import parameterized
 
+from monai.data import MetaTensor
 from monai.transforms import SqueezeDimd
-from tests.utils import TEST_NDARRAYS
+from tests.utils import TEST_NDARRAYS, assert_allclose
 
 TESTS, TESTS_FAIL = [], []
 for p in TEST_NDARRAYS:
@@ -82,6 +83,8 @@ def test_shape(self, input_param, test_data, expected_shape):
         result = SqueezeDimd(**input_param)(test_data)
         self.assertTupleEqual(result["img"].shape, expected_shape)
         self.assertTupleEqual(result["seg"].shape, expected_shape)
+        if "dim" in input_param and isinstance(result["img"], MetaTensor) and input_param["dim"] == 2:
+            assert_allclose(result["img"].affine.shape, [3, 3])
 
     @parameterized.expand(TESTS_FAIL)
     def test_invalid_inputs(self, exception, input_param, test_data):