From ee70d8a6aff3082beded189e0f5df3c6997fcec8 Mon Sep 17 00:00:00 2001
From: SWHL <liekkaskono@163.com>
Date: Thu, 17 Nov 2022 19:51:12 +0800
Subject: [PATCH] Update unit testing

---
 python/tests/base.py        | 27 -------------------
 python/tests/base_module.py | 30 +++++++++++++++++++++
 python/tests/test_cls.py    | 50 +++++++++++++++++++++++++----------
 python/tests/test_det.py    | 52 ++++++++++++++++++++++++++-----------
 python/tests/test_rec.py    | 48 ++++++++++++++++++++++++----------
 5 files changed, 139 insertions(+), 68 deletions(-)
 delete mode 100644 python/tests/base.py
 create mode 100644 python/tests/base_module.py

diff --git a/python/tests/base.py b/python/tests/base.py
deleted file mode 100644
index 274c72a..0000000
--- a/python/tests/base.py
+++ /dev/null
@@ -1,27 +0,0 @@
-# -*- encoding: utf-8 -*-
-# @Author: SWHL
-# @Contact: liekkaskono@163.com
-
-import importlib
-from pathlib import Path
-import sys
-
-import yaml
-
-package_name = 'rapidocr_onnxruntime'
-# package_name = 'rapidocr_openvino'
-
-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
diff --git a/python/tests/base_module.py b/python/tests/base_module.py
new file mode 100644
index 0000000..eccf51b
--- /dev/null
+++ b/python/tests/base_module.py
@@ -0,0 +1,30 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import importlib
+import sys
+from pathlib import Path
+
+import yaml
+
+
+class BaseModule(object):
+    def __init__(self, package_name: str = 'rapidocr_onnxruntime'):
+        self.package_name = package_name
+        self.root_dir = Path(__file__).resolve().parent.parent
+
+        sys.path.append(str(self.root_dir))
+        sys.path.append(str(self.root_dir / self.package_name))
+
+    def init_module(self, module_name: str, class_name: str = None):
+        if class_name is None:
+            module_part = importlib.import_module(f'{self.package_name}')
+            return module_part
+        module_part = importlib.import_module(f'{self.package_name}.{module_name}')
+        return getattr(module_part, class_name)
+
+    @staticmethod
+    def read_yaml(yaml_path: str):
+        with open(yaml_path, 'rb') as f:
+            data = yaml.load(f, Loader=yaml.Loader)
+        return data
diff --git a/python/tests/test_cls.py b/python/tests/test_cls.py
index d8c6dc5..84fb82a 100644
--- a/python/tests/test_cls.py
+++ b/python/tests/test_cls.py
@@ -1,23 +1,47 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-import base
+from base_module import BaseModule
 import cv2
 
 
-module_name = 'ch_ppocr_v2_cls'
-class_name = 'TextClassifier'
-TextClassifier = base.init_module(module_name, class_name)
+class TestONNXRuntime():
+    module_name = 'ch_ppocr_v2_cls'
+    class_name = 'TextClassifier'
 
-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'])
+    base = BaseModule(package_name='rapidocr_onnxruntime')
 
-text_cls = TextClassifier(config)
+    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(self):
+        img_path = 'test_files/text_cls.jpg'
+        img = cv2.imread(img_path)
+        result = self.text_cls([img])
+        assert result[1][0][0] == '180'
 
 
-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'
+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.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(self):
+        img_path = 'test_files/text_cls.jpg'
+        img = cv2.imread(img_path)
+        result = self.text_cls([img])
+        assert result[1][0][0] == '180'
diff --git a/python/tests/test_det.py b/python/tests/test_det.py
index e13a332..5033474 100644
--- a/python/tests/test_det.py
+++ b/python/tests/test_det.py
@@ -1,27 +1,49 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-import base
 import cv2
+from base_module import BaseModule
 
 
-module_name = 'ch_ppocr_v3_det'
-class_name = 'TextDetector'
-TextDetector = base.init_module(module_name, class_name)
+class TestONNXRuntime():
+    module_name = 'ch_ppocr_v3_det'
+    class_name = 'TextDetector'
 
-yaml_path = base.root_dir / base.package_name / module_name / 'config.yaml'
-print(yaml_path)
-config = base.read_yaml(str(yaml_path))
-config['model_path'] = str(base.root_dir / config['model_path'])
+    base = BaseModule(package_name='rapidocr_onnxruntime')
+    TextDetector = base.init_module(module_name, class_name)
 
-print(config)
-text_det = TextDetector(config)
+    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(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)
 
 
-def test_det():
-    img_path = 'test_files/text_det.jpg'
+class TestOpenVINO():
+    module_name = 'ch_ppocr_v3_det'
+    class_name = 'TextDetector'
 
-    img = cv2.imread(img_path)
+    base = BaseModule(package_name='rapidocr_openvino')
+    TextDetector = base.init_module(module_name, class_name)
 
-    dt_boxes, elapse = text_det(img)
-    assert dt_boxes.shape == (18, 4, 2)
+    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(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)
diff --git a/python/tests/test_rec.py b/python/tests/test_rec.py
index 790162d..71a10bc 100644
--- a/python/tests/test_rec.py
+++ b/python/tests/test_rec.py
@@ -1,23 +1,45 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: liekkaskono@163.com
-import base
 import cv2
+from base_module import BaseModule
 
 
-module_name = 'ch_ppocr_v3_rec'
-class_name = 'TextRecognizer'
-TextRecognizer = base.init_module(module_name, class_name)
+class TestONNXRuntime():
+    module_name = 'ch_ppocr_v3_rec'
+    class_name = 'TextRecognizer'
 
-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'])
+    base = BaseModule(package_name='rapidocr_onnxruntime')
+    TextRecognizer = base.init_module(module_name, class_name)
 
-text_rec = TextRecognizer(config)
+    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_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] == '韩国小馆'
 
 
-def test_rec():
-    img_path = 'test_files/text_rec.jpg'
-    img = cv2.imread(img_path)
-    rec_res, elapse = 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.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_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] == '韩国小馆'