YOLOv7训练记录-制作数据集

一、xml文件转换为txt

使用labelimg进行标注后，得到xml格式标签，需要转换成txt格式才能用于训练。
Annotations存放xml文件，JPEGImages存放jpg格式文件

import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join

# 数据标签
classes = ['insulator','self-explosion','flashover','breakage','foreign bodies']

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
    if w>=1:
        w=0.99
    if h>=1:
        h=0.99
    return (x,y,w,h)

def convert_annotation(rootpath,xmlname):
    xmlpath = rootpath + '/inAnnotations'
    xmlfile = os.path.join(xmlpath,xmlname)
    with open(xmlfile, "r", encoding='UTF-8') as in_file:
      txtname = xmlname[:-4]+'.txt'
      print(txtname)
      txtpath = rootpath + '/labels'  #生成的.txt文件会被保存在labels目录下
      if not os.path.exists(txtpath):
        os.makedirs(txtpath)
      txtfile = os.path.join(txtpath,txtname)
      with open(txtfile, "w+" ,encoding='UTF-8') as out_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)
        out_file.truncate()
        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(a) for a in bb]) + '\n')


if __name__ == "__main__":
    rootpath='D:/dataset/five-insulator'
    xmlpath=rootpath+'/inAnnotations'
    list=os.listdir(xmlpath)
    for i in range(0,len(list)) :
        path = os.path.join(xmlpath,list[i])
        if ('.xml' in path)or('.XML' in path):
            convert_annotation(rootpath,list[i])
            print('done', i)
        else:
            print('not xml file',i)

二、 划分训练、测试、验证集

新建文件夹dataset/ImageSets/main,存放划分后的标签，按照8:1:1划分。

import os
import random
import sys


root_path = 'D:/dataset/five-insulator1'
xmlfilepath = root_path + '/inAnnotations'
txtsavepath = root_path + '/imageset/main'


if not os.path.exists(txtsavepath):
    os.makedirs(txtsavepath)

train_test_percent = 0.9 # (训练集+验证集)/(训练集+验证集+测试集)
train_valid_percent = 8/9  # 训练集/(训练集+验证集)

total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * train_test_percent)   # 训练集+验证集数量
ts = int(num-tv)   # 测试集数量
tr = int(tv * train_valid_percent)  # 验证集数量
tz = int(tv-tr)   # 训练集数量
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)

print("train and valid size:", tv)
print("train size:", tz)
print("test size:", ts)
print("valid size:", tr)

