app.py

import spaces
import gradio as gr
import cv2
import numpy as np
from PIL import Image as PILImage
import os
from huggingface_hub import login

login(token=os.environ.get("HF_TOKEN"))

from models import (load_yolo, load_unet, load_efficientnet,
                    load_medgemma, load_fusion, DEVICE)
from pipeline import predict_glaucoma
from visualize import make_visualization_images
from medgemma_report import generate_report
from export import generate_pdf_report

print("Loading models...")
yolo_model                     = load_yolo("weights/best.pt")
unet_model                     = load_unet("weights/best_model_final.pth")
eff_model                      = load_efficientnet("weights/best_classifier_efficientnet.pth")
fusion_model, fusion_threshold = load_fusion("fusion_model.pkl", "fusion_threshold.pkl")
medgemma_model, processor      = load_medgemma("google/medgemma-1.5-4b-it")
print("All models loaded!")


@spaces.GPU
def run_pipeline(image_path: str):
    if image_path is None:
        return [None] * 9 + [_placeholder_diagnosis_html(), _placeholder_table_html(), "", None, gr.update(interactive=False)]

    image_bgr = cv2.imread(image_path)

    try:
        result = predict_glaucoma(
            image_path, yolo_model, unet_model, eff_model,
            fusion_model, fusion_threshold, DEVICE
        )
    except ValueError as e:
        return [None] * 9 + [f'<p id="decision-result" style="color:#F87171;">{e}</p>', _placeholder_table_html(), "", None, gr.update(interactive=False)]

    viz = make_visualization_images(image_bgr, result, eff_model, DEVICE)
    result["crop_rgb"]    = viz["crop_rgb"]
    result["overlay_rgb"] = viz["overlay_rgb"]

    report = generate_report(result, medgemma_model, processor)

    metrics = {
        "YOLO conf":    result["yolo_conf"],
        "Seg failed":   result["seg_failed"],
        "vCDR":         result["vCDR"],
        "hCDR":         result["hCDR"],
        "Inferior rim": result["SI_ratios"]["inferior"],
        "Superior rim": result["SI_ratios"]["superior"],
        "SI compliant": result["SI_compliant"],
        "NRR area":     result["NRR_area"],
        "p holistic":   result["p_holistic"],
        "p fused":      result["p_fused"],
        "Threshold":    result["fusion_threshold"],
    }

    export_data = {
        "decision":    result["decision"],
        "metrics":     metrics,
        "report":      report,
        "input_image": PILImage.open(image_path),
        "images": {
            "YOLO Detection":         PILImage.fromarray(viz["img_yolo"]),
            "Disc Crop":              PILImage.fromarray(viz["crop_rgb"]),
            "U-Net Input":            PILImage.fromarray(viz["img_proc"]),
            "U-Net Segmentation":     PILImage.fromarray(viz["seg_rgb"]),
            "Overlay":                PILImage.fromarray(viz["overlay_rgb"]),
            "EfficientNet Input":     PILImage.fromarray(viz["img_eff_input"]),
            "Preprocessed":           PILImage.fromarray(viz["img_eff_proc"]),
            "GradCAM (preprocessed)": PILImage.fromarray(viz["gradcam_proc"]),
            "GradCAM (original)":     PILImage.fromarray(viz["gradcam_orig"]),
        },
    }

    return (
        PILImage.fromarray(viz["img_yolo"]),
        PILImage.fromarray(viz["crop_rgb"]),
        PILImage.fromarray(viz["img_proc"]),
        PILImage.fromarray(viz["seg_rgb"]),
        PILImage.fromarray(viz["overlay_rgb"]),
        PILImage.fromarray(viz["img_eff_input"]),
        PILImage.fromarray(viz["img_eff_proc"]),
        PILImage.fromarray(viz["gradcam_proc"]),
        PILImage.fromarray(viz["gradcam_orig"]),
        _decision_to_html(result["decision"]),
        _table_to_html(metrics),
        _report_to_html(report),
        export_data,
        gr.update(interactive=True),
    )


