Update files

app.py

@@ -1,3 +1,5 @@
+import json
+
 import birder
 import numpy as np
 from birder.inference.classification import infer_image
@@ -12,6 +14,19 @@ def get_birder_classification_models():
     return [model.modelId.split("/")[-1] for model in models]
 
 
+def get_selected_models():
+    return [
+        "convnext_v2_tiny_intermediate-il-common",
+        "mvit_v2_t_il-all",
+        "regnet_y_8g_intermediate-eu-common",
+        "hiera_abswin_base_mim-intermediate-eu-common",
+        "focalnet_b_lrf_intermediate-arabian-peninsula",
+        "swin_transformer_v2_s_intermediate-arabian-peninsula",
+        "rope_vit_reg4_b14_capi-inat21",
+        "rope_vit_reg4_b14_capi-places365",
+    ]
+
+
 def load_model_and_predict(image, model_name):
     try:
         if len(birder.list_pretrained_models(model_name)) == 0:
@@ -22,11 +37,21 @@ def load_model_and_predict(image, model_name):
         transform = birder.classification_transform(size, rgb_stats)
         (out, _) = infer_image(net, image, transform)
 
+        if model_name.endswith("-inat21") is True:
+            with open("inat21-mapping.json", "r", encoding="utf-8") as handle:
+                class_dict = json.load(handle)
+
+            class_mapping = {k: int(v) for v, k in class_dict.items()}
+            idx_to_class = dict(zip(class_to_idx.values(), class_to_idx.keys()))
+            idx_to_class.update(dict(zip(class_mapping.values(), class_mapping.keys())))
+            class_to_idx = dict(zip(idx_to_class.values(), idx_to_class.keys()))
+
         idx_to_class = {v: k for k, v in class_to_idx.items()}
         topk_idx = np.argsort(out[0])[-3:][::-1]
         predictions = [(idx_to_class[idx], float(out[0][idx])) for idx in topk_idx]
 
         return predictions
+
     except Exception as e:
         return [(f"Error: {str(e)}", 0.0)]
 
@@ -37,18 +62,18 @@ def predict(image, model_name):
 
 
 def create_interface():
-    models = get_birder_classification_models()
+    models = get_selected_models()
 
     examples = [
         ["Common myna.jpeg", "mvit_v2_t_il-all"],
-        ["Eurasian hoopoe.jpeg", "convnext_v2_tiny_intermediate-eu-common"],
-        # ["Eurasian teal.jpeg", "iformer_s_arabian-peninsula"],
-        ["Grey heron.jpeg", "iformer_s_arabian-peninsula"],
+        ["Eurasian hoopoe.jpeg", "hiera_abswin_base_mim-intermediate-eu-common"],
+        ["Grey heron.jpeg", "swin_transformer_v2_s_intermediate-arabian-peninsula"],
     ]
 
     # Create interface
     iface = gr.Interface(
         analytics_enabled=False,
+        deep_link=False,
         fn=predict,
         inputs=[
             gr.Image(type="pil", label="Input Image"),