安全帽目标检测

安全帽数据集

这里我们使用的安全帽数据集是HelmentDetection，这是一个公开数据集，里面包含5000张voc标注格式的图像，分为三个类别，分别是 0: head 1: helmet 2: person

安全帽数据集下载地址、

我们将数据集下载后，首先需要将其转换为YOLO标注格式：

import os
import xml.etree.ElementTree as ET
import random
import shutil
def collect_classes(voc_annotations_dir):
    """
    收集 VOC 数据集中所有出现的类别。
    :param voc_annotations_dir: 包含 VOC 格式 XML 文件的目录
    :return: 类别名称的集合
    """
    classes = set()
    # 遍历所有 XML 文件
    for xml_file in os.listdir(voc_annotations_dir):
        if not xml_file.endswith('.xml'):
            continue
        # 构造输入文件路径
        voc_annotation_path = os.path.join(voc_annotations_dir, xml_file)
        # 解析 XML 文件
        tree = ET.parse(voc_annotation_path)
        root = tree.getroot()
        # 遍历每个目标并收集类别
        for obj in root.findall('object'):
            class_name = obj.find('name').text
            classes.add(class_name)
    return sorted(classes)
def create_class_mapping(classes):
    """
    创建类别名称到索引的映射字典。
    :param classes: 类别名称列表
    :return: 类别名称到索引的映射字典
    """
    return {class_name: idx for idx, class_name in enumerate(classes)}
