app.py

"""
Chest X-ray Recommender - HuggingFace Space entry point.

Loads pre-computed CLIP embeddings (embeddings.parquet, built by the companion
notebook) and serves a Gradio UI that returns 3-5 visually similar X-rays for
a given text or image query.

Educational demo only. NOT a medical device.
"""
from __future__ import annotations

import base64
import io
import os

import gradio as gr
import numpy as np
import pandas as pd
import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------
MODEL_ID        = os.environ.get("CLIP_MODEL_ID", "openai/clip-vit-base-patch32")
EMBEDDINGS_FILE = os.environ.get("EMBEDDINGS_FILE", "embeddings.parquet")
K_MIN           = int(os.environ.get("K_MIN", "3"))
K_MAX           = int(os.environ.get("K_MAX", "5"))
GAP_THRESHOLD   = float(os.environ.get("GAP_THRESHOLD", "0.02"))

# Optional walk-through video (set the env var on your Space)
VIDEO_EMBED_ID  = os.environ.get("VIDEO_EMBED_ID", "")

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[startup] device = {device}")

# ---------------------------------------------------------------------------
# Load model
# ---------------------------------------------------------------------------
print(f"[startup] loading CLIP model: {MODEL_ID}")
clip_model     = CLIPModel.from_pretrained(MODEL_ID).to(device).eval()
clip_processor = CLIPProcessor.from_pretrained(MODEL_ID)

# ---------------------------------------------------------------------------
# Load catalog (embeddings + thumbnails + reports)
# ---------------------------------------------------------------------------
print(f"[startup] loading catalog: {EMBEDDINGS_FILE}")
df = pd.read_parquet(EMBEDDINGS_FILE)
print(f"[startup] catalog rows: {len(df):,}")

EMB_MATRIX = np.vstack(df["embedding"].values).astype("float32")
# Defensive re-normalisation (cheap, idempotent)
norms = np.linalg.norm(EMB_MATRIX, axis=1, keepdims=True)
EMB_MATRIX = EMB_MATRIX / np.where(norms == 0, 1, norms)

def _b64_to_array(b64: str) -> np.ndarray:
    """Decode a base64 JPEG thumbnail to a numpy RGB array (most reliable in Gradio)."""
    img = Image.open(io.BytesIO(base64.b64decode(b64)))
    img.load()
    return np.array(img.convert("RGB"))

THUMB_ARRAYS = [_b64_to_array(b) for b in df["image_b64"]]
REPORTS      = df["report"].fillna("").tolist()
CLUSTER      = df["cluster"].astype(int).tolist() if "cluster" in df.columns else [0] * len(df)


# ---------------------------------------------------------------------------
# CLIP encoders (version-stable: bypass get_image_features quirks)
# ---------------------------------------------------------------------------
def _to_tensor(out):
    if torch.is_tensor(out):
        return out
    if hasattr(out, "image_embeds"):
        return out.image_embeds
    if hasattr(out, "text_embeds"):
        return out.text_embeds
    if hasattr(out, "pooler_output"):
        return out.pooler_output
    if hasattr(out, "last_hidden_state"):
        return out.last_hidden_state[:, 0]
    raise TypeError(f"Cannot unwrap CLIP output of type {type(out)}")


@torch.no_grad()
def _embed_image(pil_img: Image.Image) -> np.ndarray:
    inputs = clip_processor(images=pil_img.convert("RGB"), return_tensors="pt").to(device)
    vision_out = clip_model.vision_model(pixel_values=inputs["pixel_values"])
    pooled = _to_tensor(vision_out)
    emb = clip_model.visual_projection(pooled)
    emb = emb / emb.norm(p=2, dim=-1, keepdim=True)
    return emb.cpu().numpy()[0]


@torch.no_grad()
def _embed_text(text: str) -> np.ndarray:
    inputs = clip_processor(
        text=[text], return_tensors="pt",
        padding=True, truncation=True, max_length=77,
    ).to(device)
    text_out = clip_model.text_model(
        input_ids=inputs["input_ids"],
        attention_mask=inputs.get("attention_mask"),
    )
    pooled = _to_tensor(text_out)
    emb = clip_model.text_projection(pooled)
    emb = emb / emb.norm(p=2, dim=-1, keepdim=True)
    return emb.cpu().numpy()[0]


