Update app.py

app.py

@@ -1,7 +1,6 @@
 # app.py
 import io
 import os
-import base64
 
 import faiss
 import torch
@@ -21,20 +20,23 @@ app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # you can restrict later
+    allow_origins=["*"],  # later you can restrict to your frontend domain
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
 # ---------- Config ----------
-# FAISS index + radiology_metadata.csv
+# Dataset with FAISS index + radiology_metadata.csv
 EMBED_REPO_ID = "saad003/Red01"
 
-# All radiology images, filenames like ROCOv2_2023_valid_000001.jpg
+# Dataset with ALL radiology images (flat, filenames = ID + ".jpg")
 IMAGE_REPO_ID = "saad003/images"
-BASE_IMAGE_URL = f"https://huggingface.co/datasets/{IMAGE_REPO_ID}/resolve/main"
+BASE_IMAGE_URL = (
+    f"https://huggingface.co/datasets/{IMAGE_REPO_ID}/resolve/main"
+)
 
+# Optional token (if Red01 / images are private). Set HF_TOKEN in Space secrets.
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # ---------- Download index + metadata ----------
@@ -60,6 +62,7 @@ index = faiss.read_index(INDEX_PATH)
 print("Loading metadata CSV...")
 metadata = pd.read_csv(META_PATH)
 
+# We only need these columns
 required_cols = {"vec_index", "ID", "caption", "concepts_manual"}
 missing = required_cols - set(metadata.columns)
 if missing:
@@ -83,55 +86,34 @@ print("Loading BLIP radiology captioning model...")
 CAPTION_MODEL_ID = "WafaaFraih/blip-roco-radiology-captioning"
 
 caption_processor = AutoProcessor.from_pretrained(CAPTION_MODEL_ID)
-caption_model = BlipForConditionalGeneration.from_pretrained(CAPTION_MODEL_ID).to(device)
+caption_model = BlipForConditionalGeneration.from_pretrained(
+    CAPTION_MODEL_ID
+).to(device)
 caption_model.eval()
 
 print("Backend ready ✅")
 
 
-# ---------- Helpers for images ----------
+# ---------- Helpers ----------
 def id_to_image_url(image_id: str) -> str:
-    """Public HF URL (optional, for debugging/click)."""
-    if not isinstance(image_id, str):
-        return None
-    filename = f"{image_id}.jpg"
-    return f"{BASE_IMAGE_URL}/{filename}"
-
-
-def id_to_image_base64(image_id: str) -> str | None:
     """
-    Download the image from `saad003/images` (cached by hf_hub_download),
-    then return base64-encoded bytes so frontend can display directly.
+    Build raw image URL.
+
+    Example:
+      ID = "ROCOv2_2023_test_000040"
+      -> https://huggingface.co/datasets/saad003/images/resolve/main/ROCOv2_2023_test_000040.jpg
     """
     if not isinstance(image_id, str):
         return None
-
+    image_id = image_id.strip()
     filename = f"{image_id}.jpg"
-    try:
-        local_path = hf_hub_download(
-            repo_id=IMAGE_REPO_ID,
-            filename=filename,
-            repo_type="dataset",
-            token=HF_TOKEN,
-        )
-    except Exception as e:
-        print(f"Error downloading image for ID={image_id}: {e}")
-        return None
-
-    try:
-        with open(local_path, "rb") as f:
-            data = f.read()
-        return base64.b64encode(data).decode("utf-8")
-    except Exception as e:
-        print(f"Error reading image file for ID={image_id}: {e}")
-        return None
+    return f"{BASE_IMAGE_URL}/{filename}"
 
 
-# ---------- Retrieval ----------
 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 vec_index, ID, caption, concepts_manual, score, image_url, image_base64.
+    Encode query image with CLIP, search FAISS, and return top-k rows
+    with vec_index, ID, caption, concepts_manual, score, image_url.
     """
     inputs = clip_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
@@ -144,17 +126,15 @@ def search_similar_by_image(image: Image.Image, k: int = 5) -> pd.DataFrame:
 
     rows = metadata.iloc[I[0]].copy()
     rows["score"] = D[0]
-
     rows["image_url"] = rows["ID"].apply(id_to_image_url)
-    rows["image_base64"] = rows["ID"].apply(id_to_image_base64)
 
     return rows[
-        ["vec_index", "ID", "caption", "concepts_manual", "score", "image_url", "image_base64"]
+        ["vec_index", "ID", "caption", "concepts_manual", "score", "image_url"]
     ]
 
 
-# ---------- Captioning ----------
 def generate_query_caption(image: Image.Image) -> str:
+    """Generate a medical caption for the query image using BLIP."""
     inputs = caption_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
         out = caption_model.generate(**inputs, max_new_tokens=64)
@@ -163,6 +143,7 @@ def generate_query_caption(image: Image.Image) -> str:
 
 
 def infer_modality_from_caption(caption: str) -> str:
+    """Heuristic to infer modality from caption text."""
     if not caption:
         return "Unknown"
 
@@ -197,21 +178,23 @@ async def search_by_image(file: UploadFile = File(...), k: int = 5):
       - query_caption: BLIP caption for query image
       - modality: inferred imaging modality
       - results: list of similar images with
-          vec_index, ID, concepts_manual, score,
-          image_url, image_base64
+          vec_index, ID, concepts_manual, score, image_url
     """
     content = await file.read()
     image = Image.open(io.BytesIO(content)).convert("RGB")
 
+    # 1) Retrieval
     results_df = search_similar_by_image(image, k=k)
     results = results_df.to_dict(orient="records")
 
+    # 2) Caption for query image
     try:
         query_caption = generate_query_caption(image)
     except Exception as e:
         print("Error generating caption:", e)
         query_caption = None
 
+    # 3) Modality from caption
     modality = infer_modality_from_caption(query_caption or "")
 
     return JSONResponse(