Update app.py

app.py

@@ -1,26 +1,31 @@
 import torch
 from torchvision import models, transforms
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 import gradio as gr
 
 # =======================
-# Configuration
+# CONFIGURATION
 # =======================
 device = "cuda" if torch.cuda.is_available() else "cpu"
-MODEL_PATH = "cattle_breed_efficientnetb3_pytorch.pth"  # Upload this to the Space
+MODEL_PATH = "cattle_breed_efficientnetb3_pytorch.pth"  # Upload your .pth model here
 CLASS_NAMES = ["Gir", "Deoni", "Murrah"]
 
 # =======================
-# Load Model
+# MODEL: EfficientNetB3
 # =======================
-model = models.efficientnet_b3(pretrained=False)
+model = models.efficientnet_b3(weights=None)  # Do NOT load pretrained weights here
+# Update classifier for 3 classes
 model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(CLASS_NAMES))
-model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
+
+# Load checkpoint safely (ignores classifier mismatch)
+checkpoint = torch.load(MODEL_PATH, map_location=device)
+model.load_state_dict(checkpoint, strict=False)
+
 model.to(device)
 model.eval()
 
 # =======================
-# Image Preprocessing
+# IMAGE TRANSFORMS
 # =======================
 transform = transforms.Compose([
     transforms.Resize((300, 300)),
@@ -30,25 +35,38 @@ transform = transforms.Compose([
 ])
 
 # =======================
-# Prediction Function
+# PREDICTION FUNCTION
 # =======================
 def predict(image):
     image = image.convert("RGB")
     img_tensor = transform(image).unsqueeze(0).to(device)
+    
     with torch.no_grad():
         output = model(img_tensor)
-        pred_idx = torch.argmax(output, dim=1).item()
-    return CLASS_NAMES[pred_idx]
+        probs = torch.nn.functional.softmax(output, dim=1)
+        conf, pred_idx = torch.max(probs, dim=1)
+    
+    pred_label = CLASS_NAMES[pred_idx.item()]
+    confidence = conf.item() * 100
+
+    # Draw label on image
+    draw = ImageDraw.Draw(image)
+    font = ImageFont.load_default()
+    text = f"{pred_label} ({confidence:.2f}%)"
+    draw.text((10, 10), text, fill="red", font=font)
+
+    return image, text
 
 # =======================
-# Gradio Interface
+# GRADIO INTERFACE
 # =======================
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs="text",
+    outputs=[gr.Image(type="pil"), "text"],
     title="Indian Bovine Breed Classifier",
-    description="Upload an image of a cow and the model will predict its breed."
+    description="Upload an image of a cow and get the breed prediction with confidence."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()