Tonic commited on
Commit
1a87a19
·
unverified ·
1 Parent(s): 40d5755

do it normally

Browse files
Files changed (1) hide show
  1. app.py +19 -3
app.py CHANGED
@@ -5,8 +5,10 @@ import os
5
  import base64
6
  import spaces
7
  import io
8
- from PIL import Image
9
  import numpy as np
 
 
 
10
 
11
  title = """# 🙋🏻‍♂️Welcome to Tonic's🫴🏻📸GOT-OCR"""
12
  description = """"
@@ -52,12 +54,24 @@ model.config.pad_token_id = tokenizer.eos_token_id
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":
@@ -74,6 +88,8 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
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
 
5
  import base64
6
  import spaces
7
  import io
 
8
  import numpy as np
9
+ from PIL import Image
10
+ import io
11
+
12
 
13
  title = """# 🙋🏻‍♂️Welcome to Tonic's🫴🏻📸GOT-OCR"""
14
  description = """"
 
54
  # image.save(buffered, format="PNG")
55
  # return base64.b64encode(buffered.getvalue()).decode()
56
 
57
+
58
+ def numpy_to_pil(image_np):
59
+ return Image.fromarray(image_np.astype('uint8'), 'RGB')
60
+
61
  @spaces.GPU
62
+ def process_image(image_np, task, ocr_type=None, ocr_box=None, ocr_color=None):
63
  try:
64
+ if image_np is None:
65
  return "No image provided", None
66
 
67
+ image = numpy_to_pil(image_np)
68
+
69
+ with io.BytesIO() as buffer:
70
+ image.save(buffer, format="PNG")
71
+ image_path = "/tmp/temp_image.png"
72
+ with open(image_path, "wb") as f:
73
+ f.write(buffer.getvalue())
74
+
75
  if task == "Plain Text OCR":
76
  res = model.chat(tokenizer, image, ocr_type='ocr')
77
  elif task == "Format Text OCR":
 
88
  html_content = f.read()
89
  return res, html_content
90
 
91
+ os.remove(image_path)
92
+
93
  return res, None
94
  except Exception as e:
95
  return str(e), None