Fix ocrweb bug

python/onnxruntime_infer/ch_ppocr_mobile_v3_rec/utils.py

@@ -10,35 +10,47 @@ import cv2
 import numpy as np
 
 
-class BaseRecLabelDecode(object):
+class CTCLabelDecode(object):
     """ Convert between text-label and text-index """
 
-    def __init__(self,
-                 character_dict_path=None,
-                 character_type='ch',
-                 use_space_char=False):
-        self.character_str = ""
-        assert character_dict_path is not None, "character_dict_path should not be None when character_type is ch"
+    def __init__(self, character_dict_path):
+        super(CTCLabelDecode, self).__init__()
+        self.character_str = []
+        assert character_dict_path is not None, "character_dict_path should not be None"
         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 += line
-        if use_space_char:
-            self.character_str += " "
-        dict_character = list(self.character_str)
+                self.character_str.append(line)
+        self.character_str.append(' ')
+        dict_character = self.character_str
 
-        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 __call__(self, preds, label=None, *args, **kwargs):
+        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
 
-    def decode(self, text_index, text_prob=None, is_remove_duplicate=True):
+    def get_ignored_tokens(self):
+        return [0]  # for ctc blank
+
+    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()
@@ -64,34 +76,6 @@ class BaseRecLabelDecode(object):
             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):
-        preds_idx = preds.argmax(axis=2)
-        preds_prob = preds.max(axis=2)
-        text = self.decode(preds_idx, preds_prob)
-        if label is None:
-            return text
-        label = self.decode(label, is_remove_duplicate=False)
-        return text, label
-
-    def add_special_char(self, dict_character):
-        dict_character = ['blank'] + dict_character
-        return dict_character
-
 
 def get_image_file_list(img_file):
     imgs_lists = []