Merge pull request #25 from RapidAI/develop

feat: refacter code

Author: SWHL

.github/workflows/publish_whl.yml

@@ -26,14 +26,8 @@ jobs:
 
       - name: Unit testings
         run: |
-          wget $RESOURCES_URL
-          ZIP_NAME=${RESOURCES_URL##*/}
-          DIR_NAME=${ZIP_NAME%.*}
-
-          unzip $DIR_NAME
-          cp $DIR_NAME/*.onnx rapid_layout/models/
           pip install -r requirements.txt
-          pip install pytest wheel
+          pip install pytest wheel onnxruntime
           pytest tests/test*.py
 
   GenerateWHL_PushPyPi:

README.md

@@ -4,7 +4,7 @@
   </div>
 
 <a href="https://huggingface.co/spaces/SWHL/RapidLayout" target="_blank"><img src="https://img.shields.io/badge/%F0%9F%A4%97-Hugging Face Demo-blue"></a>
-<a href=""><img src="https://img.shields.io/badge/Python->=3.6,<3.13-aff.svg"></a>
+<a href=""><img src="https://img.shields.io/badge/Python->=3.6-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-layout/"><img alt="PyPI" src="https://img.shields.io/pypi/v/rapid-layout"></a>
 <a href="https://pepy.tech/project/rapid-layout"><img src="https://static.pepy.tech/personalized-badge/rapid-layout?period=total&units=abbreviation&left_color=grey&right_color=blue&left_text=Downloads"></a>
@@ -43,75 +43,45 @@ yolov8n系列来源：[360LayoutAnalysis](https://github.com/360AILAB-NLP/360Lay
 
 ### TODO
 
-- [ ] [ PP-DocLayout](https://github.com/PaddlePaddle/PaddleX/blob/release/3.0-rc/docs/module_usage/tutorials/ocr_modules/layout_detection.md)整理
+- [ ] [PP-DocLayout](https://github.com/PaddlePaddle/PaddleX/blob/release/3.0-rc/docs/module_usage/tutorials/ocr_modules/layout_detection.md)整理
 
 ### 安装
 
 由于模型较小，预先将中文版面分析模型(`layout_cdla.onnx`)打包进了whl包内，如果做中文版面分析，可直接安装使用
 
 ```bash
-pip install rapid-layout
+pip install rapid-layout onnxruntime
 ```
 
 ### 使用方式
 
 #### python脚本运行
 
 ```python
-import cv2
-from imread_from_url import imread_from_url  # pip install imread_from_url
+from rapid_layout import EngineType, ModelType, RapidLayout, RapidLayoutInput
 
-from rapid_layout import RapidLayout, VisLayout
+cfg = RapidLayoutInput()
+layout_engine = RapidLayout(cfg=cfg)
 
-# model_type类型参见上表。指定不同model_type时，会自动下载相应模型到安装目录下的。
-layout_engine = RapidLayout(model_type="doclayout_docstructbench", conf_thres=0.2)
+img_path = "https://raw.githubusercontent.com/RapidAI/RapidLayout/refs/heads/develop/tests/test_files/layout.jpg"
+results = layout_engine(img_path)
+print(results)
 
-img_url = "https://raw.githubusercontent.com/opendatalab/DocLayout-YOLO/refs/heads/main/assets/example/financial.jpg"
-img = imread_from_url(img_url)
-
-boxes, scores, class_names, elapse = layout_engine(img)
-ploted_img = VisLayout.draw_detections(img, boxes, scores, class_names)
-if ploted_img is not None:
-    cv2.imwrite("layout_res.png", ploted_img)
+results.vis("layout_res.png")
 ```
 
 ### 可视化结果
 
 <div align="center">
-    <img src="https://github.com/RapidAI/RapidLayout/releases/download/v0.0.0/layout_res.png" width="80%">
+    <img src="https://github.com/RapidAI/RapidLayout/releases/download/v0.0.0/layout_vis.jpg" width="80%" height="80%">
 </div>
 
 #### 终端运行
 
 ```bash
-$ rapid_layout -h
-usage: rapid_layout [-h] -img IMG_PATH
-                    [-m {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}]
-                    [--conf_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}]
-                    [--iou_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}]
-                    [--use_cuda] [--use_dml] [-v]
-
-options:
-  -h, --help            show this help message and exit
-  -img IMG_PATH, --img_path IMG_PATH
-                        Path to image for layout.
-  -m {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}, --model_type {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}
-                        Support model type
-  --conf_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}
-                        Box threshold, the range is [0, 1]
-  --iou_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6,doclayout_docstructbench,doclayout_d4la,doclayout_docsynth}
-                        IoU threshold, the range is [0, 1]
-  --use_cuda            Whether to use cuda.
-  --use_dml             Whether to use DirectML, which only works in Windows10+.
-  -v, --vis             Wheter to visualize the layout results.
+rapid_layout test_images/layout.png
 ```
 
-- 示例:
-
-    ```bash
-    rapid_layout -v -img test_images/layout.png
-    ```
-
 ### GPU推理
 
 - 因为版面分析模型输入图像尺寸固定，故可使用`onnxruntime-gpu`来提速。
@@ -132,25 +102,41 @@ pip install onnxruntime-gpu
 #### 使用
 
 ```python
-import cv2
-from rapid_layout import RapidLayout
-from pathlib import Path
+from rapid_layout import EngineType, ModelType, RapidLayout, RapidLayoutInput
+
+cfg = RapidLayoutInput(
+    model_type=ModelType.PP_LAYOUT_CDLA,
+    engine_type=EngineType.ONNXRUNTIME,
+    engine_cfg={"use_cuda": True, "cuda_ep_cfg.gpu_id": 1},
+)
+layout_engine = RapidLayout(cfg=cfg)
+
+img_path = "https://raw.githubusercontent.com/RapidAI/RapidLayout/refs/heads/develop/tests/test_files/layout.jpg"
+results = layout_engine(img_path)
+print(results)
 
-# 注意：这里需要使用use_cuda指定参数
-layout_engine = RapidLayout(model_type="doclayout_yolo", conf_thres=0.2, use_cuda=True)
+results.vis("layout_res.png")
+```
+
+### NPU使用
+
+详细配置参数参见：[link](https://github.com/RapidAI/RapidLayout/blob/a7ab63ff291bd72e1a98ac2bb11860575514f432/rapid_layout/configs/engine_cfg.yaml)
+
+```python
+from rapid_layout import EngineType, ModelType, RapidLayout, RapidLayoutInput
 
-# warm up
-layout_engine("images/12027_5.png")
+cfg = RapidLayoutInput(
+    model_type=ModelType.PP_LAYOUT_CDLA,
+    engine_type=EngineType.ONNXRUNTIME,
+    engine_cfg={"use_cann": True, "cann_ep_cfg.gpu_id": 0},
+)
+layout_engine = RapidLayout(cfg=cfg)
 
-elapses = []
-img_list = list(Path('images').iterdir())
-for img_path in img_list:
-    boxes, scores, class_names, elapse = layout_engine(img_path)
-    print(f"{img_path}: {elapse}s")
-    elapses.append(elapse)
+img_path = "https://raw.githubusercontent.com/RapidAI/RapidLayout/refs/heads/develop/tests/test_files/layout.jpg"
+results = layout_engine(img_path)
+print(results)
 
-avg_elapse = sum(elapses) / len(elapses)
-print(f'avg elapse: {avg_elapse:.4f}')
+results.vis("layout_res.png")
 ```
 
 ### 参考项目

demo.py

@@ -1,16 +1,12 @@
 # -*- encoding: utf-8 -*-
 # @Author: SWHL
 # @Contact: [email protected]
-import cv2
+from rapid_layout import EngineType, ModelType, RapidLayout, RapidLayoutInput
 
-from rapid_layout import RapidLayout, VisLayout
+layout_engine = RapidLayout()
 
-layout_engine = RapidLayout(model_type="doclayout_docsynth")
+img_path = "https://raw.githubusercontent.com/RapidAI/RapidLayout/refs/heads/develop/tests/test_files/layout.jpg"
+results = layout_engine(img_path)
+print(results)
 
-img_path = "tests/test_files/PMC3576793_00004.jpg"
-img = cv2.imread(img_path)
-
-boxes, scores, class_names, elapse = layout_engine(img_path)
-ploted_img = VisLayout.draw_detections(img, boxes, scores, class_names)
-if ploted_img is not None:
-    cv2.imwrite("layout_res.png", ploted_img)
+results.vis("layout_res.png")

rapid_layout/__init__.py

@@ -1,15 +1,5 @@
-# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
-#
-# 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: [email protected]
 from .main import RapidLayout
-from .utils import VisLayout
+from .utils.typings import EngineType, ModelType, RapidLayoutInput

rapid_layout/configs/__init__.py

@@ -0,0 +1,3 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: [email protected]

rapid_layout/configs/default_models.yaml

@@ -0,0 +1,39 @@
+pp_layout_cdla:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/pp_layout/layout_cdla.onnx
+  SHA256: 25b1f27ec56aa932a48f30cbd6293c358a156280f4b20b0a973bab210c39f62c
+
+pp_layout_publaynet:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/pp_layout/layout_publaynet.onnx
+  SHA256: 958aa6dcef1cc1a542d0a513b5976a3d5edbcc37d76460ec1e9f126358e4d100
+
+pp_layout_table:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/pp_layout/layout_table.onnx
+  SHA256: 5b07ba6df1d1889bed2877c9d7501235c6fb6e2212aca8f2f56f4b1b8d0e37b5
+
+yolov8n_layout_paper:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/360/yolov8n_layout_paper.onnx
+  SHA256: bc074c8d8fbe89e5d90c3e21b7e3b52f279c70fe210ae96d73b74141df64347c
+
+yolov8n_layout_report:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/360/yolov8n_layout_report.onnx
+  SHA256: 9d5ada6a69b5825eb255da2b82d2c8d11636a0adae801074d88892527b535980
+
+yolov8n_layout_publaynet:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/360/yolov8n_layout_publaynet.onnx
+  SHA256: 5304bf18e538312a1bd211eb2ad3283524dff956e5cbffcefb3ad294c6e3cba6
+
+yolov8n_layout_general6:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/360/yolov8n_layout_general6.onnx
+  SHA256: 927b6edcb268e896e6a170f7d78980591b408e04b3908f54d58eb69efd018c95
+
+doclayout_docstructbench:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/doclayout/doclayout_yolo_docstructbench_imgsz1024.onnx
+  SHA256: 3b452baef10ecabd615491bc82cc4d49475fbc2cd7a8e535044f2c6bb28fb9fe
+
+doclayout_d4la:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/doclayout/doclayout_yolo_d4la_imgsz1600_docsynth_pretrain.onnx
+  SHA256: 1c81715d45d5bee2e6b644f92563a9eaa5cb4cad3d4293f890f99c0862937e69
+
+doclayout_docsynth:
+  model_dir_or_path: https://www.modelscope.cn/models/RapidAI/RapidLayout/resolve/v1.0.0/onnx/doclayout/doclayout_yolo_doclaynet_imgsz1120_docsynth_pretrain.onnx
+  SHA256: 527e60cefc2801dec727dddbfa3a2bd225876a8c5505461c9d3b1193f50a7c84

rapid_layout/configs/engine_cfg.yaml

@@ -0,0 +1,27 @@
+onnxruntime:
+    intra_op_num_threads: -1
+    inter_op_num_threads: -1
+    enable_cpu_mem_arena: false
+
+    cpu_ep_cfg:
+        arena_extend_strategy: "kSameAsRequested"
+
+    use_cuda: false
+    cuda_ep_cfg:
+        gpu_id: 0
+        arena_extend_strategy: "kNextPowerOfTwo"
+        cudnn_conv_algo_search: "EXHAUSTIVE"
+        do_copy_in_default_stream: true
+
+    use_dml: false
+    dm_ep_cfg: null
+
+    use_cann: false
+    cann_ep_cfg:
+        gpu_id: 0
+        arena_extend_strategy: "kNextPowerOfTwo"
+        npu_mem_limit:  21474836480 # 20 * 1024 * 1024 * 1024
+        op_select_impl_mode: "high_performance"
+        optypelist_for_implmode: "Gelu"
+        enable_cann_graph: true
+

rapid_layout/inference_engine/__init__.py

@@ -0,0 +1,3 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: [email protected]

rapid_layout/inference_engine/base.py

@@ -0,0 +1,78 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: [email protected]
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Dict, List, Union
+
+import numpy as np
+from omegaconf import DictConfig, OmegaConf
+
+from ..utils.logger import Logger
+from ..utils.typings import EngineType
+from ..utils.utils import import_package
+
+logger = Logger(logger_name=__name__).get_log()
+
+
+class InferSession(ABC):
+    cur_dir = Path(__file__).resolve().parent.parent
+    MODEL_URL_PATH = cur_dir / "configs" / "default_models.yaml"
+    ENGINE_CFG_PATH = cur_dir / "configs" / "engine_cfg.yaml"
+
+    model_info = OmegaConf.load(MODEL_URL_PATH)
+    DEFAULT_MODEL_PATH = cur_dir / "models"
+
+    engine_cfg = OmegaConf.load(ENGINE_CFG_PATH)
+
+    @abstractmethod
+    def __init__(self, config):
+        pass
+
+    @abstractmethod
+    def __call__(self, input_content: np.ndarray) -> np.ndarray:
+        pass
+
+    @staticmethod
+    def _verify_model(model_path: Union[str, Path, None]):
+        if model_path is None:
+            raise ValueError("model_path is None!")
+
+        model_path = Path(model_path)
+        if not model_path.exists():
+            raise FileNotFoundError(f"{model_path} does not exists.")
+
+        if not model_path.is_file():
+            raise FileExistsError(f"{model_path} is not a file.")
+
+    @abstractmethod
+    def have_key(self, key: str = "character") -> bool:
+        pass
+
+    @property
+    def characters(self):
+        return self.get_character_list()
+
+    @abstractmethod
+    def get_character_list(self, key: str = "character") -> List[str]:
+        pass
+
+    @staticmethod
+    def update_params(cfg: DictConfig, params: Dict[str, Any]) -> DictConfig:
+        for k, v in params.items():
+            OmegaConf.update(cfg, k, v)
+        return cfg
+
+
+def get_engine(engine_type: EngineType):
+    logger.info("Using engine_name: %s", engine_type.value)
+
+    if engine_type == EngineType.ONNXRUNTIME:
+        if not import_package(engine_type.value):
+            raise ImportError(f"{engine_type.value} is not installed.")
+
+        from .onnxruntime import OrtInferSession
+
+        return OrtInferSession
+
+    raise ValueError(f"Unsupported engine: {engine_type.value}")

rapid_layout/inference_engine/onnxruntime/__init__.py

@@ -0,0 +1,4 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: [email protected]
+from .main import OrtInferSession