基于yolov8的版面分析AI能力生产全流程

zwxu_

已于 2024-10-09 16:26:55 修改

阅读量86

点赞数

分类专栏：大模型专栏文章标签：人工智能

于 2024-10-09 16:25:36 首次发布

本文链接：https://blog.csdn.net/xzw_910918/article/details/142789874

版权

大模型专栏专栏收录该内容

4 篇文章 0 订阅

订阅专栏

1.coco数据集

1.1 基本定义

COCO数据集是指Common Objects in Context数据集，是一个用于目标检测、图像分割和图像标注任务的知名数据集。COCO数据集由微软研究院创建，旨在提供更广泛的物体类别和更丰富的场景上下文，以促进计算机视觉领域的研究。

1.2应用场景

object instances（目标实例）

object keypoints（目标上的关键点）

image captions（看图说话）

1.3 数据结构

基本数据结构包含info、image、license，使用JSON文件存储。以Object Instance为例，json示例如下:

{

"info": info, # dict

"licenses": [license], # list,内部是dict

"images": [image], # list,内部是dict

"annotations": [annotation],# list,内部是dict

"categories": [category] # list,内部是dict

}

info{ # 数据集信息描述

"year": int, # 数据集年份

"version": str, # 数据集版本

"description": str, # 数据集描述

"contributor": str, # 数据集提供者

"url": str, # 数据集下载链接

"date_created": datetime, # 数据集创建日期

}

license{

"id": int,

"name": str,

"url": str,

}

image{ # images是一个list,存放所有图片(dict)信息。image是一个dict,存放单张图片信息

"id": int, # 图片的ID编号（每张图片ID唯一）

"width": int, # 图片宽

"height": int, # 图片高

"file_name": str, # 图片名字

"license": int, # 协议

"flickr_url": str, # flickr链接地址

"coco_url": str, # 网络连接地址

"date_captured": datetime, # 数据集获取日期

}

annotation{ # annotations是一个list,存放所有标注(dict)信息。annotation是一个dict,存放单个目标标注信息。

"id": int, # 目标对象ID（每个对象ID唯一），每张图片可能有多个目标

"image_id": int, # 对应图片ID

"category_id": int, # 对应类别ID，与categories中的ID对应

"segmentation": RLE or [polygon], # 实例分割，对象的边界点坐标[x1,y1,x2,y2,....,xn,yn]

"area": float, # 对象区域面积

"bbox": [xmin,ymin,width,height], # 目标检测，对象定位边框[x,y,w,h]

"iscrowd": 0 or 1, # 表示是否是人群

}

categories{ # 类别描述

"id": int, # 类别对应的ID（0默认为背景）

"name": str, # 子类别名字

"supercategory": str, # 主类别名字

}

COCO格式数据集的目录结构如下:

train2017是训练集数据

val2017是验证集数据

annotations是基于train2017、val2017原始数据的标注文件

注:标注文件是json格式，且文件名称以instances_开头

目录结构示例如下:

COCO_ROOT #根目录

├── annotations # 存放json格式的标注

│ ├── instances_train2017.json

│ └── instances_val2017.json

└── train2017 # 存放图片文件

│ ├── 000000000001.jpg

│ ├── 000000000002.jpg

│ └── 000000000003.jpg

└── val2017

├── 000000000004.jpg

└── 000000000005.jpg

2.labelme标注工具

2.1 基本定义

Labelme 是一个图形界面的图像标注软件。其的设计灵感来自于 http://labelme.csail.mit.edu/ 。它是用 Python 语言编写的，图形界面使用的是 Qt（PyQt）。

2.2 应用场景

对图像进行多边形，矩形，圆形，多段线，线段，点形式的标注（可用于目标检测，图像分割，等任务）。

对图像进行进行 flag 形式的标注（可用于图像分类和清理任务）。

视频标注

生成 VOC 格式的数据集（for semantic / instance segmentation）

生成 COCO 格式的数据集（for instance segmentation）

2.3 安装步骤

pip install labelme -i https://pypi.tuna.tsinghua.edu.cn/simple

3.模型训练

3.1数据标注

label.txt

__ignore__

_background_

Text

Title

Figure

Figure caption

Table

Table caption

Header

Footer

Reference

Equation

对图片区域进行打标签

# 格式: labelme 图片路径 --labels 标签路径

labelme image_file --labels labels.txt

生成image.json

{

"version": "5.2.1",

"flags": {},

"shapes": [

{

"label": "Title",

"points": [

[

237.36842105263167,

269.28947368421063

[

2205.789473684211,

564.0263157894738

]

"group_id": null,

"description": "",

"shape_type": "rectangle",

"flags": {}

{

"label": "Header",

"points": [

[

884.7368421052631,

58.76315789473695

[

1468.9473684210527,

216.65789473684222

]

"group_id": null,

"description": "",

"shape_type": "rectangle",

"flags": {}

{

"label": "Text",

"points": [

[

385.12195121951225,

594.9634146341465

[

1646.0975609756097,

936.4268292682927

]

"group_id": null,

"description": "",

"shape_type": "rectangle",

"flags": {}

}

"imagePath": "title2.jpg",

"imageData": "",

"imageHeight": 3507,

"imageWidth": 2480

}

3.2 环境准备

镜像制作

FROM 172.31.108.13/library/cuda:11.6_cudnn8_py3

RUN pip install torch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 opencv-python==4.10.0.84 ultralytics==8.2.66 -i https://pypi.tuna.tsinghua.edu.cn/simple

docker run --shm-size 128G --rm -it -v 执行目录:/data yolo_gpu:8 /bin/bash

3.3 数据预处理

将标注文件转coco格式

# train

python3 labelme2coco.py CDLA_dir/train train_save_path --labels labels.txt

# val

python3 labelme2coco.py CDLA_dir/val val_save_path --labels labels.txt

coco标注文件的json格式转txt

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='val',help='json path')

#设置txt文件保存地址

parser.add_argument('-save',type=str,default='val',help='save path')

#设置label类型,用“,”分隔

parser.add_argument('-classes',type=str,default='Header,Text,Reference,Figure caption,Figure,Table caption,Table,Title,Footer,Equation',help='classes')

args=parser.parse_args()

print(args.json,args.save,args.classes)

convert_label_json(args.json,args.save,args.classes)

python coco2txt.py

# train

python3 coco2txt.py -json train -save train

# val

python3 coco2txt.py -json val -save val

3.4 模型训练

from ultralytics import YOLO

from ultralytics.utils import DEFAULT_CFG

import os

import sys

sys.path.insert(0, os.path.dirname(os.getcwd()))

os.environ['CUDA_VISIBLE_DEVICES'] = '1'

DEFAULT_CFG.save_dir= "/data/model"

def train():

# 加载模型

print('model load。。。')

model = YOLO("8npt/8nbest.pt") # 加载模型

print('model load completed。。。')

#epochs=300, batch=16,save_dir=

#训练模型 model.train(data="/data/img-layout.yaml", name="layout", epochs=10, device=1,batch=16 ,save=True)

metrics = model.val() # 在验证集上评估模型性能

if __name__ == '__main__':

train()

训练产物介绍

名称	作用
args.yaml	训练参数
results.csv	训练过程效果指标
weights	权重文件 last.pt：最后一次训练迭代结束后的模型权重 best.pt：训练过程中在验证集上表现最好的模型权重