Optim unit tests

python/tests/base_module.py

@@ -13,6 +13,7 @@ class BaseModule():
         self.package_name = package_name
         self.root_dir = Path(__file__).resolve().parent.parent
         self.package_dir = self.root_dir / self.package_name
+        self.tests_dir = self.root_dir / 'tests'
 
         sys.path.append(str(self.root_dir))
         sys.path.append(str(self.package_dir))

python/tests/test_all_ort.py

@@ -0,0 +1,40 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import sys
+from pathlib import Path
+
+import numpy as np
+import cv2
+import pytest
+
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+
+from rapidocr_onnxruntime import RapidOCR
+
+rapid_ocr = RapidOCR()
+tests_dir = root_dir / 'test_files'
+
+
+def test_normal():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    img = cv2.imread(str(image_path))
+    result, _ = rapid_ocr(img)
+    assert result[0][1] == '正品促销'
+    assert len(result) == 17
+
+
+def test_empty():
+    img = None
+    with pytest.raises(AttributeError) as exc_info:
+        rapid_ocr(img)
+        raise AttributeError
+    assert exc_info.type is AttributeError
+
+
+def test_zeros():
+    img = np.zeros([640, 640, 3])
+    result, _ = rapid_ocr(img)
+    assert result is None

python/tests/test_all_vino.py

@@ -0,0 +1,41 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import sys
+from pathlib import Path
+
+import numpy as np
+import cv2
+import pytest
+
+
+root_dir = Path(__file__).resolve().parent.parent
+sys.path.append(str(root_dir))
+
+from rapidocr_openvino import RapidOCR
+
+rapid_ocr = RapidOCR()
+
+tests_dir = root_dir / 'test_files'
+
+
+def test_normal():
+    image_path = tests_dir / 'ch_en_num.jpg'
+    img = cv2.imread(str(image_path))
+    result, _ = rapid_ocr(img)
+    assert result[0][1] == '正品促销'
+    assert len(result) == 17
+
+
+def test_empty():
+    img = None
+    with pytest.raises(AttributeError) as exc_info:
+        rapid_ocr(img)
+        raise AttributeError
+    assert exc_info.type is AttributeError
+
+
+def test_zeros():
+    img = np.zeros([640, 640, 3])
+    result, _ = rapid_ocr(img)
+    assert result is None

python/tests/test_cls.py

@@ -1,15 +1,20 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-from base_module import BaseModule
 import cv2
+import pytest
+from base_module import BaseModule
 
 
-class TestONNXRuntime():
+@pytest.mark.parametrize(
+    'package_name',
+    ['rapidocr_onnxruntime', 'rapidocr_openvino']
+)
+def test_cls(package_name: str):
     module_name = 'ch_ppocr_v2_cls'
     class_name = 'TextClassifier'
 
-    base = BaseModule(package_name='rapidocr_onnxruntime')
+    base = BaseModule(package_name=package_name)
 
     TextClassifier = base.init_module(module_name, class_name)
 
@@ -19,29 +24,7 @@ class TestONNXRuntime():
 
     text_cls = TextClassifier(config)
 
-    def test_cls(self):
-        img_path = 'test_files/text_cls.jpg'
-        img = cv2.imread(img_path)
-        result = self.text_cls([img])
-        assert result[1][0][0] == '180'
-
-
-class TestOpenVINO():
-    module_name = 'ch_ppocr_v2_cls'
-    class_name = 'TextClassifier'
-
-    base = BaseModule(package_name='rapidocr_openvino')
-
-    TextClassifier = base.init_module(module_name, class_name)
-
-    yaml_path = base.package_dir / module_name / 'config.yaml'
-    config = base.read_yaml(str(yaml_path))
-    config['model_path'] = str(base.package_dir / config['model_path'])
-
-    text_cls = TextClassifier(config)
-
-    def test_cls(self):
-        img_path = 'test_files/text_cls.jpg'
-        img = cv2.imread(img_path)
-        result = self.text_cls([img])
-        assert result[1][0][0] == '180'
+    img_path = base.tests_dir / 'test_files' / 'text_cls.jpg'
+    img = cv2.imread(str(img_path))
+    result = text_cls([img])
+    assert result[1][0][0] == '180'

