feat: support det mean and std config (#312)

RapidAI · Jan 11, 2025 · f699fdd · f699fdd
1 parent fecad1f
commit f699fdd
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 19 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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