Update app.py

app.py

@@ -30,14 +30,14 @@ app.add_middleware(
 # Dataset with FAISS index + radiology_metadata.csv
 EMBED_REPO_ID = "saad003/Red01"
 
-# Dataset with ALL radiology images (flat, filenames = ID + ".jpg")
-IMAGE_REPO_ID = "saad003/images"
+# NEW dataset with images organized into subfolders
+# test, valid, train01, train02, ..., train07
+IMAGE_REPO_ID = "saad003/images04"
 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")
+HF_TOKEN = os.environ.get("HF_TOKEN")  # set in HF Space secrets if needed
 
 # ---------- Download index + metadata ----------
 print("Downloading FAISS index & metadata from Hugging Face...")
@@ -62,7 +62,6 @@ 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:
@@ -95,112 +94,157 @@ print("Backend ready ✅")
 
 
 # ---------- Helpers ----------
+def train_folder_from_id(image_id: str) -> str:
+  """
+  For IDs like 'ROCOv2_2023_train_000001', decide which trainXX folder.
+  Uses numeric ranges based on the last 6 digits.
+  """
+  try:
+    num_str = image_id.split("_")[-1]  # "000001"
+    num = int(num_str)
+  except Exception:
+    return "train01"  # safe default
+
+  if num <= 9000:
+    return "train01"
+  elif num <= 18000:
+    return "train02"
+  elif num <= 27000:
+    return "train03"
+  elif num <= 36000:
+    return "train04"
+  elif num <= 45000:
+    return "train05"
+  elif num <= 54000:
+    return "train06"
+  else:
+    return "train07"
+
+
 def id_to_image_url(image_id: str) -> str:
-    """
-    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"
+  """
+  Build raw image URL based on ID and folder structure.
+
+  Examples:
+    ROCOv2_2023_test_000001  -> test/ROCOv2_2023_test_000001.jpg
+    ROCOv2_2023_valid_000005 -> valid/ROCOv2_2023_valid_000005.jpg
+    ROCOv2_2023_train_000001 -> train01/ROCOv2_2023_train_000001.jpg
+    ROCOv2_2023_train_009001 -> train02/ROCOv2_2023_train_009001.jpg
+  """
+  if not isinstance(image_id, str):
+    return None
+
+  image_id = image_id.strip()
+
+  if "test_" in image_id:
+    folder = "test"
+  elif "valid_" in image_id:
+    folder = "valid"
+  elif "train_" in image_id:
+    folder = train_folder_from_id(image_id)
+  else:
+    # Fallback: put directly at root (in case of weird ID)
+    folder = ""
+
+  filename = f"{image_id}.jpg"
+
+  if folder:
+    return f"{BASE_IMAGE_URL}/{folder}/{filename}"
+  else:
     return f"{BASE_IMAGE_URL}/{filename}"
 
 
 def search_similar_by_image(image: Image.Image, k: int = 5) -> pd.DataFrame:
-    """
-    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():
-        feats = clip_model.get_image_features(**inputs)
+  """
+  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():
+    feats = clip_model.get_image_features(**inputs)
 
-    feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
-    feats = feats.cpu().numpy().astype("float32")
+  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)
 
-    rows = metadata.iloc[I[0]].copy()
-    rows["score"] = D[0]
-    rows["image_url"] = rows["ID"].apply(id_to_image_url)
+  rows = metadata.iloc[I[0]].copy()
+  rows["score"] = D[0]
+  rows["image_url"] = rows["ID"].apply(id_to_image_url)
 
-    return rows[
-        ["vec_index", "ID", "caption", "concepts_manual", "score", "image_url"]
-    ]
+  return rows[
+    ["vec_index", "ID", "caption", "concepts_manual", "score", "image_url"]
+  ]
 
 
 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)
-    caption = caption_processor.batch_decode(out, skip_special_tokens=True)[0]
-    return caption.strip()
+  """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)
+  caption = caption_processor.batch_decode(out, skip_special_tokens=True)[0]
+  return caption.strip()
 
 
 def infer_modality_from_caption(caption: str) -> str:
-    """Heuristic to infer modality from caption text."""
-    if not caption:
-        return "Unknown"
-
-    text = caption.lower()
-
-    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"
-
+  """Heuristic to infer modality from caption text."""
+  if not caption:
     return "Unknown"
 
+  text = caption.lower()
+
+  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"
+
+  return "Unknown"
+
 
 # ---------- Routes ----------
 @app.get("/")
 def root():
-    return {"status": "ok", "message": "Radiology retrieval + captioning API"}
+  return {"status": "ok", "message": "Radiology retrieval + captioning API"}
 
 
 @app.post("/search_by_image")
 async def search_by_image(file: UploadFile = File(...), k: int = 5):
-    """
-    Upload a radiology image.
-
-    Returns:
-      - 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
-    """
-    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(
-        {
-            "query_caption": query_caption,
-            "modality": modality,
-            "results": results,
-        }
-    )
+  """
+  Upload a radiology image.
+
+  Returns:
+    - 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
+  """
+  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
+  modality = infer_modality_from_caption(query_caption or "")
+
+  return JSONResponse(
+    {
+      "query_caption": query_caption,
+      "modality": modality,
+      "results": results,
+    }
+  )