Make NormalizeIntensity.inverse work with nonzero=True

With nonzero=True the inverse raised NotImplementedError, which made
Compose.inverse and Invertd fail for any pipeline using the transform.

Voxels that were zero on the forward pass are left at zero, so the
forward mask can be rebuilt from the normalized image as `out != 0`.
The only ambiguity is a non-zero voxel whose value equals the
subtrahend: it becomes exactly zero. _normalize now records the flat
indices of those voxels (usually an empty tensor) in the transform
meta information as `zeroed_idx`, and inverse() restores
`out * div + sub` on the rebuilt mask. Inversion is exact, including
the value-equals-mean case.

Also: all-zero inputs store identity stats (0.0/1.0) instead of None,
so extra_info stays collate-safe; NormalizeIntensityd.inverse is typed
on torch.Tensor to satisfy mypy; Google-style docstrings for the new
methods; tests restore the track_meta state they change.

Signed-off-by: Azra Bano <azrabano.work@gmail.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: azrabano23

monai/transforms/intensity/array.py

@@ -33,7 +33,7 @@
 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.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
@@ -847,8 +847,10 @@ class NormalizeIntensity(InvertibleTransform):
 
     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``. Inversion is not supported when ``nonzero=True``, because the
-    zero-voxel mask would be required to reverse the operation exactly.
+    ``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).
@@ -890,13 +892,28 @@ def _std(x):
         return x.item() if x.numel() == 1 else x
 
     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, None, None
+                # 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
@@ -917,12 +934,16 @@ def _normalize(self, img: NdarrayOrTensor, sub=None, div=None):
                 _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 the subtrahend/divisor actually used so the transform can be inverted.
-        return img, _sub, _div
+        return img, _sub, _div, zeroed_idx
 
     def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
@@ -931,9 +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().
+        # 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.")
@@ -944,31 +967,55 @@ 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], _sub, _div = 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, _sub, _div = 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]
-        out = self._push_transform_with_stats(out, subs, divs)
-        return out
+        return self._push_transform_with_stats(out, subs, divs, zeroed)
 
-    def _to_storable(self, value):
-        """Convert a subtrahend/divisor to something storable in transform meta."""
+    @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, subs: list, divs: list):
+    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 = {
@@ -977,33 +1024,51 @@ def _push_transform_with_stats(self, out, subs: list, divs: list):
             "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]
-        if info["nonzero"]:
-            raise NotImplementedError(
-                "NormalizeIntensity.inverse is not supported when nonzero=True, because the "
-                "zero-voxel mask is needed to reverse the normalization exactly."
-            )
         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):
-            sub, *_ = convert_to_dst_type(sub, x)
-            div, *_ = convert_to_dst_type(div, x)
-            return x * div + sub
+        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)):
-                if subs[i] is None or divs[i] is None:  # all-zero channel skipped on the forward pass
-                    continue
-                out[i] = _restore(out[i], subs[i], divs[i])
+                out[i] = _restore(out[i], subs[i], divs[i], None if zeroed is None else zeroed[i])
         else:
-            if subs[0] is not None and divs[0] is not None:
-                out = _restore(out, subs[0], divs[0])
+            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
@@ -796,8 +797,8 @@ 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. It is invertible via :meth:`inverse` (except when
-    ``nonzero=True``); see :py:class:`monai.transforms.NormalizeIntensity`.
+    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.
@@ -832,7 +833,20 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, N
             d[key] = self.normalizer(d[key])
         return d
 
-    def inverse(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
+    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])

tests/transforms/test_normalize_intensity.py

@@ -17,8 +17,8 @@
 import torch
 from parameterized import parameterized
 
-from monai.data import MetaTensor, set_track_meta
-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 = []
@@ -145,23 +145,45 @@ def test_value_errors(self, im_type):
             ["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)
 
-    def test_inverse_nonzero_not_implemented(self):
+    @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)
-        img = MetaTensor(torch.randn(2, 5, 5))
-        normalizer = NormalizeIntensity(nonzero=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())
-        with self.assertRaises(NotImplementedError):
-            normalizer.inverse(out)
+        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__":

tests/transforms/test_normalize_intensityd.py

@@ -17,8 +17,8 @@
 import torch
 from parameterized import parameterized
 
-from monai.data import MetaTensor, set_track_meta
-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 = []
@@ -78,17 +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}]])
+    @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()