Skip to content

Commit 66ac1ca

Browse files
committed
docs(README): Add instructions for using GPU inference
1 parent 0d88463 commit 66ac1ca

File tree

2 files changed

+45
-5
lines changed

2 files changed

+45
-5
lines changed

README.md

Lines changed: 44 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ layout_engine = RapidLayout(conf_thres=0.5, model_type="pp_layout_cdla")
5151

5252
img = cv2.imread('test_images/layout.png')
5353

54-
boxes, scores, class_names, *elapse = layout_engine(img)
54+
boxes, scores, class_names, elapse = layout_engine(img)
5555
ploted_img = VisLayout.draw_detections(img, boxes, scores, class_names)
5656
if ploted_img is not None:
5757
cv2.imwrite("layout_res.png", ploted_img)
@@ -61,10 +61,9 @@ if ploted_img is not None:
6161
- 用法:
6262
```bash
6363
$ rapid_layout -h
64-
usage: rapid_layout [-h] -img IMG_PATH
65-
[-m {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report}]
64+
usage: rapid_layout [-h] -img IMG_PATH [-m {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report}]
6665
[--conf_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report}]
67-
[--iou_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report}]
66+
[--iou_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report}] [--use_cuda] [--use_dml]
6867
[-v]
6968

7069
options:
@@ -77,13 +76,54 @@ if ploted_img is not None:
7776
Box threshold, the range is [0, 1]
7877
--iou_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report}
7978
IoU threshold, the range is [0, 1]
79+
--use_cuda Whether to use cuda.
80+
--use_dml Whether to use DirectML, which only works in Windows10+.
8081
-v, --vis Wheter to visualize the layout results.
8182
```
8283
- 示例:
8384
```bash
8485
$ rapid_layout -v -img test_images/layout.png
8586
```
8687

88+
89+
### GPU推理
90+
- 因为版面分析模型输入图像尺寸固定,故可使用`onnxruntime-gpu`来提速。
91+
- 因为`rapid_layout`库默认依赖是CPU版`onnxruntime`,如果想要使用GPU推理,需要手动安装`onnxruntime-gpu`
92+
- 详细使用和评测可参见[AI Studio](https://aistudio.baidu.com/projectdetail/8094594)
93+
94+
#### 安装
95+
```bash
96+
pip install rapid_layout
97+
pip uninstall onnxruntime
98+
99+
# 这里一定要确定onnxruntime-gpu与GPU对应
100+
# 可参见https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements
101+
pip install onnxruntime-gpu
102+
```
103+
104+
#### 使用
105+
```python
106+
import cv2
107+
from rapid_layout import RapidLayout
108+
from pathlib import Path
109+
110+
# 注意:这里需要使用use_cuda指定参数
111+
layout_engine = RapidLayout(conf_thres=0.5, model_type="pp_layout_cdla", use_cuda=True)
112+
113+
# warm up
114+
layout_engine("images/12027_5.png")
115+
116+
elapses = []
117+
img_list = list(Path('images').iterdir())
118+
for img_path in img_list:
119+
boxes, scores, class_names, elapse = layout_engine(img_path)
120+
print(f"{img_path}: {elapse}s")
121+
elapses.append(elapse)
122+
123+
avg_elapse = sum(elapses) / len(elapses)
124+
print(f'avg elapse: {avg_elapse:.4f}')
125+
```
126+
87127
### 可视化结果
88128
89129
<div align="center">

demo.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
img_path = "tests/test_files/layout.png"
1111
img = cv2.imread(img_path)
1212

13-
boxes, scores, class_names, *elapse = layout_engine(img)
13+
boxes, scores, class_names, elapse = layout_engine(img)
1414
ploted_img = VisLayout.draw_detections(img, boxes, scores, class_names)
1515
if ploted_img is not None:
1616
cv2.imwrite("layout_res.png", ploted_img)

0 commit comments

Comments
 (0)