diff --git a/python/config.yaml b/python/config.yaml index 9eb61d1..992f3ed 100644 --- a/python/config.yaml +++ b/python/config.yaml @@ -37,6 +37,7 @@ Det: max_candidates: 1000 unclip_ratio: 1.6 use_dilation: true + score_mode: "fast" Cls: module_name: ch_ppocr_v2_cls diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml index 21bc186..1952eeb 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml @@ -26,3 +26,4 @@ post_process: max_candidates: 1000 unclip_ratio: 1.6 use_dilation: true + score_mode: "fast" diff --git a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py index e5afe48..18f3a56 100644 --- a/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py +++ b/python/rapidocr_onnxruntime/ch_ppocr_v2_det/utils.py @@ -299,12 +299,14 @@ class DBPostProcess(object): box_thresh=0.7, max_candidates=1000, unclip_ratio=2.0, + score_mode="fast", use_dilation=False): self.thresh = thresh self.box_thresh = box_thresh self.max_candidates = max_candidates self.unclip_ratio = unclip_ratio self.min_size = 3 + self.score_mode = score_mode if use_dilation: self.dilation_kernel = np.array([[1, 1], [1, 1]]) @@ -337,7 +339,10 @@ class DBPostProcess(object): if sside < self.min_size: continue points = np.array(points) - score = self.box_score_fast(pred, points.reshape(-1, 2)) + if self.score_mode == "fast": + score = self.box_score_fast(pred, points.reshape(-1, 2)) + else: + score = self.box_score_slow(pred, contour) if self.box_thresh > score: continue @@ -400,6 +405,27 @@ class DBPostProcess(object): box[:, 1] = box[:, 1] - ymin cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1) return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] + + def box_score_slow(self, bitmap, contour): + ''' + box_score_slow: use polyon mean score as the mean score + ''' + h, w = bitmap.shape[:2] + contour = contour.copy() + contour = np.reshape(contour, (-1, 2)) + + xmin = np.clip(np.min(contour[:, 0]), 0, w - 1) + xmax = np.clip(np.max(contour[:, 0]), 0, w - 1) + ymin = np.clip(np.min(contour[:, 1]), 0, h - 1) + ymax = np.clip(np.max(contour[:, 1]), 0, h - 1) + + mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8) + + contour[:, 0] = contour[:, 0] - xmin + contour[:, 1] = contour[:, 1] - ymin + + cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype(np.int32), 1) + return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] def __call__(self, pred, shape_list): pred = pred[:, 0, :, :]