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__":
    demo.launch()