TOOLTIPS = {
    "YOLO conf":    "Detection confidence score from the YOLOv8 optic disc detector. Reflects the certainty of the disc localization, which affects the quality of all downstream biomarker computations. Lower values may indicate poor image quality.",
    "Seg failed":   "When flagged, U-Net disc/cup segmentation was unreliable and classifier scores carry more weight in the decision.",
    "vCDR":         "Vertical Cup-to-Disc Ratio is the ratio between the vertical diameters of the optic cup and disc, measuring the vertical extent of excavation. A key structural biomarker for glaucoma assessment.",
    "hCDR":         "Horizontal Cup-to-Disc Ratio is the ratio between the horizontal diameters of the optic cup and disc, measuring the horizontal extent of excavation. Used as a complementary structural biomarker alongside vCDR.",
    "Inferior rim": "Proportion of the inferior disc quadrant occupied by the neuroretinal rim. The inferior rim should be the thickest sector in a healthy disc. Thinning here is one of the earliest structural glaucoma signs.",
    "Superior rim": "Proportion of the superior disc quadrant occupied by the neuroretinal rim. The superior rim is the second thickest in a healthy disc. Thinning in this region, accompanied by inferior rim loss, is an indicator of glaucomatous damage.",
    "SI compliant": "Inferior-Superior rim compliance based on the ISNT rule, a clinical guideline stating that in healthy eyes neuroretinal rim thickness follows Inferior > Superior > Nasal > Temporal. Only the Inferior > Superior axis is assessed here since left/right eye laterality is unavailable in the datasets used, making the nasal/temporal distinction impossible.",
    "NRR area":     "Neuroretinal Rim Area Ratio measures the proportion of the optic disc occupied by the rim tissue (disc minus cup). A smaller ratio indicates a larger cup relative to the disc, which may suggest glaucomatous structural damage.",
    "p holistic":   "Holistic glaucoma probability estimated by an EfficientNet-B3 classifier trained on the full fundus image. Captures global image features beyond the optic disc region.",
    "p fused":      "Combined glaucoma probability integrating both the structural branch (segmentation metrics) and the holistic branch (EfficientNet). Provides a single calibrated score for the final decision.",
    "Threshold":    "The probability threshold used to classify an eye as glaucoma suspected. Predictions above this value are flagged as positive. Selected to optimize the balance between sensitivity and specificity on the validation set.",
}

LABELS = {
    "vCDR":         "Vertical CDR",
    "hCDR":         "Horizontal CDR",
    "NRR area":     "NRR Area Ratio",
    "p holistic":   "EfficientNet Probability",
    "p fused":      "Fused Glaucoma Score",
    "Threshold":    "Decision Threshold",
    "SI compliant": "Inferior > Superior Rim Rule",
    "Inferior rim": "Inferior Rim Ratio",
    "Superior rim": "Superior Rim Ratio",
    "Seg failed":   "Segmentation Failed",
    "YOLO conf":    "YOLO Confidence",
}


def _placeholder_diagnosis_html() -> str:
    return '<p id="decision-result" style="color:#0C4A6E;">—</p>'


def _decision_to_html(decision: str) -> str:
    return f'<p id="decision-result" style="color:#FFFFFF;">{decision}</p>'


def _table_to_html(m: dict) -> str:
    def fmt(v):
        if v is True:  return "Yes"
        if v is False: return "No"
        if isinstance(v, float): return f"{v:.4f}"
        return v
    rows = "".join(
        f'<tr>'
        f'<td class="metric-label">'
        f'<span class="tip-anchor">{LABELS.get(k, k)}'
        f'<span class="tip-box">{TOOLTIPS.get(k, "")}</span>'
        f'</span></td>'
        f'<td class="metric-value">{fmt(v)}</td>'
        f'</tr>'
        for k, v in m.items()
    )
    return _table_frame(rows)


