GIEFeatExtractor.h

#ifndef GIEFEATEXTRACTOR_H_
#define GIEFEATEXTRACTOR_H_
 
#include <string>
#include <math.h>
 
// OpenCV
#include <opencv2/opencv.hpp>
 
// CUDA-C includes
#include <cuda.h>
#include <cuda_runtime.h>
 
// GIE
#include "NvInfer.h"
#include "NvCaffeParser.h"
#include "cudaUtility.h"
 
using namespace cv;
 
// Logger for GIE info/warning/errors
class Logger : public nvinfer1::ILogger         
{
    void log( Severity severity, const char* msg ) override
    {
    if( severity != Severity::kINFO )
        std::cout << msg << std::endl;
    }
};
 
class GIEFeatExtractor {
 
protected:
    
    bool cudaFreeMapped(void *cpuPtr);
 
    bool cudaAllocMapped( void** cpuPtr, void** gpuPtr, size_t size );
 
    bool caffeToGIEModel( const std::string& deployFile,            // name for caffe prototxt
                          const std::string& modelFile,             // name for model
                          const std::string& binaryprotoFile,       // name for .binaryproto
                          const std::vector<std::string>& outputs,  // network outputs
                          unsigned int maxBatchSize,                // batch size - NB must be at least as large as the batch we want to run with)
                          std::ostream& gieModelStream);            // output stream for the GIE model
 
    bool init(string _caffemodel_file,
            string _binaryproto_meanfile, float meanR, float meanG, float meanB, 
            string _prototxt_file, int _resizeWidth, int _resizeHeight,
            string _blob_name);
 
    nvinfer1::IRuntime* mInfer;
    nvinfer1::ICudaEngine* mEngine;
    nvinfer1::IExecutionContext* mContext;
 
    //cv::Mat meanMat;
    float *meanData;
    vector<float> mean_values;
 
    nvinfer1::Dims4 resizeDims;
 
    uint32_t mWidth;
    uint32_t mHeight;
    uint32_t mInputSize;
    float*   mInputCPU;
    float*   mInputCUDA;
    
    uint32_t mOutputSize;
    uint32_t mOutputDims;
    float*   mOutputCPU;
    float*   mOutputCUDA;
 
    Logger  gLogger;
 
public:
 
    string prototxt_file;
    string caffemodel_file;
    string blob_name;
    string binaryproto_meanfile;
 
    bool timing;
    
    GIEFeatExtractor(string _caffemodel_file,
            string _binaryproto_meanfile, float _meanR, float _meanG, float _meanB, 
            string _prototxt_file, int _resizeWidth, int _resizeHeight,
            string _blob_name,
            bool _timing );
 
    ~GIEFeatExtractor();
 
    bool extract_singleFeat_1D(cv::Mat &image, vector<float> &features,  float (&times)[2]);
};
 
#endif /* GIEFEATEXTRACTOR_H_ */