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

# Load the YOLO model
m_raw_model = YOLO("M-Raw.pt")
n_raw_model = YOLO("N-Raw.pt")
s_raw_model = YOLO("S-Raw.pt")

def snap(image, model, conf):
    # Convert the image to a numpy array
    image = np.array(image)

    # Run the selected model
    results = None
    if model == "M-Raw":
        results = m_raw_model(image, conf=conf)
    elif model == "N-Raw":
        results = n_raw_model(image, conf=conf)
    elif model == "S-Raw":
        results = s_raw_model(image, conf=conf)

    # Draw the bounding boxes
    resulting_image = results.render()

    # Convert the resulting image to a PIL image
    resulting_image = Image.fromarray(resulting_image)

    # Get the labels
    labels = results.pandas().xyxy[0]["name"].values

    # Sort the labels by their x-value
    labels = labels[np.argsort(results.pandas().xyxy[0]["x"].values)]

    return [resulting_image, labels]


demo = gr.Interface(
    snap,
    [gr.Image(source="webcam", tool=None, streaming=True), gr.inputs.Radio(["M-Raw", "N-Raw", "S-Raw"]), gr.inputs.Slider(0.0, 1.0, 0.5, 0.1, "Confidence")],
    ["image", "labels"],
    title="Baybayin Instance Detection"
)

if __name__ == "__main__":
    demo.launch()