目录
1.coco数据集
1.1 基本定义
1.2应用场景
1.3 数据结构
2.labelme标注工具
2.1 基本定义
2.2 应用场景
2.3 安装步骤
3. 模型训练
3.1 数据标注
3.2 环境准备
3.3 数据预处理
3.4 模型训练
3.5 模型推理
4.参考链接
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：训练过程中在验证集上表现最好的模型权重 |
3.5 模型推理
| #import sys #import os import cv2 import json #sys.path.insert(0, os.path.dirname(os.getcwd())) #os.environ['CUDA_VISIBLE_DEVICES'] = '1' from ultralytics import YOLO def infer(): model = YOLO('8mpt/8mbest.pt') results = model('title2.jpg') names = results[0].names boxes = results[0].boxes for box in boxes: cls = int(box.cls[0]) x1, y1, x2, y2 = box.xyxy[0] x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) w, h = x2 - x1, y2 - y1 box = x1,y1,w,h region = {'type': names[cls],'bbox':box} print(json.dumps(region)) cv2.imwrite('result.png', results[0].plot()) if __name__ == '__main__': infer() |
4.参考链接
Labelme工具: GitHub - wkentaro/labelme: Image Polygonal Annotation with Python (polygon, rectangle, circle, line, point and image-level flag annotation).
Cdla数据集: GitHub - buptlihang/CDLA: CDLA: A Chinese document layout analysis (CDLA) dataset
Yolo8训练: GitHub - jiangnanboy/layout_analysis: 中文版面检测（Chinese layout detection），yolov8 is used to detect the layout of Chinese document images。
Yolo8训练参数: https://docs.ultralytics.com/modes/train(ultralytics/cfg/default.yaml)
Yolo8权重: Releases · ultralytics/assets · GitHub
