Yolov8部署——segmentation部署以及批量推理

Yolov8部署——segmentation部署以及批量推理

（1）参考:在windows上部署Yolov8主要参考下面两个仓库，https://github.com/xunzixunzi/tensorrt-cpp-api和https://github.com/xunzixunzi/YOLOv8-TensorRT-CPP，代码说是适合批量处理，但是代码中是以batchsize=1为例，所以需要修改一下。
（2）说明:我需要的batchsize=6,是要将一张大图，切分成小图，之后以batchsize=6进行推理.
（3）整体流程：

1.导出batch为动态的Yolov8模型:

from ultralytics import YOLO
model = YOLO("你的模型路径")
model.fuse()
model.info(verbose=False)  
model.export(format="onnx",opset=12,simplify=True,dynamic=True) 

2.修改engine.cpp中的build函数
首先我只想batch是动态的，其次直接用python导出的onnx用engine.cpp进行build是会报错的。
报错如下：（大致是这个错误）

针对这个问题，有三个地方改动，如下：

修改engine.cpp中的build函数

IOptimizationProfile* optProfile = builder->createOptimizationProfile();
    for (int32_t i = 0; i < numInputs; ++i) {
        // Must specify dimensions for all the inputs the model expects.
        const auto input = network->getInput(i);
        const auto inputName = input->getName();
        const auto inputDims = input->getDimensions();
        int32_t inputC = inputDims.d[1];
        int32_t inputH = inputDims.d[2];
        int32_t inputW = inputDims.d[3];
        //对于dynamic,固定H和W，这只是临时解决的办法，要是有更好的办法，大佬请说
        inputH = 640;
        inputW = 640;
        // Specify the optimization profile`
        optProfile->setDimensions(inputName, OptProfileSelector::kMIN, Dims4(1, inputC, inputH, inputW));
        optProfile->setDimensions(inputName, OptProfileSelector::kOPT, Dims4(engineOptions.optBatchSize, inputC, inputH, inputW));
        optProfile->setDimensions(inputName, OptProfileSelector::kMAX, Dims4(engineOptions.maxBatchSize, inputC, inputH, inputW));
    }

engine.cpp 的 runInference() 函数修改为如下：

 // Ensure all dynamic bindings have been defined.
    // TODO: Should use allInputShapesSpecified()
    if (!m_context->allInputShapesSpecified()) {
        throw std::runtime_error("Error, not all required dimensions specified.");
    }

在main.cpp中修改

 yolov8nFP16Options.optBatchSize = 6;
 yolov8nFP16Options.maxBatchSize = 6;

3.修改bool YoloV8::infer(const cv::cuda::GpuMat& inputImage, std::vector<InferenceObject>& inferenceObjects)和void YoloV8::preprocess(const cv::cuda::GpuMat& gpuImg, std::vector<std::vector<cv::cuda::GpuMat>>& inputs)这个函数，最主要是用preprocess中的inputs里存储batch数据进行推断，我是把这两个函数合起来进行批量推断修改如下：

// 处理输入图片
    std::vector<std::vector<cv::cuda::GpuMat>> inputs;   //二维向量
    const auto& inputDims = m_trtEngine->getInputDims();
    int imgWidth = gpuImage.cols;
    int imgHeight = gpuImage.rows;
    int numCols = imgWidth / blockWidth;
    int numRows = imgHeight / blockHeight;
    const int totalBlocks = numCols * numRows;//总的张数
    int blockCounter = 0;
    //分割并存储每个图像块到 input 向量中
    std::vector<cv::cuda::GpuMat> input;
    for (int y = 0; y < numRows; ++y) {
        for (int x = 0; x < numCols; ++x) {
            cv::Rect roi(x * blockWidth, y * blockHeight, blockWidth, blockHeight);
            cv::cuda::GpuMat block(gpuImage(roi));

            cv::cuda::GpuMat rgbMat;
            cv::cuda::cvtColor(block, rgbMat, cv::COLOR_BGR2RGB);

            if (rgbMat.rows != inputDims[0].d[1] || rgbMat.cols != inputDims[0].d[2]) {
                throw std::runtime_error("Error:图片尺寸不对.");
            }
            else {
                input.emplace_back(rgbMat);
            }
            blockCounter++;

            if (input.size() == 6 || (input.size() != 6 && blockCounter == totalBlocks)) {

                while (input.size() < 6) {
                    // 如果不足六张图像，则用全黑的图像补全
                    cv::cuda::GpuMat blackImage(blockHeight, blockWidth, CV_8UC3, cv::Scalar(0, 0, 0));
                    input.emplace_back(blackImage);
                }
                inputs.emplace_back(std::move(input));
                std::vector<std::vector<std::vector<float>>> featureVector;
                auto succ = m_trtEngine->runInference(inputs, featureVector);
                if (!succ) {
                    throw std::runtime_error("Error: Unable to run inference.");
                }
                input.clear();
                inputs.clear();
                featureVectors.insert(featureVectors.end(), std::make_move_iterator(featureVector.begin()), std::make_move_iterator(featureVector.end()));
            }
            
        }
    }

之后模型后处理拼接推断结果：

// 后处理阶段需要用
    m_imgHeight = static_cast<float>(blockWidth);
    m_imgWidth = static_cast<float>(blockHeight);
    m_ratio = 1.f / std::min(inputDims[0].d[2] / static_cast<float>(blockWidth), inputDims[0].d[1] / static_cast<float>(blockHeight));

    int numColsl = gpuImage.cols / blockWidth;
    int numRowsl = gpuImage.rows / blockHeight;
    cv::Mat fullMask = cv::Mat::zeros(gpuImage.size(), CV_8UC1);
    int cnt = 0;
    for (int i = 0; i < numRowsl; ++i) {
        for (int j = 0; j < numColsl; ++j) {
            std::vector<std::vector<float>> batch;
            batch = featureVectors[cnt];
            postprocessSegmentation(batch, inferenceObjects);

            if (!inferenceObjects.empty()) {
                for (int k = 0; k < inferenceObjects.size(); ++k) {
                    auto& object = inferenceObjects[k];
                    std::vector<int> objectInfo;
                    objectInfo.push_back(object.label); // 假设id是int类型
                    objectInfo.push_back(object.rect.x + j * 640); // x值
                    objectInfo.push_back(object.rect.y + i * 640); // y值
                    objectInfo.push_back(object.rect.width); // 宽度w
                    objectInfo.push_back(object.rect.height); // 高度h
                      // 将objectInfo添加到Result向量中
                    Result.push_back(objectInfo);
                    if (!object.boxMask.empty()) {
                        // 对对象的rect位置进行操作
                        object.rect.x += j * 640;
                        object.rect.y += i * 640;
                        // 在fullMask上根据修改后的对象的rect位置放置掩码
                        cv::Mat roi = fullMask(object.rect);
                        cv::Mat resizedMask;
                        cv::resize(object.boxMask, resizedMask, object.rect.size());
                        resizedMask.copyTo(roi, resizedMask);
                    }
                    fullMask.copyTo(BinMat);
                }
            }
            cnt += 1;
            inferenceObjects.clear();
        }
    }

最后得到Result(里面存储分割并分类的标签、x、y、height和width)和一个分割的掩码二值化图。