From 0d70e86a088f0d798b2ee90f4c834deb11f47ed7 Mon Sep 17 00:00:00 2001 From: jesse01 Date: Thu, 12 Dec 2024 09:31:58 +0800 Subject: [PATCH] =?UTF-8?q?feat(rapidocr=5Fonnxruntime):=20=E4=BC=98?= =?UTF-8?q?=E5=8C=96=E5=8D=95=E5=AD=97=E7=AC=A6=E5=9D=90=E6=A0=87=E5=8F=8A?= =?UTF-8?q?=E8=BF=94=E5=9B=9E=E5=8D=95=E5=AD=97=E7=AC=A6=E7=BD=AE=E4=BF=A1?= =?UTF-8?q?=E5=BA=A6=20(#293)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 1、utils.py：优化单字符坐标计算。根据CTC解码特点，增加每个字符长度列表作为判断条件，一般非中文字符2-3个解码宽度，中文3-4个解码宽度，当某个字符与上一字符距离大于4个解码宽度，表示中间有空白，此时进行字段拆分，这样可以减少文字间空白对坐标计算的影响；同时返回每个字符值信度，以方便应用使用 2、cal_rec_boxes/main.py： 配合返回单字符值信度列表的修改 修改计算单字符宽度的分母为总字符数减1，与分母一致，这样计算出的值更加准确 3.main.py： 修复主程序没有传递命令行参数的问题 * 优化文本行检测优化，当limit_type为max，根据图片大小动态设置limit_side_len，以减少漏检 --------- Co-authored-by: jesse01 --- .../rapidocr_onnxruntime/cal_rec_boxes/main.py | 12 ++++++------ .../ch_ppocr_det/text_detect.py | 18 +++++++++++++++--- .../rapidocr_onnxruntime/ch_ppocr_rec/utils.py | 13 +++++++++---- python/rapidocr_onnxruntime/main.py | 1 + 4 files changed, 31 insertions(+), 13 deletions(-) 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)