yolov5训练数据集

前提

最近要用yolov5训练看效果，发现网上很多博文都有点问题，感觉不适合我，就自己整理一下，给各位做个参考。

声明一下啊，代码有些是网上找的，然后根据我自己需求改了一下，侵权可以联系我删除。

电脑已配置conda环境，使用pycharm，仅供参考。

一、图像标注

我这里使用的Anaconda，打开Anaconda Prompt，创建一个环境（已有环境可以忽略这步）

conda create -n env_name python=x.x

env_name是环境名，随便取一个能记住的就行，python=x.x是指定python版本。

我这里使用3.9，环境名为了好记取了py39，进入环境

conda activate py39

安装labelimg

pip install labelimg

运行labelimg

labelimg

Open Dir是需要标注的图片路径，Change Save Dir是标注完图片保存的路径，第三个是指保存的格式，现在这个代表保存为VOC格式。

标注前建议打开自动保存，View-->Auto Save Mode。快捷键W为开始画框，A和D控制上下张。

二、数据处理

数据标注完成后，打开yolo v5文件夹，找到yolo v5下data文件夹，在其中新建一个文件夹，命名能记住就行，我标记了几张江西脐橙的数据，此处命名为jiangxi_qicheng。

将图片和标签丢进这个文件夹里，图片文件夹命名为images，标签文件夹命名为xml。

这里有个问题，你如果是按照我的方法来做，标签就是voc格式（文件拓展名为xml），如果你直接标注为yolo格式（拓展名为txt），按照后面的代码就运行不了，我懒得改代码了，如果你是yolo格式，可以使用下方代码，把yolo转换成VOC格式。

import os
import cv2
import pathlib
from xml.dom.minidom import Document


# 支持多种格式图片,可以自己添加
def search_file_name(file_path, img_root):
    file_format = ['.jpg', '.JPG', '.png', '.PNG', '.JPEG', '.jpeg', '.bmp']
    file_path = pathlib.Path(file_path)
    file_name = file_path.stem
    for endsw in file_format:
        file_judge = file_name + endsw
        if os.path.isfile(os.path.join(img_root, file_judge)):
            return file_name, endsw
    return 'stop', 0


def yolo2voc(src_img_root, src_label_root, dst_label_root, class_map):
    not_have_img = []
    if not os.path.exists(dst_label_root):
        os.makedirs(dst_label_root)
    # 遍历所有txt文件
    for i, src_label_name in enumerate(os.listdir(src_label_root)):
        xmlBuilder = Document()
        annotation = xmlBuilder.createElement("annotation")  # 创建annotation标签
        xmlBuilder.appendChild(annotation)
        src_label_path = os.path.join(src_label_root, src_label_name)
        file_name, endsw = search_file_name(src_label_path, src_img_root)
        if file_name == 'stop':
            not_have_img.append(os.path.basename(src_label_path))
            continue
        with open(src_label_path, 'r') as fr:
            txtlines = fr.readlines()
        src_img_name = file_name + endsw
        src_img_path = os.path.join(src_img_root, src_img_name)
        image_np = cv2.imread(src_img_path)
        Pheight, Pwidth, Pdepth = image_np.shape
        folder = xmlBuilder.createElement("folder")  # folder标签
        foldercontent = xmlBuilder.createTextNode("driving_annotation_dataset")
        folder.appendChild(foldercontent)
        annotation.appendChild(folder)  # folder标签结束

        filename = xmlBuilder.createElement("filename")  # filename标签
        filenamecontent = xmlBuilder.createTextNode(src_img_name)
        filename.appendChild(filenamecontent)
        annotation.appendChild(filename)  # filename标签结束

        size = xmlBuilder.createElement("size")  # size标签
        width = xmlBuilder.createElement("width")  # size子标签width
        widthcontent = xmlBuilder.createTextNode(str(Pwidth))
        width.appendChild(widthcontent)
        size.appendChild(width)  # size子标签width结束

        height = xmlBuilder.createElement("height")  # size子标签height
        heightcontent = xmlBuilder.createTextNode(str(Pheight))
        height.appendChild(heightcontent)
        size.appendChild(height)  # size子标签height结束

        depth = xmlBuilder.createElement("depth")  # size子标签depth
        depthcontent = xmlBuilder.createTextNode(str(Pdepth))
        depth.appendChild(depthcontent)
        size.appendChild(depth)  # size子标签depth结束
        annotation.appendChild(size)  # size标签结束
        for line in txtlines:
            oneline = line.strip().split(" ")
            object = xmlBuilder.createElement("object")  # object 标签
            picname = xmlBuilder.createElement("name")  # name标签
            namecontent = xmlBuilder.createTextNode(class_map[oneline[0]])
            picname.appendChild(namecontent)
            object.appendChild(picname)  # name标签结束

            pose = xmlBuilder.createElement("pose")  # pose标签
            posecontent = xmlBuilder.createTextNode("Unspecified")
            pose.appendChild(posecontent)
            object.appendChild(pose)  # pose标签结束

            truncated = xmlBuilder.createElement("truncated")  # truncated标签
            truncatedContent = xmlBuilder.createTextNode("0")
            truncated.appendChild(truncatedContent)
            object.appendChild(truncated)  # truncated标签结束

            difficult = xmlBuilder.createElement("difficult")  # difficult标签
            difficultcontent = xmlBuilder.createTextNode("0")
            difficult.appendChild(difficultcontent)
            object.appendChild(difficult)  # difficult标签结束

            bndbox = xmlBuilder.createElement("bndbox")  # bndbox标签
            xmin = xmlBuilder.createElement("xmin")  # xmin标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth)
            xminContent = xmlBuilder.createTextNode(str(mathData))
            xmin.appendChild(xminContent)
            bndbox.appendChild(xmin)  # xmin标签结束

            ymin = xmlBuilder.createElement("ymin")  # ymin标签
            mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight)
            yminContent = xmlBuilder.createTextNode(str(mathData))
            ymin.appendChild(yminContent)
            bndbox.appendChild(ymin)  # ymin标签结束

            xmax = xmlBuilder.createElement("xmax")  # xmax标签
            mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth)
            xmaxContent = xmlBuilder.createTextNode(str(mathData))
            xmax.appendChild(xmaxContent)
            bndbox.appendChild(xmax)  # xmax标签结束

            ymax = xmlBuilder.createElement("ymax")  # ymax标签
            mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight)
            ymaxContent = xmlBuilder.createTextNode(str(mathData))
            ymax.appendChild(ymaxContent)
            bndbox.appendChild(ymax)  # ymax标签结束
            object.appendChild(bndbox)  # bndbox标签结束
            annotation.appendChild(object)  # object标签结束
        dst_label_path = os.path.join(dst_label_root, file_name + '.xml')
        with open(dst_label_path, 'w') as f:
            xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8')
        print(src_img_path)
    print('not have img:', not_have_img, '\n', not_have_img.__len__())


