Spaces:
Running
Running
File size: 1,196 Bytes
60b22df 67dab86 60b22df 67dab86 60b22df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
# Load model directly
import os
from transformers import AutoModel, AutoTokenizer
from PIL import Image
import uuid
UPLOAD_FOLDER = "./uploads"
RESULTS_FOLDER = "./results"
for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
if not os.path.exists(folder):
os.makedirs(folder)
class OCRModel:
tokenizer: AutoTokenizer
model: AutoModel
def __init__(self):
self.tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
self.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=self.tokenizer.eos_token_id)
self.model = self.model.eval()
def chat(self, image: Image.Image) -> str:
unique_id = str(uuid.uuid4())
image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.txt")
image.save(image_path)
res = self.model.chat(self.tokenizer, image_path, ocr_type='ocr')
os.remove(image_path) # delete file create to avoid memory issue and data shared online
return res
ocr_model = OCRModel()
|