利用yolov5训练数据

YOLOv5是一种基于深度学习的目标检测算法，可以用于识别图像或视频中的物体。利用yolov5训练模型需要以下几个步骤

一、数据采集及标注

在开始前需要准备好需要的数据图片，将其存放于yolov5的data文件夹的images文件夹中

 然后在data文件夹中新建dataset、labels、两个文件夹。其中，dataset又来保存标注好的图片，labels用来存放数据标签

准备完成后就可以开始标注图片了 ，我是使用labelimg进行标注的

首先打开终端利用 pip install labelimg  安装labelimg（我已经安装好了）

 在利用命令labelimg打开labelimg

出现此界面

点击view选择自动保存 （Auto Save mode）

 

 点击Change Save Dir更改为之前新建的dataset文件夹

 点击Open Dir将路径修改为存放图片数据集的文件夹images

完成后就可以开始标注了

 

 标注完成后打开dataset就能看到与图片对应的xml文件

 二、数据处理与源码修改

在工程下新建split的python文件

split代码如下

import os
import random
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', default='data/dataset', type=str, help='input xml label path')
parser.add_argument('--txt_path', default='data/labels', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 1.0
train_percent = 0.8
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_test = open(txtsavepath + '/test.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)
    else:
        file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

 运行代码会在labels下出现4个txt文件

再在工程新建xml_to_txt的python文件

 其中代码如下

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

sets = ['train', 'val', 'test']
classes = ['naok']  # 这里改为你要训练的标签，否则会报错。比如你要识别“hand”，那这里就改为hand


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):
    # try:
    in_file = open('data/dataset/%s.xml' % (image_id), encoding='utf-8')
    out_file = open('data/labels/%s.txt' % (image_id), 'w', encoding='utf-8')
    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))
        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')


# except Exception as e:
#     print(e, image_id)

wd = getcwd()
for image_set in sets:
    if not os.path.exists('data/labels/'):
        os.makedirs('data/labels/')
    image_ids = open('data/labels/%s.txt' %
                     (image_set)).read().strip().split()
    list_file = open('data/%s.txt' % (image_set), 'w')
    for image_id in tqdm(image_ids):
        list_file.write('data/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

 

运行xml_to_txt程序会在labels出现很多txt文件

 以上步骤全部完成后在data新建myvoc.yaml文件

 代码如下

train: data/train.txt
val: data/val.txt
nc: 1
names: ["naok"]

nc为训练标签数量names是标签名

然后修改models下的yolov5s.yaml文件

 最后将train.py中的第72行和389行注释掉

 

 打开终端进入到自己yolov5的工程目录下

输入命令：python train.py --epoch 300 --batch 4 --data ./data/myvoc.yaml --cfg ./models/yolov5s.yaml --weight ./weights/yolov5s.pt --workers 0开始训练

训练完成后找到runs文件下最新的exp文件下的的last.pt文件，将last.pt文件复制粘贴到weights文件下

以上就是yolov5训练数据的全过程