xiaoyao9184 commited on
Commit
66f48e7
·
verified ·
1 Parent(s): 0dadd16

Synced repo using 'sync_with_huggingface' Github Action

Browse files
Files changed (1) hide show
  1. gradio_app.py +32 -25
gradio_app.py CHANGED
@@ -2,13 +2,17 @@ import os
2
  import sys
3
 
4
  if "APP_PATH" in os.environ:
5
- os.chdir(os.environ["APP_PATH"])
6
- # fix sys.path for import
7
- sys.path.append(os.getcwd())
 
 
 
 
 
8
 
9
  from typing import List
10
 
11
- import gradio as gr
12
  import pypdfium2
13
  from pypdfium2 import PdfiumError
14
 
@@ -52,7 +56,7 @@ def load_table_cached():
52
  def load_ocr_error_cached():
53
  return load_ocr_error_model(), load_ocr_error_processor()
54
 
55
-
56
  def run_ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
57
  # Sample the text from the middle of the PDF
58
  page_middle = page_count // 2
@@ -77,14 +81,14 @@ def run_ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pag
77
  label = "This PDF may have garbled or bad OCR text."
78
  return label, results.labels
79
 
80
-
81
  def text_detection(img) -> (Image.Image, TextDetectionResult):
82
  pred = batch_text_detection([img], det_model, det_processor)[0]
83
  polygons = [p.polygon for p in pred.bboxes]
84
  det_img = draw_polys_on_image(polygons, img.copy())
85
  return det_img, pred
86
 
87
-
88
  def layout_detection(img) -> (Image.Image, LayoutResult):
89
  pred = batch_layout_detection([img], layout_model, layout_processor)[0]
90
  polygons = [p.polygon for p in pred.bboxes]
@@ -92,7 +96,7 @@ def layout_detection(img) -> (Image.Image, LayoutResult):
92
  layout_img = draw_polys_on_image(polygons, img.copy(), labels=labels, label_font_size=18)
93
  return layout_img, pred
94
 
95
-
96
  def table_recognition(img, highres_img, filepath, page_idx: int, use_pdf_boxes: bool, skip_table_detection: bool) -> (Image.Image, List[TableResult]):
97
  if skip_table_detection:
98
  layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])]