def _placeholder_table_html() -> str:
    rows = "".join(
        f'<tr>'
        f'<td class="metric-label">'
        f'<span class="tip-anchor">{label}'
        f'<span class="tip-box">{TOOLTIPS.get(k, "")}</span>'
        f'</span></td>'
        f'<td class="metric-value placeholder-val">—</td>'
        f'</tr>'
        for k, label in LABELS.items()
    )
    return _table_frame(rows)


def _table_frame(rows: str) -> str:
    return f"""
    <div id="table-frame">
      <table id="metrics-table">
        <thead><tr>
          <th>Biomarker / Score</th>
          <th>Value</th>
        </tr></thead>
        <tbody>{rows}</tbody>
      </table>
    </div>
    """


def _report_to_html(text: str) -> str:
    import html as _html
    escaped = _html.escape(text) if text else ""
    return f'<pre id="report-pre">{escaped}</pre>'


# Accent: sky blue (#38BDF8 / #7DD3FC) on the existing dark navy/black backgrounds.
# Much higher contrast than dark purple; still clinical and calm.
CUSTOM_CSS = """
/* ── Base font scale ─────────────────────────────────────── */
html { font-size: 22px !important; }

/* ── Layout ─────────────────────────────────────────────── */
.gradio-container { width: 100% !important; max-width: 1280px !important; margin: 0 auto !important; padding: 0 1.5rem !important; box-sizing: border-box !important; }

/* ── Banner wrapper: remove Gradio block padding ─────────── */
#header-block,
#header-block > div {
    padding: 0 !important;
}

/* ── Banner ─────────────────────────────────────────────── */
#glaunet-header {
    background: #080614;
    background-image:
        radial-gradient(ellipse 80% 60% at 35% -10%, rgba(56,189,248,0.18) 0%, transparent 65%),
        radial-gradient(circle at 1px 1px, rgba(56,189,248,0.07) 1px, transparent 0);
    background-size: auto, 24px 24px;
    border: 1px solid #0C2A3E;
    border-radius: 16px;
    padding: 1.75rem 2.25rem;
    margin-bottom: 3.5rem;
    position: relative;
    overflow: hidden;
}

#glaunet-header::after {
    content: '';
    position: absolute;
    right: 8%;
    top: 50%;
    transform: translateY(-50%);
    width: 280px;
    height: 280px;
    background: radial-gradient(circle, rgba(56,189,248,0.08) 0%, transparent 70%);
    border-radius: 50%;
    pointer-events: none;
}

.hdr-inner {
    display: flex;
    align-items: center;
    justify-content: space-between;
    gap: 2rem;
    position: relative;
    z-index: 1;
}

.hdr-left { display: flex; align-items: center; gap: 1.25rem; }

.hdr-icon {
    width: 85px;
    height: 85px;
    background: rgba(56,189,248,0.10);
    border: 1px solid rgba(56,189,248,0.35);
    border-radius: 16px;
    display: flex;
    align-items: center;
    justify-content: center;
    flex-shrink: 0;
    box-shadow: 0 0 28px rgba(56,189,248,0.18), inset 0 1px 0 rgba(255,255,255,0.06);
}

.hdr-text {
    display: flex;
    flex-direction: column;
    align-items: flex-start;
    gap: 0;
}

.hdr-title {
    font-family: 'Syne', 'IBM Plex Sans', sans-serif;
    font-size: 2.6rem;
    font-weight: 800;
    line-height: 1;
    color: #F0F9FF;
    letter-spacing: -0.03em;
    margin: 0;
    transform: translateY(-0.3em);
}

.hdr-accent { color: #38BDF8; }

.hdr-sub {
    font-size: 0.68rem;
    color: #38BDF8;
    margin: 0;
    letter-spacing: 0.18em;
    text-transform: uppercase;
    font-family: 'JetBrains Mono', monospace;
    font-weight: 600;
}

.hdr-tagline {
    font-size: 0.8rem;
    color: #7DD3FC;
    max-width: 260px;
    text-align: right;
    line-height: 1.55;
    letter-spacing: 0.01em;
    font-family: 'IBM Plex Sans', sans-serif;
}

/* ── Controls ────────────────────────────────────────────── */
#btn-analyze {
    width: 100% !important;
    height: 48px !important;
    font-size: 1rem !important;
    font-weight: 600 !important;
    letter-spacing: 0.04em !important;
    margin-top: 0.75rem !important;
}

/* ── Section Labels ──────────────────────────────────────── */
.section-label {
    font-size: 0.6rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.16em;
    color: #38BDF8;
    padding: 0.3rem 0 0.3rem 0.65rem;
    margin: 0 0 1rem;
    border-left: 2px solid #38BDF8;
    font-family: 'JetBrains Mono', monospace;
    display: block;
}

/* ── Banner → input row gap ──────────────────────────────── */
#top-row {
    padding-top: 2rem !important;
}

/* ── Section spacing: padding-top on the Gradio HTML wrapper ─ */
#sec-biomarkers {
    padding-top: 2.25rem !important;
}

#sec-segmentation,
#sec-gradcam,
#sec-report {
    padding-top: 1.5rem !important;
}

/* ── Decision Result ─────────────────────────────────────── */
#decision-result {
    font-size: 1rem;
    font-weight: 700;
    margin: 0 0 0.85rem 0;
    letter-spacing: 0.02em;
    text-align: center;
}

/* ── Metrics Table ───────────────────────────────────────── */
#metrics-wrapper { width: 100%; }

#table-frame {
    width: 100% !important;
    border: 1px solid rgba(255,255,255,0.15);
    border-radius: 10px;
    overflow: visible;
}

/* Replicate rounded corners on cells so the table looks correct without overflow:hidden */
#metrics-table thead tr:first-child th:first-child { border-top-left-radius: 9px; }
#metrics-table thead tr:first-child th:last-child  { border-top-right-radius: 9px; }
#metrics-table tbody tr:last-child td:first-child  { border-bottom-left-radius: 9px; }
#metrics-table tbody tr:last-child td:last-child   { border-bottom-right-radius: 9px; }

#metrics-table {
    width: 100%;
    border-collapse: collapse;
    font-size: 0.82rem;
    background: #080614;
}

#metrics-table th {
    background: #060F1A;
    color: #38BDF8;
    text-align: left;
    padding: 0.6rem 1rem;
    font-weight: 700;
    font-size: 0.59rem;
    letter-spacing: 0.14em;
    text-transform: uppercase;
    font-family: 'JetBrains Mono', monospace;
    border-bottom: 1px solid rgba(255,255,255,0.15);
}

#metrics-table td {
    padding: 0.45rem 1rem;
    border-bottom: 1px solid rgba(255,255,255,0.15);
}

td.metric-label {
    color: #7DD3FC;
    font-size: 0.79rem;
}

td.metric-value {
    font-family: 'JetBrains Mono', monospace;
    color: #BAE6FD;
    font-size: 0.81rem;
}

#metrics-table tr:nth-child(even) td { background: rgba(6,15,26,0.55); }
#metrics-table tr:last-child td { border-bottom: none; }

#metrics-table tr:hover td.metric-label { color: #7DD3FC; }
#metrics-table tr:hover td.metric-value { color: #F0F9FF; }
#metrics-table tr:hover td { background: rgba(56,189,248,0.06) !important; }

/* ── Tooltips ────────────────────────────────────────────── */
.tip-anchor {
    position: relative;
    cursor: help;
    border-bottom: 1px dashed rgba(56,189,248,0.40);
    padding-bottom: 1px;
    display: inline-block;
}

.tip-box {
    display: none;
    position: absolute;
    left: 105%;
    top: 50%;
    transform: translateY(-50%);
    z-index: 9999;
    width: 230px;
    background: #080614;
    border: 1px solid #0C4A6E;
    border-radius: 8px;
    padding: 0.5rem 0.75rem;
    font-size: 0.71rem;
    color: #7DD3FC;
    line-height: 1.65;
    box-shadow: 0 8px 28px rgba(0,0,0,0.65), 0 0 0 1px rgba(56,189,248,0.10);
    pointer-events: none;
    white-space: normal;
    font-family: 'IBM Plex Sans', sans-serif;
    font-style: normal;
    font-weight: 400;
    border-bottom-width: 1px;
}

.tip-anchor:hover .tip-box { display: block; }


/* ── Clinical Report ─────────────────────────────────────── */
#report-pre {
    background: #080614;
    border: 1px solid #0C2A3E;
    border-radius: 10px;
    color: #7DD3FC;
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.79rem;
    line-height: 1.85;
    padding: 1.1rem 1.3rem;
    margin: 0;
    white-space: pre-wrap;
    word-break: break-word;
    min-height: 200px;
}

/* ── Image rows: CSS grid, 3 cols each so images are always the same size ── */
#seg-row,
#gradcam-row {
    display: grid !important;
    grid-template-columns: repeat(3, minmax(0, 1fr)) !important;
    gap: 0.5rem !important;
    margin-bottom: 1rem !important;
}

#seg-row .block,
#gradcam-row .block {
    min-width: 0 !important;
    max-width: 100% !important;
    width: 100% !important;
}

/* Gradio 6: override fixed block height so the square image-frame can expand */
#seg-row .block,
#gradcam-row .block {
    height: auto !important;
}

/* .image-container has height:100% of the block; free it to grow with contents */
#seg-row .image-container,
#gradcam-row .image-container {
    height: auto !important;
}

/* Make .image-frame a square: width fills the column, height matches via aspect-ratio */
#seg-row .image-frame,
#gradcam-row .image-frame {
    width: 100% !important;
    height: auto !important;
    aspect-ratio: 1 !important;
}

/* Placeholder (no image yet): .empty renders instead of .image-frame.
   Apply the same square constraint so the block is square before analysis too. */
#seg-row .empty,
#gradcam-row .empty {
    width: 100% !important;
    height: auto !important;
    min-height: 0 !important;
    aspect-ratio: 1 !important;
}

#seg-row .image-frame img,
#seg-row img,
#gradcam-row .image-frame img,
#gradcam-row img {
    object-fit: cover !important;
    width: 100% !important;
    height: 100% !important;
}

/* ── Hide placeholder icon in output image blocks ────────── */
#seg-row .svelte-p3y7hu,
#gradcam-row .svelte-p3y7hu,
#seg-row .wrap > svg,
#gradcam-row .wrap > svg,
#seg-row .icon,
#gradcam-row .icon,
#seg-row .placeholder,
#gradcam-row .placeholder {
    display: none !important;
}

/* ── Input column: same 1/3 width as one output image cell ── */
#input-col {
    flex: 0 0 calc(100% / 3) !important;
    max-width: calc(100% / 3) !important;
}

/* ── Input image: same square approach as output images ── */
#img-input {
    height: auto !important;
}
#img-input .image-container {
    height: auto !important;
}
/* Make .upload-container square in both empty and loaded states.
   It is the common ancestor of .image-frame (loaded) and UploadText (empty),
   so aspect-ratio here covers both without relying on slot-child selectors. */
#img-input .upload-container {
    width: 100% !important;
    height: auto !important;
    aspect-ratio: 1 !important;
}
#img-input .image-frame {
    width: 100% !important;
    height: 100% !important;
}
#img-input .image-frame img,
#img-input img {
    object-fit: cover !important;
    width: 100% !important;
    height: 100% !important;
}
/* Override UploadText's min-height so it doesn't stretch beyond the square */
#img-input .wrap {
    min-height: 0 !important;
    height: 100% !important;
}

/* ── Diagnosis column ────────────────────────────────────── */
#diagnosis-col {
    min-width: 0 !important;
}

/* ── Force Gradio Svelte wrappers to full width ──────────── */
div:has(> #metrics-html-block),
div:has(> #diagnosis-html-block) {
    display: block !important;
    width: 100% !important;
    min-width: 0 !important;
}

#metrics-html-block {
    width: 100% !important;
    min-width: 0 !important;
    max-width: 100% !important;
    margin-bottom: 1rem !important;
}

/* ── Placeholder values in table ─────────────────────────── */
.placeholder-val {
    color: #0C4A6E !important;
    font-style: italic;
}

/* ── Footer ──────────────────────────────────────────────── */
#footer-note {
    text-align: center;
    font-size: 0.7rem;
    color: #0C4A6E;
    padding: 1rem 0 0.5rem;
    border-top: 1px solid #060F1A;
    margin-top: 1.5rem;
    letter-spacing: 0.04em;
    font-family: 'JetBrains Mono', monospace;
}

/* ── Export button ───────────────────────────────────────────── */
#export-row {
    margin-top: 1.5rem !important;
    margin-bottom: 0.5rem !important;
}

#btn-export {
    width: 100% !important;
    height: 44px !important;
    font-size: 0.9rem !important;
    font-weight: 600 !important;
    letter-spacing: 0.06em !important;
    border: 1px solid rgba(56,189,248,0.35) !important;
    background: rgba(56,189,248,0.08) !important;
    color: #7DD3FC !important;
    transition: background 0.15s, border-color 0.15s !important;
}

#btn-export:not(:disabled):hover {
    background: rgba(56,189,248,0.16) !important;
    border-color: #38BDF8 !important;
    color: #F0F9FF !important;
}

#btn-export:disabled { opacity: 0.35 !important; cursor: not-allowed !important; }
"""

