File size: 2,496 Bytes
e7c7d09
4d657e7
 
 
 
 
d9a7327
 
 
 
 
 
 
 
 
 
 
97920f6
 
 
 
d9a7327
 
4d657e7
 
 
 
 
 
 
 
 
e7c7d09
97920f6
 
e7c7d09
 
 
 
 
 
4d657e7
e7c7d09
4d657e7
 
 
 
e7c7d09
4d657e7
 
 
 
d9a7327
 
 
4d657e7
d9a7327
 
4d657e7
 
97920f6
328661a
 
 
d9a7327
 
328661a
d9a7327
97920f6
d9a7327
 
328661a
 
 
 
 
 
97920f6
4d657e7
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import torch
import gradio as gr
import spaces
from inference_gradio import inference_one_image, model_init

MODEL_PATH = "./checkpoints/docres.pkl"
HEADER = """
<div align="center">
    <p>
        <span style="font-size: 30px; vertical-align: bottom;"> DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks </span>
    </p>
    <p style="margin-top: -15px;">
        <a href="https://arxiv.org/abs/2405.04408" target="_blank" style="color: grey;">ArXiv Paper</a>
        &nbsp;
        <a href="https://github.com/ZZZHANG-jx/DocRes" target="_blank" style="color: grey;">GitHub Repository</a>
    </p>
</div>

🖼️ Upload an image of a document (or choose one from examples below).  
✔️ Choose the tasks you want to perform on the document.  
🚀 Click "Run" and the model will enhance the document according to the selected tasks!
"""


possible_tasks = [
    "dewarping",
    "deshadowing",
    "appearance",
    "deblurring",
    "binarization",
]


@spaces.GPU(duration=60)
def run_tasks(image, tasks):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # load model
    model = model_init(MODEL_PATH, device)

    # run inference
    bgr_image = image[..., ::-1].copy()
    bgr_restored_image = inference_one_image(model, bgr_image, tasks, device)
    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:
    gr.Markdown(HEADER)

    task = gr.CheckboxGroup(choices=possible_tasks, label="Tasks", value=["appearance"])
    with gr.Row():
        input_image = gr.Image(label="Raw Image", type="numpy")
        output_image = gr.Image(label="Enhanced Image", type="numpy")

    button = gr.Button()
    button.click(run_tasks, inputs=[input_image, task], outputs=[output_image])

    gr.Examples(
        examples=[
            ["input/218_in.png", ["dewarping", "deshadowing", "appearance"]],
            ["input/151_in.png", ["dewarping", "deshadowing", "appearance"]],
            ["input/for_debluring.png", ["deblurring"]],
            ["input/for_appearance.png", ["appearance"]],
            ["input/for_deshadowing.jpg", ["deshadowing"]],
            ["input/for_dewarping.png", ["dewarping"]],
            ["input/for_binarization.png", ["binarization"]],
        ],
        inputs=[input_image, task],
        outputs=[output_image],
        fn=run_tasks,
        cache_examples="lazy",
    )

demo.launch()