Update app.py

app.py

@@ -1,9 +1,10 @@
-import os
+import os, io
 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
 
@@ -48,6 +49,7 @@ val_transform = transforms.Compose([
 # Predict
 # ----------------------------
 def predict_image(img_path: str):
+    # Keep original filename for outputs
     base_name = os.path.basename(img_path)
     stem, _ = os.path.splitext(base_name)
 
@@ -62,26 +64,14 @@ def predict_image(img_path: str):
     conf = float(top_prob.item()) * 100.0
     predicted_breed = breeds[int(top_idx.item())]
 
-    # Create a new image with extra space on top for text
-    extra_height = 60
-    new_img = Image.new("RGB", (img.width, img.height + extra_height), "white")
-    new_img.paste(img, (0, extra_height))
-
-    # Draw text at top-center
-    draw = ImageDraw.Draw(new_img)
-    try:
-        font = ImageFont.truetype("arial.ttf", 28)  # Bold font if available
-    except:
-        font = ImageFont.load_default()
-
-    text = f"{predicted_breed} ({conf:.2f}%)"
-    text_width, text_height = draw.textsize(text, font=font)
-    x = (new_img.width - text_width) // 2
-    y = (extra_height - text_height) // 2
-    draw.text((x, y), text, fill="black", font=font)
-
+    # Annotate image (title overlay)
+    fig, ax = plt.subplots()
+    ax.imshow(img)
+    ax.set_title(f"{predicted_breed} ({conf:.2f}%)")
+    ax.axis("off")
     annotated_name = f"{predicted_breed}_{conf:.2f}pct_{stem}.png"
-    new_img.save(annotated_name)
+    plt.savefig(annotated_name, format="png", bbox_inches="tight", pad_inches=0.1, dpi=150)
+    plt.close(fig)
 
     # CSV output
     df = pd.DataFrame([{
@@ -92,7 +82,12 @@ def predict_image(img_path: str):
     csv_name = f"{stem}_prediction.csv"
     df.to_csv(csv_name, index=False)
 
-    return predicted_breed, f"{conf:.2f}%", annotated_name, csv_name
+    # Return:
+    # 1) predicted breed (text)
+    # 2) confidence (%) (text)
+    # 3) file (CSV)
+    # 4) file (annotated image with breed+confidence in filename)
+    return predicted_breed, f"{conf:.2f}%", csv_name, annotated_name
 
 
 # ----------------------------
@@ -100,15 +95,15 @@ def predict_image(img_path: str):
 # ----------------------------
 demo = gr.Interface(
     fn=predict_image,
-    inputs=gr.Image(type="filepath", label="📤 Upload Cattle/Buffalo Image"),
+    inputs=gr.Image(type="filepath", label="Upload Cattle/Buffalo Image"),
     outputs=[
         gr.Textbox(label="Predicted Breed"),
-        gr.Textbox(label="Confidence (%)"),
-        gr.File(label="⬇️ Download Annotated Image", type="filepath"),
-        gr.File(label="⬇️ Download Prediction CSV", type="filepath"),
+        gr.Textbox(label="Prediction Confidence (%)"),
+        gr.File(label="Download CSV"),
+        gr.File(label="Download Annotated Image")
     ],
-    title="🐄 GoVed AI – Indian Cattle/Buffalo Breed Detection",
-    description="Upload an image → Get predicted breed, confidence, annotated image, and CSV download."
+    title="Indian Cattle/Buffalo Breed Detection",
+    description="Upload an image → get predicted breed, confidence score, CSV, and an annotated image file named with the predicted breed and confidence."
 )
 
 if __name__ == "__main__":