Spaces:
Sleeping
Sleeping
| 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 | |
| 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 calculate_image_tilt(landmarks): | |
| """Calculate image tilt angle based on lung symmetry""" | |
| RL = landmarks[0:44] # Right lung | |
| LL = landmarks[44:94] # Left lung | |
| # Find the topmost points of both lungs | |
| rl_top_idx = np.argmin(RL[:, 1]) | |
| ll_top_idx = np.argmin(LL[:, 1]) | |
| rl_top = RL[rl_top_idx] | |
| ll_top = LL[ll_top_idx] | |
| # Calculate angle between the line connecting lung tops and horizontal | |
| dx = ll_top[0] - rl_top[0] | |
| dy = ll_top[1] - rl_top[1] | |
| angle_rad = np.arctan2(dy, dx) | |
| angle_deg = np.degrees(angle_rad) | |
| return angle_deg, rl_top, ll_top | |
| def rotate_points(points, angle_deg, center): | |
| """Rotate points around a center by given angle""" | |
| angle_rad = np.radians(-angle_deg) # Negative to correct the tilt | |
| cos_a = np.cos(angle_rad) | |
| sin_a = np.sin(angle_rad) | |
| # Translate to origin | |
| translated = points - center | |
| # Rotate | |
| rotated = np.zeros_like(translated) | |
| rotated[:, 0] = translated[:, 0] * cos_a - translated[:, 1] * sin_a | |
| rotated[:, 1] = translated[:, 0] * sin_a + translated[:, 1] * cos_a | |
| # Translate back | |
| return rotated + center | |
| 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:] | |
| # Calculate image tilt and correct it for measurements | |
| tilt_angle, rl_top, ll_top = calculate_image_tilt(landmarks) | |
| image_center = np.array([w/2, h/2]) | |
| # Draw tilt reference line (green) | |
| image = cv2.line(image, (int(rl_top[0]), int(rl_top[1])), (int(ll_top[0]), int(ll_top[1])), (0, 1, 0), 1) | |
| # Add tilt angle text | |
| tilt_text = f"Tilt: {tilt_angle:.1f}°" | |
| cv2.putText(image, tilt_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 1, 0), 2) | |
| # Correct landmarks for tilt | |
| if abs(tilt_angle) > 2: # Only correct if tilt is significant | |
| RL_corrected = rotate_points(RL, tilt_angle, image_center) | |
| LL_corrected = rotate_points(LL, tilt_angle, image_center) | |
| H_corrected = rotate_points(H, tilt_angle, image_center) | |
| cv2.putText(image, "Tilt Corrected", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (1, 1, 0), 2) | |
| else: | |
| RL_corrected, LL_corrected, H_corrected = RL, LL, H | |
| # 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) | |
| # Draw horizontal lines for CTR calculation using corrected landmarks | |
| # Heart (red line) - using corrected coordinates | |
| heart_xmin = int(np.min(H_corrected[:, 0])) | |
| heart_xmax = int(np.max(H_corrected[:, 0])) | |
| heart_y = int(np.mean([H_corrected[np.argmin(H_corrected[:, 0]), 1], H_corrected[np.argmax(H_corrected[:, 0]), 1]])) | |
| image = cv2.line(image, (heart_xmin, heart_y), (heart_xmax, heart_y), (1, 0, 0), 2) | |
| # Add vertical lines at heart endpoints to verify widest points | |
| line_length = 30 # Length of vertical indicator lines | |
| image = cv2.line(image, (heart_xmin, heart_y - line_length), (heart_xmin, heart_y + line_length), (1, 0, 0), 2) | |
| image = cv2.line(image, (heart_xmax, heart_y - line_length), (heart_xmax, heart_y + line_length), (1, 0, 0), 2) | |
| # Thorax (blue line) - using corrected coordinates | |
| thorax_xmin = int(min(np.min(RL_corrected[:, 0]), np.min(LL_corrected[:, 0]))) | |
| thorax_xmax = int(max(np.max(RL_corrected[:, 0]), np.max(LL_corrected[:, 0]))) | |
| # Find y at leftmost and rightmost points | |
| if np.min(RL_corrected[:, 0]) < np.min(LL_corrected[:, 0]): | |
| thorax_ymin = RL_corrected[np.argmin(RL_corrected[:, 0]), 1] | |
| else: | |
| thorax_ymin = LL_corrected[np.argmin(LL_corrected[:, 0]), 1] | |
| if np.max(RL_corrected[:, 0]) > np.max(LL_corrected[:, 0]): | |
| thorax_ymax = RL_corrected[np.argmax(RL_corrected[:, 0]), 1] | |
| else: | |
| thorax_ymax = LL_corrected[np.argmax(LL_corrected[:, 0]), 1] | |
| thorax_y = int(np.mean([thorax_ymin, thorax_ymax])) | |
| image = cv2.line(image, (thorax_xmin, thorax_y), (thorax_xmax, thorax_y), (0, 0, 1), 2) | |
| # Add vertical lines at thorax endpoints to verify widest points | |
| image = cv2.line(image, (thorax_xmin, thorax_y - line_length), (thorax_xmin, thorax_y + line_length), (0, 0, 1), 2) | |
| image = cv2.line(image, (thorax_xmax, thorax_y - line_length), (thorax_xmax, thorax_y + line_length), (0, 0, 1), 2) | |
| # Store corrected landmarks for CTR calculation | |
| return image, (RL_corrected, LL_corrected, H_corrected, tilt_angle) | |
| 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 calculate_ctr(landmarks, corrected_landmarks=None): | |
| if corrected_landmarks is not None: | |
| RL, LL, H, tilt_angle = corrected_landmarks | |
| else: | |
| H = landmarks[94:] | |
| RL = landmarks[0:44] | |
| LL = landmarks[44:94] | |
| tilt_angle = 0 | |
| 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 round(ctr, 3), abs(tilt_angle) | |
| 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, corrected_data = 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)) | |
| ctr_value, tilt_angle = calculate_ctr(output, corrected_data) | |
| # Add tilt warning to interpretation | |
| tilt_warning = "" | |
| if tilt_angle > 5: | |
| tilt_warning = f" (⚠️ Image tilted {tilt_angle:.1f}° - measurement corrected)" | |
| elif tilt_angle > 2: | |
| tilt_warning = f" (Image tilted {tilt_angle:.1f}° - corrected)" | |
| if ctr_value < 0.5: | |
| interpretation = f"Normal{tilt_warning}" | |
| elif 0.51 <= ctr_value <= 0.55: | |
| interpretation = f"Mild Cardiomegaly (CTR 51-55%){tilt_warning}" | |
| elif 0.56 <= ctr_value <= 0.60: | |
| interpretation = f"Moderate Cardiomegaly (CTR 56-60%){tilt_warning}" | |
| elif ctr_value > 0.60: | |
| interpretation = f"Severe Cardiomegaly (CTR > 60%){tilt_warning}" | |
| else: | |
| interpretation = f"Cardiomegaly{tilt_warning}" | |
| return outseg, "tmp/overlap_segmentation.png", ctr_value, interpretation | |
| 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) | |
| with gr.Row(): | |
| ctr_output = gr.Number(label="CTR (Cardiothoracic Ratio)") | |
| ctr_interpretation = gr.Textbox(label="Interpretation", interactive=False) | |
| results = gr.File() | |
| 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) | |
| clear_button.click(lambda: None, None, ctr_interpretation, queue=False) | |
| image_button.click(segment, inputs=image_input, outputs=[image_output, results, ctr_output, ctr_interpretation], queue=False) | |
| demo.launch() |