JustinLin610 commited on
Commit
0509ee0
1 Parent(s): 8c40668
Files changed (1) hide show
  1. app.py +21 -41
app.py CHANGED
@@ -8,8 +8,6 @@ os.system('cd fairseq;'
8
  os.system('cd ezocr;'
9
  'pip install .; cd ..')
10
 
11
- os.system('pip install https://storage.googleapis.com/tensorflow/mac/cpu/tensorflow-1.14.0-py3-none-any.whl;'
12
- 'pip install "modelscope[cv]" -f https://modelscope.oss-cn-beijing.aliyuncs.com/releases/repo.html')
13
 
14
  import torch
15
  import numpy as np
@@ -24,10 +22,6 @@ from typing import List, Tuple
24
  import cv2
25
  from easyocrlite import ReaderLite
26
  import gradio as gr
27
- from modelscope.pipelines import pipeline
28
- from modelscope.utils.constant import Tasks
29
- from modelscope.outputs import OutputKeys
30
- from modelscope.preprocessors.image import load_image
31
 
32
 
33
  # Register refcoco task
@@ -120,8 +114,7 @@ def patch_resize_transform(patch_image_size=480, is_document=False):
120
  return _patch_resize_transform
121
 
122
 
123
- # reader = ReaderLite(gpu=True)
124
- ocr_detection = pipeline(Tasks.ocr_detection, model='damo/cv_resnet18_ocr-detection-line-level_damo')
125
 
126
  overrides={"eval_cider": False, "beam": 5, "max_len_b": 64, "patch_image_size": 480,
127
  "orig_patch_image_size": 224, "no_repeat_ngram_size": 0, "seed": 42}
@@ -173,51 +166,37 @@ def apply_half(t):
173
  return t
174
 
175
 
176
- def ocr(img):
177
- boxes = ocr_detection(img)[OutputKeys.POLYGONS]
178
- image = cv2.imread(img)
179
  out_img = Image.open(img)
180
- ocr_result = list()
181
- for i, box in boxes: # 因为检测结果是四边形,所以用透视变化转为长方形
182
- post1 = box.reshape((4, 2)).astype(np.float32)
183
- width = box[4] - box[0]
184
- height = box[5] - box[1]
185
- post2 = np.float32([[0, 0], [width, 0], [width, height], [0, height]])
186
- M = cv2.getPerspectiveTransform(post1, post2)
187
- new_img = cv2.warpPerspective(image, M, (width, height))
188
- new_img_pil = Image.fromarray(cv2.cvtColor(new_img, cv2.COLOR_BGR2RGB))
189
- # 开启文字识别
190
- sample = construct_sample(task, new_img_pil, cfg.task.patch_image_size)
 
 
 
 
191
  sample = utils.move_to_cuda(sample) if use_cuda else sample
192
  sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
193
 
194
  with torch.no_grad():
195
  result, scores = eval_step(task, generator, models, sample)
196
  ocr_result.append([str(i+1), result[0]['ocr'].replace(' ', '')])
 
197
  result = pd.DataFrame(ocr_result, columns=['Box ID', 'Text'])
198
- # results = get_images(img, reader, text_confidence=0.7, text_threshold=0.4,
199
- # link_threshold=0.43, slope_ths=0., add_margin=0.02)
200
- # box_list, image_list = zip(*results)
201
- draw_boxes(out_img, boxes)
202
- #
203
- # ocr_result = []
204
- # for i, (box, image) in enumerate(zip(box_list, image_list)):
205
- # image = Image.fromarray(image)
206
- # sample = construct_sample(task, image, cfg.task.patch_image_size)
207
- # sample = utils.move_to_cuda(sample) if use_cuda else sample
208
- # sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
209
- #
210
- # with torch.no_grad():
211
- # result, scores = eval_step(task, generator, models, sample)
212
- # ocr_result.append([str(i+1), result[0]['ocr'].replace(' ', '')])
213
- #
214
- # result = pd.DataFrame(ocr_result, columns=['Box ID', 'Text'])
215
 
216
  return out_img, result
217
 
218
 
219
  title = "Chinese OCR"
220
- description = "Gradio Demo for Chinese OCR based on OFA. "\
221
  "Upload your own image or click any one of the examples, and click " \
222
  "\"Submit\" and then wait for the generated OCR result. " \
223
  "\n中文OCR体验区。欢迎上传图片,静待检测文字返回~"
@@ -225,7 +204,8 @@ article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA
225
  "Repo</a></p> "
226
  examples = [['shupai.png'], ['chinese.jpg'], ['gaidao.jpeg'], ['qiaodaima.png'],
227
  ['benpao.jpeg'], ['wanli.png'], ['xsd.jpg']]
228
- io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='filepath', label='Image'),
 
