Tonic commited on
Commit
f34dca6
Β·
unverified Β·
1 Parent(s): eb34ab1

add back loadimg

Browse files
Files changed (1) hide show
  1. app.py +9 -7
app.py CHANGED
@@ -47,18 +47,19 @@ model.config.pad_token_id = tokenizer.eos_token_id
47
 
48
  @spaces.GPU
49
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
 
50
  if task == "Plain Text OCR":
51
- res = model.chat(tokenizer, image, ocr_type='ocr')
52
  elif task == "Format Text OCR":
53
- res = model.chat(tokenizer, image, ocr_type='format')
54
  elif task == "Fine-grained OCR (Box)":
55
- res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box)
56
  elif task == "Fine-grained OCR (Color)":
57
- res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color)
58
  elif task == "Multi-crop OCR":
59
- res = model.chat_crop(tokenizer, image_file=image)
60
  elif task == "Render Formatted OCR":
61
- res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./results/demo.html')
62
  with open('./results/demo.html', 'r') as f:
63
  html_content = f.read()
64
  return res, html_content
@@ -157,4 +158,5 @@ with gr.Blocks() as demo:
157
  outputs=[output_text, output_html]
158
  )
159
 
160
- demo.launch()
 
 
47
 
48
  @spaces.GPU
49
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
50
+ img = load_img(image, output_type="pil", input_type="auto")
51
  if task == "Plain Text OCR":
52
+ res = model.chat(tokenizer, img, ocr_type='ocr')
53
  elif task == "Format Text OCR":
54
+ res = model.chat(tokenizer, img, ocr_type='format')
55
  elif task == "Fine-grained OCR (Box)":
56
+ res = model.chat(tokenizer, img, ocr_type=ocr_type, ocr_box=ocr_box)
57
  elif task == "Fine-grained OCR (Color)":
58
+ res = model.chat(tokenizer, img, ocr_type=ocr_type, ocr_color=ocr_color)
59
  elif task == "Multi-crop OCR":
60
+ res = model.chat_crop(tokenizer, image_file=img)
61
  elif task == "Render Formatted OCR":
62
+ res = model.chat(tokenizer, img, ocr_type='format', render=True, save_render_file='./results/demo.html')
63
  with open('./results/demo.html', 'r') as f:
64
  html_content = f.read()
65
  return res, html_content
 
158
  outputs=[output_text, output_html]
159
  )
160
 
161
+ if __name__ == "__main__":
162
+ demo.launch()