Yolov11训练自己的数据集和ONNX/Trt部署

一、环境配置

1.1 项目下载

项目原地址：GitHub - ultralytics/ultralytics: Ultralytics YOLO11 🚀

• 打开终端，输入以下命令将YOLOv11项目源码下载到本地：

git clone https://github.com/WongKinYiu/yolov7.git
# 进入yolov7的文件夹中
cd yolov7 

1.2 环境配置

# 方案一：使用给出的requirements.txt
# pip批量安装requirements.txt文件中包含的组件依赖
pip install -r requirements.txt
# conda批量安装requirements.txt文件中包含的组件依赖
conda install --yes --file requirements.txt


# 方案二：使用官网给出的安装方案【推荐】
# Install the ultralytics package from PyPI
pip install ultralytics

# 方案三：使用git克隆
# Clone the ultralytics repository
git clone https://github.com/ultralytics/ultralytics
# Navigate to the cloned directory
cd ultralytics
# Install the package in editable mode for development
pip install -e .

二、数据准备

2.1 数据集搜寻

🐞高效搜寻数据集的宝藏平台：遇见数据集整合搜索🐞

🔗 https://www.selectdataset.com/

为什么选择遇见数据集？—— 集成搜索的核心优势

1. 一站式资源聚合
   平台将 Kaggle、UCI、Google Dataset Search 等数十个主流数据集平台的资源整合于一体，无需在多个网站间切换，输入关键词即可同步获取多源结果，大幅提升检索效率。

2. 智能关联与跳转
   搜索结果直接展示各平台的数据集链接，支持一键跳转至原始页面。例如搜索 “医学影像数据集”，可同时呈现 Kaggle 的公开数据集、UCI 的医学数据库及相关学术机构的资源入口，避免重复筛选。

3. 精准分类与筛选
   提供领域标签（计算机视觉 / 自然语言处理 / 金融等）、数据格式（CSV/JSON/ 图像）、许可类型（开源 / 学术使用）等多维筛选条件，帮助用户快速定位符合需求的数据集。

🐞总结的数据集地址：优质数据集平台全收录-CSDN博客🐞

2.2 数据标注 [lableimg和CVAT]

2.2.1 常用的Lableimg

使用命令行安装

pip install labelimg -i https://pypi.tuna.tsinghua.edu.cn/simple

# 安装完成之后输入下面命令即可启动
labelimg

安装命令行成功图：

view里面有个auto save mode 点击之后就可以自动保存标签
## 常用快捷键：
A：切换到上一张图片
D：切换到下一张图片
W：调出标注十字架
del ：删除标注框

标注报错问题解决

libpng warning: iCCP: known incorrect sRGB profile
libpng warning: iCCP: known incorrect sRGB profile
Traceback (most recent call last):
  File "C:\Users\Jorya\.conda\envs\yolov8\Lib\site-packages\libs\canvas.py", line 530, in paintEvent
    p.drawLine(self.prev_point.x(), 0, self.prev_point.x(), self.pixmap.height())
TypeError: arguments did not match any overloaded call:
  drawLine(self, l: QLineF): argument 1 has unexpected type 'float'
  drawLine(self, line: QLine): argument 1 has unexpected type 'float'
  drawLine(self, x1: int, y1: int, x2: int, y2: int): argument 1 has unexpected type 'float'
  drawLine(self, p1: QPoint, p2: QPoint): argument 1 has unexpected type 'float'
  drawLine(self, p1: Union[QPointF, QPoint], p2: Union[QPointF, QPoint]): argument 1 has unexpected type 'float'

解决方案

• 这个错误是由于 PyQt5 的drawLine方法需要整数参数，而代码中传入了浮点数导致的。在 Python 3 中，除法运算默认返回浮点数，而 PyQt5 的绘图函数要求坐标点为整数类型。

# 你需要修改canvas.py文件中第 530 行的代码，将浮点数坐标转换为整数。
# 打开文件C:\Users\Jorya\.conda\envs\yolov8\Lib\site-packages\libs\canvas.py，找到第 530 行：
原代码（错误行）：
p.drawLine(self.prev_point.x(), 0, self.prev_point.x(), self.pixmap.height())
修改后：
p.drawLine(int(self.prev_point.x()), 0, int(self.prev_point.x()), self.pixmap.height())

# 文件中第531  行的代码，将浮点数坐标转换为整数。
原代码（错误行）：
p.drawLine(0, self.prev_point.y(), self.pixmap.width(), self.prev_point.y())
修改后：
p.drawLine(0, int(self.prev_point.y()), self.pixmap.width(), int(self.prev_point.y()))

