Update app.py

app.py

@@ -1,22 +1,41 @@
-import gradio as gr
-from ultralytics import YOLO
-
-model = YOLO('best_animal_classifier.pt')
-class_names = ['butterflies', 'chickens', 'elephants', 'horses', 'spiders', 'squirrels']
-
-def predict_animal(image):
-    results = model.predict(image, verbose=False)
-    probs = results[0].probs.data.cpu().numpy()
-    return {class_names[i]: float(probs[i]) for i in range(len(class_names))}
-
-demo = gr.Interface(
-    fn=predict_animal,
-    inputs=gr.Image(type="pil", label="Upload Animal Image"),
-    outputs=gr.Label(num_top_classes=6, label="Predictions"),
-    title="🐾 Animal Type Classifier",
-    description="Upload an image to classify: butterflies, chickens, elephants, horses, spiders, or squirrels",
-    theme=gr.themes.Soft()
-)
-
-if __name__ == "__main__":
+import gradio as gr
+from ultralytics import YOLO
+import numpy as np
+
+# Load YOLO model (update path to your model file if needed)
+model = YOLO('best_animal_classifier.pt')
+
+class_names = ['butterflies', 'chickens', 'elephants', 'horses', 'spiders', 'squirrels']
+
+def predict_animal(image):
+    if image is None:
+        return {}
+    # Run prediction without verbose logging for cleaner output
+    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)
+        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()