app.py

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()