Update model.py

model.py

@@ -1,8 +1,13 @@
-from pix2tex.cli import LatexOCR
-from pix2tex.model import get_model as pix2tex_model
-from pix2tex.utils import Tokenizer
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from PIL import Image
 
-def get_model(weights='trained_model'):
-    tokenizer = Tokenizer()
-    model = pix2tex_model(checkpoint_path=weights)
-    return model, tokenizer
+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