if __name__ == "__main__":
    # 注意： opencv 不支持中文路径，路径及文件名需要是英文数字符号等
    src_img_root = r"F:\runs\detect\dakuiy2\0"             # 图片路径（需要获取图片shape）
    src_label_root = r"F:\runs\detect\dakuiy2\labels\0"    # 输入label路径（YOLO的txt格式）
    dst_label_root = r"F:\runs\detect\dakuiy2\labels\1"       # 输出label路径（PASCAL VOC的xml格式）
    # txt文件中每列第一个为类别序号，分别对应label名
    class_map = {'0': "hao", '1': "huai"}
    yolo2voc(src_img_root, src_label_root, dst_label_root, class_map)

图片有了，标签有了，开始划分训练集和验证集 ，只需要改xml和txt的path就行了，意思是将xml划分为train和val，然后把对应的名字保存为txt文件，到时候只需要读取txt文件就行了。

import os
import random
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='data/jiangxi_qicheng1/xml', type=str, help='input xml labels path')
parser.add_argument('--txt_path', default='data/jiangxi_qicheng1/dataset', type=str, help='output txt labels path')
opt = parser.parse_args()
trainval_percent = 1.0
train_percent = 0.8     #百分比，此处说明80%用作训练，20%用作测试，可修改
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_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)
file_trainval.close()
file_train.close()
file_val.close()

现在你的文件夹里有三个文件夹

然后把VOC格式转换成YOLO格式，并将dataset划分的train和val加上全路径名称。

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

sets = ['train', 'val']
classes = ["hao", "huai"]  # 改成自己的类别
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('data/jiangxi_qicheng1/xml/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('data/jiangxi_qicheng1/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.findall('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))
        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('data/jiangxi_qicheng1/labels/'):
        os.makedirs('data/jiangxi_qicheng1/labels/')
    # 修改部分：读取每一行的数据并作为 image_ids 的元素
    with open('data/jiangxi_qicheng1/dataset/%s.txt' % (image_set), 'r') as file:
	    image_ids = [line.strip() for line in file]
    if not os.path.exists('data/jiangxi_qicheng1/split/'):
        os.makedirs('data/jiangxi_qicheng1/split/')
    list_file = open('data/jiangxi_qicheng1/split/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        print(image_id)  # 打印 image_id 的值，检查是否包含完整的文件名
        list_file.write(abs_path + '/data/jiangxi_qicheng1/images/%s.bmp\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

注意：我标注的图片拓展名为bmp，要换成你对应的图片拓展名

此时文件夹里应该有这些文件夹

接下来需要制作yaml，用来指定一些参数

在data文件夹下新建一个yaml文件，内容如下，根据自己需要修改

path: data/jiangxi_qicheng1/split   #划分train和val的文件夹（有全部路径的）
train: train.txt
val: val.txt

# number of classes
nc: 4

# class names
names: ["huapi", "keshang", "kuiyang", "lan"]

三、开始训练

打开yolo下的train.py，下滑，找到需要我们修改的地方

这里解释一下常用的

weights：权重，就是在yolo官网下载的那个pt，看清啊，yolov5s-seg后面的seg代表是目标检测，cls代表是图像分类

cfg：yolo自带，根据你用的模型是 s、n、l 来选择，和权重对应

data：刚刚创建的那个yaml，会直接自动读取里面的文件

epochs：训练多少轮

batch-size：根据你电脑运算能力来设，跑不起来就设小点

imgsz：图片大小，32的倍数

修改完成后就可以开始训练了，直接运行train也行，conda命令运行也行

训练完成后，可以在runs/train文件夹下可以找到保存的文件

四、推理

跑完模型是不是想看下效果？

找到yolo下detect.py，和train一样，修改参数

改前四个够用了

weights：放上train下训练好的pt路径

 source：要预测的图片路径

data：上面写的yaml文件

imgsz：图片大小，和train一致

运行代码，然后在runs/detect里面找到推测完的数据和标签

这一步也可以用来快速增加数据，方法就是将预测出的标签和图片用labelimg打开，不过此时要在标签文件夹下新建classes.txt，里面放上类别，参考如下

hao
huai

打开后根据需要修改标注的框就行了