//g++ inferPeopleNet.cpp `pkg-config --cflags --libs opencv4` -I /usr/local/cuda-10.2/include -I /usr/include/aarch64-linux-gnu/ -L /usr/lib/aarch64-linux-gnu -lnvinfer -lnvinfer_plugin -L /usr/local/cuda-10.2/lib64 -lcudart -lcublas -lcurand
/* OpenCV headers */
#include <opencv2/core/core.hpp>
#include <opencv2/dnn/dnn.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include <opencv2/imgcodecs/imgcodecs.hpp>
#include <iostream>
#include <fstream>
#include <cudnn.h>
#include "NvInfer.h"
#include "NvInferPlugin.h"
//#define MIN(a,b) ((a) < (b) ? (a) : (b))
//#define MAX(a,b) ((a) > (b) ? (a) : (b))
#define CLIP(a,min,max) (MAX(MIN(a, max), min))
#define DIVIDE_AND_ROUND_UP(a, b) ((a + b - 1) / b)
using namespace cv;
using namespace std;
class Logger : public nvinfer1::ILogger
{
public:
void log(nvinfer1::ILogger::Severity severity, const char* msg) override
{
// suppress info-level messages
if (severity == Severity::kINFO) return;
switch (severity)
{
case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break;
case Severity::kERROR: std::cerr << "ERROR: "; break;
case Severity::kWARNING: std::cerr << "WARNING: "; break;
case Severity::kINFO: std::cerr << "INFO: "; break;
default: std::cerr << "UNKNOWN: "; break;
}
std::cerr << msg << std::endl;
}
};
int main()
{
std::string engineFilePath = "/opt/nvidia/deepstream/deepstream-5.0/samples/models/tlt_pretrained_models/peoplenet/resnet34_peoplenet_pruned.etlt_b1_gpu0_fp16.engine";
std::string imagePath = "test.jpg";
// General parameters
uint16_t m_InputH, m_OutputW;
uint16_t m_InputW, m_OutputH;
uint16_t m_InputC;
uint64_t m_InputSize, m_OutputBBoxSize, m_OutputConfidenceSize;
uint16_t m_NumOutputClasses;
// TRT specific parameters
uint16_t m_maxBatchSize = 1;
int m_InputIndex = -1;
int m_OutputBboxIndex = -1, m_OutputClassIndex = -1;
Logger m_Logger;
nvinfer1::ICudaEngine* m_Engine;
nvinfer1::IExecutionContext* m_Context;
nvinfer1::IRuntime* runtime;
std::vector<void*> m_Bindings;
std::vector<float*> m_TrtOutputBuffers;
cudaStream_t m_CudaStream;
m_InputW = 960;
m_InputH = 544;
m_InputC = 3;
m_OutputW = 60;
m_OutputH = 34;
m_NumOutputClasses = 3;
m_InputSize = m_InputW * m_InputH * m_InputC;
m_OutputBBoxSize = m_OutputW * m_OutputH * m_NumOutputClasses * 4;
m_OutputConfidenceSize = m_OutputW * m_OutputH * m_NumOutputClasses;
// Deserializing engine
// reading the model in memory
std::cout << "[Info] Loading TRT Engine...\n";
std::stringstream trtModelStream;
trtModelStream.seekg(0, trtModelStream.beg);
std::ifstream cache(engineFilePath);
assert(cache.good());
trtModelStream << cache.rdbuf();
cache.close();
// calculating model size
trtModelStream.seekg(0, std::ios::end);
const int modelSize = trtModelStream.tellg();
trtModelStream.seekg(0, std::ios::beg);
void* modelMem = malloc(modelSize);
trtModelStream.read((char*) modelMem, modelSize);
runtime = nvinfer1::createInferRuntime(m_Logger);
m_Engine = runtime->deserializeCudaEngine(modelMem, modelSize, nullptr);
free(modelMem);
runtime->destroy();
std::cout << "[Info] Loading Complete!\n";
if(m_Engine == nullptr)
{
std::cout << "[Error] TensorRT engine loading failed\n";
return -1;
}
m_Context = m_Engine->createExecutionContext();
if(m_Context == nullptr)
{
std::cout << "[Error] TensorRT getting context failed\n";
return -2;
}
// Get the bindings
std::cout << "[Info] Getting the Bindings...\n";
m_Bindings.resize(m_Engine->getNbBindings(), nullptr);
m_TrtOutputBuffers.resize(m_Engine->getNbBindings() - 1, nullptr);
m_InputIndex = m_Engine->getBindingIndex("input_1");
m_OutputBboxIndex = m_Engine->getBindingIndex("output_bbox/BiasAdd");
m_OutputClassIndex = m_Engine->getBindingIndex("output_cov/Sigmoid");
if (m_InputIndex == -1 || m_OutputBboxIndex == -1 || m_OutputClassIndex == -1)
{
std::cout << "[Error] TensorRT binding not found\n";
return -3;
}
std::cout << "[Info] Bindings size : " << m_Engine->getNbBindings() << "\n";
std::cout << "[Info] Bindings " << m_InputIndex << " " << m_OutputBboxIndex << " " << m_OutputClassIndex << "\n";
// Allocate Buffers
(cudaMalloc(&m_Bindings.at(m_InputIndex), m_maxBatchSize * m_InputSize * sizeof(float)));
(cudaMalloc(&m_Bindings.at(m_OutputBboxIndex), m_maxBatchSize * m_OutputBBoxSize * sizeof(float)));
(cudaMalloc(&m_Bindings.at(m_OutputClassIndex), m_maxBatchSize * m_OutputConfidenceSize * sizeof(float)));
(cudaMallocHost(&m_TrtOutputBuffers[0], m_OutputBBoxSize * m_maxBatchSize * sizeof(float)));
(cudaMallocHost(&m_TrtOutputBuffers[1], m_OutputConfidenceSize * m_maxBatchSize * sizeof(float)));
(cudaStreamCreate(&m_CudaStream));
// Loading input image to device
std::cout << "[Info] Loading input image\n";
Mat inputImage = imread(imagePath);
Mat inferImage = cv::dnn::blobFromImage(inputImage, 0.0039215697906911373, cv::Size(m_InputW, m_InputH), cv::Scalar(0.0, 0.0, 0.0), true, false);
//Mat inferImage = cv::dnn::blobFromImage(inputImage, 1.0, cv::Size(m_InputW, m_InputH), cv::Scalar(0.0, 0.0, 0.0), false, false);
cudaMemcpyAsync(m_Bindings.at(m_InputIndex), inputImage.data,
m_maxBatchSize * m_InputSize * sizeof(float), cudaMemcpyHostToDevice,
m_CudaStream);
// Running Inference
std::cout << "[Info] Running Inference\n";
m_Context->enqueue(m_maxBatchSize, m_Bindings.data(), m_CudaStream, nullptr);
cudaMemcpyAsync(m_TrtOutputBuffers.at(0), m_Bindings.at(m_OutputBboxIndex),
m_maxBatchSize * m_OutputBBoxSize * sizeof(float),
cudaMemcpyDeviceToHost, m_CudaStream);
cudaMemcpyAsync(m_TrtOutputBuffers.at(1), m_Bindings.at(m_OutputClassIndex),
m_maxBatchSize * m_OutputConfidenceSize * sizeof(float),
cudaMemcpyDeviceToHost, m_CudaStream);
// Decoding output buffers
std::cout << "[Info] Decoding the output Buffers\n";
int gridW = m_OutputW;
int gridH = m_OutputH;
int gridSize = gridW * gridH;
float gcCentersX[gridW];
float gcCentersY[gridH];
float bboxNormX = 35.0;
float bboxNormY = 35.0;
float* outputBboxBuf = &m_TrtOutputBuffers.at(0)[0];
float* outputCovBuf = &m_TrtOutputBuffers.at(1)[0];
int strideX = DIVIDE_AND_ROUND_UP(m_InputW, gridW);
int strideY = DIVIDE_AND_ROUND_UP(m_InputH, gridH);
for (int i = 0; i < gridW; i++)
{
gcCentersX[i] = (float)(i * strideX + 0.5);
gcCentersX[i] /= (float)bboxNormX;
}
for (int i = 0; i < gridH; i++)
{
gcCentersY[i] = (float)(i * strideY + 0.5);
gcCentersY[i] /= (float)bboxNormY;
}
for (int c = 0; c < m_NumOutputClasses; c++)
{
float *outputX1 = outputBboxBuf + (c * 4 * gridW * gridH);
float *outputY1 = outputX1 + gridSize;
float *outputX2 = outputY1 + gridSize;
float *outputY2 = outputX2 + gridSize;
float threshold = 0.1;//detectionParams.perClassPreclusterThreshold[c];
for (int h = 0; h < gridH; h++)
{
for (int w = 0; w < gridW; w++)
{
int i = w + h * gridW;
float confidence = outputCovBuf[c * gridSize + i];
if (confidence >= threshold)
{
//NvDsInferObjectDetectionInfo object;
float rectX1f, rectY1f, rectX2f, rectY2f;
rectX1f = (outputX1[w + h * gridW] - gcCentersX[w]) * -bboxNormX;
rectY1f = (outputY1[w + h * gridW] - gcCentersY[h]) * -bboxNormY;
rectX2f = (outputX2[w + h * gridW] + gcCentersX[w]) * bboxNormX;
rectY2f = (outputY2[w + h * gridW] + gcCentersY[h]) * bboxNormY;
//object.classId = c;
//object.detectionConfidence = outputCovBuf[c * gridSize + i];
/* Clip object box co-ordinates to network resolution */
rectX1f = CLIP(rectX1f, 0, m_InputW - 1);
rectY1f = CLIP(rectY1f, 0, m_InputH - 1);
rectX2f = CLIP(rectX2f, 0, m_InputW - 1);
rectY2f = CLIP(rectY2f, 0, m_InputH - 1);
//Prevent underflows
if(((rectX2f - rectX1f) < 0) || ((rectY2f - rectY1f) < 0))
continue;
// Detected boxes
std::cout << "[Info] ClassIdx : " << c << " BBox : " << rectX1f << "," << rectY1f << "," << (rectX2f) << "," << (rectY2f) << "," << confidence << "\n";
}
}
}
}
//imshow("Display", inputImage);
//waitKey();
return 0;
}