def convert_voc_to_yolo(voc_annotation_path, yolo_annotation_path, class_mapping):
    """
    将单个 VOC 格式的标注文件转换为 YOLO 格式。
    :param voc_annotation_path: VOC 格式 XML 文件路径
    :param yolo_annotation_path: 输出的 YOLO 格式 TXT 文件路径
    :param class_mapping: 类别名称到索引的映射字典
    """
    # 解析 XML 文件
    tree = ET.parse(voc_annotation_path)
    root = tree.getroot()
    # 获取图像尺寸
    size = root.find('size')
    image_width = int(size.find('width').text)
    image_height = int(size.find('height').text)
    # 存储 YOLO 格式的目标信息
    yolo_lines = []
    # 遍历每个目标
    for obj in root.findall('object'):
        # 获取类别名称
        class_name = obj.find('name').text
        if class_name not in class_mapping:
            print(f"类别 '{class_name}' 不在映射中，跳过该目标。")
            continue
        # 获取类别索引
        class_index = class_mapping[class_name]
        # 获取边界框坐标
        bbox = obj.find('bndbox')
        xmin = float(bbox.find('xmin').text)
        ymin = float(bbox.find('ymin').text)
        xmax = float(bbox.find('xmax').text)
        ymax = float(bbox.find('ymax').text)
        # 转换为 YOLO 格式的归一化坐标 (x_center, y_center, width, height)
        x_center = (xmin + xmax) / 2.0 / image_width
        y_center = (ymin + ymax) / 2.0 / image_height
        width = (xmax - xmin) / image_width
        height = (ymax - ymin) / image_height
        # 添加到 YOLO 格式的行
        yolo_lines.append(f"{class_index} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
    # 写入 YOLO 格式的文件
    with open(yolo_annotation_path, 'w') as f:
        f.write("\n".join(yolo_lines))
def split_dataset(xml_files, train_ratio=0.8, val_ratio=0.2, test_ratio=0.1):
    """
    将数据集按照指定比例划分为训练集、验证集和测试集。
    :param xml_files: 所有 XML 文件的列表
    :param train_ratio: 训练集比例，默认 0.7
    :param val_ratio: 验证集比例，默认 0.2
    :param test_ratio: 测试集比例，默认 0.1
    :return: 划分后的三个集合（训练集、验证集、测试集）
    """
    random.shuffle(xml_files)  # 随机打乱文件顺序
    total_count = len(xml_files)
    # 计算各部分的分割点
    train_end = int(total_count * train_ratio)
    val_end = train_end + int(total_count * val_ratio)
    # 划分数据集
    train_set = xml_files[:train_end]
    val_set = xml_files[train_end:val_end]
    test_set = xml_files[val_end:]
    return train_set, val_set, test_set
def organize_files(voc_images_dir, voc_annotations_dir, output_dir, train_set, val_set, test_set):
    """
    将图片和标注文件按照划分结果组织到对应的文件夹中。
    :param voc_images_dir: 包含原始图片的目录
    :param voc_annotations_dir: 包含原始标注文件的目录
    :param output_dir: 输出目录
    :param train_set: 训练集文件名列表
    :param val_set: 验证集文件名列表
    :param test_set: 测试集文件名列表
    """
    # 创建输出目录结构
    images_dir = os.path.join(output_dir, "images")
    labels_dir = os.path.join(output_dir, "labels")
    os.makedirs(os.path.join(images_dir, "train"), exist_ok=True)
    os.makedirs(os.path.join(images_dir, "val"), exist_ok=True)
    os.makedirs(os.path.join(images_dir, "test"), exist_ok=True)
    os.makedirs(os.path.join(labels_dir, "train"), exist_ok=True)
    os.makedirs(os.path.join(labels_dir, "val"), exist_ok=True)
    os.makedirs(os.path.join(labels_dir, "test"), exist_ok=True)
    def copy_files(file_list, subset):
        """
        复制文件到对应的子集目录中。
        """
        for file_name in file_list:
            base_name = os.path.splitext(file_name)[0]
            # 复制图片
            src_image_path = os.path.join(voc_images_dir, base_name + ".png")
            dst_image_path = os.path.join(images_dir, subset, base_name + ".png")
            if os.path.exists(src_image_path):
                shutil.copy(src_image_path, dst_image_path)
            else:
                print(f"警告：未找到图片文件 {src_image_path}")
            # 复制标注文件
            src_label_path = os.path.join(voc_annotations_dir, base_name + ".txt")
            dst_label_path = os.path.join(labels_dir, subset, base_name + ".txt")
            if os.path.exists(src_label_path):
                shutil.copy(src_label_path, dst_label_path)
            else:
                print(f"警告：未找到标注文件 {src_label_path}")
    # 复制训练集、验证集和测试集
    copy_files(train_set, "train")
    copy_files(val_set, "val")
    copy_files(test_set, "test")
    print("图片和标注文件已成功组织到对应的文件夹中！")
def batch_convert_voc_to_yolo(voc_images_dir, voc_annotations_dir, output_dir):
    """
    批量将 VOC 格式的标注文件转换为 YOLO 格式，并划分数据集。
    :param voc_images_dir: 包含原始图片的目录
    :param voc_annotations_dir: 包含 VOC 格式 XML 文件的目录
    :param output_dir: 输出目录
    """
    # 收集所有类别
    classes = collect_classes(voc_annotations_dir)
    class_mapping = create_class_mapping(classes)
    print("发现以下类别：", classes)
    # 获取所有 XML 文件
    xml_files = [f for f in os.listdir(voc_annotations_dir) if f.endswith('.xml')]
    # 划分数据集
    train_set, val_set, test_set = split_dataset(xml_files)
    print(f"训练集数量：{len(train_set)}，验证集数量：{len(val_set)}，测试集数量：{len(test_set)}")
    # 创建临时标注输出目录
    temp_labels_dir = os.path.join(output_dir, "temp_labels")
    os.makedirs(temp_labels_dir, exist_ok=True)
    # 转换所有标注文件为 YOLO 格式
    for xml_file in xml_files:
        # 构造输入和输出文件路径
        voc_annotation_path = os.path.join(voc_annotations_dir, xml_file)
        yolo_annotation_path = os.path.join(temp_labels_dir, os.path.splitext(xml_file)[0] + ".txt")
        # 转换单个文件
        convert_voc_to_yolo(voc_annotation_path, yolo_annotation_path, class_mapping)
    # 组织文件到对应的文件夹中
    organize_files(voc_images_dir, temp_labels_dir, output_dir, train_set, val_set, test_set)
    # 删除临时标注目录
    shutil.rmtree(temp_labels_dir)
    print("数据集转换与组织完成！")
# 示例用法
if __name__ == "__main__":
    # 输入和输出目录
    voc_images_dir = "D:/project_mine/detection/datasets/HelmetDetection/train/JPEGImages"
    voc_annotations_dir = "D:/project_mine/detection/datasets/HelmetDetection/train/Annotations"
    output_dir = "D:/project_mine/detection/datasets/anquanmao"
    # 批量转换与划分
    batch_convert_voc_to_yolo(voc_images_dir, voc_annotations_dir, output_dir)

模型训练

模型训练很简单，我们只需要修改一下数据集配置文件即可

path: ../datasets/anquanmao # dataset root dir
train: images/train # train images (relative to 'path') 4 images
val: images/val # val images (relative to 'path') 4 images
test: # test images (optional)
# Classes
names:
  0: head
  1: helmet
  2: person

随后，即可开启模型训练

from ultralytics import YOLO
model=YOLO("yolo11.yaml")#.load("yolo11n.pt")#.load("yolo11s.pt")  # build from YAML and transfer weights
# Train the model
results = model.train(data="anquanmao.yaml",
                      epochs=60,
                      batch=16,       # 根据GPU显存调整（T4建议batch=8）
                      imgsz=640,
                      device="0",     # 指定GPU ID
                      optimizer="AdamW",
                      lr0=1e-4,
                      warmup_epochs=4,
                      label_smoothing=0.1,
                      amp=True)

结果如下：

这里的mAP值不高是因为，人这个类别似乎缺失了，导致人这个类别的AP值为0，单独来看，头和安全帽的AP值均在0.5以上。

ONNX模型推理

为使其能够具有更好的扩展性，我们将原本的pt模型文件转换为ONNX格式，随后进行推理，代码如下：

import time
import cv2
import numpy as np
import onnxruntime
#from my_utils.detect.utils import detections_dog
from utils import xywh2xyxy, multiclass_nms
class SafeHat:
    #初始化YOLO模型
    def __init__(self, path, conf_thres=0.7, iou_thres=0.5):
        self.conf_threshold = conf_thres
        self.iou_threshold = iou_thres
        # Initialize model
        self.initialize_model(path)
    #调用推理
    def __call__(self, image):
        return self.detect_objects(image)
    #加载模型并获取模型的输入与输出结构
    def initialize_model(self, path):
        self.session = onnxruntime.InferenceSession(path,providers=onnxruntime.get_available_providers())
        self.get_input_details()
        self.get_output_details()
    #执行模型推理过程
    def detect_objects(self, image):
        input_tensor = self.prepare_input(image)
        # Perform inference on the image
        outputs = self.inference(input_tensor)
        self.boxes, self.scores, self.class_ids = self.process_output(outputs)
        return self.boxes, self.scores, self.class_ids
    #前处理操作
    def prepare_input(self, image):
        self.img_height, self.img_width = image.shape[:2]
        input_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        # Resize input image
        input_img = cv2.resize(input_img, (self.input_width, self.input_height))
        # Scale input pixel values to 0 to 1
        input_img = input_img / 255.0
        input_img = input_img.transpose(2, 0, 1)
        input_tensor = input_img[np.newaxis, :, :, :].astype(np.float32)
        return input_tensor
    #具体使用onnx推理
    def inference(self, input_tensor):
        outputs = self.session.run(self.output_names, {self.input_names[0]: input_tensor})
        return outputs
    #后处理操作
    def process_output(self, output):
        predictions = np.squeeze(output[0]).T
        # Filter out object confidence scores below threshold
        scores = np.max(predictions[:, 4:], axis=1)
        predictions = predictions[scores > self.conf_threshold, :]
        scores = scores[scores > self.conf_threshold]
        if len(scores) == 0:
            return [], [], []
        # Get the class with the highest confidence
        class_ids = np.argmax(predictions[:, 4:], axis=1)
        # Get bounding boxes for each object
        boxes = self.extract_boxes(predictions)
        indices = multiclass_nms(boxes, scores, class_ids, self.iou_threshold)
        return boxes[indices], scores[indices], class_ids[indices]
    #box转换，包含尺度变换与xywh转换
    def extract_boxes(self, predictions):
        # Extract boxes from predictions
        boxes = predictions[:, :4]
        boxes = self.rescale_boxes(boxes)
        boxes = xywh2xyxy(boxes)
        return boxes
    #尺度变换
    def rescale_boxes(self, boxes):
        # Rescale boxes to original image dimensions
        input_shape = np.array([self.input_width, self.input_height, self.input_width, self.input_height])
        boxes = np.divide(boxes, input_shape, dtype=np.float32)
        boxes *= np.array([self.img_width, self.img_height, self.img_width, self.img_height])
        return boxes
    def get_input_details(self):
        model_inputs = self.session.get_inputs()
        self.input_names = [model_inputs[i].name for i in range(len(model_inputs))]
        self.input_shape = model_inputs[0].shape
        self.input_height = self.input_shape[2]
        self.input_width = self.input_shape[3]
    def get_output_details(self):
        model_outputs = self.session.get_outputs()
        self.output_names = [model_outputs[i].name for i in range(len(model_outputs))]
#     def draw_detections(self, image, draw_scores=True, mask_alpha=0.4):
#
#         return detections_dog(image, self.boxes, self.scores,
#                               self.class_ids, mask_alpha)
#
#
# if __name__ == "__main__":
#     model_path = "anquanmao.onnx"  # 替换为你的 TFLite 模型路径
#     image_path = "down_head.png"          # 替换为你的测试图像路径
#
#     # 初始化模型
#     detector = YOLODet(model_path)
#
#     # 加载图像
#     image = cv2.imread(image_path)
#
#     # 检测对象
#     boxes, scores, class_ids = detector(image)
#
#     # 绘制检测结果
#     result_image = detector.draw_detections(image)
#
#     # 显示结果
#     cv2.imshow("Detections", result_image)
#     cv2.waitKey(0)
#     cv2.destroyAllWindows()