thebears/deepstream_yolo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add YOLOR-P6 support

Browse Source
This commit is contained in:
Marcos Luciano
2022-02-14 15:38:12 -03:00
parent 9d118801be
commit a82f1b8662
7 changed files with 117 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
nvdsinfer_custom_impl_Yolo/Makefile
View File

@@ -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 \

1
nvdsinfer_custom_impl_Yolo/layers/implicit_layer.cpp
View File

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

62
nvdsinfer_custom_impl_Yolo/layers/reorg_r_layer.cpp Normal file
View File

@@ -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;
}

20
nvdsinfer_custom_impl_Yolo/layers/reorg_r_layer.h Normal file
View File

@@ -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

46
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -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
{

1
nvdsinfer_custom_impl_Yolo/yolo.h
View File

@@ -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"