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

Add support to YOLOv5 v4.0 and v5.0

Browse Source
This commit is contained in:
Marcos Luciano
2022-04-08 17:13:05 -03:00
parent 2aa52a8e8c
commit 7bcc9b62fa
6 changed files with 102 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
nvdsinfer_custom_impl_Yolo/Makefile
View File

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

4
nvdsinfer_custom_impl_Yolo/layers/reorg_r_layer.cpp → nvdsinfer_custom_impl_Yolo/layers/reorgv5_layer.cpp
View File

@@ -3,9 +3,9 @@
* https://www.github.com/marcoslucianops
*/
#include "reorg_r_layer.h"
#include "reorgv5_layer.h"
nvinfer1::ILayer* reorgRLayer(
nvinfer1::ILayer* reorgV5Layer(
int layerIdx,
nvinfer1::ITensor* input,
nvinfer1::INetworkDefinition* network)

6
nvdsinfer_custom_impl_Yolo/layers/reorg_r_layer.h → nvdsinfer_custom_impl_Yolo/layers/reorgv5_layer.h
View File

@@ -3,8 +3,8 @@
* https://www.github.com/marcoslucianops
*/
#ifndef __REORG_R_LAYER_H__
#define __REORG_R_LAYER_H__
#ifndef __REORGV5_LAYER_H__
#define __REORGV5_LAYER_H__
#include <map>
#include <vector>
@@ -12,7 +12,7 @@
#include "NvInfer.h"
nvinfer1::ILayer* reorgRLayer(
nvinfer1::ILayer* reorgV5Layer(
int layerIdx,
nvinfer1::ITensor* input,
nvinfer1::INetworkDefinition* network);

6
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -301,15 +301,15 @@ NvDsInferStatus Yolo::buildYoloNetwork(
}
else if (m_ConfigBlocks.at(i).at("type") == "reorg") {
if (m_NetworkType.find("yolor") != std::string::npos) {
if (m_NetworkType.find("yolov5") != std::string::npos || m_NetworkType.find("yolor") != std::string::npos) {
std::string inputVol = dimsToString(previous->getDimensions());
nvinfer1::ILayer* out = reorgRLayer(i, previous, &network);
nvinfer1::ILayer* out = reorgV5Layer(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";
std::string layerType = "reorgV5";
printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr));
}
else {

2
nvdsinfer_custom_impl_Yolo/yolo.h
View File

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