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

Add YOLOv8 support

Browse Source
This commit is contained in:
Marcos Luciano
2023-01-27 15:56:00 -03:00
parent f1cd701247
commit f9c7a4dfca
59 changed files with 3260 additions and 2763 deletions
Download Patch File Download Diff File Expand all files Collapse all files

79
nvdsinfer_custom_impl_Yolo/layers/pooling_layer.cpp
View File

@@ -5,53 +5,50 @@
#include "pooling_layer.h"
nvinfer1::ITensor* poolingLayer(
int layerIdx,
std::map<std::string, std::string>& block,
nvinfer1::ITensor* input,
#include <cassert>
#include <iostream>
nvinfer1::ITensor*
poolingLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
nvinfer1::INetworkDefinition* network)
{
nvinfer1::ITensor* output;
nvinfer1::ITensor* output;
assert(block.at("type") == "maxpool" || block.at("type") == "avgpool");
assert(block.at("type") == "maxpool" || block.at("type") == "avgpool");
if (block.at("type") == "maxpool")
{
assert(block.find("size") != block.end());
assert(block.find("stride") != block.end());
if (block.at("type") == "maxpool") {
assert(block.find("size") != block.end());
assert(block.find("stride") != block.end());
int size = std::stoi(block.at("size"));
int stride = std::stoi(block.at("stride"));
int size = std::stoi(block.at("size"));
int stride = std::stoi(block.at("stride"));
nvinfer1::IPoolingLayer* maxpool
= network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}});
assert(maxpool != nullptr);
std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
maxpool->setName(maxpoolLayerName.c_str());
maxpool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
maxpool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
if (size == 2 && stride == 1)
{
maxpool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
maxpool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
}
output = maxpool->getOutput(0);
}
else if (block.at("type") == "avgpool")
{
nvinfer1::Dims inputDims = input->getDimensions();
nvinfer1::IPoolingLayer* avgpool = network->addPoolingNd(
*input, nvinfer1::PoolingType::kAVERAGE, nvinfer1::Dims{2, {inputDims.d[1], inputDims.d[2]}});
assert(avgpool != nullptr);
std::string avgpoolLayerName = "avgpool_" + std::to_string(layerIdx);
avgpool->setName(avgpoolLayerName.c_str());
output = avgpool->getOutput(0);
}
else
{
std::cerr << "Pooling not supported: " << block.at("type") << std::endl;
std::abort();
nvinfer1::IPoolingLayer* maxpool = network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX,
nvinfer1::Dims{2, {size, size}});
assert(maxpool != nullptr);
std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
maxpool->setName(maxpoolLayerName.c_str());
maxpool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
maxpool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
if (size == 2 && stride == 1) {
maxpool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
maxpool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
}
output = maxpool->getOutput(0);
}
else if (block.at("type") == "avgpool") {
nvinfer1::Dims inputDims = input->getDimensions();
nvinfer1::IPoolingLayer* avgpool = network->addPoolingNd(*input, nvinfer1::PoolingType::kAVERAGE,
nvinfer1::Dims{2, {inputDims.d[1], inputDims.d[2]}});
assert(avgpool != nullptr);
std::string avgpoolLayerName = "avgpool_" + std::to_string(layerIdx);
avgpool->setName(avgpoolLayerName.c_str());
output = avgpool->getOutput(0);
}
else {
std::cerr << "Pooling not supported: " << block.at("type") << std::endl;
assert(0);
}
return output;
return output;
}