sky = gr.themes.Color(
    c50="#F0F9FF", c100="#E0F2FE", c200="#BAE6FD", c300="#7DD3FC",
    c400="#38BDF8", c500="#0EA5E9", c600="#0284C7", c700="#0369A1",
    c800="#075985", c900="#0C4A6E", c950="#082F49",
)

theme = gr.themes.Base(
    primary_hue=sky,
    neutral_hue=gr.themes.colors.slate,
    font=gr.themes.GoogleFont("IBM Plex Sans"),
    font_mono=gr.themes.GoogleFont("JetBrains Mono"),
).set(
    button_primary_background_fill="#0284C7",
    button_primary_background_fill_hover="#0369A1",
    button_primary_text_color="#FFFFFF",
    button_primary_shadow="0 4px 14px rgba(2,132,199,0.45)",
    body_background_fill="#0A0716",
    background_fill_primary="#0F0B1E",
    background_fill_secondary="#160F2E",
    block_background_fill="#0F0B1E",
    block_border_color="#0C2A3E",
    block_border_width="1px",
    block_radius="12px",
    block_shadow="0 2px 12px rgba(0,0,0,0.5)",
    block_label_background_fill="#060F1A",
    block_label_text_color="#38BDF8",
    block_label_text_weight="600",
    input_background_fill="#080614",
    input_border_color="#0C2A3E",
    input_border_color_focus="#38BDF8",
    body_text_color="#BAE6FD",
    body_text_color_subdued="#0C4A6E",
)

