1.这是一个检测是否佩戴安全帽的完整项目，包含了数据集，训练代码，检测模型代码，转换模型与如何部署模型并应用到项目上的完整过程。
2.训练和开发环境是win10,显卡RTX3080;cuda10.2,cudnn7.1;OpenCV4.5;yolov5用的是5s的模型，2020年8月13日的发布v3.0这个版本; ncnn版本是20210525;C++ IDE vs2019，Anaconda 3.5。
3.实现的效果：
4.C++部署模型的工程源码地址：https://download.csdn.net/download/matt45m/85384651
5.训练代码地址：https://download.csdn.net/download/matt45m/85384016
6.标注好的数据集地址：https://download.csdn.net/download/matt45m/85384715
1.首先要安装anaconda环境，我这里用的是Anaconda 3.5。
2.启动环境，尽量以管理员身份运行环境。
3. 创建环境
conda create --name yolov5 python=3.7 activate yolov5
2.安装依赖
git clone https://github.com/ultralytics/yolov5.git cd yolov5 pip install -r requirements.txt
或者
git clone https://github.com/ultralytics/yolov5.git cd yolov5 conda install pytorch torchvision cudatoolkit=10.2 -c pytorch pip install cython matplotlib tqdm opencv-python tensorboard scipy pillow onnx pyyaml pandas seaborn二、数据处理
1.标注数据
这里要判断三种形态，[人体]、[佩戴安全帽人的头部]、[没有佩戴安全帽的人头部]三个类别，使用的标注工具是labelImg，是按VOC2007数据标注格式进行标注的。
只标注佩戴在头上的安全帽，而且连着头部一起标注，拿在手中或者放在地方的不标注。
2.处理数据
标注的数据格式是voc,但这里用的yolo,要把数据转换成yolo的txt格式并分好训练集与测试集。
使用python脚本处理数据的generate_txt.py
import os
import glob
import argparse
import random
import xml.etree.ElementTree as ET
from PIL import Image
from tqdm import tqdm
def get_all_classes(xml_path):
xml_fns = glob.glob(os.path.join(xml_path, '*.xml'))
class_names = []
for xml_fn in xml_fns:
tree = ET.parse(xml_fn)
root = tree.getroot()
for obj in root.iter('object'):
cls = obj.find('name').text
class_names.append(cls)
return sorted(list(set(class_names)))
def convert_annotation(img_path, xml_path, class_names, out_path):
output = []
im_fns = glob.glob(os.path.join(img_path, '*.jpg'))
for im_fn in tqdm(im_fns):
if os.path.getsize(im_fn) == 0:
continue
xml_fn = os.path.join(xml_path, os.path.splitext(os.path.basename(im_fn))[0] + '.xml')
if not os.path.exists(xml_fn):
continue
img = Image.open(im_fn)
height, width = img.height, img.width
tree = ET.parse(xml_fn)
root = tree.getroot()
anno = []
xml_height = int(root.find('size').find('height').text)
xml_width = int(root.find('size').find('width').text)
if height != xml_height or width != xml_width:
print((height, width), (xml_height, xml_width), im_fn)
continue
for obj in root.iter('object'):
cls = obj.find('name').text
cls_id = class_names.index(cls)
xmlbox = obj.find('bndbox')
xmin = int(xmlbox.find('xmin').text)
ymin = int(xmlbox.find('ymin').text)
xmax = int(xmlbox.find('xmax').text)
ymax = int(xmlbox.find('ymax').text)
cx = (xmax + xmin) / 2.0 / width
cy = (ymax + ymin) / 2.0 / height
bw = (xmax - xmin) * 1.0 / width
bh = (ymax - ymin) * 1.0 / height
anno.append('{} {} {} {} {}'.format(cls_id, cx, cy, bw, bh))
if len(anno) > 0:
output.append(im_fn)
with open(im_fn.replace('.jpg', '.txt'), 'w') as f:
f.write('n'.join(anno))
random.shuffle(output)
train_num = int(len(output) * 0.9)
with open(os.path.join(out_path, 'train.txt'), 'w') as f:
f.write('n'.join(output[:train_num]))
with open(os.path.join(out_path, 'val.txt'), 'w') as f:
f.write('n'.join(output[train_num:]))
def parse_args():
parser = argparse.ArgumentParser('generate annotation')
parser.add_argument('--img_path', type=str, help='input image directory')
parser.add_argument('--xml_path', type=str, help='input xml directory')
parser.add_argument('--out_path', type=str, help='output directory')
args = parser.parse_args()
return args
if __name__ == '__main__':
args = parse_args()
class_names = get_all_classes(args.xml_path)
print(class_names)
convert_annotation(args.img_path, args.xml_path, class_names, args.out_path)
运行python脚本：
python generate_txt.py --img_path data/XXXXX/JPEGImages --xml_path data/XXXXX/Annotations --out_path data/XXXXX
运行之后会自动转换成归一化后的txt格式坐标，并在指定目录下能成train.txt 和val.txt这两个训练要用到的txt文件。
三、训练模型1.这里使用的模型是比s大一些，精度更好一些的m模型，model/yolov5m.yaml,更改nc数目。
2.在data目录下添加一个复制coco.yaml重新命名为helmet.yaml的训练数据配置文件。 ```markup # download command/URL (optional) download: bash data/scripts/get_voc.sh # 训练集txt与验证集txt路径 train: data/train.txt val: data/val.txt # 总类别数 nc: 3 # 类名 names: ['person', 'head', 'helmet']
3.训练参数
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path') # 权重文件，是否在使用预训练权重文件
parser.add_argument('--cfg', type=str, default='', help='model.yaml path') # 网络配置文件
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path') # 训练数据集目录
parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path') #超参数配置文件
parser.add_argument('--epochs', type=int, default=300) # 训练迭代次数
parser.add_argument('--batch-size', type=int, default=32, help='total batch size for all GPUs') # batch-size大小
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') # 训练图像大小
parser.add_argument('--rect', action='store_true', help='rectangular training') #矩形训练
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') # 是否接着上一次的日志权重继续训练
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') # 不保存
parser.add_argument('--notest', action='store_true', help='only test final epoch') # 不测试
parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') #超参数范围
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') #是否缓存图像
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') # 用GPU或者CPU进行训练
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') #是否多尺度训练
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') # 是否一个类别
parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') # 优化器先择
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
parser.add_argument('--log-imgs', type=int, default=16, help='number of images for W&B logging, max 100')
parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers') #win不能改，win上改不改都容易崩
parser.add_argument('--project', default='runs/train', help='save to project/name')
parser.add_argument('--name', default='exp', help='save to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
opt = parser.parse_args()
4.训练命令
- 单卡：
python train.py --cfg models/yolov5s.yaml --data data/ODID.yaml --hyp data/hyps/hyp.scratch.yaml --epochs 100 --multi-scale --device 0
- 多卡：
python train.py --cfg models/yolov5s.yaml --data data/ODID.yaml --hyp data/hyps/hyp.scratch.yaml --epochs 100 --multi-scale --device 0，1四、模型测试
1.测试模型
python detect.py --source TestImagesPath --weights ./weights/yolov5m.pt
2.输出检测的结果，检测一张图像大概在1左右。
1.转换模型
部署模型有很多种选择，部署语言可选C++或者python都可以，推理框架可以选择onnxruntime，libtorch,mnn,ncnn等，这里我使用我比较熟悉的C++与NCNN。
转换模型
python models/export.py --weights weights/yolov5m.pt --img 640 --batch 1
会生成yolov5m.onnx模型，onnx模型可以使用onnxruntime进行推理，也可以使用opencv的dnn进行推理。
2.转ncnn模型。
模型简化，模型简化我这里使用的还是旧的0.36版本，我试过用最新的简化版本，简化出来的模型还是存在一堆的带onnx前缀的op,在onnxruntime和转ncnn之后完全无法进行推理，搞了半天时间没有搞定直接放弃。onnx-simplifier：https://github.com/daquexian/onnx-simplifier
onnx转ncnn模型
onnx2ncnn yolov5s-sim.onnx yolov5s.param yolov5s.bin
- onnx转为 ncnn 模型，会输出很多 Unsupported slice step！，这是focus模块转换的报错.
这里可以直接参考nihui大佬的知乎文章对着改就成了，文章地址：https://zhuanlan.zhihu.com/p/275989233 。
3.NCNN推理代码，动态注册了YoloV5Focus层。
#include "YoloV5Detect.h"
class YoloV5Focus : public ncnn::Layer
{
public:
YoloV5Focus()
{
one_blob_only = true;
}
virtual int forward(const ncnn::Mat& bottom_blob, ncnn::Mat& top_blob, const ncnn::Option& opt) const
{
int w = bottom_blob.w;
int h = bottom_blob.h;
int channels = bottom_blob.c;
int outw = w / 2;
int outh = h / 2;
int outc = channels * 4;
top_blob.create(outw, outh, outc, 4u, 1, opt.blob_allocator);
if (top_blob.empty())
return -100;
#pragma omp parallel for num_threads(opt.num_threads)
for (int p = 0; p < outc; p++)
{
const float* ptr = bottom_blob.channel(p % channels).row((p / channels) % 2) + ((p / channels) / 2);
float* outptr = top_blob.channel(p);
for (int i = 0; i < outh; i++)
{
for (int j = 0; j < outw; j++)
{
*outptr = *ptr;
outptr += 1;
ptr += 2;
}
ptr += w;
}
}
return 0;
}
};
DEFINE_LAYER_CREATOR(YoloV5Focus)
int initYolov5Net(std::string& param_path, std::string& bin_path, ncnn::Net& yolov5_net,bool use_gpu)
{
bool has_gpu = false;
yolov5_net.clear();
//CPU相关设置(只实现了安卓端)
/// 0 = all cores enabled(default)
/// 1 = only little clusters enabled
/// 2 = only big clusters enabled
//ncnn::set_cpu_powersave(2);
//ncnn::set_omp_num_threads(ncnn::get_big_cpu_count());
#if NCNN_VULKAN
ncnn::create_gpu_instance();
has_gpu = ncnn::get_gpu_count() > 0;
#endif
yolov5_net.opt.use_vulkan_compute = (use_gpu && has_gpu);
yolov5_net.opt.use_bf16_storage = true;
//动态注册层
yolov5_net.register_custom_layer("YoloV5Focus", YoloV5Focus_layer_creator);
//读取模型
int rp = yolov5_net.load_param(param_path.c_str());
int rb = yolov5_net.load_model(bin_path.c_str());
if (rp < 0 || rb < 0)
{
return -1;
}
return 0;
}
static inline float sigmoid(float x)
{
return static_cast(1.f / (1.f + exp(-x)));
}
static void generateProposals(const ncnn::Mat& anchors, int stride, const ncnn::Mat& in_pad, const ncnn::Mat& feat_blob, float prob_threshold, std::vector
运行结果：
光头：
佩戴安全帽：
带帽子也可以检测出来不是安全帽
正常状态：
1.配置包括目录
2.配置lib目录
3.添加依赖项
GenericCodeGen.lib glslang.lib MachineIndependent.lib ncnn.lib OGLCompiler.lib onnxruntime.lib opencv_world450.lib OSDependent.lib SPIRV.lib VkLayer_utils.lib vulkan-1.lib
4.配置运行环境
