【深度学习】安全帽检测，目标检测，yolov10算法，yolov10训练

寻求帮助请看这里：

https://docs.qq.com/sheet/DUEdqZ2lmbmR6UVdU?tab=BB08J2
1

一、数据集

安全帽佩戴检测
数据集：https://github.com/njvisionpower/Safety-Helmet-Wearing-Dataset
基准模型：

二、yolov10介绍

听说过yolov10吗：https://www.jiqizhixin.com/articles/2024-05-28-7

论文：

https://arxiv.org/abs/2405.14458

代码：

https://github.com/THU-MIG/yolov10

三、数据voc转换为yolo

调整一下，整成这样：

VOC2028 # tree -L 1
.
├── images
├── labels
├── test.txt
├── train.txt
├── trainval.txt
└── val.txt

2 directories, 4 files

1
2
3
4
5
6
7
8
9
10
11

写为绝对路径：

# 定义需要处理的文件名列表
file_names = ['test.txt', 'train.txt', 'trainval.txt', 'val.txt']

for file_name in file_names:
    # 打开文件用于读取
    with open(file_name, 'r') as file:
        # 读取所有行
        lines = file.readlines()
    
    # 打开（或创建）另一个文件用于写入修改后的内容，这里使用新的文件名表示已修改
    new_file_name = 'modified_' + file_name
    with open(new_file_name, 'w') as new_file:
        # 遍历每一行并进行修改
        for line in lines:
            # 删除行尾的换行符，添加'.jpg'和'images/'，然后再添加回换行符
            modified_line = '/ssd/xiedong/yolov10/VOC2028/images/' + line.strip() + '.jpg\n'
            # 将修改后的内容写入新文件
            new_file.write(modified_line)

print("所有文件处理完成。")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

转yolo txt：

import traceback
import xml.etree.ElementTree as ET
import os
import shutil
import random
import cv2
import numpy as np
from tqdm import tqdm


def convert_annotation_to_list(xml_filepath, size_width, size_height, classes):
    in_file = open(xml_filepath, encoding='UTF-8')
    tree = ET.parse(in_file)
    root = tree.getroot()
    # size = root.find('size')
    # size_width = int(size.find('width').text)
    # size_height = int(size.find('height').text)
    yolo_annotations = []
    # if size_width == 0 or size_height == 0:
    for obj in root.iter('object'):
        difficult = obj.find('difficult').text
        cls = obj.find('name').text
        if cls not in classes:
            classes.append(cls)
        cls_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)]

        # 标注越界修正
        if b[1] > size_width:
            b[1] = size_width
        if b[3] > size_height:
            b[3] = size_height

        txt_data = [((b[0] + b[1]) / 2.0) / size_width, ((b[2] + b[3]) / 2.0) / size_height,
                    (b[1] - b[0]) / size_width, (b[3] - b[2]) / size_height]
        # 标注越界修正
        if txt_data[0] > 1:
            txt_data[0] = 1
        if txt_data[1] > 1:
            txt_data[1] = 1
        if txt_data[2] > 1:
            txt_data[2] = 1
        if txt_data[3] > 1:
            txt_data[3] = 1
        yolo_annotations.append(f"{cls_id} {' '.join([str(round(a, 6)) for a in txt_data])}")

    in_file.close()
    return yolo_annotations


def main():
    classes = []

    root = r"/ssd/xiedong/yolov10/VOC2028"
    img_path_1 = os.path.join(root, "images")
    xml_path_1 = os.path.join(root, "labels")

    dst_yolo_root_txt = xml_path_1

    index = 0
    img_path_1_files = os.listdir(img_path_1)
    xml_path_1_files = os.listdir(xml_path_1)
    for img_id in tqdm(img_path_1_files):
        # 右边的.之前的部分
        xml_id = img_id.split(".")[0] + ".xml"
        if xml_id in xml_path_1_files:
            try:
                img = cv2.imdecode(np.fromfile(os.path.join(img_path_1, img_id), dtype=np.uint8), 1)  # img是矩阵
                new_txt_name = img_id.split(".")[0] + ".txt"
                yolo_annotations = convert_annotation_to_list(os.path.join(xml_path_1, img_id.split(".")[0] + ".xml"),
                                                              img.shape[1],
                                                              img.shape[0],
                                                              classes)
                with open(os.path.join(dst_yolo_root_txt, new_txt_name), 'w') as f:
                    f.write('\n'.join(yolo_annotations))
            except:
                traceback.print_exc()

    # classes
    print(f"我已经完成转换 {classes}")


if __name__ == '__main__':
    main()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89

vim voc2028x.yaml

train: /ssd/xiedong/yolov10/VOC2028/modified_train.txt
val: /ssd/xiedong/yolov10/VOC2028/modified_val.txt
test: /ssd/xiedong/yolov10/VOC2028/modified_test.txt

# Classes
names:
  0: hat
  1: person
1
2
3
4
5
6
7
8

四、训练

环境：

git clone https://github.com/THU-MIG/yolov10.git
cd yolov10
conda create -n yolov10 python=3.9 -y
conda activate yolov10
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple some-package
pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple some-package


1
2
3
4
5
6
7
8

训练

yolo detect train data="/ssd/xiedong/yolov10/voc2028x.yaml" model=yolov10s.yaml epochs=200 batch=64 imgsz=640 device=1,3
1

训练启动后：

训练完成后：

五、验证

yolo val model="/ssd/xiedong/yolov10/runs/detect/train2/weights/best.pt" data="/ssd/xiedong/yolov10/voc2028x.yaml" batch=32 imgsz=640 device=1,3
1

map50平均达到0.94，已超出基准很多了。

预测：

yolo predict model=yolov10n/s/m/b/l/x.pt
1

导出：

# End-to-End ONNX
yolo export model=yolov10n/s/m/b/l/x.pt format=onnx opset=13 simplify
# Predict with ONNX
yolo predict model=yolov10n/s/m/b/l/x.onnx

# End-to-End TensorRT
yolo export model=yolov10n/s/m/b/l/x.pt format=engine half=True simplify opset=13 workspace=16
# Or
trtexec --onnx=yolov10n/s/m/b/l/x.onnx --saveEngine=yolov10n/s/m/b/l/x.engine --fp16
# Predict with TensorRT
yolo predict model=yolov10n/s/m/b/l/x.engine
1
2
3
4
5
6
7
8
9
10
11

demo：

wget https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10s.pt
python app.py
# Please visit http://127.0.0.1:7860
1
2
3

六、数据、模型、训练后的所有文件

yolov10训练安全帽目标监测全部东西，下载看这里：

https://docs.qq.com/sheet/DUEdqZ2lmbmR6UVdU?tab=BB08J2
1

七、更大、更强

用yolov10m训练了一个模型，1280输入。