安全帽佩戴检测算法模型训练详细流程

news2025/1/11 3:21:15

一、任务描述

实际施工现场需要对每个进出的人员进行安全帽监测，对未佩戴安全帽的人员平台进行风险告警，通知工作人员并记录下来。

主要包括三类目标物体：头盔（helmet），人（person）和人头（head）

二、实现流程

1、数据集：直接用的网络数据集，下载地址：

https://aistudio.baidu.com/aistudio/datasetdetail/50329

2、算法：用的算法框架是百度的PaddleDetection，代码下载地址：

https://github.com/PaddlePaddle/PaddleDetection

3、将数据集解压保存在helmet_data文件夹下

4、将数据集拆分成训练集、测试集和验证集，用txt文件记录保存下来。代码如下：

import logging
import random
import os
import os.path as osp
import xml.etree.ElementTree as ET

def list_files(dirname):
    """ 列出目录下所有文件（包括所属的一级子目录下文件）
    Args:
        dirname: 目录路径
    """
    def filter_file(f):
        if f.startswith('.'):
            return True
        return False

    all_files = list()
    dirs = list()
    for f in os.listdir(dirname):
        if filter_file(f):
            continue
        if osp.isdir(osp.join(dirname, f)):
            dirs.append(f)
        else:
            all_files.append(f)
    for d in dirs:
        for f in os.listdir(osp.join(dirname, d)):
            if filter_file(f):
                continue
            if osp.isdir(osp.join(dirname, d, f)):
                continue
            all_files.append(osp.join(d, f))
    return all_files

def is_pic(filename):
    """ 判断文件是否为图片格式
    Args:
        filename: 文件路径
    """
    suffixes = {'JPEG', 'jpeg', 'JPG', 'jpg', 'BMP', 'bmp', 'PNG', 'png'}
    suffix = filename.strip().split('.')[-1]
    if suffix not in suffixes:
        return False
    return True

def replace_ext(filename, new_ext):
    """ 替换文件后缀
    Args:
        filename: 文件路径
        new_ext: 需要替换的新的后缀
    """
    items = filename.split(".")
    items[-1] = new_ext
    new_filename = ".".join(items)
    return new_filename

def split_voc_dataset(dataset_dir, save_dir, val_percent=0.15, test_percent=0.15):
    # 注意图片目录和标注目录名已全部修改
    if not osp.exists(osp.join(dataset_dir, "JPEGImages")):
        logging.error("\'JPEGImages\' is not found in {}!".format(dataset_dir))
    if not osp.exists(osp.join(dataset_dir, "Annotations")):
        logging.error("\'Annotations\' is not found in {}!".format(
            dataset_dir))
    all_image_files = list_files(osp.join(dataset_dir, "JPEGImages"))

    image_anno_list = list()
    label_list = list()
    for image_file in all_image_files:
        if not is_pic(image_file): # 判断是否为图片格式
            continue
        anno_name = replace_ext(image_file, "xml")
        if osp.exists(osp.join(dataset_dir, "Annotations", anno_name)):
            image_anno_list.append([image_file, anno_name])
            try:
                tree = ET.parse(osp.join(dataset_dir, "Annotations", anno_name))
            except:
                raise Exception("文件{}不是一个良构的xml文件，请检查标注文件".format(
                    osp.join(dataset_dir, "Annotations", anno_name)))
            objs = tree.findall("object")
            for i, obj in enumerate(objs):
                cname = obj.find('name').text
                if not cname in label_list:
                    label_list.append(cname)
        else:
            logging.error("The annotation file {} doesn't exist!".format(anno_name))

    random.shuffle(image_anno_list) # 随机打乱
    image_num = len(image_anno_list) # 总图片数量
    val_num = int(image_num * val_percent) # 验证集数量
    test_num = int(image_num * test_percent) # 测试集数量
    train_num = image_num - val_num - test_num # 训练集数量

    train_image_anno_list = image_anno_list[:train_num] # 训练集样本
    val_image_anno_list = image_anno_list[train_num:train_num + val_num] # 验证集样本
    test_image_anno_list = image_anno_list[train_num + val_num:] # 测试集样本

    with open(osp.join(save_dir, 'train_list.txt'), mode='w', encoding='utf-8') as f:
        for x in train_image_anno_list:
            file = osp.join("JPEGImages", x[0])
            label = osp.join("Annotations", x[1])
            f.write('{} {}\n'.format(file, label))
    with open(osp.join(save_dir, 'val_list.txt'), mode='w', encoding='utf-8') as f:
        for x in val_image_anno_list:
            file = osp.join("JPEGImages", x[0])
            label = osp.join("Annotations", x[1])
            f.write('{} {}\n'.format(file, label))
    if len(test_image_anno_list):
        with open(osp.join(save_dir, 'test_list.txt'), mode='w', encoding='utf-8') as f:
            for x in test_image_anno_list:
                file = osp.join("JPEGImages", x[0])
                label = osp.join("Annotations", x[1])
                f.write('{} {}\n'.format(file, label))
    with open(osp.join(save_dir, 'labels.txt'), mode='w', encoding='utf-8') as f:
        for l in sorted(label_list):
            f.write('{}\n'.format(l))
    return image_anno_list, label_list

if __name__ == '__main__':
    dataset_dir = "/root/bigdata/pycharm_projects/PaddleDetection/data/helmet_data/helmet_train_data"
    split_voc_dataset(dataset_dir=dataset_dir, save_dir=dataset_dir)