122 lines
3.8 KiB
C++
122 lines
3.8 KiB
C++
// trtNet.h
|
|
#ifndef __TRTNET_H__
|
|
#define __TRTNET_H__
|
|
|
|
#include <cassert>
|
|
#include <iostream>
|
|
#include <cstring>
|
|
#include <sstream>
|
|
#include <fstream>
|
|
#include <cuda_runtime_api.h>
|
|
|
|
#include "NvInfer.h"
|
|
#include "NvCaffeParser.h"
|
|
|
|
using namespace nvinfer1;
|
|
using namespace nvcaffeparser1;
|
|
|
|
#if NV_TENSORRT_MAJOR >= 8
|
|
#define NOEXCEPT noexcept
|
|
#else // NV_TENSORRT_MAJOR < 8
|
|
#define NOEXCEPT
|
|
#endif // NV_TENSORRT_MAJOR
|
|
|
|
namespace trtnet {
|
|
|
|
class Logger : public ILogger
|
|
{
|
|
void log(Severity severity, const char *msg) NOEXCEPT override
|
|
{
|
|
if (severity != Severity::kINFO)
|
|
std::cout << msg << std::endl;
|
|
}
|
|
};
|
|
|
|
class IHostMemoryFromFile : public IHostMemory
|
|
{
|
|
public:
|
|
IHostMemoryFromFile(std::string filename) {
|
|
std::ifstream infile(filename, std::ifstream::binary |
|
|
std::ifstream::ate);
|
|
_s = infile.tellg();
|
|
infile.seekg(0, std::ios::beg);
|
|
_mem = malloc(_s);
|
|
infile.read(reinterpret_cast<char*>(_mem), _s);
|
|
}
|
|
#if NV_TENSORRT_MAJOR >= 6
|
|
void* data() const noexcept { return _mem; }
|
|
std::size_t size() const noexcept { return _s; }
|
|
DataType type () const noexcept { return DataType::kFLOAT; } // not used
|
|
void destroy() noexcept { free(_mem); }
|
|
#else // NV_TENSORRT_MAJOR < 6
|
|
void* data() const { return _mem; }
|
|
std::size_t size() const { return _s; }
|
|
DataType type () const { return DataType::kFLOAT; } // not used
|
|
void destroy() { free(_mem); }
|
|
#endif // NV_TENSORRT_MAJOR
|
|
private:
|
|
void *_mem{nullptr};
|
|
std::size_t _s;
|
|
};
|
|
|
|
class TrtGooglenet
|
|
{
|
|
public:
|
|
TrtGooglenet();
|
|
// init from engine file
|
|
void initEngine(std::string filePath, int dataDims[3], int probDims[3]);
|
|
void forward(float *imgs, float *prob);
|
|
void destroy();
|
|
|
|
private:
|
|
Logger _gLogger;
|
|
IHostMemoryFromFile *_gieModelStream{nullptr};
|
|
IRuntime *_runtime;
|
|
ICudaEngine *_engine;
|
|
IExecutionContext *_context;
|
|
cudaStream_t _stream;
|
|
void *_gpu_buffers[2];
|
|
int _blob_sizes[2];
|
|
int _binding_data;
|
|
int _binding_prob;
|
|
|
|
void _initEngine(std::string filePath);
|
|
};
|
|
|
|
class TrtMtcnnDet
|
|
{
|
|
public:
|
|
TrtMtcnnDet();
|
|
// init from engine file
|
|
void initDet1(std::string filePath, int dataDims[3], int prob1Dims[3], int boxesDims[3]);
|
|
void initDet2(std::string filePath, int dataDims[3], int prob1Dims[3], int boxesDims[3]);
|
|
void initDet3(std::string filePath, int dataDims[3], int prob1Dims[3], int boxesDims[3], int marksDims[3]);
|
|
void setBatchSize(int value);
|
|
int getBatchSize();
|
|
void forward(float *imgs, float *probs, float *boxes, float *);
|
|
void destroy();
|
|
|
|
private:
|
|
Logger _gLogger;
|
|
IHostMemoryFromFile *_gieModelStream{nullptr};
|
|
IRuntime *_runtime;
|
|
ICudaEngine *_engine;
|
|
IExecutionContext *_context;
|
|
cudaStream_t _stream;
|
|
void *_gpu_buffers[4];
|
|
int _blob_sizes[4];
|
|
int _num_bindings = 0;
|
|
int _binding_data;
|
|
int _binding_prob1;
|
|
int _binding_boxes;
|
|
int _binding_marks;
|
|
int _batchsize = 0;
|
|
|
|
void _initEngine(std::string filePath, const char *dataName, const char *prob1Name, const char *boxesName, const char *marksName);
|
|
void _setBlobSizes(int dataDims[3], int prob1Dims[3], int boxesDims[3]);
|
|
};
|
|
|
|
} // namespace trtnet
|
|
|
|
#endif // __TRTNET_H__
|