|
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() |
|
|
