File size: 621 Bytes
62f52d3
 
be9c806
62f52d3
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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