app.py

"""
RadAssist-MedGemma — HuggingFace Space (Showcase Demo)
=======================================================
This is a CPU-only showcase demo with pre-computed example outputs.
For live inference, see the Kaggle notebook linked below.

MedGemma Impact Challenge Submission by AIencoder
"""

import gradio as gr
from PIL import Image
import requests
from io import BytesIO

# ============================================================
# PRE-COMPUTED DEMO RESULTS
# These were generated using MedGemma 1.5 4B on a T4 GPU
# ============================================================

DEMO_CASES = {
    "Normal Chest X-Ray": {
        "image_url": "https://upload.wikimedia.org/wikipedia/commons/c/c8/Chest_Xray_PA_3-8-2010.png",
        "severity": "NORMAL",
        "confidence": 0.92,
        "key_finding": "No acute cardiopulmonary abnormality",
        "triage_summary": """## 🟢 Triage Result: **NORMAL**

- **Confidence:** 92%
- **Key Finding:** No acute cardiopulmonary abnormality
- **Urgency:** Routine — no immediate action required""",
        "analysis": """### Systematic Chest X-Ray Analysis

**Technical Quality:** The radiograph is a well-penetrated PA view with adequate inspiration. No significant rotation is noted.

**Airway:** The trachea is midline. The main bronchi appear patent without evidence of narrowing or deviation.

**Bones & Soft Tissues:** The osseous structures of the thorax appear intact without fractures or lytic lesions. The soft tissues are unremarkable. No subcutaneous emphysema is identified.

**Cardiac:** The cardiac silhouette is within normal limits in size and configuration. The cardiothoracic ratio is within the normal range. The mediastinal contours appear normal. The aortic knob is unremarkable.

**Diaphragm:** Both hemidiaphragms are smooth and well-defined. The costophrenic angles are sharp bilaterally, without evidence of pleural effusion.

**Lungs:** The lung fields are clear bilaterally. There are no focal consolidations, masses, or nodules identified. The pulmonary vascularity appears normal. The hila are symmetric and unremarkable. No pneumothorax is seen.""",
        "report": """**CHEST X-RAY REPORT**

**Clinical Indication:** Routine screening / baseline imaging

**Technique:** PA chest radiograph

**Findings:**
The lungs are well-expanded and clear bilaterally. No focal airspace opacities, masses, or nodules are identified. The pulmonary vasculature appears normal. The cardiac silhouette is normal in size with a cardiothoracic ratio within normal limits. The mediastinal contours are unremarkable. The trachea is midline. Both hemidiaphragms are smooth, and the costophrenic angles are sharp. No pleural effusion or pneumothorax is identified. The osseous structures are intact.

**Impression:**
1. No acute cardiopulmonary abnormality identified.
2. Normal cardiac silhouette size.
3. Clear lung fields bilaterally.

**Recommendation:**
No further imaging follow-up is recommended based on these findings. Routine follow-up as clinically indicated.""",
    },
    "Pneumonia (Abnormal)": {
        "image_url": None,
        "severity": "ABNORMAL",
        "confidence": 0.87,
        "key_finding": "Right lower lobe consolidation consistent with pneumonia",
        "triage_summary": """## 🟡 Triage Result: **ABNORMAL**

- **Confidence:** 87%
- **Key Finding:** Right lower lobe consolidation consistent with pneumonia
- **Urgency:** Follow-up required — clinical correlation recommended""",
        "analysis": """### Systematic Chest X-Ray Analysis

**Technical Quality:** Adequate PA radiograph with good penetration and inspiration.

**Airway:** The trachea is midline. No airway deviation or narrowing.

**Bones & Soft Tissues:** No acute osseous abnormalities. Soft tissues are unremarkable.

**Cardiac:** The cardiac silhouette is at the upper limits of normal size. The mediastinal contours are within normal limits.

**Diaphragm:** The left hemidiaphragm is well-defined with a sharp costophrenic angle. The right costophrenic angle shows mild blunting, suggesting a small right-sided pleural effusion.

**Lungs:** There is a focal area of airspace opacity in the right lower lobe, demonstrating an air bronchogram sign, consistent with consolidation. The right middle lobe and upper lobe are clear. The left lung is clear throughout. No masses, nodules, or cavitary lesions are identified. No pneumothorax is seen.""",
        "report": """**CHEST X-RAY REPORT**

**Clinical Indication:** Cough and fever — evaluate for pneumonia

**Technique:** PA chest radiograph

**Findings:**
There is a focal area of airspace consolidation within the right lower lobe with visible air bronchograms, consistent with a pneumonic infiltrate. Mild blunting of the right costophrenic angle suggests a small associated parapneumonic effusion. The remaining lung fields are clear. The cardiac silhouette is at the upper limits of normal. The mediastinum is midline and unremarkable. No pneumothorax.

**Impression:**
1. Right lower lobe consolidation, most consistent with community-acquired pneumonia.
2. Small right-sided pleural effusion, likely parapneumonic in nature.

**Recommendation:**
Clinical correlation with laboratory findings (CBC, CRP, blood cultures) is recommended. Consider follow-up chest radiograph in 4-6 weeks to document resolution. If clinical concern for complicated effusion, lateral decubitus view or ultrasound may be considered.""",
    },
    "Critical Finding (Pneumothorax)": {
        "image_url": None,
        "severity": "CRITICAL",
        "confidence": 0.94,
        "key_finding": "Large left-sided tension pneumothorax with mediastinal shift",
        "triage_summary": """## 🔴 Triage Result: **CRITICAL**

- **Confidence:** 94%
- **Key Finding:** Large left-sided tension pneumothorax with mediastinal shift
- **Urgency:** IMMEDIATE — requires emergent chest tube placement""",
        "analysis": """### Systematic Chest X-Ray Analysis

**Technical Quality:** AP portable radiograph, slightly rotated. Despite technical limitations, critical findings are clearly identifiable.

**Airway:** The trachea is deviated to the RIGHT, indicating mediastinal shift away from the left hemithorax. This is a critical finding.

**Bones & Soft Tissues:** No acute fractures identified. Subcutaneous emphysema is noted in the left lateral chest wall.

**Cardiac:** The cardiac silhouette appears shifted to the right. Assessment of true cardiac size is limited by mediastinal displacement.

**Diaphragm:** The left hemidiaphragm appears flattened and depressed. The right hemidiaphragm appears normal.

**Lungs:** There is complete absence of lung markings in the left hemithorax with a visible visceral pleural line, consistent with a large left-sided pneumothorax. The left lung appears significantly collapsed toward the hilum. The right lung shows mildly increased vascular markings, likely from compensatory hyperinflation. No focal consolidation in the right lung.""",
        "report": """**CHEST X-RAY REPORT**

**Clinical Indication:** Acute dyspnea and chest pain — STAT portable

**Technique:** AP portable chest radiograph

**Findings:**
There is a large left-sided pneumothorax with near-complete collapse of the left lung. The visceral pleural line is clearly visible. There is rightward deviation of the trachea and mediastinal structures, consistent with tension physiology. The left hemidiaphragm is depressed and flattened. Subcutaneous emphysema is present in the left lateral chest wall. The right lung appears hyperinflated but clear without focal consolidation.

**Impression:**
1. **CRITICAL: Large left-sided tension pneumothorax with mediastinal shift** — requires emergent decompression.
2. Subcutaneous emphysema of the left chest wall.
3. Near-complete left lung collapse.

**Recommendation:**
EMERGENT needle decompression followed by chest tube placement is indicated. Immediate clinical intervention required. This finding has been communicated urgently.""",
    },
}

