Fix detections when batch_size > 1

nvdsinfer_custom_impl_Yolo/sortDetections.cu

@@ -82,8 +82,9 @@ cudaError_t sortDetections(
 
         sortOutput<<<number_of_blocks, threads_per_block, 0, stream>>>(
             _d_indexes, _d_scores, reinterpret_cast<float*>(d_boxes) + (batch * 4 * outputSize),
-            reinterpret_cast<int*>(d_classes) + (batch * outputSize), reinterpret_cast<float*>(bboxData) + (batch * topK),
+            reinterpret_cast<int*>(d_classes) + (batch * outputSize),
-            reinterpret_cast<float*>(scoreData) + (batch * topK), numOutputClasses, _topK);
+            reinterpret_cast<float*>(bboxData) + (batch * topK * 4),
+            reinterpret_cast<float*>(scoreData) + (batch * topK * numOutputClasses), numOutputClasses, _topK);
 
         cudaFree(d_keys_out);
         cudaFree(d_values_out);

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]) +
+                      '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 len(stride) == 2 and stride[0] == stride[1] else str(stride)[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 len(pad) == 2 and pad[0] == pad[1] else str(pad)[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)