Add text orientation cls module.

python/rapid_structure/README.md

@@ -1,7 +1,9 @@
 ## Rapid Structure
 - 模型和推理代码均来自于[PP-Structure](https://github.com/PaddlePaddle/PaddleOCR/blob/release%2F2.6/ppstructure/README_ch.md#32)部分。
-- 该部分的功能主要针对文档类图像，包括版面分析和表格识别两类。
+- 该部分的功能主要针对文档类图像，包括文档图像分类、版面分析和表格识别。
 
-### [Rapid Layout](./docs/README_Layout.md)
+### [版面分析](./docs/README_Layout.md)
 
-### [Rapid Table](./docs/README_Table.md)
+### [表格识别](./docs/README_Table.md)
+
+### [文档方向分类](./docs/)

python/rapid_structure/demo.py

@@ -6,7 +6,9 @@ from pathlib import Path
 
 import cv2
 import numpy as np
+
 from rapid_layout import RapidLayout
+from rapid_orientation import RapidOrientation
 from rapid_table import RapidTable
 
 
@@ -51,8 +53,7 @@ def vis_table(table_res):
 
 
 def demo_layout():
-    model_path = 'rapid_layout/models/layout_cdla.onnx'
-    layout_engine = RapidLayout(model_path)
+    layout_engine = RapidLayout()
 
     img = cv2.imread('test_images/layout.png')
 
@@ -63,14 +64,22 @@ def demo_layout():
 
 
 def demo_table():
-    rapid_table = RapidTable()
+    table_engine = RapidTable()
     img = cv2.imread('test_images/table.jpg')
-    table_html_str, _ = rapid_table(img)
+    table_html_str, _ = table_engine(img)
 
     vis_table(table_html_str)
     print(table_html_str)
 
 
+def demo_orientation():
+    orientation_engine = RapidOrientation()
+    img = cv2.imread('test_images/table.jpg')
+    cls_result, _ = orientation_engine(img)
+    print(cls_result)
+
+
 if __name__ == '__main__':
     # demo_layout()
-    demo_table()
+    # demo_table()
+    demo_orientation()

python/rapid_structure/docs/README_Orientation.md

@@ -0,0 +1,56 @@
+## Rapid Orientation
+<p align="left">
+    <a href=""><img src="https://img.shields.io/badge/Python->=3.7,<=3.10-aff.svg"></a>
+    <a href=""><img src="https://img.shields.io/badge/OS-Linux%2C%20Win%2C%20Mac-pink.svg"></a>
+    <a href="https://pypi.org/project/rapid-orientation/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-orientation"></a>
+    <a href="https://pypi.org/project/rapid-orientation/"><img src="https://img.shields.io/pypi/dm/rapid-orientation?color=9cf"></a>
+</p>
+
+#### 简介和说明
+- 该部分主要是做含文字图像方向分类模型。
+
+    | 模型类型  |        模型名称         | 模型大小 |                           支持种类                           |
+    |:---:|:---:|:---:|:---:|
+    |   四方向分类   |   `rapid_orientation.onnx`   |  6.5M | `0 90 180 270`|
+
+#### 使用方式
+1. pip安装
+   - 由于模型较小，已经将分类模型(`rapid_orientation.onnx`)打包进了whl包内：
+        ```bash
+        $ pip install rapid-orientation
+        ```
+2. python脚本运行
+   ```python
+    import cv2
+    from rapid_orientation import RapidOrientation
+
+    orientation_engine = RapidOrientation()
+
+    img = cv2.imread('test_images/layout.png')
+
+    orientation_res, elapse = orientation_engine(img)
+    print(orientation_res)
+   ```
+
+3. 终端运行
+   - 用法:
+     ```bash
+     $ rapid_orientation -h
+     usage: rapid_orientation [-h] -img IMG_PATH [-m MODEL_PATH]
+
+     optional arguments:
+     -h, --help            show this help message and exit
+     -img IMG_PATH, --img_path IMG_PATH
+                           Path to image for layout.
+     -m MODEL_PATH, --model_path MODEL_PATH
+                           The model path used for inference.
+     ```
+   - 示例:
+     ```bash
+     $ rapid_orientation -img test_images/layout.png
+     ```
+
+4. 结果
+    ```python
+    # 返回结果为str类型，有四类：0 | 90 | 180 | 270
+    ```

python/rapid_structure/rapid_layout/rapid_layout.py

@@ -83,10 +83,10 @@ def main():
                         help='The model path used for inference.')
     args = parser.parse_args()
 
