abinet_label_encode.py 1.1 KB

123456789101112131415161718192021222324252627282930313233343536
  1. import numpy as np
  2. from .ctc_label_encode import BaseRecLabelEncode
  3. class ABINetLabelEncode(BaseRecLabelEncode):
  4. """Convert between text-label and text-index."""
  5. def __init__(self,
  6. max_text_length,
  7. character_dict_path=None,
  8. use_space_char=False,
  9. ignore_index=100,
  10. **kwargs):
  11. super(ABINetLabelEncode,
  12. self).__init__(max_text_length, character_dict_path,
  13. use_space_char)
  14. self.ignore_index = ignore_index
  15. def __call__(self, data):
  16. text = data['label']
  17. text = self.encode(text)
  18. if text is None:
  19. return None
  20. if len(text) > self.max_text_len:
  21. return None
  22. data['length'] = np.array(len(text))
  23. text.append(0)
  24. text = text + [self.ignore_index] * (self.max_text_len + 1 - len(text))
  25. data['label'] = np.array(text)
  26. return data
  27. def add_special_char(self, dict_character):
  28. dict_character = ['</s>'] + dict_character
  29. return dict_character