JustinLin610 commited on
Commit
b926706
1 Parent(s): c981ca5
Files changed (1) hide show
  1. app.py +7 -14
app.py CHANGED
@@ -162,23 +162,17 @@ def apply_half(t):
162
  return t
163
 
164
 
165
- def ocr(img, task_type):
166
  out_img = Image.open(img)
167
- if task_type == "General":
168
- results = get_images(img, reader, text_confidence=0.7, text_threshold=0.43,
169
- link_threshold=0.43, slope_ths=0., add_margin=0.02)
170
- elif task_type == "Document":
171
- results = get_images(img, reader, text_threshold=0.3, sharp=1.2,
172
- contrast=1.2, space_ths=1.8)
173
- else:
174
- raise NotImplementedError
175
  box_list, image_list = zip(*results)
176
  draw_boxes(out_img, box_list)
177
 
178
  ocr_result = []
179
  for i, (box, image) in enumerate(zip(box_list, image_list)):
180
  image = Image.fromarray(image)
181
- sample = construct_sample(image, cfg.task.patch_image_size, is_document=False)
182
  sample = utils.move_to_cuda(sample) if use_cuda else sample
183
  sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
184
 
@@ -198,11 +192,10 @@ description = "Gradio Demo for Chinese OCR based on OFA-Base. "\
198
  "\n中文OCR体验区。欢迎上传图片,静待检测文字返回~"
199
  article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \
200
  "Repo</a></p> "
201
- examples = [['shupai.png', 'General'], ['chinese.jpg', 'General'], ['gaidao.jpeg', 'General'],
202
- ['qiaodaima.png', 'General'], ['benpao.jpeg', 'General'], ['wanli.png', 'General'],
203
  ['xsd.jpg', 'General']]
204
- io = gr.Interface(fn=ocr, inputs=[gr.inputs.Image(type='filepath', label='Image'),
205
- gr.inputs.Radio(choices=['General', 'Document'], type='value', default='General', label='Domain')],
206
  outputs=[gr.outputs.Image(type='pil', label='Image'),
207
  gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')],
208
  title=title, description=description, article=article, examples=examples)
 
162
  return t
163
 
164
 
165
+ def ocr(img):
166
  out_img = Image.open(img)
167
+ results = get_images(img, reader, text_confidence=0.7, text_threshold=0.43,
168
+ link_threshold=0.43, slope_ths=0., add_margin=0.02)
 
 
 
 
 
 
169
  box_list, image_list = zip(*results)
170
  draw_boxes(out_img, box_list)
171
 
172
  ocr_result = []
173
  for i, (box, image) in enumerate(zip(box_list, image_list)):
174
  image = Image.fromarray(image)
175
+ sample = construct_sample(image, cfg.task.patch_image_size)
176
  sample = utils.move_to_cuda(sample) if use_cuda else sample
177
  sample = utils.apply_to_sample(apply_half, sample) if use_fp16 else sample
178
 
 
192
  "\n中文OCR体验区。欢迎上传图片,静待检测文字返回~"
193
  article = "<p style='text-align: center'><a href='https://github.com/OFA-Sys/OFA' target='_blank'>OFA Github " \
194
  "Repo</a></p> "
195
+ examples = [['shupai.png'], ['chinese.jpg'], ['gaidao.jpeg'],
196
+ ['qiaodaima.png'], ['benpao.jpeg'], ['wanli.png'],
197
  ['xsd.jpg', 'General']]
198
+ io = gr.Interface(fn=ocr, inputs=gr.inputs.Image(type='filepath', label='Image'),
 
199
  outputs=[gr.outputs.Image(type='pil', label='Image'),
200
  gr.outputs.Dataframe(headers=['Box ID', 'Text'], type='pandas', label='OCR Results')],
201
  title=title, description=description, article=article, examples=examples)