Repository: Project-MONAI/MONAI

Issue:
divide std in NormalizeIntensity, but std could be 0
**Describe the bug**
In monai/transforms/intensity/array.py,
img[slices] = (img[slices] - np.mean(img[slices])) / np.std(img[slices])
But the np.std(img[slices]) could be 0 in some cases.
This will cause NaN.

**To Reproduce**
Steps to reproduce the behavior:

Let the input image to be all zeros

**Expected behavior**
The result is NaN

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environment (please complete the following information; e.g. using `sh runtests.sh -v`):**
 - OS
 - Python version
 - MONAI version [e.g. git commit hash]
 - CUDA/cuDNN version
 - GPU models and configuration

**Additional context**
if np.std(img[slices])==0:
img[slices] = (img[slices] - np.mean(img[slices])) 




Discussion:

Patch:
diff --git a/monai/transforms/intensity/array.py b/monai/transforms/intensity/array.py
--- a/monai/transforms/intensity/array.py
+++ b/monai/transforms/intensity/array.py
@@ -185,11 +185,11 @@ def __call__(self, img: np.ndarray) -> np.ndarray:
 
 class NormalizeIntensity(Transform):
     """
-    Normalize input based on provided args, using calculated mean and std if not provided
-    (shape of subtrahend and divisor must match. if 0, entire volume uses same subtrahend and
-    divisor, otherwise the shape can have dimension 1 for channels).
+    Normalize input based on provided args, using calculated mean and std if not provided.
     This transform can normalize only non-zero values or entire image, and can also calculate
     mean and std on each channel separately.
+    When `channel_wise` is True, the first dimension of `subtrahend` and `divisor` should
+    be the number of image channels if they are not None.
 
     Args:
         subtrahend: the amount to subtract by (usually the mean).
@@ -201,27 +201,33 @@ class NormalizeIntensity(Transform):
 
     def __init__(
         self,
-        subtrahend: Optional[np.ndarray] = None,
-        divisor: Optional[np.ndarray] = None,
+        subtrahend: Optional[Sequence] = None,
+        divisor: Optional[Sequence] = None,
         nonzero: bool = False,
         channel_wise: bool = False,
     ) -> None:
-        if subtrahend is not None or divisor is not None:
-            assert isinstance(subtrahend, np.ndarray) and isinstance(
-                divisor, np.ndarray
-            ), "subtrahend and divisor must be set in pair and in numpy array."
         self.subtrahend = subtrahend
         self.divisor = divisor
         self.nonzero = nonzero
         self.channel_wise = channel_wise
 
-    def _normalize(self, img: np.ndarray) -> np.ndarray:
+    def _normalize(self, img: np.ndarray, sub=None, div=None) -> np.ndarray:
         slices = (img != 0) if self.nonzero else np.ones(img.shape, dtype=np.bool_)
-        if np.any(slices):
-            if self.subtrahend is not None and self.divisor is not None:
-                img[slices] = (img[slices] - self.subtrahend[slices]) / self.divisor[slices]
-            else:
-                img[slices] = (img[slices] - np.mean(img[slices])) / np.std(img[slices])
+        if not np.any(slices):
+            return img
+
+        _sub = sub if sub is not None else np.mean(img[slices])
+        if isinstance(_sub, np.ndarray):
+            _sub = _sub[slices]
+
+        _div = div if div is not None else np.std(img[slices])
+        if np.isscalar(_div):
+            if _div == 0.0:
+                _div = 1.0
+        elif isinstance(_div, np.ndarray):
+            _div = _div[slices]
+            _div[_div == 0.0] = 1.0
+        img[slices] = (img[slices] - _sub) / _div
         return img
 
     def __call__(self, img: np.ndarray) -> np.ndarray:
@@ -229,10 +235,19 @@ def __call__(self, img: np.ndarray) -> np.ndarray:
         Apply the transform to `img`, assuming `img` is a channel-first array if `self.channel_wise` is True,
         """
         if self.channel_wise:
+            if self.subtrahend is not None and len(self.subtrahend) != len(img):
+                raise ValueError(f"img has {len(img)} channels, but subtrahend has {len(self.subtrahend)} components.")
+            if self.divisor is not None and len(self.divisor) != len(img):
+                raise ValueError(f"img has {len(img)} channels, but divisor has {len(self.divisor)} components.")
+
             for i, d in enumerate(img):
-                img[i] = self._normalize(d)
+                img[i] = self._normalize(
+                    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,
+                )
         else:
-            img = self._normalize(img)
+            img = self._normalize(img, self.subtrahend, self.divisor)
 
         return img
 


Test patch:
diff --git a/tests/test_normalize_intensity.py b/tests/test_normalize_intensity.py
--- a/tests/test_normalize_intensity.py
+++ b/tests/test_normalize_intensity.py
@@ -17,16 +17,43 @@
 from monai.transforms import NormalizeIntensity
 from tests.utils import NumpyImageTestCase2D
 
