yolov5测试和训练自己的数据集

ylov5测试与训练自己的数据集

项目地址：yolov5-git官方地址

1. 说明和环境配置
2. 测试自己环境
3. 制作自己的数据集
4. 开始训练和检测自己的模型
5. 过程遇到的问题解决

1.环境配置和说明

本人未做深度学习方面的研究，主要是复现代码，不涉及原理部分的讲解，若需要进一步了原理部分，可以移步较高博主的写的博客。
如果有问题欢迎在评论区讨论
环境配置：pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt（采用国内的源，安装较快）
个人建议：torch和opencv-python单独安装

2.测试自己的环境

克隆项目到本地：

git clone https://github.com/ultralytics/yolov5.git
cd yolov5
python detect.py

检测结果在run/detect/exp文件里面
如果能出现上面的图片，说明环境没有问题

3制作自己的数据集

这里我展示将kitti数据集转化为yolo格式的数据集形式进行训练
下载kitti数据集:下载连接
下载完成后，images是图片，labels是标签，但这个标签是kitti格式的，不能用于yolov5训练，进行转换
step1

# modify_annotations_txt.py
import glob
import string
txt_list = glob.glob('./kitti_labels/*.txt') 
def show_category(txt_list):
    category_list= []
    for item in txt_list:
        try:
            with open(item) as tdf:
                for each_line in tdf:
                    labeldata = each_line.strip().split(' ') 
                    category_list.append(labeldata[0]) 
        except IOError as ioerr:
            print('File error:'+str(ioerr))
    print(set(category_list)) 
def merge(line):
    each_line=''
    for i in range(len(line)):
        if i!= (len(line)-1):
            each_line=each_line+line[i]+' '
        else:
            each_line=each_line+line[i] 
    each_line=each_line+'\n'
    return (each_line)
print('before modify categories are:\n')
show_category(txt_list)
for item in txt_list:
    new_txt=[]
    try:
        with open(item, 'r') as r_tdf:
            for each_line in r_tdf:
                labeldata = each_line.strip().split(' ')
                if labeldata[0] in ['Truck','Van','Tram']:
                    labeldata[0] = labeldata[0].replace(labeldata[0],'Car')
                if labeldata[0] == 'cyclist': 
                    continue
                if labeldata[0] == 'pedestrian': 
                    continue
                if labeldata[0] == 'person_sitting': 
                    continue
                if labeldata[0] == 'dontCare': 
                    continue
                if labeldata[0] == 'misc': 
                    continue
                new_txt.append(merge(labeldata)) 
        with open(item,'w+') as w_tdf: 
            for temp in new_txt:
                w_tdf.write(temp)
    except IOError as ioerr:
        print('File error:'+str(ioerr))
print('\nafter modify categories are:\n')
show_category(txt_list) 

由于我只进行车辆的检测，所以我将’Truck’,‘Van’,‘Tram’
'Car’合并到"Car"类里面，其他的类别进行了忽略
step2
再将其转化为xml文件
这里需要创建一个Annotations文件夹用于存放xml

