import gradio as gr
from ultralytics import YOLO
from PIL import Image
import torch

# Load your model
model = YOLO('HockeyAI_model_weight.pt')

def predict(image):
    # Convert gradio image to PIL
    if isinstance(image, str):
        image = Image.open(image)
    
    # Run inference
    results = model.predict(image)
    
    # Get the plotted image with predictions
    return results[0].plot()

# Example images
examples = [
    "exm_1.jpg",  
    "exm_3.jpg" 
]

# Create Gradio interface
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(),
    outputs=gr.Image(),
    title="HockeyAI",
    description="Upload an image to detect 7 differnt objects using a finetuned YOLOv8 for Icek Hockey frames.",
    examples=examples 
)

demo.launch()