229
  outputs=[gr.outputs.Image(type='pil', label='Image'),
230
  gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')],
231
  title=title, description=description, article=article, examples=examples)
 
8
  os.system('cd ezocr;'
9
  'pip install .; cd ..')
10
 
 
 
11
 
12
  import torch
13
  import numpy as np
 
22
  import cv2
23
  from easyocrlite import ReaderLite
24
  import gradio as gr
 
 
 
 
25
 
26
 
27
  # Register refcoco task
 
114
  return _patch_resize_transform
115
 
116
 
117
+ reader = ReaderLite(gpu=True)
 
118
 
119
  overrides={"eval_cider": False, "beam": 5, "max_len_b": 64, "patch_image_size": 480,
120
  "orig_patch_image_size": 224, "no_repeat_ngram_size": 0, "seed": 42}
 
166
  return t
167
 
168
 
169
+ def ocr(img, task):
 
 
170
  out_img = Image.open(img)
171
+ if task == "General":
172
+ results = get_images(img, reader, text_confidence=0.7, text_threshold=0.43,
173
+ link_threshold=0.43, slope_ths=0., add_margin=0.02)
174
+ elif task == "Document":
175
+ results = get_images(img, reader, text_threshold=0.3, sharp=1.2,
176
+ contrast=1.2, space_ths=1.8)
177
+ else:
178
+ raise NotImplementedError
179
+ box_list, image_list = zip(*results)
180
+ draw_boxes(out_img, box_list)
181
+
182
+ ocr_result = []
183
+ for i, (box, image) in enumerate(zip(box_list, image_list)):
184
+ image = Image.fromarray(image)
185
+ sample = construct_sample(task, image, cfg.task.patch_image_size)
186
  sample = utils.move_to_cuda(sample) if use_cuda else sample
187
  sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
188
 
189
  with torch.no_grad():
190
  result, scores = eval_step(task, generator, models, sample)
191
  ocr_result.append([str(i+1), result[0]['ocr'].replace(' ', '')])
192
+
193
  result = pd.DataFrame(ocr_result, columns=['Box ID', 'Text'])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
194
 
195
  return out_img, result
196
 
197
 
198
  title = "Chinese OCR"
199
+ description = "Gradio Demo for Chinese OCR based on OFA-Base. "\
200
  "Upload your own image or click any one of the examples, and click " \
201
  "\"Submit\" and then wait for the generated OCR result. " \
202
  "\n中文OCR体验区。欢迎上传图片,静待检测文字返回~"
 
204
  "Repo</a></p> "
205
  examples = [['shupai.png'], ['chinese.jpg'], ['gaidao.jpeg'], ['qiaodaima.png'],
206
  ['benpao.jpeg'], ['wanli.png'], ['xsd.jpg']]
207
+ io = gr.Interface(fn=ocr, inputs=[gr.inputs.Image(type='filepath', label='Image'),
208
+ gr.inputs.Radio(choices=["General", "Document"], type="value", default="General", label="Domain")],
209
  outputs=[gr.outputs.Image(type='pil', label='Image'),
210
  gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')],
211
  title=title, description=description, article=article, examples=examples)