|
import gradio as gr |
|
from transformers import AutoModel, AutoTokenizer |
|
from PIL import Image |
|
import torch |
|
|
|
|
|
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, |
|
device_map='cuda' if torch.cuda.is_available() else 'cpu', |
|
use_safetensors=True, |
|
pad_token_id=tokenizer.eos_token_id |
|
) |
|
model = model.eval() |
|
if torch.cuda.is_available(): |
|
model = model.cuda() |
|
|
|
|
|
def ocr_from_image(image, ocr_type): |
|
if image is None: |
|
return "Please upload an image." |
|
image_path = "uploaded_image.jpg" |
|
image.save(image_path) |
|
res = model.chat(tokenizer, image_path, ocr_type=ocr_type) |
|
return res |
|
|
|
|
|
ocr_types = ["ocr", "format"] |
|
|
|
|
|
iface = gr.Interface( |
|
fn=ocr_from_image, |
|
inputs=[ |
|
gr.File(label="Upload Image", file_types=[".jpg", ".jpeg", ".png"]), |
|
gr.Radio(ocr_types, label="OCR Type", value="ocr") |
|
], |
|
outputs="text", |
|
title="🧠 GOT-OCR2.0 Transformer OCR", |
|
description="Upload an image file and select the OCR type: plain text (`ocr`) or formatted (`format`)." |
|
) |
|
|
|
if __name__ == "__main__": |
|
iface.launch() |
|
|