From 2b7ea77b2cc07157a2c466b88945b669729f2a5c Mon Sep 17 00:00:00 2001 From: azrabano23 Date: Sun, 7 Jun 2026 23:35:02 -0700 Subject: [PATCH 1/2] Make NormalizeIntensity invertible (#5647) Signed-off-by: azrabano23 --- monai/transforms/intensity/array.py | 76 +++++++++++++++++-- monai/transforms/intensity/dictionary.py | 12 ++- tests/transforms/test_normalize_intensity.py | 25 ++++++ tests/transforms/test_normalize_intensityd.py | 13 ++++ 4 files changed, 117 insertions(+), 9 deletions(-) diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index 23a57ae9fbe..3f4c8cd4a46 100644 --- a/monai/transforms/intensity/array.py +++ b/monai/transforms/intensity/array.py @@ -26,13 +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 MetaTensor 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.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 @@ -836,7 +838,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. @@ -846,6 +848,11 @@ 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``. Inversion is not supported when ``nonzero=True``, because the + zero-voxel mask would be required to reverse the operation exactly. + Args: subtrahend: the amount to subtract by (usually the mean). divisor: the amount to divide by (usually the standard deviation). @@ -885,14 +892,14 @@ 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): img, *_ = convert_data_type(img, dtype=torch.float32) if self.nonzero: slices = img != 0 masked_img = img[slices] if not slices.any(): - return img + return img, None, None else: slices = None masked_img = img @@ -917,7 +924,8 @@ def _normalize(self, img: NdarrayOrTensor, sub=None, div=None) -> NdarrayOrTenso img[slices] = (masked_img - _sub) / _div else: img = (img - _sub) / _div - return img + # Return the subtrahend/divisor actually used so the transform can be inverted. + return img, _sub, _div def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor: """ @@ -926,6 +934,9 @@ 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(). + subs: list = [] + divs: 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.") @@ -936,15 +947,66 @@ 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 = 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) else: - img_t = self._normalize(img_t, self.subtrahend, self.divisor) # type: ignore[assignment] + img_t, _sub, _div = self._normalize(img_t, self.subtrahend, self.divisor) # type: ignore[assignment] + subs.append(_sub) + divs.append(_div) out = convert_to_dst_type(img_t, img_t, dtype=dtype)[0] + out = self._push_transform_with_stats(out, subs, divs) + return out + + def _to_storable(self, value): + """Convert a subtrahend/divisor to something storable in transform meta.""" + 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): + 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, + } + self.push_transform(out, extra_info=extra_info) + return out + + def inverse(self, data: torch.Tensor) -> torch.Tensor: + 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"] + 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 + + 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]) + else: + if subs[0] is not None and divs[0] is not None: + out = _restore(out, subs[0], divs[0]) return out diff --git a/monai/transforms/intensity/dictionary.py b/monai/transforms/intensity/dictionary.py index 0c25d4ac994..4fa5f0e2da8 100644 --- a/monai/transforms/intensity/dictionary.py +++ b/monai/transforms/intensity/dictionary.py @@ -60,6 +60,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 +792,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` (except when + ``nonzero=True``); see :py:class:`monai.transforms.NormalizeIntensity`. Args: keys: keys of the corresponding items to be transformed. @@ -830,6 +832,12 @@ 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]: + d = dict(data) + for key in self.key_iterator(d): + d[key] = self.normalizer.inverse(d[key]) + return d + class ThresholdIntensityd(MapTransform): """ diff --git a/tests/transforms/test_normalize_intensity.py b/tests/transforms/test_normalize_intensity.py index c58bc587f2b..ee0295a1e22 100644 --- a/tests/transforms/test_normalize_intensity.py +++ b/tests/transforms/test_normalize_intensity.py @@ -17,6 +17,7 @@ import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import NormalizeIntensity from tests.test_utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose @@ -138,6 +139,30 @@ 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}], + ] + ) + def test_inverse(self, _, args): + set_track_meta(True) + img = MetaTensor(torch.randn(3, 6, 6) * 5 + 2) + normalizer = NormalizeIntensity(**args) + out = normalizer(img.clone()) + inv = normalizer.inverse(out) + assert_allclose(inv, img, type_test=False, rtol=1e-4, atol=1e-4) + + def test_inverse_nonzero_not_implemented(self): + set_track_meta(True) + img = MetaTensor(torch.randn(2, 5, 5)) + normalizer = NormalizeIntensity(nonzero=True) + out = normalizer(img.clone()) + with self.assertRaises(NotImplementedError): + normalizer.inverse(out) + if __name__ == "__main__": unittest.main() diff --git a/tests/transforms/test_normalize_intensityd.py b/tests/transforms/test_normalize_intensityd.py index b8e4c7bca88..e11cce1c620 100644 --- a/tests/transforms/test_normalize_intensityd.py +++ b/tests/transforms/test_normalize_intensityd.py @@ -14,8 +14,10 @@ import unittest import numpy as np +import torch from parameterized import parameterized +from monai.data import MetaTensor, set_track_meta from monai.transforms import NormalizeIntensityd from tests.test_utils import TEST_NDARRAYS, NumpyImageTestCase2D, assert_allclose @@ -76,6 +78,17 @@ 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}]]) + def test_inverse(self, _, args): + set_track_meta(True) + key = "img" + normalizer = NormalizeIntensityd(keys=key, **args) + data = {key: MetaTensor(torch.randn(3, 6, 6) * 4 + 1)} + 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) + if __name__ == "__main__": unittest.main() From 8e53cfb772c145e56aa9206efc90ca26d0849d7c Mon Sep 17 00:00:00 2001 From: Azra Bano Date: Fri, 4 Sep 2026 11:30:49 -0400 Subject: [PATCH 2/2] 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 Co-Authored-By: Claude Fable 5.1 --- monai/transforms/intensity/array.py | 119 ++++++++++++++---- monai/transforms/intensity/dictionary.py | 20 ++- tests/transforms/test_normalize_intensity.py | 36 ++++-- tests/transforms/test_normalize_intensityd.py | 26 +++- 4 files changed, 161 insertions(+), 40 deletions(-) diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index 35ea7944f8a..a80f94727ee 100644 --- a/monai/transforms/intensity/array.py +++ b/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 diff --git a/monai/transforms/intensity/dictionary.py b/monai/transforms/intensity/dictionary.py index 4fa5f0e2da8..a91ad4f59c6 100644 --- a/monai/transforms/intensity/dictionary.py +++ b/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]) diff --git a/tests/transforms/test_normalize_intensity.py b/tests/transforms/test_normalize_intensity.py index ee0295a1e22..53fea19d332 100644 --- a/tests/transforms/test_normalize_intensity.py +++ b/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__": diff --git a/tests/transforms/test_normalize_intensityd.py b/tests/transforms/test_normalize_intensityd.py index e11cce1c620..50a9c302db0 100644 --- a/tests/transforms/test_normalize_intensityd.py +++ b/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()