Spaces:
Running
Running
import gradio as gr | |
import tempfile | |
from PIL import Image | |
from transformers import AutoModel, AutoTokenizer | |
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', use_safetensors=True, pad_token_id=tokenizer.eos_token_id) | |
model = model.eval().cuda() | |
def OCR(image): | |
pil_image = Image.fromarray(image) | |
with tempfile.TemporaryDirectory() as temp_dir: | |
temp_image_path = f"{temp_dir}/processed_image.png" | |
pil_image.save(temp_image_path) | |
res = model.chat(tokenizer, temp_image_path, ocr_type='ocr') | |
return res | |
demo = gr.Interface(fn=OCR, inputs="image", outputs="text") | |
demo.launch() | |