YOLOv10目标检测-训练自己的数据

yolov10
https://github.com/THU-MIG/yolov10?tab=readme-ov-file

1. 数据集

模型的建立需要收集图片并且进行标注。YOLOv10标注的文件格式如下（每张图片对应一个标签文件）：

0 0.441753 0.815461 0.061021 0.042763
1 0.395895 0.759868 0.066198 0.046053
2 0.497781 0.744737 0.060651 0.039474
0 0.575629 0.787171 0.059541 0.042763

<object-class-id> <x> <y> <width> <height>

第一个数字表示目标框的类别（如类别0，1，2，…），后面四个长数字代表：标记框中心点的横纵坐标(x, y)，标记框宽高的大小(w, h)，且都是归一化后的值（图片左上角为坐标原点）。

1.1 标注工具

数据标注的工具有很多，比如LabelImg，LableMe，VIA等，在线标注工具Make Sense

1.2 格式转换

from PIL import Image

def convert_to_yolo_format(bbox, img_size, class_id):
    """
    将目标框 [xmin, ymin, xmax, ymax] 转换为 YOLO 格式
    参数:
    bbox: list，目标框的坐标 [xmin, ymin, xmax, ymax]；
    img_size: img_width: 图像的宽度, img_height: 图像的高度；
    class_id: 目标的类别 ID
    返回:
    str: YOLO 格式的标注
    """
    xmin, ymin, xmax, ymax = bbox
    img_width, img_height = img_size
    # 计算中心坐标和宽度、高度
    x_center = (xmin + xmax) / 2
    y_center = (ymin + ymax) / 2
    width = xmax - xmin
    height = ymax - ymin
    # 归一化坐标和尺寸
    x_center_normalized = x_center / img_width
    y_center_normalized = y_center / img_height
    width_normalized = width / img_width
    height_normalized = height / img_height

    # 返回 YOLO 格式的字符串
    return f"{class_id} {x_center_normalized:.6f} {y_center_normalized:.6f} {width_normalized:.6f} {height_normalized:.6f}"

def write_to_yolo_txt(img_path, bboxes, classes, file_path):
    """
    将多个目标框转换为 YOLO 格式并写入 .txt 文件
    """
    with Image.open(img_path) as img:
        # img_width, img_height = img.size
        img_size = img.size
    with open(file_path, 'w') as f:
        for i in range(len(bboxes)):
            bbox = bboxes[i]
            yolo_format = convert_to_yolo_format(bbox, img_size, classes[i])
            f.write(yolo_format + '\n')

# 示例数据
bboxes = [[404, 233, 497, 326], [308, 233, 401, 329]]
img_path = "./data/0a1a9c2e5a64f54b5f5899f5114bdc6549b6.jpg"
classes = [0, 0]
output_file = 'output.txt'
write_to_yolo_txt(img_path, bboxes, classes, output_file)

voc(xml)转yolo格式

import os
import json
import argparse
import sys
import shutil
from lxml import etree
from tqdm import tqdm

category_set = set()
image_set = set()
bbox_nums = 0


def parse_xml_to_dict(xml):
    """
    将xml文件解析成字典形式，参考tensorflow的recursive_parse_xml_to_dict
    Args:
        xml: xml tree obtained by parsing XML file contents using lxml.etree
    Returns:
        Python dictionary holding XML contents.
    """
    if len(xml) == 0:  # 遍历到底层，直接返回tag对应的信息
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = parse_xml_to_dict(child)  # 递归遍历标签信息
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:  # 因为object可能有多个，所以需要放入列表里
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


def xyxy2xywhn(bbox, size):
    bbox = list(map(float, bbox))
    size = list(map(float, size))
    xc = (bbox[0] + (bbox[2] - bbox[0]) / 2.) / size[0]
    yc = (bbox[1] + (bbox[3] - bbox[1]) / 2.) / size[1]
    wn = (bbox[2] - bbox[0]) / size[0]
    hn = (bbox[3] - bbox[1]) / size[1]
    return (xc, yc, wn, hn)


def parser_info(info: dict, only_cat=True, class_indices=None):
    filename = info['annotation']['filename']
    image_set.add(filename)
    objects = []
    width = int(info['annotation']['size']['width'])
    height = int(info['annotation']['size']['height'])
    for obj in info['annotation']['object']:
        obj_name = obj['name']
        category_set.add(obj_name)
        if only_cat:
            continue
        xmin = int(obj['bndbox']['xmin'])
        ymin = int(obj['bndbox']['ymin'])
        xmax = int(obj['bndbox']['xmax'])
        ymax = int(obj['bndbox']['ymax'])
        bbox = xyxy2xywhn((xmin, ymin, xmax, ymax), (width, height))
        if class_indices is not None:
            obj_category = class_indices[obj_name]
            object = [obj_category, bbox]
            objects.append(object)

    return filename, objects


voc_dir = "./data/Annotations"
txt_dir = "./data/txt"
image_dir = "./data/Images"

class_ind = {"person": 0}

