yolov6训练自己的数据集
发布时间
阅读量：
阅读量
YOLOv6 是美团视觉智能部研发的一款目标检测框架，专注于工业应用中的精度和推理效率。它支持多种平台部署（如 GPU、CPU、ARM 等），并提供了详细的代码示例和部署指南。此外，文本还指导了如何准备数据集（包括 XML 标签生成）以及如何修改配置文件以适应特定任务，并展示了训练后的结果展示和模型保存方法。
YOLOv6 是美团视觉智能部开发的一款目标检测框架，在工业应用领域具有显著影响力。该框架不仅注重检测精度的提升，同时也兼顾了推理效率的优化，在工业界广泛使用的统一模型尺寸中：YOLOv6-nano 模型在 COCO 数据集上实现了高达 35.0% 的平均精度指标，并能够在 T4 系列处理器上实现实时推理速度达到 1242 帧/秒；而 YOLOv6-s 模型同样在 COCO 数据集上展现出卓越的检测能力，在 T4 处理器上的推理速度更是达到了 520 帧/秒。从部署角度来看，YOLOv6 框架可支持多种不同的硬件平台配置：包括基于 GPU 的 TensorRT 高性能加速、基于 CPU 的 OpenVINO 开源推理引擎以及多种 ARM 处理器下的 MNN/TNN/NCNN 神经网络后端支持。这种多样化的设计极大简化了工程人员在实际部署过程中的适应性工作流程。
本专题主要涉及操作方面的内容，在以下部分中将详细阐述如何利用该yolov4算法进行数据集的训练与检测过程。
源码地址：
https://github.com/meituan/YOLOv6
代码结构
其中，
① assets 测试图片
配置管理模块包含相关的配置接口函数
③ data 数据相关配置文件
该模型转换接口能够将PyTorch模型部署为ONNX格式和OpenVINO格式的模型，适用于不同的计算平台。
⑤ docs 说明文档
⑥ runs 中间过程已经最终模型保存路径
⑦ tools 主要包含训练，测试，评估接口函数
⑧ yolov6 网络相关代码
1、准备数据集
需要准备VOC格式的数据集（以VOC格式数据集为例）
数据集中，Annotations为标注文件（xml格式），JPEGImages为图片文件
2、数据集预处理，
数据集划分为训练集以及验证集
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import random
from shutil import copyfile
# 根据自己的数据标签修改
classes=['mask','nomask']
def clear_hidden_files(path):
dir_list = os.listdir(path)
for i in dir_list:
abspath = os.path.join(os.path.abspath(path), i)
if os.path.isfile(abspath):
if i.startswith("._"):
os.remove(abspath)
else:
clear_hidden_files(abspath)
def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1])/2.0
y = (box[2] + box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def convert_annotation(image_id):
in_file = open('VOCdevkit/VOC2007/Annotations/%s.xml' %image_id)
out_file = open('VOCdevkit/VOC2007/YOLOLabels/%s.txt' %image_id, 'w')
tree=ET.parse(in_file)
root = tree.getroot()
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
for obj in root.iter('object'):
#difficult = obj.find('difficult').text
cls = obj.find('name').text
#if cls not in classes or int(difficult) == 1:
# continue
if cls not in classes:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
bb = convert((w,h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
in_file.close()
out_file.close()
wd = os.getcwd()
#wd = os.getcwd()
data_base_dir = os.path.join("./", "VOCdevkit/")
if not os.path.isdir(data_base_dir):
os.mkdir(data_base_dir)
work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
if not os.path.isdir(work_sapce_dir):
os.mkdir(work_sapce_dir)
annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
if not os.path.isdir(annotation_dir):
os.mkdir(annotation_dir)
clear_hidden_files(annotation_dir)
image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
if not os.path.isdir(image_dir):
os.mkdir(image_dir)
clear_hidden_files(image_dir)
yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
if not os.path.isdir(yolo_labels_dir):
os.mkdir(yolo_labels_dir)
clear_hidden_files(yolo_labels_dir)
yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(yolov5_labels_dir):
os.mkdir(yolov5_labels_dir)
clear_hidden_files(yolov5_labels_dir)
yolov5_images_train_dir = os.path.join(data_base_dir, "train/")
if not os.path.isdir(yolov5_images_train_dir):
os.mkdir(yolov5_images_train_dir)
clear_hidden_files(yolov5_images_train_dir)
yolov5_images_test_dir = os.path.join(data_base_dir, "val/")
if not os.path.isdir(yolov5_images_test_dir):
os.mkdir(yolov5_images_test_dir)
clear_hidden_files(yolov5_images_test_dir)
yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
if not os.path.isdir(yolov5_labels_train_dir):
os.mkdir(yolov5_labels_train_dir)
clear_hidden_files(yolov5_labels_train_dir)
yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
if not os.path.isdir(yolov5_labels_test_dir):
os.mkdir(yolov5_labels_test_dir)
clear_hidden_files(yolov5_labels_test_dir)
list_imgs = os.listdir(image_dir) # list image files
probo = random.randint(1, 100)
print("Probobility: %d" % probo)
for i in range(0,len(list_imgs)):
path = os.path.join(image_dir,list_imgs[i])
if os.path.isfile(path):
image_path = image_dir + list_imgs[i]
voc_path = list_imgs[i]
(nameWithoutExtention, extention) = os.path.splitext(os.path.basename(image_path))
(voc_nameWithoutExtention, voc_extention) = os.path.splitext(os.path.basename(voc_path))
annotation_name = nameWithoutExtention + '.xml'
annotation_path = os.path.join(annotation_dir, annotation_name)
label_name = nameWithoutExtention + '.txt'
label_path = os.path.join(yolo_labels_dir, label_name)
probo = random.randint(1, 100)
print("Probobility: %d" % probo)
if(probo < 80): # train dataset
if os.path.exists(annotation_path):
convert_annotation(nameWithoutExtention) # convert label
copyfile(image_path, yolov5_images_train_dir + voc_path)
copyfile(label_path, yolov5_labels_train_dir + label_name)
else: # test dataset
if os.path.exists(annotation_path):
convert_annotation(nameWithoutExtention) # convert label
copyfile(image_path, yolov5_images_test_dir + voc_path)
copyfile(label_path, yolov5_labels_test_dir + label_name)上述代码执行完毕后，将在数据集的根目录生成
注意，lables与train。val是同级目录，这点与其他版本有差异
2.2、创建分类配置文件 "data/basketball.txt" ，其中包含需要进行检测的所有类别。具体说明如下：
以上改写遵循以下原则
# COCO 2017 dataset http://cocodataset.org
train: VOCdevkit/train # 118287 images
val: VOCdevkit/val # 5000 images
# number of classes
nc: 1
# whether it is coco dataset, only coco dataset should be set to True.
is_coco: False
# class names
names: [ 'basketball']其中，is_coco需要填写"False"
3、训练
3.1 修改train.py中相关配置
主要是修改tools/train.py即可
修改如下内容
parser.add_argument('--data-path', default='./data/basketball.yaml', type=str, help='path of dataset')3.2、运行
python tools/train.py将出现如下打印
Training start...
Epoch iou_loss l1_loss obj_loss cls_loss
0%| | 0/496 [00:00<?, ?it/s] /root/anaconda3/envs/yolov6/lib/python3.8/site-packages/torch/functional.py:478: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:2895.)
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
0/399 0 0 7.595 0: 100%|██████████| 496/496 [02:28<00:00, 3.34it/s]
Inferencing model in val datasets.: 100%|█████████████████████████████████████████████████████████████| 61/61 [00:13<00:00, 4.43it/s]
Evaluating speed.
Evaluating mAP by pycocotools.
Epoch: 0 | mAP@0.5: 0.0 | mAP@0.50:0.95: 0.0
Epoch iou_loss l1_loss obj_loss cls_loss
1/399 0 0 0.0003865 0: 100%|██████████| 496/496 [02:33<00:00, 3.24it/s]
Epoch iou_loss l1_loss obj_loss cls_loss
2/399 0 0 0.0003506 0: 100%|██████████| 496/496 [02:28<00:00, 3.33it/s]
Epoch iou_loss l1_loss obj_loss cls_loss
3/399 0 0 0.0002806 0: 100%|██████████| 496/496 [02:13<00:00, 3.72it/s]
Epoch iou_loss l1_loss obj_loss cls_loss
4/399 0 0 0.0002111 0: 100%|██████████| 496/496 [02:13<00:00, 3.72it/s]
Epoch iou_loss l1_loss obj_loss cls_loss
5/399 0 0 0.0001935 0: 100%|██████████| 496/496 [02:26<00:00, 3.38it/s]并且在runs/train将会生成相应的中间文件以及模型文件
全部评论 (0)
还没有任何评论哟~