Upload 4 files

.gitattributes

@@ -1,35 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz 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
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

README.md

@@ -1,12 +1,15 @@
----
-title: ChestX Ray CTR
-emoji: 🏃
-colorFrom: pink
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.47.0
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Chest x-ray HybridGNet Segmentation
+emoji: ⚡
+colorFrom: yellow
+colorTo: blue
+sdk: gradio
+sdk_version: 5.42.0
+app_file: app.py
+pinned: false
+license: gpl-3.0
+---
+
+Demo of the HybridGNet model with 2 image-to-graph skip connections from: arxiv.org/abs/2203.10977
+Original HybridGNet model: arxiv.org/abs/2106.09832
+The training procedure was taken from: arxiv.org/abs/2211.07395

app.py

@@ -0,0 +1,303 @@
+import numpy as np
+import gradio as gr
+import cv2 
+
+from models.HybridGNet2IGSC import Hybrid 
+from utils.utils import scipy_to_torch_sparse, genMatrixesLungsHeart
+import scipy.sparse as sp
+import torch
+from zipfile import ZipFile
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+hybrid = None
+
+def getDenseMask(landmarks, h, w):
+    
+    RL = landmarks[0:44]
+    LL = landmarks[44:94]
+    H = landmarks[94:]
+    
+    img = np.zeros([h, w], dtype = 'uint8')
+    
+    RL = RL.reshape(-1, 1, 2).astype('int')
+    LL = LL.reshape(-1, 1, 2).astype('int')
+    H = H.reshape(-1, 1, 2).astype('int')
+
+    img = cv2.drawContours(img, [RL], -1, 1, -1)
+    img = cv2.drawContours(img, [LL], -1, 1, -1)
+    img = cv2.drawContours(img, [H], -1, 2, -1)
+    
+    return img
+
+def getMasks(landmarks, h, w):
+    
+    RL = landmarks[0:44]
+    LL = landmarks[44:94]
+    H = landmarks[94:]
+    
+    RL = RL.reshape(-1, 1, 2).astype('int')
+    LL = LL.reshape(-1, 1, 2).astype('int')
+    H = H.reshape(-1, 1, 2).astype('int')
+    
+    RL_mask = np.zeros([h, w], dtype = 'uint8')
+    LL_mask = np.zeros([h, w], dtype = 'uint8')
+    H_mask = np.zeros([h, w], dtype = 'uint8')
+    
+    RL_mask = cv2.drawContours(RL_mask, [RL], -1, 255, -1)
+    LL_mask = cv2.drawContours(LL_mask, [LL], -1, 255, -1)
+    H_mask = cv2.drawContours(H_mask, [H], -1, 255, -1)
+
+    return RL_mask, LL_mask, H_mask
+
+def drawOnTop(img, landmarks, original_shape):
+    h, w = original_shape
+    output = getDenseMask(landmarks, h, w)
+    
+    image = np.zeros([h, w, 3])
+    image[:,:,0] = img + 0.3 * (output == 1).astype('float') - 0.1 * (output == 2).astype('float')
+    image[:,:,1] = img + 0.3 * (output == 2).astype('float') - 0.1 * (output == 1).astype('float') 
+    image[:,:,2] = img - 0.1 * (output == 1).astype('float') - 0.2 * (output == 2).astype('float') 
+
+    image = np.clip(image, 0, 1)
+    
+    RL, LL, H = landmarks[0:44], landmarks[44:94], landmarks[94:]
+    
+    # Draw the landmarks as dots
+    
+    for l in RL:
+        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 0, 1), -1)
+    for l in LL:
+        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 0, 1), -1)
+    for l in H:
+        image = cv2.circle(image, (int(l[0]), int(l[1])), 5, (1, 1, 0), -1)
+    
+    return image
+    
+
+def loadModel(device):    
+    A, AD, D, U = genMatrixesLungsHeart()
+    N1 = A.shape[0]
+    N2 = AD.shape[0]
+
+    A = sp.csc_matrix(A).tocoo()
+    AD = sp.csc_matrix(AD).tocoo()
+    D = sp.csc_matrix(D).tocoo()
+    U = sp.csc_matrix(U).tocoo()
+
+    D_ = [D.copy()]
+    U_ = [U.copy()]
+
+    config = {}
+
+    config['n_nodes'] = [N1, N1, N1, N2, N2, N2]
+    A_ = [A.copy(), A.copy(), A.copy(), AD.copy(), AD.copy(), AD.copy()]
+    
+    A_t, D_t, U_t = ([scipy_to_torch_sparse(x).to(device) for x in X] for X in (A_, D_, U_))
+
+    config['latents'] = 64
+    config['inputsize'] = 1024
+
+    f = 32
+    config['filters'] = [2, f, f, f, f//2, f//2, f//2]
+    config['skip_features'] = f
+
+    hybrid = Hybrid(config.copy(), D_t, U_t, A_t).to(device)
+    hybrid.load_state_dict(torch.load("weights/weights.pt", map_location=torch.device(device)))
+    hybrid.eval()
+    
+    return hybrid
+
+
+def pad_to_square(img):
+    h, w = img.shape[:2]
+    
+    if h > w:
+        padw = (h - w) 
+        auxw = padw % 2
+        img = np.pad(img, ((0, 0), (padw//2, padw//2 + auxw)), 'constant')
+        
+        padh = 0
+        auxh = 0
+        
+    else:
+        padh = (w - h) 
+        auxh = padh % 2
+        img = np.pad(img, ((padh//2, padh//2 + auxh), (0, 0)), 'constant')
+
+        padw = 0
+        auxw = 0
+        
+    return img, (padh, padw, auxh, auxw)
+    
+
+def preprocess(input_img):
+    img, padding = pad_to_square(input_img)
+    
+    h, w = img.shape[:2]
+    if h != 1024 or w != 1024:
+        img = cv2.resize(img, (1024, 1024), interpolation = cv2.INTER_CUBIC)
+        
+    return img, (h, w, padding)
+
+
+def removePreprocess(output, info):
+    h, w, padding = info
+    
+    if h != 1024 or w != 1024:
+        output = output * h
+    else:
+        output = output * 1024
+    
+    padh, padw, auxh, auxw = padding
+    
+    output[:, 0] = output[:, 0] - padw//2
+    output[:, 1] = output[:, 1] - padh//2
+    
+    return output   
+
+
+def zip_files(files):
+    with ZipFile("complete_results.zip", "w") as zipObj:
+        for idx, file in enumerate(files):
+            zipObj.write(file, arcname=file.split("/")[-1])
+    return "complete_results.zip"
+
+
+def calculate_ctr(landmarks):
+    H = landmarks[94:]
+    RL = landmarks[0:44]
+    LL = landmarks[44:94]
+    cardiac_width = np.max(H[:,0]) - np.min(H[:,0])
+    thoracic_width = max(np.max(RL[:,0]), np.max(LL[:,0])) - min(np.min(RL[:,0]), np.min(LL[:,0]))
+    ctr = cardiac_width / thoracic_width if thoracic_width > 0 else 0
+    return ctr
+
+def segment(input_img):
+    global hybrid, device
+    
+    if hybrid is None:
+        hybrid = loadModel(device)
+    
+    input_img = cv2.imread(input_img, 0) / 255.0
+    original_shape = input_img.shape[:2]
+    
+    img, (h, w, padding) = preprocess(input_img)    
+        
+    data = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(device).float()
+    
+    with torch.no_grad():
+        output = hybrid(data)[0].cpu().numpy().reshape(-1, 2)
+        
+    output = removePreprocess(output, (h, w, padding))
+    
+    output = output.astype('int')
+    
+    outseg = drawOnTop(input_img, output, original_shape) 
+    
+    seg_to_save = (outseg.copy() * 255).astype('uint8')
+    cv2.imwrite("tmp/overlap_segmentation.png" , cv2.cvtColor(seg_to_save, cv2.COLOR_RGB2BGR))
+    
+    RL = output[0:44]
+    LL = output[44:94]
+    H = output[94:]
+            
+    np.savetxt("tmp/RL_landmarks.txt", RL, delimiter=" ", fmt="%d")
+    np.savetxt("tmp/LL_landmarks.txt", LL, delimiter=" ", fmt="%d")
+    np.savetxt("tmp/H_landmarks.txt", H, delimiter=" ", fmt="%d")
+    
+    RL_mask, LL_mask, H_mask = getMasks(output, original_shape[0], original_shape[1])
+    
+    cv2.imwrite("tmp/RL_mask.png", RL_mask)
+    cv2.imwrite("tmp/LL_mask.png", LL_mask)
+    cv2.imwrite("tmp/H_mask.png", H_mask)
+    
+    zip = zip_files(["tmp/RL_landmarks.txt", "tmp/LL_landmarks.txt", "tmp/H_landmarks.txt", "tmp/RL_mask.png", "tmp/LL_mask.png", "tmp/H_mask.png", "tmp/overlap_segmentation.png"])    
+    
+    ctr_value = calculate_ctr(output)
+
+    return outseg, ["tmp/RL_landmarks.txt", "tmp/LL_landmarks.txt", "tmp/H_landmarks.txt", "tmp/RL_mask.png", "tmp/LL_mask.png", "tmp/H_mask.png", "tmp/overlap_segmentation.png", zip], ctr_value
+
+
+if __name__ == "__main__":    
+    
+    with gr.Blocks() as demo:
+
+        gr.Markdown("""
+                    # Chest X-ray HybridGNet Segmentation.
+                    
+                    Demo of the HybridGNet model introduced in "Improving anatomical plausibility in medical image segmentation via hybrid graph neural networks: applications to chest x-ray analysis."
+                    
+                    Instructions:
+                    1. Upload a chest X-ray image (PA or AP) in PNG or JPEG format.
+                    2. Click on "Segment Image".
+                    
+                    Note: Pre-processing is not needed, it will be done automatically and removed after the segmentation.
+                    
+                    Please check citations below.                    
+                    """)
+
+        with gr.Tab("Segment Image"):
+            with gr.Row():
+                with gr.Column():
+                    image_input = gr.Image(type="filepath", height=750)
+                    
+                    with gr.Row():
+                        clear_button = gr.Button("Clear")
+                        image_button = gr.Button("Segment Image")
+                        
+                    gr.Examples(inputs=image_input, examples=['utils/example1.jpg','utils/example2.jpg','utils/example3.png','utils/example4.jpg'])
+                        
+                with gr.Column():
+                    image_output = gr.Image(type="filepath", height=750)
+                    results = gr.File()
+                    ctr_output = gr.Number(label="CTR (Cardiothoracic Ratio)")
+
+        gr.Markdown("""
+                    If you use this code, please cite:
+                    
+                    ```
+                    @article{gaggion2022TMI,
+                        doi = {10.1109/tmi.2022.3224660},
+                        url = {https://doi.org/10.1109%2Ftmi.2022.3224660},
+                        year = 2022,
+                        publisher = {Institute of Electrical and Electronics Engineers ({IEEE})},
+                        author = {Nicolas Gaggion and Lucas Mansilla and Candelaria Mosquera and Diego H. Milone and Enzo Ferrante},
+                        title = {Improving anatomical plausibility in medical image segmentation via hybrid graph neural networks: applications to chest x-ray analysis},
+                        journal = {{IEEE} Transactions on Medical Imaging}
+                    }
+                    ```
+                    
+                    This model was trained following the procedure explained on:
+                    
+                    ```
+                    @INPROCEEDINGS{gaggion2022ISBI,
+                        author={Gaggion, Nicolás and Vakalopoulou, Maria and Milone, Diego H. and Ferrante, Enzo},
+                        booktitle={2023 IEEE 20th International Symposium on Biomedical Imaging (ISBI)}, 
+                        title={Multi-Center Anatomical Segmentation with Heterogeneous Labels Via Landmark-Based Models}, 
+                        year={2023},
+                        volume={},
+                        number={},
+                        pages={1-5},
+                        doi={10.1109/ISBI53787.2023.10230691}
+                    }
+                    ```
+
+                    Example images extracted from Wikipedia, released under:
+                    1. CC0 Universial Public Domain. Source: https://commons.wikimedia.org/wiki/File:Normal_posteroanterior_(PA)_chest_radiograph_(X-ray).jpg
+                    2. Creative Commons Attribution-Share Alike 4.0 International. Source: https://commons.wikimedia.org/wiki/File:Chest_X-ray.jpg
+                    3. Creative Commons Attribution 3.0 Unported. Source https://commons.wikimedia.org/wiki/File:Implantable_cardioverter_defibrillator_chest_X-ray.jpg
+                    4. Creative Commons Attribution-Share Alike 3.0 Unported. Source: https://commons.wikimedia.org/wiki/File:Medical_X-Ray_imaging_PRD06_nevit.jpg
+                    
+                    Author: Nicolás Gaggion
+                    Website: [ngaggion.github.io](https://ngaggion.github.io/)
+                    
+                    """)
+        
+
+        clear_button.click(lambda: None, None, image_input, queue=False)
+        clear_button.click(lambda: None, None, image_output, queue=False)
+        clear_button.click(lambda: None, None, ctr_output, queue=False)
+
+        image_button.click(segment, inputs=image_input, outputs=[image_output, results, ctr_output], queue=False)
+
+    demo.launch()

requirements.txt

@@ -0,0 +1,7 @@
+torch==2.0.1
+numpy==1.25.0
+opencv-python==4.8.0.74
+scipy==1.10.1
+torch_geometric==2.3.0
+torchvision
+gradio==4.15.0