Updated deployment

README.md

@@ -1,8 +1,8 @@
 ---
 title: Birds classifier
 emoji: 🐦
-colorFrom: red
-colorTo: pink
+colorFrom: white
+colorTo: blue
 sdk: gradio
 sdk_version: 3.20.1
 app_file: app.py

app.py

@@ -14,6 +14,15 @@ to_tensor = ToTensor()
 norm = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 
 
+def decode_pred(pred: torch.Tensor) -> str:
+    indices = pred > 0.95
+    if indices.any():
+        # return first match
+        return vocab[indices.nonzero()[0]]
+    else:
+        return "I don't know what this is, ¡páharo!"
+
+
 def classify_image(inp):
     inp = Image.fromarray(inp)
     transformed_input = resized_crop_pad(inp, (460, 460))
@@ -23,8 +32,7 @@ def classify_image(inp):
     model.eval()
     with torch.no_grad():
         pred = model(transformed_input)
-    pred = torch.argmax(pred, dim=1)
-    return vocab[pred]
+    return decode_pred(torch.sigmoid(pred).squeeze(dim=0))
 
 
 iface = gr.Interface(

vit_saved.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:09e9c264c4b31db320125563076207358c54b1d460470f38f2ee2c6c196eafe1
+oid sha256:e8169a84843ea366cae9295d3a6b7870884fdca2169fe7b419d9552990a26942
 size 22974149