yolo训练自己的模型（附完整代码）

用python写的，期间遇到的问题解决方法在另一篇博文里，需要的话可以自行查看

1、准备图像集合

2、使用labelImg标注图像（我标注的是屋顶和车顶），标注后生成的.txt文件存到labels文件夹下（下面会给出文件目录）

3、运行split.py，把图像和标注数据放到合适的地方

split.py

# -*- coding: utf-8 -*-
"""
Created on Fri Dec 29 14:06:08 2023

@author: DIY-PC
"""

import os
import random
import shutil

# 原数据集目录
root_dir = 'D:/python/trainproject/dataset'
# 划分比例
train_ratio = 0.8
valid_ratio = 0.1
test_ratio = 0.1

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

# 拆分后数据集目录
split_dir = 'D:/python/trainproject/dataset/images'
os.makedirs(os.path.join(split_dir, 'train/images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'train/labels'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'valid/images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'valid/labels'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'test/images'), exist_ok=True)
os.makedirs(os.path.join(split_dir, 'test/labels'), exist_ok=True)

# 获取图片文件列表
image_files = os.listdir(os.path.join(root_dir, 'images'))
label_files = os.listdir(os.path.join(root_dir, 'labels'))

# 随机打乱文件列表
combined_files = list(zip(image_files, label_files))
random.shuffle(combined_files)
image_files_shuffled, label_files_shuffled = zip(*combined_files)

# 根据比例计算划分的边界索引
train_bound = int(train_ratio * len(image_files_shuffled))
valid_bound = int((train_ratio + valid_ratio) * len(image_files_shuffled))

# 将图片和标签文件移动到相应的目录
for i, (image_file, label_file) in enumerate(zip(image_files_shuffled, label_files_shuffled)):
    if i < train_bound:
        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'train/images', image_file))
        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'train/labels', label_file))
    elif i < valid_bound:
        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'valid/images', image_file))
        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'valid/labels', label_file))
    else:
        shutil.move(os.path.join(root_dir, 'images', image_file), os.path.join(split_dir, 'test/images', image_file))
        shutil.move(os.path.join(root_dir, 'labels', label_file), os.path.join(split_dir, 'test/labels', label_file))

4、创建roof.yaml文件，指定模型训练元数据的位置（上面准备好的文件夹train,valid,test下的数据）

# Helmet
path: D:/python/trainproject/dataset
train: images/train
val: images/valid
test: images/test

# number of classes
nc: 2
# class names
names: ['roof','busroof']

5、开始模型训练，创建.py文件并运行，代码如下

from ultralytics import YOLO
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
# 加载模型
# model = YOLO("yolov8n.yaml")  # 从头开始构建新模型
model = YOLO("yolov8n.pt")  # 加载预训练模型（推荐用于训练）

# Use the model
results = model.train(data="roof.yaml", epochs=20, batch=2)  # 训练模型

运行结果：

6、获得我们需要的模型

7、使用我们自己训练出的模型进行图像标注测试，创建test.py,代码如下：

from ultralytics import YOLO
from PIL import Image
#model = YOLO('yolov8n.pt')
model = YOLO('XXX/runs/detect/train6/weights/best.pt')
#image = Image.open("bus.jpg")
image = Image.open("D:/python/project/images/DJI_20230925133837_0046.JPG")
results = model.predict(source=image, save=True, save_txt=True) 

运行成功：

8、查看图像标注结果（由于训练模型使用的数据过少只有20张，所以识别不精确，但也能大致识别出来）：