Vinay15's picture
Update app.py
3534c83 verified
raw
history blame
1.37 kB
import gradio as gr
from transformers import AutoModel, AutoTokenizer
from PIL import Image
import numpy as np
# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval() # Ensure model is set to evaluation mode
# Define the OCR function with error handling
def perform_ocr(image):
try:
# Convert PIL image to RGB format (if necessary)
if image.mode != "RGB":
image = image.convert("RGB")
# Convert the image to a format suitable for the model (if needed)
image_array = np.array(image)
# Perform OCR using the model
res = model.chat(tokenizer, image_array, ocr_type='ocr') # Adjusted to pass the image array
return res
except Exception as e:
return f"An error occurred: {str(e)}"
# Define the Gradio interface
interface = gr.Interface(
fn=perform_ocr,
inputs=gr.Image(type="pil", label="Upload Image"),
outputs=gr.Textbox(label="Extracted Text"),
title="OCR and Document Search Web Application",
description="Upload an image to extract text using the GOT-OCR2_0 model."
)
# Launch the Gradio app
interface.launch()