Solo448 commited on
Commit
b0a2b83
·
verified ·
1 Parent(s): 7547a25

Delete app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -182
app.py DELETED
@@ -1,182 +0,0 @@
1
- import gradio as gr
2
- import spaces
3
- import json
4
- from transformers import AutoModel, AutoTokenizer
5
- from PIL import Image
6
- import numpy as np
7
- import os
8
- import base64
9
- import io
10
- import uuid
11
- import tempfile
12
- import time
13
- import shutil
14
- from pathlib import Path
15
-
16
- tokenizer = AutoTokenizer.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True)
17
- model = AutoModel.from_pretrained('srimanth-d/GOT_CPU', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cpu', use_safetensors=True)
18
- model = model.eval()
19
-
20
- UPLOAD_FOLDER = "./uploads"
21
- RESULTS_FOLDER = "./results"
22
-
23
- for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
24
- if not os.path.exists(folder):
25
- os.makedirs(folder)
26
-
27
- def image_to_base64(image):
28
- buffered = io.BytesIO()
29
- image.save(buffered, format="PNG")
30
- return base64.b64encode(buffered.getvalue()).decode()
31
-
32
- @spaces.GPU
33
- def run_GOT(image, got_mode, fine_grained_mode="", ocr_color="", ocr_box=""):
34
- unique_id = str(uuid.uuid4())
35
- image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
36
- result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
37
-
38
- shutil.copy(image, image_path)
39
-
40
- try:
41
- if got_mode == "plain texts OCR":
42
- res = model.chat(tokenizer, image_path, ocr_type='ocr')
43
- return res, None
44
- elif got_mode == "format texts OCR":
45
- res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
46
- elif got_mode == "plain multi-crop OCR":
47
- res = model.chat_crop(tokenizer, image_path, ocr_type='ocr')
48
- return res, None
49
- elif got_mode == "format multi-crop OCR":
50
- res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
51
- elif got_mode == "plain fine-grained OCR":
52
- res = model.chat(tokenizer, image_path, ocr_type='ocr', ocr_box=ocr_box, ocr_color=ocr_color)
53
- return res, None
54
- elif got_mode == "format fine-grained OCR":
55
- res = model.chat(tokenizer, image_path, ocr_type='format', ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file=result_path)
56
-
57
- # res_markdown = f"$$ {res} $$"
58
- res_markdown = res
59
-
60
- if "format" in got_mode and os.path.exists(result_path):
61
- with open(result_path, 'r') as f:
62
- html_content = f.read()
63
- encoded_html = base64.b64encode(html_content.encode('utf-8')).decode('utf-8')
64
- iframe_src = f"data:text/html;base64,{encoded_html}"
65
- iframe = f'<iframe src="{iframe_src}" width="100%" height="600px"></iframe>'
66
- download_link = f'<a href="data:text/html;base64,{encoded_html}" download="result_{unique_id}.html">Download Full Result</a>'
67
- return res_markdown, f"{download_link}<br>{iframe}"
68
- else:
69
- return res_markdown, None
70
- except Exception as e:
71
- return f"Error: {str(e)}", None
72
- finally:
73
- if os.path.exists(image_path):
74
- os.remove(image_path)
75
-
76
- def task_update(task):
77
- if "fine-grained" in task:
78
- return [
79
- gr.update(visible=True),
80
- gr.update(visible=False),
81
- gr.update(visible=False),
82
- ]
83
- else:
84
- return [
85
- gr.update(visible=False),
86
- gr.update(visible=False),
87
- gr.update(visible=False),
88
- ]
89
-
90
- def fine_grained_update(task):
91
- if task == "box":
92
- return [
93
- gr.update(visible=False, value = ""),
94
- gr.update(visible=True),
95
- ]
96
- elif task == 'color':
97
- return [
98
- gr.update(visible=True),
99
- gr.update(visible=False, value = ""),
100
- ]
101
-
102
- def cleanup_old_files():
103
- current_time = time.time()
104
- for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
105
- for file_path in Path(folder).glob('*'):
106
- if current_time - file_path.stat().st_mtime > 3600: # 1 hour
107
- file_path.unlink()
108
-
109
- title_html = """ OCR Multilingual """
110
-
111
- with gr.Blocks() as demo:
112
- gr.HTML(title_html)
113
- gr.Markdown("""
114
- "by Souvik Biswas"
115
-
116
- ### Guidelines
117
- Upload your image below and select your preferred mode. Note that more characters may increase wait times.
118
- - **Plain Texts OCR & Format Texts OCR:** Use these modes for basic image-level OCR.
119
- - **Plain Multi-Crop OCR & Format Multi-Crop OCR:** Ideal for images with complex content, offering higher-quality results.
120
- - **Plain Fine-Grained OCR & Format Fine-Grained OCR:** These modes allow you to specify fine-grained regions on the image for more flexible OCR. Regions can be defined by coordinates or colors (red, blue, or green).
121
-
122
- """)
123
-
124
- with gr.Row():
125
- with gr.Column():
126
- image_input = gr.Image(type="filepath", label="upload your image")
127
- task_dropdown = gr.Dropdown(
128
- choices=[
129
- "plain texts OCR",
130
- "format texts OCR",
131
- "plain multi-crop OCR",
132
- "format multi-crop OCR",
133
- "plain fine-grained OCR",
134
- "format fine-grained OCR",
135
- ],
136
- label="Choose one mode of GOT",
137
- value="plain texts OCR"
138
- )
139
- fine_grained_dropdown = gr.Dropdown(
140
- choices=["box", "color"],
141
- label="fine-grained type",
142
- visible=False
143
- )
144
- color_dropdown = gr.Dropdown(
145
- choices=["red", "green", "blue"],
146
- label="color list",
147
- visible=False
148
- )
149
- box_input = gr.Textbox(
150
- label="input box: [x1,y1,x2,y2]",
151
- placeholder="e.g., [0,0,100,100]",
152
- visible=False
153
- )
154
- submit_button = gr.Button("Submit")
155
-
156
- with gr.Column():
157
- ocr_result = gr.Textbox(label="GOT output")
158
-
159
- with gr.Column():
160
- gr.Markdown("**If you choose the mode with format, the mathpix result will be automatically rendered as follows:**")
161
- html_result = gr.HTML(label="rendered html", show_label=True)
162
-
163
- task_dropdown.change(
164
- task_update,
165
- inputs=[task_dropdown],
166
- outputs=[fine_grained_dropdown, color_dropdown, box_input]
167
- )
168
- fine_grained_dropdown.change(
169
- fine_grained_update,
170
- inputs=[fine_grained_dropdown],
171
- outputs=[color_dropdown, box_input]
172
- )
173
-
174
- submit_button.click(
175
- run_GOT,
176
- inputs=[image_input, task_dropdown, fine_grained_dropdown, color_dropdown, box_input],
177
- outputs=[ocr_result, html_result]
178
- )
179
-
180
- if __name__ == "__main__":
181
- cleanup_old_files()
182
- demo.launch()