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

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

参考:在windows上部署Yolov8主要参考下面两个仓库，https://github.com/xunzixunzi/tensorrt-cpp-api和https://github.com/xunzixunzi/YOLOv8-TensorRT-CPP，代码说是适合批量处理，但是代码中是以batchsize=1为例，所以需要修改一下。
具体修改:我需要的batchsize=6,是要将一张大图，切分成小图，之后以batchsize=6进行推理，主要修改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)这个函数，我是把这两个函数合起来进行批量推断修改如下：

// 处理输入图片
    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)和一个分割的掩码二值化图。