Make NormalizeIntensity invertible (#5647)

monai/transforms/intensity/array.py

@@ -26,14 +26,15 @@
 from monai.config import DtypeLike
 from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
 from monai.data.meta_obj import get_track_meta
-from monai.data.meta_tensor import get_spatial_ndim
+from monai.data.meta_tensor import MetaTensor, get_spatial_ndim
 from monai.data.ultrasound_confidence_map import UltrasoundConfidenceMap
 from monai.data.utils import get_random_patch, get_valid_patch_size
 from monai.networks.layers import GaussianFilter, HilbertTransform, MedianFilter, SavitzkyGolayFilter
+from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.transform import RandomizableTransform, Transform
 from monai.transforms.utils import Fourier, equalize_hist, is_positive, rescale_array, soft_clip
-from monai.transforms.utils_pytorch_numpy_unification import clip, percentile, where
-from monai.utils.enums import TransformBackends
+from monai.transforms.utils_pytorch_numpy_unification import clip, nonzero, percentile, ravel, where
+from monai.utils.enums import TraceKeys, TransformBackends
 from monai.utils.misc import ensure_tuple, ensure_tuple_rep, ensure_tuple_size, fall_back_tuple
 from monai.utils.module import min_version, optional_import
 from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_tensor, get_equivalent_dtype
@@ -834,7 +835,7 @@ def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTen
         return out
 
 
-class NormalizeIntensity(Transform):
+class NormalizeIntensity(InvertibleTransform):
     """
     Normalize input based on the `subtrahend` and `divisor`: `(img - subtrahend) / divisor`.
     Use calculated mean or std value of the input image if no `subtrahend` or `divisor` provided.
@@ -844,6 +845,13 @@ class NormalizeIntensity(Transform):
     be the number of image channels if they are not None.
     If the input is not of floating point type, it will be converted to float32
 
+    The subtrahend and divisor actually used (whether provided or computed) are stored in the
+    transform's meta information, so the transform is invertible via :meth:`inverse`, recovering
+    ``img * divisor + subtrahend``. With ``nonzero=True`` only the voxels that were non-zero on the
+    forward pass are restored: they are identified as the non-zero voxels of the normalized image,
+    plus the (usually empty) set of voxels whose value equalled the subtrahend exactly and therefore
+    became zero, whose flat indices are also stored in the meta information.
+
     Args:
         subtrahend: the amount to subtract by (usually the mean).
         divisor: the amount to divide by (usually the standard deviation).
@@ -883,14 +891,29 @@ def _std(x):
         x = torch.std(x.float(), unbiased=False)
         return x.item() if x.numel() == 1 else x
 
-    def _normalize(self, img: NdarrayOrTensor, sub=None, div=None) -> NdarrayOrTensor:
+    def _normalize(self, img: NdarrayOrTensor, sub=None, div=None):
+        """
+        Normalize ``img`` in place where possible and report what was done, for :meth:`inverse`.
+
+        Args:
+            img: image (or single channel when ``channel_wise=True``) to normalize.
+            sub: subtrahend to use; computed as the mean of the (non-zero) voxels if None.
+            div: divisor to use; computed as the std of the (non-zero) voxels if None.
+
+        Returns:
+            a tuple ``(normalized, sub, div, zeroed_idx)``: the normalized image, the subtrahend and
+            divisor actually used (identity ``0.0``/``1.0`` when ``nonzero=True`` and there is nothing to
+            normalize), and, when ``nonzero=True``, the flat indices of voxels that were non-zero before but
+            are exactly zero after normalization (``None`` when ``nonzero=False``).
+        """
         img, *_ = convert_data_type(img, dtype=torch.float32)
 
         if self.nonzero:
             slices = img != 0
             masked_img = img[slices]
             if not slices.any():
-                return img
+                # nothing was normalized: store identity stats (keeps meta collate-safe) and no indices
+                return img, 0.0, 1.0, nonzero(ravel(slices))
         else:
             slices = None
             masked_img = img
@@ -911,11 +934,16 @@ def _normalize(self, img: NdarrayOrTensor, sub=None, div=None) -> NdarrayOrTenso
                 _div = _div[slices]
             _div[_div == 0.0] = 1.0
 
+        zeroed_idx = None
         if slices is not None:
             img[slices] = (masked_img - _sub) / _div
+            # voxels that were non-zero but now equal zero (value == subtrahend) are indistinguishable
+            # from the untouched zero voxels in the output, so record them for inverse().
+            zeroed_idx = nonzero(ravel(slices & (img == 0)))
         else:
             img = (img - _sub) / _div
-        return img
+        # Return the subtrahend/divisor actually used so the transform can be inverted.
+        return img, _sub, _div, zeroed_idx
 
     def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
@@ -924,6 +952,11 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         img_t: torch.Tensor = convert_to_tensor(img, track_meta=get_track_meta())  # type: ignore[assignment]
         dtype = self.dtype or img.dtype
         img_len = len(img_t)
+        # Subtrahend/divisor used per channel (channel_wise) or once (global), kept for inverse(),
+        # plus (nonzero=True only) the indices of voxels that were zeroed by the normalization.
+        subs: list = []
+        divs: list = []
+        zeroed: list = []
         if self.channel_wise:
             if self.subtrahend is not None and len(self.subtrahend) != img_len:
                 raise ValueError(f"img has {img_len} channels, but subtrahend has {len(self.subtrahend)} components.")
@@ -934,15 +967,108 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
                 img_t, *_ = convert_data_type(img_t, dtype=torch.float32)
 
             for i, d in enumerate(img_t):
-                img_t[i] = self._normalize(  # type: ignore
+                img_t[i], _sub, _div, _idx = self._normalize(  # type: ignore
                     d,
                     sub=self.subtrahend[i] if self.subtrahend is not None else None,
                     div=self.divisor[i] if self.divisor is not None else None,
                 )
+                subs.append(_sub)
+                divs.append(_div)
+                zeroed.append(_idx)
         else:
-            img_t = self._normalize(img_t, self.subtrahend, self.divisor)  # type: ignore[assignment]
+            img_t, _sub, _div, _idx = self._normalize(img_t, self.subtrahend, self.divisor)  # type: ignore
+            subs.append(_sub)
+            divs.append(_div)
+            zeroed.append(_idx)
 
         out = convert_to_dst_type(img_t, img_t, dtype=dtype)[0]
+        return self._push_transform_with_stats(out, subs, divs, zeroed)
+
+    @staticmethod
+    def _to_storable(value):
+        """
+        Convert a value computed on the forward pass to something storable in the transform meta information.
+
+        Args:
+            value: a subtrahend, divisor or index array; a python/numpy scalar, ``np.ndarray`` or ``torch.Tensor``.
+
+        Returns:
+            a detached, plain (non-Meta) CPU ``torch.Tensor`` for array inputs, otherwise the value unchanged.
+        """
+        if isinstance(value, MetaTensor):
+            value = value.as_tensor()
+        if isinstance(value, torch.Tensor):
+            return value.detach().cpu()
+        if isinstance(value, np.ndarray):
+            return torch.as_tensor(value)
+        return value  # python/numpy scalar
+
+    def _push_transform_with_stats(self, out: NdarrayOrTensor, subs: list, divs: list, zeroed: list) -> NdarrayOrTensor:
+        """
+        Record the parameters needed by :meth:`inverse` in the transform meta information of ``out``.
+
+        Args:
+            out: the normalized image; only a ``MetaTensor`` (with meta tracking enabled) can carry the record.
+            subs: subtrahend used for each channel (``channel_wise=True``) or a single-element list.
+            divs: divisor used for each channel (``channel_wise=True``) or a single-element list.
+            zeroed: per-channel flat indices of voxels zeroed by the normalization (``nonzero=True`` only).
+
+        Returns:
+            ``out``, with the transform pushed onto its applied operations when it is a ``MetaTensor``.
+        """
+        if not isinstance(out, MetaTensor) or not get_track_meta():
+            return out
+        extra_info = {
+            "sub": [self._to_storable(s) for s in subs],
+            "div": [self._to_storable(d) for d in divs],
+            "channel_wise": self.channel_wise,
+            "nonzero": self.nonzero,
+        }
+        if self.nonzero:
+            extra_info["zeroed_idx"] = [self._to_storable(z) for z in zeroed]
+        self.push_transform(out, extra_info=extra_info)
+        return out
+
+    def inverse(self, data: torch.Tensor) -> torch.Tensor:
+        """
+        Undo the normalization recorded on ``data`` by :meth:`__call__`, i.e. ``img * divisor + subtrahend``.
+        With ``nonzero=True`` only the voxels that were normalized on the forward pass are restored.
+
+        Args:
+            data: a ``MetaTensor`` produced by this transform, with the transform still on its applied operations.
+
+        Returns:
+            the de-normalized image, of the same type as ``data``.
+
+        Raises:
+            RuntimeError: if the most recent applied operation on ``data`` was not made by this transform.
+        """
+        transform = self.pop_transform(data)
+        info = transform[TraceKeys.EXTRA_INFO]
+        subs, divs = info["sub"], info["div"]
+        zeroed = info.get("zeroed_idx") if info["nonzero"] else None
+        out: torch.Tensor = convert_to_tensor(data, track_meta=get_track_meta())  # type: ignore[assignment]
+
+        def _restore(x, sub, div, zeroed_idx=None):
+            if zeroed_idx is None:
+                sub, *_ = convert_to_dst_type(sub, x)
+                div, *_ = convert_to_dst_type(div, x)
+                return x * div + sub
+            # nonzero=True: the forward mask is the output's non-zero voxels plus the recorded zeroed ones
+            mask = x != 0
+            zeroed_idx, *_ = convert_to_dst_type(zeroed_idx, mask, dtype=torch.long)
+            mask.view(-1)[zeroed_idx] = True
+            vals = x[mask]
+            sub, *_ = convert_to_dst_type(sub, vals)
+            div, *_ = convert_to_dst_type(div, vals)
+            x[mask] = vals * div + sub
+            return x
+
+        if info["channel_wise"]:
+            for i in range(len(out)):
+                out[i] = _restore(out[i], subs[i], divs[i], None if zeroed is None else zeroed[i])
+        else:
+            out = _restore(out, subs[0], divs[0], None if zeroed is None else zeroed[0])
         return out
 
 

monai/transforms/intensity/dictionary.py

@@ -20,6 +20,7 @@
 from collections.abc import Callable, Hashable, Mapping, Sequence
 
 import numpy as np
+import torch
 
 from monai.config import DtypeLike, KeysCollection
 from monai.config.type_definitions import NdarrayOrTensor
@@ -60,6 +61,7 @@
     StdShiftIntensity,
     ThresholdIntensity,
 )
+from monai.transforms.inverse import InvertibleTransform
 from monai.transforms.transform import MapTransform, RandomizableTransform
 from monai.transforms.utils import is_positive
 from monai.utils import convert_to_tensor, ensure_tuple, ensure_tuple_rep
@@ -791,11 +793,12 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
         return d
 
 
-class NormalizeIntensityd(MapTransform):
+class NormalizeIntensityd(MapTransform, InvertibleTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.NormalizeIntensity`.
     This transform can normalize only non-zero values or entire image, and can also calculate
-    mean and std on each channel separately.
+    mean and std on each channel separately. It is invertible via :meth:`inverse`;
+    see :py:class:`monai.transforms.NormalizeIntensity`.
 
     Args:
         keys: keys of the corresponding items to be transformed.
@@ -830,6 +833,25 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
             d[key] = self.normalizer(d[key])
         return d
 
+    def inverse(self, data: Mapping[Hashable, torch.Tensor]) -> dict[Hashable, torch.Tensor]:
+        """
+        Undo the normalization of every key in ``self.keys``.
+
+        Args:
+            data: dictionary whose values for ``self.keys`` were produced by this transform.
+
+        Returns:
+            a shallow copy of ``data`` with those values de-normalized.
+
+        Raises:
+            RuntimeError: propagated from :meth:`NormalizeIntensity.inverse` if the most recent applied
+                operation on a value was not made by this transform.
+        """
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.normalizer.inverse(d[key])
+        return d
+
 
 class ThresholdIntensityd(MapTransform):
     """

tests/transforms/test_normalize_intensity.py

@@ -17,7 +17,8 @@
 import torch
 from parameterized import parameterized
 
-from monai.transforms import NormalizeIntensity
+from monai.data import MetaTensor, get_track_meta, set_track_meta
+from monai.transforms import Compose, NormalizeIntensity
 from tests.test_utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose
 
 TESTS = []
@@ -138,6 +139,52 @@ def test_value_errors(self, im_type):
         with self.assertRaises(ValueError):
             normalizer(input_data)
 
+    @parameterized.expand(
+        [
+            ["global_computed", {}],
+            ["channelwise_computed", {"channel_wise": True}],
+            ["global_explicit", {"subtrahend": 2.0, "divisor": 3.0}],
+            ["channelwise_explicit", {"subtrahend": [1.0, 2.0, 3.0], "divisor": [2.0, 3.0, 4.0], "channel_wise": True}],
+            ["nonzero", {"nonzero": True}],
+            ["channelwise_nonzero", {"nonzero": True, "channel_wise": True}],
+            ["nonzero_explicit", {"nonzero": True, "subtrahend": 2.0, "divisor": 3.0}],
+        ]
+    )
+    def test_inverse(self, _, args):
+        self.addCleanup(set_track_meta, get_track_meta())
+        set_track_meta(True)
+        img = MetaTensor(torch.randn(3, 6, 6) * 5 + 2)
+        img[0, :2] = 0  # some zero voxels, which nonzero=True must leave untouched
+        img[2] = 0  # an all-zero channel, where nonzero=True has nothing to normalize
+        normalizer = NormalizeIntensity(**args)
+        out = normalizer(img.clone())
+        inv = normalizer.inverse(out)
+        assert_allclose(inv, img, type_test=False, rtol=1e-4, atol=1e-4)
+        self.assertEqual(len(inv.applied_operations), 0)
+
+    @parameterized.expand([["global", {}], ["channelwise", {"channel_wise": True}]])
+    def test_inverse_nonzero_value_equal_to_mean(self, _, args):
+        """A non-zero voxel equal to the mean becomes exactly 0 and must still be restored."""
+        self.addCleanup(set_track_meta, get_track_meta())
+        set_track_meta(True)
+        # mean of the non-zero voxels is 2 globally and in each channel, so the voxels equal to 2 become 0
+        img = MetaTensor(torch.tensor([[0.0, 1.0, 2.0, 3.0], [0.0, 0.0, 0.0, 2.0]]))
+        normalizer = NormalizeIntensity(nonzero=True, **args)
+        out = normalizer(img.clone())
+        self.assertEqual(out[0, 2].item(), 0.0)
+        self.assertEqual(out[1, 3].item(), 0.0)
+        inv = normalizer.inverse(out)
+        assert_allclose(inv, img, type_test=False, rtol=0, atol=0)
+
+    def test_inverse_nonzero_in_compose(self):
+        self.addCleanup(set_track_meta, get_track_meta())
+        set_track_meta(True)
+        img = MetaTensor(torch.randn(2, 5, 5))
+        img[0, 0] = 0
+        transform = Compose([NormalizeIntensity(nonzero=True)])
+        inv = transform.inverse(transform(img.clone()))
+        assert_allclose(inv, img, type_test=False, rtol=1e-4, atol=1e-4)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/transforms/test_normalize_intensityd.py

@@ -14,9 +14,11 @@
 import unittest
 
 import numpy as np
+import torch
 from parameterized import parameterized
 
-from monai.transforms import NormalizeIntensityd
+from monai.data import MetaTensor, get_track_meta, set_track_meta
+from monai.transforms import Compose, Invertd, NormalizeIntensityd
 from tests.test_utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose
 
 TESTS = []
@@ -76,6 +78,37 @@ def test_channel_wise(self, im_type):
         expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]])
         assert_allclose(normalized, im_type(expected), type_test="tensor")
 