xml_files = [os.path.join(voc_dir, i) for i in os.listdir(voc_dir) if os.path.splitext(i)[-1] == '.xml']
for xml_file in xml_files:
    with open(xml_file, 'rb') as fid:
        xml_str = fid.read()
    xml = etree.fromstring(xml_str)
    info_dict = parse_xml_to_dict(xml) 
    filename, objects = parser_info(info_dict, only_cat=False, class_indices=class_ind) # 273.jpeg [[0, (0.411, 0.7448418156808804, 0.21, 0.14718019257221457)], [0, (0.688, 0.7895460797799174, 0.212, 0.1485557083906465)]]

    image_p = image_dir + "/" + filename
    txt_p = image_dir + "/" + os.path.splitext(os.path.basename(xml_file))[0] + ".txt"
    #################### 写入txt
    with open(txt_p, 'w') as f:
        for obj in objects:
            f.write("{} {:.6f} {:.6f} {:.6f} {:.6f}\n".format(obj[0], obj[1][0], obj[1][1], obj[1][2], obj[1][3]))

参考：
https://blog.csdn.net/mywhyyds/article/details/143831918
https://blog.csdn.net/festaw/article/details/138039269

1.3 数据

1.3.1 数据集

在根目录下建data文件夹，存放数据，将原本数据集按照7：2：1的比例划分成训练集、验证集和测试集三类

import os
import random
import shutil

# 创建文件夹
# dataset_path = "/data3/guoman/pay_check/seal_model/seal_detection/yolov10/data/dataset1"
# train_images_path = os.path.join(dataset_path,"train/images")
# train_labels_path = os.path.join(dataset_path,"train/labels")
# valid_images_path = os.path.join(dataset_path,"valid/images")
# valid_labels_path = os.path.join(dataset_path,"valid/labels")
# test_images_path = os.path.join(dataset_path,"test/images")
# test_labels_path = os.path.join(dataset_path,"test/labels")

# os.makedirs(train_images_path, exist_ok=True)
# os.makedirs(train_labels_path, exist_ok=True)
# os.makedirs(valid_images_path, exist_ok=True)
# os.makedirs(valid_labels_path, exist_ok=True)
# os.makedirs(test_images_path, exist_ok=True)
# os.makedirs(test_labels_path, exist_ok=True)

# 设置目录路径
image_dir = 'D:/ModelsDatas/test/allimgs'  # 改成你自己的原图片目录
label_dir = 'D:/ModelsDatas/test/alllables'  # 改成你自己的原标签目录

# 获取图片和txt文件列表
images = os.listdir(image_dir)
labels = os.listdir(label_dir)

# 设置随机种子，确保结果可复现
random.seed(2024)
# 随机打乱图片列表
random.shuffle(images)

# 计算训练集、验证集和测试集的数量
total_images = len(images)
train_count = int(total_images * 0.7)
val_count = int(total_images * 0.2)
test_count = total_images - train_count - val_count

# 分配文件到训练集、验证集和测试集
train_images = images[:train_count]
val_images = images[train_count:train_count + val_count]
test_images = images[train_count + val_count:]

# 移动文件到对应的目录
for image in train_images:
    # 移动图片和标签到训练集目录
    shutil.move(os.path.join(image_dir, image), 'D:/ModelsDatas/YOLO_datasets/train/images') # 请改成你自己的训练集存放图片的文件夹目录
    shutil.move(os.path.join(label_dir, image[:-4]+'.txt'), 'D:/ModelsDatas/YOLO_datasets/train/labels')# 请改成你自己的训练集存放标签的文件夹目录

for image in val_images:
    # 移动图片和标签到验证集目录
    shutil.move(os.path.join(image_dir, image), 'D:/ModelsDatas/YOLO_datasets/valid/images')# 请改成你自己的验证集存放图片的文件夹目录
    shutil.move(os.path.join(label_dir, image[:-4] + '.txt'), 'D:/ModelsDatas/YOLO_datasets/valid/labels')# 请改成你自己的验证集存放标签的文件夹目录

for image in test_images:
    # 移动图片和标签到测试集目录
    shutil.move(os.path.join(image_dir, image), 'D:/ModelsDatas/YOLO_datasets/test/images')# 请改成你自己的测试集存放图片的文件夹目录
    shutil.move(os.path.join(label_dir, image[:-4] + '.txt'), 'D:/ModelsDatas/YOLO_datasets/test/labels')# 请改成你自己的测试集存放标签的文件夹目录

1.3.2 数据集文件配置

准备yaml文件，放在data路径下，用于后续训练

names: # class names
  - 0:fish # 类别号: 类别名称 （需要改成自己的）
  #-1:cat #如果还有其他类别，以此往下加就行了，类别号请认真和自己当时目标框labels文件中的一一对应
  #-2:dog
nc: 1 # number of classes 数据集中一共有几个类别，参考上面说的
path: ultralytics/datasets/det/fish2 # 数据集路径（需要改成自己的，也就是train、test和valid目录的上级目录）
train: train/images # 训练集路径(相对于数据集路径)
val: valid/images # 验证集路径(相对于数据集路径)
test: test/images # 测试集路径(相对于数据集路径)

