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

New features and fixes

Browse Source
This commit is contained in:
Marcos Luciano
2023-06-05 14:48:23 -03:00
parent 3f14b0d95d
commit 66a6754b77
57 changed files with 2137 additions and 1534 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
nvdsinfer_custom_impl_Yolo/layers/implicit_layer.cpp
View File

@@ -28,7 +28,7 @@ implicitLayer(int layerIdx, std::map<std::string, std::string>& block, std::vect
convWt.values = val;
trtWeights.push_back(convWt);
nvinfer1::IConstantLayer* implicit = network->addConstant(nvinfer1::Dims{3, {filters, 1, 1}}, convWt);
nvinfer1::IConstantLayer* implicit = network->addConstant(nvinfer1::Dims{4, {1, filters, 1, 1}}, convWt);
assert(implicit != nullptr);
std::string implicitLayerName = block.at("type") + "_" + std::to_string(layerIdx);
implicit->setName(implicitLayerName.c_str());

118
nvdsinfer_custom_impl_Yolo/layers/reorg_layer.cpp
View File

@@ -14,46 +14,100 @@ reorgLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::IT
{
nvinfer1::ITensor* output;
assert(block.at("type") == "reorg3d");
assert(block.at("type") == "reorg" || block.at("type") == "reorg3d");
int stride = 1;
if(block.find("stride") != block.end()) {
stride = std::stoi(block.at("stride"));
}
nvinfer1::Dims inputDims = input->getDimensions();
nvinfer1::ISliceLayer *slice1 = network->addSlice(*input, nvinfer1::Dims{3, {0, 0, 0}},
nvinfer1::Dims{3, {inputDims.d[0], inputDims.d[1] / 2, inputDims.d[2] / 2}}, nvinfer1::Dims{3, {1, 2, 2}});
assert(slice1 != nullptr);
std::string slice1LayerName = "slice1_" + std::to_string(layerIdx);
slice1->setName(slice1LayerName.c_str());
if (block.at("type") == "reorg3d") {
nvinfer1::ISliceLayer* slice1 = network->addSlice(*input, nvinfer1::Dims{4, {0, 0, 0, 0}},
nvinfer1::Dims{4, {inputDims.d[0], inputDims.d[1], inputDims.d[2] / stride, inputDims.d[3] / stride}},
nvinfer1::Dims{4, {1, 1, stride, stride}});
assert(slice1 != nullptr);
std::string slice1LayerName = "slice1_" + std::to_string(layerIdx);
slice1->setName(slice1LayerName.c_str());
nvinfer1::ISliceLayer *slice2 = network->addSlice(*input, nvinfer1::Dims{3, {0, 1, 0}},
nvinfer1::Dims{3, {inputDims.d[0], inputDims.d[1] / 2, inputDims.d[2] / 2}}, nvinfer1::Dims{3, {1, 2, 2}});
assert(slice2 != nullptr);
std::string slice2LayerName = "slice2_" + std::to_string(layerIdx);
slice2->setName(slice2LayerName.c_str());
nvinfer1::ISliceLayer* slice2 = network->addSlice(*input, nvinfer1::Dims{4, {0, 0, 0, 1}},
nvinfer1::Dims{4, {inputDims.d[0], inputDims.d[1], inputDims.d[2] / stride, inputDims.d[3] / stride}},
nvinfer1::Dims{4, {1, 1, stride, stride}});
assert(slice2 != nullptr);
std::string slice2LayerName = "slice2_" + std::to_string(layerIdx);
slice2->setName(slice2LayerName.c_str());
nvinfer1::ISliceLayer *slice3 = network->addSlice(*input, nvinfer1::Dims{3, {0, 0, 1}},
nvinfer1::Dims{3, {inputDims.d[0], inputDims.d[1] / 2, inputDims.d[2] / 2}}, nvinfer1::Dims{3, {1, 2, 2}});
assert(slice3 != nullptr);
std::string slice3LayerName = "slice3_" + std::to_string(layerIdx);
slice3->setName(slice3LayerName.c_str());
nvinfer1::ISliceLayer* slice3 = network->addSlice(*input, nvinfer1::Dims{4, {0, 0, 1, 0}},
nvinfer1::Dims{4, {inputDims.d[0], inputDims.d[1], inputDims.d[2] / stride, inputDims.d[3] / stride}},
nvinfer1::Dims{4, {1, 1, stride, stride}});
assert(slice3 != nullptr);
std::string slice3LayerName = "slice3_" + std::to_string(layerIdx);
slice3->setName(slice3LayerName.c_str());
nvinfer1::ISliceLayer *slice4 = network->addSlice(*input, nvinfer1::Dims{3, {0, 1, 1}},
nvinfer1::Dims{3, {inputDims.d[0], inputDims.d[1] / 2, inputDims.d[2] / 2}}, nvinfer1::Dims{3, {1, 2, 2}});
assert(slice4 != nullptr);
std::string slice4LayerName = "slice4_" + std::to_string(layerIdx);
slice4->setName(slice4LayerName.c_str());
nvinfer1::ISliceLayer* slice4 = network->addSlice(*input, nvinfer1::Dims{4, {0, 0, 1, 1}},
nvinfer1::Dims{4, {inputDims.d[0], inputDims.d[1], inputDims.d[2] / stride, inputDims.d[3] / stride}},
nvinfer1::Dims{4, {1, 1, stride, stride}});
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));
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);
std::string concatLayerName = "concat_" + std::to_string(layerIdx);
concat->setName(concatLayerName.c_str());
concat->setAxis(0);
output = concat->getOutput(0);
nvinfer1::IConcatenationLayer* concat = network->addConcatenation(concatInputs.data(), concatInputs.size());
assert(concat != nullptr);
std::string concatLayerName = "concat_" + std::to_string(layerIdx);
concat->setName(concatLayerName.c_str());
concat->setAxis(0);
output = concat->getOutput(0);
}
else {
nvinfer1::IShuffleLayer* shuffle1 = network->addShuffle(*input);
assert(shuffle1 != nullptr);
std::string shuffle1LayerName = "shuffle1_" + std::to_string(layerIdx);
shuffle1->setName(shuffle1LayerName.c_str());
nvinfer1::Dims reshapeDims1{6, {inputDims.d[0], inputDims.d[1] / (stride * stride), inputDims.d[2], stride,
inputDims.d[3], stride}};
shuffle1->setReshapeDimensions(reshapeDims1);
nvinfer1::Permutation permutation1{{0, 1, 2, 4, 3, 5}};
shuffle1->setSecondTranspose(permutation1);
output = shuffle1->getOutput(0);
nvinfer1::IShuffleLayer* shuffle2 = network->addShuffle(*output);
assert(shuffle2 != nullptr);
std::string shuffle2LayerName = "shuffle2_" + std::to_string(layerIdx);
shuffle2->setName(shuffle2LayerName.c_str());
nvinfer1::Dims reshapeDims2{4, {inputDims.d[0], inputDims.d[1] / (stride * stride), inputDims.d[2] * inputDims.d[3],
stride * stride}};
shuffle2->setReshapeDimensions(reshapeDims2);
nvinfer1::Permutation permutation2{{0, 1, 3, 2}};
shuffle2->setSecondTranspose(permutation2);
output = shuffle2->getOutput(0);
nvinfer1::IShuffleLayer* shuffle3 = network->addShuffle(*output);
assert(shuffle3 != nullptr);
std::string shuffle3LayerName = "shuffle3_" + std::to_string(layerIdx);
shuffle3->setName(shuffle3LayerName.c_str());
nvinfer1::Dims reshapeDims3{4, {inputDims.d[0], inputDims.d[1] / (stride * stride), stride * stride,
inputDims.d[2] * inputDims.d[3]}};
shuffle3->setReshapeDimensions(reshapeDims3);
nvinfer1::Permutation permutation3{{0, 2, 1, 3}};
shuffle3->setSecondTranspose(permutation3);
output = shuffle3->getOutput(0);
nvinfer1::IShuffleLayer* shuffle4 = network->addShuffle(*output);
assert(shuffle4 != nullptr);
std::string shuffle4LayerName = "shuffle4_" + std::to_string(layerIdx);
shuffle4->setName(shuffle4LayerName.c_str());
nvinfer1::Dims reshapeDims4{4, {inputDims.d[0], inputDims.d[1] * stride * stride, inputDims.d[2] / stride,
inputDims.d[3] / stride}};
shuffle4->setReshapeDimensions(reshapeDims4);
output = shuffle4->getOutput(0);
}
return output;
}

