Yolov8训练自己的数据集记录

Yolov8训练自己的数据集记录

前言：仅做个人试验记录！

第一章：数据集制作及划分

一、数据集制作

1.制作软件labelImg
这个不赘述，选择yolo格式即可。
yolo格式为：

Class-id x y w h

数据经过归一化处理后，实际效果如下图：

二、数据集划分

参考：Yolov8训练自己的数据集
此处根据自己的情况进行修改。
在yolov8目录下新建两个文件夹“data”和“yolov8_dataset”,并将标注好的数据放入“data”文件夹下。“data”文件夹路径如下：

data
	├─images
	   └─all
	└─labels
	   └─all

运行脚本，需要修改自己的数据文件夹路径及图像格式等。

import os
import random
import shutil

# 设置随机数种子
random.seed(123)

# 定义文件夹路径
root_dir = 'data'
image_dir = os.path.join(root_dir, 'images', 'all')
label_dir = os.path.join(root_dir, 'labels', 'all')
output_dir = 'yolov8_dataset'

# 定义训练集、验证集和测试集比例
train_ratio = 0.7
valid_ratio = 0.15
test_ratio = 0.15

# 获取所有图像文件和标签文件的文件名（不包括文件扩展名）
image_filenames = [os.path.splitext(f)[0] for f in os.listdir(image_dir)]
label_filenames = [os.path.splitext(f)[0] for f in os.listdir(label_dir)]

# 随机打乱文件名列表
random.shuffle(image_filenames)

# 计算训练集、验证集和测试集的数量
total_count = len(image_filenames)
train_count = int(total_count * train_ratio)
valid_count = int(total_count * valid_ratio)
test_count = total_count - train_count - valid_count

# 定义输出文件夹路径
train_image_dir = os.path.join(output_dir, 'train', 'images')
train_label_dir = os.path.join(output_dir, 'train', 'labels')
valid_image_dir = os.path.join(output_dir, 'valid', 'images')
valid_label_dir = os.path.join(output_dir, 'valid', 'labels')
test_image_dir = os.path.join(output_dir, 'test', 'images')
test_label_dir = os.path.join(output_dir, 'test', 'labels')

# 创建输出文件夹
os.makedirs(train_image_dir, exist_ok=True)
os.makedirs(train_label_dir, exist_ok=True)
os.makedirs(valid_image_dir, exist_ok=True)
os.makedirs(valid_label_dir, exist_ok=True)
os.makedirs(test_image_dir, exist_ok=True)
os.makedirs(test_label_dir, exist_ok=True)

# 将图像和标签文件划分到不同的数据集中
for i, filename in enumerate(image_filenames):
    if i < train_count:
        output_image_dir = train_image_dir
        output_label_dir = train_label_dir
    elif i < train_count + valid_count:
        output_image_dir = valid_image_dir
        output_label_dir = valid_label_dir
    else:
        output_image_dir = test_image_dir
        output_label_dir = test_label_dir

    # 复制图像文件
    src_image_path = os.path.join(image_dir, filename + '.jpeg')
    dst_image_path = os.path.join(output_image_dir, filename + '.jpeg')
    shutil.copy(src_image_path, dst_image_path)

    # 复制标签文件
    src_label_path = os.path.join(label_dir, filename + '.txt')
    dst_label_path = os.path.join(output_label_dir, filename + '.txt')
    shutil.copy(src_label_path, dst_label_path)

第二章：环境配置

这一部分参考官网配置即可。

第三章：模型训练

一、数据集配置文件

在yolov8_dataset文件夹下新建data.yaml配置文件，并按自己的路径和类别写入：

train: /home/cj/chaintwork/yolov8/yolov8_dataset/train  # train images (relative to 'path') 128 images
val: /home/cj/chaintwork/yolov8/yolov8_dataset/valid # val images (relative to 'path') 128 images
test: /home/cj/chaintwork/yolov8/yolov8_dataset/test # test images (optional)

# Classes
names:
  0: a
  1: b

二、构建训练脚本

在yolov8目录下比新建train.py文件，并写入：

from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n.yaml')  # build a new model from YAML
model = YOLO('yolov8n.pt')  # load a pretrained model (recommended for training)
model = YOLO('yolov8n.yaml').load('yolov8n.pt')  # build from YAML and transfer weights

# Train the model
model.train(data='/home/cj/chaintwork/yolov8/yolov8_dataset/data.yaml', epochs=10, imgsz=640)

运行代码即可开始训练。

三、推理脚本

from ultralytics import YOLO
from PIL import Image
import cv2

model = YOLO("/home/cj/chaintwork/yolov8/runs/detect/train4/weights/best.pt")

# accepts all formats - image/dir/Path/URL/video/PIL/ndarray. 0 for webcam
#results = model.predict(source="0")
#results = model.predict(source="folder", show=True) # Display preds. Accepts all YOLO predict arguments

# from PIL
#im1 = Image.open("/home/cj/chaintwork/yolov8/001.jpeg")
#results = model.predict(source=im1, save=True)  # save plotted images

# from ndarray
im2 = cv2.imread("/home/cj/chaintwork/yolov8/001.jpeg")
results = model.predict(source=im2, save=True, save_txt=True)  # save predictions as labels