# ftrainall = open(txtsavepath + '/ftrainall.txt', 'w')
ftest = open(txtsavepath + '/test.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fvalid = open(txtsavepath + '/valid.txt', 'w')

ftestimg = open(txtsavepath + '/img_test.txt', 'w')
ftrainimg = open(txtsavepath + '/img_train.txt', 'w')
fvalidimg = open(txtsavepath + '/img_valid.txt', 'w')

for i in list:
    name = total_xml[i][:-4] + '.txt' + '\n'
    imgname = total_xml[i][:-4] + '.jpg' + '\n'
    if i in trainval:
        # ftrainall.write(name)
        if i in train:
            ftrain.write(name)
            ftrainimg.write(imgname)
        else:
            fvalid.write(name)
            fvalidimg.write(imgname)
    else:
        ftest.write(name)
        ftestimg.write(imgname)

# ftrainall.close()
ftrain.close()
fvalid.close()
ftest.close()

ftrainimg.close()
fvalidimg.close()
ftestimg.close()

三、重构数据集

重构后数据集下是

import os
import shutil

# 获取分割好的train\test\valid名称
img_txt_cg_train = []
img_txt_cg_test = []
img_txt_cg_valid = []
label_txt_cg_train = []
label_txt_cg_test = []
label_txt_cg_valid = []

path = 'D:/dataset/five-insulator1/imageset/main/'

for line in open(path+"img_train.txt"):
    line=line.strip('\n')
    img_txt_cg_train.append(line)
for line1 in open(path+"img_test.txt"):
    line1=line1.strip('\n')
    img_txt_cg_test.append(line1)
for line2 in open(path+"img_valid.txt"):
    line2=line2.strip('\n')
    img_txt_cg_valid.append(line2)

for line3 in open(path+"train.txt"):
    line3=line3.strip('\n')
    label_txt_cg_train.append(line3)
for line4 in open(path+"test.txt"):
    line4=line4.strip('\n')
    label_txt_cg_test.append(line4)
for line5 in open(path+"valid.txt"):
    line5=line5.strip('\n')
    label_txt_cg_valid.append(line5)

# 建立cg数据的文件夹
new_dataset_train = 'D:/dataset/five-insulator1/images/train/'
new_dataset_test = 'D:/dataset/five-insulator1/images/test/'
new_dataset_valid = 'D:/dataset/five-insulator1/images/valid/'

new_dataset_trainl ='D:/dataset/five-insulator1/labels/train/'
new_dataset_testl ='D:/dataset/five-insulator1/labels/test/'
new_dataset_validl ='D:/dataset/five-insulator1/labels/valid/'

if not os.path.exists(new_dataset_train):
    os.makedirs(new_dataset_train)
if not os.path.exists(new_dataset_test):
    os.makedirs(new_dataset_test)
if not os.path.exists(new_dataset_valid):
    os.makedirs(new_dataset_valid)
if not os.path.exists(new_dataset_trainl):
    os.makedirs(new_dataset_trainl)
if not os.path.exists(new_dataset_testl):
    os.makedirs(new_dataset_testl)
if not os.path.exists(new_dataset_validl):
    os.makedirs(new_dataset_validl)

# cg移动
fimg = 'D:/dataset/five-insulator1/JPGinsulator/'
flable = 'D:/dataset/five-insulator1/labels/'

# 小数据建议：copy 大数据建议：move
for i in range(len(img_txt_cg_train)):
    shutil.copy(fimg+str(img_txt_cg_train[i]),new_dataset_train)
    shutil.copy(flable+str(label_txt_cg_train[i]),new_dataset_trainl)
for j in range(len(img_txt_cg_test)):
    shutil.copy(fimg+str(img_txt_cg_test[j]),new_dataset_test)
    shutil.copy(flable+str(label_txt_cg_test[j]),new_dataset_testl)
for q in range(len(img_txt_cg_valid)):
    shutil.copy(fimg+str(img_txt_cg_valid[q]),new_dataset_valid)
    shutil.copy(flable+str(label_txt_cg_valid[q]),new_dataset_validl)

四、 数据增强

运动模糊，添加雾雨雪，但是无法同时变换标签

# ##############简易变换#################

# https://imgaug.readthedocs.io/en/latest/source/overview_of_augmenters.html
import cv2
from imgaug import augmenters as iaa
import os

# Sometimes(0.5, ...) applies the given augmenter in 50% of all cases,
# e.g. Sometimes(0.5, GaussianBlur(0.3)) would blur roughly every second image.
# sometimes = lambda aug: iaa.Sometimes(0.5, aug)

# 定义一组变换方法.
seq = iaa.Sequential([
    iaa.MotionBlur(k=15),  # 运动模糊
    # iaa.Clouds(),  # 云雾
    #iaa.imgcorruptlike.Fog(severity=1),  # 多雾/霜
    #iaa.imgcorruptlike.Snow(severity=2),  # 下雨、大雪
    iaa.Rain(drop_size=(0.10, 0.20), speed=(0.2, 0.3)),  # 雨
    #iaa.Rain(speed=(0.3, 0.5)),  # 雨
    #iaa.Snowflakes(flake_size=(0.6, 0.7), speed=(0.02, 0.03)), # 雪点
    # iaa.imgcorruptlike.Spatter(severity=2),  # 溅 123水滴、45泥
    # iaa.contrast.LinearContrast((0.5, 2.0), per_channel=0.5),# 对比度变为原来的一半或者二倍
    # iaa.imgcorruptlike.Brightness(severity=2),  # 亮度增加
    # iaa.imgcorruptlike.Saturate(severity=3),  # 色彩饱和度
    # iaa.FastSnowyLandscape(lightness_threshold=(100, 255),lightness_multiplier=(1.5, 2.0)), # 雪地   亮度阈值是从 uniform(100, 255)（每张图像）和来自 uniform(1.5, 2.0)（每张图像）的乘数采样的。
    # iaa.Cartoon(blur_ksize=3, segmentation_size=1.0, saturation=2.0, edge_prevalence=1.0), # 卡通

])

# 图片文件相关路径
path = './input/'
savedpath = './output/'

imglist = []
filelist = os.listdir(path)

# 遍历要增强的文件夹，把所有的图片保存在imglist中
for item in filelist:
    img = cv2.imread(path + item)
    # print('item is ',item)
    # print('img is ',img)
    # images = load_batch(batch_idx)
    imglist.append(img)
# print('imglist is ' ,imglist)
print('all the picture have been appent to imglist')

# 对文件夹中的图片进行增强操作，循环1次
for count in range(1):
    images_aug = seq.augment_images(imglist)
    for index in range(len(images_aug)):
        # filename = str(count) + str(index) + '.jpg'
        # 保存图片
        filename = str(filelist[index])
        cv2.imwrite(savedpath + filename, images_aug[index])
        print('image of count%s index%s has been writen' % (count, index))

翻转、亮度改变等，可同时改变标签

# 数据增强

# import xml.etree.ElementTree as ET
import pickle
import os
from os import getcwd
import numpy as np
from PIL import Image
import shutil
import matplotlib.pyplot as plt

import imgaug as ia
from imgaug import augmenters as iaa


ia.seed(1)


def read_xml_annotation(root, image_id):
    in_file = open(os.path.join(root, image_id))
    tree = ET.parse(in_file)
    root = tree.getroot()
    bndboxlist = []

    for object in root.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        xmin = int(bndbox.find('xmin').text)
        xmax = int(bndbox.find('xmax').text)
        ymin = int(bndbox.find('ymin').text)
        ymax = int(bndbox.find('ymax').text)
        # print(xmin,ymin,xmax,ymax)
        bndboxlist.append([xmin, ymin, xmax, ymax])
        # print(bndboxlist)

    bndbox = root.find('object').find('bndbox')
    return bndboxlist


# (506.0000, 330.0000, 528.0000, 348.0000) -> (520.4747, 381.5080, 540.5596, 398.6603)
def change_xml_annotation(root, image_id, new_target):
    new_xmin = new_target[0]
    new_ymin = new_target[1]
    new_xmax = new_target[2]
    new_ymax = new_target[3]

    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    xmlroot = tree.getroot()
    object = xmlroot.find('object')
    bndbox = object.find('bndbox')
    xmin = bndbox.find('xmin')
    xmin.text = str(new_xmin)
    ymin = bndbox.find('ymin')
    ymin.text = str(new_ymin)
    xmax = bndbox.find('xmax')
    xmax.text = str(new_xmax)
    ymax = bndbox.find('ymax')
    ymax.text = str(new_ymax)
    tree.write(os.path.join(root, str("%06d" % (str(id) + '.xml'))))


def change_xml_list_annotation(root, image_id, new_target, saveroot, id):
    in_file = open(os.path.join(root, str(image_id) + '.xml'))  # 这里root分别由两个意思
    tree = ET.parse(in_file)
    elem = tree.find('filename')
    elem.text = (str("%06d" % int(id)) + '.png')
    xmlroot = tree.getroot()
    index = 0

    for object in xmlroot.findall('object'):  # 找到root节点下的所有country节点
        bndbox = object.find('bndbox')  # 子节点下节点rank的值

        # xmin = int(bndbox.find('xmin').text)
        # xmax = int(bndbox.find('xmax').text)
        # ymin = int(bndbox.find('ymin').text)
        # ymax = int(bndbox.find('ymax').text)

        new_xmin = new_target[index][0]
        new_ymin = new_target[index][1]
        new_xmax = new_target[index][2]
        new_ymax = new_target[index][3]

        xmin = bndbox.find('xmin')
        xmin.text = str(new_xmin)
        ymin = bndbox.find('ymin')
        ymin.text = str(new_ymin)
        xmax = bndbox.find('xmax')
        xmax.text = str(new_xmax)
        ymax = bndbox.find('ymax')
        ymax.text = str(new_ymax)

        index = index + 1

    tree.write(os.path.join(saveroot, str("%06d" % int(id)) + '.xml'))


def mkdir(path):
    # 去除首位空格
    path = path.strip()
    # 去除尾部 \ 符号
    path = path.rstrip("\\")
    # 判断路径是否存在
    # 存在     True
    # 不存在   False
    isExists = os.path.exists(path)
    # 判断结果
    if not isExists:
        # 如果不存在则创建目录
        # 创建目录操作函数
        os.makedirs(path)
        print(path + ' 创建成功')
        return True
    else:
        # 如果目录存在则不创建，并提示目录已存在
        print(path + ' 目录已存在')
        return False


if __name__ == "__main__":

    IMG_DIR = "D:/dataset/pictures/JPGinsulator"
    #IMG_DIR = "D:/shiyan/yangben/chicun/"
    #IMG_DIR = "D:/shiyan/VOC2007/JPEGImages/"
    #XML_DIR = "D:/shiyan/yangben/xml"
    XML_DIR = "D:/dataset/pictures/Annotations"
    #AUG_XML_DIR = "D:/shiyan/yangben/Annotations"  # 存储增强后的XML文件夹路径
    AUG_XML_DIR = "D:/dataset/pictures/inAnnotations"
    try:
        shutil.rmtree(AUG_XML_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_XML_DIR)

    #AUG_IMG_DIR = "D:/shiyan/yangben/Images"  # 存储增强后的影像文件夹路径
    AUG_IMG_DIR = "D:/dataset/pictures/insulator"
    try:
        shutil.rmtree(AUG_IMG_DIR)
    except FileNotFoundError as e:
        a = 1
    mkdir(AUG_IMG_DIR)

    AUGLOOP = 10  # 每张影像增强的数量

    boxes_img_aug_list = []
    new_bndbox = []
    new_bndbox_list = []

    # 影像增强
    seq = iaa.Sequential([
        iaa.Flipud(0.5),  # vertically flip 50% of all images垂直 对50%的图像做镜像翻转
        iaa.Fliplr(0.5),  # 镜像 水平 对50%的图像进行上下翻转
        #iaa.Multiply((1.2, 1.5)),  # change brightness, doesn't affect BBs
        iaa.Multiply((0.5, 1.5)),
        #iaa.GaussianBlur(sigma=(0, 3.0)),  # iaa.GaussianBlur(0.5),
        iaa.GaussianBlur(sigma=(0, 2.0)),
        #iaa.AdditiveGaussianNoise(1.5),
        iaa.AdditiveGaussianNoise(0.5),
        #iaa.Crop(percent=(0, 0.1)),
        iaa.Affine(
            translate_px={"x": 15, "y": 15},
            scale=(0.8, 0.95),
            #rotate=(0, 15)
        )  # translate by 40/60px on x/y axis, and scale to 50-70%, affects BBs
    ])

    for root, sub_folders, files in os.walk(XML_DIR):

        for name in files:

            bndbox = read_xml_annotation(XML_DIR, name)
            shutil.copy(os.path.join(XML_DIR, name), AUG_XML_DIR)
            shutil.copy(os.path.join(IMG_DIR, name[:-4] + '.jpg'), AUG_IMG_DIR)

            for epoch in range(AUGLOOP):
                seq_det = seq.to_deterministic()  # 保持坐标和图像同步改变，而不是随机
                # 读取图片
                img = Image.open(os.path.join(IMG_DIR, name[:-4] + '.jpg'))
                # sp = img.size
                img = np.asarray(img)
                # bndbox 坐标增强
                for i in range(len(bndbox)):
                    bbs = ia.BoundingBoxesOnImage([
                        ia.BoundingBox(x1=bndbox[i][0], y1=bndbox[i][1], x2=bndbox[i][2], y2=bndbox[i][3]),
                    ], shape=img.shape)

                    bbs_aug = seq_det.augment_bounding_boxes([bbs])[0]
                    boxes_img_aug_list.append(bbs_aug)

                    # new_bndbox_list:[[x1,y1,x2,y2],...[],[]]
                    n_x1 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x1)))
                    n_y1 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y1)))
                    n_x2 = int(max(1, min(img.shape[1], bbs_aug.bounding_boxes[0].x2)))
                    n_y2 = int(max(1, min(img.shape[0], bbs_aug.bounding_boxes[0].y2)))
                    if n_x1 == 1 and n_x1 == n_x2:
                        n_x2 += 1
                    if n_y1 == 1 and n_y2 == n_y1:
                        n_y2 += 1
                    if n_x1 >= n_x2 or n_y1 >= n_y2:
                        print('error', name)
                    new_bndbox_list.append([n_x1, n_y1, n_x2, n_y2])
                # 存储变化后的图片
                image_aug = seq_det.augment_images([img])[0]
                path = os.path.join(AUG_IMG_DIR,
                                    str("%06d" % (len(files) + int(name[:-4]) + epoch * 250)) + '.jpg')
                image_auged = bbs.draw_on_image(image_aug, thickness=0)
                Image.fromarray(image_auged).save(path)

                # 存储变化后的XML
                change_xml_list_annotation(XML_DIR, name[:-4], new_bndbox_list, AUG_XML_DIR,
                                           len(files) + int(name[:-4]) + epoch * 250)
                print(str("%06d" % (len(files) + int(name[:-4]) + epoch * 250)) + '.jpg')
                new_bndbox_list = []