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

Fix route layer

Browse Source
This commit is contained in:
Marcos Luciano
2022-08-16 14:36:52 -03:00
parent ab082fc292
commit dc023b308e
6 changed files with 33 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -304,7 +304,17 @@ NvDsInferStatus Yolo::buildYoloNetwork(std::vector<float>& weights, nvinfer1::IN
else if (m_ConfigBlocks.at(i).at("type") == "reorg")
{
std::string inputVol = dimsToString(previous->getDimensions());
previous = reorgLayer(i, m_ConfigBlocks.at(i), previous, &network);
if (m_NetworkType.find("yolov2") != std::string::npos) {
nvinfer1::IPluginV2* reorgPlugin = createReorgPlugin(2);
assert(reorgPlugin != nullptr);
nvinfer1::IPluginV2Layer* reorg = network.addPluginV2(&previous, 1, *reorgPlugin);
assert(reorg != nullptr);
std::string reorglayerName = "reorg_" + std::to_string(i);
reorg->setName(reorglayerName.c_str());
previous = reorg->getOutput(0);
}
else
previous = reorgLayer(i, m_ConfigBlocks.at(i), previous, &network);
assert(previous != nullptr);
std::string outputVol = dimsToString(previous->getDimensions());
tensorOutputs.push_back(previous);