Add unet versions(v0, v1, v2, v3)

app.py

@@ -1,7 +1,216 @@
 import gradio as gr
+from gradio_client import Client
+import numpy as np
+#import torch
+import requests
+from PIL import Image
+#from torchvision import transforms
+from predict_unet import predict_model
 
-def greet(name):
-    return "Hello " + name + "!!"
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+# device = torch.device(
+#     "cuda"
+#     if torch.cuda.is_available()
+#     else "mps"
+#     if torch.backends.mps.is_available()
+#     else "cpu"
+# )
+
+title = "<center><strong><font size='8'> Medical Image Segmentation with UNet </font></strong></center>"
+
+examples = [["examples/50494616.jpg"], ["examples/50494676.jpg"], ["examples/56399783.jpg"],
+            ["examples/56399789.jpg"], ["examples/56399831.jpg"], ["examples/56399959.jpg"],
+            ["examples/56400014.jpg"], ["examples/56400119.jpg"],
+            ["examples/56481903.jpg"], ["examples/70749195.jpg"]]
+
+def run_unetv0(input):
+    output = predict_model(input, "v0")
+    normalized_output = np.clip(output, 0, 1)
+    return normalized_output
+
+def run_unetv1(input):
+    output = predict_model(input, "v1")
+    normalized_output = np.clip(output, 0, 1)
+    return normalized_output
+
+def run_unetv2(input):
+    output = predict_model(input, "v2")
+    normalized_output = np.clip(output, 0, 1)
+    return normalized_output
+
+def run_unetv3(input):
+    output = predict_model(input, "v3")
+    normalized_output = np.clip(output, 0, 1)
+    return normalized_output
+    
+
+input_img_v0 = gr.Image(label="Input", type='numpy')
+segm_img_v0 = gr.Image(label="Segmented Image")
+
+input_img_v1 = gr.Image(label="Input", type='numpy')
+segm_img_v1 = gr.Image(label="Segmented Image")
+
+input_img_v2 = gr.Image(label="Input", type='numpy')
+segm_img_v2 = gr.Image(label="Segmented Image")
+
+input_img_v3 = gr.Image(label="Input", type='numpy')
+segm_img_v3 = gr.Image(label="Segmented Image")
+
+
+with gr.Blocks(title='UNet examples') as demo:
+    # v0: regular UNet
+    with gr.Tab("Regular UNet (v0)"):
+        # display input image and segmented image
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                input_img_v0.render()
+
+            with gr.Column(scale=1):
+                segm_img_v0.render()
+        
+        # submit and clear
+        with gr.Row():
+            with gr.Column():
+                segment_btn_v0 = gr.Button("Run Segmentation", variant='primary')
+                clear_btn_v0 = gr.Button("Clear", variant="secondary")
+
+                # load examples
+                gr.Markdown("Try some of the examples below")
+                gr.Examples(examples=examples,
+                            inputs=[input_img_v0],
+                            outputs=segm_img_v0,
+                            fn=run_unetv0,
+                            cache_examples=False,
+                            examples_per_page=5)
+
+            # just a placeholder for second column
+            with gr.Column():
+                gr.Markdown("")
+
+        segment_btn_v0.click(run_unetv0,
+                            inputs=[
+                                input_img_v0,
+                            ],
+                            outputs=segm_img_v0)
+
+
+    # v1: UNet3+
+    with gr.Tab("UNet3+ (v1)"):
+        # display input image and segmented image
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                input_img_v1.render()
+
+            with gr.Column(scale=1):
+                segm_img_v1.render()
+        
+        # submit and clear
+        with gr.Row():
+            with gr.Column():
+                segment_btn_v1 = gr.Button("Run Segmentation", variant='primary')
+                clear_btn_v1 = gr.Button("Clear", variant="secondary")
+
+                # load examples
+                gr.Markdown("Try some of the examples below")
+                gr.Examples(examples=examples,
+                            inputs=[input_img_v1],
+                            outputs=segm_img_v1,
+                            fn=run_unetv1,
+                            cache_examples=False,
+                            examples_per_page=5)
+
+            # just a placeholder for second column
+            with gr.Column():
+                gr.Markdown("")
+
+        segment_btn_v1.click(run_unetv1,
+                            inputs=[
+                                input_img_v1,
+                            ],
+                            outputs=segm_img_v1)
+
+    
+    # v2: UNet3+ with deep supervision
+    with gr.Tab("UNet3+(v2) with deep supervision"):
+        # display input image and segmented image
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                input_img_v2.render()
+
+            with gr.Column(scale=1):
+                segm_img_v2.render()
+        
+        # submit and clear
+        with gr.Row():
+            with gr.Column():
+                segment_btn_v2 = gr.Button("Run Segmentation", variant='primary')
+                clear_btn_v2 = gr.Button("Clear", variant="secondary")
+
+                # load examples
+                gr.Markdown("Try some of the examples below")
+                gr.Examples(examples=examples,
+                            inputs=[input_img_v2],
+                            outputs=segm_img_v2,
+                            fn=run_unetv2,
+                            cache_examples=False,
+                            examples_per_page=5)
+
+            # just a placeholder for second column
+            with gr.Column():
+                gr.Markdown("")
+
+        segment_btn_v2.click(run_unetv2,
+                            inputs=[
+                                input_img_v2,
+                            ],
+                            outputs=segm_img_v2)
+
+
+    # v3: UNet3+ with deep supervision and cgm
+    with gr.Tab("UNet3+(v3) with deep supervision and cgm"):
+        # display input image and segmented image
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                input_img_v3.render()
+
+            with gr.Column(scale=1):
+                segm_img_v3.render()
+        
+        # submit and clear
+        with gr.Row():
+            with gr.Column():
+                segment_btn_v3 = gr.Button("Run Segmentation", variant='primary')
+                clear_btn_v3 = gr.Button("Clear", variant="secondary")
+
+                # load examples
+                gr.Markdown("Try some of the examples below")
+                gr.Examples(examples=examples,
+                            inputs=[input_img_v3],
+                            outputs=segm_img_v3,
+                            fn=run_unetv3,
+                            cache_examples=False,
+                            examples_per_page=5)
+
+            # just a placeholder for second column
+            with gr.Column():
+                gr.Markdown("")
+
+        segment_btn_v3.click(run_unetv3,
+                            inputs=[
+                                input_img_v3,
+                            ],
+                            outputs=segm_img_v3)
+
+
+    def clear():
+        return None, None
+
+    clear_btn_v0.click(clear, outputs=[input_img_v0, segm_img_v0])
+    clear_btn_v1.click(clear, outputs=[input_img_v1, segm_img_v1])
+    clear_btn_v2.click(clear, outputs=[input_img_v2, segm_img_v2])
+    clear_btn_v3.click(clear, outputs=[input_img_v3, segm_img_v3])
+
+
+demo.queue()
+demo.launch()
+

