Class TextDecoder

mrzscanner/utils.py:44–78

class TextDecoder:

    def __init__(
        self,
        *,
        chars_dict: Dict[str, int],
        decode_mode: Optional[Union[DecodeMode, str, int]] = DecodeMode.Default
    ):
        self.chars_dict = chars_dict
        self.chars = {v: k for k, v in self.chars_dict.items()}
        self.decode_mode = DecodeMode.obj_to_enum(decode_mode)

    def decode(self, encode: List[np.ndarray]) -> List[List[str]]:
        encode = np.array(encode)
        if self.decode_mode == DecodeMode.CTC:
            masks = (encode != np.roll(encode, 1)) & (encode != 0)
        elif self.decode_mode in [DecodeMode.Default, DecodeMode.Normal]:
            masks = []
            for row in encode:
                eos_index = np.where(row == self.chars_dict["<EOS>"])[0]
                if eos_index.size > 0:
                    mask = np.zeros_like(row, dtype=bool)
                    mask[:eos_index[0]] = True
                else:
                    mask = np.ones_like(row, dtype=bool)
                mask = mask & (row != self.chars_dict["<PAD>"])
                masks.append(mask)

        chars_list = [''.join([self.chars[idx] for idx in e[m]])
                      for e, m in zip(encode, masks)]

        return chars_list

    def __call__(self, *args, **kwargs) -> List[List[str]]:
        return self.decode(*args, **kwargs)