app.py

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} degrees"
    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 measurement lines that follow the image tilt for visual accuracy
    # Use corrected coordinates for accurate measurement, but draw tilted lines for visual appeal
    
    # Heart (red line) - calculate positions from corrected coordinates
    heart_xmin_corrected = np.min(H_corrected[:, 0])
    heart_xmax_corrected = np.max(H_corrected[:, 0])
    heart_y_corrected = np.mean([H_corrected[np.argmin(H_corrected[:, 0]), 1], H_corrected[np.argmax(H_corrected[:, 0]), 1]])
    
    # Rotate back to match the tilted image for display
    heart_points_corrected = np.array([[heart_xmin_corrected, heart_y_corrected], [heart_xmax_corrected, heart_y_corrected]])
    heart_points_display = rotate_points(heart_points_corrected, -tilt_angle, image_center)  # Rotate back for display
    
    heart_start = (int(heart_points_display[0, 0]), int(heart_points_display[0, 1]))
    heart_end = (int(heart_points_display[1, 0]), int(heart_points_display[1, 1]))
    image = cv2.line(image, heart_start, heart_end, (1, 0, 0), 2)
    
    # Add perpendicular lines at heart endpoints
    line_length = 30
    # Calculate perpendicular direction
    heart_dx = heart_end[0] - heart_start[0]
    heart_dy = heart_end[1] - heart_start[1]
    heart_length = np.sqrt(heart_dx**2 + heart_dy**2)
    if heart_length > 0:
        perp_x = -heart_dy / heart_length * line_length
        perp_y = heart_dx / heart_length * line_length
        
        # Perpendicular lines at start point
        image = cv2.line(image, 
                        (int(heart_start[0] + perp_x), int(heart_start[1] + perp_y)),
                        (int(heart_start[0] - perp_x), int(heart_start[1] - perp_y)), 
                        (1, 0, 0), 2)
        # Perpendicular lines at end point
        image = cv2.line(image, 
                        (int(heart_end[0] + perp_x), int(heart_end[1] + perp_y)),
                        (int(heart_end[0] - perp_x), int(heart_end[1] - perp_y)), 
                        (1, 0, 0), 2)

    # Thorax (blue line) - calculate positions from corrected coordinates
    thorax_xmin_corrected = min(np.min(RL_corrected[:, 0]), np.min(LL_corrected[:, 0]))
    thorax_xmax_corrected = max(np.max(RL_corrected[:, 0]), np.max(LL_corrected[:, 0]))
    
    # Find y at leftmost and rightmost points (corrected)
    if np.min(RL_corrected[:, 0]) < np.min(LL_corrected[:, 0]):
        thorax_ymin_corrected = RL_corrected[np.argmin(RL_corrected[:, 0]), 1]
    else:
        thorax_ymin_corrected = LL_corrected[np.argmin(LL_corrected[:, 0]), 1]
    if np.max(RL_corrected[:, 0]) > np.max(LL_corrected[:, 0]):
        thorax_ymax_corrected = RL_corrected[np.argmax(RL_corrected[:, 0]), 1]
    else:
        thorax_ymax_corrected = LL_corrected[np.argmax(LL_corrected[:, 0]), 1]
    thorax_y_corrected = np.mean([thorax_ymin_corrected, thorax_ymax_corrected])
    
    # Rotate back to match the tilted image for display
    thorax_points_corrected = np.array([[thorax_xmin_corrected, thorax_y_corrected], [thorax_xmax_corrected, thorax_y_corrected]])
    thorax_points_display = rotate_points(thorax_points_corrected, -tilt_angle, image_center)  # Rotate back for display
    
    thorax_start = (int(thorax_points_display[0, 0]), int(thorax_points_display[0, 1]))
    thorax_end = (int(thorax_points_display[1, 0]), int(thorax_points_display[1, 1]))
    image = cv2.line(image, thorax_start, thorax_end, (0, 0, 1), 2)
    
    # Add perpendicular lines at thorax endpoints
    thorax_dx = thorax_end[0] - thorax_start[0]
    thorax_dy = thorax_end[1] - thorax_start[1]
    thorax_length = np.sqrt(thorax_dx**2 + thorax_dy**2)
    if thorax_length > 0:
        perp_x = -thorax_dy / thorax_length * line_length
        perp_y = thorax_dx / thorax_length * line_length
        
        # Perpendicular lines at start point
        image = cv2.line(image, 
                        (int(thorax_start[0] + perp_x), int(thorax_start[1] + perp_y)),
                        (int(thorax_start[0] - perp_x), int(thorax_start[1] - perp_y)), 
                        (0, 0, 1), 2)
        # Perpendicular lines at end point
        image = cv2.line(image, 
                        (int(thorax_end[0] + perp_x), int(thorax_end[1] + perp_y)),
                        (int(thorax_end[0] - perp_x), int(thorax_end[1] - perp_y)), 
                        (0, 0, 1), 2)
    
    # Store corrected landmarks for CTR calculation
    return image, (RL_corrected, LL_corrected, H_corrected, tilt_angle)


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 validate_landmarks_consistency(landmarks, original_landmarks, threshold=0.05):
    """Validate that corrected landmarks maintain anatomical consistency"""
    try:
        # Check if heart is still between lungs
        RL = landmarks[0:44]
        LL = landmarks[44:94]
        H = landmarks[94:]
        
        rl_center_x = np.mean(RL[:, 0])
        ll_center_x = np.mean(LL[:, 0])
        h_center_x = np.mean(H[:, 0])
        
        # Heart should be between lung centers
        if not (min(rl_center_x, ll_center_x) <= h_center_x <= max(rl_center_x, ll_center_x)):
            print("Warning: Heart position validation failed")
            return False
            
        # Check if total change is reasonable
        total_change = np.mean(np.linalg.norm(landmarks - original_landmarks, axis=1))
        relative_change = total_change / np.mean(np.linalg.norm(original_landmarks, axis=1))
        
        if relative_change > threshold:
            print(f"Warning: Landmarks changed by {relative_change:.3f}, exceeds threshold {threshold}")
            return False
            
        return True
        
    except Exception as e:
        print(f"Error in landmark validation: {e}")
        return False

