app.py

"""
ClearScan - Patient Care Navigator for Medical Imaging
A patient-first AI care navigator that helps people understand their
medical imaging results in plain language.

Powered by MedGemma from Google Health AI Developer Foundations.
"""

import gradio as gr
from PIL import Image
import torch
from transformers import pipeline
import spaces
import os
from huggingface_hub import login

# Authenticate with HuggingFace
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("HuggingFace authentication successful")
else:
    print("Warning: HF_TOKEN not found in environment")

# Model configuration
MODEL_ID = "google/medgemma-1.5-4b-it"

# Global model variable
pipe = None

def load_model():
    """Load MedGemma model."""
    global pipe
    if pipe is not None:
        return pipe

    try:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

        print(f"Loading model on device: {device}, dtype: {dtype}")
        print(f"CUDA available: {torch.cuda.is_available()}")

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

        pipe = pipeline(
            "image-text-to-text",
            model=MODEL_ID,
            torch_dtype=dtype,
            device_map="auto" if device == "cuda" else None,
            token=token,
        )
        print("Model loaded successfully!")
        return pipe
    except Exception as e:
        print(f"Model loading error: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        return None


def build_explanation_prompt(context: str, report_text: str) -> str:
    """Build prompt for patient-friendly explanation."""
    return f"""You are a patient care navigator helping someone understand their medical imaging results.

CRITICAL RULES:
- Do NOT provide any diagnosis
- Do NOT recommend specific treatments
- Do NOT contradict what their doctor has said
- Write at an 8th-grade reading level
- Be empathetic and reassuring while being honest
- Always encourage discussion with their healthcare provider

Context from patient: {context if context else "No additional context provided."}

Radiology report or findings: {report_text if report_text else "Please analyze the image provided."}

Please provide a clear explanation with these sections:

**What This Scan Examines**: Briefly explain what type of imaging this is and what body area it looks at.

**What The Report Describes**: Explain the key findings in simple terms a patient can understand.

**What's Typically Normal vs. Notable**: Help them understand what findings are commonly seen versus what might need monitoring.

Remember: Be clear, compassionate, and always emphasize discussing with their doctor."""


def build_questions_prompt(context: str, report_text: str) -> str:
    """Build prompt for doctor questions."""
    return f"""You are helping a patient prepare for a conversation with their doctor about their imaging results.

Context: {context if context else "No additional context provided."}

Findings: {report_text if report_text else "Based on the imaging provided."}

Generate 5-8 thoughtful questions the patient could ask their doctor. Focus on:
- Understanding if findings typically progress or stay stable
- What symptoms to watch for
- Whether follow-up imaging is standard
- Alternative approaches or treatment options if applicable
- Timeline for any recommended monitoring

Format each question clearly with a bullet point. These should help the patient have a productive conversation with their healthcare provider."""


def build_decision_support_prompt(context: str, report_text: str) -> str:
    """Build prompt for non-diagnostic decision support."""
    return f"""You are a patient care navigator providing non-diagnostic guidance.

CRITICAL: You must NOT diagnose or recommend specific treatments. You can only provide general information about common care pathways.

Context: {context if context else "No additional context provided."}

Findings: {report_text if report_text else "Based on the imaging provided."}

Provide brief, non-diagnostic guidance addressing:
1. Is this type of finding commonly monitored over time, or typically addressed differently?
2. Are second opinions sometimes sought for findings like this?
3. What general categories of questions might be worth exploring with a specialist?

Be supportive but clear that all decisions should be made with their healthcare team."""


def run_inference(model_pipe, image, prompt: str) -> str:
    """Run inference with MedGemma."""
    try:
        content = []
        if image is not None:
            content.append({"type": "image", "image": image})
        content.append({"type": "text", "text": prompt})

        messages = [{"role": "user", "content": content}]
        output = model_pipe(text=messages, max_new_tokens=1500)

        if output and len(output) > 0:
            generated = output[0].get("generated_text", [])
            if generated and len(generated) > 0:
                last_message = generated[-1]
                if isinstance(last_message, dict):
                    return last_message.get("content", "Unable to generate response.")
                return str(last_message)
        return "Unable to generate response."
    except Exception as e:
        return f"Error: {str(e)}"


@spaces.GPU
def analyze_imaging(image, report_text, context):
    """Main analysis function - processes image and/or report text."""

    if image is None and not report_text:
        return (
            "⚠️ Please upload an image or paste your radiology report text.",
            "",
            ""
        )

    # Load model
    model = load_model()
    if model is None:
        error_msg = """⚠️ Unable to load MedGemma model. This could be due to:
• Need to accept model license at: https://huggingface.co/google/medgemma-1.5-4b-it
• Insufficient GPU memory
• Model download in progress (please wait and retry)"""
        return error_msg, "", ""

    # Generate explanation
    explanation_prompt = build_explanation_prompt(context, report_text)
    explanation = run_inference(model, image, explanation_prompt)

    # Generate questions
    questions_prompt = build_questions_prompt(context, report_text)
    questions = run_inference(model, image, questions_prompt)

    # Generate decision support
    support_prompt = build_decision_support_prompt(context, report_text)
    decision_support = run_inference(model, image, support_prompt)

    return explanation, questions, decision_support


# Build Gradio Interface
with gr.Blocks(
    title="ClearScan - Medical Imaging Navigator",
    theme=gr.themes.Soft(primary_hue="green")
) as demo:

    gr.Markdown("""
    # 🔬 ClearScan
    ### Your Patient Care Navigator for Medical Imaging
    """)

    gr.Markdown("""
    > ⚠️ **Important**: ClearScan helps you understand and discuss your imaging results with your doctor.
    > It does **NOT** provide medical diagnosis, treatment recommendations, or replace professional medical advice.
    > Always discuss your results with your healthcare provider.
    """)

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📤 Upload Your Information")

            image_input = gr.Image(
                label="Upload Medical Image (MRI, CT, X-ray, Ultrasound)",
                type="pil"
            )

            report_input = gr.Textbox(
                label="Paste Your Radiology Report Text",
                placeholder="Paste the text from your radiology report here...",
                lines=8
            )

            context_input = gr.Textbox(
                label="What did your doctor say this scan was for? (Optional)",
                placeholder="e.g., 'Checking for back pain' or 'Follow-up on previous finding'"
            )

            analyze_btn = gr.Button("🔍 Help Me Understand", variant="primary", size="lg")

        with gr.Column(scale=1):
            gr.Markdown("### 📋 Your Results")

            explanation_output = gr.Textbox(
                label="📖 Understanding Your Results",
                lines=10,
                interactive=False
            )

            questions_output = gr.Textbox(
                label="❓ Questions for Your Doctor",
                lines=8,
                interactive=False
            )

            decision_output = gr.Textbox(
                label="🧭 Next Steps to Consider",
                lines=6,
                interactive=False
            )

    analyze_btn.click(
        fn=analyze_imaging,
        inputs=[image_input, report_input, context_input],
        outputs=[explanation_output, questions_output, decision_output]
    )

    gr.Markdown("""
    ---
    **Disclaimer**: ClearScan is an educational tool designed to help you understand medical imaging concepts
    and prepare for conversations with your healthcare provider. It does not provide medical diagnosis,
    treatment recommendations, or clinical advice. Always seek the advice of your physician or other
    qualified health provider with any questions you may have regarding a medical condition.

    *Powered by MedGemma from Google Health AI Developer Foundations | Built for the MedGemma Impact Challenge*
    """)

# Launch
if __name__ == "__main__":
    demo.launch()