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

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

def process_image(input_image):
    if input_image is not None:
        results = model(input_image)

        for r in results:
            im_array = r.plot()
            im_array = im_array.astype(np.uint8)
            return im_array

# Create Gradio Interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(),
    outputs=gr.Image(),  # Specify output as Gradio Image
    title="YOLOv8-obb ships detection",
    description="YOLOv8-obb trained on ShipRSImageNet_BAL-2",
    live=True)

# Launch the Gradio interface
iface.launch()