-TEST_CASE_1 = [{"nonzero": True}, np.array([0.0, 3.0, 0.0, 4.0]), np.array([0.0, -1.0, 0.0, 1.0])]
-
-TEST_CASE_2 = [
-    {"subtrahend": np.array([3.5, 3.5, 3.5, 3.5]), "divisor": np.array([0.5, 0.5, 0.5, 0.5]), "nonzero": True},
-    np.array([0.0, 3.0, 0.0, 4.0]),
-    np.array([0.0, -1.0, 0.0, 1.0]),
+TEST_CASES = [
+    [{"nonzero": True}, np.array([0.0, 3.0, 0.0, 4.0]), np.array([0.0, -1.0, 0.0, 1.0])],
+    [
+        {"subtrahend": np.array([3.5, 3.5, 3.5, 3.5]), "divisor": np.array([0.5, 0.5, 0.5, 0.5]), "nonzero": True},
+        np.array([0.0, 3.0, 0.0, 4.0]),
+        np.array([0.0, -1.0, 0.0, 1.0]),
+    ],
+    [{"nonzero": True}, np.array([0.0, 0.0, 0.0, 0.0]), np.array([0.0, 0.0, 0.0, 0.0])],
+    [{"nonzero": False}, np.array([0.0, 0.0, 0.0, 0.0]), np.array([0.0, 0.0, 0.0, 0.0])],
+    [{"nonzero": False}, np.array([1, 1, 1, 1]), np.array([0.0, 0.0, 0.0, 0.0])],
+    [
+        {"nonzero": False, "channel_wise": True, "subtrahend": [1, 2, 3]},
+        np.ones((3, 2, 2)),
+        np.array([[[0.0, 0.0], [0.0, 0.0]], [[-1.0, -1.0], [-1.0, -1.0]], [[-2.0, -2.0], [-2.0, -2.0]]]),
+    ],
+    [
+        {"nonzero": True, "channel_wise": True, "subtrahend": [1, 2, 3], "divisor": [0, 0, 2]},
+        np.ones((3, 2, 2)),
+        np.array([[[0.0, 0.0], [0.0, 0.0]], [[-1.0, -1.0], [-1.0, -1.0]], [[-1.0, -1.0], [-1.0, -1.0]]]),
+    ],
+    [
+        {"nonzero": True, "channel_wise": False, "subtrahend": 2, "divisor": 0},
+        np.ones((3, 2, 2)),
+        np.ones((3, 2, 2)) * -1.0,
+    ],
+    [
+        {"nonzero": True, "channel_wise": False, "subtrahend": np.ones((3, 2, 2)) * 0.5, "divisor": 0},
+        np.ones((3, 2, 2)),
+        np.ones((3, 2, 2)) * 0.5,
+    ],
+    [
+        {"nonzero": True, "channel_wise": True, "subtrahend": np.ones((3, 2, 2)) * 0.5, "divisor": [0, 1, 0]},
+        np.ones((3, 2, 2)),
+        np.ones((3, 2, 2)) * 0.5,
+    ],
 ]
 
-TEST_CASE_3 = [{"nonzero": True}, np.array([0.0, 0.0, 0.0, 0.0]), np.array([0.0, 0.0, 0.0, 0.0])]
-
 
 class TestNormalizeIntensity(NumpyImageTestCase2D):
     def test_default(self):
@@ -35,7 +62,7 @@ def test_default(self):
         expected = (self.imt - np.mean(self.imt)) / np.std(self.imt)
         np.testing.assert_allclose(normalized, expected, rtol=1e-6)
 
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3])
+    @parameterized.expand(TEST_CASES)
     def test_nonzero(self, input_param, input_data, expected_data):
         normalizer = NormalizeIntensity(**input_param)
         np.testing.assert_allclose(expected_data, normalizer(input_data))
@@ -46,6 +73,15 @@ def test_channel_wise(self):
         expected = np.array([[0.0, -1.0, 0.0, 1.0], [0.0, -1.0, 0.0, 1.0]])
         np.testing.assert_allclose(expected, normalizer(input_data))
 
+    def test_value_errors(self):
+        input_data = np.array([[0.0, 3.0, 0.0, 4.0], [0.0, 4.0, 0.0, 5.0]])
+        normalizer = NormalizeIntensity(nonzero=True, channel_wise=True, subtrahend=[1])
+        with self.assertRaises(ValueError):
+            normalizer(input_data)
+        normalizer = NormalizeIntensity(nonzero=True, channel_wise=True, subtrahend=[1, 2], divisor=[1])
+        with self.assertRaises(ValueError):
+            normalizer(input_data)
+
 
 if __name__ == "__main__":
     unittest.main()