import gradio as gr
import spaces
from inference_gradio import inference_one_image, model_init

MODEL_PATH = "./checkpoints/docres.pkl"

model = model_init(MODEL_PATH)
possible_tasks = [
    "dewarping",
    "deshadowing",
    "appearance",
    "deblurring",
    "binarization",
]

@spaces.GPU
def run_tasks(image, tasks):
    bgr_image = image[..., ::-1].copy()
    bgr_restored_image = inference_one_image(model, bgr_image, tasks)
    if bgr_restored_image.ndim == 3:
        rgb_image = bgr_restored_image[..., ::-1]
    else:
        rgb_image = bgr_restored_image
    return rgb_image


with gr.Blocks() as demo:
    with gr.Row():
        input_image = gr.Image(type="numpy")
        output_image = gr.Image(type="numpy")

    task = gr.CheckboxGroup(choices=possible_tasks, label="Choose tasks:")
    button = gr.Button()
    button.click(
        run_tasks, inputs=[input_image, task], outputs=[output_image]
    )
    
demo.launch()