Update app.py

app.py

@@ -14,82 +14,132 @@ import os
 # Import our model architecture
 from models import create_model
 
+# Optional: Hugging Face imports (used only when evaluating HF-format checkpoints)
+try:
+    from transformers import AutoConfig, AutoModelForImageClassification
+    HF_AVAILABLE = True
+except Exception:
+    HF_AVAILABLE = False
+
 # Configuration
-MODEL_PATH = "best_model.pth"
-CLASS_NAMES_PATH = "class_names.json"
+# Default to the moved fine-tuned checkpoint if present
+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')
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load class names
-with open(CLASS_NAMES_PATH, 'r') as f:
-    class_names = json.load(f)
+if os.path.exists(CLASS_NAMES_PATH):
+    try:
+        with open(CLASS_NAMES_PATH, 'r') as f:
+            class_names = json.load(f)
+    except Exception:
+        class_names = []
+else:
+    class_names = []
 
 NUM_CLASSES = len(class_names)
 
-# Load model - detect architecture from checkpoint
-print("Loading model...")
 
-# First, try to detect the correct architecture from the model file
-if os.path.exists(MODEL_PATH):
-    checkpoint = torch.load(MODEL_PATH, map_location='cpu')
-    
-    # Detect EfficientNet variant based on feature dimensions
-    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-        state_dict = checkpoint['model_state_dict']
-    else:
-        state_dict = checkpoint
-    
-    # Check backbone head feature size to determine EfficientNet variant
-    if 'backbone._conv_head.weight' in state_dict:
-        conv_head_shape = state_dict['backbone._conv_head.weight'].shape
-        if conv_head_shape[0] == 1536:  # EfficientNet-B3
-            model_type = 'efficientnet_b3'
-        elif conv_head_shape[0] == 1408:  # EfficientNet-B2
-            model_type = 'efficientnet_b2'
-        elif conv_head_shape[0] == 1280:  # EfficientNet-B0/B1
-            model_type = 'efficientnet_b1'
-        else:
-            model_type = 'efficientnet_b2'  # Default fallback
-    else:
-        model_type = 'efficientnet_b2'  # Default fallback
-    
-    # Check actual number of classes from classifier
-    if 'classifier.9.weight' in state_dict:
-        actual_classes = state_dict['classifier.9.weight'].shape[0]
+def load_checkpoint_model(model_path, device):
+    """Attempt to load a checkpoint. Supports local create_model-based checkpoints and
+    heuristic handling for Hugging Face (Swin) checkpoints when HF_MODEL_ID is set.
+    Returns (model, actual_num_classes) or (None, None) on failure.
+    """
+    if not os.path.exists(model_path):
+        print(f"Model file not found at {model_path}")
+        # If HF_MODEL_ID is set and transformers are available, try to load from hub
+        if HF_MODEL_ID and HF_AVAILABLE:
+            try:
+                print(f"Attempting to load model from Hugging Face 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")
+                return hf_model, num_labels
+            except Exception as e:
+                print("Failed to load HF model from hub:", e)
+        return None, None
+
+    ckpt = torch.load(model_path, map_location='cpu')
+    # unwrap common dict wrapper
+    if isinstance(ckpt, dict) and 'model_state_dict' in ckpt:
+        state_dict = ckpt['model_state_dict']
     else:
