""" 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("""

🩻 Chest X-ray Recommender

AI-powered visual search across radiology studies

""") gr.HTML("""

📖 About this app

This tool helps you find chest X-rays that look similar to your query. The catalog draws on MIMIC-CXR, 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 CLIP (a multimodal AI model) and compared against pre-computed embeddings using cosine similarity.

Use cases: medical education, comparative case lookup, exploring how visual AI represents medical imagery.

""") gr.HTML("""
⚠️ Educational demo only. 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.
""") 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('
') 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('
') 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"""

🎬 Walk-through video

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