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); // 宽度wobjectInfo.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)和一个分割的掩码二值化图。