-    rapid_layout = RapidLayout(args.model_path)
+    layout_engine = RapidLayout(args.model_path)
 
     img = cv2.imread(args.img_path)
-    layout_res, elapse = rapid_layout(img)
+    layout_res, elapse = layout_engine(img)
     print(layout_res)
 
     if args.vis:

python/rapid_structure/rapid_orientation/__init__.py

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

python/rapid_structure/rapid_orientation/config.yaml

@@ -0,0 +1,16 @@
+model_path: models/rapid_orientation.onnx
+
+use_cuda: false
+CUDAExecutionProvider:
+    device_id: 0
+    arena_extend_strategy: kNextPowerOfTwo
+    cudnn_conv_algo_search: EXHAUSTIVE
+    do_copy_in_default_stream: true
+
+PreProcess:
+    - ResizeImage:
+        resize_short: 256
+    - CropImage:
+        size: 224
+    - NormalizeImage:
+    - ToCHWImage:

python/rapid_structure/rapid_orientation/rapid_orientation.py

@@ -0,0 +1,84 @@
+# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import argparse
+import time
+from pathlib import Path
+import cv2
+
+import numpy as np
+import yaml
+
+from .utils import OrtInferSession, create_operators
+
+root_dir = Path(__file__).resolve().parent
+
+
+class RapidOrientation():
+    def __init__(self, model_path: str = None):
+        config_path = str(root_dir / 'config.yaml')
+        config = self.read_yaml(config_path)
+        if model_path is None:
+            model_path = str(root_dir / 'models' / 'rapid_orientation.onnx')
+        config['model_path'] = model_path
+
+        session_instance = OrtInferSession(config)
+        self.predictor = session_instance.session
+        self.input_name = session_instance.get_input_name()
+        self.labels = session_instance.get_metadata()['character'].splitlines()
+
+        self.preprocess_ops = create_operators(config["PreProcess"])
+
+    def __call__(self, images):
+        s = time.time()
+        for ops in self.preprocess_ops:
+            images = ops(images)
+        image = np.array(images)
+
+        input_dict = {self.input_name: image}
+        pred_output = self.predictor.run(None, input_feed=input_dict)[0]
+        pred_output = pred_output.squeeze()
+        pred_idx = np.argmax(pred_output)
+        pred_txt = self.labels[pred_idx]
+        elapse = time.time() - s
+        return pred_txt, elapse
+
+    @staticmethod
+    def read_yaml(yaml_path):
+        with open(yaml_path, 'rb') as f:
+            data = yaml.load(f, Loader=yaml.Loader)
+        return data
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-img', '--img_path', type=str, required=True,
+                        help='Path to image for layout.')
+    parser.add_argument('-m', '--model_path', type=str,
+                        default=str(root_dir / 'models' /
+                                    'rapid_orientation.onnx'),
+                        help='The model path used for inference.')
+    args = parser.parse_args()
+
+    orientation_engine = RapidOrientation(args.model_path)
+
+    img = cv2.imread(args.img_path)
+    orientaion_result, _ = orientation_engine(img)
+    print(orientaion_result)
+
+
+if __name__ == '__main__':
+    main()

python/rapid_structure/rapid_orientation/utils.py

