diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index e17f738..8fb10d3 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -133,11 +133,13 @@ class RapidOCR: origin_words_points = self._get_origin_points( [one_word_points], op_record, raw_h, raw_w ) - origin_words_points = origin_words_points.astype(np.int32).tolist()[0] + origin_words_points = origin_words_points.astype(np.int32).tolist()[ + 0 + ] origin_words_item.append( (one_word[0], one_word[1], origin_words_points) ) - + if origin_words_item: origin_words.append(tuple(origin_words_item)) rec_res.word_results = tuple(origin_words) @@ -367,7 +369,7 @@ def main(arg_list: Optional[List[str]] = None): cur_dir = args.vis_save_dir if args.return_word_box: - words_results = result.word_results + words_results = sum(result.word_results, ()) words, words_scores, words_boxes = list(zip(*words_results)) vis_img = vis(args.img_path, words_boxes, words, words_scores) save_path = cur_dir / f"{Path(args.img_path).stem}_vis_single.png" diff --git a/python/rapidocr/utils/output.py b/python/rapidocr/utils/output.py index d92d31b..c97af9e 100644 --- a/python/rapidocr/utils/output.py +++ b/python/rapidocr/utils/output.py @@ -59,7 +59,7 @@ class RapidOCROutput: return vis_img # single word vis - words_results = self.word_results + words_results = sum(self.word_results, ()) words, words_scores, words_boxes = list(zip(*words_results)) vis_img = vis( self.img, diff --git a/python/tests/test_main.py b/python/tests/test_main.py index 5a7ea05..11a998f 100644 --- a/python/tests/test_main.py +++ b/python/tests/test_main.py @@ -406,44 +406,14 @@ def test_input_three_ndim_one_channel(engine): ), ( "text_vertical_words.png", - ( - "已", - "取", - "之", - "時", - "不", - "參", - "一", - "人", - "見", - "而", - "是", - "非", - "不", - "得", - "問", - "之", - "人", - "要", - "取", - "之", - "有", - "是", - "是", - "非", - "非", - "之", - "士", - "師", - "也", - ), + ("已", "取", "之", "時", "不", "參", "一", "人", "見", "而"), ), ], ) def test_cn_word_ocr(engine, img_name: str, words: List[str]): img_path = tests_dir / img_name result = engine(img_path, return_word_box=True) - txts, _, _ = list(zip(*result.word_results)) + txts, _, _ = list(zip(*result.word_results[0])) assert txts == words @@ -454,12 +424,12 @@ def test_cn_word_ocr(engine, img_name: str, words: List[str]): def test_en_word_ocr(engine, img_name: str, words: str): img_path = tests_dir / img_name result = engine(img_path, return_word_box=True) - txts, _, _ = list(zip(*result.word_results)) + txts, _, _ = list(zip(*result.word_results[0])) assert txts[0] == words def test_en_return_single_char_box(engine): img_path = tests_dir / "en.jpg" result = engine(img_path, return_word_box=True, return_single_char_box=True) - txts, _, _ = list(zip(*result.word_results)) - assert txts[:3] == ("3", "M", "o") + txts, _, _ = list(zip(*result.word_results[0])) + assert txts[:3] == ("3",)