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

Fixes

Browse Source
This commit is contained in:
Marcos Luciano
2023-06-05 16:07:37 -03:00
parent 66a6754b77
commit 79d4a0a8cd
4 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
README.md
View File

@@ -9,7 +9,6 @@ NVIDIA DeepStream SDK 6.2 / 6.1.1 / 6.1 / 6.0.1 / 6.0 / 5.1 configuration for Y
### Future updates ### Future updates
* DeepStream tutorials * DeepStream tutorials
* Updated INT8 calibration * Updated INT8 calibration
* Support for segmentation models * Support for segmentation models
* Support for classification models * Support for classification models

6
nvdsinfer_custom_impl_Yolo/yolo.cpp
View File

@@ -55,7 +55,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
assert(builder); assert(builder);
nvinfer1::NetworkDefinitionCreationFlags flags = nvinfer1::NetworkDefinitionCreationFlags flags =
(1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)); 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flags); nvinfer1::INetworkDefinition* network = builder->createNetworkV2(flags);
assert(network); assert(network);
@@ -64,7 +64,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
if (m_NetworkType == "onnx") { if (m_NetworkType == "onnx") {
parser = nvonnxparser::createParser(*network, *builder->getLogger()); parser = nvonnxparser::createParser(*network, *builder->getLogger());
if (!parser->parseFromFile(m_OnnxWtsFilePath.c_str(), static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING))) { if (!parser->parseFromFile(m_OnnxWtsFilePath.c_str(), static_cast<INT>(nvinfer1::ILogger::Severity::kWARNING))) {
std::cerr << "\nCould not parse the ONNX model\n" << std::endl; std::cerr << "\nCould not parse the ONNX model\n" << std::endl;
#if NV_TENSORRT_MAJOR >= 8 #if NV_TENSORRT_MAJOR >= 8
@@ -99,7 +99,7 @@ Yolo::createEngine(nvinfer1::IBuilder* builder, nvinfer1::IBuilderConfig* config
if (!m_ImplicitBatch && network->getInput(0)->getDimensions().d[0] == -1) { if (!m_ImplicitBatch && network->getInput(0)->getDimensions().d[0] == -1) {
nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile(); nvinfer1::IOptimizationProfile* profile = builder->createOptimizationProfile();
assert(profile); assert(profile);
for (int32_t i = 0; i < network->getNbInputs(); ++i) { for (INT i = 0; i < network->getNbInputs(); ++i) {
nvinfer1::ITensor* input = network->getInput(i); nvinfer1::ITensor* input = network->getInput(i);
nvinfer1::Dims inputDims = input->getDimensions(); nvinfer1::Dims inputDims = input->getDimensions();
nvinfer1::Dims dims = inputDims; nvinfer1::Dims dims = inputDims;

6
nvdsinfer_custom_impl_Yolo/yolo.h
View File

@@ -41,6 +41,12 @@
#include "layers/pooling_layer.h" #include "layers/pooling_layer.h"
#include "layers/reorg_layer.h" #include "layers/reorg_layer.h"
#if NV_TENSORRT_MAJOR >= 8
#define INT int32_t
#else
#define INT int
#endif
struct NetworkInfo struct NetworkInfo
{ {
std::string inputBlobName; std::string inputBlobName;

6
nvdsinfer_custom_impl_Yolo/yoloPlugins.h
View File

@@ -38,12 +38,6 @@
} \ } \
} }
#if NV_TENSORRT_MAJOR >= 8
#define INT int32_t
#else
#define INT int
#endif
namespace { namespace {
const char* YOLOLAYER_PLUGIN_VERSION {"1"}; const char* YOLOLAYER_PLUGIN_VERSION {"1"};
const char* YOLOLAYER_PLUGIN_NAME {"YoloLayer_TRT"}; const char* YOLOLAYER_PLUGIN_NAME {"YoloLayer_TRT"};