当前位置: 代码网 > 科技>人工智能>数据分析 > 目标检测:yolov5训练自己的数据集,新手小白也能学会训练模型,一看就会

目标检测:yolov5训练自己的数据集,新手小白也能学会训练模型,一看就会

2024年07月31日 数据分析 我要评论
训练自己的数据集分为4部分,先配置环境,再获取制作自己的数据集,然后修改配置训练,最后验证训练结果,可选择将结果进行可视化界面展示。yolov5训练起来较为简单,跟yolov8相差不多,如果有其他目标检测的数据集理论上可以直接拿来用,从第3训练模型开始看,新手小白0基础建议一步一步跟着来,哪里看不懂的或者遇到哪有问题可以评论区交流或者私信问~

目录

1. 环境配置

2. 数据集获取

2.1 网上搜索公开数据集

2.2 自制数据集

2.2.1 labelimg安装

2.2.2 labelimg使用

2.3 数据集转换及划分

2.3.1 数据集voc格式转yolo格式

2.3.2 数据集划分

3. 训练模型

3.1 创建data.yaml

3.2 训练模型

4. 模型测试

5. 可视化界面


训练自己的数据集分为4部分,先配置环境,再获取制作自己的数据集,然后修改配置训练,最后验证训练结果,可选择将结果进行可视化界面展示。yolov5训练起来较为简单,跟yolov8相差不多,如果有其他目标检测的数据集理论上可以直接拿来用,从第3训练模型开始看,新手小白0基础建议一步一步跟着来,哪里看不懂的或者遇到哪有问题可以评论区交流或者私信问~

1. 环境配置

在训练yolov5模型前环境必须配置完成,还不会配置环境的可以看我的这篇博客

深度学习目标检测:yolov5环境配置,适合0基础小白,超详细-csdn博客文章浏览阅读72次,点赞2次,收藏5次。对于小白也很简单的yolov5环境配置https://blog.csdn.net/qq_67105081/article/details/138232424环境配置完验证之后就可以获取数据集

2. 数据集获取

数据集可以使用网上公开的跟自己研究相契合的数据集,或者是搜索/拍摄自己研究所需要的图片进行标注制作成数据集,这里两种方法都详细介绍一下

2.1 网上搜索公开数据集

可以在搜索引擎上搜索或者在公开数据集的网站上搜索关键词,使用到的公开数据集网站为kagglekaggle: your machine learning and data science communitykaggle is the world’s largest data science community with powerful tools and resources to help you achieve your data science goals.https://www.kaggle.com/比如这里做一个安全帽检测的研究,就可以在搜索框输入安全帽的英文(因为是英文网站,都需要翻译成英文后搜索)safety helmet

搜索后就可以找到相关的内容,点击datasets筛选数据集,下载几个看一下数据集是否为目标检测的数据格式,一般文件夹为jpegimages和annotations包含这两个就可以使用

点开一个后点击download即可下载。

若是下载到分割数据集,即json格式的标注可以看我的这篇文章转为txt,其他数据格式转换也可以私聊定制

2.2 自制数据集

自制数据集需要先获取一定数量的目标图片,可以拍摄或者下载,图片足够之后使用标注工具labelimg进行标注

2.2.1 labelimg安装

使用labelimg建议使用python3.10以下的环境,这里创建一个python3.8的虚拟环境,不会创建的可以去看我这篇博客

conda create -n labelimg python=3.8

这里创建完之后进入labelimg环境

conda activate labelimg

进入labelimg环境之后通过pip下载labelimg(需要关闭加速软件)

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

安装完成之后就可以使用

2.2.2 labelimg使用

在使用labelimg之前,需要准备好数据集存放位置,这里推荐创建一个大文件夹为data,里面有jpegimages、annotations和classes.txt,其中jpegimages文件夹里面放所有的图片,annotations文件夹是将会用来对标签文件存放,classes.txt里存放所有的类别,每种一行

classes.txt里存放所有的类别,可以自己起名,需要是英文,如果有空格最好用下划线比如no_hat

上述工作准备好之后,在labelimg环境中cd到data目录下,如果不是在c盘需要先输入其他盘符+:

例如d: 回车之后再输入文件路径,接着输入以下命令打开labelimg

labelimg jpegimages classes.txt

打开软件后可以看到左侧有很多按钮,open dir是选择图片文件夹,上面选过了

点击change save dir 切换到annotations目录之中,点击save下面的图标切换到pascal voc格式

切换好之后点击软件上边的view,将 auto save mode(切换到下一张图会自动保存标签)和display labels(显示标注框和标签) 保持打开状态。

常用快捷键:

a:切换到上一张图片

d:切换到下一张图片

w:调出标注十字架

del :删除标注框

例如,按下w调出标注十字架,标注完成之后选择对应的类别,这张图全部标注完后按d下一张

所有图像标注完成后数据集即制作完成。

2.3 数据集转换及划分

2.3.1 数据集voc格式转yolo格式

暂时还没有数据集可以先跟着本文进行训练。

如何查看自己数据集格式,打开annotations文件夹,如果看到文件后缀为.xml,则为voc格式,如果文件后缀为.txt则为yolo格式,后缀名看不到请搜索 如何显示文件后缀名。yolov5训练需要转为yolo格式训练,转换代码如下:

import xml.etree.elementtree as et
import pickle
import os
from os import listdir, getcwd
from os.path import join


def convert(size, box):
    x_center = (box[0] + box[1]) / 2.0
    y_center = (box[2] + box[3]) / 2.0
    x = x_center / size[0]
    y = y_center / size[1]
    w = (box[1] - box[0]) / size[0]
    h = (box[3] - box[2]) / size[1]
    return (x, y, w, h)


