app.py

import os
import threading
import traceback
import gradio as gr
import torch
from transformers import AutoProcessor, Gemma3nForConditionalGeneration, TextIteratorStreamer
from PIL import Image
import inspect
import traceback
import spaces

# -----------------------------
# Config
# -----------------------------
MODEL_ID = "yasserrmd/GemmaECG-Vision"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32   # safe CPU dtype

# Generation defaults
GEN_KW = dict(
    max_new_tokens=1024,
    do_sample=True,
    temperature=1.0,
    top_p=0.95,
    top_k=64,
    use_cache=True,
)

# Clinical prompt
CLINICAL_PROMPT = """You are a clinical assistant specialized in ECG interpretation. 
you should not answer like below instead reply "Please provide valid ecg" if image doesnt have ECG 
Given an ECG image, generate a concise, structured, and medically accurate report.

Use this exact format:

Rhythm:
PR Interval:
QRS Duration:
Axis:
Bundle Branch Blocks:
Atrial Abnormalities:
Ventricular Hypertrophy:
Q Wave or QS Complexes:
T Wave Abnormalities:
ST Segment Changes:
Final Impression:

Guidance:
- Confirm sinus rhythm only if consistent P waves precede each QRS.
- Describe PACs only if early, ectopic P waves are visible.
- Do not diagnose myocardial infarction solely based on QS complexes unless accompanied by other signs (e.g., ST elevation, reciprocal changes, poor R wave progression).
- Only mention axis deviation if QRS axis is clearly rightward (RAD) or leftward (LAD).
- Use terms like "suggestive of" or "possible" for uncertain findings.
- Avoid repetition and keep the report clinically focused.
- Do not include external references or source citations.
- Do not diagnose left bundle branch block unless QRS duration is ≥120 ms with typical morphology in leads I, V5, V6.
- Mark T wave changes in inferior leads as “nonspecific” unless clear ST elevation or reciprocal depression is present.

Your goal is to provide a structured ECG summary useful for a cardiologist or internal medicine physician.
"""

# -----------------------------
# Load model & processor
# -----------------------------
model = Gemma3nForConditionalGeneration.from_pretrained(
    MODEL_ID, torch_dtype=DTYPE
).to(DEVICE).eval()
processor = AutoProcessor.from_pretrained(MODEL_ID)



# -----------------------------
# Inference (single-shot)
# -----------------------------
@spaces.GPU
def analyze_ecg(image: Image.Image):
    if image is None:
        return "Please upload an ECG image."

    # Chat messages (preferred) -> insert image via template
    messages = [
        {"role": "user", "content": [
            {"type": "text", "text": CLINICAL_PROMPT},
            {"type": "image"},
        ]}
    ]

    try:
        chat_text = processor.apply_chat_template(messages, add_generation_prompt=True)
        inputs = processor(text=chat_text, images=image, return_tensors="pt")
    except Exception:
        # Fallback if image-token count/template mismatches
        inputs = processor(text=CLINICAL_PROMPT, images=image, return_tensors="pt")

    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    try:
        with torch.inference_mode():
            out = model.generate(**inputs, **GEN_KW)

        # Decode only new tokens if possible
        input_len = inputs.get("input_ids", None)
        if input_len is not None:
            in_len = inputs["input_ids"].shape[-1]
            gen_ids = out[0][in_len:] if out.shape[-1] > in_len else out[0]
        else:
            gen_ids = out[0]

        text = processor.decode(gen_ids, skip_special_tokens=True).strip()
        return text if text else "[Empty output] Try a clearer ECG image."
    except Exception as e:
        return f"[Generation Error]\n{traceback.format_exc()}"

def reset():
    return None, ""

# -----------------------------
# UI (accessible contrast)
# -----------------------------
theme = gr.themes.Soft(primary_hue="indigo", neutral_hue="slate")
css = """
/* High-contrast output area */
textarea[aria-label="Generated ECG Report"] {
  background-color: #ffffff !important;
  color: #111827 !important;   /* near-black */
}

/* Card look */
.card { background:#fff; border:1px solid #e5e7eb; border-radius:14px; padding:16px; }

/* Buttons */
.gr-button { background-color:#1e3a8a !important; color:#ffffff !important; }

/* Disclaimer with strong contrast */
.disclaimer {
  margin-top:12px; padding:12px 14px; border-radius:12px;
  background:#fff7ed; color:#7c2d12; border:1px solid #fdba74; font-weight:600;
}

/* Header */
.header {
  display:flex; align-items:center; justify-content:space-between;
  padding:14px 16px; border-radius:14px; background:#0f172a; color:#e5e7eb;
}
.header .brand { font-weight:800; letter-spacing:.2px; }
"""

with gr.Blocks(theme=theme, css=css) as demo:
    gr.HTML(f"""
    <div class="header">
      <div class="brand">🫀 ECG Interpretation Assistant</div>
      <div>Model: <code>{MODEL_ID}</code></div>
    </div>
    <div class="disclaimer">
      ⚠️ Education & Research Only: This tool is not a medical device and must not be used for diagnosis or treatment.
      Always consult a licensed clinician for interpretation and clinical decisions.
    </div>
    """)

    with gr.Row(equal_height=True):
        with gr.Column(scale=1):
            with gr.Group(elem_classes="card"):
                image_input = gr.Image(type="pil", label="Upload ECG Image", height=360)
                with gr.Row():
                    submit_btn = gr.Button("Generate Report", variant="primary")
                    reset_btn = gr.Button("Reset")

        with gr.Column(scale=2):
            with gr.Group(elem_classes="card"):
                output_box = gr.Textbox(
                    label="Generated ECG Report",
                    lines=28,
                    show_copy_button=True,
                    placeholder="The model's report will appear here…",
                )

    gr.Markdown(
        "Tip: Use a high-resolution ECG with visible lead labels to improve assessment of P waves and ST segments."
    )

    submit_btn.click(analyze_ecg, inputs=image_input, outputs=output_box, queue=False)
    reset_btn.click(reset, outputs=[image_input, output_box])

demo.launch(share=False, debug=True)