import pathlib

import cv2
import gradio as gr
import numpy as np

MAX_W, MAX_H = 600, 400


def update_input_image(image: np.ndarray):
    if image is None:
        return gr.Image.update(value=None)

    h, w = image.shape[:2]
    scale = min(MAX_W / w, MAX_H / h)
    if scale < 1:
        image = cv2.resize(image, None, fx=scale, fy=scale)
    return gr.Image.update(value=image)


def set_example_image(example):
    return [
        gr.Image.update(value=example[0]),
        gr.Image.update(value=example[0])
    ]


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            input_image1 = gr.Image(label='Input',
                                    type='numpy',
                                    shape=(MAX_W, MAX_H))
        with gr.Column():
            output_image1 = gr.Image(label='Output')

    with gr.Row():
        with gr.Column():
            input_image2 = gr.Image(label='Input', type='numpy')
        with gr.Column():
            output_image2 = gr.Image(label='Output')

    button = gr.Button('Run')

    with gr.Row():
        paths = sorted(pathlib.Path('images').rglob('*.jpg'))
        examples = gr.Dataset(components=[input_image1, input_image2],
                              samples=[[path.as_posix(),
                                        path.as_posix()] for path in paths])

    input_image2.change(fn=update_input_image,
                        inputs=input_image2,
                        outputs=input_image2)
    button.click(fn=lambda x: x, inputs=input_image1, outputs=output_image1)
    button.click(fn=lambda x: x, inputs=input_image2, outputs=output_image2)
    examples.click(fn=set_example_image,
                   inputs=[examples],
                   outputs=[input_image1, input_image2])

demo.launch()