Update app.py

app.py

@@ -17,11 +17,16 @@ from models import create_model
 
 # Optional: Hugging Face imports (used only when evaluating HF-format checkpoints)
 try:
+    import transformers
     from transformers import AutoConfig, AutoModelForImageClassification
     HF_AVAILABLE = True
 except Exception:
+    transformers = None
     HF_AVAILABLE = False
 
+# HF image processor (AutoImageProcessor or AutoFeatureExtractor) will be stored here when available
+hf_processor = None
+
 # Configuration
 # Default to the moved fine-tuned checkpoint if present
 MODEL_PATH = os.environ.get('MODEL_PATH', os.path.join('best_model_finetuned.pth'))
@@ -58,6 +63,9 @@ 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.
     """
+    # Allow writing to module-level hf_processor
+    global hf_processor
+
     # 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:
@@ -65,6 +73,13 @@ def load_checkpoint_model(model_path, device):
             hf_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_ID)
             hf_model.to(device)
             hf_model.eval()
+            # Try to load a matching image processor from the hub for preprocessing
+            try:
+                if transformers is not None:
+                    hf_processor = transformers.AutoImageProcessor.from_pretrained(HF_MODEL_ID)
+                    print("Loaded HF image processor from hub (force load)")
+            except Exception:
+                print("Warning: failed to load HF image processor for forced hub model")
             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
@@ -82,6 +97,13 @@ def load_checkpoint_model(model_path, device):
                 hf_model = AutoModelForImageClassification.from_pretrained(HF_MODEL_ID)
                 hf_model.to(device)
                 hf_model.eval()
+                # Try to load image processor for preprocessing
+                try:
+                    if transformers is not None:
+                        hf_processor = transformers.AutoImageProcessor.from_pretrained(HF_MODEL_ID)
+                        print("Loaded HF image processor from hub")
+                except Exception:
+                    print("Warning: failed to load HF image processor from hub")
                 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
@@ -159,6 +181,13 @@ def load_checkpoint_model(model_path, device):
             missing, unexpected = hf_model.load_state_dict(state_dict, strict=False)
             hf_model.to(device)
             hf_model.eval()
+            # Try to fetch image processor for this hf id so we can preprocess in predict
+            try:
+                if transformers is not None:
+                    hf_processor = transformers.AutoImageProcessor.from_pretrained(hf_id_to_use)
+                    print(f"Loaded HF image processor for {hf_id_to_use}")
+            except Exception:
+                print(f"Warning: failed to load HF image processor for {hf_id_to_use}")
             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)
@@ -260,16 +289,37 @@ def predict_bird(image):
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
-        # Define preprocessing step by step to avoid namespace issues
-        resize = transforms.Resize((320, 320))
-        to_tensor = transforms.ToTensor()
-        normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        
-        # Apply transformations step by step
-        resized_image = resize(image)
-        tensor_image = to_tensor(resized_image)
-        normalized_tensor = normalize(tensor_image)
-        input_tensor = normalized_tensor.unsqueeze(0).to(DEVICE)
+        # If an HF image processor was loaded alongside a HF model, prefer it for preprocessing
+        if hf_processor is not None:
+            try:
+                # AutoImageProcessor expects PIL images or numpy arrays; return_tensors='pt' gives PyTorch tensors
+                proc = hf_processor(images=image, return_tensors='pt')
+                # Some processors return 'pixel_values', others return 'pixel_values' key
+                if 'pixel_values' in proc:
+                    input_tensor = proc['pixel_values'].to(DEVICE)
+                else:
+                    # Fall back to first tensor-like value
+                    val = next(iter(proc.values()))
+                    input_tensor = val.to(DEVICE)
+            except Exception:
+                logger.exception('HF processor failed; falling back to torchvision preprocessing')
+                hf_local_fallback = True
+            else:
+                hf_local_fallback = False
+        else:
+            hf_local_fallback = True
+
+        if hf_local_fallback:
+            # Define preprocessing step by step to avoid namespace issues
+            resize = transforms.Resize((320, 320))
+            to_tensor = transforms.ToTensor()
+            normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+            
+            # Apply transformations step by step
+            resized_image = resize(image)
+            tensor_image = to_tensor(resized_image)
+            normalized_tensor = normalize(tensor_image)
+            input_tensor = normalized_tensor.unsqueeze(0).to(DEVICE)
         
         # Prediction
         with torch.no_grad():
@@ -287,10 +337,34 @@ def predict_bird(image):
             except Exception:
                 logits = outputs
 
+
             # Ensure logits is a tensor
             if not isinstance(logits, torch.Tensor):
                 logits = torch.tensor(np.asarray(logits)).to(DEVICE)
 
+            # Handle unexpected tensor shapes:
+            # - if logits has spatial dims (e.g., 4D), average them
+            # - if logits is 1D, unsqueeze batch dim
+            try:
+                if logits.dim() > 2:
+                    # average over all dims after channel dim
+                    reduce_dims = tuple(range(2, logits.dim()))
+                    logits = logits.mean(dim=reduce_dims)
+                if logits.dim() == 1:
+                    logits = logits.unsqueeze(0)
+            except Exception:
+                # if shape ops fail, log and return safe error
+                logger.exception('Failed to normalize logits shape')
+                return {"Error": 0.0}
+
+            # If single-logit output, treat as sigmoid probability
+            if logits.size(1) == 1:
+                probs = torch.sigmoid(logits)
+                # return single-label prob mapped to first class or generic
+                prob = float(probs[0, 0].item())
+                label = class_names[0].replace('_', ' ') if class_names else 'Class_0'
+                return {label: prob}
+
             probabilities = F.softmax(logits, dim=1)
             # Get top 5 predictions
             top5_prob, top5_indices = torch.topk(probabilities, min(5, probabilities.shape[1]), dim=1)
@@ -310,7 +384,10 @@ def predict_bird(image):
         return results
         
     except Exception as e:
-        return {"Error": "Prediction failed: " + str(e)}
+        # Log exception and return a numeric-friendly error response for Gradio
+        logger.exception('Prediction failed')
+        print('Prediction failed:', e)
+        return {"Error": 0.0}
 
 # Create Gradio interface
 title = "🐦 Bird Species Classifier"