mohamed12ahmed commited on
Commit
13bb14f
·
verified ·
1 Parent(s): 0407759

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +22 -19
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
- 🖼️ Upload an image of a document (or choose one from examples below).
20
- ✔️ Choose the tasks you want to perform on the document.
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
- device = "cuda" if torch.cuda.is_available() else "cpu"
37
-
 
 
38
  # load model
39
  model = model_init(MODEL_PATH, device)
40
-
41
- # run inference
 
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
- rgb_image = bgr_restored_image
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"]],