制作自己的数据集并训练的YOLOv8模型

1.制作数据集

1.1 下载安装labelImg

LabelImg源码

下载好源码并 cd到源码路径下
安装需要的环境

# 安装lxml
pip install lxml
 
# 安装pyqt5
pip install pyqt5

# 将"resources.qrc" 文件编译成 Python 代码，并将输出到 "libs/resources.py" 文件中
pyrcc5 -o libs/resources.py resources.qrc

# 启动labelImg
python labelimg.py

1.2 开始制作数据集

打开data/predefined_classes.txt，修改默认类别，也可以直接清空，后面打标签的时候会自动生成，我这里将这个文件清空。

点击“创建区块” 开始打标签

准备训练

可以参考ultralytics官方安装教程

# 创建一个用于训练的虚拟环境
conda create --name YOLOv8_ultralytics python=3.8

#启动训练环境
conda activate YOLOv8_ultralytics

PyTorch 要求因操作系统和 CUDA 要求而异，因此建议先按照 https://pytorch.org/get-started/locally 中的说明安装 PyTorch。

#我这里cuda是11.6
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.6 -c pytorch -c nvidia

#安装ultralytics
pip install ultralytics

分配数据集

1. 新建文件夹 datasets/data/images
   然后将图片和标签全部放到这个文件夹下。

2. 在datasets下新建一个分配数据集数据的脚本文件Process.py

import os
import random
import shutil

def split_dataset(data_dir,train_val_test_dir, train_ratio, val_ratio, test_ratio):
    # 创建目标文件夹
    train_dir = os.path.join(train_val_test_dir, 'train')
    val_dir = os.path.join(train_val_test_dir, 'val')
    test_dir = os.path.join(train_val_test_dir, 'test')
    os.makedirs(train_dir, exist_ok=True)
    os.makedirs(val_dir, exist_ok=True)
    os.makedirs(test_dir, exist_ok=True)

    # 获取数据集中的所有文件
    files = os.listdir(data_dir)

    # 过滤掉非图片文件
    image_files = [f for f in files if f.endswith('.jpg') or f.endswith('.png')]
    # 随机打乱文件列表
    random.shuffle(image_files)

    # 计算切分数据集的索引
    num_files = len(image_files)
    num_train = int(num_files * train_ratio)
    num_val = int(num_files * val_ratio)
    num_test = num_files - num_train - num_val

    # 分离训练集
    train_files = image_files[:num_train]
    for file in train_files:
        src_image_path = os.path.join(data_dir, file)
        src_label_path = os.path.join(data_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
        dst_image_path = os.path.join(train_dir, file)
        dst_label_path = os.path.join(train_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
        shutil.copy(src_image_path, dst_image_path)
        shutil.copy(src_label_path, dst_label_path)

    # 分离验证集
    val_files = image_files[num_train:num_train+num_val]
    for file in val_files:
        src_image_path = os.path.join(data_dir, file)
        src_label_path = os.path.join(data_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
        dst_image_path = os.path.join(val_dir, file)
        dst_label_path = os.path.join(val_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
        shutil.copy(src_image_path, dst_image_path)
        shutil.copy(src_label_path, dst_label_path)

    # 分离测试集
    test_files = image_files[num_train+num_val:]
    for file in test_files:
        src_image_path = os.path.join(data_dir, file)
        src_label_path = os.path.join(data_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
        dst_image_path = os.path.join(test_dir, file)
        dst_label_path = os.path.join(test_dir, file.replace('.jpg', '.txt').replace('.png', '.txt'))
        shutil.copy(src_image_path, dst_image_path)
        shutil.copy(src_label_path, dst_label_path)

    print("数据集分离完成！")
    print(f"训练集数量：{len(train_files)}")
    print(f"验证集数量：{len(val_files)}")
    print(f"测试集数量：{len(test_files)}")

def move_files(data_dir):
    # 创建目标文件夹
    images_dir = os.path.join(data_dir, 'images')
    labels_dir = os.path.join(data_dir, 'labels')
    os.makedirs(images_dir, exist_ok=True)
    os.makedirs(labels_dir, exist_ok=True)

    # 获取数据集中的所有文件
    files = os.listdir(data_dir)

    # 移动PNG文件到images文件夹
    png_files = [f for f in files if f.endswith('.png')]
    for file in png_files:
        src_path = os.path.join(data_dir, file)
        dst_path = os.path.join(images_dir, file)
        shutil.move(src_path, dst_path)

    # 移动TXT文件到labels文件夹
    txt_files = [f for f in files if f.endswith('.txt')]
    for file in txt_files:
        src_path = os.path.join(data_dir, file)
        dst_path = os.path.join(labels_dir, file)
        shutil.move(src_path, dst_path)

    print(f"{data_dir}文件移动完成！")
    print(f"总共移动了 {len(png_files)} 个PNG文件到images文件夹")
    print(f"总共移动了 {len(txt_files)} 个TXT文件到labels文件夹")


# 设置数据集路径和切分比例
data_dir = './data/images'      # 图片和标签路径
train_val_test_dir= './data'    # 目标文件夹
train_ratio = 0.7               # 训练集比例
val_ratio = 0.2                 # 验证集比例
test_ratio = 0.1                # 测试集比例

# 调用函数分离数据集
split_dataset(data_dir, train_val_test_dir,train_ratio, val_ratio, test_ratio)
# 调用函数移动文件
move_files(os.path.join(train_val_test_dir, 'train'))
move_files(os.path.join(train_val_test_dir, 'val'))
move_files(os.path.join(train_val_test_dir, 'test'))

3. 在datasets下运行这个脚本，然后生成这样的文件结构
   
4. 在datasets同级目录下新建一个mydata.yaml文件

train: ./data/train/images
val: ./data/val/images
test: ./data/test/images

# 类别数
nc: 4

# 类别名称
names: ['CCTV2','fivestar','CCTV3',"CCTV1"]

训练

在datasets同级路径下新建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='./mydata.yaml', epochs=100, imgsz=640)

测试

在datasets同级路径下新建test.py测试脚本

from PIL import Image
from ultralytics import YOLO

# Load a pretrained YOLOv8n model
model = YOLO('./runs/detect/train/weights/last.pt')

# Run inference on 'bus.jpg'
results = model('./test1.png')  # results list

# Show the results
for r in results:
    im_array = r.plot()  # plot a BGR numpy array of predictions
    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL image
    im.save('results.jpg')  # save image