+    @parameterized.expand(
+        [
+            ["global", {}],
+            ["channelwise", {"channel_wise": True}],
+            ["nonzero", {"nonzero": True}],
+            ["channelwise_nonzero", {"nonzero": True, "channel_wise": True}],
+        ]
+    )
+    def test_inverse(self, _, args):
+        self.addCleanup(set_track_meta, get_track_meta())
+        set_track_meta(True)
+        key = "img"
+        normalizer = NormalizeIntensityd(keys=key, **args)
+        data = {key: MetaTensor(torch.randn(3, 6, 6) * 4 + 1)}
+        data[key][0, :2] = 0
+        original = data[key].clone()
+        out = normalizer(dict(data))
+        inv = normalizer.inverse(out)
+        assert_allclose(inv[key], original, type_test=False, rtol=1e-4, atol=1e-4)
+
+    def test_invertd_nonzero(self):
+        self.addCleanup(set_track_meta, get_track_meta())
+        set_track_meta(True)
+        key = "img"
+        transform = Compose([NormalizeIntensityd(keys=key, nonzero=True)])
+        original = MetaTensor(torch.randn(2, 5, 5))
+        original[0, 0] = 0
+        out = transform({key: original.clone()})
+        inv = Invertd(keys=key, transform=transform, orig_keys=key)(out)
+        assert_allclose(inv[key], original, type_test=False, rtol=1e-4, atol=1e-4)
+
 
 if __name__ == "__main__":
     unittest.main()