Tonic commited on
Commit
40d5755
Β·
unverified Β·
1 Parent(s): 8e49713

do it normally

Browse files
Files changed (1) hide show
  1. app.py +29 -23
app.py CHANGED
@@ -41,37 +41,43 @@ GOT-OCR-2.0 can handle:
41
 
42
  model_name = 'ucaslcl/GOT-OCR2_0'
43
 
 
44
  tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
45
- config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
46
  model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
47
  model = model.eval().cuda()
48
  model.config.pad_token_id = tokenizer.eos_token_id
49
 
50
- def image_to_base64(image):
51
- buffered = io.BytesIO()
52
- image.save(buffered, format="PNG")
53
- return base64.b64encode(buffered.getvalue()).decode()
54
 
55
  @spaces.GPU
56
- def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
57
- if task == "Plain Text OCR":
58
- res = model.chat(tokenizer, image, ocr_type='ocr')
59
- elif task == "Format Text OCR":
60
- res = model.chat(tokenizer, image, ocr_type='format')
61
- elif task == "Fine-grained OCR (Box)":
62
- res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box)
63
- elif task == "Fine-grained OCR (Color)":
64
- res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color)
65
- elif task == "Multi-crop OCR":
66
- res = model.chat_crop(tokenizer, image_file=image)
67
- elif task == "Render Formatted OCR":
68
- res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./demo.html')
69
- with open('./demo.html', 'r') as f:
70
- html_content = f.read()
71
- return res, html_content
72
-
73
- return res, None
74
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
75
  def update_inputs(task):
76
  if task == "Plain Text OCR" or task == "Format Text OCR" or task == "Multi-crop OCR":
77
  return [gr.update(visible=False)] * 4
 
41
 
42
  model_name = 'ucaslcl/GOT-OCR2_0'
43
 
44
+
45
  tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
 
46
  model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
47
  model = model.eval().cuda()
48
  model.config.pad_token_id = tokenizer.eos_token_id
49
 
50
+ # def image_to_base64(image):
51
+ # buffered = io.BytesIO()
52
+ # image.save(buffered, format="PNG")
53
+ # return base64.b64encode(buffered.getvalue()).decode()
54
 
55
  @spaces.GPU
56
+ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
57
+ try:
58
+ if image is None:
59
+ return "No image provided", None
 
 
 
 
 
 
 
 
 
 
 
 
 
 
60
 
61
+ if task == "Plain Text OCR":
62
+ res = model.chat(tokenizer, image, ocr_type='ocr')
63
+ elif task == "Format Text OCR":
64
+ res = model.chat(tokenizer, image, ocr_type='format')
65
+ elif task == "Fine-grained OCR (Box)":
66
+ res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box)
67
+ elif task == "Fine-grained OCR (Color)":
68
+ res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color)
69
+ elif task == "Multi-crop OCR":
70
+ res = model.chat_crop(tokenizer, image_file=image)
71
+ elif task == "Render Formatted OCR":
72
+ res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file='./results/demo.html')
73
+ with open('./results/demo.html', 'r') as f:
74
+ html_content = f.read()
75
+ return res, html_content
76
+
77
+ return res, None
78
+ except Exception as e:
79
+ return str(e), None
80
+
81
  def update_inputs(task):
82
  if task == "Plain Text OCR" or task == "Format Text OCR" or task == "Multi-crop OCR":
83
  return [gr.update(visible=False)] * 4