【YOLO11】简单详细的自制数据集并训练YOLO11目标检测

YOLO11新发布，本篇将详细介绍如何用自制的数据集来训练YOLO11模型。

1 数据集制作

数据集的制作跟YOLO系列是完全一样的，使用Labelme进行标注，会生成 .json文件。

YOLO接受的标签数据是包含id和归一化后坐标的 .txt文件，所以自己写一个 .json转 .txt的python 脚本。以下程序首先将没有标签的数据删除，将剩余的图片和标签保存到新的一个文件夹内(/selec_data)；然后将.json标签转为.txt标签；最后随机划分为训练集和验证集保存在两个子文件夹内（/train_data/train和/train_data/val）

import os 
import glob
import cv2
import json
import numpy as np
import random
import shutil

## Labelme保存的图片和标签的文件夹：
in_dir = './custom_data/data/'

## 按匹配对整理，去除没有标签的图片
img_pths = glob.glob(in_dir + "/*.bmp")	# 自制数据的图片后缀
for img_pth in img_pths:
    json_pth = img_pth.replace(".bmp", ".json")
    print(json_pth)
    if os.path.exists(json_pth):
        shutil.copy(img_pth, img_pth.replace('/data', '/select_data'))
        shutil.copy(json_pth, json_pth.replace('/data', '/select_data'))


## 将labelme_json标注转_txt
def convert(size, box):
    """
    convert [xmin, xmax, ymin, ymax] to [x_centre, y_centre, w, h]
    """
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

class_names = ["pupil", 'iris'] #数组元素替换为自己数据的类别

out_dir = "./custom_data/select_data/"
txt_dir = "./custom_data/select_data/"
os.makedirs(txt_dir, exist_ok=True)

json_pths = glob.glob(out_dir + "/*.json")

for json_pth in json_pths:
    f1 = open(json_pth, "r")
    json_data = json.load(f1)

    img_pth = os.path.join(json_pth.replace("json", "bmp"))
    img = cv2.imread(img_pth)
    h, w = img.shape[:2]

    tag = os.path.basename(json_pth)
    out_file = open(os.path.join(txt_dir, tag.replace("json", "txt")), "w")
    # print(json_data)
    label_infos = json_data["shapes"]
    for label_info in label_infos:
        label = label_info["label"]
        points = label_info["points"]
        print("+++", len(points))
        if len(points) >= 3:
            points = np.array(points)
            print(points.shape)
            xmin, xmax = max(0, min(np.unique(points[:, 0]))), min(w, max(np.unique(points[:, 0])))
            ymin, ymax = max(0, min(np.unique(points[:, 1]))), min(h, max(np.unique(points[:, 1])))
            print("++++", ymin, ymax)
        elif len(points) == 2:
            x1, y1 = points[0]
            x2, y2 = points[1]
            xmin, xmax = min(x1, x2), max(x1, x2)
            ymin, ymax = min(y1, y2), max(y1, y2)
        else:
            continue
        bbox = [xmin, xmax, ymin, ymax]
        bbox_ = convert((w,h), bbox)
        cls_id = class_names.index(label)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bbox_]) + '\n')


def mymovefile(srcfile, dstpath):                       # 移动函数
    if not os.path.isfile(srcfile):
        print ("%s not exist!"%(srcfile))
    else:
        fpath,fname=os.path.split(srcfile)             # 分离文件名和路径
        if not os.path.exists(dstpath):
            os.makedirs(dstpath)                       # 创建路径
        shutil.move(srcfile, dstpath + fname)          # 移动文件
        print ("move %s -> %s"%(srcfile, dstpath + fname))

# 划分训练集和验证集
out_dir10 = "./custom_data/train_data/train/images"
out_dir11= "./custom_data/train_data/train/labels"
out_dir20 = "./custom_data/train_data/val/images"
out_dir21 = "./custom_data/train_data/val/labels"

