From 02fcf705c522590f7e2aced4ee2a54851515cc1d Mon Sep 17 00:00:00 2001
From: pme0 <12113751+pme0@users.noreply.github.com>
Date: Thu, 22 Jun 2023 12:39:14 +0100
Subject: [PATCH] single-channel support

---
 deeplightning/viz/image/image.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/deeplightning/viz/image/image.py b/deeplightning/viz/image/image.py
index dce46fd..532faa0 100644
--- a/deeplightning/viz/image/image.py
+++ b/deeplightning/viz/image/image.py
@@ -13,6 +13,8 @@ def plot_resized_image(
     resize: Tuple[int,int] = None,
     display_image: bool = True,
     display_size: Tuple[int,int] = None,
+    channel: int = None,
+    channel_cmap: str = None,
     save_fp : str = None,
 ):
     """Plot resized image & saves with new exact pixel size.
@@ -26,8 +28,10 @@ def plot_resized_image(
     display_size : size of displayed image, in inches (width, height)
     save_fp : output image save filepath
     """
-    assert isinstance(resize, tuple) and len(resize) == 2
-    
+    assert resize is None or isinstance(resize, tuple) and len(resize) == 2
+    if resize is None:
+        resize = imagesize.get(image_fp)  #(w,h)
+
     image = Image.open(image_fp)
     image = image.convert('RGB')
     image = image.resize(resize)
@@ -40,7 +44,12 @@ def plot_resized_image(
     
     if display_image:
         fig = plt.figure(figsize=display_size)
-        plt.imshow(new_image)
+        if channel is not None:
+            new_image = np.array(new_image)
+            #new_image[:,:,(channel!=0, channel!=1, channel!=2)] *= 0
+            new_image = new_image[:,:,channel]
+            new_image = Image.fromarray(new_image)
+        plt.imshow(new_image, cmap=channel_cmap)
         plt.axis("off")
         plt.show()