File size: 3,577 Bytes
e7c7d09 4d657e7 13bb14f d9a7327 13bb14f 4d657e7 97920f6 13bb14f e7c7d09 13bb14f 4d657e7 13bb14f e7c7d09 13bb14f 4d657e7 13bb14f 4d657e7 d9a7327 4d657e7 d9a7327 13bb14f 4d657e7 97920f6 13bb14f 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 81 82 83 |
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>
<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):
# يُرجى ملاحظة أنني سأفرض استخدام CPU هنا لتجنب المشاكل في بيئات التشغيل التي لا تدعم GPU.
# إذا كنت تريد استخدام GPU، يجب أن تضمن أن البيئة تدعمها.
device = "cpu" # تم تعيينه لـ CPU بشكل إجباري
# load model
model = model_init(MODEL_PATH, device)
# Gradio يتعامل مع RGB، ولكن CV2 (المستخدم في ملفات التقييم) يتعامل مع BGR
# التحويل من Gradio (RGB) إلى BGR للنموذج
bgr_image = image[..., ::-1].copy()
# run inference (توقع أن تكون الدالة قادرة على التعامل مع BGR وتحويلها إلى RGB للإخراج)
bgr_restored_image = inference_one_image(model, bgr_image, tasks, device)
# التحويل من BGR إلى RGB قبل الإخراج إلى Gradio
if bgr_restored_image.ndim == 3:
rgb_image = bgr_restored_image[..., ::-1]
else:
# إذا كانت الصورة المُعالجة هي صورة ثنائية (مثل Binarization)، قد تكون قناة واحدة
# يجب التعامل معها حسب ما تُعيده inference_one_image
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])
# لاحظ: يجب أن تكون ملفات الأمثلة (مثل input/218_in.png) متاحة في مجلد 'input' ليتمكن Gradio من تحميلها.
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()
|