38
nvdsinfer_custom_impl_Yolo/layers/route_layer.cpp
View File

@@ -24,29 +24,36 @@ routeLayer(int layerIdx, std::string& layers, std::map<std::string, std::string>
}
if (lastPos < strLayers.length()) {
std::string lastV = trim(strLayers.substr(lastPos));
if (!lastV.empty())
if (!lastV.empty()) {
idxLayers.push_back(std::stoi(lastV));
}
}
assert (!idxLayers.empty());
assert(!idxLayers.empty());
std::vector<nvinfer1::ITensor*> concatInputs;
for (uint i = 0; i < idxLayers.size(); ++i) {
if (idxLayers[i] < 0)
if (idxLayers[i] < 0) {
idxLayers[i] = tensorOutputs.size() + idxLayers[i];
assert (idxLayers[i] >= 0 && idxLayers[i] < (int)tensorOutputs.size());
}
assert(idxLayers[i] >= 0 && idxLayers[i] < (int)tensorOutputs.size());
concatInputs.push_back(tensorOutputs[idxLayers[i]]);
if (i < idxLayers.size() - 1)
if (i < idxLayers.size() - 1) {
layers += std::to_string(idxLayers[i]) + ", ";
}
}
layers += std::to_string(idxLayers[idxLayers.size() - 1]);
if (concatInputs.size() == 1)
if (concatInputs.size() == 1) {
output = concatInputs[0];
}
else {
int axis = 0;
if (block.find("axis") != block.end())
axis = std::stoi(block.at("axis"));
if (axis < 0)
axis = concatInputs[0]->getDimensions().nbDims + axis;
int axis = 1;
if (block.find("axis") != block.end()) {
axis += std::stoi(block.at("axis"));
std::cout << axis << std::endl;
}
if (axis < 0) {
axis += concatInputs[0]->getDimensions().nbDims;
}
nvinfer1::IConcatenationLayer* concat = network->addConcatenation(concatInputs.data(), concatInputs.size());
assert(concat != nullptr);
@@ -60,10 +67,11 @@ routeLayer(int layerIdx, std::string& layers, std::map<std::string, std::string>
nvinfer1::Dims prevTensorDims = output->getDimensions();
int groups = stoi(block.at("groups"));
int group_id = stoi(block.at("group_id"));
int startSlice = (prevTensorDims.d[0] / groups) * group_id;
int channelSlice = (prevTensorDims.d[0] / groups);
nvinfer1::ISliceLayer* slice = network->addSlice(*output, nvinfer1::Dims{3, {startSlice, 0, 0}},
nvinfer1::Dims{3, {channelSlice, prevTensorDims.d[1], prevTensorDims.d[2]}}, nvinfer1::Dims{3, {1, 1, 1}});
int startSlice = (prevTensorDims.d[1] / groups) * group_id;
int channelSlice = (prevTensorDims.d[1] / groups);
nvinfer1::ISliceLayer* slice = network->addSlice(*output, nvinfer1::Dims{4, {0, startSlice, 0, 0}},
nvinfer1::Dims{4, {prevTensorDims.d[0], channelSlice, prevTensorDims.d[2], prevTensorDims.d[3]}},
nvinfer1::Dims{4, {1, 1, 1, 1}});
assert(slice != nullptr);
std::string sliceLayerName = "slice_" + std::to_string(layerIdx);
slice->setName(sliceLayerName.c_str());

4
nvdsinfer_custom_impl_Yolo/layers/shortcut_layer.cpp
View File

@@ -17,8 +17,8 @@ shortcutLayer(int layerIdx, std::string activation, std::string inputVol, std::s
assert(block.at("type") == "shortcut");
if (inputVol != shortcutVol) {
nvinfer1::ISliceLayer* slice = network->addSlice(*shortcutInput, nvinfer1::Dims{3, {0, 0, 0}}, input->getDimensions(),
nvinfer1::Dims{3, {1, 1, 1}});
nvinfer1::ISliceLayer* slice = network->addSlice(*shortcutInput, nvinfer1::Dims{4, {0, 0, 0, 0}}, input->getDimensions(),
nvinfer1::Dims{4, {1, 1, 1, 1}});
assert(slice != nullptr);
std::string sliceLayerName = "slice_" + std::to_string(layerIdx);
slice->setName(sliceLayerName.c_str());

4
nvdsinfer_custom_impl_Yolo/layers/upsample_layer.cpp
View File

@@ -18,14 +18,14 @@ upsampleLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1:
int stride = std::stoi(block.at("stride"));
float scale[3] = {1, static_cast<float>(stride), static_cast<float>(stride)};
float scale[4] = {1, 1, static_cast<float>(stride), static_cast<float>(stride)};
nvinfer1::IResizeLayer* resize = network->addResize(*input);
assert(resize != nullptr);
std::string resizeLayerName = "upsample_" + std::to_string(layerIdx);
resize->setName(resizeLayerName.c_str());
resize->setResizeMode(nvinfer1::ResizeMode::kNEAREST);
resize->setScales(scale, 3);
resize->setScales(scale, 4);
output = resize->getOutput(0);
return output;