Add YOLOR-P6 support

nvdsinfer_custom_impl_Yolo/Makefile

@@ -61,6 +61,7 @@ SRCFILES:= nvdsinfer_yolo_engine.cpp \
            layers/upsample_layer.cpp \
            layers/maxpool_layer.cpp \
            layers/activation_layer.cpp \
+           layers/reorg_r_layer.cpp \
            utils.cpp \
            yolo.cpp \
            yoloForward.cu \

nvdsinfer_custom_impl_Yolo/layers/implicit_layer.cpp

@@ -3,7 +3,6 @@
  * https://www.github.com/marcoslucianops
  */
 
-#include <math.h>
 #include "implicit_layer.h"
 
 nvinfer1::ILayer* implicitLayer(

nvdsinfer_custom_impl_Yolo/layers/reorg_r_layer.cpp

@@ -0,0 +1,62 @@
+/*
+ * Created by Marcos Luciano
+ * https://www.github.com/marcoslucianops
+ */
+
+#include "reorg_r_layer.h"
+
+nvinfer1::ILayer* reorgRLayer(
+    int layerIdx,
+    nvinfer1::ITensor* input,
+    nvinfer1::INetworkDefinition* network)
+{
+    nvinfer1::Dims prevTensorDims = input->getDimensions();
+
+    nvinfer1::ISliceLayer *slice1 = network->addSlice(
+        *input,
+        nvinfer1::Dims3{0, 0, 0},
+        nvinfer1::Dims3{prevTensorDims.d[0], prevTensorDims.d[1] / 2, prevTensorDims.d[2] / 2},
+        nvinfer1::Dims3{1, 2, 2});
+    assert(slice1 != nullptr);
+    std::string slice1LayerName = "slice1_" + std::to_string(layerIdx);
+    slice1->setName(slice1LayerName.c_str());
+
+    nvinfer1::ISliceLayer *slice2 = network->addSlice(
+        *input,
+        nvinfer1::Dims3{0, 1, 0},
+        nvinfer1::Dims3{prevTensorDims.d[0], prevTensorDims.d[1] / 2, prevTensorDims.d[2] / 2},
+        nvinfer1::Dims3{1, 2, 2});
+    assert(slice2 != nullptr);
+    std::string slice2LayerName = "slice2_" + std::to_string(layerIdx);
+    slice2->setName(slice2LayerName.c_str());
+
+    nvinfer1::ISliceLayer *slice3 = network->addSlice(
+        *input,
+        nvinfer1::Dims3{0, 0, 1},
+        nvinfer1::Dims3{prevTensorDims.d[0], prevTensorDims.d[1] / 2, prevTensorDims.d[2] / 2},
+        nvinfer1::Dims3{1, 2, 2});
+    assert(slice3 != nullptr);
+    std::string slice3LayerName = "slice3_" + std::to_string(layerIdx);
+    slice3->setName(slice3LayerName.c_str());
+
+    nvinfer1::ISliceLayer *slice4 = network->addSlice(
+        *input,
+        nvinfer1::Dims3{0, 1, 1},
+        nvinfer1::Dims3{prevTensorDims.d[0], prevTensorDims.d[1] / 2, prevTensorDims.d[2] / 2},
+        nvinfer1::Dims3{1, 2, 2});
+    assert(slice4 != nullptr);
+    std::string slice4LayerName = "slice4_" + std::to_string(layerIdx);
+    slice4->setName(slice4LayerName.c_str());
+
+    std::vector<nvinfer1::ITensor*> concatInputs;
+    concatInputs.push_back (slice1->getOutput(0));
+    concatInputs.push_back (slice2->getOutput(0));
+    concatInputs.push_back (slice3->getOutput(0));
+    concatInputs.push_back (slice4->getOutput(0));
+
+    nvinfer1::IConcatenationLayer* concat =
+        network->addConcatenation(concatInputs.data(), concatInputs.size());
+    assert(concat != nullptr);
+
+    return concat;
+}

nvdsinfer_custom_impl_Yolo/layers/reorg_r_layer.h

@@ -0,0 +1,20 @@
+/*
+ * Created by Marcos Luciano
+ * https://www.github.com/marcoslucianops
+ */
+
+#ifndef __REORG_R_LAYER_H__
+#define __REORG_R_LAYER_H__
+
+#include <map>
+#include <vector>
+#include <cassert>
+
+#include "NvInfer.h"
+
+nvinfer1::ILayer* reorgRLayer(
+    int layerIdx,
+    nvinfer1::ITensor* input,
+    nvinfer1::INetworkDefinition* network);
+
+#endif

nvdsinfer_custom_impl_Yolo/yolo.cpp

@@ -299,6 +299,36 @@ NvDsInferStatus Yolo::buildYoloNetwork(
             printLayerInfo(layerIndex, "maxpool", inputVol, outputVol, std::to_string(weightPtr));
         }
 
+        else if (m_ConfigBlocks.at(i).at("type") == "reorg") {
+            if (m_NetworkType.find("yolor") != std::string::npos) {
+                std::string inputVol = dimsToString(previous->getDimensions());
+                nvinfer1::ILayer* out = reorgRLayer(i, previous, &network);
+                previous = out->getOutput(0);
+                assert(previous != nullptr);
+                channels = getNumChannels(previous);
+                std::string outputVol = dimsToString(previous->getDimensions());
+                tensorOutputs.push_back(previous);
+                std::string layerType = "reorgR";
+                printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr));
+            }
+            else {
+                std::string inputVol = dimsToString(previous->getDimensions());
+                nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2);
+                assert(reorgPlugin != nullptr);
+                nvinfer1::IPluginV2Layer* reorg =
+                    network.addPluginV2(&previous, 1, *reorgPlugin);
+                assert(reorg != nullptr);
+                std::string layerName = "reorg_" + std::to_string(i);
+                reorg->setName(layerName.c_str());
+                previous = reorg->getOutput(0);
+                assert(previous != nullptr);
+                std::string outputVol = dimsToString(previous->getDimensions());
+                channels = getNumChannels(previous);
+                tensorOutputs.push_back(reorg->getOutput(0));
+                printLayerInfo(layerIndex, "reorg", inputVol, outputVol, std::to_string(weightPtr));
+            }
+        }
+
         else if (m_ConfigBlocks.at(i).at("type") == "yolo") {
             uint model_type;
             if (m_NetworkType.find("yolor") != std::string::npos) {
@@ -391,22 +421,6 @@ NvDsInferStatus Yolo::buildYoloNetwork(
             printLayerInfo(layerIndex, "region", inputVol, outputVol, std::to_string(weightPtr));
             ++outputTensorCount;
         }
-        else if (m_ConfigBlocks.at(i).at("type") == "reorg") {
-            std::string inputVol = dimsToString(previous->getDimensions());
-            nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2);
-            assert(reorgPlugin != nullptr);
-            nvinfer1::IPluginV2Layer* reorg =
-                network.addPluginV2(&previous, 1, *reorgPlugin);
-            assert(reorg != nullptr);
-            std::string layerName = "reorg_" + std::to_string(i);
-            reorg->setName(layerName.c_str());
-            previous = reorg->getOutput(0);
-            assert(previous != nullptr);
-            std::string outputVol = dimsToString(previous->getDimensions());
-            channels = getNumChannels(previous);
-            tensorOutputs.push_back(reorg->getOutput(0));
-            printLayerInfo(layerIndex, "reorg", inputVol, outputVol, std::to_string(weightPtr));
-        }
 
         else
         {

nvdsinfer_custom_impl_Yolo/yolo.h

@@ -34,6 +34,7 @@
 #include "layers/route_layer.h"
 #include "layers/upsample_layer.h"
 #include "layers/maxpool_layer.h"
+#include "layers/reorg_r_layer.h"
 
 #include "nvdsinfer_custom_impl.h"
 

readme.md

@@ -20,9 +20,9 @@ NVIDIA DeepStream SDK 6.0 configuration for YOLO models
 * Support for convolutional groups
 * Support for INT8 calibration
 * Support for non square models
-* Support for implicit and channel layers (YOLOR)
+* Support for reorg, implicit and channel layers (YOLOR)
 * YOLOv5 6.0 native support
-* Initial YOLOR native support
+* YOLOR native support
 * **Models benchmarks**
 
 ##
@@ -99,6 +99,7 @@ NOTE: Used maintain-aspect-ratio=1 in config_infer file for YOLOv4 (with letter_
 
 | DeepStream         | Precision | Resolution | IoU=0.5:0.95 | IoU=0.5 | IoU=0.75 | FPS<br />(without display) |
 |:------------------:|:---------:|:----------:|:------------:|:-------:|:--------:|:--------------------------:|
+| YOLOR-P6           | FP32      | 1280       | 0.478        | 0.663   | 0.519    | 5.53                       |
 | YOLOR-CSP-X*       | FP32      | 640        | 0.473        | 0.664   | 0.513    | 7.59                       |
 | YOLOR-CSP-X        | FP32      | 640        | 0.470        | 0.661   | 0.507    | 7.52                       |
 | YOLOR-CSP*         | FP32      | 640        | 0.459        | 0.652   | 0.496    | 13.28                      |
@@ -465,8 +466,6 @@ deepstream-app -c deepstream_app_config.txt
 
 ### YOLOR usage
 
-**NOTE**: For now, available only for YOLOR-CSP, YOLOR-CSP*, YOLOR-CSP-X and YOLOR-CSP-X*.
-
 #### 1. Copy gen_wts_yolor.py from DeepStream-Yolo/utils to [yolor](https://github.com/WongKinYiu/yolor) folder
 
 #### 2. Open the yolor folder