-        actual_classes = NUM_CLASSES
-    
-    print("Detected model: {} with {} classes".format(model_type, actual_classes))
-    
-else:
-    model_type = 'efficientnet_b2'
+        # if checkpoint is a state dict directly
+        state_dict = ckpt if isinstance(ckpt, dict) else {}
+
+    # 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
+
+    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)
+            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)")
+            num_labels = getattr(hf_model.config, 'num_labels', NUM_CLASSES)
+            return hf_model, num_labels
+        except Exception as e:
+            print("HF load failed:", e)
+            print("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
     actual_classes = NUM_CLASSES
-    print("Model file not found, using default: {}".format(model_type))
+    for k, v in state_dict.items():
+        if k.endswith('classifier.9.weight') or k.endswith('classifier.weight'):
+            try:
+                actual_classes = v.shape[0]
+                break
+            except Exception:
+                pass
 
-model = create_model(
-    num_classes=actual_classes,
-    model_type=model_type,
-    pretrained=False,  # We're loading trained weights
-    dropout_rate=0.3
-)
+    # Heuristic to choose an EfficientNet variant based on conv head size
+    model_type = 'efficientnet_b2'
+    if state_dict:
+        if 'backbone._conv_head.weight' in state_dict:
+            try:
+                conv_head_shape = state_dict['backbone._conv_head.weight'].shape
+                if conv_head_shape[0] == 1536:
+                    model_type = 'efficientnet_b3'
+                elif conv_head_shape[0] == 1408:
+                    model_type = 'efficientnet_b2'
+                elif conv_head_shape[0] == 1280:
+                    model_type = 'efficientnet_b1'
+            except Exception:
+                pass
 
-# Load trained weights
-if os.path.exists(MODEL_PATH):
+    print(f"Creating local model {model_type} with {actual_classes} classes (fallback)")
+    model = create_model(num_classes=actual_classes, model_type=model_type, pretrained=False, dropout_rate=0.3)
+    # Try to load state dict
     try:
-        checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
-        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
-            model.load_state_dict(checkpoint['model_state_dict'])
-            print("✅ Model loaded successfully! ({}, {} classes)".format(model_type, actual_classes))
-        else:
-            model.load_state_dict(checkpoint)
-            print("✅ Model loaded successfully! ({}, {} classes)".format(model_type, actual_classes))
+        # if ckpt was a dict without model_state_dict, attempt to load directly
+        to_load = state_dict if state_dict else ckpt
+        model.load_state_dict(to_load, strict=False)
+        model.to(device)
+        model.eval()
+        print("✅ Local model loaded (non-strict).")
+        return model, actual_classes
     except Exception as e:
-        print("❌ Error loading model: {}".format(str(e)))
-        print("Please ensure the model architecture matches the saved weights.")
-else:
-    print("⚠️ Model file not found. Please ensure best_model.pth is in the repository.")
+        print("Failed to load local model:", e)
+        return None, None
 
-model.to(DEVICE)
-model.eval()
+
+# Load model
+print("Loading model...", MODEL_PATH)
+model, actual_classes = load_checkpoint_model(MODEL_PATH, DEVICE)
+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}")
 
 def predict_bird(image):
     """
@@ -118,24 +168,39 @@ def predict_bird(image):
         # Prediction
         with torch.no_grad():
             outputs = model(input_tensor)
-            probabilities = F.softmax(outputs, dim=1)
-            confidence, predicted = torch.max(probabilities, 1)
-            
+
+            # Handle Hugging Face ModelOutput objects
+            try:
+                # HF ModelOutput may be dict-like with a 'logits' attribute
+                if hasattr(outputs, 'logits'):
+                    logits = outputs.logits
+                elif isinstance(outputs, (tuple, list)):
+                    logits = outputs[0]
+                else:
+                    logits = outputs
+            except Exception:
+                logits = outputs
+
+            # Ensure logits is a tensor
+            if not isinstance(logits, torch.Tensor):
+                logits = torch.tensor(np.asarray(logits)).to(DEVICE)
+
+            probabilities = F.softmax(logits, dim=1)
             # Get top 5 predictions
-            top5_prob, top5_indices = torch.topk(probabilities, 5)
-            
+            top5_prob, top5_indices = torch.topk(probabilities, min(5, probabilities.shape[1]), dim=1)
+
             # Format results
             results = {}
-            for i in range(5):
-                class_idx = top5_indices[0][i].item()
-                prob = top5_prob[0][i].item()
+            for i in range(top5_indices.shape[1]):
+                class_idx = int(top5_indices[0][i].item())
+                prob = float(top5_prob[0][i].item())
                 # Handle potential class index mismatch
                 if class_idx < len(class_names):
                     class_name = class_names[class_idx].replace('_', ' ')
                 else:
                     class_name = "Class_" + str(class_idx)
-                results[class_name] = float(prob)
-        
+                results[class_name] = prob
+
         return results
         
     except Exception as e: