yolov5训练自己的数据集(一文搞定训练)

目录

  • 1 yolov5
    • 1.1 环境配置
    • 1.2 测试代码
  • 2 数据集制作
    • 2.1 官网数据集结构
    • 2.2 训练数据集制作
  • 3 开始训练
  • 4 测试效果
  • 5 补充训练过程的可视化

1 yolov5

下载yolov5:github链接。

1.1 环境配置

下载好yolov5文件后,cd到文件路径创建yolov5的环境,终端输入conda create -n yolov5 python==3.7,进入环境conda activate yolov5,安装requirements中的环境pip install -r requirements.txt

1.2 测试代码

按照需求更改测试代码中的参数,这里我直接测试几张照片。
yolov5训练自己的数据集(一文搞定训练)_第1张图片

python detect.py

输出结果保存在runs文件夹下
测试成功,说明环境建立正确,接着制作自己的数据集进行训练。

2 数据集制作

2.1 官网数据集结构

我们可以看下官方给出的训练数据的传入方式,有两种,如图,第一种直接将训练文件的路径写入txt文件传入。第二种直接传入训练文件所在文件夹。

第一种:
yolov5训练自己的数据集(一文搞定训练)_第2张图片第二种:
yolov5训练自己的数据集(一文搞定训练)_第3张图片这里我们按照第二种建立数据集,官网提供了数据集的tree:
yolov5训练自己的数据集(一文搞定训练)_第4张图片
也就是下面的树状结构:

mytrain
├── mycoco
│   ├── images
│   │   ├── train
│   │   └── val
│   └── labels
│       ├── train
│       └── val
└── yolov5

2.2 训练数据集制作

首先建立一个自己的数据文件夹:mycoco
在这里插入图片描述

其目录结构如下:(剩下的文件架构代码生成)

mycoco
├── all_images
├── all_xml
├── make_txt.py
└── train_val.py

其中all_images文件夹下放置所有图片,all_xml文件夹下放置所有与之对应的xml文件。

make_txt.py文件是用来划分数据集使用,内容如下:

import os
import random
trainval_percent = 0.1  
train_percent = 0.9    
xmlfilepath = 'all_images'
txtsavepath = 'ImageSets'
total_xml = os.listdir(xmlfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list, tv) #从所有list中返回tv个数量的项目
train = random.sample(trainval, tr)
if not os.path.exists('ImageSets/'):
    os.makedirs('ImageSets/')
ftrainval = open('ImageSets/trainval.txt', 'w')
ftest = open('ImageSets/test.txt', 'w')
ftrain = open('ImageSets/train.txt', 'w')
fval = open('ImageSets/val.txt', 'w')
for i in list:
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftest.write(name)
        else:
            fval.write(name)
    else:
        ftrain.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

划分完的数据分布为:

├── train                  占90%
└── trainval               占10%
    ├── test               占90%*10%
    └── val                占10%*10%

该目录终端下运行:

python make_txt.py

运行结果,生成四个只包含图片名称的txt文件:
yolov5训练自己的数据集(一文搞定训练)_第5张图片

接着运行train_val.py,该文件一方面将all_xml中xml文件转为txt文件存于all_labels文件夹中,另一方面生成训练所需数据存放架构。(这里如果你的数据直接是txt的标签的话将标签转化的功能注释掉即可)代码如下:

import xml.etree.ElementTree as ET
import pickle
import os
import shutil
from os import listdir, getcwd
from os.path import join
sets = ['train', 'trainval']
classes = ['car','chemicals vehicle','truck','bus','triangle warning sign','warning sign','warning slogan']
def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / 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_id):
    in_file = open('all_xml/%s.xml' % (image_id))
    out_file = open('all_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.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')
wd = getcwd()
print(wd)
for image_set in sets:
    if not os.path.exists('all_labels/'):
        os.makedirs('all_labels/')
    image_ids = open('ImageSets/%s.txt' % (image_set)).read().strip().split()
    image_list_file = open('images_%s.txt' % (image_set), 'w')
    labels_list_file=open('labels_%s.txt'%(image_set),'w')
    for image_id in image_ids:
        image_list_file.write('%s.jpg\n' % (image_id))
        labels_list_file.write('%s.txt\n'%(image_id))
        convert_annotation(image_id) #如果标签已经是txt格式,将此行注释掉,所有的txt存放到all_labels文件夹。
    image_list_file.close()
    labels_list_file.close()


