From 9b753352c810903369763dffb71f527668e9990c Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Sun, 17 Jul 2022 11:25:34 -0300 Subject: [PATCH] Fix detections when batch_size > 1 --- nvdsinfer_custom_impl_Yolo/sortDetections.cu | 5 +++-- utils/gen_wts_yoloV5.py | 6 +++--- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/nvdsinfer_custom_impl_Yolo/sortDetections.cu b/nvdsinfer_custom_impl_Yolo/sortDetections.cu index 8525ad8..cda816a 100644 --- a/nvdsinfer_custom_impl_Yolo/sortDetections.cu +++ b/nvdsinfer_custom_impl_Yolo/sortDetections.cu @@ -82,8 +82,9 @@ cudaError_t sortDetections( sortOutput<<>>( _d_indexes, _d_scores, reinterpret_cast(d_boxes) + (batch * 4 * outputSize), - reinterpret_cast(d_classes) + (batch * outputSize), reinterpret_cast(bboxData) + (batch * topK), - reinterpret_cast(scoreData) + (batch * topK), numOutputClasses, _topK); + reinterpret_cast(d_classes) + (batch * outputSize), + reinterpret_cast(bboxData) + (batch * topK * 4), + reinterpret_cast(scoreData) + (batch * topK * numOutputClasses), numOutputClasses, _topK); cudaFree(d_keys_out); cudaFree(d_values_out); diff --git a/utils/gen_wts_yoloV5.py b/utils/gen_wts_yoloV5.py index 9d20859..1cb12d5 100644 --- a/utils/gen_wts_yoloV5.py +++ b/utils/gen_wts_yoloV5.py @@ -173,9 +173,9 @@ class Layers(object): self.fc.write('\n[convolutional]\n' + b + 'filters=%d\n' % filters + - 'size=%s\n' % (size[0] if len(size) == 2 and size[0] == size[1] else str(size)[1:-1]) + - 'stride=%s\n' % (stride[0] if len(stride) == 2 and stride[0] == stride[1] else str(stride)[1:-1]) + - 'pad=%s\n' % (pad[0] if len(pad) == 2 and pad[0] == pad[1] else str(pad)[1:-1]) + + 'size=%s\n' % (size[0] if type(size) != int and size[0] == size[1] else str(size)[1:-1]) + + 'stride=%s\n' % (stride[0] if type(stride) != int and stride[0] == stride[1] else str(stride)[1:-1]) + + 'pad=%s\n' % (pad[0] if type(pad) != int and pad[0] == pad[1] else str(pad)[1:-1]) + g + w + 'activation=%s\n' % act)