import gradio as gr
import yolov5

# Load models
model_paths = {
    "YOLOv5 Large 640p": "weights/megafishdetector_v0_yolov5l_640p.pt",
    "YOLOv5 Medium 1280p": "weights/megafishdetector_v0_yolov5m_1280p.pt",
    "YOLOv5 Small 640p": "weights/megafishdetector_v0_yolov5s_640p.pt",
}

models = {name: yolov5.load(path) for name, path in model_paths.items()}


def detect_objects(image, model_name):
    model = models[model_name]
    results = model(image)
    return results.render()[0]


# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# megafishdetector")

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="numpy")
            model_selector = gr.Dropdown(
                choices=list(model_paths.keys()), label="Select Model"
            )
            submit_button = gr.Button("Submit")

        with gr.Column():
            image_output = gr.Image(type="numpy")

    submit_button.click(
        fn=detect_objects, inputs=[image_input, model_selector], outputs=image_output
    )

demo.launch()