diff --git a/python/rapidocr_onnxruntime/cal_rec_boxes/main.py b/python/rapidocr_onnxruntime/cal_rec_boxes/main.py index cbc529493..4071633de 100644 --- a/python/rapidocr_onnxruntime/cal_rec_boxes/main.py +++ b/python/rapidocr_onnxruntime/cal_rec_boxes/main.py @@ -28,14 +28,14 @@ def __call__( rec_txt, rec_conf, rec_word_info = rec_res[0], rec_res[1], rec_res[2] h, w = img.shape[:2] img_box = np.array([[0, 0], [w, 0], [w, h], [0, h]]) - word_box_content_list, word_box_list = self.cal_ocr_word_box( + word_box_content_list, word_box_list, conf_list = self.cal_ocr_word_box( rec_txt, img_box, rec_word_info ) word_box_list = self.adjust_box_overlap(copy.deepcopy(word_box_list)) word_box_list = self.reverse_rotate_crop_image( copy.deepcopy(box), word_box_list, direction ) - res.append([rec_txt, rec_conf, word_box_list, word_box_content_list]) + res.append([rec_txt, rec_conf, word_box_list, word_box_content_list, conf_list]) return res @staticmethod @@ -60,13 +60,13 @@ def get_box_direction(box: np.ndarray) -> str: @staticmethod def cal_ocr_word_box( rec_txt: str, box: np.ndarray, rec_word_info: List[Tuple[str, List[int]]] - ) -> Tuple[List[str], List[List[int]]]: + ) -> Tuple[List[str], List[List[int]], List[float]]: """Calculate the detection frame for each word based on the results of recognition and detection of ocr 汉字坐标是单字的 英语坐标是单词级别的 """ - col_num, word_list, word_col_list, state_list = rec_word_info + col_num, word_list, word_col_list, state_list, conf_list = rec_word_info box = box.tolist() bbox_x_start = box[0][0] bbox_x_end = box[1][0] @@ -84,7 +84,7 @@ def cal_ocr_word_box( def cal_char_width(width_list, word_col_): if len(word_col_) == 1: return - char_total_length = (word_col_[-1] - word_col_[0] + 1) * cell_width + char_total_length = (word_col_[-1] - word_col_[0]) * cell_width char_width = char_total_length / (len(word_col_) - 1) width_list.append(char_width) @@ -124,7 +124,7 @@ def cal_box(col_list, width_list, word_box_list_): cal_box(cn_col_list, cn_width_list, word_box_list) cal_box(en_col_list, en_width_list, word_box_list) sorted_word_box_list = sorted(word_box_list, key=lambda box: box[0][0]) - return word_box_content_list, sorted_word_box_list + return word_box_content_list, sorted_word_box_list, conf_list @staticmethod def adjust_box_overlap( diff --git a/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py b/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py index ce3f630bc..41c60bacc 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py @@ -26,9 +26,9 @@ class TextDetector: def __init__(self, config: Dict[str, Any]): - limit_side_len = config.get("limit_side_len", 736) - limit_type = config.get("limit_type", "min") - self.preprocess_op = DetPreProcess(limit_side_len, limit_type) + self.limit_type = config.get("limit_type", "min") + self.limit_side_len = config.get("limit_side_len", 736) + self.preprocess_op = None post_process = { "thresh": config.get("thresh", 0.3), @@ -49,6 +49,7 @@ def __call__(self, img: np.ndarray) -> Tuple[Optional[np.ndarray], float]: raise ValueError("img is None") ori_img_shape = img.shape[0], img.shape[1] + self.preprocess_op = self.get_preprocess(max(img.shape[0], img.shape[1])) prepro_img = self.preprocess_op(img) if prepro_img is None: return None, 0 @@ -59,6 +60,17 @@ def __call__(self, img: np.ndarray) -> Tuple[Optional[np.ndarray], float]: elapse = time.perf_counter() - start_time return dt_boxes, elapse + def get_preprocess(self, max_wh): + if self.limit_type == 'min': + limit_side_len = self.limit_side_len + elif max_wh < 960: + limit_side_len = 960 + elif max_wh < 1500: + limit_side_len = 1500 + else: + limit_side_len = 2000 + return DetPreProcess(limit_side_len, self.limit_type) + def filter_tag_det_res( self, dt_boxes: np.ndarray, image_shape: Tuple[int, int] ) -> np.ndarray: diff --git a/python/rapidocr_onnxruntime/ch_ppocr_rec/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_rec/utils.py index 83b89518d..ace70ad44 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_rec/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_rec/utils.py @@ -92,7 +92,7 @@ def decode( selection &= text_index[batch_idx] != ignored_token if text_prob is not None: - conf_list = text_prob[batch_idx][selection] + conf_list = np.array(text_prob[batch_idx][selection]).tolist() else: conf_list = [1] * len(selection) @@ -116,6 +116,7 @@ def decode( word_list, word_col_list, state_list, + conf_list ], ) ) @@ -147,7 +148,11 @@ def get_word_info( word_list = [] word_col_list = [] state_list = [] - valid_col = np.where(selection == True)[0] + valid_col = np.where(selection)[0] + col_width = np.zeros(valid_col.shape) + if len(valid_col) > 0: + col_width[1:] = valid_col[1:] - valid_col[:-1] + col_width[0] = min(3 if "\u4e00" <= text[0] <= "\u9fff" else 2, int(valid_col[0])) for c_i, char in enumerate(text): if "\u4e00" <= char <= "\u9fff": @@ -155,10 +160,10 @@ def get_word_info( else: c_state = "en&num" - if state == None: + if state is None: state = c_state - if state != c_state: + if state != c_state or col_width[c_i] > 4: if len(word_content) != 0: word_list.append(word_content) word_col_list.append(word_col_content) diff --git a/python/rapidocr_onnxruntime/main.py b/python/rapidocr_onnxruntime/main.py index ae405ee02..3c0300e6b 100644 --- a/python/rapidocr_onnxruntime/main.py +++ b/python/rapidocr_onnxruntime/main.py @@ -334,6 +334,7 @@ def main(): use_det=use_det, use_cls=use_cls, use_rec=use_rec, + **vars(args) ) logger.info(result)