# 文件中第 526 行的代码，将浮点数坐标转换为整数。
原代码（错误行）：
p.drawRect(left_top.x(), left_top.y(), rect_width, rect_height)
修改后：
p.drawRect(int(left_top.x()), int(left_top.y()), int(rect_width), int(rect_height))

• 标注完整之后会有两个主要的文件夹，一个是图片文件夹另外一个是标签文件夹
  

2.2.2 私有化部署和协同CVAT

官网地址

2.3 数据划分

针对图片和标签都在一个文件夹中

-*- coding: utf-8 -*-
import os
import random
import shutil
from tqdm import tqdm
import re
# 其中数据集和标签都在一个文件夹中，通过读取后缀进行替换
"""
最终划分成yolo可训练数据集格式:
|——images
    |-test
    |-train
    |-valid
|——labels
    |-test
    |-train
    |-valid
|——test.txt
|——train.txt
|——val.txt
## 遇到的yolov7训练问题，已经把预训练模型给下载下来了，进行train训练的时候。
## 一直提示git问题：subprocess.CalledProcessError: Command 'git tag' returned non-zero exit status 128.
## 解决方法：使用git拉取yolov7框架，在进行训练
# nohup python   train.py  --batch 48 --epochs 150 --cfg cfg/deploy/yolov7.yaml --weights yolov7.pt --device 0,1,2 > train.log 2>&1 &
"""
# 数据集文件夹路径
data_dir = "datasets/to/your/path"
# 创建一个空列表来存储有效图片和标签的路径
valid_images = []
valid_labels = []  
# 遍历 data 文件夹下的所有文件
for file_name in os.listdir(data_dir):
    # 获取文件的完整路径
    file_path = os.path.join(data_dir, file_name)
    # 获取文件的扩展名
    ext = os.path.splitext(file_name)[-1].lower()
    if ext.lower()  in ['.jpg', '.png' ,'.jpeg']:  # 处理图片文件
        # 根据扩展名替换成对应的 label 文件名
        label_name = file_name.replace(ext, ".txt")
        label_path = os.path.join(data_dir, label_name)
        # 判断 label 是否存在
        if not os.path.exists(label_path):
            # 删除图片
            os.remove(file_path)
            print("Deleted:", file_path)
        else:
            # 将图片路径添加到列表中
            valid_images.append(file_path)
            # 将 label 路径添加到列表中
            valid_labels.append(label_path)
# 确保每个标签都有对应的图片
valid_labels = [label for label in valid_labels if os.path.exists(label)]
# 创建目标目录结构
base_dir = "newdatasets/to/your/path"
for split in ["train", "valid", "test"]:
    os.makedirs(os.path.join(base_dir, "images", split), exist_ok=True)
    os.makedirs(os.path.join(base_dir, "labels", split), exist_ok=True)
# 打开三个txt文件用于存储分配的路径
train_txt = open(os.path.join(base_dir, "train.txt"), "w")
valid_txt = open(os.path.join(base_dir, "val.txt"), "w")
test_txt = open(os.path.join(base_dir, "test.txt"), "w")
# 初始化计数器
train_count = 0
valid_count = 0
test_count = 0
# 遍历每个有效图片路径
for i in tqdm(range(len(valid_images))):
    image_path = valid_images[i]
    label_path = re.sub(r'\.(jpg|jpeg|png)$', '.txt', image_path, flags=re.IGNORECASE)
    # 随机生成一个概率
    r = random.random()
    # 判断图片和标签应该移动到哪个文件夹 (train:valid:test = 7:2:1)
    if r < 0.1:
        # 移动到 test 文件夹
        destination = "test"
        image_destination = os.path.join(base_dir, "images", destination, os.path.basename(image_path))
        test_txt.write(f"{image_destination}\n")  # 写入到 test.txt
        test_count += 1
    elif r < 0.2:
        # 移动到 valid 文件夹
        destination = "valid"
        image_destination = os.path.join(base_dir, "images", destination, os.path.basename(image_path))
        valid_txt.write(f"{image_destination}\n")  # 写入到 valid.txt
        valid_count += 1
    else:
        # 移动到 train 文件夹
        destination = "train"
        image_destination = os.path.join(base_dir, "images", destination, os.path.basename(image_path))
        train_txt.write(f"{image_destination}\n")  # 写入到 train.txt
        train_count += 1
    # 生成目标文件夹中 images 和 labels 的路径
    image_destination_path = os.path.join(base_dir, "images", destination, os.path.basename(image_path))
    label_destination_path = os.path.join(base_dir, "labels", destination, os.path.basename(label_path))
    # 移动图片到目标文件夹
    shutil.move(image_path, image_destination_path)
    # 移动标签到目标文件夹
    if os.path.exists(label_path):
        shutil.move(label_path, label_destination_path)
