基于Yolov5的烟火检测——模型训练与C++实现部署

前言

1.系统环境是win10,显卡RTX3080;cuda10.2,cudnn7.1;OpenCV4.5;yolov5用的是5s的模型，2020年8月13日的发布v3.0这个版本; ncnn版本是20210525;C++ IDE vs2019.
2.使用NCNN作模型推理加速库，能更容易的使用GPU进行加速，代码不需要改动很大就可以移植到边缘设备或者移动端上。

一、数据集

1.数据集是使用网上收集的数据，大概2000多张左右。

2.标注工具用的是labelimg,标注了浓烟和火焰的位置。


3.手工标注是很费时费力的一件事情，如果是数据集一大这就很麻烦，但我一般的处理是先标注一批小的数据集之后，用yolov5l去训练一个初始模型用来标注，之后用初始模型去自动标注，按标注工具的格式来就可以了。

4.标注好的数据分xml和jpg两个目录。

5.标注好的数据集下载地址：https://download.csdn.net/download/matt45m/84044933 。这里只标注了2000多张，还有两万多张样本没有标注，所以就没有上传，如果想把模型用于项目上，可以私聊我要没有标注的数据集进行标注训练。

二、模型训练

1.我这里用的yolov5是2020年8月13日的发布v3.0这个版本，现在的yolov5已经到了m6了，但我的生产的环境一直是旧的版本，就没有改，模型的步骤可以参考我之前写的：https://blog.csdn.net/matt45m/article/details/118674051?spm=1001.2014.3001.5501 。
2.yolov5 v3的github地址：https://github.com/ultralytics/yolov5/tree/v3.0 。

三.模型转换

1.pt模型转onnx模型

# --weights: 训练得到的模型
python models/export.py --weights runs/exp/weights/best.pt

运行后，onnx模型保存为了runs/exp/weights/best.onnx，这个模型可以用OpenCV DNN理或者是onnxruntime推理。
但如果用DNN推理就要改下export.py源码，可参考之前我之前写的：https://blog.csdn.net/matt45m/article/details/118674207?spm=1001.2014.3001.5502

2.转ncnn模型可以参考nihui大佬的知乎文章 《详细记录u版YOLOv5目标检测ncnn实现》 。
3.但文章给出添加层的方式不适用于iOS系统，如果模型想在iOS上使用，则参考 https://github.com/Tencent/ncnn/wiki/add-custom-layer.zh 。

四.模型部署

1.如果要使用GPU加速，要安装Vulkan,步骤如下：
1.1.安装CMake，并把xxx\xxx\CMake\bin添加系统环境变量。
1.2.安装Vulkan各它的依赖库。

Vulkan
https://vulkan.lunarg.com/sdk/home
版本：VulkanSDK-1.2.141.2
直接点击安装，把之后验证是否安装成功，运行xxx\VulkanSDK\1.2.141.2\Bin\vkcube.exe，出现下面图像代表安装成功。

glfw
https://www.glfw.org/
把glfw-3.3.2.bin.WIN64复制到VulkanSDK\1.2.141.2\Third-Party

GLM
https://github.com/g-truc/glm/
把GLM复制到VulkanSDK\1.2.141.2\Third-Party
添加系统环境变量:

2.C++代码：

void drawObjectsYolo(const cv::Mat& bgr, cv::Mat& cv_dst, std::vector<ObjectFlag>& objects)
{
    static const char* class_names[] = {
        "fire", "smoke"
    };

    cv_dst = bgr.clone();


    for (size_t i = 0; i < objects.size(); i++)
    {
        const ObjectFlag& obj = objects[i];

        std::cout << "Object name: " << class_names[obj.label] << "  Proba = " << obj.prob * 100 << std::endl;

        cv::rectangle(cv_dst, obj.rect, cv::Scalar(255, 0, 0));

        char text[256];
        sprintf(text, "%s %.1f%%", class_names[obj.label], obj.prob * 100);

        int baseLine = 0;
        cv::Size label_size = cv::getTextSize(text, cv::FONT_HERSHEY_SIMPLEX, 1, 2, &baseLine);

        int x = obj.rect.x;
        int y = obj.rect.y - label_size.height - baseLine;
        if (y < 0)
            y = 0;
        if (x + label_size.width > cv_dst.cols)
            x = cv_dst.cols - label_size.width;


        cv::putText(cv_dst, text, cv::Point(x, y + label_size.height),
            cv::FONT_HERSHEY_SIMPLEX, 0.5, cv::Scalar(0, 0, 255));
    }
}

