YOLOv8制作自定义数据集并训练

东东不熬夜

已于 2024-02-15 20:09:35 修改

阅读量3.1k

点赞数 22

分类专栏：机器人 # YOLOv8 文章标签： YOLO python 人工智能深度学习

于 2024-02-12 22:03:40 首次发布

本文链接：https://blog.csdn.net/lyddsg520/article/details/136098019

版权

机器人同时被 2 个专栏收录

13 篇文章 3 订阅

订阅专栏

YOLOv8

4 篇文章 1 订阅

订阅专栏

前言

本文中的数据集以制作一个鼠标和手机的自定义为例，大家在制作自定义数据集时方法类似。

一、制作自定义数据集

1、建立相应文件夹

在桌面建立一个YOLO_phonemouse（该名字自定义，自由设置即可），点击打开,在其中新建两个文件夹分别命名为images和labels，再新建一个文本文件命名为classes，里面写上需要标记的类的名称，博主要标记手机和鼠标，所以在里面写上phone、mouse
YOLO_phonemouse的目录结构如下所示：

YOLO_phonemouse
├─ images
├─ labels
├─ classes.txt

2、下载图片

在网络上下载关于鼠标和手机的照片，下载的照片放在YOLO_phonemouse的images文件夹中
注：下载的时候要把下载名称的后缀都改为.jpg
在这里插入图片描述

3、为图片打标签

（1）安装labelimg

用快捷键：win+r，输入cmd，进入cmd命令行控制台。输入如下的命令：

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

（2）打开labelimg

用快捷键：win+r，输入cmd，进入cmd命令行控制台。输入如下命令：

cd "C:\Users\86139\Desktop\YOLO_phonemouse"  //这里的路径根据自己文件夹所在的位置进行修改

在这里插入图片描述
输入以下命令：

labelimg images classes.txt

指打开labelimg工具，打开imagines文件夹，初始化classes.txt里面定义的类。

（3）标记图片

打开labelimg后，labelimg中的一些设置、标注方法和快捷键可以参考以下文章的3.2之后的部分：
目标检测—利用labelimg制作自己的深度学习目标检测数据集

在这里插入图片描述

标注完成后，打开labels文件夹，可以看到里面都是txt文件。

在这里插入图片描述

到这里，自定义的数据集就制作好了。

二、按比例移动自定义数据集中的内容

在pycharm中运行以下代码即可，需要根据自己的情况修改相应文件的路径。

# -*- coding: gbk -*-
import os
import random
import shutil
from tqdm import tqdm

# 数据集images 文件夹路径
image_dir = "C:/Users/86139/Desktop/YOLO_phonemouse/images"
# 数据集labels 文件夹路径
label_dir = "C:/Users/86139/Desktop/YOLO_phonemouse/labels"

# 创建一个空列表来存储有效图片的路径
valid_images = []
# 创建一个空列表来存储有效 label 的路径
valid_labels = []

# 遍历 images 文件夹下的所有图片
for image_name in os.listdir(image_dir):
    # 获取图片的完整路径
    image_path = os.path.join(image_dir, image_name)
    # 获取图片文件的扩展名
    ext = os.path.splitext(image_name)[-1]
    # 根据扩展名替换成对应的 label 文件名
    label_name = image_name.replace(ext, ".txt")
    # 获取对应 label 的完整路径
    label_path = os.path.join(label_dir, label_name)
    # 判断 label 是否存在
    if not os.path.exists(label_path):
        # 删除图片
        os.remove(image_path)
        print("Deleted:", image_path)
    else:
        # 将图片路径添加到列表中
        valid_images.append(image_path)
        # 将 label 路径添加到列表中
        valid_labels.append(label_path)

