ChinmoyDutta commited on
Commit
029eb57
1 Parent(s): 4163c88

Upload 3 files

Browse files
Files changed (3) hide show
  1. app.py +142 -0
  2. utils/__init__.py +0 -0
  3. utils/processimage.py +70 -0
app.py ADDED
@@ -0,0 +1,142 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import io
2
+ import os
3
+ import time
4
+ import base64
5
+ # import spaces
6
+ import gradio as gr
7
+ from pathlib import Path
8
+ from transformers import AutoModel, AutoTokenizer
9
+
10
+
11
+ # ........................................................................................................
12
+
13
+
14
+ from utils.processimage import run_GOT
15
+
16
+
17
+ UPLOAD_FOLDER = "./uploads"
18
+ RESULTS_FOLDER = "./results"
19
+
20
+ for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
21
+ if not os.path.exists(folder):
22
+ os.makedirs(folder)
23
+
24
+ def image_to_base64(image):
25
+ buffered = io.BytesIO()
26
+ image.save(buffered, format="PNG")
27
+ return base64.b64encode(buffered.getvalue()).decode()
28
+
29
+ # ........................................................................................................
30
+
31
+ def task_update(task):
32
+ if "fine-grained" in task:
33
+ return [
34
+ gr.update(visible=True),
35
+ gr.update(visible=False),
36
+ gr.update(visible=False),
37
+ ]
38
+ else:
39
+ return [
40
+ gr.update(visible=False),
41
+ gr.update(visible=False),
42
+ gr.update(visible=False),
43
+ ]
44
+
45
+ def fine_grained_update(task):
46
+ if task == "box":
47
+ return [
48
+ gr.update(visible=False, value = ""),
49
+ gr.update(visible=True),
50
+ ]
51
+ elif task == 'color':
52
+ return [
53
+ gr.update(visible=True),
54
+ gr.update(visible=False, value = ""),
55
+ ]
56
+
57
+ def cleanup_old_files():
58
+ current_time = time.time()
59
+ for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
60
+ for file_path in Path(folder).glob('*'):
61
+ if current_time - file_path.stat().st_mtime > 3600: # 1 hour
62
+ file_path.unlink()
63
+
64
+
65
+ title_html = """
66
+ <h2> <span class="gradient-text" id="text">General OCR Theory</span><span class="plain-text">Implementation for Demo purposes </span></h2>
67
+ """
68
+
69
+ with gr.Blocks() as demo:
70
+ gr.HTML(title_html)
71
+ gr.Markdown("""
72
+ "This is a demo using G.0.T for Optical Character Recognition "
73
+
74
+ ### Demo Guidelines
75
+ You need to upload your image below and choose one mode of GOT, then click "Submit" to run GOT model. More characters will result in longer wait times.
76
+ - **plain texts OCR & format texts OCR**: The two modes are for the image-level OCR.
77
+ - **plain multi-crop OCR & format multi-crop OCR**: For images with more complex content, you can achieve higher-quality results with these modes.
78
+ - **plain fine-grained OCR & format fine-grained OCR**: In these modes, you can specify fine-grained regions on the input image for more flexible OCR. Fine-grained regions can be coordinates of the box, red color, blue color, or green color.
79
+ - **Warning: Please upload the file .jpeg, .jpg, .png format only. Other Format like PDF, .hvec etc do not work and will result in error.**
80
+ """)
81
+
82
+ with gr.Row():
83
+ with gr.Column():
84
+ image_input = gr.Image(type="filepath", label="upload the image in .jpeg, .jpg, .png format only")
85
+ task_dropdown = gr.Dropdown(
86
+ choices=[
87
+ "plain texts OCR",
88
+ "format texts OCR",
89
+ "plain multi-crop OCR",
90
+ "format multi-crop OCR",
91
+ "plain fine-grained OCR",
92
+ "format fine-grained OCR",
93
+ ],
94
+ label="Choose one mode of GOT",
95
+ value="plain texts OCR"
96
+ )
97
+ fine_grained_dropdown = gr.Dropdown(
98
+ choices=["box", "color"],
99
+ label="fine-grained type",
100
+ visible=False
101
+ )
102
+ color_dropdown = gr.Dropdown(
103
+ choices=["red", "green", "blue"],
104
+ label="color list",
105
+ visible=False
106
+ )
107
+ box_input = gr.Textbox(
108
+ label="input box: [x1,y1,x2,y2]",
109
+ placeholder="e.g., [0,0,100,100]",
110
+ visible=False
111
+ )
112
+ submit_button = gr.Button("Submit-Image")
113
+
114
+ with gr.Column():
115
+ ocr_result = gr.Textbox(label="GOT-OCR output")
116
+
117
+ with gr.Column():
118
+ gr.Markdown("**The mathpix result will be automatically rendered here:**")
119
+ html_result = gr.HTML(label="rendered html", show_label=True)
120
+
121
+
122
+ task_dropdown.change(
123
+ task_update,
124
+ inputs=[task_dropdown],
125
+ outputs=[fine_grained_dropdown, color_dropdown, box_input]
126
+ )
127
+ fine_grained_dropdown.change(
128
+ fine_grained_update,
129
+ inputs=[fine_grained_dropdown],
130
+ outputs=[color_dropdown, box_input]
131
+ )
132
+
133
+ submit_button.click(
134
+ run_GOT,
135
+ inputs=[image_input, task_dropdown, fine_grained_dropdown, color_dropdown, box_input],
136
+ outputs=[ocr_result, html_result]
137
+ )
138
+
139
+
140
+ if __name__ == "__main__":
141
+ cleanup_old_files()
142
+ demo.launch()
utils/__init__.py ADDED
File without changes
utils/processimage.py ADDED
@@ -0,0 +1,70 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import uuid
3
+ import base64
4
+ import numpy as np
5
+ from PIL import Image
6
+ from transformers import AutoModel, AutoTokenizer
7
+ import shutil
8
+
9
+ # ........................................................................................................
10
+
11
+ tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
12
+ model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True)
13
+ model = model.eval().cuda()
14
+
15
+ UPLOAD_FOLDER = "./uploads"
16
+ RESULTS_FOLDER = "./results"
17
+
18
+ # for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
19
+ # if not os.path.exists(folder):
20
+ # os.makedirs(folder)
21
+
22
+ # ........................................................................................................
23
+
24
+ from dotenv import load_dotenv
25
+ token = os.environ.get("HF_TOKEN1")
26
+
27
+ # ........................................................................................................
28
+
29
+ def run_GOT(image, got_mode, fine_grained_mode="", ocr_color="", ocr_box=""):
30
+ unique_id = str(uuid.uuid4())
31
+ image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
32
+ result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
33
+
34
+ shutil.copy(image, image_path)
35
+
36
+ try:
37
+ if got_mode == "plain texts OCR":
38
+ res = model.chat(tokenizer, image_path, ocr_type='ocr')
39
+ return res, None
40
+ elif got_mode == "format texts OCR":
41
+ res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
42
+ elif got_mode == "plain multi-crop OCR":
43
+ res = model.chat_crop(tokenizer, image_path, ocr_type='ocr')
44
+ return res, None
45
+ elif got_mode == "format multi-crop OCR":
46
+ res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
47
+ elif got_mode == "plain fine-grained OCR":
48
+ res = model.chat(tokenizer, image_path, ocr_type='ocr', ocr_box=ocr_box, ocr_color=ocr_color)
49
+ return res, None
50
+ elif got_mode == "format fine-grained OCR":
51
+ res = model.chat(tokenizer, image_path, ocr_type='format', ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file=result_path)
52
+
53
+ # res_markdown = f"$$ {res} $$"
54
+ res_markdown = res
55
+
56
+ if "format" in got_mode and os.path.exists(result_path):
57
+ with open(result_path, 'r') as f:
58
+ html_content = f.read()
59
+ encoded_html = base64.b64encode(html_content.encode('utf-8')).decode('utf-8')
60
+ iframe_src = f"data:text/html;base64,{encoded_html}"
61
+ iframe = f'<iframe src="{iframe_src}" width="100%" height="600px"></iframe>'
62
+ download_link = f'<a href="data:text/html;base64,{encoded_html}" download="result_{unique_id}.html">Download Full Result</a>'
63
+ return res_markdown, f"{download_link}<br>{iframe}"
64
+ else:
65
+ return res_markdown, None
66
+ except Exception as e:
67
+ return f"Error: {str(e)}", None
68
+ finally:
69
+ if os.path.exists(image_path):
70
+ os.remove(image_path)