Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 from ultralytics import YOLO
+from PIL import Image
 import numpy as np
 
-# Load YOLO model (update path to your model file if needed)
+# Load YOLO model (update path if needed)
 model = YOLO('best_animal_classifier.pt')
 
 class_names = ['butterflies', 'chickens', 'elephants', 'horses', 'spiders', 'squirrels']
@@ -10,32 +11,21 @@ class_names = ['butterflies', 'chickens', 'elephants', 'horses', 'spiders', 'squ
 def predict_animal(image):
     if image is None:
         return {}
-    # Run prediction without verbose logging for cleaner output
+
+    # Convert numpy array input to PIL Image if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image)
+
+    # Run prediction quietly
     results = model.predict(image, verbose=False)
 
-    # Extract the probabilities; fallback if attribute unavailable
     try:
         probs = results[0].probs.data.cpu().numpy()
     except AttributeError:
-        # If 'probs' not available, generate dummy equal probabilities (prevent crash)
+        # Fallback to uniform probabilities if probs unavailable
         probs = np.ones(len(class_names)) / len(class_names)
 
-    # Map class names to probability scores
     return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
 
-# Enhanced UI with modern theme and layout
 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.")
-
-if __name__ == "__main__":
-    demo.launch()
+    gr.Markdown("# 🐾 Animal Type Class