Fixed label

app.py

@@ -1,47 +1,73 @@
 import gradio as gr
 import torch
-from PIL import Image
 from transformers import ViTForImageClassification, ViTImageProcessor
+from PIL import Image
 import requests
-import json
+from io import BytesIO
 import os
 
-hf_token = os.getenv("HF_TOKEN")
+# Authenticate with Hugging Face Hub for private model access
+from huggingface_hub import login
+login(token=os.environ.get("HF_TOKEN"))  # Set this in your Space's Secrets tab
 
-
-# Load model and processor from Hugging Face Hub
-model = ViTForImageClassification.from_pretrained("patcdaniel/phytoViT_508k_20250611",token=hf_token)
-processor = ViTImageProcessor.from_pretrained("patcdaniel/phytoViT_508k_20250611",token=hf_token)
+# Load model and processor
+model = ViTForImageClassification.from_pretrained("patcdaniel/phytoViT_508k_20250611")
+processor = ViTImageProcessor.from_pretrained("patcdaniel/phytoViT_508k_20250611")
 model.eval()
 
-# Load class labels from hosted file
-LABELS_URL = "https://huggingface.co/patcdaniel/phytoViT_508k_20250611/resolve/main/label_names.json"
-headers = {"Authorization": f"Bearer {hf_token}"}
-class_labels = requests.get(LABELS_URL, headers=headers).json()
+# Use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
+# Inference function
 def predict(image):
-    image = image.convert("RGB")
-    inputs = processor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = torch.nn.functional.softmax(logits, dim=-1).squeeze()
-
-    # Get top 2 predictions
-    topk = torch.topk(probs, k=2)
-    top_scores = topk.values.tolist()
-    top_labels = [class_labels[i] for i in topk.indices.tolist()]
-    
-    # Format output
-    output = {label: round(score, 4) for label, score in zip(top_labels, top_scores)}
-    return output
-
-# Gradio interface
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload or Paste an Image"),
-    outputs=gr.Label(num_top_classes=2, label="Top Predictions"),
-    title="PhytoViT Classifier",
-    description="Upload an IFCB phytoplankton image or paste an image URL to classify it using a ViT model trained on 508k examples."
-)
+    try:
+        image = image.convert("RGB")
+        inputs = processor(images=image, return_tensors="pt").to(device)
+
+        with torch.no_grad():
+            logits = model(**inputs).logits
+            probs = torch.nn.functional.softmax(logits, dim=-1).squeeze()
+
+        topk = torch.topk(probs, k=2)
+        top_indices = topk.indices.tolist()
+        top_scores = topk.values.tolist()
+
+        id2label = model.config.id2label
+        top_labels = [id2label[str(i)] for i in top_indices]
+
+        return {label: round(score, 4) for label, score in zip(top_labels, top_scores)}
+
+    except Exception as e:
+        import traceback
+        print(traceback.format_exc())
+        return {"Error": str(e)}
+
+# Optional: allow input via URL
+def classify_from_url(url):
+    try:
+        response = requests.get(url)
+        image = Image.open(BytesIO(response.content))
+        return predict(image)
+    except Exception as e:
+        return {"Error": f"Could not load image from URL. {e}"}
+
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# PhytoViT - IFCB Phytoplankton Classifier")
+    gr.Markdown("Upload an image or paste a URL. Model: `phytoViT_508k_20250611`")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Image")
+            url_input = gr.Textbox(label="...or paste image URL")
+            predict_btn = gr.Button("Classify")
+
+        with gr.Column():
+            image_output = gr.Image(label="Input Image")
+            label_output = gr.Label(label="Top 2 Predictions")
+
+    predict_btn.click(fn=predict, inputs=image_input, outputs=label_output)
+    url_input.change(fn=classify_from_url, inputs=url_input, outputs=label_output)
 
 demo.launch()