Spaces:
Sleeping
Sleeping
File size: 803 Bytes
7a67d04 dd6ab3d e9e571c 7a67d04 dd6ab3d e9e571c 7a67d04 dd6ab3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
import gradio as gr
import tempfile
from PIL import Image
import numpy as np
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='cpu', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
model = model.eval().cpu()
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()
|