diff --git a/yolov8/README.md b/yolov8/README.md index b945d4d..0d32f8b 100644 --- a/yolov8/README.md +++ b/yolov8/README.md @@ -8,6 +8,7 @@ The tensorrt code is derived from [xiaocao-tian/yolov8_tensorrt](https://github. + ## Requirements @@ -26,7 +27,7 @@ Currently, we support yolov8 - Choose the model n/s/m/l/x from command line arguments. - Check more configs in [include/config.h](./include/config.h) -## How to Run, yolov8-tiny as example +## How to Run, yolov8n as example 1. generate .wts from pytorch with .pt, or download .wts from model zoo @@ -35,7 +36,7 @@ Currently, we support yolov8 cp {tensorrtx}/yolov8/gen_wts.py {ultralytics}/ultralytics cd {ultralytics}/ultralytics python gen_wts.py -// a file 'yolov8.wts' will be generated. +// a file 'yolov8n.wts' will be generated. ``` 2. build tensorrtx/yolov8 and run @@ -49,12 +50,13 @@ cp {ultralytics}/ultralytics/yolov8.wts {tensorrtx}/yolov8/build cmake .. make sudo ./yolov8 -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file -sudo ./yolov8 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. +sudo ./yolov8 -d [.engine] [image folder] [c/g] // deserialize and run inference, the images in [image folder] will be processed. // For example yolov8 sudo ./yolov8 -s yolov8n.wts yolov8.engine n -sudo ./yolov8 -d yolov8n.engine ../images -``` +sudo ./yolov8 -d yolov8n.engine ../images c //cpu postprocess +sudo ./yolov8 -d yolov8n.engine ../images g //gpu postprocess +``` 3. check the images generated, as follows. _zidane.jpg and _bus.jpg 4. optional, load and run the tensorrt model in python diff --git a/yolov8/images b/yolov8/images deleted file mode 100644 index 02cc755..0000000 --- a/yolov8/images +++ /dev/null @@ -1 +0,0 @@ -../yolov3-spp/samples \ No newline at end of file diff --git a/yolov8/include/model.h b/yolov8/include/model.h index 2e2ea1f..5e15390 100644 --- a/yolov8/include/model.h +++ b/yolov8/include/model.h @@ -3,17 +3,17 @@ #include #include -nvinfer1::IHostMemory* buildEngineYolov8n(const int& kBatchSize, nvinfer1::IBuilder* builder, +nvinfer1::IHostMemory* buildEngineYolov8n(const int& batchsize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path); -nvinfer1::IHostMemory* buildEngineYolov8s(const int& kBatchSize, nvinfer1::IBuilder* builder, +nvinfer1::IHostMemory* buildEngineYolov8s(const int& batchsize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path); -nvinfer1::IHostMemory* buildEngineYolov8m(const int& kBatchSize, nvinfer1::IBuilder* builder, +nvinfer1::IHostMemory* buildEngineYolov8m(const int& batchsize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path); -nvinfer1::IHostMemory* buildEngineYolov8l(const int& kBatchSize, nvinfer1::IBuilder* builder, +nvinfer1::IHostMemory* buildEngineYolov8l(const int& batchsize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path); -nvinfer1::IHostMemory* buildEngineYolov8x(const int& kBatchSize, nvinfer1::IBuilder* builder, +nvinfer1::IHostMemory* buildEngineYolov8x(const int& batchsize, nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config, nvinfer1::DataType dt, const std::string& wts_path); diff --git a/yolov8/include/postprocess.h b/yolov8/include/postprocess.h index d19f9ff..66522af 100644 --- a/yolov8/include/postprocess.h +++ b/yolov8/include/postprocess.h @@ -11,3 +11,4 @@ void batch_nms(std::vector>& batch_res, float *output, in void draw_bbox(std::vector& img_batch, std::vector>& res_batch); +void batch_process(std::vector>& res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector& img_batch); diff --git a/yolov8/include/preprocess.h b/yolov8/include/preprocess.h index 12b170e..e1aea64 100644 --- a/yolov8/include/preprocess.h +++ b/yolov8/include/preprocess.h @@ -6,10 +6,20 @@ #include +struct AffineMatrix { + float value[6]; +}; + +const int bbox_element = sizeof(AffineMatrix) / sizeof(float)+1; // left, top, right, bottom, confidence, class, keepflag + void cuda_preprocess_init(int max_image_size); void cuda_preprocess_destroy(); -void cuda_preprocess(uint8_t *src, int src_width, int src_height,float *dst, int dst_width, int dst_height,cudaStream_t stream); +void cuda_preprocess(uint8_t *src, int src_width, int src_height, float *dst, int dst_width, int dst_height, cudaStream_t stream); -void cuda_batch_preprocess(std::vector &img_batch,float *dst, int dst_width, int dst_height,cudaStream_t stream); \ No newline at end of file +void cuda_batch_preprocess(std::vector &img_batch, float *dst, int dst_width, int dst_height, cudaStream_t stream); + +void cuda_decode(float* predict, int num_bboxes, float confidence_threshold,float* parray,int max_objects, cudaStream_t stream); + +void cuda_nms(float* parray, float nms_threshold, int max_objects, cudaStream_t stream); \ No newline at end of file diff --git a/yolov8/main.cpp b/yolov8/main.cpp index c3ef3f0..c134e0d 100644 --- a/yolov8/main.cpp +++ b/yolov8/main.cpp @@ -68,7 +68,7 @@ void deserialize_engine(std::string &engine_name, IRuntime **runtime, ICudaEngin } void prepare_buffer(ICudaEngine *engine, float **input_buffer_device, float **output_buffer_device, - float **output_buffer_host) { + float **output_buffer_host, float **decode_ptr_host, float **decode_ptr_device, std::string cuda_post_process) { assert(engine->getNbBindings() == 2); // 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() @@ -79,26 +79,46 @@ void prepare_buffer(ICudaEngine *engine, float **input_buffer_device, float **ou // Create GPU buffers on device CUDA_CHECK(cudaMalloc((void **) input_buffer_device, kBatchSize * 3 * kInputH * kInputW * sizeof(float))); CUDA_CHECK(cudaMalloc((void **) output_buffer_device, kBatchSize * kOutputSize * sizeof(float))); - - *output_buffer_host = new float[kBatchSize * kOutputSize]; + if (cuda_post_process == "c") { + *output_buffer_host = new float[kBatchSize * kOutputSize]; + } else if (cuda_post_process == "g") { + // Allocate memory for decode_ptr_host and copy to device + *decode_ptr_host = new float[1 + kMaxNumOutputBbox * bbox_element]; + CUDA_CHECK(cudaMalloc((void **)decode_ptr_device, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element))); + } } -void infer(IExecutionContext &context, cudaStream_t &stream, void **buffers, float *output, int batchSize) { +void infer(IExecutionContext &context, cudaStream_t &stream, void **buffers, float *output, int batchSize, float* decode_ptr_host, float* decode_ptr_device, int batchSize_in, int model_bboxes, std::string cuda_post_process) { // infer on the batch asynchronously, and DMA output back to host + auto start = std::chrono::system_clock::now(); context.enqueue(batchSize, buffers, stream, nullptr); - CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost,stream)); + if (cuda_post_process == "c") { + CUDA_CHECK(cudaMemcpyAsync(output, buffers[1], batchSize * kOutputSize * sizeof(float), cudaMemcpyDeviceToHost,stream)); + auto end = std::chrono::system_clock::now(); + std::cout << "inference time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + } else if (cuda_post_process == "g") { + CUDA_CHECK(cudaMemsetAsync(decode_ptr_device, 0, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element), stream)); + cuda_decode((float *)buffers[1], model_bboxes, kConfThresh, decode_ptr_device, kMaxNumOutputBbox, stream); + cuda_nms(decode_ptr_device, kNmsThresh, kMaxNumOutputBbox, stream);//cuda nms + CUDA_CHECK(cudaMemcpyAsync(decode_ptr_host, decode_ptr_device, sizeof(float) * (1 + kMaxNumOutputBbox * bbox_element), cudaMemcpyDeviceToHost, stream)); + auto end = std::chrono::system_clock::now(); + std::cout << "inference and gpu postprocess time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + } + CUDA_CHECK(cudaStreamSynchronize(stream)); } -bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, std::string &img_dir, std::string &sub_type) { + +bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, std::string &img_dir, std::string &sub_type, std::string &cuda_post_process) { if (argc < 4) return false; if (std::string(argv[1]) == "-s" && argc == 5) { wts = std::string(argv[2]); engine = std::string(argv[3]); sub_type = std::string(argv[4]); - } else if (std::string(argv[1]) == "-d" && argc == 4) { + } else if (std::string(argv[1]) == "-d" && argc == 5) { engine = std::string(argv[2]); img_dir = std::string(argv[3]); + cuda_post_process = std::string(argv[4]); } else { return false; } @@ -111,11 +131,13 @@ int main(int argc, char **argv) { std::string engine_name = ""; std::string img_dir; std::string sub_type = ""; + std::string cuda_post_process=""; + int model_bboxes; - if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type)) { + if (!parse_args(argc, argv, wts_name, engine_name, img_dir, sub_type, cuda_post_process)) { std::cerr << "Arguments not right!" << std::endl; std::cerr << "./yolov8 -s [.wts] [.engine] [n/s/m/l/x] // serialize model to plan file" << std::endl; - std::cerr << "./yolov8 -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + std::cerr << "./yolov8 -d [.engine] ../samples [c/g]// deserialize plan file and run inference" << std::endl; return -1; } @@ -132,13 +154,14 @@ int main(int argc, char **argv) { deserialize_engine(engine_name, &runtime, &engine, &context); cudaStream_t stream; CUDA_CHECK(cudaStreamCreate(&stream)); - cuda_preprocess_init(kMaxInputImageSize); - + auto out_dims = engine->getBindingDimensions(1); + model_bboxes = out_dims.d[0]; // Prepare cpu and gpu buffers float *device_buffers[2]; float *output_buffer_host = nullptr; - prepare_buffer(engine, &device_buffers[0], &device_buffers[1], &output_buffer_host); + float *decode_ptr_host=nullptr; + float *decode_ptr_device=nullptr; // Read images from directory std::vector file_names; @@ -147,6 +170,8 @@ int main(int argc, char **argv) { return -1; } + prepare_buffer(engine, &device_buffers[0], &device_buffers[1], &output_buffer_host, &decode_ptr_host, &decode_ptr_device, cuda_post_process); + // batch predict for (size_t i = 0; i < file_names.size(); i += kBatchSize) { // Get a batch of images @@ -157,23 +182,20 @@ int main(int argc, char **argv) { img_batch.push_back(img); img_name_batch.push_back(file_names[j]); } - // Preprocess cuda_batch_preprocess(img_batch, device_buffers[0], kInputW, kInputH, stream); - // Run inference - auto start = std::chrono::system_clock::now(); - infer(*context, stream, (void **) device_buffers, output_buffer_host, kBatchSize); - auto end = std::chrono::system_clock::now(); - std::cout << "inference time: " << std::chrono::duration_cast(end - start).count()<< "ms" << std::endl; - - // NMS + infer(*context, stream, (void **)device_buffers, output_buffer_host, kBatchSize, decode_ptr_host, decode_ptr_device, img_batch.size(), model_bboxes, cuda_post_process); std::vector> res_batch; - batch_nms(res_batch, output_buffer_host, img_batch.size(), kOutputSize, kConfThresh, kNmsThresh); - + if (cuda_post_process == "c") { + // NMS + batch_nms(res_batch, output_buffer_host, img_batch.size(), kOutputSize, kConfThresh, kNmsThresh); + } else if (cuda_post_process == "g") { + //Process gpu decode and nms results + batch_process(res_batch, decode_ptr_host, img_batch.size(), bbox_element, img_batch); + } // Draw bounding boxes draw_bbox(img_batch, res_batch); - // Save images for (size_t j = 0; j < img_batch.size(); j++) { cv::imwrite("_" + img_name_batch[j], img_batch[j]); @@ -184,6 +206,8 @@ int main(int argc, char **argv) { cudaStreamDestroy(stream); CUDA_CHECK(cudaFree(device_buffers[0])); CUDA_CHECK(cudaFree(device_buffers[1])); + CUDA_CHECK(cudaFree(decode_ptr_device)); + delete[] decode_ptr_host; delete[] output_buffer_host; cuda_preprocess_destroy(); // Destroy the engine diff --git a/yolov8/src/postprocess.cpp b/yolov8/src/postprocess.cpp index c3fda62..d09bf79 100644 --- a/yolov8/src/postprocess.cpp +++ b/yolov8/src/postprocess.cpp @@ -28,7 +28,6 @@ cv::Rect get_rect(cv::Mat &img, float bbox[4]) { return cv::Rect(round(l), round(t), round(r - l), round(b - t)); } - static float iou(float lbox[4], float rbox[4]) { float interBox[] = { (std::max)(lbox[0] - lbox[2] / 2.f, rbox[0] - rbox[2] / 2.f), //left @@ -51,6 +50,7 @@ static bool cmp(const Detection &a, const Detection &b) { void nms(std::vector &res, float *output, float conf_thresh, float nms_thresh) { int det_size = sizeof(Detection) / sizeof(float); std::map> m; + for (int i = 0; i < output[0]; i++) { if (output[1 + det_size * i + 4] <= conf_thresh) continue; Detection det; @@ -82,6 +82,35 @@ void batch_nms(std::vector> &res_batch, float *output, in } } + +void process_decode_ptr_host(std::vector &res, const float* decode_ptr_host, int bbox_element, cv::Mat& img, int count) { + Detection det; + for (int i = 0; i < count; i++) { + int basic_pos = 1 + i * bbox_element; + int keep_flag = decode_ptr_host[basic_pos + 6]; + if (keep_flag == 1) { + det.bbox[0] = decode_ptr_host[basic_pos + 0]; + det.bbox[1] = decode_ptr_host[basic_pos + 1]; + det.bbox[2] = decode_ptr_host[basic_pos + 2]; + det.bbox[3] = decode_ptr_host[basic_pos + 3]; + det.conf = decode_ptr_host[basic_pos + 4]; + det.class_id = decode_ptr_host[basic_pos + 5]; + res.push_back(det); + } + } +} + +void batch_process(std::vector> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector& img_batch) { + res_batch.resize(batch_size); + int count = static_cast(*decode_ptr_host); + count = std::min(count, kMaxNumOutputBbox); + for (int i = 0; i < batch_size; i++) { + auto& img = const_cast(img_batch[i]); + process_decode_ptr_host(res_batch[i], &decode_ptr_host[i * count], bbox_element, img, count); + } +} + + void draw_bbox(std::vector &img_batch, std::vector> &res_batch) { for (size_t i = 0; i < img_batch.size(); i++) { auto &res = res_batch[i]; @@ -94,4 +123,3 @@ void draw_bbox(std::vector &img_batch, std::vector= edge) return; @@ -45,10 +41,10 @@ __global__ void warpaffine_kernel( float hy = 1 - ly; float hx = 1 - lx; float w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx; - uint8_t* v1 = const_value; - uint8_t* v2 = const_value; - uint8_t* v3 = const_value; - uint8_t* v4 = const_value; + uint8_t *v1 = const_value; + uint8_t *v2 = const_value; + uint8_t *v3 = const_value; + uint8_t *v4 = const_value; if (y_low >= 0) { if (x_low >= 0) @@ -83,18 +79,94 @@ __global__ void warpaffine_kernel( // rgbrgbrgb to rrrgggbbb int area = dst_width * dst_height; - float* pdst_c0 = dst + dy * dst_width + dx; - float* pdst_c1 = pdst_c0 + area; - float* pdst_c2 = pdst_c1 + area; + float *pdst_c0 = dst + dy * dst_width + dx; + float *pdst_c1 = pdst_c0 + area; + float *pdst_c2 = pdst_c1 + area; *pdst_c0 = c0; *pdst_c1 = c1; *pdst_c2 = c2; } -void cuda_preprocess( - uint8_t* src, int src_width, int src_height, - float* dst, int dst_width, int dst_height, - cudaStream_t stream) { +static __global__ void +decode_kernel(float *predict, int num_bboxes, float confidence_threshold, float *parray, int max_objects) { + + float count = predict[0]; + int position = (blockDim.x * blockIdx.x + threadIdx.x); + if (position >= count) + return; + float *pitem = predict + 1 + position * 6; + int index = atomicAdd(parray, 1); + if (index >= max_objects) + return; + float confidence = pitem[4]; + if (confidence < confidence_threshold) + return; + float left = pitem[0]; + float top = pitem[1]; + float right = pitem[2]; + float bottom = pitem[3]; + float label = pitem[5]; + float *pout_item = parray + 1 + index * bbox_element; + *pout_item++ = left; + *pout_item++ = top; + *pout_item++ = right; + *pout_item++ = bottom; + *pout_item++ = confidence; + *pout_item++ = label; + *pout_item++ = 1; // 1 = keep, 0 = ignore +} + +static __device__ float +box_iou(float aleft, float atop, float aright, float abottom, float bleft, float btop, float bright, float bbottom) { + + float cleft = max(aleft, bleft); + float ctop = max(atop, btop); + float cright = min(aright, bright); + float cbottom = min(abottom, bbottom); + + float c_area = max(cright - cleft, 0.0f) * max(cbottom - ctop, 0.0f); + if (c_area == 0.0f) + return 0.0f; + + float a_area = max(0.0f, aright - aleft) * max(0.0f, abottom - atop); + float b_area = max(0.0f, bright - bleft) * max(0.0f, bbottom - btop); + return c_area / (a_area + b_area - c_area); +} + +static __global__ void nms_kernel(float *bboxes, int max_objects, float threshold) { + + int position = (blockDim.x * blockIdx.x + threadIdx.x); + int count = bboxes[0]; + +// float count = 0.0f; + if (position >= count) + return; + + float *pcurrent = bboxes + 1 + position * bbox_element; + for (int i = 1; i < count; ++i) { + float *pitem = bboxes + 1 + i * bbox_element; + if (i == position || pcurrent[5] != pitem[5]) continue; + + if (pitem[4] >= pcurrent[4]) { + if (pitem[4] == pcurrent[4] && i < position) + continue; + + float iou = box_iou( + pcurrent[0], pcurrent[1], pcurrent[2], pcurrent[3], + pitem[0], pitem[1], pitem[2], pitem[3] + ); + + if (iou > threshold) { + pcurrent[6] = 0; + return; + } + } + } +} + + +void cuda_preprocess(uint8_t *src, int src_width, int src_height, float *dst, int dst_width, int dst_height, + cudaStream_t stream) { int img_size = src_width * src_height * 3; // copy data to pinned memory memcpy(img_buffer_host, src, img_size); @@ -102,7 +174,7 @@ void cuda_preprocess( CUDA_CHECK(cudaMemcpyAsync(img_buffer_device, img_buffer_host, img_size, cudaMemcpyHostToDevice, stream)); AffineMatrix s2d, d2s; - float scale = std::min(dst_height / (float)src_height, dst_width / (float)src_width); + float scale = std::min(dst_height / (float) src_height, dst_width / (float) src_width); s2d.value[0] = scale; s2d.value[1] = 0; @@ -118,7 +190,7 @@ void cuda_preprocess( int jobs = dst_height * dst_width; int threads = 256; - int blocks = ceil(jobs / (float)threads); + int blocks = ceil(jobs / (float) threads); warpaffine_kernel<<>>( img_buffer_device, src_width * 3, src_width, src_height, dst, dst_width, @@ -126,21 +198,40 @@ void cuda_preprocess( } -void cuda_batch_preprocess(std::vector& img_batch, - float* dst, int dst_width, int dst_height, +void cuda_batch_preprocess(std::vector &img_batch, + float *dst, int dst_width, int dst_height, cudaStream_t stream) { int dst_size = dst_width * dst_height * 3; for (size_t i = 0; i < img_batch.size(); i++) { - cuda_preprocess(img_batch[i].ptr(), img_batch[i].cols, img_batch[i].rows, &dst[dst_size * i], dst_width, dst_height, stream); + cuda_preprocess(img_batch[i].ptr(), img_batch[i].cols, img_batch[i].rows, &dst[dst_size * i], dst_width, + dst_height, stream); CUDA_CHECK(cudaStreamSynchronize(stream)); } } + +void cuda_decode(float *predict, int num_bboxes, float confidence_threshold, float *parray, int max_objects, + cudaStream_t stream) { + int block = 256; + int grid = ceil(num_bboxes / (float) block); + decode_kernel << < + grid, block, 0, stream >> > ((float *) predict, num_bboxes, confidence_threshold, parray, max_objects); + +} + +void cuda_nms(float *parray, float nms_threshold, int max_objects, cudaStream_t stream) { + int block = max_objects < 256 ? max_objects : 256; + int grid = ceil(max_objects / (float) block); + nms_kernel << < grid, block, 0, stream >> > (parray, max_objects, nms_threshold); + +} + + void cuda_preprocess_init(int max_image_size) { // prepare input data in pinned memory - CUDA_CHECK(cudaMallocHost((void**)&img_buffer_host, max_image_size * 3)); + CUDA_CHECK(cudaMallocHost((void **) &img_buffer_host, max_image_size * 3)); // prepare input data in device memory - CUDA_CHECK(cudaMalloc((void**)&img_buffer_device, max_image_size * 3)); + CUDA_CHECK(cudaMalloc((void **) &img_buffer_device, max_image_size * 3)); } void cuda_preprocess_destroy() {