SavlonBhai's picture
Upload app.py
942e7c0 verified
raw
history blame
802 Bytes
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()