From 801e094321193ff444e2dbad32274b3f353287f9 Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Sun, 14 May 2023 15:07:22 -0300 Subject: [PATCH] Add YOLO-NAS and ONNX support --- config_infer_primary_yoloV5_onnx.txt | 2 +- config_infer_primary_yoloV6_onnx.txt | 2 +- config_infer_primary_yoloV7_onnx.txt | 2 +- config_infer_primary_yoloV8_onnx.txt | 2 +- config_infer_primary_yolox_legacy_onnx.txt | 2 +- config_infer_primary_yolox_onnx.txt | 2 +- .../nvdsparsebbox_Yolo_cuda.cu | 30 +++++++++---------- 7 files changed, 21 insertions(+), 21 deletions(-) diff --git a/config_infer_primary_yoloV5_onnx.txt b/config_infer_primary_yoloV5_onnx.txt index 7313f42..a059d17 100644 --- a/config_infer_primary_yoloV5_onnx.txt +++ b/config_infer_primary_yoloV5_onnx.txt @@ -16,7 +16,7 @@ network-type=0 cluster-mode=2 maintain-aspect-ratio=1 symmetric-padding=1 -parse-bbox-func-name=NvDsInferParseYolo_ONNX +parse-bbox-func-name=NvDsInferParse_YOLO_ONNX custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so [class-attrs-all] diff --git a/config_infer_primary_yoloV6_onnx.txt b/config_infer_primary_yoloV6_onnx.txt index 92b151f..7b0dde6 100644 --- a/config_infer_primary_yoloV6_onnx.txt +++ b/config_infer_primary_yoloV6_onnx.txt @@ -16,7 +16,7 @@ network-type=0 cluster-mode=2 maintain-aspect-ratio=1 symmetric-padding=1 -parse-bbox-func-name=NvDsInferParseYolo_ONNX +parse-bbox-func-name=NvDsInferParse_YOLO_ONNX custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so [class-attrs-all] diff --git a/config_infer_primary_yoloV7_onnx.txt b/config_infer_primary_yoloV7_onnx.txt index fab408b..c940736 100644 --- a/config_infer_primary_yoloV7_onnx.txt +++ b/config_infer_primary_yoloV7_onnx.txt @@ -16,7 +16,7 @@ network-type=0 cluster-mode=2 maintain-aspect-ratio=1 symmetric-padding=1 -parse-bbox-func-name=NvDsInferParseYolo_ONNX +parse-bbox-func-name=NvDsInferParse_YOLO_ONNX custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so [class-attrs-all] diff --git a/config_infer_primary_yoloV8_onnx.txt b/config_infer_primary_yoloV8_onnx.txt index b18a3b2..2d85b28 100644 --- a/config_infer_primary_yoloV8_onnx.txt +++ b/config_infer_primary_yoloV8_onnx.txt @@ -16,7 +16,7 @@ network-type=0 cluster-mode=2 maintain-aspect-ratio=1 symmetric-padding=1 -parse-bbox-func-name=NvDsInferParseYoloV8_ONNX +parse-bbox-func-name=NvDsInferParse_YOLOV8_ONNX custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so [class-attrs-all] diff --git a/config_infer_primary_yolox_legacy_onnx.txt b/config_infer_primary_yolox_legacy_onnx.txt index 1d9f410..521a59c 100644 --- a/config_infer_primary_yolox_legacy_onnx.txt +++ b/config_infer_primary_yolox_legacy_onnx.txt @@ -17,7 +17,7 @@ network-type=0 cluster-mode=2 maintain-aspect-ratio=1 symmetric-padding=0 -parse-bbox-func-name=NvDsInferParseYoloX_ONNX +parse-bbox-func-name=NvDsInferParse_YOLOX_ONNX custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so [class-attrs-all] diff --git a/config_infer_primary_yolox_onnx.txt b/config_infer_primary_yolox_onnx.txt index 49fb4aa..a7120e3 100644 --- a/config_infer_primary_yolox_onnx.txt +++ b/config_infer_primary_yolox_onnx.txt @@ -16,7 +16,7 @@ network-type=0 cluster-mode=2 maintain-aspect-ratio=1 symmetric-padding=0 -parse-bbox-func-name=NvDsInferParseYoloX_ONNX +parse-bbox-func-name=NvDsInferParse_YOLOX_ONNX custom-lib-path=nvdsinfer_custom_impl_Yolo/libnvdsinfer_custom_impl_Yolo.so [class-attrs-all] diff --git a/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo_cuda.cu b/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo_cuda.cu index f25c0ec..0b57043 100644 --- a/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo_cuda.cu +++ b/nvdsinfer_custom_impl_Yolo/nvdsparsebbox_Yolo_cuda.cu @@ -31,7 +31,7 @@ #include "utils.h" #include "yoloPlugins.h" -__global__ void decodeTensorYolo_ONNX(NvDsInferParseObjectInfo *binfo, const float* detections, const int numClasses, +__global__ void decodeTensor_YOLO_ONNX(NvDsInferParseObjectInfo *binfo, const float* detections, const int numClasses, const int outputSize, float netW, float netH) { uint x_id = blockIdx.x * blockDim.x + threadIdx.x; @@ -74,7 +74,7 @@ __global__ void decodeTensorYolo_ONNX(NvDsInferParseObjectInfo *binfo, const flo binfo[x_id].classId = maxIndex; } -__global__ void decodeTensorYoloV8_ONNX(NvDsInferParseObjectInfo *binfo, const float* detections, const int numClasses, +__global__ void decodeTensor_YOLOV8_ONNX(NvDsInferParseObjectInfo *binfo, const float* detections, const int numClasses, const int outputSize, float netW, float netH) { uint x_id = blockIdx.x * blockDim.x + threadIdx.x; @@ -115,7 +115,7 @@ __global__ void decodeTensorYoloV8_ONNX(NvDsInferParseObjectInfo *binfo, const f binfo[x_id].classId = maxIndex; } -__global__ void decodeTensorYoloX_ONNX(NvDsInferParseObjectInfo *binfo, const float* detections, const int numClasses, +__global__ void decodeTensor_YOLOX_ONNX(NvDsInferParseObjectInfo *binfo, const float* detections, const int numClasses, const int outputSize, float netW, float netH, const int *grid0, const int *grid1, const int *strides) { uint x_id = blockIdx.x * blockDim.x + threadIdx.x; @@ -233,7 +233,7 @@ __global__ void decodeTensor_PPYOLOE_ONNX(NvDsInferParseObjectInfo *binfo, const } static bool -NvDsInferParseCustomYolo_ONNX(std::vector const& outputLayersInfo, +NvDsInferParseCustom_YOLO_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, NvDsInferParseDetectionParams const& detectionParams, std::vector& objectList) { @@ -257,7 +257,7 @@ NvDsInferParseCustomYolo_ONNX(std::vector const& outputLayer int threads_per_block = 1024; int number_of_blocks = ((outputSize - 1) / threads_per_block) + 1; - decodeTensorYolo_ONNX<<>>( + decodeTensor_YOLO_ONNX<<>>( thrust::raw_pointer_cast(objects.data()), (const float*) (layer.buffer), numClasses, outputSize, static_cast(networkInfo.width), static_cast(networkInfo.height)); @@ -268,7 +268,7 @@ NvDsInferParseCustomYolo_ONNX(std::vector const& outputLayer } static bool -NvDsInferParseCustomYoloV8_ONNX(std::vector const& outputLayersInfo, +NvDsInferParseCustom_YOLOV8_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, NvDsInferParseDetectionParams const& detectionParams, std::vector& objectList) { @@ -292,7 +292,7 @@ NvDsInferParseCustomYoloV8_ONNX(std::vector const& outputLay int threads_per_block = 1024; int number_of_blocks = ((outputSize - 1) / threads_per_block) + 1; - decodeTensorYoloV8_ONNX<<>>( + decodeTensor_YOLOV8_ONNX<<>>( thrust::raw_pointer_cast(objects.data()), (const float*) (layer.buffer), numClasses, outputSize, static_cast(networkInfo.width), static_cast(networkInfo.height)); @@ -303,7 +303,7 @@ NvDsInferParseCustomYoloV8_ONNX(std::vector const& outputLay } static bool -NvDsInferParseCustomYoloX_ONNX(std::vector const& outputLayersInfo, +NvDsInferParseCustom_YOLOX_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, NvDsInferParseDetectionParams const& detectionParams, std::vector& objectList) { @@ -349,7 +349,7 @@ NvDsInferParseCustomYoloX_ONNX(std::vector const& outputLaye int threads_per_block = 1024; int number_of_blocks = ((outputSize - 1) / threads_per_block) + 1; - decodeTensorYoloX_ONNX<<>>( + decodeTensor_YOLOX_ONNX<<>>( thrust::raw_pointer_cast(objects.data()), (const float*) (layer.buffer), numClasses, outputSize, static_cast(networkInfo.width), static_cast(networkInfo.height), thrust::raw_pointer_cast(d_grid0.data()), thrust::raw_pointer_cast(d_grid1.data()), @@ -434,24 +434,24 @@ NvDsInferParseCustom_PPYOLOE_ONNX(std::vector const& outputL } extern "C" bool -NvDsInferParseYolo_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, +NvDsInferParse_YOLO_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, NvDsInferParseDetectionParams const& detectionParams, std::vector& objectList) { - return NvDsInferParseCustomYolo_ONNX(outputLayersInfo, networkInfo, detectionParams, objectList); + return NvDsInferParseCustom_YOLO_ONNX(outputLayersInfo, networkInfo, detectionParams, objectList); } extern "C" bool -NvDsInferParseYoloV8_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, +NvDsInferParse_YOLOV8_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, NvDsInferParseDetectionParams const& detectionParams, std::vector& objectList) { - return NvDsInferParseCustomYoloV8_ONNX(outputLayersInfo, networkInfo, detectionParams, objectList); + return NvDsInferParseCustom_YOLOV8_ONNX(outputLayersInfo, networkInfo, detectionParams, objectList); } extern "C" bool -NvDsInferParseYoloX_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, +NvDsInferParse_YOLOX_ONNX(std::vector const& outputLayersInfo, NvDsInferNetworkInfo const& networkInfo, NvDsInferParseDetectionParams const& detectionParams, std::vector& objectList) { - return NvDsInferParseCustomYoloX_ONNX(outputLayersInfo, networkInfo, detectionParams, objectList); + return NvDsInferParseCustom_YOLOX_ONNX(outputLayersInfo, networkInfo, detectionParams, objectList); } extern "C" bool