@@ -0,0 +1,148 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import importlib
+import warnings
+
+import cv2
+import numpy as np
+from onnxruntime import (GraphOptimizationLevel, InferenceSession,
+                         SessionOptions, get_available_providers, get_device)
+
+
+class OrtInferSession():
+    def __init__(self, config):
+        sess_opt = SessionOptions()
+        sess_opt.log_severity_level = 4
+        sess_opt.enable_cpu_mem_arena = False
+        sess_opt.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_ALL
+
+        cuda_ep = 'CUDAExecutionProvider'
+        cpu_ep = 'CPUExecutionProvider'
+        cpu_provider_options = {"arena_extend_strategy": "kSameAsRequested", }
+
+        EP_list = []
+        if config['use_cuda'] and get_device() == 'GPU' \
+                and cuda_ep in get_available_providers():
+            EP_list = [(cuda_ep, config[cuda_ep])]
+        EP_list.append((cpu_ep, cpu_provider_options))
+
+        self.session = InferenceSession(config['model_path'],
+                                        sess_options=sess_opt,
+                                        providers=EP_list)
+
+        has_cuda_ep = cuda_ep not in self.session.get_providers()
+        if config['use_cuda'] and has_cuda_ep:
+            warnings.warn(f'{cuda_ep} is not avaiable for current env,'
+                          f'the inference part is automatically shifted to '
+                          f'be executed under {cpu_ep}. '
+                          f'Please ensure the installed onnxruntime-gpu '
+                          f' version matches your cuda and cudnn version, '
+                          f'you can check their relations from the offical web site: '
+                          f'https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html',
+                          RuntimeWarning)
+
+    def get_input_name(self, input_idx=0):
+        return self.session.get_inputs()[input_idx].name
+
+    def get_output_name(self, output_idx=0):
+        return self.session.get_outputs()[output_idx].name
+
+    def get_metadata(self):
+        meta_dict = self.session.get_modelmeta().custom_metadata_map
+        return meta_dict
+
+
+def create_operators(params):
+    """
+    create operators based on the config
+
+    Args:
+        params(list): a dict list, used to create some operators
+    """
+    assert isinstance(params, list), ('operator config should be a list')
+    mod = importlib.import_module(__name__)
+    ops = []
+    for operator in params:
+        assert isinstance(operator,
+                          dict) and len(operator) == 1, "yaml format error"
+        op_name = list(operator)[0]
+        param = {} if operator[op_name] is None else operator[op_name]
+        op = getattr(mod, op_name)(**param)
+        ops.append(op)
+    return ops
+
+
+class ResizeImage():
+    def __init__(self, size=None, resize_short=None):
+        if resize_short is not None and resize_short > 0:
+            self.resize_short = resize_short
+            self.w, self.h = None, None
+        elif size is not None:
+            self.resize_short = None
+            self.w = size if isinstance(size, int) else size[0]
+            self.h = size if isinstance(size, int) else size[1]
+        else:
+            raise ValueError("invalid params for ReisizeImage for '\
+                'both 'size' and 'resize_short' are None")
+
+    def __call__(self, img: np.ndarray):
+        img_h, img_w = img.shape[:2]
+
+        if self.resize_short:
+            percent = float(self.resize_short) / min(img_w, img_h)
+            w = int(round(img_w * percent))
+            h = int(round(img_h * percent))
+        else:
+            w = self.w
+            h = self.h
+        return cv2.resize(img, (w, h))
+
+
+class NormalizeImage():
+    def __init__(self,):
+        self.scale = np.float32(1.0 / 255.0)
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+
+        shape = 1, 1, 3
+        self.mean = np.array(mean).reshape(shape).astype('float32')
+        self.std = np.array(std).reshape(shape).astype('float32')
+
+    def __call__(self, img):
+        img = np.array(img).astype(np.float32)
+        img = (img * self.scale - self.mean) / self.std
+        return img.astype(np.float32)
+
+
+class ToCHWImage():
+    def __init__(self):
+        pass
+
+    def __call__(self, img):
+        img = np.array(img)
+        return img.transpose((2, 0, 1))
+
+
+class CropImage():
+    def __init__(self, size):
+        self.size = size
+        if isinstance(size, int):
+            self.size = (size, size)
+
+    def __call__(self, img):
+        w, h = self.size
+        img_h, img_w = img.shape[:2]
+
+        if img_h < h or img_w < w:
+            raise ValueError(
+                f"The size({h}, {w}) of CropImage must be greater than "
+                f"size({img_h}, {img_w}) of image."
+            )
+
+        w_start = (img_w - w) // 2
+        h_start = (img_h - h) // 2
+
+        w_end = w_start + w
+        h_end = h_start + h
+        return img[h_start:h_end, w_start:w_end, :]

python/rapid_structure/requirements.txt

@@ -3,4 +3,5 @@ opencv_python>=4.6.0.66
 setuptools>=61.2.0
 PyYAML
 pytest
-pip==22.3.1
+pip==22.3.1
+onnxruntime

python/rapid_structure/setup_orientation.py

