Tonic commited on
Commit
71a766f
1 Parent(s): 79e6016

add loadimg !!!!!!

Browse files
Files changed (1) hide show
  1. app.py +8 -6
app.py CHANGED
@@ -4,6 +4,7 @@ from transformers import AutoModel, AutoTokenizer, AutoConfig
4
  import os
5
  import base64
6
  import spaces
 
7
 
8
  title = """# 🙋🏻‍♂️Welcome to Tonic's🫴🏻📸GOT-OCR"""
9
  description = """"
@@ -44,9 +45,10 @@ model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True,
44
  model = model.eval().cuda()
45
  model.config.pad_token_id = tokenizer.eos_token_id
46
 
47
-
48
  @spaces.GPU
49
- def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=True):
 
 
50
  if task == "Plain Text OCR":
51
  res = model.chat(tokenizer, image, ocr_type='ocr')
52
  elif task == "Format Text OCR":
@@ -58,8 +60,8 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, rend
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='./demo.html')
62
- with open('./demo.html', 'r') as f:
63
  html_content = f.read()
64
  return res, html_content
65
 
@@ -85,8 +87,8 @@ def update_inputs(task):
85
  elif task == "Render Formatted OCR":
86
  return [gr.update(visible=False)] * 3 + [gr.update(visible=True)]
87
 
88
- def ocr_demo(image, task, ocr_type, ocr_box, ocr_color, render):
89
- res, html_content = process_image(image, task, ocr_type, ocr_box, ocr_color, render)
90
  if html_content:
91
  return res, html_content
92
  return res, None
 
4
  import os
5
  import base64
6
  import spaces
7
+ from loadimg import load_img
8
 
9
  title = """# 🙋🏻‍♂️Welcome to Tonic's🫴🏻📸GOT-OCR"""
10
  description = """"
 
45
  model = model.eval().cuda()
46
  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
+ # Load the image using loadimg
51
+ img = load_img(image, output_type="pil", input_type="auto")
52
  if task == "Plain Text OCR":
53
  res = model.chat(tokenizer, image, ocr_type='ocr')
54
  elif task == "Format Text OCR":
 
60
  elif task == "Multi-crop OCR":
61
  res = model.chat_crop(tokenizer, image_file=image)
62
  elif task == "Render Formatted OCR":
63
+ res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./results/demo.html')
64
+ with open('./results/demo.html', 'r') as f:
65
  html_content = f.read()
66
  return res, html_content
67
 
 
87
  elif task == "Render Formatted OCR":
88
  return [gr.update(visible=False)] * 3 + [gr.update(visible=True)]
89
 
90
+ def ocr_demo(image, task, ocr_type, ocr_box, ocr_color):
91
+ res, html_content = process_image(image, task, ocr_type, ocr_box, ocr_color)
92
  if html_content:
93
  return res, html_content
94
  return res, None