Upload 3 files

app.py

@@ -0,0 +1,331 @@
+"""
+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()

requirements.txt

@@ -0,0 +1,7 @@
+gradio>=5.0.0
+transformers>=4.45.0
+torch>=2.2.0
+Pillow>=10.4.0
+numpy>=1.26.0
+pandas>=2.2.0
+pyarrow>=17.0.0