diff --git a/python/rapidocr_onnxruntime/main.py b/python/rapidocr_onnxruntime/main.py index 82e9175..feb63e2 100644 --- a/python/rapidocr_onnxruntime/main.py +++ b/python/rapidocr_onnxruntime/main.py @@ -2,7 +2,6 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import copy -import logging from pathlib import Path from typing import List, Optional, Tuple, Union @@ -114,8 +113,8 @@ class RapidOCR: use_limit_ratio = w / h > self.width_height_ratio if h <= self.min_height or use_limit_ratio: - new_h = w - padding_h = int((new_h - h) / 2) + new_h = max(int(w / self.width_height_ratio), self.min_height) * 2 + padding_h = int(abs(new_h - h) / 2) block_img = cv2.copyMakeBorder( img, padding_h, padding_h, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0) ) diff --git a/python/rapidocr_openvino/main.py b/python/rapidocr_openvino/main.py index 269b952..feb63e2 100644 --- a/python/rapidocr_openvino/main.py +++ b/python/rapidocr_openvino/main.py @@ -2,7 +2,6 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import copy -import logging from pathlib import Path from typing import List, Optional, Tuple, Union @@ -83,6 +82,7 @@ class RapidOCR: det_elapse, cls_elapse, rec_elapse = 0.0, 0.0, 0.0 if use_det: + img, padding_h = self.maybe_add_letterbox(img) dt_boxes, det_elapse = self.auto_text_det(img) if dt_boxes is None: return None, None @@ -95,28 +95,36 @@ class RapidOCR: if use_rec: rec_res, rec_elapse = self.text_rec(img) + if dt_boxes is not None and padding_h > 0: + for box in dt_boxes: + box[:, 1] -= padding_h + ocr_res = self.get_final_res( dt_boxes, cls_res, rec_res, det_elapse, cls_elapse, rec_elapse ) return ocr_res - def auto_text_det( - self, - img: np.ndarray, - ) -> Tuple[Optional[np.ndarray], float, Optional[List[np.ndarray]]]: + def maybe_add_letterbox(self, img: np.ndarray) -> Tuple[np.ndarray, int]: h, w = img.shape[:2] + if self.width_height_ratio == -1: use_limit_ratio = False else: use_limit_ratio = w / h > self.width_height_ratio if h <= self.min_height or use_limit_ratio: - logging.warning( - "Because the aspect ratio of the current image exceeds the limit (min_height or width_height_ratio), the program will skip the detection step." + new_h = max(int(w / self.width_height_ratio), self.min_height) * 2 + padding_h = int(abs(new_h - h) / 2) + block_img = cv2.copyMakeBorder( + img, padding_h, padding_h, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0) ) - dt_boxes = self.get_boxes_img_without_det(h, w) - return dt_boxes, 0.0 + return block_img, padding_h + return img, 0 + def auto_text_det( + self, + img: np.ndarray, + ) -> Tuple[Optional[np.ndarray], float, Optional[List[np.ndarray]]]: dt_boxes, det_elapse = self.text_det(img) if dt_boxes is None or len(dt_boxes) < 1: return None, 0.0 @@ -124,12 +132,6 @@ class RapidOCR: dt_boxes = self.sorted_boxes(dt_boxes) return dt_boxes, det_elapse - def get_boxes_img_without_det(self, h, w): - x0, y0, x1, y1 = 0, 0, w, h - dt_boxes = np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) - dt_boxes = dt_boxes[np.newaxis, ...] - return dt_boxes - def get_crop_img_list(self, img, dt_boxes): def get_rotate_crop_image(img, points): img_crop_width = int( @@ -241,7 +243,10 @@ def main(): use_cls = not args.no_cls use_rec = not args.no_rec result, elapse_list = ocr_engine( - args.img_path, use_det=use_det, use_cls=use_cls, use_rec=use_rec + args.img_path, + use_det=use_det, + use_cls=use_cls, + use_rec=use_rec, ) print(result) diff --git a/python/rapidocr_paddle/main.py b/python/rapidocr_paddle/main.py index 269b952..feb63e2 100644 --- a/python/rapidocr_paddle/main.py +++ b/python/rapidocr_paddle/main.py @@ -2,7 +2,6 @@ # @Author: SWHL # @Contact: liekkaskono@163.com import copy -import logging from pathlib import Path from typing import List, Optional, Tuple, Union @@ -83,6 +82,7 @@ class RapidOCR: det_elapse, cls_elapse, rec_elapse = 0.0, 0.0, 0.0 if use_det: + img, padding_h = self.maybe_add_letterbox(img) dt_boxes, det_elapse = self.auto_text_det(img) if dt_boxes is None: return None, None @@ -95,28 +95,36 @@ class RapidOCR: if use_rec: rec_res, rec_elapse = self.text_rec(img) + if dt_boxes is not None and padding_h > 0: + for box in dt_boxes: + box[:, 1] -= padding_h + ocr_res = self.get_final_res( dt_boxes, cls_res, rec_res, det_elapse, cls_elapse, rec_elapse ) return ocr_res - def auto_text_det( - self, - img: np.ndarray, - ) -> Tuple[Optional[np.ndarray], float, Optional[List[np.ndarray]]]: + def maybe_add_letterbox(self, img: np.ndarray) -> Tuple[np.ndarray, int]: h, w = img.shape[:2] + if self.width_height_ratio == -1: use_limit_ratio = False else: use_limit_ratio = w / h > self.width_height_ratio if h <= self.min_height or use_limit_ratio: - logging.warning( - "Because the aspect ratio of the current image exceeds the limit (min_height or width_height_ratio), the program will skip the detection step." + new_h = max(int(w / self.width_height_ratio), self.min_height) * 2 + padding_h = int(abs(new_h - h) / 2) + block_img = cv2.copyMakeBorder( + img, padding_h, padding_h, 0, 0, cv2.BORDER_CONSTANT, value=(0, 0, 0) ) - dt_boxes = self.get_boxes_img_without_det(h, w) - return dt_boxes, 0.0 + return block_img, padding_h + return img, 0 + def auto_text_det( + self, + img: np.ndarray, + ) -> Tuple[Optional[np.ndarray], float, Optional[List[np.ndarray]]]: dt_boxes, det_elapse = self.text_det(img) if dt_boxes is None or len(dt_boxes) < 1: return None, 0.0 @@ -124,12 +132,6 @@ class RapidOCR: dt_boxes = self.sorted_boxes(dt_boxes) return dt_boxes, det_elapse - def get_boxes_img_without_det(self, h, w): - x0, y0, x1, y1 = 0, 0, w, h - dt_boxes = np.float32([[x0, y0], [x1, y0], [x1, y1], [x0, y1]]) - dt_boxes = dt_boxes[np.newaxis, ...] - return dt_boxes - def get_crop_img_list(self, img, dt_boxes): def get_rotate_crop_image(img, points): img_crop_width = int( @@ -241,7 +243,10 @@ def main(): use_cls = not args.no_cls use_rec = not args.no_rec result, elapse_list = ocr_engine( - args.img_path, use_det=use_det, use_cls=use_cls, use_rec=use_rec + args.img_path, + use_det=use_det, + use_cls=use_cls, + use_rec=use_rec, ) print(result) diff --git a/python/tests/test_ort.py b/python/tests/test_ort.py index 247bd9e..3e70514 100644 --- a/python/tests/test_ort.py +++ b/python/tests/test_ort.py @@ -26,7 +26,7 @@ package_name = "rapidocr_onnxruntime" ( "test_letterbox_like.jpg", 2, - "A:：取决于所使用的执行提供者，它可能没有完全支持模型中的所有操作。回落到CPU操作可能会导致性能速度的下降。此外，即使一个操作是由CUDAexecution", + "A：：取决于所使用的执行提供者，它可能没有完全支持模型中的所有操作。回落到CPU操作可能会导致性能速度的下降。此外，即使一个操作是由CUDAeXecution", ), ("test_without_det.jpg", 1, "在中国作家协会第三届儿童文学"), ], diff --git a/python/tests/test_paddle.py b/python/tests/test_paddle.py index af344ce..c370001 100644 --- a/python/tests/test_paddle.py +++ b/python/tests/test_paddle.py @@ -20,11 +20,23 @@ img_path = tests_dir / "ch_en_num.jpg" package_name = "rapidocr_paddle" -def test_without_det(): - img_path = tests_dir / "test_without_det.png" +@pytest.mark.parametrize( + "img_name,gt_len,gt_first_len", + [ + ( + "test_letterbox_like.jpg", + 2, + "A：：取决于所使用的执行提供者，它可能没有完全支持模型中的所有操作。回落到CPU操作可能会导致性能速度的下降。此外，即使一个操作是由CUDAeXecution", + ), + ("test_without_det.jpg", 1, "在中国作家协会第三届儿童文学"), + ], +) +def test_letterbox_like(img_name, gt_len, gt_first_len): + img_path = tests_dir / img_name result, _ = engine(img_path) - assert result[0][1] == "の持、持场所无。" - assert len(result) == 1 + + assert len(result) == gt_len + assert result[0][1] == gt_first_len def test_only_det(): @@ -114,7 +126,7 @@ def test_input_parameters(): def test_input_det_parameters(): with pytest.raises(FileNotFoundError) as exc_info: - engine = RapidOCR(det_model_path="1") + engine = RapidOCR(det_model_path="1.onnx") result, _ = engine(img_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError @@ -122,7 +134,7 @@ def test_input_det_parameters(): def test_input_cls_parameters(): with pytest.raises(FileNotFoundError) as exc_info: - engine = RapidOCR(cls_model_path="1") + engine = RapidOCR(cls_model_path="1.onnx") result, _ = engine(img_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError @@ -130,7 +142,7 @@ def test_input_cls_parameters(): def test_input_rec_parameters(): with pytest.raises(FileNotFoundError) as exc_info: - engine = RapidOCR(rec_model_path="1") + engine = RapidOCR(rec_model_path="1.onnx") result, _ = engine(img_path) raise FileNotFoundError() assert exc_info.type is FileNotFoundError @@ -142,7 +154,7 @@ def test_input_three_ndim_two_channel(): result, _ = engine(image_array) assert len(result) == 1 - assert result[0][1] == " TREND PLOT REPORT" + assert result[0][1] == "TREND PLOT REPORT" def test_input_three_ndim_one_channel(): diff --git a/python/tests/test_vino.py b/python/tests/test_vino.py index 5f5f619..db8e2a0 100644 --- a/python/tests/test_vino.py +++ b/python/tests/test_vino.py @@ -20,11 +20,23 @@ img_path = tests_dir / "ch_en_num.jpg" package_name = "rapidocr_openvino" -def test_without_det(): - img_path = tests_dir / "test_without_det.png" +@pytest.mark.parametrize( + "img_name,gt_len,gt_first_len", + [ + ( + "test_letterbox_like.jpg", + 2, + "A：：取决于所使用的执行提供者，它可能没有完全支持模型中的所有操作。回落到CPU操作可能会导致性能速度的下降。此外，即使一个操作是由CUDAeXecution", + ), + ("test_without_det.jpg", 1, "在中国作家协会第三届儿童文学"), + ], +) +def test_letterbox_like(img_name, gt_len, gt_first_len): + img_path = tests_dir / img_name result, _ = engine(img_path) - assert result[0][1] == "の持、持场所无。" - assert len(result) == 1 + + assert len(result) == gt_len + assert result[0][1] == gt_first_len def test_only_det(): @@ -142,7 +154,7 @@ def test_input_three_ndim_two_channel(): result, _ = engine(image_array) assert len(result) == 1 - assert result[0][1] == " TREND PLOT REPORT" + assert result[0][1] == "TREND PLOT REPORT" def test_input_three_ndim_one_channel():