diff --git a/python/rapidocr/cal_rec_boxes/main.py b/python/rapidocr/cal_rec_boxes/main.py index ed905a5..963f3cd 100644 --- a/python/rapidocr/cal_rec_boxes/main.py +++ b/python/rapidocr/cal_rec_boxes/main.py @@ -107,14 +107,24 @@ class CalRecBoxes: is_all_en_num = all(v is WordType.EN_NUM for v in word_info.word_types) - line_cols, char_widths, word_contents = [], [], [] - for word, word_col in zip(word_info.words, word_info.word_cols): + # A WordInfo built without char_confs gives every word the line's score. + char_confs_list = word_info.char_confs or [ + [np.mean(word_info.confs).round(5).tolist()] * len(word) + for word in word_info.words + ] + + line_cols, char_widths, word_contents, word_confs = [], [], [], [] + for word, word_col, char_confs in zip( + word_info.words, word_info.word_cols, char_confs_list + ): if is_all_en_num and not return_single_char_box: line_cols.append(word_col) word_contents.append("".join(word)) + word_confs.append(np.mean(char_confs).round(5).tolist()) else: line_cols.extend(word_col) word_contents.extend(word) + word_confs.extend(char_confs) if len(word_col) == 1: continue @@ -134,7 +144,7 @@ class CalRecBoxes: word_boxes = self.calc_box( line_cols, avg_char_width, avg_col_width, bbox_points ) - return word_contents, word_boxes, word_info.confs + return word_contents, word_boxes, word_confs def calc_en_num_box( self, diff --git a/python/rapidocr/ch_ppocr_rec/typings.py b/python/rapidocr/ch_ppocr_rec/typings.py index a61952a..1bfaea7 100644 --- a/python/rapidocr/ch_ppocr_rec/typings.py +++ b/python/rapidocr/ch_ppocr_rec/typings.py @@ -75,3 +75,4 @@ class WordInfo: word_types: List[WordType] = field(default_factory=list) line_txt_len: float = 0.0 confs: List[float] = field(default_factory=list) + char_confs: List[List[float]] = field(default_factory=list) diff --git a/python/rapidocr/ch_ppocr_rec/utils.py b/python/rapidocr/ch_ppocr_rec/utils.py index 0f8ba15..76c5dd0 100644 --- a/python/rapidocr/ch_ppocr_rec/utils.py +++ b/python/rapidocr/ch_ppocr_rec/utils.py @@ -123,6 +123,14 @@ class CTCLabelDecode: len(token_indices) * wh_ratio_list[batch_idx] / max_wh_ratio ) rec_word_info.confs = conf_list + + # words holds every character of text except spaces, in order. + char_confs = iter( + conf for char, conf in zip(text, conf_list) if not char.isspace() + ) + rec_word_info.char_confs = [ + [next(char_confs) for _ in word] for word in rec_word_info.words + ] result_words_list.append(rec_word_info) return result_list, result_words_list diff --git a/python/tests/test_return_word.py b/python/tests/test_return_word.py index 3328e73..4f44582 100644 --- a/python/tests/test_return_word.py +++ b/python/tests/test_return_word.py @@ -3,8 +3,9 @@ # @Contact: liekkaskono@163.com import sys from pathlib import Path -from typing import List +from typing import List, Tuple +import numpy as np import pytest from pytest import mark @@ -12,6 +13,9 @@ root_dir = Path(__file__).resolve().parent.parent sys.path.append(str(root_dir)) from rapidocr import EngineType, LangRec, ModelType, OCRVersion, RapidOCR +from rapidocr.cal_rec_boxes import CalRecBoxes +from rapidocr.ch_ppocr_rec import TextRecOutput +from rapidocr.ch_ppocr_rec.utils import CTCLabelDecode test_dir = root_dir / "tests" / "test_files" @@ -78,3 +82,37 @@ def test_en_return_single_char_box(engine): result = engine(img_path, return_word_box=True, return_single_char_box=True) txts, _, _ = list(zip(*result.word_results[0])) assert txts[0] == "3" + + +@mark.parametrize( + "chars,return_single_char_box,expected", + [ + ("ab cd", False, [("ab", 0.75), ("cd", 0.65)]), + ("ab cd", True, [("a", 0.9), ("b", 0.6), ("c", 0.8), ("d", 0.5)]), + ("中 文", False, [("中", 0.9), ("文", 0.6)]), + ], +) +def test_word_conf_from_own_chars( + chars: str, return_single_char_box: bool, expected: List[Tuple[str, float]] +): + decoder = CTCLabelDecode(character=["a", "b", "c", "d", "中", "文"]) + + # One character every other step, blanks between. Spaces score 0.99. + probs = iter([0.9, 0.6, 0.8, 0.5]) + steps = [(decoder.dict[c], 0.99 if c == " " else next(probs)) for c in chars] + preds = np.zeros((1, 2 * len(steps), len(decoder.character)), dtype=np.float32) + preds[0, :, 0] = 0.95 + for i, (char_idx, prob) in enumerate(steps): + preds[0, 2 * i, 0] = 0.0 + preds[0, 2 * i, char_idx] = prob + + line_results, word_infos = decoder(preds, return_word_box=True) + assert line_results[0][0] == chars + + img = np.zeros((48, 320, 3), dtype=np.uint8) + box = np.array([[0, 0], [320, 0], [320, 48], [0, 48]], dtype=np.float32) + rec_res = TextRecOutput(imgs=[img], txts=(chars,), word_results=tuple(word_infos)) + rec_res = CalRecBoxes()([img], [box], rec_res, return_single_char_box) + + words = [(txt, conf) for txt, conf, _ in rec_res.word_results[0]] + assert words == expected