Update app.py
Browse files
app.py
CHANGED
@@ -5,36 +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 |
-
|
10 |
-
|
11 |
-
|
12 |
-
trust_remote_code=True,
|
13 |
-
low_cpu_mem_usage=True,
|
14 |
-
pad_token_id=tokenizer.eos_token_id).eval()
|
15 |
|
16 |
# Check if GPU is available and use it, else use CPU
|
17 |
-
|
18 |
-
|
19 |
|
20 |
# Function to perform OCR on the image
|
21 |
def perform_ocr(image):
|
|
|
22 |
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
|
23 |
-
image.save(temp_file.name)
|
24 |
-
temp_image_path = temp_file.name
|
25 |
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
|
30 |
-
# Create the Gradio interface
|
31 |
interface = gr.Interface(
|
32 |
-
fn=perform_ocr,
|
33 |
-
inputs=gr.Image(type="pil"),
|
34 |
-
outputs=gr.Textbox(),
|
35 |
title="OCR Web App",
|
36 |
description="Upload an image to extract text using the GOT-OCR2.0 model."
|
37 |
)
|
38 |
|
39 |
# Launch the app
|
40 |
-
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', trust_remote_code=True,
|
10 |
+
low_cpu_mem_usage=True,
|
11 |
+
pad_token_id=tokenizer.eos_token_id).eval()
|
|
|
|
|
|
|
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) # Save the image
|
22 |
+
temp_image_path = temp_file.name # Get the file path for the saved image
|
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 using the new syntax
|
29 |
interface = gr.Interface(
|
30 |
+
fn=perform_ocr,
|
31 |
+
inputs=gr.Image(type="pil"), # Updated to gr.Image
|
32 |
+
outputs=gr.Textbox(), # Updated to 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()
|