yolov5训练自己的数据集

1. 获取自己的数据集

1.1 爬虫

先爬虫爬取自己想要的图片：

import re
import requests
from urllib import error
from bs4 import BeautifulSoup
import os


num = 0
numPicture = 0
file = ''
List = []


def Find(url, A):
    global List
    print('正在检测图片总数，请稍等.....')
    t = 0
    i = 1
    s = 0
    while t < 1000:
        Url = url + str(t)
        try:
            # 这里搞了下
            Result = A.get(Url, timeout=7, allow_redirects=False)
        except BaseException:
            t = t + 60
            continue
        else:
            result = Result.text
            pic_url = re.findall('"objURL":"(.*?)",', result, re.S)  # 先利用正则表达式找到图片url
            s += len(pic_url)
            if len(pic_url) == 0:
                break
            else:
                List.append(pic_url)
                t = t + 60
    return s


def recommend(url):
    Re = []
    try:
        html = requests.get(url, allow_redirects=False)
    except error.HTTPError as e:
        return
    else:
        html.encoding = 'utf-8'
        bsObj = BeautifulSoup(html.text, 'html.parser')
        div = bsObj.find('div', id='topRS')
        if div is not None:
            listA = div.findAll('a')
            for i in listA:
                if i is not None:
                    Re.append(i.get_text())
        return Re


def dowmloadPicture(html, keyword):
    global num
    # t =0
    pic_url = re.findall('"objURL":"(.*?)",', html, re.S)  # 先利用正则表达式找到图片url
    print('找到关键词:' + keyword + '的图片，即将开始下载图片...')
    for each in pic_url:
        print('正在下载第' + str(num + 1) + '张图片，图片地址:' + str(each))
        try:
            if each is not None:
                pic = requests.get(each, timeout=7)
            else:
                continue
        except BaseException:
            print('错误，当前图片无法下载')
            continue
        else:
            string = file + r'\\' + keyword + '_' + str(num) + '.jpg'
            fp = open(string, 'wb')
            fp.write(pic.content)
            fp.close()
            num += 1
        if num >= numPicture:
            return


if __name__ == '__main__':  # 主函数入口

##############################
    # 这里加了点
    headers = {
        'Accept-Language': 'zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2',
        'Connection': 'keep-alive',
        'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:60.0) Gecko/20100101 Firefox/60.0',
        'Upgrade-Insecure-Requests': '1'
    }

    A = requests.Session()
    A.headers = headers
###############################

    word = input("请输入搜索关键词(可以是人名，地名等): ")
    # add = 'http://image.baidu.com/search/flip?tn=baiduimage&ie=utf-8&word=%E5%BC%A0%E5%A4%A9%E7%88%B1&pn=120'
    url = 'https://image.baidu.com/search/flip?tn=baiduimage&ie=utf-8&word=' + word + '&pn='

    # 这里搞了下
    tot = Find(url, A)
    Recommend = recommend(url)  # 记录相关推荐
    print('经过检测%s类图片共有%d张' % (word, tot))
    numPicture = int(input('请输入想要下载的图片数量 '))
    file = input('请建立一个存储图片的文件夹，输入文件夹名称即可')
    y = os.path.exists(file)
    if y == 1:
        print('该文件已存在，请重新输入')
        file = input('请建立一个存储图片的文件夹，)输入文件夹名称即可')
        os.mkdir(file)
    else:
        os.mkdir(file)
    t = 0
    tmp = url
    while t < numPicture:
        try:
            url = tmp + str(t)

            # 这里搞了下
            result = A.get(url, timeout=10, allow_redirects=False)
        except error.HTTPError as e:
            print('网络错误，请调整网络后重试')
            t = t + 60
        else:
            dowmloadPicture(result.text, word)
            t = t + 60

    print('当前搜索结束')
    print('猜你喜欢')
    for re in Recommend:
        print(re, end='  ')

1.2 标注图片

用labelimg标注图片中要检测的东西：
https://blog.csdn.net/Cupid_kl/article/details/135216888

2. 下载安装yolov5

