Update app.py

app.py

@@ -7,6 +7,8 @@ from transformers import (
 from peft import PeftModel
 from PIL import Image
 import torch
+import json
+from huggingface_hub import hf_hub_download
 
 # Global model (all species)
 GLOBAL_MODEL_ID = "juppy44/plant-identification-2m-vit-b"
@@ -56,6 +58,7 @@ def get_wa_model():
     - Same base checkpoint as the global model
     - Classifier resized to WA_NUM_LABELS
     - WA LoRA applied on top
+    - id2label loaded primarily from labels.json in the adapter repo
     """
     global wa_model, wa_id2label
 
@@ -73,18 +76,47 @@ def get_wa_model():
     ).to(device)
     wa_base.eval()
 
-    # Try to get proper labels from the adapter config first
+    # --- PRIMARY: load labels from labels.json in the LoRA repo root ---
     local_id2label = {}
     try:
-        wa_cfg = AutoConfig.from_pretrained(WA_LORA_ID, subfolder=WA_LORA_SUBFOLDER)
-        if getattr(wa_cfg, "id2label", None):
-            local_id2label = normalize_id2label(wa_cfg.id2label)
-    except Exception:
-        pass
-
-    # Fallback to base config if adapter config doesn't have them
+        labels_path = hf_hub_download(
+            repo_id=WA_LORA_ID,
+            filename="labels.json",
+        )
+        with open(labels_path, "r") as f:
+            labels_data = json.load(f)
+
+        label2id = labels_data.get("label2id", {})
+        # invert label2id: {species_name: idx} -> {idx: species_name}
+        for label, idx in label2id.items():
+            try:
+                idx_int = int(idx)
+                local_id2label[idx_int] = label
+            except (TypeError, ValueError):
+                continue
+
+        if local_id2label:
+            print("Loaded id2label from labels.json in WA adapter repo.")
+    except Exception as e:
+        print(f"Could not load labels.json from WA adapter repo: {e}")
+
+    # --- SECONDARY: adapter config id2label (if labels.json missing/empty) ---
+    if not local_id2label:
+        try:
+            wa_cfg = AutoConfig.from_pretrained(
+                WA_LORA_ID,
+                subfolder=WA_LORA_SUBFOLDER,
+            )
+            if getattr(wa_cfg, "id2label", None):
+                local_id2label = normalize_id2label(wa_cfg.id2label)
+                print("Loaded id2label from WA adapter config.")
+        except Exception as e:
+            print(f"Could not load id2label from adapter config: {e}")
+
+    # --- TERTIARY: WA base config id2label ---
     if not local_id2label and getattr(wa_base.config, "id2label", None):
         local_id2label = normalize_id2label(wa_base.config.id2label)
+        print("Fallback: using id2label from WA base config.")
 
     print(
         f"Loading WA adapter {WA_LORA_ID} (subfolder='{WA_LORA_SUBFOLDER}')..."