Update app.py

app.py

@@ -1,15 +1,14 @@
 # app.py
 import os
 import json
+import torch
 import gradio as gr
-from transformers import pipeline
+from PIL import Image
 from huggingface_hub import login
+from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
 
 # ============= 🔐 AUTHENTICATION =============
-# Hugging Face token (set in your Space settings)
-# Go to your Space → Settings → Repository secrets → add key: HF_TOKEN, value: your access token
 hf_token = os.getenv("HF_TOKEN")
-
 if hf_token:
     login(token=hf_token)
 else:
@@ -17,26 +16,45 @@ else:
 
 # ============= 📦 LOAD LABELS =============
 with open("labels.json", "r", encoding="utf-8") as f:
-    candidate_labels = json.load(f)
-
-# ============= 🧠 MODEL PIPELINE =============
-# Use the authenticated pipeline for MedSigLIP (gated model)
-pipe = pipeline(
-    "zero-shot-image-classification",
-    model="google/medsiglip-448",
-    use_auth_token=hf_token
-)
+    all_labels = json.load(f)
+
+# ============= 🧠 MODEL & PROCESSOR =============
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# ============= ⚙️ PREDICTION FUNCTION =============
-def classify_medical_image(image):
-    """
-    Run zero-shot classification using MedSigLIP.
-    """
+processor = AutoProcessor.from_pretrained("google/medsiglip-448", use_auth_token=hf_token)
+model = AutoModelForZeroShotImageClassification.from_pretrained(
+    "google/medsiglip-448",
+    use_auth_token=hf_token,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+).to(device)
+model.eval()
+
+# ============= ⚙️ INFERENCE FUNCTION =============
+def classify_medical_image(image_path):
     try:
-        results = pipe(image, candidate_labels=candidate_labels)
-        formatted = {r["label"]: round(r["score"] * 100, 2) for r in results}
-        sorted_result = dict(sorted(formatted.items(), key=lambda x: x[1], reverse=True))
-        return sorted_result
+        image = Image.open(image_path).convert("RGB")
+
+        # 🔹 Split labels into small batches to avoid memory overload
+        batch_size = 50  # adjust if needed
+        results = []
+
+        for i in range(0, len(all_labels), batch_size):
+            batch = all_labels[i:i+batch_size]
+            inputs = processor(text=batch, images=image, return_tensors="pt", padding=True).to(device)
+
+            with torch.no_grad():
+                outputs = model(**inputs)
+                logits_per_image = outputs.logits_per_image
+                probs = torch.softmax(logits_per_image, dim=1)[0]
+
+            for label, score in zip(batch, probs.tolist()):
+                results.append((label, round(score * 100, 2)))
+
+        # 🔹 Sort final results
+        results.sort(key=lambda x: x[1], reverse=True)
+        top5 = dict(results[:5])
+        return top5
+
     except Exception as e:
         return {"Error": str(e)}
 
@@ -47,9 +65,8 @@ demo = gr.Interface(
     outputs=gr.Label(num_top_classes=5, label="🧠 Top Predictions"),
     title="🩻 MedSigLIP Zero-Shot Medical Classifier",
     description=(
-        "This demo uses Google's **MedSigLIP (448x448)** model for zero-shot medical image understanding. "
-        "Upload a medical image (e.g., chest X-ray, skin lesion, fundus photo, histopathology slide) and the model "
-        "will compute similarity against known medical findings loaded from `labels.json`."
+        "Efficient version using Google's **MedSigLIP (448x448)** model for medical image understanding. "
+        "Optimized for Hugging Face CPU Spaces. Uses batched label processing to reduce memory load."
     ),
     examples=[
         ["https://storage.googleapis.com/dx-scin-public-data/dataset/images/3445096909671059178.png"],
@@ -60,4 +77,4 @@ demo = gr.Interface(
 
 # ============= 🚀 RUN APP =============
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860, queue=True)