fix memory leak, #27

UNeedCryDear · Dec 29, 2023 · cc89fe3 · cc89fe3
1 parent 67c0dfe
commit cc89fe3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 11 deletions.
diff --git a/yolov5_onnx.h b/yolov5_onnx.h
@@ -14,7 +14,11 @@ class Yolov5Onnx {
 public:
 	Yolov5Onnx() :_OrtMemoryInfo(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeCPUOutput)) {};
 
-	~Yolov5Onnx() {};
+	~Yolov5Onnx() {
+		if (_OrtSession != nullptr) {
+			delete _OrtSession;
+	}
+	};
 
 
 

diff --git a/yolov5_seg_onnx.h b/yolov5_seg_onnx.h
@@ -13,8 +13,11 @@
 class Yolov5SegOnnx {
 public:
 	Yolov5SegOnnx() :_OrtMemoryInfo(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeCPUOutput)) {};
-
-	~Yolov5SegOnnx() {};
+	~Yolov5SegOnnx() {
+		if (_OrtSession != nullptr) {
+			delete _OrtSession;
+		};
+	}
 
 
 
@@ -25,27 +28,27 @@ class Yolov5SegOnnx {
 	* \param[in] cudaID:if isCuda==true,run Ort-GPU on cudaID.
 	* \param[in] warmUp:if isCuda==true,warm up GPU-model.
 	*/
-	bool ReadModel(const std::string& modelPath, bool isCuda = false, int cudaID = 0, bool warmUp = true);
+	bool ReadModel(const std::string & modelPath, bool isCuda = false, int cudaID = 0, bool warmUp = true);
 
 	/** \brief  detect.
 	* \param[in] srcImg:a 3-channels image.
 	* \param[out] output:detection results of input image.
 	*/
-	bool OnnxDetect(cv::Mat& srcImg, std::vector<OutputSeg>& output);
+	bool OnnxDetect(cv::Mat & srcImg, std::vector<OutputSeg>&output);
 	/** \brief  detect,batch size= _batchSize
 	* \param[in] srcImg:A batch of images.
 	* \param[out] output:detection results of input images.
 	*/
-	bool OnnxBatchDetect(std::vector<cv::Mat>& srcImg, std::vector<std::vector<OutputSeg>>& output);
+	bool OnnxBatchDetect(std::vector<cv::Mat>&srcImg, std::vector<std::vector<OutputSeg>>&output);
 
 private:
 
 	template <typename T>
-	T VectorProduct(const std::vector<T>& v)
+	T VectorProduct(const std::vector<T>&v)
 	{
 		return std::accumulate(v.begin(), v.end(), 1, std::multiplies<T>());
 	};
-	int Preprocessing(const std::vector<cv::Mat>& SrcImgs, std::vector<cv::Mat>& OutSrcImgs, std::vector<cv::Vec4d>& params);
+	int Preprocessing(const std::vector<cv::Mat>&SrcImgs, std::vector<cv::Mat>&OutSrcImgs, std::vector<cv::Vec4d>&params);
 
 
 	const int _netWidth = 640;   //ONNX-net-input-width
@@ -55,15 +58,15 @@ class Yolov5SegOnnx {
 
 	float _classThreshold = 0.5;
 	float _nmsThreshold = 0.45;
-	float _maskThreshold = 0.5; 
+	float _maskThreshold = 0.5;
 
 	Ort::Env _OrtEnv = Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_ERROR, "Yolov5-Seg");
 	Ort::SessionOptions _OrtSessionOptions = Ort::SessionOptions();
 	Ort::Session* _OrtSession = nullptr;
 	Ort::MemoryInfo _OrtMemoryInfo;
 	Ort::AllocatorWithDefaultOptions _OrtAllocator;
 	Ort::RunOptions _OrtRunOptions{ nullptr };
-	OrtStatus* _OrtStatus{nullptr};
+	OrtStatus* _OrtStatus{ nullptr };
 #if ORT_API_VERSION < ORT_OLD_VISON
 
 	char* _inputName, * _output_name0, * _output_name1;
@@ -97,4 +100,4 @@ class Yolov5SegOnnx {
 
 
 
-};
+	};