Update app.py

app.py

@@ -10,6 +10,7 @@ import json
 import numpy as np
 from torchvision import transforms
 import os
+import logging
 
 # Import our model architecture
 from models import create_model
@@ -23,12 +24,22 @@ except Exception:
 
 # Configuration
 # Default to the moved fine-tuned checkpoint if present
-MODEL_PATH = os.environ.get('MODEL_PATH', os.path.join('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')
+FORCE_HF_LOAD = os.environ.get('FORCE_HF_LOAD', '0').lower() in ('1', 'true', 'yes')
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+# Default HF model id to try when checkpoint looks HF-like and HF_MODEL_ID not set
+DEFAULT_HF_ID = 'Emiel/cub-200-bird-classifier-swin'
+
+# Setup file logger for traceability in Spaces
+LOG_FILE = os.environ.get('APP_LOG_PATH', 'app.log')
+logging.basicConfig(level=logging.INFO, filename=LOG_FILE, filemode='a',
+                    format='%(asctime)s %(levelname)s: %(message)s')
+logger = logging.getLogger(__name__)
+
 # Load class names
 if os.path.exists(CLASS_NAMES_PATH):
     try:
@@ -47,8 +58,23 @@ def load_checkpoint_model(model_path, device):
     heuristic handling for Hugging Face (Swin) checkpoints when HF_MODEL_ID is set.
     Returns (model, actual_num_classes) or (None, None) on failure.
     """
+    # If user wants to force HF loading from hub, try that first (useful in Spaces)
+    if FORCE_HF_LOAD and HF_MODEL_ID and HF_AVAILABLE:
+        try:
+            print(f"FORCE_HF_LOAD enabled: loading HF model from hub: {HF_MODEL_ID}")
+            hf_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_ID)
+            hf_model.to(device)
+            hf_model.eval()
+            num_labels = getattr(hf_model.config, 'num_labels', NUM_CLASSES)
+            print(f"Loaded HF model from hub with {num_labels} labels (force)")
+            return hf_model, num_labels
+        except Exception as e:
+            print("Forced HF hub load failed:", e)
+
     if not os.path.exists(model_path):
-        print(f"Model file not found at {model_path}")
+        msg = f"Model file not found at {model_path}"
+        print(msg)
+        logger.info(msg)
         # If HF_MODEL_ID is set and transformers are available, try to load from hub
         if HF_MODEL_ID and HF_AVAILABLE:
             try:
@@ -63,7 +89,14 @@ def load_checkpoint_model(model_path, device):
                 print("Failed to load HF model from hub:", e)
         return None, None
 
-    ckpt = torch.load(model_path, map_location='cpu')
+    print(f"Loading checkpoint from: {model_path}")
+    logger.info(f"Loading checkpoint from: {model_path}")
+    try:
+        ckpt = torch.load(model_path, map_location='cpu')
+    except Exception as e:
+        print("Failed to load checkpoint file:", e)
+        logger.exception("Failed to load checkpoint file:")
+        ckpt = {}
     # unwrap common dict wrapper
     if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
         state_dict = ckpt['model_state_dict']
@@ -71,25 +104,50 @@ def load_checkpoint_model(model_path, device):
         # if checkpoint is a state dict directly
         state_dict = ckpt if isinstance(ckpt, dict) else {}
 
+    # Diagnostic: print a few state_dict keys so we can tell checkpoint format
+    try:
+        sample_keys = list(state_dict.keys())[:8]
+        print("Checkpoint sample keys:", sample_keys)
+        logger.info(f"Checkpoint sample keys: {sample_keys}")
+    except Exception:
+        print("No state_dict keys to sample")
+        logger.info("No state_dict keys to sample")
+
     # Heuristic: detect HF-style Swin checkpoint by looking for keys that start with 'swin.'
     hf_like = any(k.startswith('swin.') or 'swin.embeddings' in k for k in state_dict.keys()) if state_dict else False
+    hf_msg = f"hf_like_checkpoint_detected={hf_like} HF_AVAILABLE={HF_AVAILABLE} HF_MODEL_ID={'set' if HF_MODEL_ID else 'not-set'}"
+    print(hf_msg)
+    logger.info(hf_msg)
+
+    if hf_like and HF_AVAILABLE:
+        # choose which HF id to use: env var or default
+        hf_id_to_use = HF_MODEL_ID or DEFAULT_HF_ID
+        if HF_MODEL_ID is None:
+            info_msg = f"HF_MODEL_ID not set; using DEFAULT_HF_ID='{DEFAULT_HF_ID}' to attempt hub load"
+            print(info_msg)
+            logger.info(info_msg)
 
-    if hf_like and HF_AVAILABLE and HF_MODEL_ID:
-        # Try to instantiate HF model from the hub config to match architecture
         try:
-            print(f"Attempting to load Hugging Face model '{HF_MODEL_ID}' and apply checkpoint weights...")
-            config = AutoConfig.from_pretrained(HF_MODEL_ID)
+            msg = f"Attempting to load Hugging Face model '{hf_id_to_use}' and apply checkpoint weights..."
+            print(msg)
+            logger.info(msg)
+            # prefer using the hub config to instantiate exact architecture
+            config = AutoConfig.from_pretrained(hf_id_to_use)
             hf_model = AutoModelForImageClassification.from_config(config)
             # load weights non-strictly: match shapes
             missing, unexpected = hf_model.load_state_dict(state_dict, strict=False)
             hf_model.to(device)
             hf_model.eval()
-            print(f"Loaded HF model with non-strict state_dict (missing {len(missing)} keys, unexpected {len(unexpected)} keys)")
+            ok_msg = f"Loaded HF model with non-strict state_dict (missing {len(missing)} keys, unexpected {len(unexpected)} keys)"
+            print(ok_msg)
+            logger.info(ok_msg)
             num_labels = getattr(hf_model.config, 'num_labels', NUM_CLASSES)
             return hf_model, num_labels
         except Exception as e:
             print("HF load failed:", e)
+            logger.exception("HF load failed")
             print("Falling back to local model loader...")
+            logger.info("Falling back to local model loader")
 
     # Fallback: try to detect EfficientNet-like shapes and create local model
     # Determine actual num classes by inspecting a likely classifier weight key
@@ -149,11 +207,11 @@ else:
                 # 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)
+                hf_class_names = [""] * (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]
+                # Filter out empty entries
+                hf_class_names = [c for c in hf_class_names if c]
                 if len(hf_class_names) > 0:
                     class_names = hf_class_names
                     NUM_CLASSES = len(class_names)