text: # copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #    http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  import numpy as np
  import string
  import paddle
  from paddle.nn import functional as F
  
  
  class BaseRecLabelDecode(object):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False):
          support_character_type = [
              'ch', 'en', 'EN_symbol', 'french', 'german', 'japan', 'korean',
              'it', 'xi', 'pu', 'ru', 'ar', 'ta', 'ug', 'fa', 'ur', 'rs', 'oc',
              'rsc', 'bg', 'uk', 'be', 'te', 'ka', 'chinese_cht', 'hi', 'mr',
              'ne', 'EN', 'latin', 'arabic', 'cyrillic', 'devanagari'
          ]
          assert character_type in support_character_type, "Only {} are supported now but get {}".format(
              support_character_type, character_type)
  
          self.beg_str = "sos"
          self.end_str = "eos"
  
          if character_type == "en":
              self.character_str = "0123456789abcdefghijklmnopqrstuvwxyzabcdefghijklmnopqrstuvwxyz"
              dict_character = list(self.character_str)
          elif character_type == "EN_symbol":
              # same with ASTER setting (use 94 char).
              self.character_str = string.printable[:-6]
              dict_character = list(self.character_str)
          elif character_type in support_character_type:
              self.character_str = []
              assert character_dict_path is not None, "character_dict_path should not be None when character_type is {}".format(
                  character_type)
              with open(character_dict_path, "rb") as fin:
                  lines = fin.readlines()
                  for line in lines:
                      line = line.decode('utf-8').strip("\n").strip("\r\n")
                      self.character_str.append(line)
              if use_space_char:
                  self.character_str.append(" ")
              dict_character = list(self.character_str)
  
          else:
              raise NotImplementedError
          self.character_type = character_type
          dict_character = self.add_special_char(dict_character)
          self.dict = {}
          for i, char in enumerate(dict_character):
              self.dict[char] = i
          self.character = dict_character
  
      def add_special_char(self, dict_character):
          return dict_character
  
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          ignored_tokens = self.get_ignored_tokens()
          batch_size = len(text_index)
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] in ignored_tokens:
                      continue
                  if is_remove_duplicate:
                      # only for predict
                      if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                              batch_idx][idx]:
                          continue
                  try:
                      char = self.character[int(text_index[batch_idx][
                          idx])]
                      if char is '!':
                          char = ' '
                      char_list.append(char)
                  except:
                      pass
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
              text = ''.join(char_list)
              result_list.append((text, np.mean(conf_list)))
          return result_list
  
      def get_ignored_tokens(self):
          return [0]  # for ctc blank
  
  
  class CTCLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False,
                   **kwargs):
          super(CTCLabelDecode, self).__init__(character_dict_path,
                                               character_type, use_space_char)
  
      def __call__(self, preds, label=None, *args, **kwargs):
          if isinstance(preds, paddle.Tensor):
              preds = preds.numpy()
          preds_idx = preds.argmax(axis=2)
          preds_prob = preds.max(axis=2)
          text = self.decode(preds_idx, preds_prob, is_remove_duplicate=True)
          if label is None:
              return text
          label = self.decode(label)
          return text, label
  
      def add_special_char(self, dict_character):
          dict_character = ['blank'] + dict_character
          return dict_character
  
  
  class DistillationCTCLabelDecode(CTCLabelDecode):
      """
      Convert 
      Convert between text-label and text-index
      """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False,
                   model_name=["student"],
                   key=None,
                   **kwargs):
          super(DistillationCTCLabelDecode, self).__init__(
              character_dict_path, character_type, use_space_char)
          if not isinstance(model_name, list):
              model_name = [model_name]
          self.model_name = model_name
  
          self.key = key
  
      def __call__(self, preds, label=None, *args, **kwargs):
          output = dict()
          for name in self.model_name:
              pred = preds[name]
              if self.key is not None:
                  pred = pred[self.key]
              output[name] = super().__call__(pred, label=label, *args, **kwargs)
          return output
  
  
  class NRTRLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='EN_symbol',
                   use_space_char=True,
                   **kwargs):
          super(NRTRLabelDecode, self).__init__(character_dict_path,
                                               character_type, use_space_char)
  
      def __call__(self, preds, label=None, *args, **kwargs):
          if len(preds) == 2:
              preds_id = preds[0]
              preds_prob = preds[1]
              if isinstance(preds_id, paddle.Tensor):
                  preds_id = preds_id.numpy()
              if isinstance(preds_prob, paddle.Tensor):
                  preds_prob = preds_prob.numpy()
              if preds_id[0][0] == 2:
                  preds_idx = preds_id[:, 1:]
                  preds_prob = preds_prob[:, 1:]
              else:
                  preds_idx = preds_id
              text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
              if label is None:
                  return text
              label = self.decode(label[:,1:])
          else:
              if isinstance(preds, paddle.Tensor):
                  preds = preds.numpy()
              preds_idx = preds.argmax(axis=2)
              preds_prob = preds.max(axis=2)
              text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
              if label is None:
                  return text
              label = self.decode(label[:,1:])
          return text, label
  
      def add_special_char(self, dict_character):
          dict_character = ['blank','<unk>','<s>','</s>'] + dict_character
          return dict_character
      
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          batch_size = len(text_index)
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] == 3: # end
                      break
                  try:
                      char_list.append(self.character[int(text_index[batch_idx][idx])])
                  except:
                      continue
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
              text = ''.join(char_list)
              result_list.append((text.lower(), np.mean(conf_list)))
          return result_list
  
  
  
  class AttnLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='ch',
                   use_space_char=False,
                   **kwargs):
          super(AttnLabelDecode, self).__init__(character_dict_path,
                                                character_type, use_space_char)
  
      def add_special_char(self, dict_character):
          self.beg_str = "sos"
          self.end_str = "eos"
          dict_character = dict_character
          dict_character = [self.beg_str] + dict_character + [self.end_str]
          return dict_character
  
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          ignored_tokens = self.get_ignored_tokens()
          [beg_idx, end_idx] = self.get_ignored_tokens()
          batch_size = len(text_index)
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] in ignored_tokens:
                      continue
                  if int(text_index[batch_idx][idx]) == int(end_idx):
                      break
                  if is_remove_duplicate:
                      # only for predict
                      if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                              batch_idx][idx]:
                          continue
                  char_list.append(self.character[int(text_index[batch_idx][idx])])
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
              text = ''.join(char_list)
              result_list.append((text, np.mean(conf_list)))
          return result_list
  
      def __call__(self, preds, label=None, *args, **kwargs):
          """
          text = self.decode(text)
          if label is None:
              return text
          else:
              label = self.decode(label, is_remove_duplicate=False)
              return text, label
          """
          if isinstance(preds, paddle.Tensor):
              preds = preds.numpy()
  
          preds_idx = preds.argmax(axis=2)
          preds_prob = preds.max(axis=2)
          text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
          if label is None:
              return text
          label = self.decode(label, is_remove_duplicate=False)
          return text, label
  
      def get_ignored_tokens(self):
          beg_idx = self.get_beg_end_flag_idx("beg")
          end_idx = self.get_beg_end_flag_idx("end")
          return [beg_idx, end_idx]
  
      def get_beg_end_flag_idx(self, beg_or_end):
          if beg_or_end == "beg":
              idx = np.array(self.dict[self.beg_str])
          elif beg_or_end == "end":
              idx = np.array(self.dict[self.end_str])
          else:
              assert False, "unsupport type %s in get_beg_end_flag_idx" \
                            % beg_or_end
          return idx
  
  
  class SRNLabelDecode(BaseRecLabelDecode):
      """ Convert between text-label and text-index """
  
      def __init__(self,
                   character_dict_path=None,
                   character_type='en',
                   use_space_char=False,
                   **kwargs):
          super(SRNLabelDecode, self).__init__(character_dict_path,
                                               character_type, use_space_char)
          self.max_text_length = kwargs.get('max_text_length', 25)
  
      def __call__(self, preds, label=None, *args, **kwargs):
          pred = preds['predict']
          char_num = len(self.character_str) + 2
          if isinstance(pred, paddle.Tensor):
              pred = pred.numpy()
          pred = np.reshape(pred, [-1, char_num])
  
          preds_idx = np.argmax(pred, axis=1)
          preds_prob = np.max(pred, axis=1)
  
          preds_idx = np.reshape(preds_idx, [-1, self.max_text_length])
  
          preds_prob = np.reshape(preds_prob, [-1, self.max_text_length])
  
          text = self.decode(preds_idx, preds_prob)
  
          if label is None:
              text = self.decode(preds_idx, preds_prob, is_remove_duplicate=False)
              return text
          label = self.decode(label)
          return text, label
  
      def decode(self, text_index, text_prob=None, is_remove_duplicate=False):
          """ convert text-index into text-label. """
          result_list = []
          ignored_tokens = self.get_ignored_tokens()
          batch_size = len(text_index)
  
          for batch_idx in range(batch_size):
              char_list = []
              conf_list = []
              for idx in range(len(text_index[batch_idx])):
                  if text_index[batch_idx][idx] in ignored_tokens:
                      continue
                  if is_remove_duplicate:
                      # only for predict
                      if idx > 0 and text_index[batch_idx][idx - 1] == text_index[
                              batch_idx][idx]:
                          continue
                  char_list.append(self.character[int(text_index[batch_idx][
                      idx])])
                  if text_prob is not None:
                      conf_list.append(text_prob[batch_idx][idx])
                  else:
                      conf_list.append(1)
  
              text = ''.join(char_list)
              result_list.append((text, np.mean(conf_list)))
          return result_list
  
      def add_special_char(self, dict_character):
          dict_character = dict_character + [self.beg_str, self.end_str]
          return dict_character
  
      def get_ignored_tokens(self):
          beg_idx = self.get_beg_end_flag_idx("beg")
          end_idx = self.get_beg_end_flag_idx("end")
          return [beg_idx, end_idx]
  
      def get_beg_end_flag_idx(self, beg_or_end):
          if beg_or_end == "beg":
              idx = np.array(self.dict[self.beg_str])
          elif beg_or_end == "end":
              idx = np.array(self.dict[self.end_str])
          else:
              assert False, "unsupport type %s in get_beg_end_flag_idx" \
                            % beg_or_end
          return idx
  
  
  class TableLabelDecode(object):
      """  """
  
      def __init__(self,
                   character_dict_path,
                   **kwargs):
          list_character, list_elem = self.load_char_elem_dict(character_dict_path)
          list_character = self.add_special_char(list_character)
          list_elem = self.add_special_char(list_elem)
          self.dict_character = {}
          self.dict_idx_character = {}
          for i, char in enumerate(list_character):
              self.dict_idx_character[i] = char
              self.dict_character[char] = i
          self.dict_elem = {}
          self.dict_idx_elem = {}
          for i, elem in enumerate(list_elem):
              self.dict_idx_elem[i] = elem
              self.dict_elem[elem] = i
  
      def load_char_elem_dict(self, character_dict_path):
          list_character = []
          list_elem = []
          with open(character_dict_path, "rb") as fin:
              lines = fin.readlines()
              substr = lines[0].decode('utf-8').strip("\n").strip("\r\n").split("\t")
              character_num = int(substr[0])
              elem_num = int(substr[1])
              for cno in range(1, 1 + character_num):
                  character = lines[cno].decode('utf-8').strip("\n").strip("\r\n")
                  list_character.append(character)
              for eno in range(1 + character_num, 1 + character_num + elem_num):
                  elem = lines[eno].decode('utf-8').strip("\n").strip("\r\n")
                  list_elem.append(elem)
          return list_character, list_elem
  
      def add_special_char(self, list_character):
          self.beg_str = "sos"
          self.end_str = "eos"
          list_character = [self.beg_str] + list_character + [self.end_str]
          return list_character
  
      def __call__(self, preds):
          structure_probs = preds['structure_probs']
          loc_preds = preds['loc_preds']
          if isinstance(structure_probs,paddle.Tensor):
              structure_probs = structure_probs.numpy()
          if isinstance(loc_preds,paddle.Tensor):
              loc_preds = loc_preds.numpy()
          structure_idx = structure_probs.argmax(axis=2)
          structure_probs = structure_probs.max(axis=2)
          structure_str, structure_pos, result_score_list, result_elem_idx_list = self.decode(structure_idx,
                                                                                              structure_probs, 'elem')
          res_html_code_list = []
          res_loc_list = []
          batch_num = len(structure_str)
          for bno in range(batch_num):
              res_loc = []
              for sno in range(len(structure_str[bno])):
                  text = structure_str[bno][sno]
                  if text in ['<td>', '<td']:
                      pos = structure_pos[bno][sno]
                      res_loc.append(loc_preds[bno, pos])
              res_html_code = ''.join(structure_str[bno])
              res_loc = np.array(res_loc)
              res_html_code_list.append(res_html_code)
              res_loc_list.append(res_loc)
          return {'res_html_code': res_html_code_list, 'res_loc': res_loc_list, 'res_score_list': result_score_list,
                  'res_elem_idx_list': result_elem_idx_list,'structure_str_list':structure_str}
  
      def decode(self, text_index, structure_probs, char_or_elem):
          """convert text-label into text-index.
          """
          if char_or_elem == "char":
              current_dict = self.dict_idx_character
          else:
              current_dict = self.dict_idx_elem
              ignored_tokens = self.get_ignored_tokens('elem')
              beg_idx, end_idx = ignored_tokens
  
          result_list = []
          result_pos_list = []
          result_score_list = []
          result_elem_idx_list = []
          batch_size = len(text_index)
          for batch_idx in range(batch_size):
              char_list = []
              elem_pos_list = []
              elem_idx_list = []
              score_list = []
              for idx in range(len(text_index[batch_idx])):
                  tmp_elem_idx = int(text_index[batch_idx][idx])
                  if idx > 0 and tmp_elem_idx == end_idx:
                      break
                  if tmp_elem_idx in ignored_tokens:
                      continue
  
                  char_list.append(current_dict[tmp_elem_idx])
                  elem_pos_list.append(idx)
                  score_list.append(structure_probs[batch_idx, idx])
                  elem_idx_list.append(tmp_elem_idx)
              result_list.append(char_list)
              result_pos_list.append(elem_pos_list)
              result_score_list.append(score_list)
              result_elem_idx_list.append(elem_idx_list)
          return result_list, result_pos_list, result_score_list, result_elem_idx_list
  
      def get_ignored_tokens(self, char_or_elem):
          beg_idx = self.get_beg_end_flag_idx("beg", char_or_elem)
          end_idx = self.get_beg_end_flag_idx("end", char_or_elem)
          return [beg_idx, end_idx]
  
      def get_beg_end_flag_idx(self, beg_or_end, char_or_elem):
          if char_or_elem == "char":
              if beg_or_end == "beg":
                  idx = self.dict_character[self.beg_str]
              elif beg_or_end == "end":
                  idx = self.dict_character[self.end_str]
              else:
                  assert False, "Unsupport type %s in get_beg_end_flag_idx of char" \
                                % beg_or_end
          elif char_or_elem == "elem":
              if beg_or_end == "beg":
                  idx = self.dict_elem[self.beg_str]
              elif beg_or_end == "end":
                  idx = self.dict_elem[self.end_str]
              else:
                  assert False, "Unsupport type %s in get_beg_end_flag_idx of elem" \
                                % beg_or_end
          else:
              assert False, "Unsupport type %s in char_or_elem" \
                            % char_or_elem
          return idx
  