Update app.py

app.py

@@ -3,42 +3,42 @@ from ultralytics import YOLO
 from PIL import Image
 import numpy as np
 
-# Load YOLO model (update path if needed)
-model = YOLO('best_animal_classifier.pt')
+# Load your pretrained YOLOv8 model (replace with your model path if customized)
+model = YOLO("yolov8n.pt")  # Using tiny YOLOv8 pretrained weights for example
 
-class_names = ['butterflies', 'chickens', 'elephants', 'horses', 'spiders', 'squirrels']
+# Define the class names according to your trained model classes
+class_names = ["butterfly", "chicken", "elephant", "horse", "spider", "squirrel"]
 
-def predict_animal(image):
-    if image is None:
-        return {}
-
-    # Convert numpy array input to PIL Image if needed
+def classify_animal(image):
+    # Convert numpy array to PIL Image
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image)
-
-    # Run prediction quietly
-    results = model.predict(image, verbose=False)
-
-    try:
-        probs = results[0].probs.data.cpu().numpy()
-    except AttributeError:
-        # Fallback to uniform probabilities if probs unavailable
-        probs = np.ones(len(class_names)) / len(class_names)
-
-    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🐾 Animal Type Classifier")
-    gr.Markdown("Upload an image of an animal below and get predictions for butterflies, chickens, elephants, horses, spiders, or squirrels.")
-
-    with gr.Row():
-        img_input = gr.Image(type="pil", label="Upload Animal Image")
-        label_output = gr.Label(num_top_classes=6, label="Prediction Scores")
-
-    predict_button = gr.Button("Classify Animal")
-    predict_button.click(fn=predict_animal, inputs=img_input, outputs=label_output)
-
-    gr.Markdown("Developed with Ultralytics YOLO and Gradio framework.")
+    # YOLO expects numpy array input, convert back to np array
+    img_np = np.array(image)
+
+    # Perform prediction
+    results = model(img_np)
+
+    # Extract top prediction
+    if results and len(results) > 0:
+        boxes = results[0].boxes
+        if boxes is not None and len(boxes) > 0:
+            # Get the highest confidence prediction
+            best_box = boxes[0]
+            conf = best_box.conf[0].item()
+            class_id = int(best_box.cls[0].item())
+            class_name = class_names[class_id] if class_id < len(class_names) else "Unknown"
+            return f"Predicted: {class_name} (Confidence: {conf:.2f})"
+    return "No animal detected."
+
+# Build Gradio interface
+iface = gr.Interface(
+    fn=classify_animal, 
+    inputs=gr.Image(type="numpy", label="Upload Animal Image"),
+    outputs="text",
+    title="Real-Time Animal Type Classification",
+    description="Upload an image of an animal to classify it among butterfly, chicken, elephant, horse, spider, and squirrel."
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch()