【代码复现】YOLO11复现全流程+自定义数据集训练测试

【代码复现】YOLO11复现全流程+自定义数据集训练

一、YOLO11复现全流程

1、环境配置

我的环境：windows11 cuda12.4 python3.9 pytorch2.6.0

1.1 环境建立

conda create -n yolo11 python==3.9
conda activate yolo11

1.2 安装pytorch

通过nvidia-smi查看cuda版本，去pytorch官网查找可用的pytorch版本，要确保python版本、torch版本、cuda版本、torchvision版本对应，然后安装：

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124

1.3 安装ultralytics

pip install ultralytics

2、代码下载

到yolo11的官方代码网址下载源代码：https://github.com/ultralytics/ultralytics

打开项目后配置python解释器：

3、权重下载及复现

3.1 权重下载

在根目录创建防止权重的文件夹，并命名为weights；

在官方代码链接下，找到模型这一块，包括检测、分割、分类等任务的预训练权重；

直接点击需要的权重，下载，放置到weights文件夹下：

3.2 对检测任务进行复现

在终端执行：

yolo predict model=yolo11n.pt source=./ultralytics/assets/bus.jpg

结果：

3.3 对姿态任务进行复现

在终端执行：

yolo predict model=yolo11n-pose.pt source=./ultralytics/assets/bus.jpg

结果：

3.4 对分割任务进行复现

在终端执行：

yolo predict model=yolo11n-seg.pt source=./ultralytics/assets/bus.jpg

结果：

二、自定义数据集训练

1、数据标注（labelimg）

我使用labelimg进行数据标注，具体做法如下：

先建立环境，然后安装labelimg，打开Anaconda Prompt，执行以下代码：

conda create -n objectdetect python==3.8
conda activate objectdetect
pip install labelimg

然后再输入labelimg就会自动弹出标注界面：

labelimg

点击“Open Dir”，找到你要标注的数据的文件夹，点击“选择文件夹”，再打开“Change Save Dir”，找到你要保存的路径，点击“选择文件夹”：

界面中的“Next Image”表示下一张，“Prev Image”表示上一张，快捷键分别是“D”和“A”。

点击“PascalVOC”可以切换标注的格式，有voc、yolo和json三种。

“Create RectBox”用来创建矩形框，点击即可创建，快捷键是“W”。

此时就可以对自己的数据进行标注了。

2、VOC数据集格式（xml）与YOLO格式（txt）转换

代码如下：

# 图像和xml修改分辨率
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import joindef convert(size, box):x_center = (box[0] + box[1]) / 2.0y_center = (box[2] + box[3]) / 2.0x = x_center / size[0]y = y_center / size[1]w = (box[1] - box[0]) / size[0]h = (box[3] - box[2]) / size[1]return (x, y, w, h)def convert_annotation(xml_files_path, save_txt_files_path, classes):xml_files = os.listdir(xml_files_path)print(xml_files)for xml_name in xml_files:print(xml_name)xml_file = os.path.join(xml_files_path, xml_name)out_txt_path = os.path.join(save_txt_files_path, xml_name.split('.')[0] + '.txt')out_txt_f = open(out_txt_path, 'w')tree = ET.parse(xml_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))# b=(xmin, xmax, ymin, ymax)print(w, h, b)bb = convert((w, h), b)out_txt_f.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')if __name__ == "__main__":# 需要转换的类别，需要一一对应classes1 = ['pos1', 'pos2', 'pos3', 'pos4', 'pos5', 'pos6']# 2、voc格式的xml标签文件路径xml_files1 = r'你的voc格式标注路径'# 3、转化为yolo格式的txt标签文件存储路径save_txt_files1 = r'你要保存的路径'convert_annotation(xml_files1, save_txt_files1, classes1)

为检查转换后的标注是否出错，可以在原图上画出标注的矩形框，可执行以下代码：

# 对修改分辨率且转换成txt格式的标注信息在原图上画出来，检测转换完的txt与图片是否对应import os
import cv2data_path = r'你的原图数据集路径'
txt_label_path = r'原图对应的YOLO格式标注路径'class_names = ["pos1", "pos2", "pos3", "pos4", 'pos5', 'pos6']  # 类别for file in os.listdir(data_path):filepath = os.path.join(data_path, file)temp = file.split('.jpg')[0]temp = temp + '.txt'txtpath = os.path.join(txt_label_path, temp)if os.path.exists(txtpath):image = cv2.imread(filepath)height, width, _ = image.shape# 读取YOLO标注文件with open(txtpath, "r") as file:lines = file.readlines()for line in lines:parts = line.strip().split()if len(parts) != 5:continueclass_id = int(parts[0])x_center, y_center, w, h = map(float, parts[1:])# 计算像素坐标x1 = int(x_center*width - w*width/2)x1 = int((x_center - w / 2) * width)y1 = int((y_center - h / 2) * height)x2 = int((x_center + w / 2) * width)y2 = int((y_center + h / 2) * height)# 画矩形框color = (0, 255, 0)  # 绿色thickness = 2cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)# 画类别标签label = class_names[class_id] if class_id < len(class_names) else str(class_id)font = cv2.FONT_HERSHEY_SIMPLEXcv2.putText(image, label, (x1, y1 - 10), font, 0.5, color, 1, cv2.LINE_AA)# 显示图像cv2.imshow("YOLO Annotations", image)cv2.waitKey(0)  # 按任意键关闭窗口cv2.destroyAllWindows()

