Update app.py
Browse files
app.py
CHANGED
|
@@ -4,8 +4,7 @@ import spaces
|
|
| 4 |
from inference_gradio import inference_one_image, model_init
|
| 5 |
|
| 6 |
MODEL_PATH = "./checkpoints/docres.pkl"
|
| 7 |
-
HEADER = """
|
| 8 |
-
<div align="center">
|
| 9 |
<p>
|
| 10 |
<span style="font-size: 30px; vertical-align: bottom;"> DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks </span>
|
| 11 |
</p>
|
|
@@ -15,12 +14,9 @@ HEADER = """
|
|
| 15 |
<a href="https://github.com/ZZZHANG-jx/DocRes" target="_blank" style="color: grey;">GitHub Repository</a>
|
| 16 |
</p>
|
| 17 |
</div>
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
🚀 Click "Run" and the model will enhance the document according to the selected tasks!
|
| 22 |
-
"""
|
| 23 |
-
|
| 24 |
|
| 25 |
possible_tasks = [
|
| 26 |
"dewarping",
|
|
@@ -30,36 +26,43 @@ possible_tasks = [
|
|
| 30 |
"binarization",
|
| 31 |
]
|
| 32 |
|
| 33 |
-
|
| 34 |
@spaces.GPU(duration=60)
|
| 35 |
def run_tasks(image, tasks):
|
| 36 |
-
|
| 37 |
-
|
|
|
|
|
|
|
| 38 |
# load model
|
| 39 |
model = model_init(MODEL_PATH, device)
|
| 40 |
-
|
| 41 |
-
#
|
|
|
|
| 42 |
bgr_image = image[..., ::-1].copy()
|
|
|
|
|
|
|
| 43 |
bgr_restored_image = inference_one_image(model, bgr_image, tasks, device)
|
|
|
|
|
|
|
| 44 |
if bgr_restored_image.ndim == 3:
|
| 45 |
rgb_image = bgr_restored_image[..., ::-1]
|
| 46 |
else:
|
| 47 |
-
|
| 48 |
-
|
|
|
|
|
|
|
| 49 |
return rgb_image
|
| 50 |
|
| 51 |
-
|
| 52 |
with gr.Blocks() as demo:
|
| 53 |
gr.Markdown(HEADER)
|
| 54 |
-
|
| 55 |
task = gr.CheckboxGroup(choices=possible_tasks, label="Tasks", value=["appearance"])
|
| 56 |
with gr.Row():
|
| 57 |
input_image = gr.Image(label="Raw Image", type="numpy")
|
| 58 |
output_image = gr.Image(label="Enhanced Image", type="numpy")
|
| 59 |
-
|
| 60 |
button = gr.Button()
|
| 61 |
button.click(run_tasks, inputs=[input_image, task], outputs=[output_image])
|
| 62 |
-
|
|
|
|
| 63 |
gr.Examples(
|
| 64 |
examples=[
|
| 65 |
["input/218_in.png", ["dewarping", "deshadowing", "appearance"]],
|
|
|
|
| 4 |
from inference_gradio import inference_one_image, model_init
|
| 5 |
|
| 6 |
MODEL_PATH = "./checkpoints/docres.pkl"
|
| 7 |
+
HEADER = """<div align="center">
|
|
|
|
| 8 |
<p>
|
| 9 |
<span style="font-size: 30px; vertical-align: bottom;"> DocRes: A Generalist Model Toward Unifying Document Image Restoration Tasks </span>
|
| 10 |
</p>
|
|
|
|
| 14 |
<a href="https://github.com/ZZZHANG-jx/DocRes" target="_blank" style="color: grey;">GitHub Repository</a>
|
| 15 |
</p>
|
| 16 |
</div>
|
| 17 |
+
🖼️ Upload an image of a document (or choose one from examples below).
|
| 18 |
+
✔️ Choose the tasks you want to perform on the document.
|
| 19 |
+
🚀 Click "Run" and the model will enhance the document according to the selected tasks!"""
|
|
|
|
|
|
|
|
|
|
| 20 |
|
| 21 |
possible_tasks = [
|
| 22 |
"dewarping",
|
|
|
|
| 26 |
"binarization",
|
| 27 |
]
|
| 28 |
|
|
|
|
| 29 |
@spaces.GPU(duration=60)
|
| 30 |
def run_tasks(image, tasks):
|
| 31 |
+
# يُرجى ملاحظة أنني سأفرض استخدام CPU هنا لتجنب المشاكل في بيئات التشغيل التي لا تدعم GPU.
|
| 32 |
+
# إذا كنت تريد استخدام GPU، يجب أن تضمن أن البيئة تدعمها.
|
| 33 |
+
device = "cpu" # تم تعيينه لـ CPU بشكل إجباري
|
| 34 |
+
|
| 35 |
# load model
|
| 36 |
model = model_init(MODEL_PATH, device)
|
| 37 |
+
|
| 38 |
+
# Gradio يتعامل مع RGB، ولكن CV2 (المستخدم في ملفات التقييم) يتعامل مع BGR
|
| 39 |
+
# التحويل من Gradio (RGB) إلى BGR للنموذج
|
| 40 |
bgr_image = image[..., ::-1].copy()
|
| 41 |
+
|
| 42 |
+
# run inference (توقع أن تكون الدالة قادرة على التعامل مع BGR وتحويلها إلى RGB للإخراج)
|
| 43 |
bgr_restored_image = inference_one_image(model, bgr_image, tasks, device)
|
| 44 |
+
|
| 45 |
+
# التحويل من BGR إلى RGB قبل الإخراج إلى Gradio
|
| 46 |
if bgr_restored_image.ndim == 3:
|
| 47 |
rgb_image = bgr_restored_image[..., ::-1]
|
| 48 |
else:
|
| 49 |
+
# إذا كانت الصورة المُعالجة هي صورة ثنائية (مثل Binarization)، قد تكون قناة واحدة
|
| 50 |
+
# يجب التعامل معها حسب ما تُعيده inference_one_image
|
| 51 |
+
rgb_image = bgr_restored_image
|
| 52 |
+
|
| 53 |
return rgb_image
|
| 54 |
|
|
|
|
| 55 |
with gr.Blocks() as demo:
|
| 56 |
gr.Markdown(HEADER)
|
|
|
|
| 57 |
task = gr.CheckboxGroup(choices=possible_tasks, label="Tasks", value=["appearance"])
|
| 58 |
with gr.Row():
|
| 59 |
input_image = gr.Image(label="Raw Image", type="numpy")
|
| 60 |
output_image = gr.Image(label="Enhanced Image", type="numpy")
|
| 61 |
+
|
| 62 |
button = gr.Button()
|
| 63 |
button.click(run_tasks, inputs=[input_image, task], outputs=[output_image])
|
| 64 |
+
|
| 65 |
+
# لاحظ: يجب أن تكون ملفات الأمثلة (مثل input/218_in.png) متاحة في مجلد 'input' ليتمكن Gradio من تحميلها.
|
| 66 |
gr.Examples(
|
| 67 |
examples=[
|
| 68 |
["input/218_in.png", ["dewarping", "deshadowing", "appearance"]],
|