七. 部署YOLOv8检测器-deploy-yolov8-basic

发布于:2024-09-19 ⋅ 阅读:(8) ⋅ 点赞:(0)

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》,链接。记录下个人学习笔记,仅供自己参考

本次课程我们来学习课程第七章—部署YOLOv8检测器,一起来学习 YOLOv8 的部署

课程大纲可以看下面的思维导图

在这里插入图片描述

0. 简述

本小节目标:学习 YOLOv8 的部署

这节我们主要学习 YOLOv8 模型的部署

下面我们开始本次课程的学习🤗

1. 案例运行

在正式开始课程之前,博主先带大家跑通 7.3-deploy-yolo-basic 这个小节的案例

源代码获取地址:https://github.com/kalfazed/tensorrt_starter

首先大家需要把 tensorrt_starter 这个项目给 clone 下来,指令如下:

git clone https://github.com/kalfazed/tensorrt_starter.git

也可手动点击下载,点击右上角的 Code 按键,将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码(注意代码下载于 2024/7/14 日,若有改动请参考最新

整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV,大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装,博主这里不再赘述

假设你的项目、环境准备完成,下面我们一起来运行下 7.3-deploy-yolo-basic 小节案例代码

开始之前我们需要创建几个文件夹,在 tensorrt_starter/chapter7-deploy-yolo-detection/7.3-deploy-yolo-basic 小节中创建一个 models 文件夹,接着在 models 文件夹下创建一个 onnx 和 engine 文件夹,总共三个文件夹需要创建

创建完后 7.3 小节整个目录结构如下:

在这里插入图片描述

接着我们要导出 YOLOv8 的 ONNX 模型,流程如下:

1. 准备一个虚拟环境(如果你已经有 ultralytics 的虚拟环境,可以跳过)

conda create -n yolov8 python=3.9
conda activate yolov8
git clone https://github.com/ultralytics/ultralytics.git
cd ultralytics
pip install -e .

2. 下载 YOLOv8

git clone https://github.com/ultralytics/ultralytics.git

3. 修改代码,添加 transpose 节点

# ========== head.py ==========

# ultralytics/nn/modules/head.py第61行,forward函数
# return y if self.export else (y, x)
# 修改为:

return y.permute(0, 2, 1) if self.export else (y, x)

Note:如果不添加 tranpose 节点,原始 yolov8 的输出格式是 [n, feature, bbox],这个格式不方便我们后续在 C++ 中做处理,我们希望对每一个 bbox 做处理,所以希望每一个 bbox 内部的数据在内存上是连续的,也就是说我们希望 yolov8 的输出格式是 [n, bbox, feature],而这个变换我们可以通过上面的代码添加 tranpose 节点完成

4. 导出 onnx 模型,在 ultralytics-main 新建导出文件 export.py 内容如下:

from ultralytics import YOLO

model = YOLO("yolov8x.pt")

success = model.export(format="onnx", dynamic=False, simplify=True)

5. 执行 export.py 导出 yolov8x 的 onnx 模型

cd ultralytics-main
conda activate yolov8
python export.py

执行后输出如下图所示:

在这里插入图片描述

导出的 ONNX 如下图所示:

在这里插入图片描述

大家也可以点击 here 下载博主准备好的 ONNX 模型

接着大家需要把刚导出好的 ONNX 模型放在 tensorrt_starter/chapter7-deploy-yolo-detection/7.3-deploy-yolo-basic/models/onnx 文件夹下

ONNX 模型准备好后,我们把韩君老师生成的推理结果给删除,方便后续查看我们自己推理的结果,指令如下:

cd 7.3-deploy-yolo-basic/data/result
rm -rf *

然后我们需要利用 ONNX 生成对应的 engine 完成推理,在此之前我们需要修改下整体的 Makefile.config,指定一些库的路径:

# tensorrt_starter/config/Makefile.config
# CUDA_VER                    :=  11
CUDA_VER                    :=  11.6
    
# opencv和TensorRT的安装目录
OPENCV_INSTALL_DIR          :=  /usr/local/include/opencv4
# TENSORRT_INSTALL_DIR        :=  /mnt/packages/TensorRT-8.4.1.5
TENSORRT_INSTALL_DIR        :=  /home/jarvis/lean/TensorRT-8.6.1.6

Note:大家查看自己的 CUDA 是多少版本,修改为对应版本即可,另外 OpenCV 和 TensorRT 修改为你自己安装的路径即可

接着我们就可以来执行编译,指令如下:

make -j64

输出如下:

在这里插入图片描述

接着执行:

./bin/trt-infer

输出如下:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

我们这里可以看到每张图片的一个推理结果,并且推理后的图片保存在 data/result 文件夹下,大家可以查看,如下图所示:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

如果大家能够看到上述输出结果,那就说明本小节案例已经跑通,下面我们就来看看具体的代码实现

2. 补充说明

在分析代码之前我们先来看下韩君老师在这小节中写的 README 文档

这节案例实现了一个基本的 yolov8 的部署,整个代码是 6.2-deploy-classification-advanced 小节的扩展,大家感兴趣的可以看看:六. 部署分类器-deploy-classification-advanced, 由于我们创建的推理框架是模块化的,所以我们在 6.2 小节代码的基础上需要扩展的内容包括:

  • src/cpp/trt_detector.cpp
  • src/cpp/trt_worker.cpp
  • src/cpp/trt_model.cpp
  • include/trt_detector.hpp
  • include/trt_worker.hpp

对于 worker 的修改,我们只需要修改一下构造函数,添加一个针对 detection task 的成员变量 m_detector,之后通过 m_detector 来做推理就可以了

对于 detector 我们只需要关注 preprocess 和 postprocess 的实现就可以了,preprocess 部分我们在 7.2 小节中已经实现好了一个生成 letterbox 的方法,也就是上节课程提到的 warpAffine。因此 postprocess 才是我们这节课真正需要关注的,我们需要自己实现一个 decode 和 nms 的算法

yolov8 的 postprocess 过程主要包含:

  • 1. decode
    • 把 bbox 从输出 tensor 拿出来并进行 decode,把获取的 bbox 放入 m_bboxes 中
  • 2. nms
    • 把 decode 得到的 m_bboxes 根据 nms threshold 进行 NMS 处理
  • 3. draw_bbox
    • 把最终得到的 bbox 绘制在原图中

下面我们来看每个部分具体的实现过程:

1. decode 部分

我们需要做的就是将 [batch, bboxes, ch] 转换为 vector

  • 1. 从每一个 bbox 对应的 ch 中获取 cx, cy, width, height
  • 2. 对每一个 bbox 对应的 ch 中找到最大的 class label,可以使用 std::max_element
  • 3. 将 cx, cy, width, height 转换为 x0, y0, x1, y1
  • 4. 根据 preprocess 中的 affine matrix 逆变换将 bbox 映射回原图尺寸
  • 5. 将转换好的 x0, y0, x1, y1 以及 confidence 和 class label 存入 box 中,并 push 到 m_bboxes 中准备下面的 NMS 处理

2. NMS 部分

  • 1. 写一个 IoU 计算的 lambda 函数
  • 2. 将 m_bboxes 中所有的数据按照 confidence 从高到低进行排序
  • 3. 应用 NMS 算法过滤重叠框

3. draw_bbox

  • 1. 通过 class label 获取 name
  • 2. 通过 class label 获取 color
  • 3. 调用 opencv 的 cv::rectangle 绘制 bbox 的框
  • 4. 调用 opencv 的 cv::putText 绘制 bbox 的类别和置信度

OK,以上就是 yolov8 整个后处理过程了,理解了算法之后实现起来就很容易了,更多细节我们在后续代码分析中再说

Note:不论是什么模型的部署,我们首先需要去梳理它的前后处理,理解它的前后处理后我们才能够更好的完成后续的部署工作。因此在这个小节的 YOLOv8 模型部署中大家需要提前对 YOLOv8 的前后处理熟悉,这里博主之前有简单分析过,大家感兴趣的可以看看:YOLOv8推理详解及部署实现

3. 代码分析

3.1 main.cpp

我们先从 main.cpp 看起:

#include "trt_model.hpp"
#include "trt_logger.hpp"
#include "trt_worker.hpp"
#include "utils.hpp"

using namespace std;

int main(int argc, char const *argv[])
{
    /*这么实现目的在于让调用的整个过程精简化*/
    string onnxPath    = "models/onnx/yolov8x.onnx";

    auto level         = logger::Level::VERB;
    auto params        = model::Params();

    params.img         = {640, 640, 3};
    params.task        = model::task_type::DETECTION;
    params.dev         = model::device::GPU;
    params.prec        = model::precision::FP16;

    // 创建一个worker的实例, 在创建的时候就完成初始化
    auto worker   = thread::create_worker(onnxPath, level, params);

    // 根据worker中的task类型进行推理
    worker->inference("data/source/car.jpg");
    worker->inference("data/source/crowd.jpg");
    worker->inference("data/source/crossroad.jpg");
    worker->inference("data/source/airport.jpg");
    worker->inference("data/source/bedroom.jpg");

    return 0;
}

相比于 6.2 小节的 classification 而言,主要的不同点在于 params 的 task 设置的是 DETECTION,其它的都和之前的一样,创建一个 worker 实例然后调用 inference 方法去进行推理

3.2 trt_detector.hpp

在 main.cpp 中我们指定 task 为 DETECTION 后在 create_worker 时创建的是一个 m_detector 类的 shared_ptr 对象,和前面的 m_classifier 类似:

Worker::Worker(string onnxPath, logger::Level level, model::Params params) {
    m_logger = logger::create_logger(level);

    // 这里根据task_type选择创建的trt_model的子类,今后会针对detection, segmentation扩充
    if (params.task == model::task_type::CLASSIFICATION) 
        m_classifier = model::classifier::make_classifier(onnxPath, level, params);
    else if (params.task == model::task_type::DETECTION) 
        m_detector = model::detector::make_detector(onnxPath, level, params);
}

我们来看下具体的 Detector 类的定义:

#ifndef __TRT_DETECTOR_HPP__
#define __TRT_DETECTOR_HPP__

#include <memory>
#include <vector>
#include <string>
#include "NvInfer.h"
#include "trt_logger.hpp"
#include "trt_model.hpp"

namespace model{

namespace detector {

enum model {
    YOLOV5,
    YOLOV8
};

struct bbox {
    float x0, x1, y0, y1;
    float confidence;
    bool  flg_remove;
    int   label;
    
    bbox() = default;
    bbox(float x0, float y0, float x1, float y1, float conf, int label) : 
        x0(x0), y0(y0), x1(x1), y1(y1), 
        confidence(conf), flg_remove(false), 
        label(label){};
};

class Detector : public Model{

public:
    // 这个构造函数实际上调用的是父类的Model的构造函数
    Detector(std::string onnx_path, logger::Level level, Params params) : 
        Model(onnx_path, level, params) {};

public:
    // 这里detection自己实现了一套前处理/后处理,以及内存分配的初始化
    virtual void setup(void const* data, std::size_t size) override;
    virtual void reset_task() override;
    virtual bool preprocess_cpu() override;
    virtual bool preprocess_gpu() override;
    virtual bool postprocess_cpu() override;
    virtual bool postprocess_gpu() override;

private:
    std::vector<bbox> m_bboxes;
    int m_inputSize; 
    int m_imgArea;
    int m_outputSize;
};

// 外部调用的接口
std::shared_ptr<Detector> make_detector(
    std::string onnx_path, logger::Level level, Params params);

}; // namespace detector
}; // namespace model

#endif //__TRT_DETECTOR_HPP__

和 Classifier 类似通过 make_detector 函数接口创建一个 Detector 的实例对象,Detector 继承自 Model 类,自己实现了一套前处理和后处理以及内存分配的初始化。此外这里还定义了一个 bbox 的结构体对象,用来存储 detection 模型输出的每一个检测框的结果

这里有一个点需要大家注意,在 Detector 的成员变量中有一个 m_bboxes,我们希望每次推理完一张图片利用 m_bboxes 可视化后就将它清空,方便下一张图片的推理结果保存,因此我们在 Model 类的 init_model 方法中如果 model 已经创建则我们会调用 reset_task 方法来情况一些类的成员变量:

void Model::init_model() {
    /* 一个model的engine, context这些一旦创建好了,当多次调用这个模型的时候就没必要每次都初始化了*/
    if (m_context == nullptr){
        if (!fileExists(m_enginePath)){
            LOG("%s not found. Building trt engine...", m_enginePath.c_str());
            build_engine();
        } else {
            LOG("%s has been generated! loading trt engine...", m_enginePath.c_str());
            load_engine();
        }
    }else{
        m_timer->init();
        reset_task();
    }
}

reset_task 是一个纯虚函数,所以继承 Model 类的子类都需要实现该函数,而 Detector 类中实现的主要是清空 m_bboxes,如下所示:

void Detector::reset_task(){
    m_bboxes.clear();
}

3.2 trt_detector.cpp

main.cpp 中创建完 worker 对象后会进行 inference 推理,代码如下所示:

void Worker::inference(string imagePath) {
    if (m_classifier != nullptr) {
        m_classifier->init_model();
        m_classifier->load_image(imagePath);
        m_classifier->inference();
    }

    if (m_detector != nullptr) {
        m_detector->init_model();
        m_detector->load_image(imagePath);
        m_detector->inference();
    }
}

init_model 会创建 engine 并进行一些初始化,而 load_image 则主要是加载图片,inference 方法也是一样根据 dev 参数的设置选择在 CPU 或 GPU 上进行前后处理操作:

/* 
    可以根据情况选择是否在CPU上跑pre/postprocess
    对于一些edge设备,为了最大化GPU利用效率,我们可以考虑让CPU做一些pre/postprocess,让其执行与GPU重叠
*/
void Model::inference() {
    if (m_params->dev == CPU) {
        preprocess_cpu();
    } else {
        preprocess_gpu();
    }

    enqueue_bindings();

    if (m_params->dev == CPU) {
        postprocess_cpu();
    } else {
        postprocess_gpu();
    }
}

主要的不同点就在于不同的 task 的前后处理函数的实现不一样,我们先来看下 Detection task 前处理函数:

bool Detector::preprocess_gpu() {
    /*Preprocess -- yolo的预处理并没有mean和std,所以可以直接skip掉mean和std的计算 */

    /*Preprocess -- 读取数据*/
    m_inputImage = cv::imread(m_imagePath);
    if (m_inputImage.data == nullptr) {
        LOGE("ERROR: file not founded! Program terminated"); return false;
    }
    
    /*Preprocess -- 测速*/
    m_timer->start_gpu();

    /*Preprocess -- 使用GPU进行warpAffine, 并将结果返回到m_inputMemory中*/
    preprocess::preprocess_resize_gpu(m_inputImage, m_inputMemory[1],
                                   m_params->img.h, m_params->img.w, 
                                   preprocess::tactics::GPU_WARP_AFFINE);

    m_timer->stop_gpu("preprocess(GPU)");
    return true;
}

这里 preprocess 指定的 tactics 是 GPU_WARP_AFFINE,也就是我们上节课提到的 warpAffine 仿射变换,具体的实现细节我们之前分析过,这边就不再赘述了

我们重点来看下后处理部分的实现:

bool Detector::postprocess_cpu() {
    m_timer->start_cpu();

    /*Postprocess -- 将device上的数据移动到host上*/
    int output_size = m_outputDims.d[1] * m_outputDims.d[2] * sizeof(float);
    CUDA_CHECK(cudaMemcpyAsync(m_outputMemory[0], m_outputMemory[1], output_size, cudaMemcpyKind::cudaMemcpyDeviceToHost, m_stream));
    CUDA_CHECK(cudaStreamSynchronize(m_stream));

    /*Postprocess -- yolov8的postprocess需要做的事情*/
    /*
     * 1. 把bbox从输出tensor拿出来,并进行decode,把获取的bbox放入到m_bboxes中
     * 2. 把decode得到的m_bboxes根据nms threshold进行NMS处理
     * 3. 把最终得到的bbox绘制到原图中
     */

    float conf_threshold = 0.25; //用来过滤decode时的bboxes
    float nms_threshold  = 0.45;  //用来过滤nms时的bboxes

    /*Postprocess -- 1. decode*/
    /*
     * 我们需要做的就是将[batch, bboxes, ch]转换为vector<bbox>
     * 几个步骤:
     * 1. 从每一个bbox中对应的ch中获取cx, cy, width, height
     * 2. 对每一个bbox中对应的ch中,找到最大的class label, 可以使用std::max_element
     * 3. 将cx, cy, width, height转换为x0, y0, x1, y1
     * 4. 因为图像是经过resize了的,所以需要根据resize的scale和shift进行坐标的转换(这里面可以根据preprocess中的到的affine matrix来进行逆变换)
     * 5. 将转换好的x0, y0, x1, y1,以及confidence和classness给存入到box中,并push到m_bboxes中,准备接下来的NMS处理
     */
    int    boxes_count = m_outputDims.d[1];
    int    class_count = m_outputDims.d[2] - 4;
    float* tensor;

    float  cx, cy, w, h, obj, prob, conf;
    float  x0, y0, x1, y1;
    int    label;

    for (int i = 0; i < boxes_count; i ++){
        tensor = m_outputMemory[0] + i * m_outputDims.d[2];
        label  = max_element(tensor + 4, tensor + 4 + class_count) - (tensor + 4);
        conf   = tensor[4 + label];
        if (conf < conf_threshold) 
            continue;

        cx = tensor[0];
        cy = tensor[1];
        w  = tensor[2];
        h  = tensor[3];
        
        x0 = cx - w / 2;
        y0 = cy - h / 2;
        x1 = x0 + w;
        y1 = y0 + h;

        // 通过warpaffine的逆变换得到yolo feature中的x0, y0, x1, y1在原图上的坐标
        preprocess::affine_transformation(preprocess::affine_matrix.reverse, x0, y0, &x0, &y0);
        preprocess::affine_transformation(preprocess::affine_matrix.reverse, x1, y1, &x1, &y1);
        
        bbox yolo_box(x0, y0, x1, y1, conf, label);
        m_bboxes.emplace_back(yolo_box);
    }
    LOGD("the count of decoded bbox is %d", m_bboxes.size());
    

    /*Postprocess -- 2. NMS*/
    /* 
     * 几个步骤:
     * 1. 做一个IoU计算的lambda函数
     * 2. 将m_bboxes中的所有数据,按照confidence从高到低进行排序
     * 3. 最终希望是对于每一个class,我们都只有一个bbox,所以对同一个class的所有bboxes进行IoU比较,
     *    选取confidence最大。并与其他的同类bboxes的IoU的重叠率最大的同时IoU > IoU threshold
     */

    vector<bbox> final_bboxes;
    final_bboxes.reserve(m_bboxes.size());
    std::sort(m_bboxes.begin(), m_bboxes.end(), 
              [](bbox& box1, bbox& box2){return box1.confidence > box2.confidence;});

    /*
     * nms在网上有很多实现方法,其中有一些是根据nms的值来动态改变final_bboex的大小(resize, erease)
     * 这里需要注意的是,频繁的对vector的大小的更改的空间复杂度会比较大,所以尽量不要这么做
     * 可以通过给bbox设置skip计算的flg来调整。
    */
    for(int i = 0; i < m_bboxes.size(); i ++){
        if (m_bboxes[i].flg_remove)
            continue;
        
        final_bboxes.emplace_back(m_bboxes[i]);
        for (int j = i + 1; j < m_bboxes.size(); j ++) {
            if (m_bboxes[j].flg_remove)
                continue;

            if (m_bboxes[i].label == m_bboxes[j].label){
                if (iou_calc(m_bboxes[i], m_bboxes[j]) > nms_threshold)
                    m_bboxes[j].flg_remove = true;
            }
        }
    }
    LOGD("the count of bbox after NMS is %d", final_bboxes.size());


    /*Postprocess -- draw_bbox*/
    /*
     * 几个步骤
     * 1. 通过label获取name
     * 2. 通过label获取color
     * 3. cv::rectangle
     * 4. cv::putText
     */
    string tag   = "detect-" + getPrec(m_params->prec);
    m_outputPath = changePath(m_imagePath, "../result", ".png", tag);

    int   font_face  = 0;
    float font_scale = 0.001 * MIN(m_inputImage.cols, m_inputImage.rows);
    int   font_thick = 2;
    int   baseline;
    CocoLabels labels;

    LOG("\tResult:");
    for (int i = 0; i < final_bboxes.size(); i ++){
        auto box = final_bboxes[i];
        auto name = labels.coco_get_label(box.label);
        auto rec_color = labels.coco_get_color(box.label);
        auto txt_color = labels.get_inverse_color(rec_color);
        auto txt = cv::format({"%s: %.2f%%"}, name.c_str(), box.confidence * 100);
        auto txt_size = cv::getTextSize(txt, font_face, font_scale, font_thick, &baseline);

        int txt_height = txt_size.height + baseline + 10;
        int txt_width  = txt_size.width + 3;

        cv::Point txt_pos(round(box.x0), round(box.y0 - (txt_size.height - baseline + font_thick)));
        cv::Rect  txt_rec(round(box.x0 - font_thick), round(box.y0 - txt_height), txt_width, txt_height);
        cv::Rect  box_rec(round(box.x0), round(box.y0), round(box.x1 - box.x0), round(box.y1 - box.y0));

        cv::rectangle(m_inputImage, box_rec, rec_color, 3);
        cv::rectangle(m_inputImage, txt_rec, rec_color, -1);
        cv::putText(m_inputImage, txt, txt_pos, font_face, font_scale, txt_color, font_thick, 16);

        LOG("%+20s detected. Confidence: %.2f%%. Cord: (x0, y0):(%6.2f, %6.2f), (x1, y1)(%6.2f, %6.2f)", 
            name.c_str(), box.confidence * 100, box.x0, box.y0, box.x1, box.y1);

    }
    LOG("\tSummary:");
    LOG("\t\tDetected Objects: %d", final_bboxes.size());
    LOG("");

    m_timer->stop_cpu<timer::Timer::ms>("postprocess(CPU)");

    cv::imwrite(m_outputPath, m_inputImage);
    LOG("\tsave image to %s", m_outputPath.c_str());

    m_timer->show();
    printf("\n");

    return true;
}

上述代码展示了将 yolov8 模型在 gpu 上推理得到的结果迁移到 cpu 上,并进行一系列后处理操作,包括解码、非极大值抑制(NMS)以及绘制检测结果到图像中,该函数的最终结果是将带有检测边界框的图像保存到指定路径

下面是该函数的具体实现分析:(from ChatGPT)

1. 将数据从 GPU 传输到 CPU

m_timer->start_cpu();
int output_size = m_outputDims.d[1] * m_outputDims.d[2] * sizeof(float);
CUDA_CHECK(cudaMemcpyAsync(m_outputMemory[0], m_outputMemory[1], output_size, cudaMemcpyKind::cudaMemcpyDeviceToHost, m_stream));
CUDA_CHECK(cudaStreamSynchronize(m_stream));
  • 计算输出大小
    • output_size 计算了输出 tensor 的数据大小
  • 数据传输
    • 使用 cudaMemcpyAsync 将数据从 device(m_outputMemory[1]) 复制到 host(m_outputMemory[0]) 上
    • 使用 cudaStreamSynchronize 确保数据传输完成,防止后续操作在数据未完全传输时进行。

2. 解码(Decode)过程

float conf_threshold = 0.25;
float nms_threshold  = 0.45;

int boxes_count = m_outputDims.d[1];
int class_count = m_outputDims.d[2] - 4;
float* tensor;

float cx, cy, w, h, obj, prob, conf;
float x0, y0, x1, y1;
int label;

for (int i = 0; i < boxes_count; i ++){
    tensor = m_outputMemory[0] + i * m_outputDims.d[2];
    label  = max_element(tensor + 4, tensor + 4 + class_count) - (tensor + 4);
    conf   = tensor[4 + label];
    if (conf < conf_threshold) 
        continue;

    cx = tensor[0];
    cy = tensor[1];
    w  = tensor[2];
    h  = tensor[3];
    
    x0 = cx - w / 2;
    y0 = cy - h / 2;
    x1 = x0 + w;
    y1 = y0 + h;

    preprocess::affine_transformation(preprocess::affine_matrix.reverse, x0, y0, &x0, &y0);
    preprocess::affine_transformation(preprocess::affine_matrix.reverse, x1, y1, &x1, &y1);
    
    bbox yolo_box(x0, y0, x1, y1, conf, label);
    m_bboxes.emplace_back(yolo_box);
}
LOGD("the count of decoded bbox is %d", m_bboxes.size());

详细步骤如下:

2.1 阈值设定

  • conf_threshold:用于过滤低置信度的边界框
  • nms_threshold:用于 NMS 过滤重叠边界框

2.2 初始化变量

  • boxes_count:边界框总数量
  • class_count:类别数量,对于 coco 数据集而言是 80 个类别

2.3 循环遍历每个边界框

  • 获取当前边界框的指针:tensor 指向当前边界框的数据
  • 确定类别标签:
    • 使用 std::max_element 在类别置信度区域找到最大值,确定边界框的类别 label
  • 置信度过滤:
    • 如果当前边界框的置信度 conf 低于 conf_threshold,则跳过该边界框

2.4 坐标转换

  • 从中心坐标转换为左上和右下坐标
    • x0 = cx - w / 2
    • y0 = cy - h / 2
    • x1 = x0 + w
    • y1 = y0 + h

2.5 逆仿射变换

  • 使用预处理阶段得到的逆仿射矩阵,将边界框的坐标从 640x640 尺寸映射回原图尺寸,确保绘制到原图时位置准确

2.6 存储一个边界框

  • 创建一个 bbox 对象,包含转换后的坐标、置信度和类别标签,并将其加入到 m_bboxes 容器中

2.7 日志记录

  • 打印解码后边界框的数量,便于调试和性能监控

3. 非极大值抑制(NMS)

vector<bbox> final_bboxes;
final_bboxes.reserve(m_bboxes.size());
std::sort(m_bboxes.begin(), m_bboxes.end(), 
          [](bbox& box1, bbox& box2){return box1.confidence > box2.confidence;});

for(int i = 0; i < m_bboxes.size(); i ++){
    if (m_bboxes[i].flg_remove)
        continue;
    
    final_bboxes.emplace_back(m_bboxes[i]);
    for (int j = i + 1; j < m_bboxes.size(); j ++) {
        if (m_bboxes[j].flg_remove)
            continue;

        if (m_bboxes[i].label == m_bboxes[j].label){
            if (iou_calc(m_bboxes[i], m_bboxes[j]) > nms_threshold)
                m_bboxes[j].flg_remove = true;
        }
    }
}
LOGD("the count of bbox after NMS is %d", final_bboxes.size());

详细步骤如下:

3.1 初始化最终边界框容器

  • final_bboxes 预先分配内存,避免动态扩展带来的性能损耗

3.2 排序

  • m_bboxes 按照置信度从高到低排序,确保优先保留高置信度的边界框

3.3 遍历并应用 NMS

  • 外层循环: 遍历所有边界框
    • 跳过已标记移除的边界框: flg_remove 标记为 true 的边界框被跳过
    • 保留当前边界框: 将当前边界框加入到 final_bboxes
  • 内层循环: 对于当前边界框,遍历后续的所有边界框
    • 同类别比较: 仅比较同一类别的边界框
    • 计算IoU: 如果两个边界框的 IoU 大于 nms_threshold,则将后续边界框标记为 flg_remove,避免重复检测

3.4 日志记录

  • 打印经过 NMS 后的边界框数量

Note:关于 NMS 可能大家看到过很多不同的实现方式,博主这里还是比较喜欢韩君老师的这个实现,简单高效,那其实最早在杜老师的课程中 CPU 版本的 NMS 实现就是这么做的,大家感兴趣的可以看看:6.3.tensorRT高级(1)-yolov5模型导出、编译到推理(无封装)

4. 绘制边界框

string tag   = "detect-" + getPrec(m_params->prec);
m_outputPath = changePath(m_imagePath, "../result", ".png", tag);

int   font_face  = 0;
float font_scale = 0.001 * MIN(m_inputImage.cols, m_inputImage.rows);
int   font_thick = 2;
int   baseline;
CocoLabels labels;

LOG("\tResult:");
for (int i = 0; i < final_bboxes.size(); i ++){
    auto box = final_bboxes[i];
    auto name = labels.coco_get_label(box.label);
    auto rec_color = labels.coco_get_color(box.label);
    auto txt_color = labels.get_inverse_color(rec_color);
    auto txt = cv::format({"%s: %.2f%%"}, name.c_str(), box.confidence * 100);
    auto txt_size = cv::getTextSize(txt, font_face, font_scale, font_thick, &baseline);

    int txt_height = txt_size.height + baseline + 10;
    int txt_width  = txt_size.width + 3;

    cv::Point txt_pos(round(box.x0), round(box.y0 - (txt_size.height - baseline + font_thick)));
    cv::Rect  txt_rec(round(box.x0 - font_thick), round(box.y0 - txt_height), txt_width, txt_height);
    cv::Rect  box_rec(round(box.x0), round(box.y0), round(box.x1 - box.x0), round(box.y1 - box.y0));

    cv::rectangle(m_inputImage, box_rec, rec_color, 3);
    cv::rectangle(m_inputImage, txt_rec, rec_color, -1);
    cv::putText(m_inputImage, txt, txt_pos, font_face, font_scale, txt_color, font_thick, 16);

    LOG("%+20s detected. Confidence: %.2f%%. Cord: (x0, y0):(%6.2f, %6.2f), (x1, y1)(%6.2f, %6.2f)", 
        name.c_str(), box.confidence * 100, box.x0, box.y0, box.x1, box.y1);
}
LOG("\tSummary:");
LOG("\t\tDetected Objects: %d", final_bboxes.size());
LOG("");

详细步骤如下:

4.1 生成输出路径

  • 根据检测精度参数生成一个标签 tag,并通过 changePath 函数构造最终保存的图像路径 m_outputPath

4.2 设置绘制参数

  • font_face: 字体类型,使用默认字体(OpenCV 中 0 通常表示 FONT_HERSHEY_SIMPLEX
  • font_scale: 字体缩放比例,基于图像的最小尺寸动态调整,确保在不同分辨率下字体大小适中
  • font_thick: 字体厚度
  • baseline: 用于计算文本大小

4.3 初始化类别标签和颜色

  • CocoLabels labels: 一个用于获取类别名称和颜色的辅助类

4.4 遍历最终边界框并绘制

  • 获取类别信息

    • name:获取当前边界框的类别名称
    • rec_color:获取用于绘制边界框的颜色
    • txt_color:获取文本颜色,通常是边界框颜色的反色,以确保文本可读性
  • 格式化文本

    • 使用 cv::format 生成显示的文本,例如 "person: 99.99%"
    • 使用 cv::getTextSize 计算文本的尺寸,确保绘制文本背景框时大小合适
  • 计算文本位置和背景框

    • txt_pos:文本的位置,位于边界框的左上角上方
    • txt_rec:文本背景框的位置和大小,确保文本不覆盖检测区域
    • box_rec:边界框的位置和大小
  • 绘制边界框和文本

    • 使用 cv::rectangle 绘制边界框和文本背景框
    • 使用 cv::putText 绘制文本内容
  • 日志记录

    • 打印每个检测到的对象的详细信息,包括类别、置信度和坐标

4.5 总结信息

  • 打印检测到的总对象数量,提供整体检测结果的概览

5. 保存图像

m_timer->stop_cpu<timer::Timer::ms>("postprocess(CPU)");

cv::imwrite(m_outputPath, m_inputImage);
LOG("\tsave image to %s", m_outputPath.c_str());

m_timer->show();
printf("\n");

return true;
  • 停止计时器:记录后处理的总耗时
  • 保存图像:使用 cv::imwrite 将带有绘制边界框的图像保存到 m_outputPath
  • 显示计时信息:打印计时器统计信息,帮助分析性能
  • 返回结果:函数成功完成,返回 true

该函数总体流程总结如下:

  • 1. 数据迁移:将模型输出从 GPU 传输到 CPU
  • 2. 解码:将模型输出的 tensor 数据解码成具体的边界框坐标,并应用逆仿射变换以匹配原图像坐标
  • 3. NMS:通过非极大值抑制算法过滤重叠度高的边界框,保留最具代表性的检测结果
  • 4. 绘制:使用 OpenCV 在原图上绘制最终的检测边界框和标签,并保存结果图像
  • 5. 性能监控:通过计时器记录后处理的耗时,便于性能优化

值得注意的是当前解码和 NMS 都是串行执行的,若处理的边界框数量较大,可以考虑使用多线程或并行算法加速(CUDA),大家感兴趣的可以看看:app_yolo/yolo_decode.cu#L64

4. INT8量化前瞻

最后我们这里再看下 YOLOv8 的 INT8 量化的结果,那这里博主校准数据集选择的是 COCO2017 中训练集的 1000 张图片,大家可以自行选择,也可以点击 here 下载博主准备好的校准图片

接着我们需要像 6.3 小节案例一样,生成一个 .txt 文件,里面保存着校准图片的路径,这里校准流程博主就不再赘述了,大家感兴趣的可以看看:六. 部署分类器-int8-calibration

校准过程如下图所示:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

接着我们再看下 INT8 模型推理的图片结果:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

大家从上面几张推理图片可以看到 INT8 模型推理出来的图片的置信度相比 FP32 和 FP16 低很多,甚至很多都没有检测出目标,精度掉点比较严重,那为什么 yolov8 的 INT8 模型掉点会这么严重呢?另外虽然 INT8 模型的置信度比较低但是它每个目标框的位置还是比较准确的,回归得比较好,这个又是为什么呢?还有就是 INT8 模型推理的速度相比 FP16 也没有一个线性的增长,这是为什么呢?那大家可以自己思考下,这些问题我们会在 7.4-quantization-analysis 小节中详细讲解

那对于第一个问题博主思考过 INT8 校准算法是 tensorRT 官方提供的,没有啥问题,接着影响精度的就只有两个方面,一个是校准数据集,第二个是 INT8 的预处理函数,校准数据集肯定没问题,直接从训练集中随机选择的 1000 张,那最终只有一个可能那就是校准数据经过预处理的实现方法存在问题

我们在 trt_model.cpp 中可以发现如下的代码:

shared_ptr<Int8EntropyCalibrator> calibrator(new Int8EntropyCalibrator(
    64, 
    "calibration/calibration_list_coco.txt", 
    "calibration/calibration_table.txt",
    3 * 224 * 224, 224, 224));

可以看到韩君老师这边传递进去的维度是 classfication 任务的 224x224,在 detection 任务中我们需要修改为 640x640,修改后的代码如下所示:

shared_ptr<Int8EntropyCalibrator> calibrator(new Int8EntropyCalibrator(
    64, 
    "calibration/calibration_list_coco.txt", 
    "calibration/calibration_table.txt",
    3 * 640 * 640, 640, 640));

再次执行 INT8 量化,如下图所示:

在这里插入图片描述

可以看到校准图片缩放后的 size 修改成了 640x640,符合我们的预期

接着我们来看下推理结果的置信度是不是提高了,如下图所示:

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

可以看到效果好了不少,另外这也解释了第二个问题,为什么回归效果好而分类效果不好,因为原本需要缩放的 size 是 640x640,修改成了 224x224,缩放的宽高比是一致的,只是倍数不一致,导致缩放到 224x224 经过 32 倍下采样后特征图的分辨率非常小,很多目标无法检测出来,所以它更多影响的是分类效果,而不是回归效果

关于第三个问题为什么 INT8 的推理速度相比于 FP16 不是线性增长的,我们通过 trt-engine-explore 工具分析后也许可以找到答案,关于 trt-engine-explore 工具的使用我们前面有提到过,这边博主不再赘述,大家感兴趣的可以看看:六. 部署分类器-trt-engine-explorer

在这里插入图片描述

在这里插入图片描述

在这里插入图片描述

上面是博主通过 trt-engine-explore 工具可视化 yolov8 INT8 engine 的结构,从中我们可以发现几个问题,第一个是 engine 中存在着一些 Reformat 节点,这些可能是影响推理速度的一个原因,第二个是我们能够发现这个 engine 并不是每个 layer 跑的都是 INT8 精度,特别是在 concat 节点部分以及最后的 head 部分全部退回到 FP32 精度,这个可能与 tensorRT 内部 PTQ 量化的机制有关,我们无法操控每一层的精度,这个可能是导致推理速度没有成线性变化的一个主因

那当然这些都是博主自己的猜测,下节课程我们会一起来跟随韩君老师学习当量化出现问题时该如何分析并解决

总结

本次课程我们学习了 yolov8 模型的部署,由于我们在 6.2 小节中已经将基础的推理框架搭建好了,因此新添加的 Detection task 只需要新增前后处理的实现即可,其中前处理我们可以用 7.1 小节的 warpAffine 来实现,后处理主要包括 decode 和 nms 等步骤。总的来说,如果大家把 yolov8 的前后处理梳理清楚之后,要实现其部署相对来说还是比较简单的

OK,以上就是 7.3 小节案例的全部内容了,下节我们来学习 7.4 小节 yolov8 的 quantization 量化,敬请期待😄

下载链接

参考


网站公告

今日签到

点亮在社区的每一天
去签到