Add pytest

python/README.md

@@ -89,6 +89,11 @@
    - Note: 在Windows端，Shapely库可能自动安装会有问题，解决方案参见[Q15](../docs/FAQ.md#q15-装完环境之后运行python-mainpy之后报错oserror-winerror-126-找不到指定的模組)
 
 4. 运行示例
+    - 运行单元测试
+        ```bash
+        cd tests
+        pytest test_*.py
+        ```
     - 接口调用
         ```python
         import cv2
@@ -129,7 +134,7 @@
 |Global|  `text_score`  |    [0, 1]    |   0.5   |       文本识别结果置信度，值越大，把握越大       |
 |Global|  `use_angle_cls`  |  `bool`      |   `true`   |       是否使用文本行的方向分类       |
 |Global|  `print_verbose`  |    `bool`    |   `true`   |       是否打印各个部分耗时信息       |
-|Global|  `min_heigh`  |    `int`    |   50   |       图像高度，单位是像素       |
+|Global|  `min_heigh`  |    `int`    |   30   |       图像最小高度（单位是像素），低于这个值，会跳过文本检测阶段，直接进行后续识别       |
 ||||||
 |Det|  `box_thresh`  |    [0, 1]    |   0.5   | 文本检测所得框是否保留的阈值，值越大，召回率越低 |
 |Det| `unclip_ratio` |  [1.6, 2.0]  |   1.6   |   控制文本检测框的大小，值越大，检测框整体越大   |

python/rapidocr_onnxruntime/__init__.py

@@ -1,4 +1,4 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-from .rapid_ocr_api import TextSystem
+from .rapid_ocr_api import TextSystem

python/rapidocr_onnxruntime/ch_ppocr_v2_cls/text_cls.py

@@ -19,7 +19,6 @@ from typing import List
 
 import cv2
 import numpy as np
-import onnxruntime as ort
 
 try:
     from .utils import ClsPostProcess, read_yaml, OrtInferSession
@@ -38,29 +37,6 @@ class TextClassifier(object):
         self.session = session_instance.session
         self.input_name = session_instance.get_input_name()
 
-    def resize_norm_img(self, img):
-        img_c, img_h, img_w = self.cls_image_shape
-        h, w = img.shape[:2]
-        ratio = w / float(h)
-        if math.ceil(img_h * ratio) > img_w:
-            resized_w = img_w
-        else:
-            resized_w = int(math.ceil(img_h * ratio))
-
-        resized_image = cv2.resize(img, (resized_w, img_h))
-        resized_image = resized_image.astype('float32')
-        if img_c == 1:
-            resized_image = resized_image / 255
-            resized_image = resized_image[np.newaxis, :]
-        else:
-            resized_image = resized_image.transpose((2, 0, 1)) / 255
-
-        resized_image -= 0.5
-        resized_image /= 0.5
-        padding_im = np.zeros((img_c, img_h, img_w), dtype=np.float32)
-        padding_im[:, :, :resized_w] = resized_image
-        return padding_im
-
     def __call__(self, img_list: List[np.ndarray]):
         if isinstance(img_list, np.ndarray):
             img_list = [img_list]
@@ -106,6 +82,29 @@ class TextClassifier(object):
                         img_list[indices[beg_img_no + rno]], 1)
         return img_list, cls_res, elapse
 
+    def resize_norm_img(self, img):
+        img_c, img_h, img_w = self.cls_image_shape
+        h, w = img.shape[:2]
+        ratio = w / float(h)
+        if math.ceil(img_h * ratio) > img_w:
+            resized_w = img_w
+        else:
+            resized_w = int(math.ceil(img_h * ratio))
+
+        resized_image = cv2.resize(img, (resized_w, img_h))
+        resized_image = resized_image.astype('float32')
+        if img_c == 1:
+            resized_image = resized_image / 255
+            resized_image = resized_image[np.newaxis, :]
+        else:
+            resized_image = resized_image.transpose((2, 0, 1)) / 255
+
+        resized_image -= 0.5
+        resized_image /= 0.5
+        padding_im = np.zeros((img_c, img_h, img_w), dtype=np.float32)
+        padding_im[:, :, :resized_w] = resized_image
+        return padding_im
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

python/rapidocr_onnxruntime/ch_ppocr_v2_det/config.yaml

@@ -1,11 +1,9 @@
-model_path: G:\ProgramFiles\_self\RapidOCR\python\resources\models\ch_PP-OCRv2_det_infer.onnx
+model_path: resources/models/ch_PP-OCRv2_det_infer.onnx
 
 use_cuda: false
-# Details of the params: https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html
 CUDAExecutionProvider:
     device_id: 0
     arena_extend_strategy: kNextPowerOfTwo
-    gpu_mem_limit: 2 * 1024 * 1024 * 1024
     cudnn_conv_algo_search: EXHAUSTIVE
     do_copy_in_default_stream: true
 

python/rapidocr_onnxruntime/ch_ppocr_v2_det/text_detect.py

@@ -20,7 +20,6 @@ import time
 import cv2
 import numpy as np
 
-
 try:
     from .utils import (DBPostProcess, create_operators,
                         draw_text_det_res, transform, read_yaml, OrtInferSession)

python/rapidocr_openvino/ch_ppocr_v2_cls/text_cls.py

@@ -40,29 +40,6 @@ class TextClassifier(object):
         self.session = ort.InferenceSession(config['model_path'], sess_opt)
         self.input_name = self.session.get_inputs()[0].name
 
-    def resize_norm_img(self, img):
-        img_c, img_h, img_w = self.cls_image_shape
-        h, w = img.shape[:2]
-        ratio = w / float(h)
-        if math.ceil(img_h * ratio) > img_w:
-            resized_w = img_w
-        else:
-            resized_w = int(math.ceil(img_h * ratio))
-
-        resized_image = cv2.resize(img, (resized_w, img_h))
-        resized_image = resized_image.astype('float32')
-        if img_c == 1:
-            resized_image = resized_image / 255
-            resized_image = resized_image[np.newaxis, :]
-        else:
-            resized_image = resized_image.transpose((2, 0, 1)) / 255
-
-        resized_image -= 0.5
-        resized_image /= 0.5
-        padding_im = np.zeros((img_c, img_h, img_w), dtype=np.float32)
-        padding_im[:, :, :resized_w] = resized_image
-        return padding_im
-
     def __call__(self, img_list: List[np.ndarray]):
         if isinstance(img_list, np.ndarray):
             img_list = [img_list]
@@ -108,6 +85,29 @@ class TextClassifier(object):
                         img_list[indices[beg_img_no + rno]], 1)
         return img_list, cls_res, elapse
 
+    def resize_norm_img(self, img):
+        img_c, img_h, img_w = self.cls_image_shape
+        h, w = img.shape[:2]
+        ratio = w / float(h)
+        if math.ceil(img_h * ratio) > img_w:
+            resized_w = img_w
+        else:
+            resized_w = int(math.ceil(img_h * ratio))
+
+        resized_image = cv2.resize(img, (resized_w, img_h))
+        resized_image = resized_image.astype('float32')
+        if img_c == 1:
+            resized_image = resized_image / 255
+            resized_image = resized_image[np.newaxis, :]
+        else:
+            resized_image = resized_image.transpose((2, 0, 1)) / 255
+
+        resized_image -= 0.5
+        resized_image /= 0.5
+        padding_im = np.zeros((img_c, img_h, img_w), dtype=np.float32)
+        padding_im[:, :, :resized_w] = resized_image
+        return padding_im
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

python/requirements.txt

@@ -4,4 +4,5 @@ opencv_python==4.5.1.48
 six>=1.15.0
 numpy>=1.19.2
 Pillow==8.3.0
-PyYAML
+PyYAML
+pytest

python/tests/base.py

@@ -0,0 +1,26 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+
+import importlib
+from pathlib import Path
+import sys
+
+import yaml
+
+package_name = 'rapidocr_onnxruntime'
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+sys.path.append(str(root_dir / package_name))
+
+
+def init_module(module_name, class_name):
+    module_part = importlib.import_module(module_name)
+    return getattr(module_part, class_name)
+
+
+def read_yaml(yaml_path):
+    with open(yaml_path, 'rb') as f:
+        data = yaml.load(f, Loader=yaml.Loader)
+    return data

python/tests/test_cls.py

@@ -0,0 +1,23 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import base
+import cv2
+
+
+module_name = 'ch_ppocr_v2_cls'
+class_name = 'TextClassifier'
+TextClassifier = base.init_module(module_name, class_name)
+
+yaml_path = base.root_dir / base.package_name / module_name / 'config.yaml'
+config = base.read_yaml(str(yaml_path))
+config['model_path'] = str(base.root_dir / config['model_path'])
+
+text_cls = TextClassifier(config)
+
+
+def test_cls():
+    img_path = 'test_files/text_cls.jpg'
+    img = cv2.imread(img_path)
+    result = text_cls([img])
+    assert result[1][0][0] == '180'

python/tests/test_det.py

@@ -0,0 +1,25 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import base
+import cv2
+
+
+module_name = 'ch_ppocr_v2_det'
+class_name = 'TextDetector'
+TextDetector = base.init_module(module_name, class_name)
+
+yaml_path = base.root_dir / base.package_name / module_name / 'config.yaml'
+config = base.read_yaml(str(yaml_path))
+config['model_path'] = str(base.root_dir / config['model_path'])
+
+text_det = TextDetector(config)
+
+
+def test_det():
+    img_path = 'test_files/ch_en_num.jpg'
+
+    img = cv2.imread(img_path)
+
+    dt_boxes, elapse = text_det(img)
+    assert dt_boxes.shape == (18, 4, 2)