# ============================================================
# SEVERITY BADGE HTML
# ============================================================
SEVERITY_STYLES = {
    "NORMAL": {"emoji": "🟢", "color": "#22c55e", "label": "Normal"},
    "ABNORMAL": {"emoji": "🟡", "color": "#f59e0b", "label": "Abnormal — Follow Up"},
    "CRITICAL": {"emoji": "🔴", "color": "#ef4444", "label": "Critical — Immediate"},
}


def make_badge(severity, confidence):
    s = SEVERITY_STYLES.get(severity, SEVERITY_STYLES["ABNORMAL"])
    return f"""
    <div style="
        background: {s['color']}15;
        border: 2px solid {s['color']};
        border-radius: 16px;
        padding: 20px 32px;
        text-align: center;
        margin: 12px 0;
    ">
        <span style="font-size: 56px;">{s['emoji']}</span>
        <h2 style="margin: 8px 0 4px; color: {s['color']}; font-size: 1.5em;">{s['label']}</h2>
        <p style="margin: 0; opacity: 0.7; font-size: 1.1em;">Confidence: {confidence}%</p>
    </div>
    """


# ============================================================
# DEMO HANDLER
# ============================================================
def run_demo(case_name):
    """Load pre-computed results for the selected demo case."""
    if case_name not in DEMO_CASES:
        return (
            "<p>Select a demo case above.</p>",
            "", "", "",
        )

    case = DEMO_CASES[case_name]
    badge = make_badge(case["severity"], int(case["confidence"] * 100))

    return (
        badge,
        case["triage_summary"],
        case["analysis"],
        case["report"],
    )


