Repository: Project-MONAI/MONAI

Issue:
add StatsTransform to debug data in transform chain
**Is your feature request related to a problem? Please describe.**
Currently, our transforms compose a chain and hard to know the data after every transform and debug, need to add StatsTransform which can be inserted into the chain to print expected info.


Discussion:
Nic-Ma: /black

Patch:
diff --git a/monai/transforms/utility/array.py b/monai/transforms/utility/array.py
index c64ae164..861907d3 100644
--- a/monai/transforms/utility/array.py
+++ b/monai/transforms/utility/array.py
@@ -208,17 +208,18 @@ class DataStats(Transform):
         if logger_handler is not None:
             self._logger.addHandler(logger_handler)
 
-    def __call__(self, img):
-        lines = [f"{self.prefix} statistics:"]
+    def __call__(self, img, prefix=None, data_shape=None, intensity_range=None, data_value=None, additional_info=None):
+        lines = [f"{prefix or self.prefix} statistics:"]
 
-        if self.data_shape:
+        if self.data_shape if data_shape is None else data_shape:
             lines.append(f"Shape: {img.shape}")
-        if self.intensity_range:
+        if self.intensity_range if intensity_range is None else intensity_range:
             lines.append(f"Intensity range: ({np.min(img)}, {np.max(img)})")
-        if self.data_value:
+        if self.data_value if data_value is None else data_value:
             lines.append(f"Value: {img}")
-        if self.additional_info:
-            lines.append(f"Additional info: {self.additional_info(img)}")
+        additional_info = self.additional_info if additional_info is None else additional_info
+        if additional_info is not None:
+            lines.append(f"Additional info: {additional_info(img)}")
         separator = "\n"
         self.output = f"{separator.join(lines)}"
         self._logger.debug(self.output)
diff --git a/monai/transforms/utility/dictionary.py b/monai/transforms/utility/dictionary.py
index 2553690c..fc59e911 100644
--- a/monai/transforms/utility/dictionary.py
+++ b/monai/transforms/utility/dictionary.py
@@ -222,22 +222,36 @@ class DataStatsd(MapTransform):
         Args:
             keys (hashable items): keys of the corresponding items to be transformed.
                 See also: :py:class:`monai.transforms.compose.MapTransform`
-            prefix (string): will be printed in format: "{prefix} statistics".
-            data_shape (bool): whether to show the shape of input data.
-            intensity_range (bool): whether to show the intensity value range of input data.
-            data_value (bool): whether to show the raw value of input data.
+            prefix (string or list of string): will be printed in format: "{prefix} statistics".
+            data_shape (bool or list of bool): whether to show the shape of input data.
+            intensity_range (bool or list of bool): whether to show the intensity value range of input data.
+            data_value (bool or list of bool): whether to show the raw value of input data.
                 a typical example is to print some properties of Nifti image: affine, pixdim, etc.
-            additional_info (Callable): user can define callable function to extract additional info from input data.
+            additional_info (Callable or list of Callable): user can define callable function to extract
+                additional info from input data.
             logger_handler (logging.handler): add additional handler to output data: save to file, etc.
                 add existing python logging handlers: https://docs.python.org/3/library/logging.handlers.html
         """
         super().__init__(keys)
-        self.printer = DataStats(prefix, data_shape, intensity_range, data_value, additional_info, logger_handler)
+        self.prefix = ensure_tuple_rep(prefix, len(self.keys))
+        self.data_shape = ensure_tuple_rep(data_shape, len(self.keys))
+        self.intensity_range = ensure_tuple_rep(intensity_range, len(self.keys))
+        self.data_value = ensure_tuple_rep(data_value, len(self.keys))
+        self.additional_info = ensure_tuple_rep(additional_info, len(self.keys))
+        self.logger_handler = logger_handler
+        self.printer = DataStats(logger_handler=logger_handler)
 
     def __call__(self, data):
         d = dict(data)
-        for key in self.keys:
-            d[key] = self.printer(d[key])
+        for idx, key in enumerate(self.keys):
+            d[key] = self.printer(
+                d[key],
+                self.prefix[idx],
+                self.data_shape[idx],
+                self.intensity_range[idx],
+                self.data_value[idx],
+                self.additional_info[idx],
+            )
         return d
 
 


Test patch:
diff --git a/tests/test_data_statsd.py b/tests/test_data_statsd.py
index 9466584a..e26dfdb7 100644
--- a/tests/test_data_statsd.py
+++ b/tests/test_data_statsd.py
@@ -83,19 +83,32 @@ TEST_CASE_5 = [
 ]
 
 TEST_CASE_6 = [
+    {
+        "keys": ("img", "affine"),
+        "prefix": ("image", "affine"),
+        "data_shape": True,
+        "intensity_range": (True, False),
+        "data_value": (False, True),
+        "additional_info": (lambda x: np.mean(x), None),
+    },
+    {"img": np.array([[0, 1], [1, 2]]), "affine": np.eye(2, 2)},
+    "affine statistics:\nShape: (2, 2)\nValue: [[1. 0.]\n [0. 1.]]",
+]
+
+TEST_CASE_7 = [
     {"img": np.array([[0, 1], [1, 2]])},
     "test data statistics:\nShape: (2, 2)\nIntensity range: (0, 2)\nValue: [[0 1]\n [1 2]]\nAdditional info: 1.0\n",
 ]
 
 
 class TestDataStatsd(unittest.TestCase):
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6])
     def test_value(self, input_param, input_data, expected_print):
         transform = DataStatsd(**input_param)
         _ = transform(input_data)
         self.assertEqual(transform.printer.output, expected_print)
 
-    @parameterized.expand([TEST_CASE_6])
+    @parameterized.expand([TEST_CASE_7])
     def test_file(self, input_data, expected_print):
         with tempfile.TemporaryDirectory() as tempdir:
             filename = os.path.join(tempdir, "test_stats.log")