from mrz_decoder.pipeline import MRZDecoder
from mrz_decoder.model import CharacterModel
import torch
import importlib.resources
import PIL.Image


class MrzReader:
    def __init__(self, model_path=None):
        self.device = torch.device('cpu')
        self.model = CharacterModel()
        if model_path is None:
            path = importlib.resources.open_binary('mrz_decoder', 'char_model.pytorch.model')
            self.model.load_model(path)
        else:
            self.model.load_model(model_path)
        self.model = self.model.model
        self.model.eval()
        self.pipeline = MRZDecoder(model=self.model)
    
    def decode(self, image: PIL.Image.Image, fuzzy=False):
        return self.pipeline.decode_mrz(image, fuzzy=fuzzy)