3、数据集yaml配置文件

打开ultralytics/cfg/datasets目录，在该目录新建一个yaml文件，例如mydata_peigu.yaml，放置你的数据集配置参数：

path: E:\Instrument-development\ultralytics-main\datasets\pei_gu\new-train1
train: train/images # train images
val: val/images # val images# Classes
names:0: pos11: pos22: pos33: pos44: pos55: pos6

4、编写训练测试代码

4.1 编写训练代码

在根目录下创建train.py：

from ultralytics import YOLO# 加载预训练的模型
model = YOLO("yolo11s.yaml").load("weights/yolo11s.pt")# 定义训练参数，添加默认值、范围和中文注释
train_params = {'data': "mydata_peigu.yaml",  # 数据集配置文件路径，需要自定义修改'epochs': 200,  # 总训练轮次，默认值 100，范围 >= 1'imgsz': 640,  # 输入图像大小，默认值 640，范围 >= 32'batch': 8,  # 批次大小，默认值 16，范围 >= 1'save': True,  # 是否保存训练结果和模型，默认值 True'save_period': -1,  # 模型保存频率，默认值 -1，表示只保存最终结果'cache': False,  # 是否缓存数据集，默认值 False'device': None,  # 训练设备，默认值 None，支持 "cpu", "gpu"(device=0,1), "mps"'workers': 6,  # 数据加载线程数，默认值 8，影响数据预处理速度'project': None,  # 项目名称，保存训练结果的目录，默认值 None'name': None,  # 训练运行的名称，用于创建子目录保存结果，默认值 None'exist_ok': False,  # 是否覆盖已有项目/名称目录，默认值 False'optimizer': 'auto',  # 优化器，默认值 'auto'，支持 'SGD', 'Adam', 'AdamW''verbose': True,  # 是否启用详细日志输出，默认值 False'seed': 0,  # 随机种子，确保结果的可重复性，默认值 0'deterministic': True,  # 是否强制使用确定性算法，默认值 True'single_cls': False,  # 是否将多类别数据集视为单一类别，默认值 False'rect': False,  # 是否启用矩形训练（优化批次图像大小），默认值 False'cos_lr': False,  # 是否使用余弦学习率调度器，默认值 False'close_mosaic': 10,  # 在最后 N 轮次中禁用 Mosaic 数据增强，默认值 10'resume': False,  # 是否从上次保存的检查点继续训练，默认值 False'amp': True,  # 是否启用自动混合精度（AMP）训练，默认值 True'fraction': 1.0,  # 使用数据集的比例，默认值 1.0'profile': False,  # 是否启用 ONNX 或 TensorRT 模型优化分析，默认值 False'freeze': None,  # 冻结模型的前 N 层，默认值 None'lr0': 0.01,  # 初始学习率，默认值 0.01，范围 >= 0'lrf': 0.01,  # 最终学习率与初始学习率的比值，默认值 0.01'momentum': 0.937,  # SGD 或 Adam 的动量因子，默认值 0.937，范围 [0, 1]'weight_decay': 0.0005,  # 权重衰减，防止过拟合，默认值 0.0005'warmup_epochs': 3.0,  # 预热学习率的轮次，默认值 3.0'warmup_momentum': 0.8,  # 预热阶段的初始动量，默认值 0.8'warmup_bias_lr': 0.1,  # 预热阶段的偏置学习率，默认值 0.1'box': 7.5,  # 边框损失的权重，默认值 7.5'cls': 0.5,  # 分类损失的权重，默认值 0.5'dfl': 1.5,  # 分布焦点损失的权重，默认值 1.5'pose': 12.0,  # 姿态损失的权重，默认值 12.0'kobj': 1.0,  # 关键点目标损失的权重，默认值 1.0'label_smoothing': 0.0,  # 标签平滑处理，默认值 0.0'nbs': 64,  # 归一化批次大小，默认值 64'overlap_mask': True,  # 是否在训练期间启用掩码重叠，默认值 True'mask_ratio': 4,  # 掩码下采样比例，默认值 4'dropout': 0.0,  # 随机失活率，用于防止过拟合，默认值 0.0'val': True,  # 是否在训练期间启用验证，默认值 True'plots': True,  # 是否生成训练曲线和验证指标图，默认值 True# 数据增强相关参数'hsv_h': 0.2,  # 色相变化范围 (0.0 - 1.0)，默认值 0.015'hsv_s': 0.7,  # 饱和度变化范围 (0.0 - 1.0)，默认值 0.7'hsv_v': 0.4,  # 亮度变化范围 (0.0 - 1.0)，默认值 0.4'degrees': 30.0,  # 旋转角度范围 (-180 - 180)，默认值 0.0'translate': 0.1,  # 平移范围 (0.0 - 1.0)，默认值 0.1'scale': 0.5,  # 缩放比例范围 (>= 0.0)，默认值 0.5'shear': 0.0,  # 剪切角度范围 (-180 - 180)，默认值 0.0'perspective': 0.0,  # 透视变化范围 (0.0 - 0.001)，默认值 0.0'flipud': 0.0,  # 上下翻转概率 (0.0 - 1.0)，默认值 0.0'fliplr': 0.5,  # 左右翻转概率 (0.0 - 1.0)，默认值 0.5'bgr': 0.0,  # BGR 色彩顺序调整概率 (0.0 - 1.0)，默认值 0.0'mosaic': 0.5,  # Mosaic 数据增强 (0.0 - 1.0)，默认值 1.0'mixup': 0.0,  # Mixup 数据增强 (0.0 - 1.0)，默认值 0.0'copy_paste': 0.0,  # Copy-Paste 数据增强 (0.0 - 1.0)，默认值 0.0'copy_paste_mode': 'flip',  # Copy-Paste 增强模式 ('flip' 或 'mixup')，默认值 'flip''auto_augment': 'randaugment',  # 自动增强策略 ('randaugment', 'autoaugment', 'augmix')，默认值 'randaugment''erasing': 0.4,  # 随机擦除增强比例 (0.0 - 0.9)，默认值 0.4'crop_fraction': 1.0,  # 裁剪比例 (0.1 - 1.0)，默认值 1.0}# 进行训练
if __name__ == '__main__':results = model.train(**train_params)

4.2 编写测试代码

在根目录下创建infer.py：

#  默认检测的结果保存在runs/detect/XXXX下，包括图像和labels
import os.pathfrom ultralytics import YOLO# 加载预训练的YOLOv11n模型
model = YOLO(r"runs/detect/配固/new-train1/weights/best.pt")# 对指定的图像文件夹进行推理，并设置各种参数
if __name__ == '__main__':results = model.predict(source=r"E:\Instrument-development\ultralytics-main\datasets\pei_gu\test\img",  # 数据来源，可以是文件夹、图片路径、视频、URL，或设备ID（如摄像头）conf=0.45,  # 置信度阈值iou=0.6,  # IoU 阈值imgsz=640,  # 图像大小half=False,  # 使用半精度推理device=None,  # 使用设备，None 表示自动选择，比如'cpu','0'max_det=300,  # 最大检测数量vid_stride=1,  # 视频帧跳跃设置stream_buffer=False,  # 视频流缓冲visualize=False,  # 可视化模型特征augment=False,  # 启用推理时增强agnostic_nms=False,  # 启用类无关的NMSclasses=None,  # 指定要检测的类别retina_masks=False,  # 使用高分辨率分割掩码embed=None,  # 提取特征向量层show=False,  # 是否显示推理图像save=True,  # 保存推理结果save_frames=False,  # 保存视频的帧作为图像save_txt=True,  # 保存检测结果到文本文件save_conf=False,  # 保存置信度到文本文件save_crop=False,  # 保存裁剪的检测对象图像show_labels=True,  # 显示检测的标签show_conf=True,  # 显示检测置信度show_boxes=True,  # 显示检测框line_width=None  # 设置边界框的线条宽度，比如2，4)# a = results[0]# for i in range(len(results)):#     p = results[i]#     object = p.names#     filename = p.path#     print(filename + " : " + str(object))# print(1)# save_path = r'E:\Instrument-development\ultralytics-main\runs\detect\predict\labels'# for file in save_path:#     filepath = os.path.join(save_path, file)#     with open(filepath, 'r', encoding='utf-8') as file:#         for line in file:for r in results:print(r.boxes.cls)

5、模型训练和测试

运行train.py，开始训练：

结束后可以看到：

运行infer.py，检测的结果会保存在默认路径runs/detect/XXXX下，包括图像和labels

至此完成源代码复现及使用自己的数据集进行训练。

该文章记录了个人使用YOLO11进行目标检测的流程，若有问题，欢迎批评指正