def copy_file(new_path,path_txt,search_path):#参数1:存放新文件的位置  参数2:为上一步建立好的train,val训练数据的路径txt文件  参数3:为搜索的文件位置
    if not os.path.exists(new_path):
        os.makedirs(new_path)
    with open(path_txt, 'r') as lines:
        filenames_to_copy = set(line.rstrip() for line in lines)
        # print('filenames_to_copy:',filenames_to_copy)
        # print(len(filenames_to_copy))
    for root, _, filenames in os.walk(search_path):
        # print('root',root)
        # print(_)
        # print(filenames)
        for filename in filenames:
            if filename in filenames_to_copy:
                shutil.copy(os.path.join(root, filename), new_path)

#按照划分好的训练文件的路径搜索目标,并将其复制到yolo格式下的新路径
copy_file('./images/train/','./images_train.txt','./all_images')
copy_file('./images/val/','./images_trainval.txt','./all_images')
copy_file('./labels/train/','./labels_train.txt','./all_labels')
copy_file('./labels/val/','./labels_trainval.txt','./all_labels')

该目录终端下运行:

python train_val.py

运行结果:

mytrain
├── mycoco
│   ├── all_images
│   ├── all_labels
│   ├── all_xml
│   ├── ImageSets
│   │   ├── train.txt
│   │   ├── test.txt
│   │   ├── trainval.txt
│   │   └── val.txt
│   ├── images
│   │   ├── train
│   │   └── val
│   ├── labels
│        ├── train
│        └── val
│   ├── images_train.txt
│   ├── images_trainval.txt
│   ├── labels_train.txt
│   ├── labels_trainval.txt
│   ├── make_txt.py
│   └── train_val.py
└── yolov5

至此数据集已经完全建好。将制作好的mycoco文件夹与下载好的yolov5文件夹放入同一级文件夹中:
yolov5训练自己的数据集(一文搞定训练)_第6张图片

接着按照yolov5-master/data/coco128.yaml文件,制作mycoco.yaml文件(与coco128.yaml文件同目录):

# Default dataset location is next to /yolov5:
#   /parent_folder
#     /mycoco
#     /yolov5


# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: ../mycoco/images/train/  
val: ../mycoco/images/val/  

# number of classes
nc: 7

# class names
names: ['car','chemicals vehicle','truck','bus','triangle warning sign','warning sign','warning slogan' ]

3 开始训练

接着打开train.py文件,train.py文件详解,更改相应的参数(预训练模型,训练的数据集),这里使用yolov5s.pt为预训练模型,更改yolov5s.yaml文件中的参数(自己的类别数)。
yolov5训练自己的数据集(一文搞定训练)_第7张图片

yolov5训练自己的数据集(一文搞定训练)_第8张图片然后就可以开始训练了。

yolov5训练自己的数据集(一文搞定训练)_第9张图片

4 测试效果

训练完以后可以看到在runs/train/exp/weights下生成了训练好的权重文件best.pt和last.pt,接着我们用训练好的权重文件进行测试,打开detect.py文件
修改权重文件路径和输入测试文件,然后run
在这里插入图片描述runs/detect/exp下我们看看测试效果(这里为了走流程,我的训练参数设置的简单,所以精度肯定不行)

5 补充训练过程的可视化

首先在该环境的终端下安装工具:

pip install wandb

终端下输入:

然后训练,如果报错:wandb.errors.error.UsageError: api_key not configured (no-tty). Run wandb login
那么需要终端初始化以下:

wandb init

yolov5训练自己的数据集(一文搞定训练)_第10张图片按照提示复制api
然后启动重新训练,我们可以再下面的连接中观察训练的动态信息:

yolov5训练自己的数据集(一文搞定训练)_第11张图片
yolov5训练自己的数据集(一文搞定训练)_第12张图片或者直接上https://wandb.ai/home观看自己的训练。
关于tensorboard的的使用可以看Github.
关于yolo的参数设置可以参看link.
关于训练可视化可参看link.

你可能感兴趣的:(环境配置,深度学习,神经网络,python)