diff --git a/python/rapidocr_onnxruntime/cal_rec_boxes/main.py b/python/rapidocr_onnxruntime/cal_rec_boxes/main.py index a1975d0..cbc5294 100644 --- a/python/rapidocr_onnxruntime/cal_rec_boxes/main.py +++ b/python/rapidocr_onnxruntime/cal_rec_boxes/main.py @@ -59,7 +59,7 @@ class CalRecBoxes: @staticmethod def cal_ocr_word_box( - rec_txt: str, box: np.ndarray, rec_word_info: List[Tuple[str, List[int]]] + rec_txt: str, box: np.ndarray, rec_word_info: List[Tuple[str, List[int]]] ) -> Tuple[List[str], List[List[int]]]: """Calculate the detection frame for each word based on the results of recognition and detection of ocr 汉字坐标是单字的 @@ -82,7 +82,7 @@ class CalRecBoxes: en_col_list = [] def cal_char_width(width_list, word_col_): - if len(word_col) == 1: + if len(word_col_) == 1: return char_total_length = (word_col_[-1] - word_col_[0] + 1) * cell_width char_width = char_total_length / (len(word_col_) - 1) @@ -100,8 +100,8 @@ class CalRecBoxes: center_x = (center_idx + 0.5) * cell_width cell_x_start = max(int(center_x - avg_char_width / 2), 0) + bbox_x_start cell_x_end = ( - min(int(center_x + avg_char_width / 2), bbox_x_end - bbox_x_start) - + bbox_x_start + min(int(center_x + avg_char_width / 2), bbox_x_end - bbox_x_start) + + bbox_x_start ) cell = [ [cell_x_start, bbox_y_start], diff --git a/python/rapidocr_openvino/cal_rec_boxes/main.py b/python/rapidocr_openvino/cal_rec_boxes/main.py index 5e0d71b..cbc5294 100644 --- a/python/rapidocr_openvino/cal_rec_boxes/main.py +++ b/python/rapidocr_openvino/cal_rec_boxes/main.py @@ -77,34 +77,26 @@ class CalRecBoxes: word_box_list = [] word_box_content_list = [] cn_width_list = [] + en_width_list = [] cn_col_list = [] - for word, word_col, state in zip(word_list, word_col_list, state_list): - if state == "cn": - if len(word_col) != 1: - char_seq_length = (word_col[-1] - word_col[0] + 1) * cell_width - char_width = char_seq_length / (len(word_col) - 1) - cn_width_list.append(char_width) - cn_col_list += word_col - word_box_content_list += word - else: - cell_x_start = bbox_x_start + int(word_col[0] * cell_width) - cell_x_end = bbox_x_start + int((word_col[-1] + 1) * cell_width) - cell = [ - [cell_x_start, bbox_y_start], - [cell_x_end, bbox_y_start], - [cell_x_end, bbox_y_end], - [cell_x_start, bbox_y_end], - ] - word_box_list.append(cell) - word_box_content_list.append("".join(word)) + en_col_list = [] - if len(cn_col_list) != 0: - if len(cn_width_list) != 0: - avg_char_width = np.mean(cn_width_list) + 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_width = char_total_length / (len(word_col_) - 1) + width_list.append(char_width) + + def cal_box(col_list, width_list, word_box_list_): + if len(col_list) == 0: + return + if len(width_list) != 0: + avg_char_width = np.mean(width_list) else: avg_char_width = (bbox_x_end - bbox_x_start) / len(rec_txt) - for center_idx in cn_col_list: + for center_idx in col_list: center_x = (center_idx + 0.5) * cell_width cell_x_start = max(int(center_x - avg_char_width / 2), 0) + bbox_x_start cell_x_end = ( @@ -117,7 +109,20 @@ class CalRecBoxes: [cell_x_end, bbox_y_end], [cell_x_start, bbox_y_end], ] - word_box_list.append(cell) + word_box_list_.append(cell) + + for word, word_col, state in zip(word_list, word_col_list, state_list): + if state == "cn": + cal_char_width(cn_width_list, word_col) + cn_col_list += word_col + word_box_content_list += word + else: + cal_char_width(en_width_list, word_col) + en_col_list += word_col + word_box_content_list += word + + 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 @@ -256,5 +261,9 @@ class CalRecBoxes: p23[np.where(p23[:, 1] == np.min(p23[:, 1]))], p23[np.where(p23[:, 1] == np.max(p23[:, 1]))], ) - + # 解决单字矩形框重叠导致多个相同框的情况 + p1 = p1[:1, :] + p2 = p2[:1, :] + p3 = p3[:1, :] + p4 = p4[:1, :] return np.array([p1, p2, p3, p4]).reshape((-1, 2)).tolist() diff --git a/python/tests/test_vino.py b/python/tests/test_vino.py index 67c4bce..30b1ca0 100644 --- a/python/tests/test_vino.py +++ b/python/tests/test_vino.py @@ -3,6 +3,7 @@ # @Contact: liekkaskono@163.com import sys from pathlib import Path +from typing import List import cv2 import numpy as np @@ -202,20 +203,24 @@ def test_input_three_ndim_one_channel(): assert len(result) >= 17 -# @pytest.mark.parametrize( -# "img_name,words", -# [ -# ( -# "black_font_color_transparent.png", -# ["我", "是", "中", "国", "人"], -# ), -# ( -# "text_vertical_words.png", -# ["已", "取", "之", "時", "不", "參", "一", "人", "見", "而"], -# ), -# ], -# ) -# def test_word_ocr(img_name: str, words: List[str]): -# img_path = tests_dir / img_name -# result, _ = engine(img_path, return_word_box=True) -# assert result[0][4] == words +@pytest.mark.parametrize( + "img_name,words", + [ + ( + "black_font_color_transparent.png", + ["我", "是", "中", "国", "人"], + ), + ( + "text_vertical_words.png", + ["已", "取", "之", "時", "不", "參", "一", "人", "見", "而"], + ), + ( + "issue_170.png", + ["T", "E", "S", "T"], + ), + ], +) +def test_word_ocr(img_name: str, words: List[str]): + img_path = tests_dir / img_name + result, _ = engine(img_path, return_word_box=True) + assert result[0][4] == words