diff --git a/ocrweb/resources/ch_ppocr_mobile_v2_det/text_detect.py b/ocrweb/resources/ch_ppocr_mobile_v2_det/text_detect.py index b6f9f51..bb29e38 100644 --- a/ocrweb/resources/ch_ppocr_mobile_v2_det/text_detect.py +++ b/ocrweb/resources/ch_ppocr_mobile_v2_det/text_detect.py @@ -11,7 +11,6 @@ # 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. -# !/usr/bin/env python # -*- encoding: utf-8 -*- import argparse import time diff --git a/ocrweb/resources/ch_ppocr_mobile_v2_det/utils.py b/ocrweb/resources/ch_ppocr_mobile_v2_det/utils.py index 048e762..1235d2b 100644 --- a/ocrweb/resources/ch_ppocr_mobile_v2_det/utils.py +++ b/ocrweb/resources/ch_ppocr_mobile_v2_det/utils.py @@ -25,8 +25,7 @@ import numpy as np import pyclipper import six from shapely.geometry import Polygon - from PIL import Image - +from PIL import Image class DecodeImage(object): diff --git a/python/onnxruntime_infer/ch_ppocr_mobile_v3_rec/utils.py b/python/onnxruntime_infer/ch_ppocr_mobile_v3_rec/utils.py index 3e063ba..bf9e731 100644 --- a/python/onnxruntime_infer/ch_ppocr_mobile_v3_rec/utils.py +++ b/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 = []