YOLOv5训练自己的数据集(VOC格式的数据集和txt格式的数据集)

本文章主要介绍了两种数据集的准备,第一种是我只有VOC格式,这是需转txt。第二种是我只有txt格式的。
*
准备工作
环境:ubuntu18,虚拟环境python=3.8,CUDA=11.3(这是我的环境,其实根据yolov5官方需求就可以)
先把项目克隆:

git clone https://github.com/ultralytics/yolov5.git

在YOLOv5文件里打开终端,创建虚拟环境:

conda create -n yolov5 python=3.8

创建完成后:

conda activate yolov5

进入虚拟环境安装yolov5所需要的环境

 pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt```

OK环境没问题了
开始准备数据咯

*一.第一种数据保存(VOC转txt):

**
1.在yolov5文件夹中创建my_data文件夹,目录结构如下

--my_data
  --Annotations#保存的是用labelimage标注工具标注出来的.xml文件
  --images#我们的数据集图片 .jpg
  --ImageSets#之后会在该目录下创建Main文件,并自动生成train.txt,val.txt,test.txt和trainval.txt四个文件,存放训练集、验证集、测试集图片的名字(无后缀.jpg)
    --Main#z在这里创建一个空的文件夹

开始创建Main下的txt文件路径
在my_data文件夹下出创建一个split_train_val.py文件,
记得把所有路径改成自己的。代码如下:

# coding:utf-8

import os
import random
import argparse

parser = argparse.ArgumentParser()
#xml文件的地址,根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='/home/lad/yolov5/my_data/Annotations', type=str, help='input xml label path')
#数据集的划分,地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='/home/lad/yolov5/my_data/ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 1.0
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)

file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')

for i in list_index:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        file_trainval.write(name)
        if i in train:
            file_train.write(name)
        else:
            file_val.write(name)
    else:
        file_test.write(name)

file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

运行结束后,不出意外的在Main文件下生成四个txt文件

--Main
  --test.txt#这个应该是空的,毕竟一般只用训练集和验证集(又名测试集),我不用这个test测试集。为了让数据更加多的参与训练
  --train.txt
  --trainval.txt
  --val.txt

2.开始准备标签,把VOC的格式转成yolo的需要的格式,即xml转txt
创建voc_label.py文件,将训练集、验证集、测试集生成label标签(训练中要用到),同时将数据集路径导入txt文件中,代码内容如下:

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwd

sets = ['train', 'val', 'test']
classes = ["crack0"]   # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)

def convert(size, box):
    dw = 1. / (size[0])
    dh = 1. / (size[1])
    x = (box[0] + box[1]) / 2.0 - 1
    y = (box[2] + box[3]) / 2.0 - 1
    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('/home/lad/yolov5_hcb/my_data/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('/home/lad/yolov5_hcb/my_data/labels/%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
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            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))
        b1, b2, b3, b4 = b
        # 标注越界修正
        if b2 > w:
            b2 = w
        if b4 > h:
            b4 = h
        b = (b1, b2, b3, b4)
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

wd = getcwd()
for image_set in sets:
    if not os.path.exists('/home/lad/yolov5_hcb/my_data/labels/'):
        os.makedirs('/home/lad/yolov5_hcb/my_data/labels/')
    image_ids = open('/home/lad/yolov5_hcb/my_data/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()
    list_file = open('paper_data/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write(abs_path + '/my_data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

运行后会生成如下labels文件夹和三个包含数据集的txt文件,其中labels中为不同图像的标注文件,train.txt等txt文件为划分后图像所在位置的绝对路径,如train.txt就含有所有训练集图像的绝对路径。
最后的数据集目录应该是:

--my_data
  --Annotations
  --images
  --ImageSets
    --Main#这里面有第一步生成的四个文件夹
  --labels#刚生成的
  --test.txt#刚生成的
  --train.txt#刚生成的
  --val.txt#刚生成的

3.配置文件
1.在yolov5目录下的data文件夹下新建一个ab.yaml文件(可以自定义命名),用来存放训练集和验证集的划分文件(train.txt和val.txt),这两个文件是通过运行voc_label.py代码生成的(就是上面备注:#刚生成的那三个),然后是目标的类别数目和具体类别列表,ab.yaml内容如下:
(不过注意发现冒号后面需要加空格,否则会被认为是字符串而不是字典)

train: /home/lad/yolov5/my_data/train.txt # voclabel.py生成的train.txt的路径
val: /home/lad/yolov5/my_data/val.txt   # voclabel.py生成的val.txt的路径
 
# number of classes
nc: 5
 
# class names
names: ['a','b','c','d','e']

2.选择一个你需要的模型:比如就yolov5s吧
在models/yolov5s.yaml修改类别nc:
在这里插入图片描述
3.修改train.py
把这个对应的修改成我们为刚才的路径
在这里插入图片描述
以上参数解释如下:
epochs:指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。
batch-size:一次看完多少张图片才进行权重更新,梯度下降的mini-batch,显卡不行你就调小点。
cfg:存储模型结构的配置文件
data:存储训练、测试数据的文件
img-size:输入图片宽高,显卡不行你就调小点。
rect:进行矩形训练
resume:恢复最近保存的模型开始训练
nosave:仅保存最终checkpoint
notest:仅测试最后的epoch
evolve:进化超参数
bucket:gsutil bucket
cache-images:缓存图像以加快训练速度
weights:权重文件路径
name: 重命名results.txt to results_name.txt
device:cuda device, i.e. 0 or 0,1,2,3 or cpu
adam:使用adam优化
multi-scale:多尺度训练,img-size +/- 50%
single-cls:单类别的训练集
之后运行训练命令如下:
(路径依然改哦)

python train.py --img 640 --batch 16 --epoch 300 --data data/ab.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device '0'     # 0号GPU

4.训练过程可视化
利用tensorboard可视化训练过程,训练开始会在yolov5目录生成一个runs文件夹,利用tensorboard打开即可查看训练日志,打开新的终端,命令如下:

tensorboard --logdir= 路径/runs/expx --host=127.0.0.1

就可以可视化啦。
ok,VOC格式的已经完成。再来看看txt格式的吧。

二.第一种数据保存(txt)

1.在yolov5文件夹中创建my_data文件夹,目录结构如下

--yolov5
--其他clone的文件
--my_data
	--train#训练数据集的文件夹
		--images#训练数据集的图片 .jpg
		--labels#训练数据集的txt标注标签
	--val#验证数据集的文件夹
		--images#验证数据集的图片 .jpg
		--labels#验证数据集的txt标注标签
	--data.yaml

–data.yaml里的内容如下:

train: /home/lad/yolov5/my_data/train
val: /home/lad/yolov5/my_data/val

nc: [你的检测数据集的类别数]
names: [你的检测数据集的类别名称列表]

然后数据处理就结束了
2.选择一个你需要的模型:比如就yolov5s吧
在models/yolov5s.yaml修改类别nc:
在这里插入图片描述
3.修改train.py
把这个对应的修改成我们为刚才的路径
在这里插入图片描述
以上参数解释如下:
epochs:指的就是训练过程中整个数据集将被迭代多少次,显卡不行你就调小点。
batch-size:一次看完多少张图片才进行权重更新,梯度下降的mini-batch,显卡不行你就调小点。
cfg:存储模型结构的配置文件
data:存储训练、测试数据的文件
img-size:输入图片宽高,显卡不行你就调小点。
rect:进行矩形训练
resume:恢复最近保存的模型开始训练
nosave:仅保存最终checkpoint
notest:仅测试最后的epoch
evolve:进化超参数
bucket:gsutil bucket
cache-images:缓存图像以加快训练速度
weights:权重文件路径
name: 重命名results.txt to results_name.txt
device:cuda device, i.e. 0 or 0,1,2,3 or cpu
adam:使用adam优化
multi-scale:多尺度训练,img-size +/- 50%
single-cls:单类别的训练集
之后运行训练命令如下:
(路径依然改哦)

python train.py --img 640 --batch 16 --epoch 300 --data data/ab.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --device '0'     # 0号GPU

4.训练过程可视化
利用tensorboard可视化训练过程,训练开始会在yolov5目录生成一个runs文件夹,利用tensorboard打开即可查看训练日志,打开新的终端,命令如下:

tensorboard --logdir= 路径/runs/expx --host=127.0.0.1

就可以可视化啦。
ok,txt格式的已经完成。

测试

评估模型好坏就是在有标注的测试集或者验证集上进行模型效果的评估,在目标检测中最常使用的评估指标为mAP。在test.py文件中指定数据集配置文件和训练结果模型,之后运行代码如下:(改路径)

python test.py  --data data/ab.yaml --weights runs/exp1/weights/best.pt --augment

训练完成后可以运行yolov5/detect.py文件,指定一个文件夹,推理其中的图像(无标注)。就像该train.py一样该相应的参数,权重改成我们runs文件里训练出的。

python detect.py --weights runs/exp1/weights/best.pt --source inference/images/ --device 0 --save-txt

这里主要供自己自学使用,参考的文章如下;
https://blog.csdn.net/weixin_44966641/article/details/118605063
https://blog.csdn.net/qq_36756866/article/details/109111065
记录一下。
不然老忘记。好了,干饭干饭

版权声明:本文为CSDN博主「cheng-Lad」的原创文章,遵循CC 4.0 BY-SA版权协议,转载请附上原文出处链接及本声明。
原文链接:https://blog.csdn.net/weixin_48695781/article/details/121677849

cheng-Lad

我还没有学会写个人说明!

暂无评论

发表评论

相关推荐

Pytorch—万字入门SSD物体检测

前言 由于初入物体检测领域,我在学习SSD模型的时候遇到了很多的困难。一部分困难在于相关概念不清楚,专业词汇不知其意,相关文章不知所云;另一部分困难在于网上大部分文章要么只是简要介绍了SS