import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np


model = YOLO('mouseFeetDetection.pt')  

# Define the prediction function for YOLO model
def predict(image):
    if image is None:
        return None
    image = Image.fromarray(np.array(image))  
    results = model(image) 
    annotated_image = results[0].plot()  
    return Image.fromarray(annotated_image) 

description = """
Welcome to the Gradio interface! This simple app detects mouse feet observing from below.
Upload an image to see the prediction results!
"""

demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(label="Upload an Image"), 
    outputs=gr.Image(label="Prediction Result"),
    live=True,
    title="Mouse Feet Detection",
    description=description
)

demo.launch()