人工智能图像识别分析之——Yolov5模型训练

上一课讲述了Yolov5模型环境搭建的过程

这一课讲Yolov5模型训练的过程

进行模型训练前，首先要先进行样本标注，标注后产生标注文件，将图片源文件和标注文件进行文件划分，本文以2000张负样本进行训练。

1.新建三级目录datasets/images/train、datasets/images/val

2.新建三级目录datasets/labels/train、datasets/labels/val，存放标注文件（VOC、XML、txt）格式

注意：二级目录images下面存放图片源文件，80%负样本存放train中用于训练，20%正样本图片存放val中用于验证。二级目录label中文件夹train和val分别和二级目录image文件夹下的train和val相对应，必须一致，否则会报错找不到标注文件。

3.将图片源文件和标注文件按照1，2要求分别划分存放,，我是按照xml格式进行标注保存的，当然可以进行格式转换，后面会讲到。

4.数据集划分号以后，修改配置文件

1）复制源码文件夹下/data/coco128.yaml文件到源码文件夹下
2）修改coco128.yaml文件名为yolov5_test.yaml

除path,train,val,nc,names以外全部删除
path:#数据集(原图)的根目录
train:#填写标注的原图相对路径
val:#和train相同
nc:标签数量
names:数组,里面是按照标注输出文件夹内xml格式

5.修改train.py 中--data,--epochs,--batch-size

找到def parse_opt(known=False)函数，修改parse_opt函数内的:
1) parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')#配置文件,改成刚刚修改过的yolov5_test.yaml文件
2)parser.add_argument('--epochs', type=int, default=5, help='total training epochs')#迭代次数
3)parser.add_argument('--batch-size', type=int, default=1, help='total batch size for all GPUs, -1 for autobatch')#训练时占用的内存

也可以在命令行进行配置

注意：以上三个配置由GPU性能决定，如果GPU性能比较好的话，可以设置--epochs=200，--batch-size=16。

6.开始训练

python train.py --data yolov5_test.yaml --weights weights/yolov5s.pt --batch-size 1 --epochs 5 --cache --name yolov5_test

1报错：All git commands will error until this is rectified.解决办法：在train.py中加入os.environ["GIT_PYTHON_REFRESH"] = "quiet"

2报错：AssertionError: train: No labels found in F:\RGZN\Module\datasets\labels\train.cache, can not start training. See https://docs.ultralytics.com/yolov5/tutorials/train_custom_data

解决办法：提示没有找到标签。

将xml格式转换成txt文件

转换脚本如下：

import xml.etree.ElementTree as ET
 
import pickle
import os
from os import listdir , getcwd
from os.path import join
import glob
 
classes = ["021_td_hj_pfw_dxhzgtsxgyw"]
 
def convert(size, box):
 
    dw = 1.0/size[0]
    dh = 1.0/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_annotation(image_name):
    in_file = open('./val/'+image_name[:-3]+'xml') #xml文件路径
    out_file = open('./val/'+image_name[:-3]+'txt', 'w') #转换后的txt文件存放路径
    f = open('./val/'+image_name[:-3]+'xml')
    xml_text = f.read()
    root = ET.fromstring(xml_text)
    f.close()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

 
 
 
 
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classes:
            print(cls)
            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')
 
wd = getcwd()
 
if __name__ == '__main__':
 
    for image_path in glob.glob("F:\\RGZN\\Module\\datasets\\images\\val\\*.jpg"): #每一张图片都对应一个xml文件这里写xml对应的图片的路径
        image_name = image_path.split('\\')[-1]
        convert_annotation(image_name)

训练过程：

等到数据训练好了以后，就会在runs/train/exp/weights目录下会产生两个权重文件，一个是最后一轮的权重文件，一个是最好的权重文件。我们要利用这个最好的权重文件来做推理测试。