Update app.py

app.py

@@ -1,5 +1,7 @@
 # app.py
 import io
+import os
+
 import faiss
 import torch
 import pandas as pd
@@ -24,20 +26,34 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ---------- Load index + metadata + model at startup ----------
+# ---------- Config ----------
+# Dataset with FAISS index + metadata
+EMBED_REPO_ID = "saad003/Red01"
+# Dataset with raw radiology images
+IMAGE_REPO_ID = "saad003/images"
+
+# Base URL for images (you uploaded ROCOv2_*.jpg directly in the root)
+BASE_IMAGE_URL = f"https://huggingface.co/datasets/{IMAGE_REPO_ID}/resolve/main"
 
-REPO_ID = "saad003/Red01"  # your dataset repo
+# Optional: token if EMBED_REPO_ID is private (set HF_TOKEN secret in Space)
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN is None:
+    print("⚠️  No HF_TOKEN env var found. If the dataset is private, this may fail.")
 
+# ---------- Load index + metadata ----------
 print("Downloading FAISS index & metadata...")
 INDEX_PATH = hf_hub_download(
-    repo_id=REPO_ID,
+    repo_id=EMBED_REPO_ID,
     filename="radiology_index.faiss",
     repo_type="dataset",
+    token=HF_TOKEN,
 )
+
 META_PATH = hf_hub_download(
-    repo_id=REPO_ID,
+    repo_id=EMBED_REPO_ID,
     filename="radiology_metadata.csv",
     repo_type="dataset",
+    token=HF_TOKEN,
 )
 
 print("Loading FAISS index...")
@@ -46,9 +62,11 @@ index = faiss.read_index(INDEX_PATH)
 print("Loading metadata CSV...")
 metadata = pd.read_csv(META_PATH)
 
+# ---------- Load CLIP model ----------
 print("Loading CLIP model...")
 MODEL_NAME = "flaviagiammarino/pubmed-clip-vit-base-patch32"
 device = "cuda" if torch.cuda.is_available() else "cpu"
+print("Using device:", device)
 
 clip_model = CLIPModel.from_pretrained(MODEL_NAME).to(device)
 clip_processor = CLIPProcessor.from_pretrained(MODEL_NAME)
@@ -58,28 +76,44 @@ print("Backend ready ✅")
 
 
 # ---------- Helper: search by image ----------
-
-def _search_similar_by_image(image: Image.Image, k: int = 5):
+def _search_similar_by_image(image: Image.Image, k: int = 5) -> pd.DataFrame:
+    """
+    Encode query image with CLIP, search FAISS, return top-k rows
+    with ID, split, caption, concepts, score, and image_url.
+    """
     inputs = clip_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
         feats = clip_model.get_image_features(**inputs)
 
+    # L2-normalize to match index normalization
     feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
     feats = feats.cpu().numpy().astype("float32")
 
-    D, I = index.search(feats, k)
+    D, I = index.search(feats, k)  # D: similarity scores, I: indices
 
     rows = metadata.iloc[I[0]].copy()
     rows["score"] = D[0]
-    # Only send useful columns
-    return rows[["ID", "split", "img_path", "caption", "concepts_manual", "score"]]
 
+    # Build image URL for each retrieved ID
+    # Files are named like ROCOv2_2023_test_000001.jpg in saad003/images
+    rows["image_url"] = rows["ID"].apply(
+        lambda id_str: f"{BASE_IMAGE_URL}/{id_str}.jpg"
+    )
 
-# ---------- API endpoint ----------
+    return rows[["ID", "split", "caption", "concepts_manual", "score", "image_url"]]
+
+
+# ---------- Routes ----------
+@app.get("/")
+def root():
+    return {"status": "ok", "message": "Radiology retrieval API"}
 
 @app.post("/search_by_image")
 async def search_by_image(file: UploadFile = File(...), k: int = 5):
-    # read image from request
+    """
+    Accepts an uploaded radiology image.
+    Returns top-k similar images (ID, caption, concepts, score, image_url).
+    """
     content = await file.read()
     image = Image.open(io.BytesIO(content)).convert("RGB")
 
@@ -87,8 +121,3 @@ async def search_by_image(file: UploadFile = File(...), k: int = 5):
     results = results_df.to_dict(orient="records")
 
     return JSONResponse({"results": results})
-
-
-@app.get("/")
-def root():
-    return {"status": "ok", "message": "Radiology retrieval API"}