2. 模型

2.1 模型文件配置

在ultralytics/cfg/models/v10文件夹下存放的是YOLOv10的各个版本的模型配置文件，检测的类别是coco数据的80类。在训练自己数据集的时候，只需要将其中的类别数修改成自己的大小。在根目录文件夹下新建yolov10n-test.yaml文件，此处以yolov10n.yaml文件中的模型为例，将其中的内容复制到yolov10n-test.yaml文件中 ，并将nc: 1 # number of classes 修改类别数` 修改成自己的类别数

2.2 训练文件配置

在进行模型训练之前，需要到官网下载预训练权重，权重地址为：Releases · THU-MIG/yolov10 · GitHub

根据所选择的模型下载相应的权重，比如yolov10n.pt，放在根目录weights/yolov10n.pt路径下。

YOLOv10的超参数配置在ultralytics/cfg文件夹下的default.yaml文件中。

YOLOv10有多种模型，可满足不同的应用需求：

• YOLOv10-N：用于资源极其有限的环境的纳米版本。
• YOLOv10-S：兼顾速度和精度的小型版本。
• YOLOv10-M：通用中型版本。
• YOLOv10-B：平衡型，宽度增加，精度更高。
• YOLOv10-L：大型版本，精度更高，但计算资源增加。
• YOLOv10-X：超大型版本可实现高精度和性能。

3. 模型训练

3.1 train（不对）

python train.py --weights weights/yolov5s.pt  --cfg models/yolov5s.yaml  --data data/myvoc.yaml --epoch 200 --batch-size 8 --img 640   --device 0

参数说明如下（注意epochs及device）：

• weights：指定预训练模型权重文件的路径。
• cfg：模型配置文件的路径，定义网络结构。
• data：数据集配置文件的路径。
• epochs：训练的轮数，即模型将完整遍历数据集的次数（例如：200）。
• batch-size：批次大小，即每次训练中处理的图片数量（例如：8）。
• img：输入图像的尺寸，指定训练时图像的宽和高（例如：640）。
• device：指定使用的设备，输入 cpu 表示在 CPU 上进行训练；输入数字如 0, 1, 2, 3 则表示对应的 GPU 编号。例如：
  （1）device 0 表示使用默认第一个 GPU；
  （2）device 0,1 表示同时使用第一个和第二个 GPU（适合多 GPU 训练）。
  （3）当有多个 GPU 时，可以通过 nvidia-smi 命令查看每个 GPU 的编号。

3.2 train

train.py

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLOv10
 
model_yaml_path = r'D:\project\yolov10-main\datasets\yolov10n.yaml'
 
data_yaml_path = r'D:\project\yolov10-main\datasets\data.yaml'
#预训练模型
pre_model_name = 'yolov10n.pt'

if __name__ == '__main__':
    model = YOLOv10(model_yaml_path).load(pre_model_name)
    results = model.train(data=data_yaml_path,
                          epochs=100,
                          batch=32,
                          workers=0,
                          device='0',
                          project='runs/train',
                          name='exp',
                          )
 

4. 验证

from ultralytics import YOLOv10

model = YOLOv10('./runs/train/exp5/weights/best.pt')
model.val(data='./data/dataset1/myvoc.yaml', batch=256)

5. 测试

from ultralytics import YOLOv10

yolo = YOLOv10("./runs/train/exp5/weights/best.pt",task="detect")
result = yolo(source="./data/dataset1/test/images",save=True,save_conf=True,save_txt=True,name='output')

source后为要预测的图片数据集的的路径
save=True为保存预测结果
save_conf=True为保存坐标信息
save_txt=True为保存txt结果

6. 推理

from ultralytics import YOLOv10

# Load a pretrained YOLOv10n model
model = YOLOv10("./runs/train/exp5/weights/best.pt")

# Perform object detection on an image
# results = model.predict("data/dataset1/test/images/f14e45acc1484e3a7b0ee890196e5bfd.jpg",save=True)
results = model.predict("data/dataset1/test/images/f14e45acc1484e3a7b0ee890196e5bfd.jpg")
# Display the results
# results[0].show()

# 预测的位置
predict_bbox = results[0].to("cpu").numpy().boxes.xyxy.tolist()
# 预测的置信度
predict_confidence = results[0].to("cpu").numpy().boxes.conf.tolist()
# 预测的类别
predict_class = (results[0].to("cpu").numpy().boxes.cls).astype(int).tolist()
print(predict_bbox,predict_confidence,predict_class)

参考：
https://blog.csdn.net/tqh267/article/details/139283703
https://openatomworkshop.csdn.net/6743dd743a01316874d7605d.html
https://blog.csdn.net/Natsuago/article/details/143647392
https://cloud.tencent.com/developer/article/2425073
https://blog.csdn.net/weixin_43694096/article/details/139422902
https://blog.csdn.net/c858845275/article/details/141089142
https://blog.csdn.net/tsg6698/article/details/139586476
https://developer.aliyun.com/article/1536979
https://zhuanlan.zhihu.com/p/700013499