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

Add YOLOv6 support

Browse Source
This commit is contained in:
Marcos Luciano
2023-02-01 02:52:01 -03:00
parent 69f29f8934
commit 087a41acf6
19 changed files with 982 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
nvdsinfer_custom_impl_Yolo/layers/shuffle_layer.cpp
View File

@@ -6,7 +6,7 @@
#include "shuffle_layer.h"
nvinfer1::ITensor*
shuffleLayer(int layerIdx, std::string& layer, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
shuffleLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
std::vector<nvinfer1::ITensor*> tensorOutputs, nvinfer1::INetworkDefinition* network)
{
nvinfer1::ITensor* output;
@@ -18,16 +18,8 @@ shuffleLayer(int layerIdx, std::string& layer, std::map<std::string, std::string
std::string shuffleLayerName = "shuffle_" + std::to_string(layerIdx);
shuffle->setName(shuffleLayerName.c_str());
int from = -1;
if (block.find("from") != block.end())
from = std::stoi(block.at("from"));
if (from < 0)
from = tensorOutputs.size() + from;
layer = std::to_string(from);
if (block.find("reshape") != block.end()) {
nvinfer1::Dims inputTensorDims = tensorOutputs[from]->getDimensions();
nvinfer1::Dims inputTensorDims = input->getDimensions();
std::string strReshape = block.at("reshape");
std::vector<int32_t> reshape;