From f699fdd3a9c6320e02015f16105deb494c5253e7 Mon Sep 17 00:00:00 2001
From: Joker1212 <519548295@qq.com>
Date: Sat, 11 Jan 2025 17:11:28 +0800
Subject: [PATCH] feat: support det mean and std config (#312)

---
 .../rapidocr_onnxruntime/ch_ppocr_det/text_detect.py  |  8 +++++---
 python/rapidocr_onnxruntime/ch_ppocr_det/utils.py     | 11 ++++++++---
 python/rapidocr_onnxruntime/config.yaml               |  7 ++++++-
 python/rapidocr_openvino/ch_ppocr_det/text_detect.py  |  8 +++++---
 python/rapidocr_openvino/ch_ppocr_det/utils.py        | 11 ++++++++---
 python/rapidocr_paddle/ch_ppocr_det/text_detect.py    |  8 +++++---
 python/rapidocr_paddle/ch_ppocr_det/utils.py          | 11 ++++++++---
 7 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py b/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py
index 41c60bacc..1e97b9302 100644
--- a/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py
+++ b/python/rapidocr_onnxruntime/ch_ppocr_det/text_detect.py
@@ -26,8 +26,10 @@
 
 class TextDetector:
     def __init__(self, config: Dict[str, Any]):
-        self.limit_type = config.get("limit_type", "min")
-        self.limit_side_len = config.get("limit_side_len", 736)
+        self.limit_side_len = config.get("limit_side_len")
+        self.limit_type = config.get("limit_type")
+        self.mean = config.get("mean")
+        self.std = config.get("std")
         self.preprocess_op = None
 
         post_process = {
@@ -69,7 +71,7 @@ def get_preprocess(self, max_wh):
             limit_side_len = 1500
         else:
             limit_side_len = 2000
-        return DetPreProcess(limit_side_len, self.limit_type)
+        return DetPreProcess(limit_side_len, self.limit_type,self.mean, self.std)
 
     def filter_tag_det_res(
         self, dt_boxes: np.ndarray, image_shape: Tuple[int, int]
diff --git a/python/rapidocr_onnxruntime/ch_ppocr_det/utils.py b/python/rapidocr_onnxruntime/ch_ppocr_det/utils.py
index 9f1b88113..bc5ccc11b 100644
--- a/python/rapidocr_onnxruntime/ch_ppocr_det/utils.py
+++ b/python/rapidocr_onnxruntime/ch_ppocr_det/utils.py
@@ -10,9 +10,14 @@
 
 
 class DetPreProcess:
-    def __init__(self, limit_side_len: int = 736, limit_type: str = "min"):
-        self.mean = np.array([0.5, 0.5, 0.5])
-        self.std = np.array([0.5, 0.5, 0.5])
+    def __init__(self, limit_side_len: int = 736, limit_type: str = "min", mean=None,
+                 std=None):
+        if mean is None:
+            mean = [0.5, 0.5, 0.5]
+        if std is None:
+            std = [0.5, 0.5, 0.5]
+        self.mean = np.array(mean)
+        self.std = np.array(std)
         self.scale = 1 / 255.0
 
         self.limit_side_len = limit_side_len
diff --git a/python/rapidocr_onnxruntime/config.yaml b/python/rapidocr_onnxruntime/config.yaml
index 0ee6fbef3..eaa7a7df5 100644
--- a/python/rapidocr_onnxruntime/config.yaml
+++ b/python/rapidocr_onnxruntime/config.yaml
@@ -24,7 +24,12 @@ Det:
 
     limit_side_len: 736
     limit_type: min
-
+    std: [ 0.5, 0.5, 0.5 ]
+    mean: [ 0.5, 0.5, 0.5 ]
+#    limit_side_len: 960
+#    limit_type: max
+#    std: [ 0.229, 0.224, 0.225 ]
+#    mean: [ 0.485, 0.456, 0.406 ]
     thresh: 0.3
     box_thresh: 0.5
     max_candidates: 1000
diff --git a/python/rapidocr_openvino/ch_ppocr_det/text_detect.py b/python/rapidocr_openvino/ch_ppocr_det/text_detect.py
index 2fdb4e9ac..575db3cb0 100644
--- a/python/rapidocr_openvino/ch_ppocr_det/text_detect.py
+++ b/python/rapidocr_openvino/ch_ppocr_det/text_detect.py
@@ -26,8 +26,10 @@
 
 class TextDetector:
     def __init__(self, config: Dict[str, Any]):
-        self.limit_type = config.get("limit_type", "min")
-        self.limit_side_len = config.get("limit_side_len", 736)
+        self.limit_side_len = config.get("limit_side_len")
+        self.limit_type = config.get("limit_type")
+        self.mean = config.get("mean")
+        self.std = config.get("std")
         self.preprocess_op = None
 
         post_process = {
@@ -69,7 +71,7 @@ def get_preprocess(self, max_wh):
             limit_side_len = 1500
         else:
             limit_side_len = 2000
-        return DetPreProcess(limit_side_len, self.limit_type)
+        return DetPreProcess(limit_side_len, self.limit_type, self.mean, self.std)
 
     def filter_tag_det_res(
         self, dt_boxes: np.ndarray, image_shape: Tuple[int, int]
diff --git a/python/rapidocr_openvino/ch_ppocr_det/utils.py b/python/rapidocr_openvino/ch_ppocr_det/utils.py
index b274abd4a..014145c0f 100644
--- a/python/rapidocr_openvino/ch_ppocr_det/utils.py
+++ b/python/rapidocr_openvino/ch_ppocr_det/utils.py
@@ -10,9 +10,14 @@
 
 
 class DetPreProcess:
-    def __init__(self, limit_side_len: int = 736, limit_type: str = "min"):
-        self.mean = np.array([0.485, 0.456, 0.406])
-        self.std = np.array([0.229, 0.224, 0.225])
+    def __init__(self, limit_side_len: int = 736, limit_type: str = "min",mean=None,
+                 std=None):
+        if mean is None:
+            mean = [0.485, 0.456, 0.406]
+        if std is None:
+            std = [0.229, 0.224, 0.225]
+        self.mean = np.array(mean)
+        self.std = np.array(std)
         self.scale = 1 / 255.0
 
         self.limit_side_len = limit_side_len
diff --git a/python/rapidocr_paddle/ch_ppocr_det/text_detect.py b/python/rapidocr_paddle/ch_ppocr_det/text_detect.py
index 4a3e19849..b48ff2dde 100644
--- a/python/rapidocr_paddle/ch_ppocr_det/text_detect.py
+++ b/python/rapidocr_paddle/ch_ppocr_det/text_detect.py
@@ -26,8 +26,10 @@
 
 class TextDetector:
     def __init__(self, config: Dict[str, Any]):
-        self.limit_type = config.get("limit_type", "min")
-        self.limit_side_len = config.get("limit_side_len", 736)
+        self.limit_side_len = config.get("limit_side_len")
+        self.limit_type = config.get("limit_type")
+        self.mean = config.get("mean")
+        self.std = config.get("std")
         self.preprocess_op = None
 
         post_process = {
@@ -69,7 +71,7 @@ def get_preprocess(self, max_wh):
             limit_side_len = 1500
         else:
             limit_side_len = 2000
-        return DetPreProcess(limit_side_len, self.limit_type)
+        return DetPreProcess(limit_side_len, self.limit_type, self.mean, self.std)
 
     def filter_tag_det_res(
         self, dt_boxes: np.ndarray, image_shape: Tuple[int, int]
diff --git a/python/rapidocr_paddle/ch_ppocr_det/utils.py b/python/rapidocr_paddle/ch_ppocr_det/utils.py
index b274abd4a..9302d4a64 100644
--- a/python/rapidocr_paddle/ch_ppocr_det/utils.py
+++ b/python/rapidocr_paddle/ch_ppocr_det/utils.py
@@ -10,9 +10,14 @@
 
 
 class DetPreProcess:
-    def __init__(self, limit_side_len: int = 736, limit_type: str = "min"):
-        self.mean = np.array([0.485, 0.456, 0.406])
-        self.std = np.array([0.229, 0.224, 0.225])
+    def __init__(self, limit_side_len: int = 736, limit_type: str = "min", mean=None,
+                 std=None):
+        if mean is None:
+            mean = [0.485, 0.456, 0.406]
+        if std is None:
+            std = [0.229, 0.224, 0.225]
+        self.mean = np.array(mean)
+        self.std = np.array(std)
         self.scale = 1 / 255.0
 
         self.limit_side_len = limit_side_len