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

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

# Function for image prediction
def predict_image(image):
    # Convert PIL image to OpenCV format (numpy array)
    image_cv = np.array(image)
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)  # Convert to BGR format for YOLO

    # Perform inference
    results = model(image_cv)

    # Get the annotated image with bounding boxes
    annotated_image = results[0].plot()  # Use the first result and plot annotations

    # Convert annotated image (BGR) back to RGB format for Gradio
    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
    return annotated_image_rgb

# Function for webcam/live video
def predict_webcam(frame):
    # Perform inference
    results = model(frame)

    # Get the annotated frame
    annotated_frame = results[0].plot()  # Use the first result and plot annotations
    return annotated_frame

# Create Gradio Interface
iface = gr.Interface(
    fn=predict_image,  # Function for image upload
    inputs=gr.Image(type="pil", label="Upload an Image"),  # Image input
    outputs=gr.Image(type="numpy", label="Detected Image"),  # Annotated output
    live=False,  # Disable live video for image upload interface
)

iface.launch()