with gr.Blocks(title="GlauNet — Glaucoma Screening", theme=theme) as demo:

    gr.HTML("""
    <link href="https://fonts.googleapis.com/css2?family=Syne:wght@700;800&family=IBM+Plex+Sans:wght@400;600;700&display=swap" rel="stylesheet">
    <div id="glaunet-header">
      <div class="hdr-inner">
        <div class="hdr-left">
          <div class="hdr-icon">
            <svg width="32" height="32" viewBox="0 0 30 30" fill="none" xmlns="http://www.w3.org/2000/svg">
              <path d="M1.5 15C1.5 15 7 4.5 15 4.5C23 4.5 28.5 15 28.5 15C28.5 15 23 25.5 15 25.5C7 25.5 1.5 15 1.5 15Z"
                    stroke="#38BDF8" stroke-width="1.4" fill="none" stroke-linejoin="round"/>
              <circle cx="15" cy="15" r="5" stroke="#38BDF8" stroke-width="1.4" fill="rgba(56,189,248,0.10)"/>
              <circle cx="15" cy="15" r="1.8" fill="#38BDF8"/>
              <line x1="15" y1="8.5" x2="15" y2="11"  stroke="#38BDF8" stroke-width="1" opacity="0.45"/>
              <line x1="15" y1="19" x2="15" y2="21.5" stroke="#38BDF8" stroke-width="1" opacity="0.45"/>
              <line x1="8.5" y1="15" x2="11" y2="15"  stroke="#38BDF8" stroke-width="1" opacity="0.45"/>
              <line x1="19" y1="15" x2="21.5" y2="15" stroke="#38BDF8" stroke-width="1" opacity="0.45"/>
            </svg>
          </div>
          <div class="hdr-text">
            <h1 class="hdr-title">Glau<span class="hdr-accent">NET</span></h1>
            <p class="hdr-sub">Glaucoma Screening</p>
          </div>
        </div>
        <p class="hdr-tagline">Retinal image analysis<br>for early glaucoma detection</p>
      </div>
    </div>
    """, elem_id="header-block")

    with gr.Row(equal_height=False, elem_id="top-row"):
        with gr.Column(min_width=260, scale=1, elem_id="input-col"):
            img_input = gr.Image(type="filepath", label="Fundus image (JPG / PNG)", elem_id="img-input", sources=["upload"], height=260)
            btn_analyze = gr.Button("Analyze", variant="primary", size="lg", elem_id="btn-analyze")

        with gr.Column(scale=3, elem_id="diagnosis-col"):
            gr.HTML('<span class="section-label" style="margin-top:0.5rem;">Diagnosis</span>')
            diagnosis_out = gr.HTML(value=_placeholder_diagnosis_html(), elem_id="diagnosis-html-block")

    gr.HTML('<span class="section-label">Clinical Biomarkers</span>', elem_id="sec-biomarkers")
    metrics_out = gr.HTML(value=_placeholder_table_html(), elem_id="metrics-html-block")

    gr.HTML('<span class="section-label">Segmentation Pipeline</span>', elem_id="sec-segmentation")
    with gr.Row(elem_id="seg-row"):
        out_yolo    = gr.Image(label="YOLO Detection",     height=260)
        out_crop    = gr.Image(label="Disc Crop",          height=260)
        out_proc    = gr.Image(label="U-Net Input",        height=260)
        out_seg     = gr.Image(label="U-Net Segmentation", height=260)
        out_overlay = gr.Image(label="Overlay",            height=260)

    gr.HTML('<span class="section-label">GradCAM Analysis</span>', elem_id="sec-gradcam")
    with gr.Row(elem_id="gradcam-row"):
        out_eff_input = gr.Image(label="EfficientNet Input",     height=260)
        out_eff_proc  = gr.Image(label="Preprocessed",           height=260)
        out_gc_proc   = gr.Image(label="GradCAM (preprocessed)", height=260)
        out_gc_orig   = gr.Image(label="GradCAM (original)",     height=260)

    gr.HTML('<span class="section-label">MedGemma Clinical Report</span>', elem_id="sec-report")
    report_out = gr.HTML()

    result_state = gr.State(None)

    with gr.Row(elem_id="export-row"):
        btn_export = gr.DownloadButton(
            label="Export Report (PDF)",
            elem_id="btn-export",
            interactive=False,
        )

    gr.HTML('<div id="footer-note">Research prototype. Not a substitute for diagnosis by a qualified ophthalmologist.</div>')

    btn_analyze.click(
        fn=run_pipeline,
        inputs=[img_input],
        outputs=[
            out_yolo, out_crop, out_proc, out_seg, out_overlay,
            out_eff_input, out_eff_proc, out_gc_proc, out_gc_orig,
            diagnosis_out, metrics_out, report_out,
            result_state, btn_export,
        ],
    )

    def _do_export(export_data):
        if export_data is None:
            return None
        return generate_pdf_report(export_data)

    btn_export.click(fn=_do_export, inputs=[result_state], outputs=[btn_export])

if __name__ == "__main__":
    demo.launch(css=CUSTOM_CSS)