@@ -0,0 +1,87 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+import setuptools
+
+
+def get_latest_version(package_name):
+    output = subprocess.run(["pip", "index", "versions", package_name],
+                            capture_output=True)
+    output = output.stdout.decode('utf-8')
+    if output:
+        return extract_version(output)
+    return None
+
+
+def version_add_one(version, add_loc=-1):
+    if version:
+        version_list = version.split('.')
+        mini_version = str(int(version_list[add_loc]) + 1)
+        version_list[add_loc] = mini_version
+        new_version = '.'.join(version_list)
+        return new_version
+    return '0.0.1'
+
+
+def get_readme():
+    root_dir = Path(__file__).resolve().parent.parent.parent
+    readme_path = str(root_dir / 'docs' / 'doc_whl_rapid_layout.md')
+    with open(readme_path, 'r', encoding='utf-8') as f:
+        readme = f.read()
+    return readme
+
+
+def extract_version(message: str) -> str:
+    pattern = r'\d+\.(?:\d+\.)*\d+'
+    matched_versions = re.findall(pattern, message)
+    if matched_versions:
+        return matched_versions[0]
+    return ''
+
+
+MODULE_NAME = 'rapid_orientation'
+latest_version = get_latest_version(MODULE_NAME)
+VERSION_NUM = version_add_one(latest_version)
+
+# 优先提取commit message中的语义化版本号，如无，则自动加1
+if len(sys.argv) > 2:
+    match_str = ' '.join(sys.argv[2:])
+    matched_versions = extract_version(match_str)
+    if matched_versions:
+        VERSION_NUM = matched_versions
+sys.argv = sys.argv[:2]
+
+setuptools.setup(
+    name=MODULE_NAME,
+    version=VERSION_NUM,
+    platforms="Any",
+    long_description=get_readme(),
+    long_description_content_type='text/markdown',
+    description='Tools for classifying the direction of images containing text based ONNXRuntime.',
+    author="SWHL",
+    author_email="liekkaskono@163.com",
+    url="https://github.com/RapidAI/RapidOCR",
+    license='Apache-2.0',
+    include_package_data=True,
+    install_requires=["onnxruntime>=1.7.0", "PyYAML>=6.0",
+                      "opencv_python>=4.5.1.48", "numpy>=1.21.6"],
+    packages=[MODULE_NAME, f'{MODULE_NAME}.models'],
+    package_data={'': ['*.onnx', '*.yaml']},
+    keywords=[
+        'ppstructure,layout,rapidocr,rapid_orientation'
+    ],
+    classifiers=[
+        'Programming Language :: Python :: 3.7',
+        'Programming Language :: Python :: 3.8',
+        'Programming Language :: Python :: 3.9',
+        'Programming Language :: Python :: 3.10',
+    ],
+    entry_points={
+        'console_scripts': [f'{MODULE_NAME}={MODULE_NAME}.{MODULE_NAME}:main']
+    }
+)

python/rapid_structure/tests/test_layout.py

@@ -6,7 +6,6 @@ from pathlib import Path
 
 import cv2
 
-# rapid_structure路径
 root_dir = Path(__file__).resolve().parent.parent
 
 sys.path.append(str(root_dir))

python/rapid_structure/tests/test_orientation.py

@@ -0,0 +1,28 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import sys
+from pathlib import Path
+
+import cv2
+import pytest
+
+root_dir = Path(__file__).resolve().parent
+sys.path.append(str(root_dir.parent))
+
+from rapid_orientation import RapidOrientation
+text_orientation = RapidOrientation()
+
+
+@pytest.mark.parametrize(
+    'img_path, result',
+    [
+        (f'{root_dir}/test_files/img_rot0_demo.jpg', '0'),
+        (f'{root_dir}/test_files/img_rot180_demo.jpg', '180'),
+    ]
+)
+def test_img(img_path, result):
+    img = cv2.imread(str(img_path))
+    pred_result, _ = text_orientation(img)
+
+    assert pred_result == result

python/rapid_structure/tests/test_table.py

@@ -6,7 +6,6 @@ from pathlib import Path
 
 import cv2
 
-# rapid_structure路径
 root_dir = Path(__file__).resolve().parent.parent
 
 sys.path.append(str(root_dir))