diff --git a/python/rapidocr/ch_ppocr_det/main.py b/python/rapidocr/ch_ppocr_det/main.py index 8c8f950..eb2988b 100644 --- a/python/rapidocr/ch_ppocr_det/main.py +++ b/python/rapidocr/ch_ppocr_det/main.py @@ -53,7 +53,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])) + self.preprocess_op = self.get_preprocess() prepro_img = self.preprocess_op(img) if prepro_img is None: return TextDetOutput() @@ -67,16 +67,8 @@ class TextDetector: elapse = time.perf_counter() - start_time return TextDetOutput(img, boxes, scores, elapse=elapse) - def get_preprocess(self, max_wh: int) -> DetPreProcess: - 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, self.mean, self.std) + def get_preprocess(self) -> DetPreProcess: + return DetPreProcess(self.limit_side_len, self.limit_type, self.mean, self.std) @staticmethod def sorted_boxes(dt_boxes: np.ndarray) -> np.ndarray: diff --git a/python/tests/test_det_preprocess.py b/python/tests/test_det_preprocess.py new file mode 100644 index 0000000..0e8d184 --- /dev/null +++ b/python/tests/test_det_preprocess.py @@ -0,0 +1,43 @@ +# -*- encoding: utf-8 -*- +# @Author: Mohammad Hijjawi +import numpy as np +import pytest +from rapidocr.ch_ppocr_det import TextDetector + + +def get_detector(limit_side_len: int, limit_type: str) -> TextDetector: + detector = TextDetector.__new__(TextDetector) + detector.limit_side_len = limit_side_len + detector.limit_type = limit_type + detector.mean = [0.5, 0.5, 0.5] + detector.std = [0.5, 0.5, 0.5] + detector.postprocess_op = lambda preds, ori_shape: ([], []) + return detector + + +@pytest.mark.parametrize( + "limit_side_len,limit_type,img_hw,expected_hw", + [ + # issue #740: limit_side_len must be honoured when limit_type is "max" + (960, "max", (2000, 2000), (960, 960)), + (960, "max", (1000, 2000), (480, 960)), + (640, "max", (1200, 800), (640, 416)), + (960, "max", (320, 640), (320, 640)), + (736, "min", (500, 1000), (736, 1472)), + (736, "min", (1000, 800), (992, 800)), + (736, "min", (800, 1600), (800, 1600)), + ], +) +def test_det_preprocess_respects_limit_side_len( + limit_side_len, limit_type, img_hw, expected_hw +): + detector = get_detector(limit_side_len, limit_type) + img = np.zeros((*img_hw, 3), dtype=np.uint8) + + det_inputs = [] + detector.session = det_inputs.append + + detector(img) + + assert len(det_inputs) == 1 + assert det_inputs[0].shape == (1, 3, *expected_hw)