Fixed inference

app.py

@@ -5,15 +5,39 @@ from PIL import Image
 import requests
 from io import BytesIO
 import os
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
+from torchvision import transforms
+import json
 
-# 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
-model = ViTForImageClassification.from_pretrained("patcdaniel/phytoViT_508k_20250611")
-processor = ViTImageProcessor.from_pretrained("patcdaniel/phytoViT_508k_20250611")
-model.eval()
+# Download the file from your model repo (replace with your actual token if private)
+model_path = hf_hub_download(
+    repo_id="patcdaniel/phytoViT_508k_20250611",
+    filename="model.safetensors",
+    token=os.environ.get("HF_TOKEN")  # omit this line if public
+)
+state_dict = load_file(model_path)
+
+model = ViTForImageClassification.from_pretrained(
+    "google/vit-base-patch16-224-in21k",
+    num_labels=95  # this must match your training
+)
+model.load_state_dict(state_dict)
+
+model_path = hf_hub_download(
+    repo_id="patcdaniel/phytoViT_508k_20250611",
+    filename="label_names.json",
+    token=os.environ.get("HF_TOKEN"),
+    local_dir="."
+)
+
+# Load class label dictionary (label -> index)
+with open(model_path, "r") as f:
+    id2label = {int(k): v for k, v in json.load(f).items()}
+
+# Convert to id -> label
+# id2label = {v: k for k, v in label2id.items()}
 
 # Use GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -22,19 +46,27 @@ model.to(device)
 # Inference function
 def predict(image):
     try:
-        image = image.convert("RGB")
-        inputs = processor(images=image, return_tensors="pt").to(device)
+        transform = transforms.Compose([
+            transforms.Resize(256),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),  # Converts to [0, 1] range
+            transforms.Normalize(
+                mean=[0.485, 0.456, 0.406],
+                std=[0.229, 0.224, 0.225]
+            )
+        ])
+
+        pixel_values = transform(image).unsqueeze(0).to(device)
 
         with torch.no_grad():
-            logits = model(**inputs).logits
+            logits = model(pixel_values).logits
             probs = torch.nn.functional.softmax(logits, dim=-1).squeeze()
 
-        topk = torch.topk(probs, k=2)
+        topk = torch.topk(probs, k=5)
         top_indices = topk.indices.tolist()
         top_scores = topk.values.tolist()
 
-        id2label = model.config.id2label
-        top_labels = [id2label[str(i)] for i in top_indices]
+        top_labels = [id2label[i] for i in top_indices]
 
         return {label: round(score, 4) for label, score in zip(top_labels, top_scores)}
 
@@ -62,10 +94,7 @@ with gr.Blocks() as demo:
             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")
+            label_output = gr.Label(label="Top 5 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)

label_names.json

@@ -0,0 +1,91 @@
+{
+  "0": "Akashiwo",
+  "1": "Alexandrium",
+  "2": "Amylax_Gonyaulax_Protoceratium",
+  "3": "Asterionellopsis",
+  "4": "Asterionellopsis_chain",
+  "5": "Asteromphalus",
+  "6": "Bad_Beads",
+  "7": "Bad_blurred",
+  "8": "Bad_mixed_phyto",
+  "9": "Bad_setae",
+  "10": "Centric",
+  "11": "Centric_fuzzy",
+  "12": "Ceratium_divaricatum",
+  "13": "Ceratium_furca",
+  "14": "Ceratium_lineatum",
+  "15": "Chaetoceros",
+  "16": "Ciliate_cutoff",
+  "17": "Ciliate_large",
+  "18": "Ciliate_large_2",
+  "19": "Ciliate_other_morpho_1",
+  "20": "Clusterflagellate_morpho_1",
+  "21": "Clusterflagellate_morpho_2",
+  "22": "Corethron",
+  "23": "Cryptophyte",
+  "24": "Cylindrotheca_Nitzschia",
+  "25": "Detonula_Cerataulina_Lauderia",
+  "26": "Detritus",
+  "27": "Detritus_infection",
+  "28": "Dictyocha",
+  "29": "Dinoflagellate_morpho_1",
+  "30": "Dinoflagellate_morpho_2",
+  "31": "Dinoflagellate_morpho_3",
+  "32": "Dinophysis",
+  "33": "Ditylum",
+  "34": "Entomoneis",
+  "35": "Eucampia",
+  "36": "Euglenoid",
+  "37": "Flagellate_morpho_1",
+  "38": "Flagellate_morpho_2",
+  "39": "Flagellate_morpho_3",
+  "40": "Flagellate_nano_1",
+  "41": "Flagellate_nano_2",
+  "42": "Fragilariopsis",
+  "43": "Guinardia_Dactyliosolen",
+  "44": "Gymnodinium",
+  "45": "Gyrodinium",
+  "46": "Gyrosigma",
+  "47": "Haptophyte_prymnesium",
+  "48": "Hemiaulus",
+  "49": "Hemiselmis",
+  "50": "Heterocapsa_morpho_1",
+  "51": "Heterocapsa_morpho_2",
+  "52": "Heterosigma_akashiwo",
+  "53": "Laboea",
+  "54": "Leptocylindrus",
+  "55": "Lingulodinium",
+  "56": "Margalefidinium",
+  "57": "Mesodinium",
+  "58": "Nano_cluster",
+  "59": "Nano_p_white",
+  "60": "Pennate_med",
+  "61": "Pennate_morpho_1",
+  "62": "Pennate_short",
+  "63": "Pennate_wide",
+  "64": "Peridinium",
+  "65": "Phaeocystis",
+  "66": "Pleurosigma",
+  "67": "Polykrikos",
+  "68": "Proboscia",
+  "69": "Prorocentrum_narrow",
+  "70": "Prorocentrum_wide",
+  "71": "Pseudo-nitzschia",
+  "72": "Pyramimonas",
+  "73": "Rhizosolenia",
+  "74": "Scrippsiella",
+  "75": "Skeleonema",
+  "76": "Skeletonema",
+  "77": "Spiky_pacman_circular",
+  "78": "Stombidinium_morpho_1",
+  "79": "Strombidium_morpho_2",
+  "80": "Thalassionema",
+  "81": "Thalassiosira",
+  "82": "Tiarina",
+  "83": "Tintinnid",
+  "84": "Tontonia",
+  "85": "Torodinium",
+  "86": "Tropidoneis",
+  "87": "Unknown_morpho_1",
+  "88": "Vicicitus"
+}