Update app.py
Browse files
app.py
CHANGED
@@ -5,34 +5,34 @@ import gradio as gr
|
|
5 |
import tempfile
|
6 |
|
7 |
# Load the OCR model and tokenizer
|
8 |
-
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-
|
9 |
-
model = AutoModel.from_pretrained('ucaslcl/GOT-
|
10 |
-
|
11 |
-
|
|
|
12 |
|
13 |
# Check if GPU is available and use it, else use CPU
|
14 |
-
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
15 |
-
model = model.to(device)
|
16 |
|
17 |
# Function to perform OCR on the image
|
18 |
def perform_ocr(image):
|
19 |
-
# Save the image to a temporary file
|
20 |
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
|
21 |
-
image.save(temp_file.name)
|
22 |
-
temp_image_path = temp_file.name
|
23 |
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
|
28 |
-
# Create the Gradio interface
|
29 |
interface = gr.Interface(
|
30 |
-
fn=perform_ocr,
|
31 |
-
inputs=gr.Image(type="pil"),
|
32 |
-
outputs=gr.Textbox(),
|
33 |
title="OCR Web App",
|
34 |
description="Upload an image to extract text using the GOT-OCR2.0 model."
|
35 |
)
|
36 |
|
37 |
# Launch the app
|
38 |
-
interface.launch()
|
|
|
5 |
import tempfile
|
6 |
|
7 |
# Load the OCR model and tokenizer
|
8 |
+
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2.0', trust_remote_code=True)
|
9 |
+
model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2.0',
|
10 |
+
trust_remote_code=True,
|
11 |
+
low_cpu_mem_usage=True,
|
12 |
+
pad_token_id=tokenizer.eos_token_id).eval()
|
13 |
|
14 |
# Check if GPU is available and use it, else use CPU
|
15 |
+
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
16 |
+
# model = model.to(device) # Uncomment if you want to use GPU
|
17 |
|
18 |
# Function to perform OCR on the image
|
19 |
def perform_ocr(image):
|
|
|
20 |
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
|
21 |
+
image.save(temp_file.name)
|
22 |
+
temp_image_path = temp_file.name
|
23 |
|
24 |
+
# Perform OCR using the model
|
25 |
+
result = model.chat(tokenizer, temp_image_path, ocr_type='ocr')
|
26 |
+
return result
|
27 |
|
28 |
+
# Create the Gradio interface
|
29 |
interface = gr.Interface(
|
30 |
+
fn=perform_ocr,
|
31 |
+
inputs=gr.Image(type="pil"),
|
32 |
+
outputs=gr.Textbox(),
|
33 |
title="OCR Web App",
|
34 |
description="Upload an image to extract text using the GOT-OCR2.0 model."
|
35 |
)
|
36 |
|
37 |
# Launch the app
|
38 |
+
interface.launch()
|