# ---------------------------------------------------------------------------
# Adaptive top-K
# ---------------------------------------------------------------------------
def _top_k(query_vec: np.ndarray, k_min: int = K_MIN, k_max: int = K_MAX,
           gap_threshold: float = GAP_THRESHOLD):
    """Return 3-5 results: top-3 baseline, expand if consecutive scores are close."""
    scores = EMB_MATRIX @ query_vec.astype("float32")
    order  = np.argsort(-scores)

    selected = [int(i) for i in order[:k_min]]
    for i in range(k_min, min(k_max, len(order))):
        prev_score = scores[order[i - 1]]
        cand_score = scores[order[i]]
        if (prev_score - cand_score) <= gap_threshold:
            selected.append(int(order[i]))
        else:
            break

    return [
        {
            "index"   : i,
            "score"   : float(scores[i]),
            "image"   : THUMB_ARRAYS[i],
            "report"  : REPORTS[i],
            "cluster" : CLUSTER[i],
        }
        for i in selected
    ]


# ---------------------------------------------------------------------------
# Gradio handler
# ---------------------------------------------------------------------------
def recommend(text_query: str, image_query):
    if image_query is not None:
        q = _embed_image(image_query)
        used = "uploaded image"
    elif text_query and text_query.strip():
        q = _embed_text(text_query.strip())
        used = f'text query: "{text_query.strip()}"'
    else:
        return [], ("### ⚠️ No input provided\n\n"
                    "Please **upload a chest X-ray** or **type a description** "
                    "in the box on the left.")

    results = _top_k(q)
    gallery = [
        (r["image"], f"Match #{n+1} (catalog #{r['index']}) - score {r['score']:.3f}")
        for n, r in enumerate(results)
    ]

    header = f"_Query: **{used}** - showing **{len(results)}** matches"
    if len(results) > 3:
        header += " (extras included because scores are very close)_"
    else:
        header += "_"

    details = header + "\n\n" + "\n\n".join(
        f"#### Match {n+1} -  similarity {r['score']:.3f}  (cluster {r['cluster']})\n\n"
        f"```\n{r['report'][:600]}{'...' if len(r['report']) > 600 else ''}\n```"
        for n, r in enumerate(results)
    )
    return gallery, details


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CUSTOM_CSS = """
.gradio-container { max-width: 1200px !important; margin: 0 auto !important; }
.main-title {
    text-align: center;
    background: linear-gradient(135deg, #4a90e2 0%, #5e72e4 100%);
    color: white;
    padding: 30px 20px;
    border-radius: 16px;
    margin-bottom: 24px;
    box-shadow: 0 4px 12px rgba(0,0,0,0.1);
}
.main-title h1 { margin: 0; font-size: 2.2em; font-weight: 700; }
.main-title p  { margin: 8px 0 0 0; opacity: 0.95; font-size: 1.1em; }
.info-card {
    background: #f8f9fc;
    border-left: 4px solid #4a90e2;
    padding: 16px 20px;
    border-radius: 8px;
    margin: 16px 0;
}
.disclaimer-card {
    background: #fff7e6;
    border-left: 4px solid #ff9800;
    padding: 12px 16px;
    border-radius: 8px;
    margin: 16px 0;
    font-size: 0.95em;
}
.section-divider {
    border: none;
    height: 1px;
    background: linear-gradient(90deg, transparent, #d0d7de, transparent);
    margin: 24px 0;
}
"""