# encoding:utf-8
# # txt_to_xml.py
# 根据一个给定的XML Schema，使用DOM树的形式从空白文件生成一个XML
import os
from xml.dom.minidom import Document
import cv2
def generate_xml(name,split_lines,img_size,class_ind):
    doc = Document()  # 创建DOM文档对象

    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)

    title = doc.createElement('folder')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    annotation.appendChild(title)

    img_name=name+'.png'   ##和自己数据格式对应的格式，我这里所png

    title = doc.createElement('filename')
    title_text = doc.createTextNode(img_name)
    title.appendChild(title_text)
    annotation.appendChild(title)

    source = doc.createElement('source')
    annotation.appendChild(source)

    title = doc.createElement('database')
    title_text = doc.createTextNode('The KITTI Database')
    title.appendChild(title_text)
    source.appendChild(title)

    title = doc.createElement('annotation')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    source.appendChild(title)

    size = doc.createElement('size')
    annotation.appendChild(size)

    title = doc.createElement('width')
    title_text = doc.createTextNode(str(img_size[1]))
    title.appendChild(title_text)
    size.appendChild(title)

    title = doc.createElement('height')
    title_text = doc.createTextNode(str(img_size[0]))
    title.appendChild(title_text)
    size.appendChild(title)

    title = doc.createElement('depth')
    title_text = doc.createTextNode(str(img_size[2]))
    title.appendChild(title_text)
    size.appendChild(title)

    for split_line in split_lines:
        line=split_line.strip().split()
        if line[0] in class_ind:
            object = doc.createElement('object')
            annotation.appendChild(object)

            title = doc.createElement('name')
            title_text = doc.createTextNode(line[0])
            title.appendChild(title_text)
            object.appendChild(title)

            title = doc.createElement('difficult')
            title_text = doc.createTextNode('0')
            title.appendChild(title_text)
            object.appendChild(title)

            bndbox = doc.createElement('bndbox')
            object.appendChild(bndbox)
            title = doc.createElement('xmin')
            title_text = doc.createTextNode(str(int(float(line[4]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymin')
            title_text = doc.createTextNode(str(int(float(line[5]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('xmax')
            title_text = doc.createTextNode(str(int(float(line[6]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymax')
            title_text = doc.createTextNode(str(int(float(line[7]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)

    # 将DOM对象doc写入文件
    f = open('./Annotations/'+name+'.xml','w') # create a new xml file
    f.write(doc.toprettyxml(indent = ''))
    f.close()
# #source code
if __name__ == '__main__':
    class_ind=( 'Car')  #自己的类别
    #cur_dir=os.getcwd()  # current path
    #labels_dir=os.path.join(cur_dir,'labels') # get the current path and build a new path.and the result is'../yolo_learn/labels'
    labels_dir='./kitti_labels'   ##自己标签的路径
    for parent, dirnames, filenames in os.walk(labels_dir): # 分别得到根目录，子目录和根目录下文件
        for file_name in filenames:
            full_path=os.path.join(parent, file_name) # 获取文件全路径
            f=open(full_path)
            split_lines = f.readlines()
            name= file_name[:-4] # 后四位是扩展名.txt，只取前面的文件名
            img_name=name+'.png'
            img_path=os.path.join('./images/train',img_name) # 路径需要自行修改
            print(img_path)
            img_size =cv2.imread(img_path).shape
            generate_xml(name,split_lines,img_size,class_ind)
print('all txts has converted into xmls')

step3
经过上一步 已经将.txt文件转化为.xml文件
还需要将.xml文件转化为我们所需要的yolo格式的.txt
将下面代码放置在你的根目录即可
注意不要将图片文件夹和要生成文件夹的名称重复

"""
xml_to_txt_yolo.py 
将此文件放置在你的数据集根目录下即可
"""
import xml.etree.ElementTree as ET
import os
import shutil
import random

xml_file_path = 'Annotations/'       # 检查和自己的xml文件夹名称是否一致
images_file_path = 'image/train/'  # 检查和自己的图像文件夹名称是否一致
# 改成自己的类别名称
classes = ["Car"]
# 数据集划分比例，训练集75%，验证集15%，测试集15%
train_percent = 0.7
val_percent = 0.15
test_percent = 0.15
# 此处不要改动，只是创一个临时文件夹
if not os.path.exists('temp_labels/'):
    os.makedirs('temp_labels/')
txt_file_path = 'temp_labels/'


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_annotations(image_name):
    in_file = open(xml_file_path + image_name + '.xml')
    out_file = open(txt_file_path + image_name + '.txt', '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 == 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(a) for a in bb]) + '\n')


total_xml = os.listdir(xml_file_path)
num_xml = len(total_xml)  # XML文件总数

for i in range(num_xml):
    name = total_xml[i][:-4]
    convert_annotations(name)


# *********************************************** #
#  parent folder
#  --data
#  ----images
#       ----train
#       ----val
#       ----test
#  ----labels
#       ----train
#       ----val
#       ----test
def create_dir():
    if not os.path.exists('images/'):
        os.makedirs('images/')
    if not os.path.exists('labels/'):
        os.makedirs('labels/')
    if not os.path.exists('images/train'):
        os.makedirs('images/train')
    if not os.path.exists('images/val'):
        os.makedirs('images/val')
    if not os.path.exists('images/test'):
        os.makedirs('images/test/')
    if not os.path.exists('labels/train'):
        os.makedirs('labels/train')
    if not os.path.exists('labels/val'):
        os.makedirs('labels/val')
    if not os.path.exists('labels/test'):
        os.makedirs('labels/test')

    return


# *********************************************** #
# 读取所有的txt文件
create_dir()
total_txt = os.listdir(txt_file_path)
num_txt = len(total_txt)
list_all_txt = range(num_txt)  # 范围 range(0, num)

num_train = int(num_txt * train_percent)
num_val = int(num_txt * val_percent)
num_test = num_txt - num_train - num_val

train = random.sample(list_all_txt, num_train)
# train从list_all_txt取出num_train个元素
# 所以list_all_txt列表只剩下了这些元素：val_test
val_test = [i for i in list_all_txt if not i in train]
# 再从val_test取出num_val个元素，val_test剩下的元素就是test
val = random.sample(val_test, num_val)
# 检查两个列表元素是否有重合的元素
# set_c = set(val_test) & set(val)
# list_c = list(set_c)
# print(list_c)
# print(len(list_c))

print("训练集数目：{}, 验证集数目：{},测试集数目：{}".format(len(train), len(val), len(val_test) - len(val)))
for i in list_all_txt:
    name = total_txt[i][:-4]

    srcImage = images_file_path + name + '.png'
    srcLabel = txt_file_path + name + '.txt'

    if i in train:
        dst_train_Image = 'images/train/' + name + '.png'
        dst_train_Label = 'labels/train/' + name + '.txt'
        shutil.copyfile(srcImage, dst_train_Image)
        shutil.copyfile(srcLabel, dst_train_Label)
    elif i in val:
        dst_val_Image = 'images/val/' + name + '.jpg'
        dst_val_Label = 'labels/val/' + name + '.txt'
        shutil.copyfile(srcImage, dst_val_Image)
        shutil.copyfile(srcLabel, dst_val_Label)
    else:
        dst_test_Image = 'images/test/' + name + '.png'
        dst_test_Label = 'labels/test/' + name + '.txt'
        shutil.copyfile(srcImage, dst_test_Image)
        shutil.copyfile(srcLabel, dst_test_Label)
shutil.rmtree(txt_file_path)

运行代码便会生成以下文件夹，且将xml格式转化为yolo所需要txt格式

下面是kitti原始的txt
到此数据集就算制作完成了。

4 开始训练自己的数据集和测试

在yolov5的目录下建立一个kitti_data的文件夹
里面放置我们开始生成的Annotations,images和labels
images和labels都包括train,val,test三个文件夹，是将kitti数据集按7：1.5：1.5分的

开始训练

1.修改yaml文件
修改yaml文件,将dat目录下的coco128.yaml复制一份修改为my_obj.yaml

train: ./kitti_data/images/train/  # train images (relative to 'path') 128 images
val: ./kitti_data/images/val/  # val images (relative to 'path') 128 images
test: ./kitti_data/images/test/ # test images (optional)
nc:  1 # number of classes
names: ["Car"]  # class names

2.修改models文件
我这里用的yolov5s的模型，这个可以进行一个选择，主要修改里面的类别参数

# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 1  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
anchors:
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 v6.0 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Conv, [64, 6, 2, 2]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 6, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 3, C3, [1024]],
   [-1, 1, SPPF, [1024, 5]],  # 9
  ]

# YOLOv5 v6.0 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

3 修改train.py
以下是我的训练路径，到时候要根据自己的路径来改

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
    parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
    parser.add_argument('--data', type=str, default='data/my_obj.yaml', help='dataset.yaml path')
    parser.add_argument('--hyp', type=str, default= 'data/hyps/hyp.scratch.yaml', help='hyperparameters path')
    parser.add_argument('--epochs', type=int, default=10)
    parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch')

下面是训练参数的解释

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

训练过程可视化
利用tensorboard可视化训练过程，训练开始会在yolov5目录生成一个runs文件夹，利用tensorboard打开即可查看训练日志，命令如下：

tensorboard --logdir=runs

训练完成后在，runs文件夹里面会生成一个train文件夹里面有exp文件夹，权重文件就在里面，里面有一个最好的和最后的
下面我贴一张生成的结果图

下面命令进行测试：

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

下面命令检测自己的图片(将detect.py中的权重文件和测试图片改成自己)

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

到这里就算完成了YOLOv5训练自己的数据集整个过程：环境配置—制作数据集----模型训练----模型测试----模型推理阶段已全部完成。

参考连接1
参考连接2：https://blog.csdn.net/a18838956649/article/details/118639052

----------------------------end------------------------------------

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