Initial Submit

.ipynb_checkpoints/app-checkpoint.py

@@ -0,0 +1,130 @@
+import os
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+
+import gradio as gr
+import torch
+import cv2
+import numpy as np
+from preprocess import unsharp_masking
+import glob
+import time
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(
+    "torch: ", torch.__version__,
+)
+
+def filesort(img, model):
+    # img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+    ori = img.copy()
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    h, w = img.shape
+    img_out = preprocessing(img, model)
+    return img_out, h, w, img, ori
+
+def preprocessing(img, model='SE-RegUNet 4GF'):
+    # print(img.shape, img.dtype)
+    # img = cv2.resize(img, (512, 512))
+    img = unsharp_masking(img).astype(np.uint8)
+    if model == 'AngioNet' or model == 'UNet3+':
+        img = np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
+        img_out = np.expand_dims(img, axis=0)
+    elif model == 'SE-RegUNet 4GF':
+        clahe1 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+        clahe2 = cv2.createCLAHE(clipLimit=8.0, tileGridSize=(8,8))
+        image1 = clahe1.apply(img)
+        image2 = clahe2.apply(img)
+        img = np.float32((img - img.min()) / (img.max() - img.min() + 1e-6))
+        image1 = np.float32((image1 - image1.min()) / (image1.max() - image1.min() + 1e-6))
+        image2 = np.float32((image2 - image2.min()) / (image2.max() - image2.min() + 1e-6))
+        img_out = np.stack((img, image1, image2), axis=0)
+    else:
+        clahe1 = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
+        image1 = clahe1.apply(img)
+        image1 = np.float32((image1 - image1.min()) / (image1.max() - image1.min() + 1e-6))
+        img_out = np.stack((image1,)*3, axis=0)
+    return img_out
+
+def process_input_image(img, model, rescale):
+    ori_img = img.copy()
+    h, w, _ = ori_img.shape
+    pad_h = h % 32
+    pad_w = w % 32
+    if pad_h == 0 and pad_w > 0:
+        img = ori_img[:, pad_w//2:-pad_w//2]
+    elif pad_h > 0 and pad_w == 0:
+        img = ori_img[pad_h//2:-pad_h//2, :]
+    elif pad_h > 0 and pad_w > 0:
+        img = ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2]
+    
+    if model == 'SE-RegUNet 4GF':
+        pipe = torch.jit.load('./model/SERegUNet4GF.pt')
+    elif model == 'SE-RegUNet 16GF':
+        pipe = torch.jit.load('./model/SERegUNet16GF.pt')
+    elif model == 'AngioNet':
+        pipe = torch.jit.load('./model/AngioNet.pt')
+    elif model == 'EffUNet++ B5':
+        pipe = torch.jit.load('./model/EffUNetppb5.pt')
+    elif model == 'Reg-SA-UNet++':
+        pipe = torch.jit.load('./model/RegSAUnetpp.pt')
+    elif model == 'UNet3+':
+        pipe = torch.jit.load('./model/UNet3plus.pt')
+    pipe = pipe.to(device).eval()
+    
+    start = time.time()
+    img, h, w, ori_gray, ori = filesort(img, model)
+    img = torch.FloatTensor(img).unsqueeze(0).to(device)
+    with torch.no_grad():
+        if model == 'AngioNet':
+            img = torch.cat([img, img], dim=0)
+        logit = np.round(torch.softmax(pipe.forward(img), dim=1).detach().cpu().numpy()[0, 0]).astype(np.uint8)
+    spent = time.time() - start
+    spent = f"{spent:.3f} seconds"
+    
+    logit = logit.astype(bool)
+    # img_out = cv2.cvtColor(ori, cv2.COLOR_GRAY2RGB)
+    img_out = ori.copy()
+    img_out[logit, 0] = 255
+    if pad_h == 0 and pad_w == 0:
+        ori_img = img_out
+    elif pad_h == 0 and pad_w > 0:
+        ori_img[:, pad_w//2:-pad_w//2] = img_out
+    elif pad_h > 0 and pad_w == 0:
+        ori_img[pad_h//2:-pad_h//2, :] = img_out
+    elif pad_h > 0 and pad_w > 0:
+        ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2] = img_out
+    return spent, ori_img
+
+    
+my_app = gr.Blocks()
+with my_app:
+    gr.Markdown("Coronary Angiogram Segmentation with Gradio.")
+    gr.Markdown("Author: Ching-Ting Lin, Artificial Intelligence Center, China Medical University Hospital, Taichung City, Taiwan.")
+    with gr.Tabs():
+        with gr.TabItem("Select your image"):
+            with gr.Row():
+                with gr.Column():
+                    img_source = gr.Image(label="Please select angiogram.", value='./example/angio.png', shape=(512, 512))
+                    model_choice = gr.Dropdown(['SE-RegUNet 4GF', 'SE-RegUNet 16GF', 'AngioNet', 'EffUNet++ B5', 
+                                                'Reg-SA-UNet++', 'UNet3+'], label='Model', info='Which model to infer?')
+                    model_rescale = gr.Dropdown(['2x2', '4x4', '8x8', '16x16'], label='Rescale', info='How many batches?')
+                    source_image_loader = gr.Button("Vessel Segment")
+                with gr.Column():
+                    time_spent = gr.Label(label="Time Spent (Preprocessing + Inference)")
+                    img_output = gr.Image(label="Output Mask")
+                    
+        source_image_loader.click(
+            process_input_image,
+            [
+                img_source,
+                model_choice,
+                model_rescale
+            ],
+            [
+                time_spent,
+                img_output
+            ]
+        )
+    
+my_app.launch(debug=True)

app.py

@@ -55,7 +55,7 @@ def process_input_image(img, model, rescale):
         img = ori_img[:, pad_w//2:-pad_w//2]
     elif pad_h > 0 and pad_w == 0:
         img = ori_img[pad_h//2:-pad_h//2, :]
-    if pad_h == 0 and pad_w == 0:
+    elif pad_h > 0 and pad_w > 0:
         img = ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2]
     
     if model == 'SE-RegUNet 4GF':
@@ -86,11 +86,13 @@ def process_input_image(img, model, rescale):
     # img_out = cv2.cvtColor(ori, cv2.COLOR_GRAY2RGB)
     img_out = ori.copy()
     img_out[logit, 0] = 255
-    if pad_h == 0 and pad_w > 0:
+    if pad_h == 0 and pad_w == 0:
+        ori_img = img_out
+    elif pad_h == 0 and pad_w > 0:
         ori_img[:, pad_w//2:-pad_w//2] = img_out
     elif pad_h > 0 and pad_w == 0:
         ori_img[pad_h//2:-pad_h//2, :] = img_out
-    if pad_h == 0 and pad_w == 0:
+    elif pad_h > 0 and pad_w > 0:
         ori_img[pad_h//2:-pad_h//2, pad_w//2:-pad_w//2] = img_out
     return spent, ori_img