Repository: Project-MONAI/MONAI

Issue:
Support 3D RGB image in `matshow3d`
**Is your feature request related to a problem? Please describe.**
As @rijobro requested in the TensorBoard ticket, we should also add `RGB` image support in the `matshow3d` API.



Discussion:

Resolving patch:
diff --git a/monai/visualize/utils.py b/monai/visualize/utils.py
--- a/monai/visualize/utils.py
+++ b/monai/visualize/utils.py
@@ -33,6 +33,7 @@ def matshow3d(
     figsize=(10, 10),
     frames_per_row: Optional[int] = None,
     frame_dim: int = -3,
+    channel_dim: Optional[int] = None,
     vmin=None,
     vmax=None,
     every_n: int = 1,
@@ -47,15 +48,20 @@ def matshow3d(
     Create a 3D volume figure as a grid of images.
 
     Args:
-        volume: 3D volume to display. Higher dimensional arrays will be reshaped into (-1, H, W).
+        volume: 3D volume to display. data shape can be `BCHWD`, `CHWD` or `HWD`.
+            Higher dimensional arrays will be reshaped into (-1, H, W, [C]), `C` depends on `channel_dim` arg.
             A list of channel-first (C, H[, W, D]) arrays can also be passed in,
             in which case they will be displayed as a padded and stacked volume.
         fig: matplotlib figure to use. If None, a new figure will be created.
         title: title of the figure.
         figsize: size of the figure.
         frames_per_row: number of frames to display in each row. If None, sqrt(firstdim) will be used.
-        frame_dim: for higher dimensional arrays, which dimension (`-1`, `-2`, `-3`) is moved to the `-3`
-            dim and reshape to (-1, H, W) shape to construct frames, default to `-3`.
+        frame_dim: for higher dimensional arrays, which dimension from (`-1`, `-2`, `-3`) is moved to
+            the `-3` dimension. dim and reshape to (-1, H, W) shape to construct frames, default to `-3`.
+        channel_dim: if not None, explicitly specify the channel dimension to be transposed to the
+            last dimensionas shape (-1, H, W, C). this can be used to plot RGB color image.
+            if None, the channel dimension will be flattened with `frame_dim` and `batch_dim` as shape (-1, H, W).
+            note that it can only support 3D input image. default is None.
         vmin: `vmin` for the matplotlib `imshow`.
         vmax: `vmax` for the matplotlib `imshow`.
         every_n: factor to subsample the frames so that only every n-th frame is displayed.
@@ -88,6 +94,10 @@ def matshow3d(
 
     """
     vol: np.ndarray = convert_data_type(data=volume, output_type=np.ndarray)[0]  # type: ignore
+    if channel_dim is not None:
+        if channel_dim not in [0, 1] or vol.shape[channel_dim] not in [1, 3, 4]:
+            raise ValueError("channel_dim must be: None, 0 or 1, and channels of image must be 1, 3 or 4.")
+
     if isinstance(vol, (list, tuple)):
         # a sequence of channel-first volumes
         if not isinstance(vol[0], np.ndarray):
@@ -97,10 +107,13 @@ def matshow3d(
         vol = np.concatenate([pad(v) for v in vol], axis=0)
     else:  # ndarray
         while len(vol.shape) < 3:
-            vol = np.expand_dims(vol, 0)  # so that we display 1d and 2d as well
+            vol = np.expand_dims(vol, 0)  # so that we display 2d as well
 
-    vol = np.moveaxis(vol, frame_dim, -3)  # move the expected dim to construct frames with `B` or `C` dims
-    if len(vol.shape) > 3:
+    if channel_dim is not None:
+        vol = np.moveaxis(vol, frame_dim, -4)  # move the expected dim to construct frames with `B` dim
+        vol = vol.reshape((-1, vol.shape[-3], vol.shape[-2], vol.shape[-1]))
+    else:
+        vol = np.moveaxis(vol, frame_dim, -3)
         vol = vol.reshape((-1, vol.shape[-2], vol.shape[-1]))
     vmin = np.nanmin(vol) if vmin is None else vmin
     vmax = np.nanmax(vol) if vmax is None else vmax
@@ -112,9 +125,15 @@ def matshow3d(
     # create the grid of frames
     cols = max(min(len(vol), frames_per_row), 1)
     rows = int(np.ceil(len(vol) / cols))
-    width = [[0, cols * rows - len(vol)]] + [[margin, margin]] * (len(vol.shape) - 1)
+    width = [[0, cols * rows - len(vol)]]
+    if channel_dim is not None:
+        width += [[0, 0]]  # add pad width for the channel dim
+    width += [[margin, margin]] * 2
     vol = np.pad(vol.astype(dtype, copy=False), width, mode="constant", constant_values=fill_value)
     im = np.block([[vol[i * cols + j] for j in range(cols)] for i in range(rows)])
+    if channel_dim is not None:
+        # move channel dim to the end
+        im = np.moveaxis(im, 0, -1)
 
     # figure related configurations
     if fig is None:
@@ -124,6 +143,7 @@ def matshow3d(
     ax = fig.axes[0]
     ax.matshow(im, vmin=vmin, vmax=vmax, interpolation=interpolation, **kwargs)
     ax.axis("off")
+
     if title is not None:
         ax.set_title(title)
     if figsize is not None:


Test patch:
diff --git a/matshow3d_rgb_test.png b/matshow3d_rgb_test.png
new file mode 100644
Binary files /dev/null and b/matshow3d_rgb_test.png differ
diff --git a/tests/test_matshow3d.py b/tests/test_matshow3d.py
--- a/tests/test_matshow3d.py
+++ b/tests/test_matshow3d.py
@@ -15,7 +15,7 @@
 
 import numpy as np
 
-from monai.transforms import AddChanneld, Compose, LoadImaged, RandSpatialCropSamplesd, ScaleIntensityd
+from monai.transforms import AddChanneld, Compose, LoadImaged, RandSpatialCropSamplesd, RepeatChanneld, ScaleIntensityd
 from monai.utils import optional_import
 from monai.visualize.utils import matshow3d
 from tests.utils import SkipIfNoModule
@@ -71,6 +71,39 @@ def test_samples(self):
             else:
                 self.assertIsNone(comp, f"value of comp={comp}")  # None indicates test passed
 
+    def test_3d_rgb(self):
+        testing_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data")
+        keys = "image"
+        xforms = Compose(
+            [
+                LoadImaged(keys=keys),
+                AddChanneld(keys=keys),
+                ScaleIntensityd(keys=keys),
+                # change to RGB color image
+                RepeatChanneld(keys=keys, repeats=3),
+            ]
+        )
+        image_path = os.path.join(testing_dir, "anatomical.nii")
+        ims = xforms({keys: image_path})
+
+        fig = pyplot.figure()  # external figure
+        fig, _ = matshow3d(
+            volume=ims[keys],
+            fig=fig,
+            figsize=(2, 2),
+            frames_per_row=5,
+            every_n=2,
+            frame_dim=-1,
+            channel_dim=0,
+            show=False,
+        )
+
+        with tempfile.TemporaryDirectory() as tempdir:
+            tempimg = f"{tempdir}/matshow3d_rgb_test.png"
+            fig.savefig(tempimg)
+            comp = compare_images(f"{testing_dir}/matshow3d_rgb_test.png", tempimg, 5e-2)
+            self.assertIsNone(comp, f"value of comp={comp}")  # None indicates test passed
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/testing_data/matshow3d_rgb_test.png b/tests/testing_data/matshow3d_rgb_test.png
new file mode 100644
Binary files /dev/null and b/tests/testing_data/matshow3d_rgb_test.png differ