diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index a23c867d60..a80f94727e 100644 --- a/monai/transforms/intensity/array.py +++ b/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 diff --git a/monai/transforms/intensity/dictionary.py b/monai/transforms/intensity/dictionary.py index 0c25d4ac99..a91ad4f59c 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 @@ -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): """ diff --git a/tests/transforms/test_normalize_intensity.py b/tests/transforms/test_normalize_intensity.py index c58bc587f2..53fea19d33 100644 --- a/tests/transforms/test_normalize_intensity.py +++ b/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() diff --git a/tests/transforms/test_normalize_intensityd.py b/tests/transforms/test_normalize_intensityd.py index b8e4c7bca8..50a9c302db 100644 --- a/tests/transforms/test_normalize_intensityd.py +++ b/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()