📖 Rapid Layout

简介

主要是做文档类图像的版面分析。具体来说，就是分析给定的文档类别图像（论文截图、研报等），定位其中类别和位置，如标题、段落、表格和图片等各个部分。

⚠️注意：需要说明的是，由于不同场景下的版面差异较大，现阶段不存在一个模型可以搞定所有场景。如果实际业务需要，以下模型效果不好的话，建议构建自己的训练集微调。

目前支持以下场景的版面分析：

`model_type`	版面类型	模型名称	支持类别
`pp_layout_table`	表格	`layout_table.onnx`	`["table"]`
`pp_layout_publaynet`	英文	`layout_publaynet.onnx`	`["text", "title", "list", "table", "figure"]`
`pp_layout_cdla`	中文	`layout_cdla.onnx`	`['text', 'title', 'figure', 'figure_caption', 'table', 'table_caption', 'header', 'footer', 'reference', 'equation']`
`yolov8n_layout_paper`	论文	`yolov8n_layout_paper.onnx`	`['Text', 'Title', 'Header', 'Footer', 'Figure', 'Table', 'Toc', 'Figure caption', 'Table caption']`
`yolov8n_layout_report`	研报	`yolov8n_layout_report.onnx`	`['Text', 'Title', 'Header', 'Footer', 'Figure', 'Table', 'Toc', 'Figure caption', 'Table caption']`
`yolov8n_layout_publaynet`	英文	`yolov8n_layout_publaynet.onnx`	`["Text", "Title", "List", "Table", "Figure"]`
`yolov8n_layout_general6`	通用	`yolov8n_layout_general6.onnx`	`["Text", "Title", "Figure", "Table", "Caption", "Equation"]`

PP模型来源：PaddleOCR 版面分析

yolov8n系列来源：360LayoutAnalysis

模型下载地址为：link

安装

由于模型较小，预先将中文版面分析模型(layout_cdla.onnx)打包进了whl包内，如果做中文版面分析，可直接安装使用

$ pip install rapid-layout

使用方式

python脚本运行

import cv2
from rapid_layout import RapidLayout, VisLayout

# model_type类型参见上表。指定不同model_type时，会自动下载相应模型到安装目录下的。
layout_engine = RapidLayout(conf_thres=0.5, model_type="pp_layout_cdla")

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

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)

终端运行

$ 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}]
                [--conf_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6}]
                [--iou_thres {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6}]
                [--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}, --model_type {pp_layout_cdla,pp_layout_publaynet,pp_layout_table,yolov8n_layout_paper,yolov8n_layout_report,yolov8n_layout_publaynet,yolov8n_layout_general6}
                        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}
                        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}
                        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 -v -img test_images/layout.png

GPU推理

因为版面分析模型输入图像尺寸固定，故可使用onnxruntime-gpu来提速。
因为rapid_layout库默认依赖是CPU版onnxruntime，如果想要使用GPU推理，需要手动安装onnxruntime-gpu。
详细使用和评测可参见AI Studio

安装

pip install rapid_layout
pip uninstall onnxruntime

# 这里一定要确定onnxruntime-gpu与GPU对应
# 可参见https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#requirements
pip install onnxruntime-gpu

使用

import cv2
from rapid_layout import RapidLayout
from pathlib import Path

# 注意：这里需要使用use_cuda指定参数
layout_engine = RapidLayout(conf_thres=0.5, model_type="pp_layout_cdla", use_cuda=True)

# warm up
layout_engine("images/12027_5.png")

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)

avg_elapse = sum(elapses) / len(elapses)
print(f'avg elapse: {avg_elapse:.4f}')

WilliamQf-AI/RapidLayout

📖 Rapid Layout

简介

安装

使用方式

python脚本运行

终端运行

GPU推理

安装

使用

可视化结果

参考项目