HockeyAI / app.py
MehdiH7's picture
Create app.py
82b42af verified
raw
history blame
632 Bytes
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()
# Create Gradio interface
demo = gr.Interface(
fn=predict,
inputs=gr.Image(),
outputs=gr.Image(),
title="YOLOv8 Object Detection",
description="Upload an image to detect objects using YOLOv8"
)
demo.launch()