from pathlib import Path

import gradio as gr

from yolo import YOLO


yolo = YOLO()


def predict(image):
    if image is None:
        return None
    return yolo.detect_image(image)


title = "RDFNet: Real-time Object Detection Framework for Foggy Scenes"

example_images = sorted(str(path) for path in Path("img").glob("*") if path.is_file())


with gr.Blocks() as demo:
    gr.Markdown(f"### {title}")

    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Upload an Image")
            submit_btn = gr.Button("Submit")
        with gr.Column():
            output = gr.Image(type="pil", label="Prediction Result")

    submit_btn.click(fn=predict, inputs=img_input, outputs=output)
    if example_images:
        gr.Examples(examples=example_images, inputs=img_input)


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