From 6eb847801e6d1d66105a8ec2fe8c449c9e24da4c Mon Sep 17 00:00:00 2001 From: 孙天宇 <suntianyu@aiotlink.com> Date: 星期二, 12 七月 2022 15:52:27 +0800 Subject: [PATCH] Merge branch '3090' of http://192.168.5.5:10010/r/Algorithm/baseDetector into 3090 --- src/detecter_tools/model.h | 13 ++++++++++++- 1 files changed, 12 insertions(+), 1 deletions(-) diff --git a/src/detecter_tools/model.h b/src/detecter_tools/model.h index eb3adff..e993f44 100644 --- a/src/detecter_tools/model.h +++ b/src/detecter_tools/model.h @@ -3,7 +3,7 @@ #include "plugin_factory.h" #include "trt_utils.h" - +#include "calibrator.h" #include "NvInfer.h" #include "NvInferPlugin.h" #include "NvInferRuntimeCommon.h" @@ -16,6 +16,7 @@ #include "../utils/time_util.h" #include "../config.h" #include "opencv2/opencv.hpp" +#include <numeric> struct NetworkInfo { @@ -78,6 +79,7 @@ const std::string m_DeviceType; const std::string m_InputBlobName; std::vector<TensorInfo> m_OutputTensors; + std::vector<std::map<std::string, std::string>> m_configBlocks; uint32_t m_InputH; uint32_t m_InputW; uint32_t m_InputC; @@ -172,12 +174,21 @@ void setOutput(int type); private: Logger m_Logger; + void createYOLOEngine(const nvinfer1::DataType dataType = nvinfer1::DataType::kFLOAT, + Int8EntropyCalibrator* calibrator = nullptr); + void writePlanFileToDisk(); + std::vector<std::map<std::string, std::string>> parseConfigFile(const std::string cfgFilePath); + void parseConfigBlocks(); void allocateBuffers(); bool verifyEngine(); void destroyNetworkUtils(std::vector<nvinfer1::Weights>& trtWeights); +protected: + const std::string m_NetworkType; + std::unique_ptr<YoloTinyMaxpoolPaddingFormula> m_TinyMaxpoolPaddingFormula; private: Timer _timer; + int _n_yolo_ind = 0; }; #endif -- Gitblit v1.8.0