Update app.py

app.py

@@ -23,7 +23,7 @@ except Exception:
 
 # Configuration
 # Default to the moved fine-tuned checkpoint if present
-MODEL_PATH = os.environ.get('MODEL_PATH', 'best_model_finetuned.pth')
+MODEL_PATH = os.environ.get('MODEL_PATH', os.path.join('results', 'fine_tune', 'best_model_finetuned.pth'))
 # Optional: if your HF model id is known (e.g. Emiel/cub-200-bird-classifier-swin), set HF_MODEL_ID env var
 HF_MODEL_ID = os.environ.get('HF_MODEL_ID', None)
 CLASS_NAMES_PATH = os.environ.get('CLASS_NAMES_PATH', 'class_names.json')
@@ -140,6 +140,33 @@ if model is None:
     print("No model available. The app will still launch but predictions will fail.")
 else:
     print(f"Model ready. Classes={actual_classes}")
+    # If this is a Hugging Face model with id2label, prefer that mapping
+    try:
+        hf_config = getattr(model, 'config', None)
+        if hf_config is not None:
+            id2label = getattr(hf_config, 'id2label', None)
+            if id2label:
+                # id2label keys may be strings or ints
+                # Build ordered class_names list by index
+                max_idx = max(int(k) for k in id2label.keys())
+                hf_class_names = [None] * (max_idx + 1)
+                for k, v in id2label.items():
+                    hf_class_names[int(k)] = v.replace(' ', '_') if isinstance(v, str) else str(v)
+                # Filter out None at end if any
+                hf_class_names = [c for c in hf_class_names if c is not None]
+                if len(hf_class_names) > 0:
+                    class_names = hf_class_names
+                    NUM_CLASSES = len(class_names)
+                    print(f"Using Hugging Face id2label mapping with {NUM_CLASSES} classes")
+    except Exception as e:
+        print("Warning: failed to extract id2label from HF model config:", e)
+
+    # Warn if class_names.json doesn't match model classes
+    if class_names and actual_classes and len(class_names) != actual_classes:
+        print(f"Warning: class_names.json has {len(class_names)} entries but model expects {actual_classes} classes.")
+        # If HF labels exist and match expected size, prefer them
+        if len(class_names) < actual_classes:
+            print("Note: consider updating class_names.json to match the model's label order or set HF_MODEL_ID to use id2label mapping.")
 
 def predict_bird(image):
     """