diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py index a80f94727e..d7b405d096 100644 --- a/monai/transforms/intensity/array.py +++ b/monai/transforms/intensity/array.py @@ -1055,7 +1055,7 @@ def _restore(x, sub, div, zeroed_idx=None): 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 + mask = (x != 0).contiguous() zeroed_idx, *_ = convert_to_dst_type(zeroed_idx, mask, dtype=torch.long) mask.view(-1)[zeroed_idx] = True vals = x[mask] diff --git a/tests/transforms/test_normalize_intensity.py b/tests/transforms/test_normalize_intensity.py index 53fea19d33..49d9491ac0 100644 --- a/tests/transforms/test_normalize_intensity.py +++ b/tests/transforms/test_normalize_intensity.py @@ -176,6 +176,23 @@ def test_inverse_nonzero_value_equal_to_mean(self, _, args): inv = normalizer.inverse(out) assert_allclose(inv, img, type_test=False, rtol=0, atol=0) + @parameterized.expand([[False], [True]]) + def test_inverse_nonzero_noncontiguous(self, channel_wise): + """Restore transposed inputs, including non-zero voxels normalized to zero. + + Args: + channel_wise: whether each channel is normalized independently. + """ + self.addCleanup(set_track_meta, get_track_meta()) + set_track_meta(True) + img = MetaTensor(torch.tensor([[[0.0, 1.0], [2.0, 3.0]], [[0.0, 0.0], [0.0, 2.0]]])).transpose(1, 2) + normalizer = NormalizeIntensity(nonzero=True, channel_wise=channel_wise) + out = normalizer(img.clone()) + self.assertFalse(out.is_contiguous()) + inv = normalizer.inverse(out) + assert_allclose(inv, img, type_test=False, rtol=1e-6, atol=1e-6) + self.assertEqual(len(inv.applied_operations), 0) + def test_inverse_nonzero_in_compose(self): self.addCleanup(set_track_meta, get_track_meta()) set_track_meta(True) diff --git a/tests/transforms/test_normalize_intensityd.py b/tests/transforms/test_normalize_intensityd.py index 50a9c302db..992f8db67c 100644 --- a/tests/transforms/test_normalize_intensityd.py +++ b/tests/transforms/test_normalize_intensityd.py @@ -98,6 +98,23 @@ def test_inverse(self, _, args): inv = normalizer.inverse(out) assert_allclose(inv[key], original, type_test=False, rtol=1e-4, atol=1e-4) + @parameterized.expand([[False], [True]]) + def test_inverse_nonzero_noncontiguous(self, channel_wise): + """Restore transposed dictionary values without losing their zeroed voxels. + + Args: + channel_wise: whether each channel is normalized independently. + """ + self.addCleanup(set_track_meta, get_track_meta()) + set_track_meta(True) + img = MetaTensor(torch.tensor([[[0.0, 1.0], [2.0, 3.0]], [[0.0, 0.0], [0.0, 2.0]]])).transpose(1, 2) + normalizer = NormalizeIntensityd(keys="img", nonzero=True, channel_wise=channel_wise) + out = normalizer({"img": img.clone()}) + self.assertFalse(out["img"].is_contiguous()) + inv = normalizer.inverse(out) + assert_allclose(inv["img"], img, type_test=False, rtol=1e-6, atol=1e-6) + self.assertEqual(len(inv["img"].applied_operations), 0) + def test_invertd_nonzero(self): self.addCleanup(set_track_meta, get_track_meta()) set_track_meta(True)