2023-07-19-Yolov5模型训练自己的数据集

下载yolov5

Clone repo and install requirements.txt in a Python>=3.7.0 environment, including PyTorch>=1.7.

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

下载模型

在https://github.com/ultralytics/yolov5往下翻，找到Pretrained Checkpoints

下载自己需要的模型,放到weight目录下(没有就新建一个)

测试自带模型

                            ①要测试的图片位置         ②权重文件的位置
python .\detect.py --source .\data\images --weights .\weight\yolov5s.pt

生成的结果如下，找到runs\detect\exp16 目录

可以看到这里的两张图片都已经标注完成了，但是这里多了一个labels目录，这个先不用管它，后面会讲。

标记数据集

用labelimg 标记数据集，注意格式要保存为YOLO支持的txt格式

Change save format改为YOLO

**如果已经将标注设置为PascalVOC**，也就是xml格式，可以使用以下代码转换

import xml.etree.ElementTree as ET
import os
from os import getcwd
import glob

# 1.
# 自己创建文件夹,例如：label_mal label_txt  也可以修改别的
image_set = 'folder_xml'  # 需要转换的文件夹名称（文件夹内放xml标签文件）
imageset2 = 'folder_txt'  # 保存txt的文件夹
# 2.
# 换成你的类别 当前的顺序，就txt 0,1,2,3,4 五个类别
classes = ['XB', 'TL', 'KL', 'XBHLF', 'SS']  # 标注时的标签 注意顺序一定不要错。

# 3.
# # 转换文件夹的绝对路径
data_dir = 'F:'
# 或者 读取当前路径
# data_dir = getcwd()  # 当前路径


'''
xml中框的左上角坐标和右下角坐标(x1,y1,x2,y2)
》》txt中的中心点坐标和宽和高(x,y,w,h)，并且归一化
'''


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(data_dir, imageset1, imageset2, image_id):
    in_file = open(data_dir + '/%s/%s.xml' % (imageset1, image_id), 'r', encoding='utf-8')   # 读取xml
    out_file = open(data_dir + '/%s/%s.txt' % (imageset2, image_id), 'w', encoding='utf-8')  # 保存txt
    print(in_file)

    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
        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('%.6f' % a) for a in bb]) + '\n')


image_ids = []
for x in glob.glob(data_dir + '/%s' % image_set + '/*.xml'):
    image_ids.append(os.path.basename(x)[:-4])
print('\n%s数量:' % image_set, len(image_ids))  # 确认数量
i = 0
for image_id in image_ids:
    i = i + 1
    convert_annotation(data_dir, image_set, imageset2, image_id)
    print("%s 数据:%s/%s文件完成！" % (image_set, i, len(image_ids)))

print("Done!!!")

训练数据集

我是在yolov5的主目录下新建了一个sample/tunnel目录，专门用来存放数据集，可以自定义目录

注意这个data.yaml ，names 就是你在labelImg 标注时的标签，nc 是names的数量, train是用来训练的数据集,val 是用来验证的数据集,如果验证数据的误差过大，会重新计算。

train和val目录下都有images和labels两个目录，其中images目录是用来存放图片的，labels是用来存放标注文件的。

                            ①单次训练的数量
                                      ②训练轮次     ③yaml文件位置                      ④训练的权重文件位置
python train.py --batch-size 2 --epochs 200 --data sample/tunnel/data.yaml --weights weight/yolov5s.pt

运行之后就是漫长的等待过程。

训练好之后，会在runs/train目录下生成训练好的文件

权重文件目录,这里的best.pt 就是训练批次最好的模型，last.pt 是最后一次的结果

测试训练好模型

python .\detect.py --source F:\images --weights runs/train/exp15/weights/best.pt

结果会在runs\detect目录下

同时这里生成了一个runs\detect\exp17\labels，这个问题在上面提出过

这是因为我的需求是要求测试的模型，同时生成txt格式的标注

在参数--save-txt 添加一个默认值

parser.add_argument('--save-txt', action='store_true', help='save results to *.txt',default= True)

这样运行detect.py 就会同时生成标注文件。

其他设置

同时如果觉得测试模型的命令太长了

python .\detect.py --source F:\images --weights runs/train/exp15/weights/best.pt

可以手动修改detect.py ，可以修改--source 和--weights 为自定义的值

这样就只需要运行

python .\detect.py

修改cuda device

parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')