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

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

example_list = [["examples/" + example] for example in os.listdir("examples")]

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",
    examples=example_list)

# Launch the Gradio interface
iface.launch()