mirror of
https://github.com/RapidAI/RapidOCR.git
synced 2026-11-10 21:10:55 +08:00
feat(rapidocr_onnxruntime): 优化单字符坐标及返回单字符置信度 (#293)
* 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 <jesse01.163.com>
This commit is contained in:
1 parent
5d0d35fc4e
commit
0d70e86a08
4 files changed
+31
-13
No files matched your search
@@ -28,14 +28,14 @@ class CalRecBoxes:
|
||||
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 @@ class CalRecBoxes:
|
||||
@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 @@ class CalRecBoxes:
|
||||
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 @@ class CalRecBoxes:
|
||||
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(
|
||||
|
||||
@@ -26,9 +26,9 @@ from .utils import DBPostProcess, DetPreProcess
|
||||
|
||||
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 @@ class TextDetector:
|
||||
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 @@ class TextDetector:
|
||||
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:
|
||||
|
||||
@@ -92,7 +92,7 @@ class CTCLabelDecode:
|
||||
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 @@ class CTCLabelDecode:
|
||||
word_list,
|
||||
word_col_list,
|
||||
state_list,
|
||||
conf_list
|
||||
],
|
||||
)
|
||||
)
|
||||
@@ -147,7 +148,11 @@ class CTCLabelDecode:
|
||||
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 @@ class CTCLabelDecode:
|
||||
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)
|
||||
|
||||
@@ -334,6 +334,7 @@ def main():
|
||||
use_det=use_det,
|
||||
use_cls=use_cls,
|
||||
use_rec=use_rec,
|
||||
**vars(args)
|
||||
)
|
||||
logger.info(result)
|
||||
|
||||
|
||||
Reference in new issue
Block a user