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

model = YOLO("./best.pt")


def process_img(img: gr.Image):
    result = model.predict(img)
    for r in result:
        im_bgr = r.plot()
        return gr.Image(
            label="Output Image with labels", value=Image.fromarray(im_bgr[..., ::-1])
        )


with gr.Blocks() as demo:
    gr.Markdown(value="Port Classification App")

    with gr.Row():
        with gr.Column():
            upload_img = gr.Image(label="Upload Image", type="pil")
            classify_img_button = gr.Button(value="Process Image")

        with gr.Column():
            output_img = gr.Image(label="Output Image with labels")
    with gr.Row():
        gr.Examples(
            examples=[
                "./examples/01.jpg",
                "./examples/02.jpg",
                "./examples/03.jpg",
                "./examples/04.jpg",
                "./examples/05.jpg",
                "./examples/06.jpg",
            ],
            inputs=upload_img
        )
    classify_img_button.click(fn=process_img, inputs=upload_img, outputs=output_img)

    demo.launch(server_name="0.0.0.0", server_port=7777)