int initDetectionNet(std::string param_path,std::string model_path, ncnn::Net& td_net, bool use_gpu)
{
    bool has_gpu = false;

#if NCNN_VULKAN
    ncnn::create_gpu_instance();
    has_gpu = ncnn::get_gpu_count() > 0;
#endif

    bool to_use_gpu = has_gpu && use_gpu;
    td_net.opt.use_vulkan_compute = to_use_gpu;

    int rp = -1;
    int rm = -1;

    td_net.opt.use_fp16_arithmetic = true;
    rp = td_net.load_param(param_path.c_str());
    rm = td_net.load_model(model_path.c_str());

    if (rp < 0 || rm < 0)
    {
        return -70;
    }
    if (to_use_gpu)
    {
        return 1;
    }
    return 0;
}

std::string param_path = "models/fire_smoke.param";
std::string model_path = "models/fire_smoke.bin";

//path 图像目录
void detectionImages(std::string path = "fire_test_images")
{
    std::vector<std::string> filenames;
    cv::glob(path, filenames, false);
    ncnn::Net td_net;
    int rn = initDetectionNet(param_path, model_path, td_net, true);

    for (auto name : filenames)
    {
        cv::Mat cv_src = cv::imread(name);
        if (cv_src.empty())
        {
            continue;
        }
        std::vector<ObjectFlag> objects;
        double start = static_cast<double>(cv::getTickCount());
        detectYolov5(cv_src, td_net, objects, 4, 640, 0.25f, 0.45f);
        double time = ((double)cv::getTickCount() - start) / cv::getTickFrequency();
        std::cout << "Target detection time: " << time << "(s)" << std::endl;

        cv::Mat cv_dst;
        drawObjectsYolo(cv_src, cv_dst, objects);
        cv::imshow("cv_dst", cv_dst);
        cv::waitKey();
    }
}

void detectionVideo(std::string path = "fire.mp4")
{
    ncnn::Net td_net;
    int rn = initDetectionNet(param_path, model_path, td_net, true);

    cv::VideoCapture cap;
    //如果想要打开摄像头
    //cap.open(0);//摄像头索引
    cap.open(path);
    if (!cap.isOpened())
    {
        return;
    }
    cv::Mat cv_src;
    while (1)
    {
        cap >> cv_src;
        if (cv_src.empty())
        {
            break;
        }
        std::vector<ObjectFlag> objects;
        cv::Mat cv_dst;
        double start = static_cast<double>(cv::getTickCount());
        detectYolov5(cv_src, td_net, objects, 4, 640, 0.25f, 0.45f);
        double time = ((double)cv::getTickCount() - start) / cv::getTickFrequency();
        std::cout << "Target detection time: " << time << "(s)" << std::endl;
        drawObjectsYolo(cv_src, cv_dst, objects);
        cv::imshow("cv_dst", cv_dst);
        cv::waitKey(24);
    }
    cap.release();
}

int main(void)
{
    detectionVideo();
}

2.图像测试效果


3.视频检测：

4.源码: https://download.csdn.net/download/matt45m/84050225

五、备注

1.当前训练的图像只有2000多张，如果想要提高精度可以加数据样本，我这里还有几万张没有标注的数据，有兴趣的小伙伴可以问我要没有标注的数据集。
2.yolov5s的NCNN模型是可以部署到安卓或iOS上的。