Update app.py

app.py

@@ -1,6 +1,7 @@
 # app.py
 import io
 import os
+import base64
 
 import faiss
 import torch
@@ -18,7 +19,6 @@ from transformers import BlipForConditionalGeneration, AutoProcessor
 # ---------- FastAPI app ----------
 app = FastAPI()
 
-# Allow your React app to call this API
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],  # you can restrict later
@@ -28,15 +28,13 @@ app.add_middleware(
 )
 
 # ---------- Config ----------
-# Dataset with FAISS index + radiology_metadata.csv
+# FAISS index + radiology_metadata.csv
 EMBED_REPO_ID = "saad003/Red01"
 
-# Dataset with ALL radiology images (flat: ID + ".jpg" in root)
-# e.g. ROCOv2_2023_valid_000001.jpg
+# All radiology images, filenames like ROCOv2_2023_valid_000001.jpg
 IMAGE_REPO_ID = "saad003/images"
-BASE_IMAGE_URL = f"https://huggingface.co/datasets/saad003/images"
+BASE_IMAGE_URL = f"https://huggingface.co/datasets/{IMAGE_REPO_ID}/resolve/main"
 
-# Optional: token if Red01 is private (set HF_TOKEN secret on the Space)
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # ---------- Download index + metadata ----------
@@ -62,7 +60,6 @@ index = faiss.read_index(INDEX_PATH)
 print("Loading metadata CSV...")
 metadata = pd.read_csv(META_PATH)
 
-# We will only rely on: vec_index, ID, caption, concepts_manual
 required_cols = {"vec_index", "ID", "caption", "concepts_manual"}
 missing = required_cols - set(metadata.columns)
 if missing:
@@ -92,53 +89,72 @@ caption_model.eval()
 print("Backend ready ✅")
 
 
-# ---------- Helper: build image URL from ID ----------
+# ---------- Helpers for images ----------
 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:
     """
-    Your images dataset `saad003/images` has files like:
-      ROCOv2_2023_valid_000001.jpg
-    where filename = ID + ".jpg".
+    Download the image from `saad003/images` (cached by hf_hub_download),
+    then return base64-encoded bytes so frontend can display directly.
     """
     if not isinstance(image_id, str):
         return None
+
     filename = f"{image_id}.jpg"
-    return f"{BASE_IMAGE_URL}/{filename}"
+    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
 
 
-# ---------- Helper: search by image ----------
+# ---------- 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
-    containing vec_index, ID, caption, concepts_manual, score, image_url.
+    with vec_index, ID, caption, concepts_manual, score, image_url, image_base64.
     """
-    # Encode image
     inputs = clip_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
         feats = clip_model.get_image_features(**inputs)
 
-    # Normalize (must match how index was built)
     feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
     feats = feats.cpu().numpy().astype("float32")
 
-    # Search FAISS
-    D, I = index.search(feats, k)  # D: distances, I: indices
+    D, I = index.search(feats, k)
 
-    # Get rows for top-k indices
     rows = metadata.iloc[I[0]].copy()
     rows["score"] = D[0]
 
-    # Build URL from ID only
     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"]]
+    return rows[
+        ["vec_index", "ID", "caption", "concepts_manual", "score", "image_url", "image_base64"]
+    ]
 
 
-# ---------- Helper: generate caption for query image ----------
+# ---------- Captioning ----------
 def generate_query_caption(image: Image.Image) -> str:
-    """
-    Generate a medical radiology caption for the query image using BLIP
-    fine-tuned on ROCO.
-    """
     inputs = caption_processor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
         out = caption_model.generate(**inputs, max_new_tokens=64)
@@ -146,11 +162,7 @@ def generate_query_caption(image: Image.Image) -> str:
     return caption.strip()
 
 
-# ---------- Helper: infer modality from caption ----------
 def infer_modality_from_caption(caption: str) -> str:
-    """
-    Simple heuristic to infer imaging modality (CT, MRI, X-ray, etc.).
-    """
     if not caption:
         return "Unknown"
 
@@ -158,16 +170,12 @@ def infer_modality_from_caption(caption: str) -> str:
 
     if any(w in text for w in ["ct scan", "ct of", "computed tomography"]):
         return "CT"
-
     if any(w in text for w in ["mri", "magnetic resonance"]):
         return "MRI"
-
     if any(w in text for w in ["x-ray", "x ray", "radiograph", "chest xray", "chest x-ray"]):
         return "X-ray"
-
     if any(w in text for w in ["ultrasound", "sonography", "sonogram"]):
         return "Ultrasound"
-
     if any(w in text for w in ["pet-ct", "pet ct", "pet scan", "positron emission tomography"]):
         return "PET/CT"
 
@@ -175,7 +183,6 @@ def infer_modality_from_caption(caption: str) -> str:
 
 
 # ---------- Routes ----------
-
 @app.get("/")
 def root():
     return {"status": "ok", "message": "Radiology retrieval + captioning API"}
@@ -187,25 +194,24 @@ async def search_by_image(file: UploadFile = File(...), k: int = 5):
     Upload a radiology image.
 
     Returns:
-      - query_caption: generated caption for the query image (BLIP)
-      - modality: inferred imaging modality from the caption
-      - results: list of similar images with vec_index, ID, concepts_manual, score, image_url
+      - 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
     """
     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) Captioning for the query image
     try:
-      query_caption = generate_query_caption(image)
+        query_caption = generate_query_caption(image)
     except Exception as e:
-      print("Error generating caption:", e)
-      query_caption = None
+        print("Error generating caption:", e)
+        query_caption = None
 
-    # 3) Infer modality
     modality = infer_modality_from_caption(query_caption or "")
 
     return JSONResponse(