@@ -143,6 +147,16 @@ def table_recognition(img, highres_img, filepath, page_idx: int, use_pdf_boxes:
143
  table_img = draw_bboxes_on_image(adjusted_bboxes, highres_img, labels=labels, label_font_size=18, color=colors)
144
  return table_img, table_preds
145
 
 
 
 
 
 
 
 
 
 
 
146
  def open_pdf(pdf_file):
147
  return pypdfium2.PdfDocument(pdf_file)
148
 
@@ -164,22 +178,13 @@ def get_page_image(pdf_file, page_num, dpi=96):
164
  def get_uploaded_image(in_file):
165
  return Image.open(in_file).convert("RGB")
166
 
167
- # Function for OCR
168
- def ocr(img, highres_img, langs: List[str]) -> (Image.Image, OCRResult):
169
- replace_lang_with_code(langs)
170
- img_pred = run_ocr([img], [langs], det_model, det_processor, rec_model, rec_processor, highres_images=[highres_img])[0]
171
-
172
- bboxes = [l.bbox for l in img_pred.text_lines]
173
- text = [l.text for l in img_pred.text_lines]
174
- rec_img = draw_text_on_image(bboxes, text, img.size, langs, has_math="_math" in langs)
175
- return rec_img, img_pred
176
-
177
-
178
- det_model, det_processor = load_det_cached()
179
- rec_model, rec_processor = load_rec_cached()
180
- layout_model, layout_processor = load_layout_cached()
181
- table_model, table_processor = load_table_cached()
182
- ocr_error_model, ocr_error_processor = load_ocr_error_cached()
183
 
184
  with gr.Blocks(title="Surya") as demo:
185
  gr.Markdown("""
@@ -272,6 +277,7 @@ with gr.Blocks(title="Surya") as demo:
272
  inputs=[in_img, in_file, in_num, lang_dd],
273
  outputs=[result_img, result_json]
274
  )
 
275
  def table_rec_img(pil_image, in_file, page_number, use_pdf_boxes, skip_table_detection):
276
  if in_file.endswith('.pdf'):
277
  pil_image_highres = get_page_image(in_file, page_number, dpi=settings.IMAGE_DPI_HIGHRES)
@@ -298,4 +304,5 @@ with gr.Blocks(title="Surya") as demo:
298
  outputs=[result_json]
299
  )
300
 
301
- demo.launch()
 
 
2
  import sys
3
 
4
  if "APP_PATH" in os.environ:
5
+ app_path = os.path.abspath(os.environ["APP_PATH"])
6
+ if os.getcwd() != app_path:
7
+ # fix sys.path for import
8
+ os.chdir(app_path)
9
+ if app_path not in sys.path:
10
+ sys.path.append(app_path)
11
+
12
+ import gradio as gr
13
 
14
  from typing import List
15
 
 
16
  import pypdfium2
17
  from pypdfium2 import PdfiumError
18
 
 
56
  def load_ocr_error_cached():
57
  return load_ocr_error_model(), load_ocr_error_processor()
58
 
59
+ #
60
  def run_ocr_errors(pdf_file, page_count, sample_len=512, max_samples=10, max_pages=15):
61
  # Sample the text from the middle of the PDF
62
  page_middle = page_count // 2
 
81
  label = "This PDF may have garbled or bad OCR text."
82
  return label, results.labels
83
 
84
+ #
85
  def text_detection(img) -> (Image.Image, TextDetectionResult):
86
  pred = batch_text_detection([img], det_model, det_processor)[0]
87
  polygons = [p.polygon for p in pred.bboxes]
88
  det_img = draw_polys_on_image(polygons, img.copy())
89
  return det_img, pred
90
 
91
+ #
92
  def layout_detection(img) -> (Image.Image, LayoutResult):
93
  pred = batch_layout_detection([img], layout_model, layout_processor)[0]
94
  polygons = [p.polygon for p in pred.bboxes]
 
96
  layout_img = draw_polys_on_image(polygons, img.copy(), labels=labels, label_font_size=18)
97
  return layout_img, pred
98
 
99
+ #
100
  def table_recognition(img, highres_img, filepath, page_idx: int, use_pdf_boxes: bool, skip_table_detection: bool) -> (Image.Image, List[TableResult]):
101
  if skip_table_detection:
102
  layout_tables = [(0, 0, highres_img.size[0], highres_img.size[1])]
 
147
  table_img = draw_bboxes_on_image(adjusted_bboxes, highres_img, labels=labels, label_font_size=18, color=colors)
148
  return table_img, table_preds
149
 
150
+ # Function for OCR
151
+ def ocr(img, highres_img, langs: List[str]) -> (Image.Image, OCRResult):
152
+ replace_lang_with_code(langs)
153
+ img_pred = run_ocr([img], [langs], det_model, det_processor, rec_model, rec_processor, highres_images=[highres_img])[0]
154
+
155
+ bboxes = [l.bbox for l in img_pred.text_lines]
156
+ text = [l.text for l in img_pred.text_lines]
157
+ rec_img = draw_text_on_image(bboxes, text, img.size, langs, has_math="_math" in langs)
158
+ return rec_img, img_pred
159
+
160
  def open_pdf(pdf_file):
161
  return pypdfium2.PdfDocument(pdf_file)
162
 
 
178
  def get_uploaded_image(in_file):
179
  return Image.open(in_file).convert("RGB")
180
 
181
+ # Load models if not already loaded in reload mode
182
+ if 'det_model' not in globals():
183
+ det_model, det_processor = load_det_cached()
184
+ rec_model, rec_processor = load_rec_cached()
185
+ layout_model, layout_processor = load_layout_cached()
186
+ table_model, table_processor = load_table_cached()
187
+ ocr_error_model, ocr_error_processor = load_ocr_error_cached()
 
 
 
 
 
 
 
 
 
188
 
189
  with gr.Blocks(title="Surya") as demo:
190
  gr.Markdown("""
 
277
  inputs=[in_img, in_file, in_num, lang_dd],
278
  outputs=[result_img, result_json]
279
  )
280
+ # Run Table Recognition
281
  def table_rec_img(pil_image, in_file, page_number, use_pdf_boxes, skip_table_detection):
282
  if in_file.endswith('.pdf'):
283
  pil_image_highres = get_page_image(in_file, page_number, dpi=settings.IMAGE_DPI_HIGHRES)
 
304
  outputs=[result_json]
305
  )
306
 
307
+ if __name__ == "__main__":
308
+ demo.launch()