Tonic commited on
Commit
7a2763f
·
unverified ·
1 Parent(s): f34dca6

add back loadimg

Browse files
Files changed (1) hide show
  1. app.py +11 -2
app.py CHANGED
@@ -5,6 +5,7 @@ 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 = """"
@@ -47,7 +48,12 @@ 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
- 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":
@@ -64,6 +70,10 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, rend
64
  html_content = f.read()
65
  return res, html_content
66
 
 
 
 
 
67
  return res, None
68
 
69
  def update_inputs(task):
@@ -94,7 +104,6 @@ def ocr_demo(image, task, ocr_type, ocr_box, ocr_color):
94
 
95
  with gr.Blocks() as demo:
96
  gr.Markdown(title)
97
- # gr.Markdown()
98
  gr.Markdown(description)
99
  with gr.Row():
100
  with gr.Column():
 
5
  import base64
6
  import spaces
7
  from loadimg import load_img
8
+ from PIL import Image
9
 
10
  title = """# 🙋🏻‍♂️Welcome to Tonic's🫴🏻📸GOT-OCR"""
11
  description = """"
 
48
 
49
  @spaces.GPU
50
  def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
51
+ if isinstance(image, str):
52
+ img_path = image
53
+ else:
54
+ img = Image.open(image)
55
+ img_path = "/tmp/temp_image.png"
56
+ img.save(img_path)
57
  if task == "Plain Text OCR":
58
  res = model.chat(tokenizer, img, ocr_type='ocr')
59
  elif task == "Format Text OCR":
 
70
  html_content = f.read()
71
  return res, html_content
72
 
73
+ # Clean up
74
+ if img_path.startswith("/tmp/"):
75
+ os.remove(img_path)
76
+
77
  return res, None
78
 
79
  def update_inputs(task):
 
104
 
105
  with gr.Blocks() as demo:
106
  gr.Markdown(title)
 
107
  gr.Markdown(description)
108
  with gr.Row():
109
  with gr.Column():