见另一篇博文：
https://blog.csdn.net/Cupid_kl/article/details/135284546
如果还没有配置神经网络的环境，则先看博文：
https://blog.csdn.net/Cupid_kl/article/details/133428337
安装后，运行detect.py，如果成功运行，则安装完毕。

3. 处理数据集

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

在yolo的文件夹内创建文件夹VOCData放置数据：

在VOCData中新建文件夹Annotations，放置.xml的标签文件。新建文件夹images，放置图片。
之后在目录中创建split_train_val.py，划分数据集：

import os
import random
import argparse

parser = argparse.ArgumentParser()
# xml文件的地址，根据自己的数据进行修改 xml一般存放在Annotations下，注意以下为相对路径
parser.add_argument('--xml_path', default='Annotations', type=str, help='input xml label path')
# 数据集的划分，地址选择自己数据下的ImageSets/Main，注意以下为相对路径
parser.add_argument('--txt_path', default='ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()

trainval_percent = 0.8  # 训练集和验证集所占比例。 这里没有划分测试集
train_percent = 0.9  # 训练集所占比例，可自己进行调整
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()

运行后目录中多了一个文件夹ImageSets。

3.2 XML格式转yolo_txt格式

在目录中创建text_to_yolo.py，划分数据集：

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

sets = ['train', 'val', 'test']
classes = ["bird"]  # 改为自己的类别
abs_path = os.getcwd()
print(abs_path)


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):
    in_file = open('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/Annotations/%s.xml' % (image_id), encoding='UTF-8')
    out_file = open('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/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
        # 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')


wd = getcwd()
for image_set in sets:
    # 这里是绝对路径，需要根据自己的情况修改
    if not os.path.exists('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/labels/'):
        os.makedirs('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/labels/')
    image_ids = open('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()

    if not os.path.exists('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/dataSet_path/'):
        os.makedirs('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/dataSet_path/')

    list_file = open('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/dataSet_path/%s.txt' % (image_set), 'w')
    for image_id in image_ids:
        list_file.write('D:/JetBrains/pythonProject/yolov5-master/yolov5-master/VOCData/images/%s.jpg\n' % (image_id))
        convert_annotation(image_id)
    list_file.close()

这边代码中的目录需要做对应修改。

4. 训练数据集

（1）打开Data文件夹：

在其中添加一个myvoc.yaml文件，记事本打开，写入：

其中，train和val后面的内容就是上文中分配的数据集目录；
nc表示类别的数量；
names表示类别的名字，如果有多个，就都写入[ ]中，用逗号隔开。
（冒号后面一定要有一个空格）

（2）在pycharm中打开yolo的项目，打开yolov5s.yaml
把nc后面的数字改为自己的类别数目，其余不变。

（3）打开train.py，改变参数：

epochs和batch-size按照自己的需要设置，刚开始试运行的时候，可以调整的低一点。

（4）训练结束之后，可以看到目录中runs/train中有跟新的结果，我这边是exp5

（5）打开detect.py，修改参数：
weights：对应路径为自己的训练结果。
source：输入源的目录，可以是图像或视频文件。此处0表示打开默认摄像头。
data：修改为自己的数据位置。
conf-thres：置信度阈值，只有置信度高于此阈值的对象才会被检测出来，可以根据自己的结果修改。
iou-thres：非极大值抑制（NMS）的交并比（IoU）阈值。用于在重叠较大的候选框中选择最佳的检测结果。

5. 查看训练

5.1 在文件夹中查看

训练结果在目录文件的runs中

5.2 在tensorboard中查看

conda激活yolov5所在的环境，命令参考：
https://blog.csdn.net/Cupid_kl/article/details/131754832
通过“pip install tensorboard”安装
之后启动并查看yolo训练的结果：

tensorboard --logdir=训练结果保存路径

#指定到YOLOv5训练的根目录，里面有包含每一次训练的结果
tensorboard --logdir=D:\JetBrains\pythonProject\yolov5-master\yolov5-master\runs\train\
 
#指定到YOLOv5训练的根目录下面的某一次训练的结果
tensorboard --logdir=D:\JetBrains\pythonProject\yolov5-master\yolov5-master\runs\train\exp5

将 http://localhost:6006/复制到浏览器，打开面板即可看见训练结果。