목차
0 머리말
Ultralytics는 github에서 빠른 분류, 대상 감지 및 인스턴스 분할을 실현할 수 있는 yolov8 모델을 출시했습니다.공식 yolov8s-seg.pt를 사용하는 효과는 아래 그림과 같습니다.
.
이 기사는 여전히 인스턴스 분할 모델에 대한 가속 추론을 수행하고 너무 많은 파일 종속성이 없는 C++ 버전의 tensorrt 추론 코드를 개발합니다.개인 상품이 없는 3개의 cpp 프로그램 파일만 가장 간단한 추론 버전으로 직접 간주할 수 있습니다.위 링크: Yolov8-instance-seg-tensorrt , 내 환경은 cuda10.2, cudnn8.2.4, Tensorrt8.0.1.6, Opencv4.5.4입니다. yolov8[nsmlx]-seg.pt는 프로그램에서 테스트를 거쳤으며 정상적으로 사용할 수 있습니다.코드 목록은 다음과 같습니다.
├── CMakeLists.txt
├── images
│ ├── bus.jpg
│ └── zidane.jpg
├── logging.h
├── main1_onnx2trt.cpp
├── main2_trt_infer.cpp
├── models
│ ├── yolov8s-seg.engine
│ └── yolov8s-seg.onnx│
├── yolov8n-seg.engine
│ └── yolov8n-seg.onnx
├── output.jpg
├── README.md
└── utils.h
1 onnx 모델 생성
yolov8은 설치 방법과 해당 사용 코드를 제공합니다 웹 사이트에서 해당 모델을 다운로드한 후 다음 코드를 사용하여 필요한 onnx 모델을 생성합니다.
pip install ultralytics
yolo task=segment mode=export model=yolov8[n s m l x]-seg.pt format=onnx opset=12
내 onnx 버전은 최신 버전이 아니므로 공식 모델의 opset=12 및 기본 opset=17은 ultralytics/yolo/configs/default.yaml에서 찾을 수 있습니다.
2 onnx에서 tensorrt로의 엔진 모델
공식코드에서는 엔진을 직접 생성하는 방법을 제공하고 있으나 생성된 엔진은 컴퓨터 환경과 관련이 있기 때문에 직접 사용하는 것은 추천하지 않습니다. 이전 컴퓨터는 불가능합니다. 두 컴퓨터의 환경이 정확히 동일하지 않는 한 필요한 Python 라이브러리를 설치해야 하므로 onnx 모델만 생성한 다음 tensorrt의 api를 통해 모델을 변환합니다.
먼저 내 repo를 복제한 후 다음 문장을 사용합니다.
1. 먼저 Yolov8-instance-seg-tensorrt 디렉토리인 클론의 repo 디렉토리를 찾습니다.
2. yolov8[nslmx]-seg.onnx를 models/ 디렉토리에 복사합니다.
3. 다음 코드를 실행하여 변환 및 추론을 위한 실행 파일 생성 -->onnx2trt, trt_infer
mkdir build
cd build
cmake ..
make
sudo ./onnx2trt ../models/yolov8s-seg.onnx ../models/yolov8s-seg.engine
onnx2trt를 생성하는 코드는 다음과 같습니다. 우리가 사용하는 주요 API가 onnxparser임을 알 수 있습니다.
#include <iostream>
#include "logging.h"
#include "NvOnnxParser.h"
#include "NvInfer.h"
#include <fstream>
using namespace nvinfer1;
using namespace nvonnxparser;
static Logger gLogger;
int main(int argc,char** argv) {
if (argc < 2) {
argv[1] = "../../models/yolov8n-seg.onnx";
argv[2] = "../../models/yolov8n-seg.engine";
}
// 1 onnx解析器
IBuilder* builder = createInferBuilder(gLogger);
const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, gLogger);
const char* onnx_filename = argv[1];
parser->parseFromFile(onnx_filename, static_cast<int>(Logger::Severity::kWARNING));
for (int i = 0; i < parser->getNbErrors(); ++i)
{
std::cout << parser->getError(i)->desc() << std::endl;
}
std::cout << "successfully load the onnx model" << std::endl;
// 2build the engine
unsigned int maxBatchSize = 1;
builder->setMaxBatchSize(maxBatchSize);
IBuilderConfig* config = builder->createBuilderConfig();
config->setMaxWorkspaceSize(1 << 20);
//config->setMaxWorkspaceSize(128 * (1 << 20)); // 16MB
config->setFlag(BuilderFlag::kFP16);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
// 3serialize Model
IHostMemory *gieModelStream = engine->serialize();
std::ofstream p(argv[2], std::ios::binary);
if (!p)
{
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(gieModelStream->data()), gieModelStream->size());
gieModelStream->destroy();
std::cout << "successfully generate the trt engine model" << std::endl;
return 0;
}
위 작업을 통해 해당 yolov8[nsmlx]-seg.engine 모델을 얻을 수 있습니다.
3 TensorRT 추론
이전의 yolov5 인스턴스 분할과 추론 코드가 다르며, 가장 큰 차이점은 아래 그림에서 왼쪽이 v5-seg, 오른쪽이 v8-seg입니다.
v8에서 output0은 8400개의 결과를 출력하고 각 결과의 차원은 116인 반면 v5는 25200개의 결과를 가지며 각 결과의 차원은 117입니다. v8 출력이 무엇인지 간단히 설명하면 116은 4+80+32, 4는 상자의 cx cy wh, 80은 각 클래스의 신뢰 수준, 32는 세분화에 사용해야 하는 것, v5와의 차이점은 목표 신뢰도가 낮고, v5는 4+1+80+32이며, 이 1은 목표 여부에 대한 신뢰도입니다.
추론 코드는 다음과 같습니다.
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "NvInferPlugin.h"
#include "logging.h"
#include <opencv2/opencv.hpp>
#include "utils.h"
#include <string>
using namespace nvinfer1;
using namespace cv;
// stuff we know about the network and the input/output blobs
static const int INPUT_H = 640;
static const int INPUT_W = 640;
static const int _segWidth = 160;
static const int _segHeight = 160;
static const int _segChannels = 32;
static const int CLASSES = 80;
static const int Num_box = 8400;
static const int OUTPUT_SIZE = Num_box * (CLASSES+4 + _segChannels);//output0
static const int OUTPUT_SIZE1 = _segChannels * _segWidth * _segHeight ;//output1
static const float CONF_THRESHOLD = 0.1;
static const float NMS_THRESHOLD = 0.5;
static const float MASK_THRESHOLD = 0.5;
const char* INPUT_BLOB_NAME = "images";
const char* OUTPUT_BLOB_NAME = "output0";//detect
const char* OUTPUT_BLOB_NAME1 = "output1";//mask
struct OutputSeg {
int id; //结果类别id
float confidence; //结果置信度
cv::Rect box; //矩形框
cv::Mat boxMask; //矩形框内mask,节省内存空间和加快速度
};
void DrawPred(Mat& img,std:: vector<OutputSeg> result) {
//生成随机颜色
std::vector<Scalar> color;
srand(time(0));
for (int i = 0; i < CLASSES; i++) {
int b = rand() % 256;
int g = rand() % 256;
int r = rand() % 256;
color.push_back(Scalar(b, g, r));
}
Mat mask = img.clone();
for (int i = 0; i < result.size(); i++) {
int left, top;
left = result[i].box.x;
top = result[i].box.y;
int color_num = i;
rectangle(img, result[i].box, color[result[i].id], 2, 8);
mask(result[i].box).setTo(color[result[i].id], result[i].boxMask);
char label[100];
sprintf(label, "%d:%.2f", result[i].id, result[i].confidence);
//std::string label = std::to_string(result[i].id) + ":" + std::to_string(result[i].confidence);
int baseLine;
Size labelSize = getTextSize(label, FONT_HERSHEY_SIMPLEX, 0.5, 1, &baseLine);
top = max(top, labelSize.height);
putText(img, label, Point(left, top), FONT_HERSHEY_SIMPLEX, 1, color[result[i].id], 2);
}
addWeighted(img, 0.5, mask, 0.8, 1, img); //将mask加在原图上面
}
static Logger gLogger;
void doInference(IExecutionContext& context, float* input, float* output, float* output1, int batchSize)
{
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 3);
void* buffers[3];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
const int outputIndex1 = engine.getBindingIndex(OUTPUT_BLOB_NAME1);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));//
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex1], batchSize * OUTPUT_SIZE1 * sizeof(float)));
// cudaMalloc分配内存 cudaFree释放内存 cudaMemcpy或 cudaMemcpyAsync 在主机和设备之间传输数据
// cudaMemcpy cudaMemcpyAsync 显式地阻塞传输 显式地非阻塞传输
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr);
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
CHECK(cudaMemcpyAsync(output1, buffers[outputIndex1], batchSize * OUTPUT_SIZE1 * sizeof(float), cudaMemcpyDeviceToHost, stream));
cudaStreamSynchronize(stream);
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
CHECK(cudaFree(buffers[outputIndex1]));
}
int main(int argc, char** argv)
{
if (argc < 2) {
argv[1] = "../models/yolov8n-seg.engine";
argv[2] = "../images/bus.jpg";
}
// create a model using the API directly and serialize it to a stream
char* trtModelStream{ nullptr }; //char* trtModelStream==nullptr; 开辟空指针后 要和new配合使用,比如89行 trtModelStream = new char[size]
size_t size{ 0 };//与int固定四个字节不同有所不同,size_t的取值range是目标平台下最大可能的数组尺寸,一些平台下size_t的范围小于int的正数范围,又或者大于unsigned int. 使用Int既有可能浪费,又有可能范围不够大。
std::ifstream file(argv[1], std::ios::binary);
if (file.good()) {
std::cout << "load engine success" << std::endl;
file.seekg(0, file.end);//指向文件的最后地址
size = file.tellg();//把文件长度告诉给size
//std::cout << "\nfile:" << argv[1] << " size is";
//std::cout << size << "";
file.seekg(0, file.beg);//指回文件的开始地址
trtModelStream = new char[size];//开辟一个char 长度是文件的长度
assert(trtModelStream);//
file.read(trtModelStream, size);//将文件内容传给trtModelStream
file.close();//关闭
}
else {
std::cout << "load engine failed" << std::endl;
return 1;
}
Mat src = imread(argv[2], 1);
if (src.empty()) { std::cout << "image load faild" << std::endl; return 1; }
int img_width = src.cols;
int img_height = src.rows;
std::cout << "宽高:" << img_width << " " << img_height << std::endl;
// Subtract mean from image
static float data[3 * INPUT_H * INPUT_W];
Mat pr_img0, pr_img;
std::vector<int> padsize;
pr_img = preprocess_img(src, INPUT_H, INPUT_W, padsize); // Resize
int newh = padsize[0], neww = padsize[1], padh = padsize[2], padw = padsize[3];
float ratio_h = (float)src.rows / newh;
float ratio_w = (float)src.cols / neww;
int i = 0;// [1,3,INPUT_H,INPUT_W]
//std::cout << "pr_img.step" << pr_img.step << std::endl;
for (int row = 0; row < INPUT_H; ++row) {
uchar* uc_pixel = pr_img.data + row * pr_img.step;//pr_img.step=widthx3 就是每一行有width个3通道的值
for (int col = 0; col < INPUT_W; ++col)
{
data[i] = (float)uc_pixel[2] / 255.0;
data[i + INPUT_H * INPUT_W] = (float)uc_pixel[1] / 255.0;
data[i + 2 * INPUT_H * INPUT_W] = (float)uc_pixel[0] / 255.;
uc_pixel += 3;
++i;
}
}
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
bool didInitPlugins = initLibNvInferPlugins(nullptr, "");
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// Run inference
static float prob[OUTPUT_SIZE];
static float prob1[OUTPUT_SIZE1];
//for (int i = 0; i < 10; i++) {//计算10次的推理速度
// auto start = std::chrono::system_clock::now();
// doInference(*context, data, prob, prob1, 1);
// auto end = std::chrono::system_clock::now();
// std::cout << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
// }
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, prob1, 1);
auto end = std::chrono::system_clock::now();
std::cout << "推理时间:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
std::vector<int> classIds;//结果id数组
std::vector<float> confidences;//结果每个id对应置信度数组
std::vector<cv::Rect> boxes;//每个id矩形框
std::vector<cv::Mat> picked_proposals; //后续计算mask
// 处理box
int net_length = CLASSES + 4 + _segChannels;
cv::Mat out1 = cv::Mat(net_length, Num_box, CV_32F, prob);
start = std::chrono::system_clock::now();
for (int i = 0; i < Num_box; i++) {
//输出是1*net_length*Num_box;所以每个box的属性是每隔Num_box取一个值,共net_length个值
cv::Mat scores = out1(Rect(i, 4, 1, CLASSES)).clone();
Point classIdPoint;
double max_class_socre;
minMaxLoc(scores, 0, &max_class_socre, 0, &classIdPoint);
max_class_socre = (float)max_class_socre;
if (max_class_socre >= CONF_THRESHOLD) {
cv::Mat temp_proto = out1(Rect(i, 4 + CLASSES, 1, _segChannels)).clone();
picked_proposals.push_back(temp_proto.t());
float x = (out1.at<float>(0, i) - padw) * ratio_w; //cx
float y = (out1.at<float>(1, i) - padh) * ratio_h; //cy
float w = out1.at<float>(2, i) * ratio_w; //w
float h = out1.at<float>(3, i) * ratio_h; //h
int left = MAX((x - 0.5 * w), 0);
int top = MAX((y - 0.5 * h), 0);
int width = (int)w;
int height = (int)h;
if (width <= 0 || height <= 0) { continue; }
classIds.push_back(classIdPoint.y);
confidences.push_back(max_class_socre);
boxes.push_back(Rect(left, top, width, height));
}
}
//执行非最大抑制以消除具有较低置信度的冗余重叠框(NMS)
std::vector<int> nms_result;
cv::dnn::NMSBoxes(boxes, confidences, CONF_THRESHOLD, NMS_THRESHOLD, nms_result);
std::vector<cv::Mat> temp_mask_proposals;
std::vector<OutputSeg> output;
Rect holeImgRect(0, 0, src.cols, src.rows);
for (int i = 0; i < nms_result.size(); ++i) {
int idx = nms_result[i];
OutputSeg result;
result.id = classIds[idx];
result.confidence = confidences[idx];
result.box = boxes[idx]&holeImgRect;
output.push_back(result);
temp_mask_proposals.push_back(picked_proposals[idx]);
}
// 处理mask
Mat maskProposals;
for (int i = 0; i < temp_mask_proposals.size(); ++i)
maskProposals.push_back(temp_mask_proposals[i]);
Mat protos = Mat(_segChannels, _segWidth * _segHeight, CV_32F, prob1);
Mat matmulRes = (maskProposals * protos).t();//n*32 32*25600 A*B是以数学运算中矩阵相乘的方式实现的,要求A的列数等于B的行数时
Mat masks = matmulRes.reshape(output.size(), { _segWidth,_segHeight });//n*160*160
std::vector<Mat> maskChannels;
cv::split(masks, maskChannels);
Rect roi(int((float)padw / INPUT_W * _segWidth), int((float)padh / INPUT_H * _segHeight), int(_segWidth - padw / 2), int(_segHeight - padh / 2));
for (int i = 0; i < output.size(); ++i) {
Mat dest, mask;
cv::exp(-maskChannels[i], dest);//sigmoid
dest = 1.0 / (1.0 + dest);//160*160
dest = dest(roi);
resize(dest, mask, cv::Size(src.cols, src.rows), INTER_NEAREST);
//crop----截取box中的mask作为该box对应的mask
Rect temp_rect = output[i].box;
mask = mask(temp_rect) > MASK_THRESHOLD;
output[i].boxMask = mask;
}
end = std::chrono::system_clock::now();
std::cout << "后处理时间:" << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
DrawPred(src, output);
cv::imshow("output.jpg", src);
char c = cv::waitKey(0);
// Destroy the engine
context->destroy();
engine->destroy();
runtime->destroy();
system("pause");
return 0;
}
각 네트워크의 최종 결과는 아래 그림과 같습니다.
3.1 yolov8n-seg 세분화 결과
3.2 yolov8s-seg 세분화 결과
3.3 yolov8m-seg 세분화 결과
3.4 yolov8l-seg 세분화 결과
3.5 yolov8x-seg 세분화 결과
후처리 시간은 검출 프레임 처리와 분할 결과 처리로 나뉜다. 모두 opencv 자체 함수나 opencv 매트릭스로 처리하기 때문에 시간이 많이 걸리므 로 속도 개선의 여지가 많습니다 . 매우 빠른 wangxinyu의 tensorrtx 프로젝트에서 yolov5-7.0의 후처리 방법을 참조할 수 있습니다.