fixed unexpected errors

app.py

@@ -3,8 +3,7 @@ import torch
 import torch.nn as nn
 import torchvision.models as models
 import torchvision.transforms as transforms
-from PIL import Image
-import matplotlib.pyplot as plt
+from PIL import Image, ImageDraw, ImageFont
 import pandas as pd
 import gradio as gr
 
@@ -46,15 +45,28 @@ val_transform = transforms.Compose([
 ])
 
 # ----------------------------
-# Predict
+# Helper: safe font load
+# ----------------------------
+def _get_font(size):
+    try:
+        # common font on many systems; may exist in HF images
+        return ImageFont.truetype("DejaVuSans-Bold.ttf", size=size)
+    except Exception:
+        return ImageFont.load_default()
+
+# ----------------------------
+# Predict function (returns: breed, confidence_float, preview_path, download_image_path, csv_path)
 # ----------------------------
 def predict_image(img_path: str):
+    # base filename
     base_name = os.path.basename(img_path)
     stem, _ = os.path.splitext(base_name)
 
+    # open and preprocess image
     img = Image.open(img_path).convert("RGB")
     input_tensor = val_transform(img).unsqueeze(0).to(device)
 
+    # inference
     with torch.no_grad():
         outputs = model(input_tensor)
         probs = torch.nn.functional.softmax(outputs, dim=1)[0]
@@ -63,14 +75,41 @@ def predict_image(img_path: str):
     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)
+    # annotate image using PIL (avoid matplotlib dependency)
+    annotated = img.copy()
+    draw = ImageDraw.Draw(annotated)
+
+    # choose font size relative to image width
+    font_size = max(16, int(annotated.width * 0.04))
+    font = _get_font(font_size)
+
+    text = f"{predicted_breed} ({conf:.2f}%)"
+    text_w, text_h = draw.textsize(text, font=font)
+
+    # rectangle background for text (semi-opaque)
+    padding_x = 12
+    padding_y = 8
+    rect_w = text_w + padding_x * 2
+    rect_h = text_h + padding_y * 2
+
+    # place rectangle centered at top
+    rect_x0 = max(0, (annotated.width - rect_w) // 2)
+    rect_y0 = 0
+    rect_x1 = rect_x0 + rect_w
+    rect_y1 = rect_y0 + rect_h
+
+    # draw rectangle (black)
+    draw.rectangle([(rect_x0, rect_y0), (rect_x1, rect_y1)], fill=(0, 0, 0, 200))
+    # draw text (white) centered
+    text_x = rect_x0 + padding_x
+    text_y = rect_y0 + padding_y // 2
+    draw.text((text_x, text_y), text, fill=(255, 255, 255), font=font)
+
+    # Save annotated image (unique filename)
+    # sanitize breed name for filename
+    safe_breed = "".join(c if c.isalnum() or c in (' ', '_', '-') else '_' for c in predicted_breed).replace(' ', '_')
+    annotated_name = f"{safe_breed}_{conf:.2f}pct_{stem}.png"
+    annotated.save(annotated_name)
 
     # CSV output
     df = pd.DataFrame([{
@@ -81,43 +120,46 @@ 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, annotated_name, csv_name
+    # Return values: breed string, confidence (float), preview image path, file path for image download, csv path
+    return predicted_breed, round(conf, 2), annotated_name, annotated_name, csv_name
 
 # ----------------------------
-# UI (Modern Layout, Fixed Downloads)
+# UI: Gradio Blocks - clean 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(
-        "<p style='text-align:center;'>Upload an image → Detect the breed → View prediction confidence → Download results as CSV or Annotated Image.</p>"
-    )
+css = """
+/* center buttons and give nicer spacing */
+#detect-btn { width: 100%; }
+.gradio-container { max-width: 1100px; margin: auto; }
+.gr-box { min-height: auto !important; }
+"""
 
-    with gr.Row():
-        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 (%)")
+with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+    gr.Markdown("<h1 style='text-align:center; margin-bottom:0.2rem;'>🐄 GoVed AI — Indian Cattle/Buffalo Breed Detection</h1>")
+    gr.Markdown("<p style='text-align:center; color:var(--muted); margin-top:0.2rem;'>Upload an image → click Detect → see breed, confidence, annotated preview and download results.</p>")
 
     with gr.Row():
         with gr.Column(scale=1):
-            img_preview = gr.Image(type="filepath", label="Annotated Image Preview", height=300)
-            img_download = gr.File(label="⬇️ Download Annotated Image", type="file")
-
+            img_input = gr.Image(type="filepath", label="📤 Upload Cattle/Buffalo Image", height=340)
+            detect_btn = gr.Button("🔍 Detect Breed", elem_id="detect-btn", variant="primary")
         with gr.Column(scale=1):
-            csv_download = gr.File(label="⬇️ Download CSV", type="file")
-
-    # Button Click Logic
-    submit_btn.click(
+            breed_output = gr.Textbox(label="Predicted Breed", interactive=False)
+            confidence_output = gr.Number(label="Prediction Confidence (%)", interactive=False, precision=2)
+            # preview of annotated image
+            annotated_preview = gr.Image(type="filepath", label="Annotated Image Preview", height=340)
+            # download buttons (use type='filepath' so gradio serves the saved file)
+            with gr.Row():
+                image_download = gr.File(label="⬇️ Download Annotated Image", type="filepath")
+                csv_download = gr.File(label="⬇️ Download CSV", type="filepath")
+
+    # Hook up the button
+    detect_btn.click(
         fn=predict_image,
         inputs=img_input,
-        outputs=[breed_output, confidence_output, img_preview, img_download, csv_download],
+        outputs=[breed_output, confidence_output, annotated_preview, image_download, csv_download],
     )
 
 # ----------------------------
-# Run
+# Launch
 # ----------------------------
 if __name__ == "__main__":
     demo.launch()