Initial Submit

.gitattributes

@@ -32,3 +32,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model/RegSAUnetpp.pt filter=lfs diff=lfs merge=lfs -text
+model/SERegUNet16GF.pt filter=lfs diff=lfs merge=lfs -text
+model/SERegUNet4GF.pt filter=lfs diff=lfs merge=lfs -text
+model/UNet3plus.pt filter=lfs diff=lfs merge=lfs -text
+model/AngioNet.pt filter=lfs diff=lfs merge=lfs -text
+model/EffUNetppb5.pt filter=lfs diff=lfs merge=lfs -text
+example/angio.png filter=lfs diff=lfs merge=lfs -text
+example/@eaDir filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,118 @@
+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
+    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
+    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)

model/AngioNet.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60f54fb6481e59b61e421df5a0e59842f07fbbd89f0f7f2535ee5a4c0a104cdb
+size 148997984

model/EffUNetppb5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1501967cb37b8d3ec025484eedccb9ac8d343dc75886a978acd3ada22a9e1f7
+size 120253012

model/RegSAUnetpp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0fee28a388c9481c1c32c9eb04247a35474d468e4fbab562d0781112ca7f12ed
+size 128479489

model/SERegUNet16GF.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42421187a29f42d0007d686d575e86ae5a0a72c5059fbaecd142af04765f9d7e
+size 811937473

model/SERegUNet4GF.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:758b4d71740d4fa9b692532a71f99db5a8095abd40eb818b90b0eab6641e191b
+size 125795146

model/UNet3plus.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:331e34c0cd21c4e84fe4e06fb31267c3c0612cae69f4de350d7563151577d6d7
+size 108264749

preprocess.py

@@ -0,0 +1,6 @@
+import cv2
+
+def unsharp_masking(img):
+    gaussian = cv2.GaussianBlur(img, (0, 0), 2.0)
+    img = cv2.addWeighted(img, 2.0, gaussian, -1.0, 0)
+    return img

requirements.txt

@@ -0,0 +1,5 @@
+timm == 0.8.0.dev0
+torch == 1.13.1
+numpy == 1.23.5
+opencv-python == 4.5.5.64
+matplotlib == 3.5.1