Update app.py

app.py

@@ -1,15 +1,21 @@
+import os
 import torch
 import torch.nn as nn
 import torchvision.models as models
 import torchvision.transforms as transforms
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image
+import matplotlib.pyplot as plt
+import pandas as pd
 import gradio as gr
-import csv
-import os
 
-# -------------------------------
-# Load Breeds
-# -------------------------------
+# ----------------------------
+# Device
+# ----------------------------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ----------------------------
+# Labels
+# ----------------------------
 breeds = [
     "Alambadi", "Amritmahal", "Ayrshire", "Banni", "Bargur", "Bhadawari", "Brown_Swiss",
     "Dangi", "Deoni", "Gir", "Guernsey", "Hallikar", "Hariana", "Holstein_Friesian",
@@ -18,101 +24,96 @@ breeds = [
     "Nili_Ravi", "Nimari", "Ongole", "Pulikulam", "Rathi", "Red_Dane", "Red_Sindhi",
     "Sahiwal", "Surti", "Tharparkar", "Toda", "Umblachery", "Vechur"
 ]
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# -------------------------------
-# Load Model
-# -------------------------------
 num_classes = len(breeds)
+
+# ----------------------------
+# Model
+# ----------------------------
 model = models.efficientnet_b0(weights=None)
 model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)
-model.load_state_dict(torch.load("bovine_model.pth", map_location=device))
-model.to(device)
-model.eval()
-
-# -------------------------------
-# Image Transform
-# -------------------------------
-transform = transforms.Compose([
+state = torch.load("bovine_model.pth", map_location=device)
+model.load_state_dict(state)
+model.to(device).eval()
+
+# ----------------------------
+# Preprocessing
+# ----------------------------
+val_transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
-    transforms.Normalize([0.485, 0.456, 0.406], 
+    transforms.Normalize([0.485, 0.456, 0.406],
                          [0.229, 0.224, 0.225])
 ])
 
-# -------------------------------
-# Prediction Function
-# -------------------------------
-def predict(img):
-    img = img.convert("RGB")
-    input_tensor = transform(img).unsqueeze(0).to(device)
+# ----------------------------
+# Predict
+# ----------------------------
+def predict_image(img_path: str):
+    base_name = os.path.basename(img_path)
+    stem, _ = os.path.splitext(base_name)
+
+    img = Image.open(img_path).convert("RGB")
+    input_tensor = val_transform(img).unsqueeze(0).to(device)
 
     with torch.no_grad():
         outputs = model(input_tensor)
         probs = torch.nn.functional.softmax(outputs, dim=1)[0]
 
     top_prob, top_idx = torch.max(probs, dim=0)
-    predicted_breed = breeds[top_idx]
-    confidence = top_prob.item() * 100
-
-    # -------------------------------
-    # Annotate image with prediction
-    # -------------------------------
-    annotated_img = img.copy()
-    draw = ImageDraw.Draw(annotated_img)
-    font = ImageFont.load_default()
-
-    text = f"{predicted_breed} ({confidence:.2f}%)"
-    text_w, text_h = draw.textsize(text, font=font)
-
-    # Draw rectangle background
-    draw.rectangle([(0, 0), (text_w + 10, text_h + 10)], fill="black")
-    draw.text((5, 5), text, fill="white", font=font)
-
-    # Save annotated image
-    img_filename = f"{predicted_breed}_{int(confidence)}.png"
-    annotated_img.save(img_filename)
-
-    # Save CSV file
-    csv_filename = "prediction.csv"
-    with open(csv_filename, "w", newline="") as f:
-        writer = csv.writer(f)
-        writer.writerow(["Breed", "Confidence"])
-        writer.writerow([predicted_breed, f"{confidence:.2f}%"])
-
-    return (
-        f"{predicted_breed} ({confidence:.2f}%)",
-        annotated_img,
-        img_filename,
-        csv_filename
-    )
-
-# -------------------------------
-# Gradio Interface
-# -------------------------------
-with gr.Blocks(theme="default") as demo:
-    gr.Markdown("<h1 style='text-align:center;'>GoVed AI 🐄</h1>")
+    conf = float(top_prob.item()) * 100.0
+    predicted_breed = breeds[int(top_idx.item())]
+
+    # Annotate image with title
+    fig, ax = plt.subplots()
+    ax.imshow(img)
+    ax.set_title(f"{predicted_breed} ({conf:.2f}%)", fontsize=14, fontweight="bold")
+    ax.axis("off")
+    annotated_name = f"{predicted_breed}_{conf:.2f}pct_{stem}.png"
+    plt.savefig(annotated_name, format="png", bbox_inches="tight", pad_inches=0.1, dpi=150)
+    plt.close(fig)
+
+    # CSV output
+    df = pd.DataFrame([{
+        "breed": predicted_breed,
+        "confidence_percent": f"{conf:.2f}%",
+        "filename": base_name
+    }])
+    csv_name = f"{stem}_prediction.csv"
+    df.to_csv(csv_name, index=False)
+
+    return predicted_breed, f"{conf:.2f}", annotated_name, csv_name
+
+# ----------------------------
+# UI (Modern Layout)
+# ----------------------------
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("<h1 style='text-align:center;'>🐄 GoVed AI - Indian Cattle/Buffalo Breed Detection</h1>")
     gr.Markdown(
-        "Upload a cattle image to detect the breed, view prediction confidence, "
-        "and download results as **annotated image + CSV**."
+        "<p style='text-align:center;'>Upload an image → Detect the breed → View prediction confidence → Download results as CSV or Annotated Image.</p>"
     )
 
     with gr.Row():
-        with gr.Column():
-            img_input = gr.Image(type="pil", label="Upload Cattle Image")
-            submit_btn = gr.Button("🔍 Detect Breed")
-        with gr.Column():
-            breed_output = gr.Textbox(label="Predicted Breed")
-            img_output = gr.Image(type="pil", label="Image with Breed Info")
-            img_download = gr.File(label="Download Annotated Image")
-            csv_download = gr.File(label="Download CSV File")
+        with gr.Column(scale=1):
+            img_input = gr.Image(type="filepath", label="Upload Cattle/Buffalo Image", height=300)
+            submit_btn = gr.Button("🔍 Detect Breed", elem_id="detect-btn")
+
+        with gr.Column(scale=2):
+            breed_output = gr.Textbox(label="Predicted Breed", interactive=False)
+            confidence_output = gr.Label(label="Prediction Confidence (%)")
+            img_output = gr.Image(type="filepath", label="Annotated Image Preview", height=300)
+            
+            with gr.Row():
+                img_download = gr.File(label="⬇️ Download Annotated Image")
+                csv_download = gr.File(label="⬇️ Download CSV")
 
     submit_btn.click(
-        fn=predict,
+        fn=predict_image,
         inputs=img_input,
-        outputs=[breed_output, img_output, img_download, csv_download]
+        outputs=[breed_output, confidence_output, img_output, csv_download],
     )
 
+# ----------------------------
+# Run
+# ----------------------------
 if __name__ == "__main__":
     demo.launch()