predict_unet.py

@@ -0,0 +1,69 @@
+
+import os
+import numpy as np
+import skimage.io as skio
+import skimage.transform as trans
+from skimage.color import rgb2gray
+from matplotlib import pyplot as plt
+import sys
+
+sys.path.append("/panfs/jay/groups/29/umii/mo000007/zooniverse/UNet")
+
+from utils import *
+from unet import unet
+from unet_3plus import UNet_3Plus, UNet_3Plus_DeepSup, UNet_3Plus_DeepSup_CGM
+
+
+def predict_model(input, unet_type):
+    model_path = "/home/umii/mo000007/zooniverse/UNet/trained_models"
+    h, w = 256, 256
+    input_shape = [h, w, 1]
+    output_channels = 1
+    batch_size = 1
+
+    # convert image into numpy array and reshape it into model's input size
+    img = trans.resize(input, (w, h))
+    result_img = img.copy()
+    img = rgb2gray(img).reshape(1, h, w, 1)
+
+    # Load U-net model based on version: UNet type vo:unet, v1:unet3+, v2:unet3+ with deep supervision, v3:unet3+ with cgm
+    if unet_type == 'v0':
+        # load the pretrained model
+        model_name = "unetv0_sgd500_neptune"
+        model_file = os.path.join(model_path, model_name + ".hdf5")
+        model = unet(model_file)
+    elif unet_type == 'v1':
+        # load the pretrained model
+        model_name = "unetv1_sgd500_neptune"
+        model_file = os.path.join(model_path, model_name + ".hdf5")
+        model = UNet_3Plus(input_shape, output_channels, model_file)
+    elif unet_type == 'v2':
+        # load the pretrained model
+        model_name = "unetv2_sgd500_neptune"
+        model_file = os.path.join(model_path, model_name + ".hdf5")
+        model = UNet_3Plus_DeepSup(input_shape, output_channels, model_file)
+    else:
+        # load the pretrained model
+        model_name = "unetv3_sgd500_neptune"
+        model_file = os.path.join(model_path, model_name + ".hdf5")
+        model = UNet_3Plus_DeepSup_CGM(input_shape, output_channels, model_file)
+
+    # Predict and save the results as numpy array
+    results = model.predict(img)
+
+    # Preprocess the prediction from the model depending on the model
+    if unet_type == 'v2' or unet_type == 'v3':
+        pred = np.copy(results[0])
+    else:
+        pred = np.copy(results)
+    pred[pred >= 0.5] = 1
+    pred[pred < 0.5] = 0
+    output =  np.array(pred[0][:,:,0])
+
+    # visualize the output mask
+    seg_color = [0, 0, 255]
+    masked = output != 0
+    result_img[masked] = seg_color
+
+    return result_img
+