with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft(primary_hue="blue"),
               title="Chest X-ray Recommender") as demo:

    gr.HTML("""
    <div class="main-title">
        <h1>🩻 Chest X-ray Recommender</h1>
        <p>AI-powered visual search across radiology studies</p>
    </div>
    """)

    gr.HTML("""
    <div class="info-card">
      <h3 style="margin-top:0;">📖 About this app</h3>
      <p>This tool helps you find chest X-rays that look similar to your query.
      The catalog draws on <b>MIMIC-CXR</b>, a real dataset of 30,000+ chest X-rays
      paired with radiology reports. Each query — whether an uploaded image or a
      text description — is encoded with <b>CLIP</b> (a multimodal AI model) and
      compared against pre-computed embeddings using cosine similarity.</p>
      <p><b>Use cases:</b> medical education, comparative case lookup, exploring
      how visual AI represents medical imagery.</p>
    </div>
    """)

    gr.HTML("""
    <div class="disclaimer-card">
      ⚠️ <b>Educational demo only.</b> This is not a medical device and must not be
      used for clinical decisions. The recommendations reflect visual similarity in
      a general-purpose AI model — not medical diagnosis.
    </div>
    """)

    gr.Markdown("## 🔍 How to use this app")
    gr.Markdown("""
You have **two ways** to query the system:

**🖼️ Option A — Upload an X-ray image:** Drag and drop or click the image upload
area on the left to provide a chest X-ray. The app will encode your image and
find visually similar studies.

**📝 Option B — Describe a finding in English:** Type a clinical description in
the text box (e.g. *"right lower lobe pneumonia"*, *"pneumothorax"*,
*"clear lungs"*). The app uses CLIP's text encoder so words map to the same
vector space as the images.

Then click **Find Similar X-rays**. The app returns the **3 closest matches**,
plus up to **2 extra results** (5 total) when the scores are tightly clustered —
giving you "second opinions" when the model is uncertain.
""")

    gr.HTML('<hr class="section-divider">')

    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            gr.Markdown("### 📥 Your query")
            text_in = gr.Textbox(
                lines=3,
                label="📝 Describe a finding",
                placeholder='e.g. "bilateral pleural effusion with cardiomegaly"',
            )
            image_in = gr.Image(
                type="pil",
                label="🖼️ Upload a chest X-ray here (PNG / JPG)",
                height=300,
            )
            btn = gr.Button("🔍 Find Similar X-rays", variant="primary", size="lg")
            gr.Examples(
                examples=[
                    ["bilateral pleural effusion with cardiomegaly", None],
                    ["clear lungs, no acute cardiopulmonary process", None],
                    ["right lower lobe pneumonia", None],
                    ["pneumothorax", None],
                    ["pulmonary edema with vascular congestion", None],
                    ["enlarged cardiac silhouette", None],
                ],
                inputs=[text_in, image_in],
                label="💡 Click an example to try it",
            )

        with gr.Column(scale=2):
            gr.Markdown("### 📤 Recommended X-rays")
            gallery = gr.Gallery(
                label="Top matches (most similar first)",
                columns=3,
                height=380,
                object_fit="contain",
                show_label=True,
            )
            gr.Markdown("### 📋 Radiology reports for the matches")
            details = gr.Markdown()

    btn.click(recommend, inputs=[text_in, image_in], outputs=[gallery, details])

    gr.HTML('<hr class="section-divider">')
    gr.Markdown(f"""
### 🔬 Under the hood

- **Model:** `{MODEL_ID}` (CLIP ViT-B/32, 512-dim embeddings)
- **Catalog:** {len(df):,} X-rays from `MLforHealthcare/mimic-cxr`
- **Similarity:** Cosine similarity (via dot product of L2-normalized vectors)
- **Adaptive top-K:** {K_MIN} baseline matches, expands to {K_MAX} if score gaps ≤ {GAP_THRESHOLD}
""")

    if VIDEO_EMBED_ID:
        gr.HTML(f"""
        <hr class="section-divider">
        <h3 style="text-align:center;">🎬 Walk-through video</h3>
        <div style="display:flex; justify-content:center;">
          <iframe width="720" height="405"
                  src="https://www.youtube.com/embed/{VIDEO_EMBED_ID}"
                  title="Assignment walk-through" frameborder="0"
                  allow="autoplay; encrypted-media; picture-in-picture" allowfullscreen>
          </iframe>
        </div>
        """)

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