Update app.py

app.py

@@ -7,18 +7,21 @@ import gradio as gr
 # CONFIGURATION
 # =======================
 device = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_PATH = "cattle_breed_efficientnetb3_pytorch.pth"  # Upload your .pth model here
+MODEL_PATH = "best_model.pth"  # Upload your .pth model here
 CLASS_NAMES = ["Gir", "Deoni", "Murrah"]
 
 # =======================
 # MODEL: EfficientNetB3
 # =======================
-model = models.efficientnet_b3(weights=None)  # Do NOT load pretrained weights here
-# Update classifier for 3 classes
+model = models.efficientnet_b3(weights=None)
 model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
 
-# Load checkpoint safely (ignores classifier mismatch)
+# =======================
+# LOAD CHECKPOINT (Feature Extractor Only)
+# =======================
 checkpoint = torch.load(MODEL_PATH, map_location=device)
+# Remove classifier weights from checkpoint
+checkpoint = {k: v for k, v in checkpoint.items() if "classifier" not in k}
 model.load_state_dict(checkpoint, strict=False)
 
 model.to(device)