def convert_annotation(xml_files_path, save_txt_files_path, classes):
    xml_files = os.listdir(xml_files_path)
    print(xml_files)
    for xml_name in xml_files:
        print(xml_name)
        xml_file = os.path.join(xml_files_path, xml_name)
        out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')
        out_txt_f = open(out_txt_path, 'w')
        tree = et.parse(xml_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))
            # b=(xmin, xmax, ymin, ymax)
            print(w, h, b)
            bb = convert((w, h), b)
            out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


if __name__ == "__main__":
    # 需要转换的类别,需要一一对应
    classes1 = ['hat','nohat']
    # 2、voc格式的xml标签文件路径
    xml_files1 = r'f:\data\annotations'
    # 3、转化为yolo格式的txt标签文件存储路径
    save_txt_files1 = r'f:\data\labels'

    convert_annotation(xml_files1, save_txt_files1, classes1)

需要自行将类别以及文件路径替换,这里顺序要记住,文件夹也对应替换,不知道有哪些类别可以看这篇博客统计标注文件xml文件类别-csdn博客

2.3.2 数据集划分

训练自己的yolov5检测模型,数据集需要划分为训练集、验证集,这里提供一个参考代码,划分比例为8:2(这里不再划分测试集),也可以按照自己的比例划分修改代码,三者加起来为1即可

import os
import shutil
import random

random.seed(0)


def split_data(file_path,xml_path, new_file_path, train_rate, val_rate, test_rate):
    each_class_image = []
    each_class_label = []
    for image in os.listdir(file_path):
        each_class_image.append(image)
    for label in os.listdir(xml_path):
        each_class_label.append(label)
    data=list(zip(each_class_image,each_class_label))
    total = len(each_class_image)
    random.shuffle(data)
    each_class_image,each_class_label=zip(*data)
    train_images = each_class_image[0:int(train_rate * total)]
    val_images = each_class_image[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_images = each_class_image[int((train_rate + val_rate) * total):]
    train_labels = each_class_label[0:int(train_rate * total)]
    val_labels = each_class_label[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_labels = each_class_label[int((train_rate + val_rate) * total):]

    for image in train_images:
        print(image)
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'train' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in train_labels:
        print(label)
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'train' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

    for image in val_images:
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'val' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in val_labels:
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'val' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

    for image in test_images:
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'test' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in test_labels:
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'test' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)


if __name__ == '__main__':
    file_path = r"f:\data\jpegimages"
    txt_path = r'f:\data\labels'
    new_file_path = r"f:\vocdevkit"
    split_data(file_path,txt_path, new_file_path, train_rate=0.8, val_rate=0.2, test_rate=0.2)

划分完成之后数据集的准备工作就好了。

3. 训练模型

需要下载源码,这里选择的是yolov5 v7.0版本的,最新版本随时更新容易出bug

不会下载源码的可以看我的这篇博客,也可以点击链接获取,我这里把源码和预训练权重yolov5n.pt和yolov5s.pt一起打包上传了。

有了源码之后需要修改里面的参数,导入自己的数据集

3.1 创建data.yaml

在yolov5/data目录下(也就是本文所用的yolov5-7.0目录下)创建一个新的data.yaml文件,也可以是其他名字的例如hat.yaml文件,文件名可以变但是后缀需要为.yaml,内容如下

train: f:/vocdevkit/train/images  # train images (relative to 'path') 128 images
val: f:/vocdevkit/val/images  # val images (relative to 'path') 128 images
 
nc: 2
 
# classes
names: ['hat','nohat']

其他路径和类别自己替换,需要和上面数据集转换那里类别顺序一致。

3.2 训练模型

这是使用官方提供的预训练权重进行训练,使用yolov5s.pt,也可以使用yolov5n.pt,模型大小n<s<m<l<x,训练时长成倍增加。

下载完成之后放入yolov5-7.0根目录中,打开train.py文件,修改方框中框选部分内容

其中weights为导入模型,这里修改为yolov5s.pt或者其他模型,cdg为改为yolov5s.yaml或者pt对应的yaml文件名,data改为data/data.yaml,修改为自己新建的文件名,epochs可以先设为10,能够正常训练并且保存结果再修改为100或者更多,batch-size根据自己的电脑性能设置大小,建议为2的n次方,可以多次修改找到最适合自己电脑的参数大小。

训练过程如图,耐心等待训练完成即可,训练完成后会生成best.pt权重文件,可以用来验证训练效果或者部署到可视化界面之中展示。

训练过程中遇到报错可以评论区留言,看到都会及时回复

4. 模型测试

找到之前训练的结果保存路径,打开一个detect.py文件,修改图中yolov5s.pt为训练好的权重文件路径,如runs/train/exp4/weights/best.pt,然后将待检测图片或视频放入data/images文件夹中

运行后就会得到预测模型结果

可以打开对应路径下查看预测的图片效果,模型就训练好啦~

5. 可视化界面

部分同学的需求是制作出一个可视化图形界面来展示实时预测的效果,yolov5的可视化界面可以私聊定制

(0)

相关文章:

版权声明:本文内容由互联网用户贡献,该文观点仅代表作者本人。本站仅提供信息存储服务,不拥有所有权,不承担相关法律责任。 如发现本站有涉嫌抄袭侵权/违法违规的内容, 请发送邮件至 2386932994@qq.com 举报,一经查实将立刻删除。

发表评论

验证码:
Copyright © 2017-2025  代码网 保留所有权利. 粤ICP备2024248653号
站长QQ:2386932994 | 联系邮箱:2386932994@qq.com