def load_demo_image(case_name):
    """Load the demo image for the selected case."""
    if case_name not in DEMO_CASES:
        return None
    url = DEMO_CASES[case_name].get("image_url")
    if url:
        try:
            resp = requests.get(url, headers={"User-Agent": "RadAssist-MedGemma"}, timeout=10)
            return Image.open(BytesIO(resp.content))
        except Exception:
            return None
    return None


# ============================================================
# GRADIO UI
# ============================================================
with gr.Blocks(
    title="RadAssist-MedGemma",
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="slate"),
    css="""
    .hero-section { text-align: center; padding: 20px 0; }
    .pipeline-diagram {
        background: linear-gradient(135deg, #1e3a5f20, #3b82f620);
        border-radius: 12px; padding: 20px; margin: 16px 0;
        font-family: monospace; font-size: 0.95em;
    }
    .info-badge {
        display: inline-block; background: #3b82f615; border: 1px solid #3b82f650;
        border-radius: 8px; padding: 4px 12px; margin: 2px 4px; font-size: 0.9em;
    }
    """,
) as demo:

    # --- Hero Section ---
    gr.Markdown("""
    <div class="hero-section">

    # 🏥 RadAssist-MedGemma
    ### AI Radiology Triage Assistant — MedGemma Impact Challenge

    <span class="info-badge">🧠 MedGemma 1.5 4B</span>
    <span class="info-badge">🔗 3-Step Agentic Pipeline</span>
    <span class="info-badge">🔒 Privacy-Preserving (Local)</span>
    <span class="info-badge">⚡ Edge-Deployable (4-bit)</span>

    </div>
    """)

    gr.Markdown("""
    > **How it works:** RadAssist-MedGemma uses a 3-step agentic pipeline where MedGemma 1.5 4B
    > assumes three distinct clinical roles in sequence — Triage Specialist → Diagnostic Radiologist → Report Generator —
    > with context passing between steps for coherent, comprehensive analysis.

    <div class="pipeline-diagram">
    📤 Upload CXR → 🚨 <b>Step 1: TRIAGE</b> (severity classification) → 🔍 <b>Step 2: ANALYZE</b> (systematic findings) → 📋 <b>Step 3: REPORT</b> (formal radiology report)
    </div>

    ⬇️ **Select a demo case below** to see the pipeline output. For live inference, see the
    [Kaggle Notebook — Live Inference](https://www.kaggle.com/code/ary5272/radassist-medgemma-ai-radiology-triage-assistant).
    """)

    gr.Markdown("---")

    # --- Demo Selection ---
    with gr.Row():
        with gr.Column(scale=1):
            # Using gr.Dropdown instead of gr.Radio to avoid Gradio 5.x SSR bug
            # (TypeError: argument of type 'bool' is not iterable in json_schema_to_python_type)
            case_selector = gr.Dropdown(
                choices=list(DEMO_CASES.keys()),
                value="Normal Chest X-Ray",
                label="📂 Select Demo Case",
                info="Pre-computed results from MedGemma 1.5 4B on T4 GPU",
            )
            demo_image = gr.Image(
                label="🩻 Chest X-Ray",
                height=380,
                interactive=False,
            )
            run_btn = gr.Button(
                "🔬 View Pipeline Results",
                variant="primary",
                size="lg",
            )

            gr.Markdown("""
            <div style="background: #fef3c7; border: 1px solid #f59e0b;
                 border-radius: 8px; padding: 12px; margin: 12px 0; font-size: 0.9em;">
            ⚠️ <b>Research prototype only.</b> Not for clinical diagnosis.
            All AI outputs require review by qualified medical professionals.
            </div>
            """)

        # --- Results Panel ---
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("🚨 Triage"):
                    triage_badge = gr.HTML(
                        "<p style='text-align:center; opacity:0.5; padding: 40px;'>Select a case and click 'View Pipeline Results'</p>"
                    )
                    triage_text = gr.Markdown("")

                with gr.Tab("🔍 Detailed Analysis"):
                    analysis_output = gr.Markdown(
                        "*Select a case to view the systematic analysis*"
                    )

                with gr.Tab("📋 Radiology Report"):
                    report_output = gr.Markdown(
                        "*Select a case to view the generated report*"
                    )

    # --- Event Handlers ---
    run_btn.click(
        fn=run_demo,
        inputs=[case_selector],
        outputs=[triage_badge, triage_text, analysis_output, report_output],
    )

    case_selector.change(
        fn=load_demo_image,
        inputs=[case_selector],
        outputs=[demo_image],
    )

    # --- About Section ---
    gr.Markdown("---")
    gr.Markdown("""
    ## 🏗️ Architecture & Technical Details

    | Component | Details |
    |-----------|---------|
    | **Model** | [MedGemma 1.5 4B IT](https://huggingface.co/google/medgemma-1.5-4b-it) (4.3B params, multimodal) |
    | **Pipeline** | 3-step agentic: Triage → Analysis → Report with context chaining |
    | **Edge Support** | 4-bit NF4 quantization via bitsandbytes (~3GB VRAM) |
    | **Inference** | Deterministic (`do_sample=False`) for clinical reproducibility |
    | **UI** | Gradio with progressive result streaming |

    ### Why This Matters
    - **2/3 of the world** lacks access to diagnostic imaging interpretation (WHO)
    - **30% of critical findings** experience reporting delays in ERs
    - MedGemma 1.5 4B runs **locally on a single GPU**, keeping patient data private
    - The agentic pipeline mimics real clinical workflows for higher quality output

    ### Links
    - 📓 [Kaggle Notebook — Live Inference](https://www.kaggle.com/code/ary5272/radassist-medgemma-ai-radiology-triage-assistant)
    - 🏆 [MedGemma Impact Challenge](https://www.kaggle.com/competitions/med-gemma-impact-challenge)
    - 🤗 [MedGemma 1.5 4B on HuggingFace](https://huggingface.co/google/medgemma-1.5-4b-it)
    - 📄 [MedGemma Technical Report](https://arxiv.org/abs/2507.05201)

    ---
    *Built by [AIencoder](https://huggingface.co/AIencoder) for the MedGemma Impact Challenge 2026*
    """)

    # Load initial demo image on startup
    demo.load(
        fn=load_demo_image,
        inputs=[case_selector],
        outputs=[demo_image],
    )

demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)