Add 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)