diff --git a/.github/workflows/gen_whl_to_pypi_rapid_orientation.yml b/.github/workflows/gen_whl_to_pypi_rapid_orientation.yml
new file mode 100644
index 0000000..f5f4847
--- /dev/null
+++ b/.github/workflows/gen_whl_to_pypi_rapid_orientation.yml
@@ -0,0 +1,89 @@
+name: Push rapidocr_orientation to pypi
+
+on:
+ push:
+ branches: [ main ]
+ paths:
+ - 'python/rapid_structure/rapid_orientation/**'
+ - 'docs/doc_whl_rapid_orientation.md'
+ - 'python/rapid_structure/setup_orientation.py'
+ - '.github/workflows/gen_whl_to_pypi_rapid_orientation.yml'
+ - 'python/rapid_structure/tests/test_orientation.py'
+
+jobs:
+ UnitTesting:
+ runs-on: windows-latest
+ steps:
+ - name: Install 7Zip PowerShell Module
+ shell: powershell
+ run: Install-Module 7Zip4PowerShell -Force -Verbose
+
+ - name: Pull latest code
+ uses: actions/checkout@v3
+
+ - name: Download rapid orientation model
+ uses: suisei-cn/actions-download-file@v1
+ id: downloadfile3
+ with:
+ url: "https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/rapid_orientation_models.zip"
+ target: ./
+
+ - name: Decompress
+ shell: powershell
+ run: |
+ Expand-7Zip -ArchiveFileName ./rapid_orientation_models.zip -TargetPath .\
+
+ - name: List all file
+ shell: cmd
+ run: |
+ dir .
+
+ - name: Unit testings
+ shell: cmd
+ run: |
+ set git_keep_path=%cur_dir%\python\rapid_structure\rapid_orientation\models\.gitkeep
+ del %git_keep_path%
+ move rapid_orientation_models\* python\rapid_structure\rapid_orientation\models\
+ cd python/rapid_structure
+ pip install -r requirements.txt
+ pip install onnxruntime
+ cd tests
+ pytest test_orientation.py
+
+ GenerateWHL_PushPyPi:
+ needs: UnitTesting
+ runs-on: ubuntu-latest
+
+ steps:
+ - uses: actions/checkout@v3
+
+ - name: Set SSH Environment
+ env:
+ DEPLOY_KEYS: ${{ secrets.GEN_PYTHON_SDK }}
+ run: |
+ mkdir -p ~/.ssh/
+ echo "$DEPLOY_KEYS" > ~/.ssh/id_rsa
+ chmod 600 ~/.ssh/id_rsa
+ chmod 700 ~/.ssh && chmod 600 ~/.ssh/*
+
+ - name: Run setup
+ run: |
+ cd python/rapid_structure
+ python -m pip install --upgrade pip
+ wget https://github.com/RapidAI/RapidOCR/releases/download/v1.1.0/rapid_orientation_models.zip
+ unzip rapid_orientation_models.zip
+ mv rapid_orientation_models/*.onnx rapid_orientation/models/
+ python setup_orientation.py bdist_wheel ${{ github.event.head_commit.message }}
+
+ - name: Publish distribution 📦 to Test PyPI
+ uses: pypa/gh-action-pypi-publish@v1.5.0
+ with:
+ password: ${{ secrets.TEST_PYPI_API_TOKEN }}
+ repository_url: https://test.pypi.org/legacy/
+ packages_dir: python/rapid_structure/dist/
+
+ # - name: Publish distribution 📦 to PyPI
+ # uses: pypa/gh-action-pypi-publish@v1.5.0
+ # with:
+ # password: ${{ secrets.RAPIDOCR_OPENVINO }}
+ # packages_dir: python/rapid_structure/dist/
diff --git a/docs/doc_whl_rapid_orientation.md b/docs/doc_whl_rapid_orientation.md
new file mode 100644
index 0000000..dd751ed
--- /dev/null
+++ b/docs/doc_whl_rapid_orientation.md
@@ -0,0 +1,51 @@
+## rapid-orientation Package
+
+ 
+ 
+ 
+
+
+### 1. Install package by pypi.
+```bash
+$ pip install rapid-orientation
+```
+
+### 2. Run by script.
+- RapidOrientation has the default `model_path` value, you can set the different value of `model_path` to use different models, e.g. `orientation_engine = RapidOrientation(model_path='rapid_orientation.onnx')`
+- See details, for [README_Layout](https://github.com/RapidAI/RapidOCR/blob/f133ff008a1c60edd6e0ed882da83873aa7b113a/python/rapid_structure/docs/README_Layout.md) .
+- 📌 `layout.png` source: [link](https://github.com/RapidAI/RapidOCR/blob/f133ff008a1c60edd6e0ed882da83873aa7b113a/python/rapid_structure/test_images/layout.png)
+
+```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. Run by command line.
+- Usage:
+ ```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
+ ```
+- Example:
+ ```bash
+ $ rapid_orientation -img layout.png
+ ```
+
+### 4. Result.
+```python
+# Return str, four types:：0 | 90 | 180 | 270
+```
\ No newline at end of file
diff --git a/python/rapid_structure/README.md b/python/rapid_structure/README.md
index ccc5ce5..ce22b80 100644
--- a/python/rapid_structure/README.md
+++ b/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/)
\ No newline at end of file
diff --git a/python/rapid_structure/demo.py b/python/rapid_structure/demo.py
index 913b202..d69a5d5 100644
--- a/python/rapid_structure/demo.py
+++ b/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()
diff --git a/python/rapid_structure/docs/README_Orientation.md b/python/rapid_structure/docs/README_Orientation.md
new file mode 100644
index 0000000..8ec70ca
--- /dev/null
+++ b/python/rapid_structure/docs/README_Orientation.md
@@ -0,0 +1,56 @@
+## Rapid Orientation
+
+
+
+
+ 
+
+
+#### 简介和说明
+- 该部分主要是做含文字图像方向分类模型。
+
+ | 模型类型 | 模型名称 | 模型大小 | 支持种类 |
+ |:---:|:---:|:---:|:---:|
+ | 四方向分类 | `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
+ ```
diff --git a/python/rapid_structure/rapid_layout/rapid_layout.py b/python/rapid_structure/rapid_layout/rapid_layout.py
index 41fc71f..9c6319e 100644
--- a/python/rapid_structure/rapid_layout/rapid_layout.py
+++ b/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:
diff --git a/python/rapid_structure/rapid_orientation/__init__.py b/python/rapid_structure/rapid_orientation/__init__.py
new file mode 100644
index 0000000..336271c
--- /dev/null
+++ b/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
diff --git a/python/rapid_structure/rapid_orientation/config.yaml b/python/rapid_structure/rapid_orientation/config.yaml
new file mode 100644
index 0000000..78847b2
--- /dev/null
+++ b/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:
diff --git a/python/rapid_structure/rapid_orientation/models/.gitkeep b/python/rapid_structure/rapid_orientation/models/.gitkeep
new file mode 100644
index 0000000..e69de29
diff --git a/python/rapid_structure/rapid_orientation/rapid_orientation.py b/python/rapid_structure/rapid_orientation/rapid_orientation.py
new file mode 100644
index 0000000..3fd712a
--- /dev/null
+++ b/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()
diff --git a/python/rapid_structure/rapid_orientation/utils.py b/python/rapid_structure/rapid_orientation/utils.py
new file mode 100644
index 0000000..65300c3
--- /dev/null
+++ b/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, :]
diff --git a/python/rapid_structure/requirements.txt b/python/rapid_structure/requirements.txt
index 4d4831f..660bf76 100644
--- a/python/rapid_structure/requirements.txt
+++ b/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
\ No newline at end of file
+pip==22.3.1
+onnxruntime
\ No newline at end of file
diff --git a/python/rapid_structure/setup_orientation.py b/python/rapid_structure/setup_orientation.py
new file mode 100644
index 0000000..18442c1
--- /dev/null
+++ b/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']
+ }
+)
diff --git a/python/rapid_structure/test_images/layout_result.jpg b/python/rapid_structure/test_images/layout_result.jpg
deleted file mode 100644
index 4d34ac1..0000000
Binary files a/python/rapid_structure/test_images/layout_result.jpg and /dev/null differ
diff --git a/python/rapid_structure/tests/test_files/img_rot0_demo.jpg b/python/rapid_structure/tests/test_files/img_rot0_demo.jpg
new file mode 100644
index 0000000..412d419
Binary files /dev/null and b/python/rapid_structure/tests/test_files/img_rot0_demo.jpg differ
diff --git a/python/rapid_structure/tests/test_files/img_rot180_demo.jpg b/python/rapid_structure/tests/test_files/img_rot180_demo.jpg
new file mode 100644
index 0000000..f4725b9
Binary files /dev/null and b/python/rapid_structure/tests/test_files/img_rot180_demo.jpg differ
diff --git a/python/rapid_structure/tests/test_layout.py b/python/rapid_structure/tests/test_layout.py
index 636c51e..757c427 100644
--- a/python/rapid_structure/tests/test_layout.py
+++ b/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))
diff --git a/python/rapid_structure/tests/test_orientation.py b/python/rapid_structure/tests/test_orientation.py
new file mode 100644
index 0000000..f1c1dca
--- /dev/null
+++ b/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
diff --git a/python/rapid_structure/tests/test_table.py b/python/rapid_structure/tests/test_table.py
index d3e9200..a73e203 100644
--- a/python/rapid_structure/tests/test_table.py
+++ b/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))