def calculate_ctr_robust(landmarks, corrected_landmarks=None):
    """Calculate CTR with multiple validation steps"""
    try:
        original_landmarks = landmarks.copy()
        
        if corrected_landmarks is not None:
            RL, LL, H, tilt_angle = corrected_landmarks
            
            # Validate correction
            corrected_all = np.vstack([RL, LL, H])
            if validate_landmarks_consistency(corrected_all, original_landmarks):
                landmarks_to_use = corrected_all
                correction_applied = True
            else:
                # Use original landmarks if validation fails
                H = landmarks[94:]
                RL = landmarks[0:44]
                LL = landmarks[44:94]
                landmarks_to_use = landmarks
                correction_applied = False
                tilt_angle = 0
        else:
            H = landmarks[94:]
            RL = landmarks[0:44]
            LL = landmarks[44:94]
            landmarks_to_use = landmarks
            tilt_angle = 0
            correction_applied = False
        
        # Method 1: Traditional width measurement
        cardiac_width_1 = np.max(H[:, 0]) - np.min(H[:, 0])
        thoracic_width_1 = max(np.max(RL[:, 0]), np.max(LL[:, 0])) - min(np.min(RL[:, 0]), np.min(LL[:, 0]))
        
        # Method 2: Centroid-based measurement (more robust to outliers)
        h_centroid = np.mean(H, axis=0)
        rl_centroid = np.mean(RL, axis=0)
        ll_centroid = np.mean(LL, axis=0)
        
        # Find widest points from centroids
        h_distances = np.linalg.norm(H - h_centroid, axis=1)
        cardiac_width_2 = 2 * np.max(h_distances)
        
        thoracic_width_2 = max(np.max(RL[:, 0]), np.max(LL[:, 0])) - min(np.min(RL[:, 0]), np.min(LL[:, 0]))
        
        # Method 3: Percentile-based measurement (removes extreme outliers)
        cardiac_x_coords = H[:, 0]
        cardiac_width_3 = np.percentile(cardiac_x_coords, 95) - np.percentile(cardiac_x_coords, 5)
        
        lung_x_coords = np.concatenate([RL[:, 0], LL[:, 0]])
        thoracic_width_3 = np.percentile(lung_x_coords, 95) - np.percentile(lung_x_coords, 5)
        
        # Calculate CTR for each method
        ctr_1 = cardiac_width_1 / thoracic_width_1 if thoracic_width_1 > 0 else 0
        ctr_2 = cardiac_width_2 / thoracic_width_2 if thoracic_width_2 > 0 else 0
        ctr_3 = cardiac_width_3 / thoracic_width_3 if thoracic_width_3 > 0 else 0
        
        # Validate consistency between methods
        ctr_values = [ctr_1, ctr_2, ctr_3]
        ctr_std = np.std(ctr_values)
        
        if ctr_std > 0.05:  # High variance between methods
            print(f"Warning: CTR calculation methods show high variance (std: {ctr_std:.3f})")
            confidence = "Low"
        elif ctr_std > 0.02:
            confidence = "Medium"
        else:
            confidence = "High"
        
        # Use median of methods for final result
        final_ctr = np.median(ctr_values)
        
        return {
            'ctr': round(final_ctr, 3),
            'tilt_angle': abs(tilt_angle),
            'correction_applied': correction_applied,
            'confidence': confidence,
            'method_variance': round(ctr_std, 4),
            'individual_results': {
                'traditional': round(ctr_1, 3),
                'centroid': round(ctr_2, 3),
                'percentile': round(ctr_3, 3)
            }
        }
        
    except Exception as e:
        print(f"Error in robust CTR calculation: {e}")
        return {
            'ctr': 0,
            'tilt_angle': 0,
            'correction_applied': False,
            'confidence': 'Error',
            'method_variance': 0,
            'individual_results': {}
        }


