Make MetaTensor optional printed in DataStats and DataStatsd #5905

monai/transforms/utility/array.py

@@ -654,6 +654,7 @@ def __init__(
         data_shape: bool = True,
         value_range: bool = True,
         data_value: bool = False,
+        meta_tensor: bool = False,
         additional_info: Callable | None = None,
         name: str = "DataStats",
     ) -> None:
@@ -665,6 +666,7 @@ def __init__(
             value_range: whether to show the value range of input data.
             data_value: whether to show the raw value of input data.
                 a typical example is to print some properties of Nifti image: affine, pixdim, etc.
+            meta_tensor: whether to show the data of MetaTensor.
             additional_info: user can define callable function to extract additional info from input data.
             name: identifier of `logging.logger` to use, defaulting to "DataStats".
 
@@ -679,6 +681,7 @@ def __init__(
         self.data_shape = data_shape
         self.value_range = value_range
         self.data_value = data_value
+        self.meta_tensor = meta_tensor
         if additional_info is not None and not callable(additional_info):
             raise TypeError(f"additional_info must be None or callable but is {type(additional_info).__name__}.")
         self.additional_info = additional_info
@@ -705,6 +708,7 @@ def __call__(
         data_shape: bool | None = None,
         value_range: bool | None = None,
         data_value: bool | None = None,
+        meta_tensor: bool | None = None,
         additional_info: Callable | None = None,
     ) -> NdarrayOrTensor:
         """
@@ -725,6 +729,8 @@ def __call__(
                 lines.append(f"Value range: (not a PyTorch or Numpy array, type: {type(img)})")
         if self.data_value if data_value is None else data_value:
             lines.append(f"Value: {img}")
+        if self.meta_tensor if meta_tensor is None else meta_tensor:
+            lines.append(f"MetaTensor: {MetaTensor(img).meta.__repr__()}")
         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)}")

monai/transforms/utility/dictionary.py

@@ -789,6 +789,7 @@ def __init__(
         data_shape: Sequence[bool] | bool = True,
         value_range: Sequence[bool] | bool = True,
         data_value: Sequence[bool] | bool = False,
+        meta_tensor: Sequence[bool] | bool = False,
         additional_info: Sequence[Callable] | Callable | None = None,
         name: str = "DataStats",
         allow_missing_keys: bool = False,
@@ -808,6 +809,8 @@ def __init__(
             data_value: whether to show the raw value of input data.
                 it also can be a sequence of bool, each element corresponds to a key in ``keys``.
                 a typical example is to print some properties of Nifti image: affine, pixdim, etc.
+            meta_tensor: whether to show the data of MetaTensor.
+                it also can be a sequence of bool, each element corresponds to a key in ``keys``.
             additional_info: user can define callable function to extract
                 additional info from input data. it also can be a sequence of string, each element
                 corresponds to a key in ``keys``.
@@ -821,15 +824,16 @@ def __init__(
         self.data_shape = ensure_tuple_rep(data_shape, len(self.keys))
         self.value_range = ensure_tuple_rep(value_range, len(self.keys))
         self.data_value = ensure_tuple_rep(data_value, len(self.keys))
+        self.meta_tensor = ensure_tuple_rep(meta_tensor, len(self.keys))
         self.additional_info = ensure_tuple_rep(additional_info, len(self.keys))
         self.printer = DataStats(name=name)
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> dict[Hashable, NdarrayOrTensor]:
         d = dict(data)
-        for key, prefix, data_type, data_shape, value_range, data_value, additional_info in self.key_iterator(
-            d, self.prefix, self.data_type, self.data_shape, self.value_range, self.data_value, self.additional_info
+        for key, prefix, data_type, data_shape, value_range, data_value, meta_tensor, additional_info in self.key_iterator(
+            d, self.prefix, self.data_type, self.data_shape, self.value_range, self.data_value, self.meta_tensor, self.additional_info
         ):
-            d[key] = self.printer(d[key], prefix, data_type, data_shape, value_range, data_value, additional_info)
+            d[key] = self.printer(d[key], prefix, data_type, data_shape, value_range, data_value, meta_tensor, additional_info)
         return d
 
 

tests/test_data_stats.py

@@ -24,6 +24,7 @@
 from parameterized import parameterized
 
 from monai.transforms import DataStats
+from monai.data import MetaTensor
 
 TEST_CASE_1 = [
     {
@@ -130,20 +131,50 @@
 ]
 
 TEST_CASE_8 = [
+    {
+        "prefix": "test data",
+        "data_type": True,
+        "data_shape": True,
+        "value_range": True,
+        "data_value": True,
+        "additional_info": np.mean,
+        "name": "DataStats",
+    },
     np.array([[0, 1], [1, 2]]),
     "test data statistics:\nType: <class 'numpy.ndarray'> int64\nShape: (2, 2)\nValue range: (0, 2)\n"
     "Value: [[0 1]\n [1 2]]\nAdditional info: 1.0\n",
 ]
 
+TEST_CASE_9 = [
+    np.array([[0, 1], [1, 2]]),
+    "test data statistics:\nType: <class 'numpy.ndarray'> int64\nShape: (2, 2)\nValue range: (0, 2)\n"
+    "Value: [[0 1]\n [1 2]]\n"
+    "MetaTensor: {affine: tensor([[1., 0., 0., 0.],\n"
+    "        [0., 1., 0., 0.],\n"
+    "        [0., 0., 1., 0.],\n"
+    "        [0., 0., 0., 1.]], dtype=torch.float64), space: RAS}\n"
+    "Additional info: 1.0\n",
+]
+
+TEST_CASE_10 = [
+    MetaTensor(torch.tensor([[0, 1], [1, 2]]), affine=torch.as_tensor([[2,0,0,0],[0,2,0,0],[0,0,2,0],[0,0,0,1]], dtype=torch.float64), meta={"some": "info"}),
+    "test data statistics:\nType: <class 'monai.data.meta_tensor.MetaTensor'> torch.int64\nShape: torch.Size([2, 2])\nValue range: (0, 2)\n"
+    "Value: tensor([[0, 1],\n        [1, 2]])\n"
+    "MetaTensor: {'some': 'info', affine: tensor([[2., 0., 0., 0.],\n"
+    "        [0., 2., 0., 0.],\n"
+    "        [0., 0., 2., 0.],\n"
+    "        [0., 0., 0., 1.]], dtype=torch.float64), space: RAS}\n"
+    "Additional info: 1.0\n",
+]
 
 class TestDataStats(unittest.TestCase):
 
-    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7])
+    @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8])
     def test_value(self, input_param, input_data, expected_print):
         transform = DataStats(**input_param)
         _ = transform(input_data)
 
-    @parameterized.expand([TEST_CASE_8])
+    @parameterized.expand([TEST_CASE_9, TEST_CASE_10])
     def test_file(self, input_data, expected_print):
         with tempfile.TemporaryDirectory() as tempdir:
             filename = os.path.join(tempdir, "test_data_stats.log")
@@ -158,6 +189,7 @@ def test_file(self, input_data, expected_print):
                 "data_shape": True,
                 "value_range": True,
                 "data_value": True,
+                "meta_tensor": True,
                 "additional_info": np.mean,
                 "name": name,
             }

tests/test_data_statsd.py

@@ -22,6 +22,7 @@
 from parameterized import parameterized
 
 from monai.transforms import DataStatsd
+from monai.data import MetaTensor
 
 TEST_CASE_1 = [
     {
@@ -150,22 +151,53 @@
 ]
 
 TEST_CASE_9 = [
+    {
+        "keys": "img",
+        "prefix": "test data",
+        "data_shape": True,
+        "value_range": True,
+        "data_value": True,
+        "meta_tensor": True,
+        "additional_info": np.mean,
+        "name": "DataStats",
+    },
     {"img": np.array([[0, 1], [1, 2]])},
     "test data statistics:\nType: <class 'numpy.ndarray'> int64\nShape: (2, 2)\nValue range: (0, 2)\n"
     "Value: [[0 1]\n [1 2]]\nAdditional info: 1.0\n",
 ]
 
+TEST_CASE_10 = [
+    {"img": np.array([[0, 1], [1, 2]])},
+    "test data statistics:\nType: <class 'numpy.ndarray'> int64\nShape: (2, 2)\nValue range: (0, 2)\n"
+    "Value: [[0 1]\n [1 2]]\n"
+    "MetaTensor: {affine: tensor([[1., 0., 0., 0.],\n"
+    "        [0., 1., 0., 0.],\n"
+    "        [0., 0., 1., 0.],\n"
+    "        [0., 0., 0., 1.]], dtype=torch.float64), space: RAS}\n"
+    "Additional info: 1.0\n",
+]
+
+TEST_CASE_11 = [
+    {"img": (MetaTensor(torch.tensor([[0, 1], [1, 2]]), affine=torch.as_tensor([[2,0,0,0],[0,2,0,0],[0,0,2,0],[0,0,0,1]], dtype=torch.float64), meta={"some": "info"}))},
+    "test data statistics:\nType: <class 'monai.data.meta_tensor.MetaTensor'> torch.int64\nShape: torch.Size([2, 2])\nValue range: (0, 2)\n"
+    "Value: tensor([[0, 1],\n        [1, 2]])\n"
+    "MetaTensor: {'some': 'info', affine: tensor([[2., 0., 0., 0.],\n"
+    "        [0., 2., 0., 0.],\n"
+    "        [0., 0., 2., 0.],\n"
+    "        [0., 0., 0., 1.]], dtype=torch.float64), space: RAS}\n"
+    "Additional 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, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8]
+        [TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8, TEST_CASE_9]
     )
     def test_value(self, input_param, input_data, expected_print):
         transform = DataStatsd(**input_param)
         _ = transform(input_data)
 
-    @parameterized.expand([TEST_CASE_9])
+    @parameterized.expand([TEST_CASE_10, TEST_CASE_11])
     def test_file(self, input_data, expected_print):
         with tempfile.TemporaryDirectory() as tempdir:
             filename = os.path.join(tempdir, "test_stats.log")
@@ -180,6 +212,7 @@ def test_file(self, input_data, expected_print):
                 "data_shape": True,
                 "value_range": True,
                 "data_value": True,
+                "meta_tensor": True,
                 "additional_info": np.mean,
                 "name": name,
             }