update

app.py

@@ -1,6 +1,7 @@
 """
 Simple Demo for Pest and Disease Classification
 For Hugging Face Space Deployment
+Supports both single model and ensemble prediction
 """
 
 import torch
@@ -8,6 +9,8 @@ from PIL import Image
 import json
 import gradio as gr
 from torchvision import transforms
+import numpy as np
+from pathlib import Path
 
 from model import create_model
 
@@ -62,8 +65,94 @@ class PestDiseasePredictor:
         return dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
 
 
+class EnsemblePredictor:
+    """Ensemble predictor using weighted soft voting"""
+
+    def __init__(self, checkpoint_paths, weights, label_mapping_path, backbone='efficientnet_b3', device='cuda'):
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+
+        # Normalize weights to sum to 1
+        weights = np.array(weights)
+        self.weights = weights / weights.sum()
+
+        # Load label mapping
+        with open(label_mapping_path, 'r', encoding='utf-8') as f:
+            mapping = json.load(f)
+            self.id_to_label = {int(k): v for k, v in mapping['id_to_label'].items()}
+            self.num_classes = mapping['num_classes']
+
+        # Load all models
+        self.models = []
+        print(f"\n{'='*80}")
+        print("Loading Ensemble Models")
+        print(f"{'='*80}")
+
+        for i, checkpoint_path in enumerate(checkpoint_paths):
+            print(f"\nModel {i+1}/{len(checkpoint_paths)}")
+            print(f"  Checkpoint: {checkpoint_path}")
+            print(f"  Weight: {self.weights[i]:.4f}")
+
+            # Create model
+            model = create_model(
+                num_classes=self.num_classes,
+                backbone=backbone,
+                pretrained=False
+            )
+
+            # Load checkpoint
+            if Path(checkpoint_path).exists():
+                checkpoint = torch.load(checkpoint_path, map_location=self.device)
+                model.load_state_dict(checkpoint['model_state_dict'])
+                model = model.to(self.device)
+                model.eval()
+                self.models.append(model)
+                print(f"  ✅ Loaded successfully")
+            else:
+                print(f"  ❌ Checkpoint not found: {checkpoint_path}")
+                raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+
+        print(f"\n{'='*80}")
+        print(f"✅ Ensemble loaded: {len(self.models)} models")
+        print(f"💻 Device: {self.device}")
+        print(f"📚 Classes: {self.num_classes}")
+        print(f"{'='*80}\n")
+
+        # Image transforms
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                               std=[0.229, 0.224, 0.225])
+        ])
+
+    def predict(self, image):
+        """Predict using weighted soft voting"""
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        img_tensor = self.transform(image).unsqueeze(0).to(self.device)
+
+        # Get predictions from all models
+        ensemble_probs = np.zeros(self.num_classes)
+
+        with torch.no_grad():
+            for model, weight in zip(self.models, self.weights):
+                outputs = model(img_tensor)
+                probabilities = torch.nn.functional.softmax(outputs, dim=1)
+                probs = probabilities[0].cpu().numpy()
+                ensemble_probs += weight * probs
+
+        # Create results dictionary
+        results = {}
+        for idx, prob in enumerate(ensemble_probs):
+            class_name = self.id_to_label[idx]
+            results[class_name] = float(prob)
+
+        return dict(sorted(results.items(), key=lambda x: x[1], reverse=True))
+
+
 # ========== For Hugging Face Space ==========
-checkpoint_path = "checkpoints/best_model_993.pth"
+checkpoint_path = "checkpoints/best_model_fold1.pth"
 label_mapping_path = "label_mapping.json"
 backbone = 'efficientnet_b3'
 device = "cuda"

checkpoints/best_model_fold2.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:caeb0e81ac149d12a107acfa87268d90ac5f7711f9a8cb531c50f3ea04748ae3
+size 138717293

checkpoints/best_model_fold3.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ac3889dde1936a9c9c3ea18b41802a352adf8849df1b5b568d694e8684c2586
+size 138717293

checkpoints/best_model_fold4.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4de4227b21d8470070432370a27bb99551813dd77d729b2a44f11c64ef209e43
+size 138717293

checkpoints/best_model_fold5.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf4e72ccb78e6d2849f67ab552abc28d7ff07fab87957cb34d290fe00b8831fc
+size 138717293