def detect_image_rotation_advanced(img):
    """Enhanced rotation detection using multiple methods"""
    try:
        angles = []
        
        # Method 1: Edge-based detection with focus on spine/mediastinum
        edges = cv2.Canny((img * 255).astype(np.uint8), 50, 150)
        h, w = img.shape
        
        # Focus on central region where spine should be
        spine_region = edges[h//4:3*h//4, w//3:2*w//3]
        
        # Find strong vertical lines (spine alignment)
        lines = cv2.HoughLines(spine_region, 1, np.pi/180, threshold=50)
        if lines is not None:
            for line in lines[:5]:  # Top 5 lines
                rho, theta = line[0]
                angle = np.degrees(theta) - 90
                if abs(angle) < 30:  # Near vertical lines
                    angles.append(angle)
        
        # Method 2: Chest boundary detection
        # Find chest outline using contours
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            # Get largest contour (chest boundary)
            largest_contour = max(contours, key=cv2.contourArea)
            
            # Fit ellipse to chest boundary
            if len(largest_contour) >= 5:
                ellipse = cv2.fitEllipse(largest_contour)
                chest_angle = ellipse[2] - 90  # Convert to rotation angle
                if abs(chest_angle) < 45:
                    angles.append(chest_angle)
        
        # Method 3: Template-based symmetry detection
        # Check left-right symmetry
        left_half = img[:, :w//2]
        right_half = np.fliplr(img[:, w//2:])
        
        # Try different rotation angles to find best symmetry
        best_angle = 0
        best_correlation = 0
        
        for test_angle in range(-15, 16, 2):
            if test_angle == 0:
                test_left = left_half
            else:
                center = (left_half.shape[1]//2, left_half.shape[0]//2)
                rotation_matrix = cv2.getRotationMatrix2D(center, test_angle, 1.0)
                test_left = cv2.warpAffine(left_half, rotation_matrix, 
                                         (left_half.shape[1], left_half.shape[0]))
            
            # Calculate correlation
            correlation = cv2.matchTemplate(test_left, right_half, cv2.TM_CCOEFF_NORMED).max()
            if correlation > best_correlation:
                best_correlation = correlation
                best_angle = test_angle
        
        if best_correlation > 0.3:  # Good symmetry found
            angles.append(best_angle)
        
        # Combine all methods
        if angles:
            # Remove outliers using IQR
            angles = np.array(angles)
            Q1, Q3 = np.percentile(angles, [25, 75])
            IQR = Q3 - Q1
            filtered_angles = angles[(angles >= Q1 - 1.5*IQR) & (angles <= Q3 + 1.5*IQR)]
            
            if len(filtered_angles) > 0:
                final_angle = np.median(filtered_angles)
                return final_angle if abs(final_angle) > 1 else 0
        
        return 0
        
    except Exception as e:
        print(f"Error in advanced rotation detection: {e}")
        return 0

def rotate_image(img, angle):
    """Rotate image by given angle"""
    try:
        if abs(angle) < 1:
            return img, 0
        
        h, w = img.shape[:2]
        center = (w // 2, h // 2)
        
        # Get rotation matrix
        rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)
        
        # Calculate new dimensions
        cos_angle = abs(rotation_matrix[0, 0])
        sin_angle = abs(rotation_matrix[0, 1])
        new_w = int((h * sin_angle) + (w * cos_angle))
        new_h = int((h * cos_angle) + (w * sin_angle))
        
        # Adjust translation
        rotation_matrix[0, 2] += (new_w / 2) - center[0]
        rotation_matrix[1, 2] += (new_h / 2) - center[1]
        
        # Rotate image
        rotated = cv2.warpAffine(img, rotation_matrix, (new_w, new_h), 
                                borderMode=cv2.BORDER_CONSTANT, borderValue=0)
        
        return rotated, angle
    except Exception as e:
        print(f"Error in image rotation: {e}")
        return img, 0

def segment(input_img):
    global hybrid, device

    try:
        if hybrid is None:
            hybrid = loadModel(device)

        original_img = cv2.imread(input_img, 0) / 255.0
        original_shape = original_img.shape[:2]
        
        # Step 1: Enhanced rotation detection (re-enabled)
        detected_rotation = detect_image_rotation_advanced(original_img)
        was_rotated = False
        processing_img = original_img
        
        # Step 2: Rotate image if significant rotation detected
        if abs(detected_rotation) > 3:
            processing_img, actual_rotation = rotate_image(original_img, -detected_rotation)
            was_rotated = True
            print(f"Applied rotation correction: {detected_rotation:.1f}°")
        else:
            actual_rotation = 0

        # Step 3: Preprocess the image
        img, (h, w, padding) = preprocess(processing_img)

        # Step 4: AI segmentation
        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)

        # Step 5: Remove preprocessing
        output = removePreprocess(output, (h, w, padding))
        
        # Step 6: Rotate landmarks back if image was rotated
        if was_rotated:
            center = np.array([original_shape[1]/2, original_shape[0]/2])
            output = rotate_points(output, actual_rotation, center)
        
        # Step 7: Convert output to int
        output = output.astype('int')

        # Step 8: Draw results on original image
        outseg, corrected_data = drawOnTop(original_img, output, original_shape)
        
    except Exception as e:
        print(f"Error in segmentation: {e}")
        # Return a basic error response
        return None, None, 0, f"Error: {str(e)}"

    seg_to_save = (outseg.copy() * 255).astype('uint8')
    cv2.imwrite("tmp/overlap_segmentation.png", cv2.cvtColor(seg_to_save, cv2.COLOR_RGB2BGR))

    # Step 9: Robust CTR calculation
    ctr_result = calculate_ctr_robust(output, corrected_data)
    ctr_value = ctr_result['ctr']
    tilt_angle = ctr_result['tilt_angle']
    
    # Enhanced interpretation with quality indicators
    interpretation_parts = []
    
    # CTR interpretation
    if ctr_value < 0.5:
        base_interpretation = "Normal"
    elif 0.50 <= ctr_value <= 0.55:
        base_interpretation = "Mild Cardiomegaly (CTR 50-55%)"
    elif 0.56 <= ctr_value <= 0.60:
        base_interpretation = "Moderate Cardiomegaly (CTR 56-60%)"
    elif ctr_value > 0.60:
        base_interpretation = "Severe Cardiomegaly (CTR > 60%)"
    else:
        base_interpretation = "Cardiomegaly"
    
    interpretation_parts.append(base_interpretation)
    
    # Add quality indicators
    if was_rotated:
        interpretation_parts.append(f"Image rotation corrected ({detected_rotation:.1f}°)")
    
    if tilt_angle > 3 and not ctr_result['correction_applied']:
        interpretation_parts.append(f"Residual tilt detected ({tilt_angle:.1f}°)")
    
    final_interpretation = " | ".join(interpretation_parts)

    return outseg, "tmp/overlap_segmentation.png", ctr_value, final_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()