# 关闭txt文件
train_txt.close()
valid_txt.close()
test_txt.close()
# 输出有效图片和标签的数量
print(f"有效图片数: {len(valid_images)}")
print(f"有效标签数: {len(valid_labels)}")
# 输出 train, valid, test 集合中的图片数量
print(f"训练集图片数: {train_count}")
print(f"验证集图片数: {valid_count}")
print(f"测试集图片数: {test_count}")

三、模型训练

3.1 创建yaml文件

创建一个hat.yaml文件在ultralytics-main文件夹中的data文件夹中

3.2 模型训练代码和参数解释【单卡训练和多卡训练】

单卡训练：

from ultralytics import YOLO

# Load a model
model = YOLO("/PATH/to/your/yolo11n.pt")  # load a pretrained model (recommended for training)
# Train the model
results = model.train(data="hat.yaml", epochs=100, imgsz=640 , device=0)

多卡训练

from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt")  # load a pretrained model (recommended for training)
# Train the model with 2 GPUs
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])
# Train the model with the two most idle GPUs
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[-1, -1])

3.3 模型验证

from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt")
# Customize validation settings
metrics = model.val(data="hat.yaml", imgsz=640, batch=16, conf=0.25, iou=0.6, device="0")

3.4 模型测试

from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt")  # pretrained YOLO11n model
# Run batched inference on a list of images
results = model(["image1.jpg", "image2.jpg"])  # return a list of Results objects
# Process results list
for result in results:
    boxes = result.boxes  # Boxes object for bounding box outputs
    masks = result.masks  # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs  # Probs object for classification outputs
    obb = result.obb  # Oriented boxes object for OBB outputs
    result.show()  # display to screen
    result.save(filename="result.jpg")  # save to disk

四、部署ONNX和TRT模式的代码

4.1 模型导出

from ultralytics import YOLO
# Load a model
model = YOLO("yolo11n.pt")  # load an official model
model = YOLO("path/to/best.pt")  # load a custom trained model
# Export the model
model.export(format="onnx")

from ultralytics import YOLO
model = YOLO("yolo11n.pt")  # Load a model
model.export(format="engine", int8=True)

4.2 ONNX模型使用代码

import sys
import torch
print(f"Python version: {sys.version}, {sys.version_info} ")
print(f"Pytorch version: {torch.__version__} ")
# Inference for ONNX model
import cv2
cuda = True
w = "/content/yolov7/yolov7-tiny.onnx"
img = cv2.imread('/content/yolov7/inference/images/horses.jpg')
import cv2
import time
import requests
import random
import numpy as np
import onnxruntime as ort
from PIL import Image
from pathlib import Path
from collections import OrderedDict,namedtuple
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if cuda else ['CPUExecutionProvider']
session = ort.InferenceSession(w, providers=providers)
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleup=True, stride=32):
    # Resize and pad image while meeting stride-multiple constraints
    shape = im.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better val mAP)
        r = min(r, 1.0)
    # Compute padding
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    dw /= 2  # divide padding into 2 sides
    dh /= 2
    if shape[::-1] != new_unpad:  # resize
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return im, r, (dw, dh)
names = ['hat']
colors = {name:[random.randint(0, 255) for _ in range(3)] for i,name in enumerate(names)}
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image = img.copy()
image, ratio, dwdh = letterbox(image, auto=False)
image = image.transpose((2, 0, 1))
image = np.expand_dims(image, 0)
image = np.ascontiguousarray(image)
im = image.astype(np.float32)
im /= 255
im.shape
outname = [i.name for i in session.get_outputs()]
outname
inname = [i.name for i in session.get_inputs()]
inname
inp = {inname[0]:im}
# ONNX inference
outputs = session.run(outname, inp)[0]
outputs
ori_images = [img.copy()]
for i,(batch_id,x0,y0,x1,y1,cls_id,score) in enumerate(outputs):
    image = ori_images[int(batch_id)]
    box = np.array([x0,y0,x1,y1])
    box -= np.array(dwdh*2)
    box /= ratio
    box = box.round().astype(np.int32).tolist()
    cls_id = int(cls_id)
    score = round(float(score),3)
    name = names[cls_id]
    color = colors[name]
    name += ' '+str(score)
    cv2.rectangle(image,box[:2],box[2:],color,2)
    cv2.putText(image,name,(box[0], box[1] - 2),cv2.FONT_HERSHEY_SIMPLEX,0.75,[225, 255, 255],thickness=2)  
Image.fromarray(ori_images[0])

---

推荐博客

YOLOv11来了，使用YOLOv11训练自己的数据集和推理(附YOLOv11网络结构图)-CSDN博客