Yolov8 adds cuda post-processing function (#1326)
* Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Update postprocess.cpp * fix user id * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Yolov8 adds cuda post-processing function * Delete images * Update postprocess.h * Update postprocess.cpp --------- Co-authored-by: Wang Xinyu <shaywxy@gmail.com>
This commit is contained in:
lindsayshuo
GitHub
parent
bc4f854dc3
commit
70d0adc9ec
@ -8,6 +8,7 @@ The tensorrt code is derived from [xiaocao-tian/yolov8_tensorrt](https://github.
|
||||
|
||||
<a href="https://github.com/xiaocao-tian"><img src="https://avatars.githubusercontent.com/u/65889782?v=4?s=48" width="40px;" alt=""/></a>
|
||||
<a href="https://github.com/lindsayshuo"><img src="https://avatars.githubusercontent.com/u/45239466?v=4?s=48" width="40px;" alt=""/></a>
|
||||
<a href="https://github.com/xinsuinizhuan"><img src="https://avatars.githubusercontent.com/u/40679769?v=4?s=48" width="40px;" alt=""/></a>
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
@ -1 +0,0 @@
|
||||
../yolov3-spp/samples
|
||||
@ -3,17 +3,17 @@
|
||||
#include <string>
|
||||
#include <assert.h>
|
||||
|
||||
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);
|
||||
|
||||
@ -11,3 +11,4 @@ void batch_nms(std::vector<std::vector<Detection>>& batch_res, float *output, in
|
||||
|
||||
void draw_bbox(std::vector<cv::Mat>& img_batch, std::vector<std::vector<Detection>>& res_batch);
|
||||
|
||||
void batch_process(std::vector<std::vector<Detection>>& res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector<cv::Mat>& img_batch);
|
||||
|
||||
@ -6,10 +6,20 @@
|
||||
#include <map>
|
||||
|
||||
|
||||
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<cv::Mat> &img_batch,float *dst, int dst_width, int dst_height,cudaStream_t stream);
|
||||
void cuda_batch_preprocess(std::vector<cv::Mat> &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);
|
||||
@ -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<std::chrono::milliseconds>(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<std::chrono::milliseconds>(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<std::string> 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<std::chrono::milliseconds>(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<std::vector<Detection>> 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
|
||||
|
||||
@ -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<Detection> &res, float *output, float conf_thresh, float nms_thresh) {
|
||||
int det_size = sizeof(Detection) / sizeof(float);
|
||||
std::map<float, std::vector<Detection>> 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<std::vector<Detection>> &res_batch, float *output, in
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void process_decode_ptr_host(std::vector<Detection> &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<std::vector<Detection>> &res_batch, const float* decode_ptr_host, int batch_size, int bbox_element, const std::vector<cv::Mat>& img_batch) {
|
||||
res_batch.resize(batch_size);
|
||||
int count = static_cast<int>(*decode_ptr_host);
|
||||
count = std::min(count, kMaxNumOutputBbox);
|
||||
for (int i = 0; i < batch_size; i++) {
|
||||
auto& img = const_cast<cv::Mat&>(img_batch[i]);
|
||||
process_decode_ptr_host(res_batch[i], &decode_ptr_host[i * count], bbox_element, img, count);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void draw_bbox(std::vector<cv::Mat> &img_batch, std::vector<std::vector<Detection>> &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<cv::Mat> &img_batch, std::vector<std::vector<Detectio
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -1,17 +1,13 @@
|
||||
#include "preprocess.h"
|
||||
#include "cuda_utils.h"
|
||||
static uint8_t* img_buffer_host = nullptr;
|
||||
static uint8_t* img_buffer_device = nullptr;
|
||||
|
||||
struct AffineMatrix{
|
||||
float value[6];
|
||||
};
|
||||
static uint8_t *img_buffer_host = nullptr;
|
||||
static uint8_t *img_buffer_device = nullptr;
|
||||
|
||||
__global__ void warpaffine_kernel(
|
||||
uint8_t* src, int src_line_size, int src_width,
|
||||
int src_height, float* dst, int dst_width,
|
||||
int dst_height, uint8_t const_value_st,
|
||||
AffineMatrix d2s, int edge) {
|
||||
|
||||
__global__ void
|
||||
warpaffine_kernel(uint8_t *src, int src_line_size, int src_width, int src_height, float *dst, int dst_width,
|
||||
int dst_height, uint8_t const_value_st, AffineMatrix d2s, int edge) {
|
||||
int position = blockDim.x * blockIdx.x + threadIdx.x;
|
||||
if (position >= 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<<<blocks, threads, 0, stream>>>(
|
||||
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<cv::Mat>& img_batch,
|
||||
float* dst, int dst_width, int dst_height,
|
||||
void cuda_batch_preprocess(std::vector<cv::Mat> &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() {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user