通往RAG之路（二）：版面结构检测方法介绍

news2024/11/10 11:02:15

一、基于yolov5的版面结构检测

AG系统搭建过程中，版面分析是不可缺少的一个步骤，本文介绍用yolov5进行版面结构信息识别，后续再搭配表格识别、公式识别、文字识别等模块进行版面还原，完成PDF结构化输出。

添加图片注释，不超过 140 字（可选）

1.1、环境搭建

conda create -n yolov5-pdf python=3.10
conda activate yolov5-pdf
git clone https://github.com/ultralytics/yolov5.git
cd yolov5
pip   install  -r requirements.txt

等待安装完毕。

1.2、预训练权重下载

下载yolov5预训练权重放在根目录下。
添加图片注释，不超过 140 字（可选）

1.3、准备数据集

本次训练采用CDLA数据集来进行模型训练。
数据来源：https://github.com/buptlihang/CDLA
数据简介：CDLA数据集的训练集包含5000张图像，验证集包含1000张图像。该数据集包括10个类别，分别是：Text, Title, Figure, Figure caption, Table, Table caption, Header, Footer, Reference, Equation。
备注：基于PaddleDetection套件，在该数据集上训练目标检测模型时，在转换label时，需要将label.txt中的__ignore__与_background_去除。
原版数据百度云下载：https://pan.baidu.com/s/1449mhds2ze5JLk-88yKVAA, 提取码: tp0d
Google Drive Download：https://drive.google.com/file/d/14SUsp_TG8OPdK0VthRXBcAbYzIBjSNLm/view?usp=sharing
原版数据下载以后需要对格式进行转换，才能用yolov5进行训练，json转换代码如下：

import json 
import os 
import argparse
from tqdm import tqdm
import glob
import cv2 
import numpy as np
 
def convert_label_json(json_dir,save_dir,classes):
    files=os.listdir(json_dir)
    #删选出json文件
    jsonFiles=[]
    for file in files:
        if os.path.splitext(file)[1]==".json":
            jsonFiles.append(file)
    #获取类型        
    classes=classes.split(',')
 
    #获取json对应中对应元素
    for json_path in tqdm(jsonFiles):
        path=os.path.join(json_dir,json_path)
        with open(path,'r') as loadFile:
            print(loadFile)
            json_dict=json.load(loadFile)
        h,w=json_dict['imageHeight'],json_dict['imageWidth']
        txt_path=os.path.join(save_dir,json_path.replace('json','txt'))
        txt_file=open(txt_path,'w')
 
        for shape_dict in json_dict['shapes']:
            label=shape_dict['label'] 
            label_index=classes.index(label)
            points=shape_dict['points']
            points_nor_list=[]
            for point in points:
                points_nor_list.append(point[0]/w)
                points_nor_list.append(point[1]/h)
            points_nor_list=list(map(lambda x:str(x),points_nor_list))
            points_nor_str=' '.join(points_nor_list)
            label_str=str(label_index)+' '+points_nor_str+'\n'
            txt_file.writelines(label_str)
 
 
if __name__=="__main__":
    parser=argparse.ArgumentParser(description="json convert to txt params")
    #设json文件所在地址
    parser.add_argument('-json',type=str,default='/CDLA_DATASET/train',help='json path')
    #设置txt文件保存地址
    parser.add_argument('-save',type=str,default='/train/labels',help='save path')
    #设置label类型,用“,”分隔
    parser.add_argument('-classes',type=str,default='Text,Title,Figure,Figure caption,Table,Table caption,Header,Footer,Reference,Equation',help='classes')
    args=parser.parse_args()
    print(args.json,args.save,args.classes)
    convert_label_json(args.json,args.save,args.classes)

这里有我转换为yolov5的数据集可以直接下载使用
yolo数据下载链接: https://pan.baidu.com/s/1Du2_TifmUlwMkug8F3HUYA?pwd=pbyd 提取码: pbyd

1.4、训练参数修改

a、data里面.ymal文件修改

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: /yolodata # dataset root dir
train: train/images # train images (relative to 'path') 118287 images
val: val/images # val images (relative to 'path') 5000 images
test: val/images # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794


# Classes
names:
   0: Text
   1: Title
   2: Figure
   3: Figure caption
   4: Table
   5: Table caption
   6: Header
   7: Footer
   8: Reference
   9: Equation

b、yolov5s.yaml文件修改（预训练模型选择s就改s，选哪个就改那个）

只需要将nc：80改成nc:10;

c、train.py文件修改
添加图片注释，不超过 140 字（可选）

需要根据自己电脑配置情况修改–weights、–cfg、–data、–epochs、–batch-size。

1.5、开始训练

python train.py

1.6、训练完毕

添加图片注释，不超过 140 字（可选）

模型权重在weights文件夹内部，一般使用best.pt。

1.7、模型效果测试

使用detect.py进行测试。
需要修改–weights、–data。
在这里插入图片描述

结果如下：

添加图片注释，不超过 140 字（可选）

二、360LayoutAnalysis

在文档版式分析中，精细化的标注非常有必要，其中：段落的标注尤其关键，因为它直接影响到文本的语义理解和信息提取。当前，在版式分析领域，在论文场景中，以往的开源数据集如：CDLA（A Chinese document layout analysis），缺乏对段落信息的标注；在研报场景中的版式分析模型还相对空缺。
基于yolov5自己训练的版面结构检测模型，由于数据集采用了数据集CDLA，本身对段落没法检测。
段落检测还是很关键的，在后续RAG处理过程中会对长文档进行切块，有段落信息可以保留完整语义。
需要检测段落信息则需要自己重新标注数据集进行训练，或者用开源的数据集或者模型。

2.1、360LayoutAnalysis开源模型介绍

github链接：https://github.com/360AILAB-NLP/360LayoutAnalysis
主要特点：
1)涵盖中文论文、英文论文、中文研报三个垂直领域及1个通用场景模型；
2)轻量化推理快速【基于yolov8训练，单模型6.23MB】；
3)中文论文场景包含段落信息【CLDA不具备段落信息，360LayoutAnalysis开源独有】；
4)中文研报场景/通用场景【基于数万级别高质量数据训练，360LayoutAnalysis开源独有】

2.2、模型下载使用

模型huggingface下载链接：https://huggingface.co/qihoo360/360LayoutAnalysis
模型百度网盘链接：https://pan.baidu.com/s/1hp3mYWGbAXbvmiDKPlDUrg?pwd=ugy7 提取码: ugy7
环境配置

conda create -n yolov8 python=3.10
conda activate yolov8
conda install conda-forge::ultralytics

#torch需要单独安装

https://pytorch.org/get-started/locally/