From 8983e7e99ea6ff5dea0c71eff3003ed155874fdb Mon Sep 17 00:00:00 2001
From: Marcos Luciano <marcoslucianops@gmail.com>
Date: Thu, 14 Jul 2022 13:58:55 -0300
Subject: [PATCH] New YOLOv5 conversion and support (>= v2.0)

---
 nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp | 10 +++-------
 utils/gen_wts_yoloV5.py                             |  5 +----
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp b/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
index c4e6db9..df81c99 100644
--- a/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
+++ b/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp
@@ -18,18 +18,14 @@ nvinfer1::ILayer* maxpoolLayer(
     int size = std::stoi(block.at("size"));
     int stride = std::stoi(block.at("stride"));
 
-    int pad = 0;
-    if (block.find("pad") != block.end())
-    {
-        pad = std::stoi(block.at("pad"));
-    }
-
     nvinfer1::IPoolingLayer* pool
         = network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}});
     assert(pool);
     std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
     pool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
-    pool->setPaddingNd(nvinfer1::Dims{2, {pad, pad}});
+    if (stride == 1) {
+        pool->setPaddingNd(nvinfer1::Dims{2, {size / 2, size / 2}});
+    }
     pool->setName(maxpoolLayerName.c_str());
 
     return pool;
diff --git a/utils/gen_wts_yoloV5.py b/utils/gen_wts_yoloV5.py
index b0c112b..9d20859 100644
--- a/utils/gen_wts_yoloV5.py
+++ b/utils/gen_wts_yoloV5.py
@@ -210,16 +210,13 @@ class Layers(object):
 
         stride = m.stride
         size = m.kernel_size
-        pad = m.padding
         mode = m.ceil_mode
 
         m = 'maxpool_up' if mode else 'maxpool'
-        p = 'pad=%d\n' % pad if pad > 0 else ''
 
         self.fc.write('\n[%s]\n' % m +
                       'stride=%d\n' % stride +
-                      'size=%d\n' % size +
-                      p)
+                      'size=%d\n' % size)
 
     def upsample(self, child):
         self.blocks[self.current] += 1