Replaced raw with smart pointers to solve the memory leak issue. Problem was on ros_segment_inference

Author: IasonTheodorou

inc/yolo_inference.h

@@ -80,7 +80,7 @@ class YOLO_V8
 
 private:
     Ort::Env env;
-    Ort::Session* session;
+    std::unique_ptr<Ort::Session> session;
     bool cudaEnable;
     Ort::RunOptions options;
     std::vector<const char*> inputNodeNames;

src/main.cpp

@@ -5,7 +5,7 @@
 #include <fstream>
 #include <random>
 
-void Detector(YOLO_V8*& p) {
+void Detector(std::unique_ptr<YOLO_V8>& p) {
     std::filesystem::path current_path = std::filesystem::current_path();
     std::filesystem::path imgs_path = current_path / "images";
     for (auto& i : std::filesystem::directory_iterator(imgs_path))
@@ -58,7 +58,7 @@ void Detector(YOLO_V8*& p) {
 }
 
 
-void Classifier(YOLO_V8*& p)
+void Classifier(std::unique_ptr<YOLO_V8>& p)
 {
     std::filesystem::path current_path = std::filesystem::current_path();
     std::filesystem::path imgs_path = current_path;// / "images"
@@ -97,7 +97,7 @@ void Classifier(YOLO_V8*& p)
 
 
 
-int ReadCocoYaml(YOLO_V8*& p) {
+int ReadCocoYaml(std::unique_ptr<YOLO_V8>& p) {
     // Open the YAML file
     std::ifstream file("coco.yaml");
     if (!file.is_open())
@@ -148,7 +148,9 @@ int ReadCocoYaml(YOLO_V8*& p) {
 
 void DetectTest()
 {
-    YOLO_V8* yoloDetector = new YOLO_V8;
+    //YOLO_V8* yoloDetector = new YOLO_V8;
+    std::unique_ptr<YOLO_V8> yoloDetector = std::make_unique<YOLO_V8>();
+
     ReadCocoYaml(yoloDetector);
     DL_INIT_PARAM params;
     params.rectConfidenceThreshold = 0.1;
@@ -177,7 +179,7 @@ void DetectTest()
 
 void ClsTest()
 {
-    YOLO_V8* yoloDetector = new YOLO_V8;
+    std::unique_ptr<YOLO_V8> yoloDetector = std::make_unique<YOLO_V8>();
     std::string model_path = "cls.onnx";
     ReadCocoYaml(yoloDetector);
     DL_INIT_PARAM params{ model_path, YOLO_CLS, {224, 224} };

src/yolo_inference.cpp

@@ -9,7 +9,12 @@ YOLO_V8::YOLO_V8() {
 
 
 YOLO_V8::~YOLO_V8() {
-    delete session;
+    for (auto& name : inputNodeNames) {
+        delete[] name;
+    }
+    for (auto& name : outputNodeNames) {
+        delete[] name;
+    }
 }
 
 #ifdef USE_CUDA
@@ -95,6 +100,19 @@ char* YOLO_V8::PreProcess(const cv::Mat& iImg, std::vector<int> iImgSize, cv::Ma
 
 const char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) {
     const char* Ret = RET_OK;
+    if (session) {
+
+        // Clear node names
+        for (auto& name : inputNodeNames) {
+            delete[] name;
+        }
+        inputNodeNames.clear();
+
+        for (auto& name : outputNodeNames) {
+            delete[] name;
+        }
+        outputNodeNames.clear();
+    }
     std::regex pattern("[\u4e00-\u9fa5]");
     bool result = std::regex_search(iParams.modelPath, pattern);
     if (result)
@@ -132,7 +150,7 @@ const char* YOLO_V8::CreateSession(DL_INIT_PARAM& iParams) {
         const char* modelPath = iParams.modelPath.c_str();
 #endif // _WIN32
 
-        session = new Ort::Session(env, modelPath, sessionOption);
+        session = std::make_unique<Ort::Session>(env, modelPath, sessionOption);
         Ort::AllocatorWithDefaultOptions allocator;
         size_t inputNodesNum = session->GetInputCount();
         for (size_t i = 0; i < inputNodesNum; i++)