Updated initializer and SegmentAnything modules to store the data to the custom result structs properly

Author: IasonTheodorou

CMakeLists.txt

@@ -19,7 +19,7 @@ include_directories(${OpenCV_INCLUDE_DIRS})
 
 # -------------- ONNXRuntime  ------------------#
 set(ONNXRUNTIME_VERSION 1.21.0)
-set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../hero_sam/onnxruntime-linux-x64-gpu-1.21.1")
+set(ONNXRUNTIME_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../hero_sam.bak/onnxruntime-linux-x64-gpu-1.21.1")
 include_directories(${ONNXRUNTIME_ROOT}/include)
 
 # -------------- Cuda ------------------#
@@ -84,8 +84,8 @@ add_executable(${PROJECT_NAME} src/main.cpp)
 target_link_libraries(${PROJECT_NAME} sam_onnx_ros_core)
 
 # Copy sam_<model>.onnx file to the same folder of the executable file
-configure_file(~/Documents/repos/hero_sam/sam_inference/model/SAM_mask_decoder.onnx ${CMAKE_CURRENT_BINARY_DIR}/SAM_mask_decoder.onnx  COPYONLY)
-configure_file(~/Documents/repos/hero_sam/sam_inference/model/SAM_encoder.onnx ${CMAKE_CURRENT_BINARY_DIR}/SAM_encoder.onnx COPYONLY)
+configure_file(~/Documents/repos/hero_sam.bak/sam_inference/model/SAM_mask_decoder.onnx ${CMAKE_CURRENT_BINARY_DIR}/SAM_mask_decoder.onnx  COPYONLY)
+configure_file(~/Documents/repos/hero_sam.bak/sam_inference/model/SAM_encoder.onnx ${CMAKE_CURRENT_BINARY_DIR}/SAM_encoder.onnx COPYONLY)
 
 # Create folder name images in the same folder of the executable file
 add_custom_command(TARGET ${PROJECT_NAME} POST_BUILD

include/segmentation.h

@@ -4,7 +4,9 @@
 #include <tuple>
 
 #include "sam_inference.h"
-std::tuple<std::vector<std::unique_ptr<SAM>>, SEG::_DL_INIT_PARAM, SEG::_DL_INIT_PARAM> Initializer();
-std::vector<cv::Mat> SegmentAnything(std::vector<std::unique_ptr<SAM>>& samSegmentors, const SEG::_DL_INIT_PARAM& params_encoder, const SEG::_DL_INIT_PARAM& params_decoder, cv::Mat& img);
+std::tuple<std::vector<std::unique_ptr<SAM>>, SEG::_DL_INIT_PARAM, SEG::_DL_INIT_PARAM, SEG::DL_RESULT, std::vector<SEG::DL_RESULT>> Initializer();
+void SegmentAnything(std::vector<std::unique_ptr<SAM>>& samSegmentors, const SEG::_DL_INIT_PARAM& params_encoder, const SEG::_DL_INIT_PARAM& params_decoder, const cv::Mat& img,
+std::vector<SEG::DL_RESULT> &resSam,
+  SEG::DL_RESULT &res);
 
 #endif // SEGMENTATION_H

src/main.cpp

@@ -10,7 +10,9 @@ int main()
     std::vector<std::unique_ptr<SAM>> samSegmentors;
     SEG::DL_INIT_PARAM params_encoder;
     SEG::DL_INIT_PARAM params_decoder;
-    std::tie(samSegmentors, params_encoder, params_decoder) = Initializer();
+    std::vector<SEG::DL_RESULT> resSam;
+    SEG::DL_RESULT res;
+    std::tie(samSegmentors, params_encoder, params_decoder, res, resSam) = Initializer();
     std::filesystem::path current_path = std::filesystem::current_path();
     std::filesystem::path imgs_path =  "/home/amigo/Documents/repos/hero_sam/sam_inference/build/images"; // current_path / <- you could use
     for (auto &i : std::filesystem::directory_iterator(imgs_path))
@@ -19,8 +21,8 @@ int main()
         {
             std::string img_path = i.path().string();
             cv::Mat img = cv::imread(img_path);
-            std::vector<cv::Mat> masks;
-            masks = SegmentAnything(samSegmentors, params_encoder, params_decoder, img);
+
+            SegmentAnything(samSegmentors, params_encoder, params_decoder, img, resSam, res);
 
         }
     }

src/segmentation.cpp

@@ -1,7 +1,7 @@
 #include "segmentation.h"
 
 std::tuple<std::vector<std::unique_ptr<SAM>>, SEG::DL_INIT_PARAM,
-           SEG::DL_INIT_PARAM>
+           SEG::DL_INIT_PARAM, SEG::DL_RESULT, std::vector<SEG::DL_RESULT>>
 Initializer() {
   std::vector<std::unique_ptr<SAM>> samSegmentors;
   samSegmentors.push_back(std::make_unique<SAM>());
@@ -11,7 +11,8 @@ Initializer() {
   std::unique_ptr<SAM> samSegmentorDecoder = std::make_unique<SAM>();
   SEG::DL_INIT_PARAM params_encoder;
   SEG::DL_INIT_PARAM params_decoder;
-
+  SEG::DL_RESULT res;
+  std::vector<SEG::DL_RESULT> resSam;
   params_encoder.rectConfidenceThreshold = 0.1;
   params_encoder.iouThreshold = 0.5;
   params_encoder.modelPath = "/home/amigo//Documents/repos/sam_onnx_ros/build/SAM_encoder.onnx";
@@ -31,22 +32,21 @@ Initializer() {
   samSegmentorDecoder->CreateSession(params_decoder);
   samSegmentors[0] = std::move(samSegmentorEncoder);
   samSegmentors[1] = std::move(samSegmentorDecoder);
-  return {std::move(samSegmentors), params_encoder, params_decoder};
+  return {std::move(samSegmentors), params_encoder, params_decoder, res, resSam};
 }
 
-std::vector<cv::Mat>
-SegmentAnything(std::vector<std::unique_ptr<SAM>> &samSegmentors,
+void SegmentAnything(std::vector<std::unique_ptr<SAM>> &samSegmentors,
                 const SEG::DL_INIT_PARAM &params_encoder,
-                const SEG::DL_INIT_PARAM &params_decoder, cv::Mat &img) {
+                const SEG::DL_INIT_PARAM &params_decoder, const cv::Mat &img, std::vector<SEG::DL_RESULT> &resSam,
+  SEG::DL_RESULT &res) {
+
 
-  std::vector<SEG::DL_RESULT> resSam;
-  SEG::DL_RESULT res;
 
   SEG::MODEL_TYPE modelTypeRef = params_encoder.modelType;
   samSegmentors[0]->RunSession(img, resSam, modelTypeRef, res);
 
   modelTypeRef = params_decoder.modelType;
   samSegmentors[1]->RunSession(img, resSam, modelTypeRef, res);
 
-  return std::move(res.masks);
+  // return std::move(res.masks);
 }

src/utils.cpp

@@ -1,6 +1,6 @@
 #include "utils.h"
 #include <opencv2/ximgproc/edge_filter.hpp> // for guided filter
-#define LOGGING
+//#define LOGGING
 
 // Constructor
 Utils::Utils()