# 遍历每个有效图片路径
for i in tqdm(range(len(valid_images))):
    image_path = valid_images[i]
    label_path = valid_labels[i]
    # 随机生成一个概率
    r = random.random()
    # 判断图片应该移动到哪个文件夹
    # train：valid：test = 7:2:1
    if r < 0.1:
        # 移动到 test 文件夹
        destination = "E:/Robot/YOLO/datasets/phoneandmouse/test"
    elif r < 0.2:
        # 移动到 valid 文件夹
        destination = "E:/Robot/YOLO/datasets/phoneandmouse/valid"
    else:
        # 移动到 train 文件夹
        destination = "E:/Robot/YOLO/datasets/phoneandmouse/train"
    # 创建目标文件夹中 images 和 labels 子文件夹
    os.makedirs(os.path.join(destination, "images"), exist_ok=True)
    os.makedirs(os.path.join(destination, "labels"), exist_ok=True)
    # 生成目标文件夹中图片的新路径
    image_destination_path = os.path.join(destination, "images", os.path.basename(image_path))
    # 移动图片到目标文件夹
    shutil.move(image_path, image_destination_path)
    # 生成目标文件夹中 label 的新路径
    label_destination_path = os.path.join(destination, "labels", os.path.basename(label_path))
    # 移动 label 到目标文件夹
    shutil.move(label_path, label_destination_path)

# 输出有效图片路径列表
print("Valid images:", valid_images)
# 输出有效 label 路径列表
print("Valid labels:", valid_labels)

三、建立数据集测试、训练、验证的的文件夹

在所在的项目的文件中新建一个文件夹名称为datasets，点击进入datasets文件夹，在datasets中新建一个文件夹，名称自定义，这个文件夹即为自己制作的一个数据集，博主要制作一个鼠标和手机的数据集，这里就命名为phoneandmouse了，打开这个文件夹，在其中再新建3个文件夹如下图所示
在这里插入图片描述
再在每个文件夹中再新建两个文件夹，datasets的目录结构如下所示:

datasets
├─ phoneandmouse
│    ├─ test
│    │    ├─ images
│    │    ├─  labels
│    ├─ train
│    │    ├─ images
│    │    ├─  labels
│    ├─ valid
│    │    ├─ images
│    │    ├─  labels

在这里插入图片描述

四、在文件夹中添加yaml文件

1、下载yolov8项目

注：如果使用的模型是yolov8官方给的模型执行这一步，如果用的自己的模型不用执行这一步
具体方法可以参考以下文章中的第一点内容：
YOLOv8推理使用（指令讲解+报错解决方法）
下载好之后在pycharm中打开该项目

2、添加模型配置文件

博主以使用yolov8n模型为例
先找到yolov8模型所在的文件夹，博主的路径为ultralytics/cfg/models/v8

在这里插入图片描述

在该文件夹下找到使用的模型（博主使用的yolov8n模型包含在yolov8中，所以找到yolov8.yaml即可），将选择使用的模型的yaml文件复制粘贴到datasets文件夹下

在这里插入图片描述

3、添加数据集配置文件

在phoneandmouse中新建一个data.yaml文件，打开后将以下内容复制进去并根据注释内容进行相应的修改，一般只要修改path、nc、names

train: E:/Robot/YOLO/datasets/phoneandmouse/train/images   #修改成train中images的绝对路径
val: E:/Robot/YOLO/datasets/phoneandmouse/valid/images     #修改成valid中images的绝对路径
test: E:/Robot/YOLO/datasets/phoneandmouse/test/images     #修改成test中images的绝对路径

nc: 2                                          #类的数目
names:                                        #每个类对应的名称
  0: phone
  1: mouse

注：配置文件中所有的冒号：后面都必须要有一个空格，注意是所有冒号，只要有冒号，后面就必须要有一个空格，否则在后面训练的时候会报错，path、train、valid后面必须是绝对路径

五、训练、验证数据集

在pycharm新建一个python，输入以下代码：

# -*- coding: utf-8 -*-
from ultralytics import YOLO
import torch

if __name__ == '__main__':
    # 加载预训练模型，在该模型基础上，训练目标检测的模型
    model=YOLO('yolov8n.pt')

    #训练自定义数据集，数据配置保存在data.yaml中
    model.train(data='datasets/phoneandmouse/data.yaml',epochs=500,batch=2,workers=0)

    #使用验证集验证效果
    model.val()