python/tests/test_det.py

@@ -2,14 +2,19 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 import cv2
+import pytest
 from base_module import BaseModule
 
 
-class TestONNXRuntime():
+@pytest.mark.parametrize(
+    'package_name',
+    ['rapidocr_onnxruntime', 'rapidocr_openvino']
+)
+def test_det(package_name):
     module_name = 'ch_ppocr_v3_det'
     class_name = 'TextDetector'
 
-    base = BaseModule(package_name='rapidocr_onnxruntime')
+    base = BaseModule(package_name)
     TextDetector = base.init_module(module_name, class_name)
 
     yaml_path = base.package_dir / module_name / 'config.yaml'
@@ -17,33 +22,7 @@ class TestONNXRuntime():
     config['model_path'] = str(base.package_dir / config['model_path'])
 
     text_det = TextDetector(config)
-
-    def test_det(self):
-        img_path = 'test_files/text_det.jpg'
-
-        img = cv2.imread(img_path)
-
-        dt_boxes, elapse = self.text_det(img)
-        assert dt_boxes.shape == (18, 4, 2)
-
-
-class TestOpenVINO():
-    module_name = 'ch_ppocr_v3_det'
-    class_name = 'TextDetector'
-
-    base = BaseModule(package_name='rapidocr_openvino')
-    TextDetector = base.init_module(module_name, class_name)
-
-    yaml_path = base.package_dir / module_name / 'config.yaml'
-    config = base.read_yaml(str(yaml_path))
-    config['model_path'] = str(base.package_dir / config['model_path'])
-
-    text_det = TextDetector(config)
-
-    def test_det(self):
-        img_path = 'test_files/text_det.jpg'
-
-        img = cv2.imread(img_path)
-
-        dt_boxes, elapse = self.text_det(img)
-        assert dt_boxes.shape == (18, 4, 2)
+    img_path = base.tests_dir / 'test_files' / 'text_det.jpg'
+    img = cv2.imread(str(img_path))
+    dt_boxes, elapse = text_det(img)
+    assert dt_boxes.shape == (18, 4, 2)

python/tests/test_rec.py

@@ -2,14 +2,19 @@
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
 import cv2
+import pytest
 from base_module import BaseModule
 
 
-class TestONNXRuntime():
+@pytest.mark.parametrize(
+    'package_name',
+    ['rapidocr_onnxruntime', 'rapidocr_openvino']
+)
+def test_det(package_name):
     module_name = 'ch_ppocr_v3_rec'
     class_name = 'TextRecognizer'
 
-    base = BaseModule(package_name='rapidocr_onnxruntime')
+    base = BaseModule(package_name)
     TextRecognizer = base.init_module(module_name, class_name)
 
     yaml_path = base.package_dir / module_name / 'config.yaml'
@@ -18,28 +23,7 @@ class TestONNXRuntime():
 
     text_rec = TextRecognizer(config)
 
-    def test_rec(self):
-        img_path = 'test_files/text_rec.jpg'
-        img = cv2.imread(img_path)
-        rec_res, elapse = self.text_rec([img])
-        assert rec_res[0][0] == '韩国小馆'
-
-
-class TestOpenVINO():
-    module_name = 'ch_ppocr_v3_rec'
-    class_name = 'TextRecognizer'
-
-    base = BaseModule(package_name='rapidocr_openvino')
-    TextRecognizer = base.init_module(module_name, class_name)
-
-    yaml_path = base.package_dir / module_name / 'config.yaml'
-    config = base.read_yaml(str(yaml_path))
-    config['model_path'] = str(base.package_dir / config['model_path'])
-
-    text_rec = TextRecognizer(config)
-
-    def test_rec(self):
-        img_path = 'test_files/text_rec.jpg'
-        img = cv2.imread(img_path)
-        rec_res, elapse = self.text_rec([img])
-        assert rec_res[0][0] == '韩国小馆'
+    img_path = base.tests_dir / 'test_files' / 'text_rec.jpg'
+    img = cv2.imread(str(img_path))
+    rec_res, elapse = text_rec([img])
+    assert rec_res[0][0] == '韩国小馆'