diff --git a/python/demo.py b/python/demo.py index e115f06..4d456d6 100644 --- a/python/demo.py +++ b/python/demo.py @@ -19,28 +19,19 @@ with open(image_path, "rb") as f: img = f.read() # result, elapse_list = engine(img, use_det=True, use_cls=False, use_rec=False) -result, elapse_list = engine(img, return_word_box=True) +result = engine(img, return_word_box=True) print(result) -print(elapse_list) +print(result.elapse) -# result, elapse = engine(image_path, use_det=False, use_cls=False, use_rec=True) +boxes = result.boxes +txts = result.txts +scores = result.scores -(boxes, txts, scores, words_boxes, words, words_scores) = list(zip(*result)) font_path = "resources/fonts/FZYTK.TTF" +vis_img = vis(img, result.boxes, result.txts, result.scores, font_path) +cv2.imwrite("vis.png", vis_img) -words_boxes = sum(words_boxes, []) -words_all = sum(words, []) -words_scores = sum(words_scores, []) -vis_img = vis(img, words_boxes, words_all, words_scores, font_path) +words_results = result.word_results +words, words_scores, words_boxes = list(zip(*words_results)) +vis_img = vis(img, words_boxes, words, words_scores, font_path) cv2.imwrite("vis_single.png", vis_img) - -# (boxes, txts, scores, words_boxes, words) = list(zip(*result)) - -# vis_img = vis(img, boxes, txts, scores, font_path) -# cv2.imwrite("vis.png", vis_img) - -# words_boxes = sum(words_boxes, []) -# words_all = sum(words, []) -# words_scores = [1.0] * len(words_boxes) -# vis_img = vis(img, words_boxes, words_all, words_scores, font_path) -# cv2.imwrite("vis_single.png", vis_img) diff --git a/python/rapidocr/cal_rec_boxes/main.py b/python/rapidocr/cal_rec_boxes/main.py index 6d86bd9..dc37119 100644 --- a/python/rapidocr/cal_rec_boxes/main.py +++ b/python/rapidocr/cal_rec_boxes/main.py @@ -28,7 +28,7 @@ class CalRecBoxes: for idx, (img, box) in enumerate(zip(imgs, dt_boxes)): direction = self.get_box_direction(box) - rec_txt = rec_res.line_txts[idx] + rec_txt = rec_res.txts[idx] rec_word_info = rec_res.word_results[idx] h, w = img.shape[:2] diff --git a/python/rapidocr/ch_ppocr_cls/utils.py b/python/rapidocr/ch_ppocr_cls/utils.py index c4e96b2..fc82ffb 100644 --- a/python/rapidocr/ch_ppocr_cls/utils.py +++ b/python/rapidocr/ch_ppocr_cls/utils.py @@ -23,6 +23,11 @@ class TextClsOutput: cls_res: Optional[List[Tuple[str, float]]] = None elapse: Optional[float] = None + def __len__(self): + if self.img_list is None: + return 0 + return len(self.img_list) + class ClsPostProcess: def __init__(self, label_list: List[str]): diff --git a/python/rapidocr/ch_ppocr_det/utils.py b/python/rapidocr/ch_ppocr_det/utils.py index 377e9f2..8210428 100644 --- a/python/rapidocr/ch_ppocr_det/utils.py +++ b/python/rapidocr/ch_ppocr_det/utils.py @@ -16,6 +16,11 @@ class TextDetOutput: scores: Optional[Tuple[float]] = None elapse: float = 0.0 + def __len__(self): + if self.boxes is None: + return 0 + return len(self.boxes) + class DetPreProcess: def __init__( diff --git a/python/rapidocr/ch_ppocr_rec/main.py b/python/rapidocr/ch_ppocr_rec/main.py index a67955c..44ea52b 100644 --- a/python/rapidocr/ch_ppocr_rec/main.py +++ b/python/rapidocr/ch_ppocr_rec/main.py @@ -90,8 +90,8 @@ class TextRecognizer: elapse += time.perf_counter() - start_time all_line_results, all_word_results = list(zip(*rec_res)) - line_txts, line_scores = list(zip(*all_line_results)) - return TextRecOutput(line_txts, line_scores, all_word_results, elapse) + txts, scores = list(zip(*all_line_results)) + return TextRecOutput(txts, scores, all_word_results, elapse) def resize_norm_img(self, img: np.ndarray, max_wh_ratio: float) -> np.ndarray: img_channel, img_height, img_width = self.rec_image_shape diff --git a/python/rapidocr/ch_ppocr_rec/utils.py b/python/rapidocr/ch_ppocr_rec/utils.py index 4070593..d69410e 100644 --- a/python/rapidocr/ch_ppocr_rec/utils.py +++ b/python/rapidocr/ch_ppocr_rec/utils.py @@ -29,13 +29,18 @@ class TextRecInput: @dataclass class TextRecOutput: - line_txts: Optional[Tuple[str]] = None - line_scores: Tuple[float] = (1.0,) + txts: Optional[Tuple[str]] = None + scores: Tuple[float] = (1.0,) word_results: Tuple[Tuple[str, float, Optional[List[List[int]]]]] = ( ("", 1.0, None), ) elapse: Optional[float] = None + def __len__(self): + if self.txts is None: + return 0 + return len(self.txts) + class CTCLabelDecode: def __init__( diff --git a/python/rapidocr/main.py b/python/rapidocr/main.py index bd2451c..3cceb99 100644 --- a/python/rapidocr/main.py +++ b/python/rapidocr/main.py @@ -64,7 +64,7 @@ class RapidOCR: self.cal_rec_boxes = CalRecBoxes() - self.return_paddleocr_format = True + self.return_paddleocr_format = False def __call__( self, @@ -255,7 +255,7 @@ class RapidOCR: self, det_res: TextDetOutput, cls_res: TextClsOutput, rec_res: TextRecOutput ) -> Union[TextDetOutput, TextClsOutput, TextRecOutput, RapidOCROutput]: dt_boxes = det_res.boxes - txt_res = rec_res.line_txts + txt_res = rec_res.txts if dt_boxes is None and txt_res is None and cls_res.cls_res is not None: return cls_res @@ -266,13 +266,13 @@ class RapidOCR: if dt_boxes is None and txt_res is not None: return rec_res - if dt_boxes is not None and rec_res is None: + if dt_boxes is not None and txt_res is None: return det_res ocr_res = RapidOCROutput( boxes=det_res.boxes, - txts=rec_res.line_txts, - scores=rec_res.line_scores, + txts=rec_res.txts, + scores=rec_res.scores, word_results=rec_res.word_results, elapse_list=[det_res.elapse, cls_res.elapse, rec_res.elapse], ) diff --git a/python/rapidocr/utils/typings.py b/python/rapidocr/utils/typings.py index 6847986..0ac02cd 100644 --- a/python/rapidocr/utils/typings.py +++ b/python/rapidocr/utils/typings.py @@ -17,13 +17,19 @@ class RapidOCROutput: elapse: float = field(init=False) def __post_init__(self): - self.elapse = sum(self.elapse_list) + self.elapse = sum(v for v in self.elapse_list if isinstance(v, float)) + + def __len__(self): + if self.txts is None: + return 0 + return len(self.txts) def to_json(self): pass def to_paddleocr_format(self): - """Return format like: + """Refer: https://pypi.org/project/paddleocr/2.7.3/ + Return format like: [ [[[6.0, 2.0], [322.0, 9.0], [320.0, 104.0], [4.0, 97.0]], ['正品促销', 0.99893]], [[[70.0, 98.0], [252.0, 98.0], [252.0, 125.0], [70.0, 125.0]], ['大桶装更划算', 0.9843]] diff --git a/python/tests/test_main.py b/python/tests/test_main.py index b2fc534..b46c578 100644 --- a/python/tests/test_main.py +++ b/python/tests/test_main.py @@ -27,10 +27,10 @@ def test_long_img(): img_url = "https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/long.jpeg" img_path = tests_dir / "long.jpeg" download_file(img_url, save_path=img_path) - result, _ = engine(img_path) + result = engine(img_path) assert result is not None - assert len(result) >= 53 + assert len(result.boxes) >= 53 img_path.unlink() @@ -53,8 +53,8 @@ def test_ort_dml_warning(caplog): def test_mode_one_img(): img_path = tests_dir / "issue_170.png" - result, _ = engine(img_path) - assert result[0][1] == "TEST" + result = engine(img_path) + assert result.txts[0] == "TEST" @pytest.mark.parametrize( @@ -72,8 +72,8 @@ def test_mode_one_img(): ) def test_transparent_img(img_name: str, gt: str): img_path = tests_dir / img_name - result, _ = engine(img_path) - assert result[0][1] == gt + result = engine(img_path) + assert result.txts[0] == gt @pytest.mark.parametrize( @@ -89,54 +89,53 @@ def test_transparent_img(img_name: str, gt: str): ) def test_letterbox_like(img_name, gt_len, gt_first_len): img_path = tests_dir / img_name - result, _ = engine(img_path) + result = engine(img_path) assert len(result) == gt_len - assert result[0][1].lower() == gt_first_len.lower() + assert result.txts[0].lower() == gt_first_len.lower() def test_only_det(): - result, _ = engine(img_path, use_det=True, use_cls=False, use_rec=False) - + result = engine(img_path, use_det=True, use_cls=False, use_rec=False) assert len(result) == 18 def test_only_cls(): img_path = tests_dir / "text_cls.jpg" - result, _ = engine(img_path, use_det=False, use_cls=True, use_rec=False) + result = engine(img_path, use_det=False, use_cls=True, use_rec=False) assert len(result) == 1 - assert result[0][0] == "0" + assert result.cls_res[0][0] == "0" def test_only_rec(): img_path = tests_dir / "text_rec.jpg" - result, _ = engine(img_path, use_det=False, use_cls=False, use_rec=True) + result = engine(img_path, use_det=False, use_cls=False, use_rec=True) assert len(result) == 1 - assert result[0][0] == "韩国小馆" + assert result.txts[0] == "韩国小馆" def test_det_rec(): - result, _ = engine(img_path, use_det=True, use_cls=False, use_rec=True) - assert result[0][1] == "正品促销" + result = engine(img_path, use_det=True, use_cls=False, use_rec=True) assert len(result) == 18 + assert result.txts[0] == "正品促销" def test_cls_rec(): img_path = tests_dir / "text_cls.jpg" - result, _ = engine(img_path, use_det=False, use_cls=True, use_rec=True) + result = engine(img_path, use_det=False, use_cls=True, use_rec=True) assert result is not None assert len(result) == 1 - assert result[0][0] == "韩国小馆" + assert result.txts[0] == "韩国小馆" def test_det_cls_rec(): img = cv2.imread(str(img_path)) - result, _ = engine(img) + result = engine(img) assert result is not None - assert result[0][1] == "正品促销" assert len(result) == 18 + assert result.txts[0] == "正品促销" def test_empty(): @@ -149,41 +148,40 @@ def test_empty(): def test_zeros(): img = np.zeros([640, 640, 3], np.uint8) - result, _ = engine(img) - assert result is None + result = engine(img) + assert result.boxes is None def test_input_str(): - result, _ = engine(str(img_path)) - assert result[0][1] == "正品促销" + result = engine(str(img_path)) assert len(result) == 18 + assert result.txts[0] == "正品促销" def test_input_bytes(): with open(img_path, "rb") as f: - result, _ = engine(f.read()) - assert result[0][1] == "正品促销" + result = engine(f.read()) assert len(result) == 18 + assert result.txts[0] == "正品促销" def test_input_path(): - result, _ = engine(img_path) - assert result[0][1] == "正品促销" + result = engine(img_path) assert len(result) == 18 + assert result.txts[0] == "正品促销" def test_input_parameters(): img_path = tests_dir / "ch_en_num.jpg" engine = RapidOCR(text_score=1) - result, _ = engine(img_path) - - assert result is None + result = engine(img_path) + assert result.boxes is None def test_input_det_parameters(): with pytest.raises(FileNotFoundError) as exc_info: engine = RapidOCR(det_model_path="1.onnx") - result, _ = engine(img_path) + result = engine(img_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError @@ -191,7 +189,7 @@ def test_input_det_parameters(): def test_input_cls_parameters(): with pytest.raises(FileNotFoundError) as exc_info: engine = RapidOCR(cls_model_path="1.onnx") - result, _ = engine(img_path) + result = engine(img_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError @@ -199,7 +197,7 @@ def test_input_cls_parameters(): def test_input_rec_parameters(): with pytest.raises(FileNotFoundError) as exc_info: engine = RapidOCR(rec_model_path="1.onnx") - result, _ = engine(img_path) + result = engine(img_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError @@ -207,11 +205,11 @@ def test_input_rec_parameters(): def test_input_three_ndim_two_channel(): img_npy = tests_dir / "two_dim_image.npy" image_array = np.load(str(img_npy)) - result, _ = engine(image_array) + result = engine(image_array) assert result is not None assert len(result) == 1 - assert result[0][1] == "TREND PLOT REPORT" + assert result.txts[0] == "TREND PLOT REPORT" def test_input_three_ndim_one_channel(): @@ -220,30 +218,59 @@ def test_input_three_ndim_one_channel(): img = img[:, :, 0] img = img[..., None] # (H, W, 1) - result, _ = engine(img) - + result = engine(img) assert len(result) >= 17 -@pytest.mark.skip(reason="no way of currently testing this") @pytest.mark.parametrize( "img_name,words", [ ( "black_font_color_transparent.png", - ["我", "是", "中", "国", "人"], + ("我", "是", "中", "国", "人"), ), ( "text_vertical_words.png", - ["已", "取", "之", "時", "不", "參", "一", "人", "見", "而"], + ( + "已", + "取", + "之", + "時", + "不", + "參", + "一", + "人", + "見", + "而", + "是", + "非", + "不", + "得", + "問", + "之", + "人", + "要", + "取", + "之", + "有", + "是", + "是", + "非", + "非", + "之", + "士", + "師", + "也", + ), ), ( "issue_170.png", - ["T", "E", "S", "T"], + ("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 + result = engine(img_path, return_word_box=True) + txts, _, _ = list(zip(*result.word_results)) + assert txts == words