os.makedirs(out_dir10, exist_ok=True)
os.makedirs(out_dir11, exist_ok=True)
os.makedirs(out_dir20, exist_ok=True)
os.makedirs(out_dir21, exist_ok=True)

txt_dir = "./custom_data/select_data/"
img_dir = "./custom_data/select_data/"
txt_names = glob.glob(txt_dir+'*.txt')

l = len(txt_names)
sel_names = random.sample(txt_names, int(l*0.8)) #0.90
# print(sel_names)
for txt_name in txt_names:
    img_name = txt_name.replace('txt', 'bmp')
    if txt_name in sel_names:  
        print("---", txt_name)
        shutil.copy(txt_name, txt_name.replace('/select_data', '/train_data/train/labels'))          # 复制文件
        shutil.copy(img_name, img_name.replace('/select_data', '/train_data/train/images'))
    else:
        shutil.copy(txt_name, txt_name.replace('/select_data', '/train_data/val/labels'))          # 复制文件
        shutil.copy(img_name, img_name.replace('/select_data', '/train_data/val/images'))

这时候，数据集文件夹的文件结构如下：

-- custom_data
	-- select_data
	-- train_data
		-- train
			-- images
				img1.bmp
				...
			-- labels
				img1.txt
				...
		-- val
			-- images
			-- labels

2 准备训练

训练环境

Python >=3.8
PyTorch >= 1.8
再安装其他必须的库即可

在ultralytics官方仓库下载YOLO11源码，下载后打开整个工程，在根目录下新建文件夹，命名为datasets，把刚刚准备好的数据集复制到./datasets文件夹下。

如上面图片所示，在./ultraltics/cfg/datasets/下新建一个custom_data.yaml文件，打开文件复制以下代码：

# my custom dataset
path: custom_data/train_data
train: # train images (relative to 'path') 
  - train
val: # val images (relative to 'path')
  - val

# number of classes
nc: 2

# class names
names: ["pupil", "iris"]

path处填写./datasets文件夹下的数据路径，分别有train和val两个子文件夹；类别数量和类别名称按自己的填写，这里的名称要按照你标注时的id顺序填写，否则后面验证和预测的时候会出现混乱。

3 开始训练

如上图所示，在文件夹根目录下新建python文件./train.py：

from ultralytics import YOLO

if __name__ =='__main__':
	# 模型大小从小到大有n, s, m, l, x，需要哪一个就改11后面的字母
    model = YOLO('ultralytics/cfg/models/v8/yolo11n.yaml')

	# 这里的data必须设置为我们刚才创建的.yaml文件
    model.train(data = 'ultralytics/cfg/datasets/custom_data.yaml', 
                epochs = 200, 
                batch = 32, 
                imgsz = 640,
                device = 0,  # 如果是CPU训练就写'cpu'，否则表示GPU的编号
                workers = 0 # 这里推荐workers就写0，否则可能占用C盘空间和虚拟内存，导致出现页面文件太小的报错
                )  
                
    metrics = model.val()

因为我们下载了ultralytics的源码，所以不用安装ultralytics库就可以直接from ultralytics import YOLO。

运行train.py即可开始训练。

训模型文件都保存在./runs/detect/train/weights/，这里还能看到验证结果的图片。注意一下：文件夹./runs/detect/train名称会每运行一次就累加一次，train2、train3…

4 测试

在文件夹根目录下新建python文件./test.py：

from ultralytics import YOLO

if __name__ =='__main__':
	# 读取模型权重，在刚才说的run文件夹里
    model = YOLO('runs/detect/train/weights/best.pt')
	
	# 预测，路径是待预测图片的路径，我放在了训练数据的datasets文件夹里
    results = model("datasets/test_data/tImg1.bmp")
    results[0].show()

直接显示预测结果图

整个过程跟YOLOv8的过程是一模一样的，ultralytics也直接在YOLOv8仓库中更新了YOLO11，所以没有遇到什么新的bug。

有任何问题欢迎在评论区里询问，我每天都会看。