YOLOV5、V7 训练格式转换-训练自己的数据集-猛男技术控

一般我们自己标注的训练数据，一开始会有两个文件夹：

• 存放图片的：images
• 存放xml标注信息的：Annotations

yOLO训练最终需要的文件路径格式为：

images里面放的就是训练和验证的图片

labels里面放的是对应的训练和验证所用的txt文件，其内容如下：

0 1 是类别，后面是坐标位置

注：两个train和两个val文件夹中文件名称是对应的，只是文件名的后缀不同

---

我们需要做的就是

1. 将xml转换为txt
2. 将图片划分为训练集和测试集
3. 修改yaml文件

---

xml2txt.py ：xml转换成txt

# 作者：李富贵
# 公众号：猛男技术控
# 输入：xml的文件夹路径
# 输出：在txt路径下创建对应的txt文件
import xml.etree.ElementTree as ET
import os

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(xml_file, txt_path, filename):
    in_file = open('{}/{}.xml'.format(xml_file, filename), encoding='UTF-8')
    out_file = open('{}/{}.txt'.format(txt_path, filename), '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
        if obj.find('difficult'):
            difficult = obj.find('difficult').text
        else:
            difficult = 0
        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')


# 改成自己的类别
classes = ["C_wirerope_notbinding", "C_wirerope_n"]
# 改成自己xml文件夹和要保存的txt文件夹的路径
xml_path = 'Annotations'
txt_path = 'txt'
if not os.path.exists(txt_path):
    os.makedirs(txt_path)

xml_files = os.listdir(xml_path)
for xml_file in xml_files:
    print(xml_path + '/' + xml_file)
    convert_annotation(xml_path, txt_path, xml_file[:-4])
    

---

split.py 将图片和对应的txt分别放到训练集和验证集文件夹中

import os
import random
import shutil

# 这里不用更改，会在当前py文件路径下生成下面四个文件夹
set1 = ['images','labels']
set2 = ['train','val']
for s1 in set1:
    if not os.path.exists(s1):
        os.mkdir(s1)
    for s2 in set2:
        if not os.path.exists(s1+'/'+s2):
            os.mkdir(s1+'/'+s2)


# 这是原始图片路径
img_path = 'JPEGImages'
# 这是生成的txt路径
txt_path = 'txt'
file_names = os.listdir(img_path)
l = 0.8
n = len(file_names)
train_files = random.sample(file_names, int(n*l))
for file in file_names:
    if file in train_files:
        shutil.copy(img_path+'/'+file,'images/train/'+file)
        shutil.copy(txt_path+'/'+file[:-3]+'txt','labels/train/'+file[:-3]+'txt')
    else:
        shutil.copy(img_path+'/'+file,'images/val/'+file)
        shutil.copy(txt_path+'/'+file[:-3]+'txt','labels/val/'+file[:-3]+'txt')

print(train_files)

---

创建data路径下的yaml文件：myvoc.yaml

修改对应的路径，类别数目以及类别名称

# 训练和测试文件的路径
train: visdronedata/images/train/
val: visdronedata/images/val/

# number of classes
nc: 2

# class names
names: ['dog', 'cat']

---

修改model文件夹下对应的模型的类别数目，比如我要用yolov5s，就修改yolov5s.yaml下的nc的数目

---

修改train.py 修改

weight修改为对应的，用yolov5s就写yolov5s.py，用x就写x

data改成刚才创建的myvoc.yaml

epochs为要训练的轮数

batchsize为一次训练多少张图

device为用哪个GPU训练