change detection_lidar_PointPillars_MultiHead. for use engine file.

src/detection/detection_lidar_PointPillars_MultiHead/main.cpp

@@ -264,6 +264,25 @@ void exitfunc()
     std::cout<<"exit func complete"<<std::endl;
 }
 
+#include <QDir>
+#include <QFile>
+
+bool trtisexist(std::string strpfe,std::string strbackbone)
+{
+    QFile xFile;
+    xFile.setFileName(strpfe.data());
+    if(xFile.exists() == false)
+    {
+        return false;
+    }
+    xFile.setFileName(strbackbone.data());
+    if(xFile.exists() == false)
+    {
+        return false;
+    }
+    return true;
+}
+
 int main(int argc, char *argv[])
 {
     QCoreApplication a(argc, argv);
@@ -278,8 +297,15 @@ int main(int argc, char *argv[])
     char * strhome = getenv("HOME");
     std::string pfe_file = strhome;
     pfe_file += "/models/lidar/cbgs_pp_multihead_pfe.onnx";
+    std::string pfe_trt_file = strhome;
+    pfe_trt_file += "/models/lidar/cbgs_pp_multihead_pfe.trt";
+
     std::string backbone_file = strhome;
     backbone_file += "/models/lidar/cbgs_pp_multihead_backbone.onnx";
+    std::string backbone_trt_file = strhome;
+    backbone_trt_file += "/models/lidar/cbgs_pp_multihead_backbone.trt";
+
+    bool btrtexist = trtisexist(pfe_trt_file,backbone_trt_file);
 
 
     QString strpath = QCoreApplication::applicationDirPath();
@@ -297,6 +323,10 @@ int main(int argc, char *argv[])
     backbone_file = xparam.GetParam("backbone_file",backbone_file.data());
     gstrinput = xparam.GetParam("input","lidar_pc");
     gstroutput = xparam.GetParam("output","lidar_pointpillar");
+    if(btrtexist == false)
+    {
+
+       std::cout<<"use onnx model."<<std::endl;
     pPillars = new PointPillars(
       0.1,
       0.2,
@@ -305,6 +335,19 @@ int main(int argc, char *argv[])
       backbone_file,
       pp_config
     );
+    }
+    else
+    {
+        std::cout<<"use engine mode."<<std::endl;
+    pPillars = new PointPillars(
+      0.1,
+      0.2,
+      false,
+      pfe_trt_file,
+      backbone_trt_file,
+      pp_config
+    );
+    }
     std::cout<<"PointPillars Init OK."<<std::endl;
 
     gpa = iv::modulecomm::RegisterRecv(gstrinput.data(),ListenPointCloud);

src/detection/detection_lidar_PointPillars_MultiHead/pointpillars.cc

@@ -284,14 +284,32 @@ void PointPillars::SetDeviceMemoryToZero() {
 }
 
 
+#include <fstream>
+#include <iostream>
+#
 
+void PointPillars::SaveEngine(nvinfer1::ICudaEngine* pengine,std::string strpath)
+{
+    nvinfer1::IHostMemory* data = pengine->serialize();
+    std::ofstream file;
+    file.open(strpath,std::ios::binary | std::ios::out);
+    if(!file.is_open())
+    {
+        std::cout << "read create engine file" << strpath <<" failed" << std::endl;
+        return;
+    }
 
+    file.write((const char*)data->data(), data->size());
+    file.close();
+}
 
 void PointPillars::InitTRT(const bool use_onnx) {
   if (use_onnx_) {
     // create a TensorRT model from the onnx model and load it into an engine
     OnnxToTRTModel(pfe_file_, &pfe_engine_);
+    SaveEngine(pfe_engine_,"/home/nvidia/models/lidar/cbgs_pp_multihead_pfe.trt");
     OnnxToTRTModel(backbone_file_, &backbone_engine_);
+    SaveEngine(backbone_engine_,"/home/nvidia/models/lidar/cbgs_pp_multihead_backbone.trt");
   }else {
     EngineToTRTModel(pfe_file_, &pfe_engine_);
     EngineToTRTModel(backbone_file_, &backbone_engine_);
@@ -345,6 +363,59 @@ void PointPillars::OnnxToTRTModel(
     builder->destroy();
 }
 
+void PointPillars::LoadEngineModel(const string &engine_file, nvinfer1::ICudaEngine **engine_ptr)
+{
+//    using namespace std;
+//    fstream file;
+
+//    file.open(engine_file,ios::binary | ios::in);
+//    if(!file.is_open())
+//    {
+//        cout << "read engine file" << engine_file <<" failed" << endl;
+//        return;
+//    }
+//    file.seekg(0, ios::end);
+//    int length = file.tellg();
+//    file.seekg(0, ios::beg);
+//    std::unique_ptr<char[]> data(new char[length]);
+//    file.read(data.get(), length);
+
+//    file.close();
+
+//    std::cout << "deserializing" << std::endl;
+//    mTrtRunTime = createInferRuntime(gLogger);
+//    assert(mTrtRunTime != nullptr);
+//    mTrtEngine= mTrtRunTime->deserializeCudaEngine(data.get(), length, &mTrtPluginFactory);
+//    assert(mTrtEngine != nullptr);
+
+//    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(g_logger_);
+
+//    if (runtime == nullptr) {
+//        std::string msg("failed to build runtime parser");
+//        g_logger_.log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());
+//        exit(EXIT_FAILURE);
+//    }
+
+
+
+//    std::cout << "                                                                  "<< std::endl;
+//    std::cout << "------------------------------------------------------------------"<< std::endl;
+//    std::cout << ">>>>                                                          >>>>"<< std::endl;
+//    std::cout << "                                                                  "<< std::endl;
+//    std::cout << "Input filename:   " << engine_file << std::endl;
+//    std::cout << "                                                                  "<< std::endl;
+//    std::cout << ">>>>                                                          >>>>"<< std::endl;
+//    std::cout << "------------------------------------------------------------------"<< std::endl;
+//    std::cout << "                                                                  "<< std::endl;
+
+//    nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(modelMem, modelSize, NULL);
+//    if (engine == nullptr) {
+//        std::string msg("failed to build engine parser");
+//        g_logger_.log(nvinfer1::ILogger::Severity::kERROR, msg.c_str());
+//        exit(EXIT_FAILURE);
+//    }
+//    *engine_ptr = engine;
+}
 
 void PointPillars::EngineToTRTModel(
     const std::string &engine_file ,     

src/detection/detection_lidar_PointPillars_MultiHead/pointpillars.h

@@ -243,6 +243,9 @@ class PointPillars {
     void EngineToTRTModel(const std::string &engine_file ,     
                         nvinfer1::ICudaEngine** engine_ptr) ;
 
+    void LoadEngineModel(const std::string &engine_file ,
+                         nvinfer1::ICudaEngine** engine_ptr) ;
+
     /**
      * @brief Preproces points
      * @param[in] in_points_array Point cloud array
@@ -269,6 +272,8 @@ class PointPillars {
                 const std::string pp_config);
     ~PointPillars();
 
+    void SaveEngine(nvinfer1::ICudaEngine* pengine,std::string strpath);
+
     /**
      * @brief Call PointPillars for the inference
      * @param[in] in_points_array Point cloud array