PolSOL / model.py
MasteredUltraInstinct's picture
Update model.py
62f52d3 verified
raw
history blame contribute delete
621 Bytes
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
class HandwrittenOCR:
def __init__(self, weights="trained_model"):
self.processor = TrOCRProcessor.from_pretrained(weights)
self.model = VisionEncoderDecoderModel.from_pretrained(weights)
def __call__(self, image: Image.Image) -> str:
pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
generated_ids = self.model.generate(pixel_values)
generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text