update test.cpp CMakeLists.txt

vision_test/CMakeLists.txt

@@ -5,7 +5,10 @@ if(CMAKE_COMPILER_IS_GNUCXX OR CMAKE_CXX_COMPILER_ID MATCHES "Clang")
   add_compile_options(-Wall -Wextra -Wpedantic)
 endif()
 
-set(OpenCV_DIR /usr/local/lib/cmake/opencv4)
+# Set OpenCV
+set(OpenCV_DIR /usr/local/opencv/current/lib/cmake/opencv4)
+# Set ONNXRUNTIME
+set(ONNXRUNTIME_ROOT /usr/local/onnxruntime/current)
 
 # find dependencies
 find_package(ament_cmake REQUIRED)
@@ -37,12 +40,15 @@ add_executable(
 
 target_include_directories(
   detect_node 
+
   PRIVATE
-  $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include/eigen-3.4.1>
-  ${OpenCV_INCLUDE_DIRS}
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include/eigen-3.4.1>
+    ${OpenCV_INCLUDE_DIRS}
+    ${ONNXRUNTIME_ROOT}/include
+
   PUBLIC 
-  $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
-  $<INSTALL_INTERFACE:include>
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/include>
+    $<INSTALL_INTERFACE:include>
   )
 
 ament_target_dependencies(
@@ -58,13 +64,21 @@ ament_target_dependencies(
   OpenCV
 )
 
+target_link_libraries(
+  detect_node
+
+  PRIVATE
+    ${ONNXRUNTIME_ROOT}/lib/libonnxruntime.so
+)
+
 target_compile_features(
   detect_node 
-  PRIVATE c_std_99 cxx_std_17
+
+  PRIVATE 
+    c_std_99 
+    cxx_std_17
   )  
 
-
-
 install(
   TARGETS detect_node 
   DESTINATION lib/${PROJECT_NAME}

vision_test/src/test.cpp

@@ -1,20 +1,83 @@
 #include <opencv2/opencv.hpp>
+#include <onnxruntime_cxx_api.h>
+#include <vector>
+#include <iostream>
+
 
 int main() {
+    // 1、读取图片
     cv::Mat img = cv::imread("/home/lyx/ROS2/hivecore_part_test/src/hivecore_robot_vision/test/color_image.png");
+    if (img.empty()) {
+        std::cerr << "Failed to read image!" << std::endl;
+        return -1;
+    }
 
-    cv::dnn::Net net = cv::dnn::readNetFromONNX("/home/lyx/ROS2/hivecore_part_test/src/hivecore_robot_vision/vision_test/checkpoints/yolo11n-seg.onnx");
+    // 2、预处理：resize + float + normalize
+    cv::resize(img, img, cv::Size(640, 640));
+    img.convertTo(img, CV_32F, 1.0 / 255.0);
+
+    // HWC -> CHW 并存到连续内存
+    std::vector<float> input_tensor_values(1 * 3 * 640 * 640);
+    int idx = 0;
+    for (int c = 0; c < 3; ++c)
+        for (int h = 0; h < 640; ++h)
+            for (int w = 0; w < 640; ++w)
+                input_tensor_values[idx++] = img.at<cv::Vec3f>(h, w)[c];
+
+
+    // 3、初始化 ONNX Runtime
+    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "yolo11n-seg");
+    Ort::SessionOptions session_options;
+    session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
+
+    Ort::Session session(env, 
+        "/home/lyx/ROS2/hivecore_part_test/src/hivecore_robot_vision/vision_test/checkpoints/yolo11n-seg.onnx",
+        session_options
+    );  
     
-    net.setPreferableBackend(cv::dnn::DNN_BACKEND_OPENCV);
-    net.setPreferableTarget(cv::dnn::DNN_TARGET_CPU);
 
-    cv::Mat blob = cv::dnn::blobFromImage(img,
-                                        1/255.0,
-                                        cv::Size(640, 640),
-                                        cv::Scalar(),
-                                        true, false);
-    net.setInput(blob);
-    cv::Mat out = net.forward();
+    // 4、获取输入信息
+    Ort::AllocatorWithDefaultOptions allocator;
+    Ort::AllocatedStringPtr input_name_ptr = session.GetInputNameAllocated(0, allocator);
+    const char* input_name = input_name_ptr.get();
+
+    std::vector<int64_t> input_shape = {1, 3, 640, 640};
+
+    Ort::MemoryInfo mem_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
+
+    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
+        mem_info,
+        input_tensor_values.data(),
+        input_tensor_values.size(),
+        input_shape.data(),
+        input_shape.size()
+    );
+
+    // 5、获取输出名称
+    Ort::AllocatedStringPtr output_name_ptr = session.GetOutputNameAllocated(0, allocator);
+    const char* output_name = output_name_ptr.get();
+
+    // 6、推理
+    auto output_tensors = session.Run(
+        Ort::RunOptions{nullptr},
+        &input_name,
+        &input_tensor,
+        1,
+        &output_name,
+        1
+    );
+
+    // 7、输出信息
+    float* output_data = output_tensors[0].GetTensorMutableData<float>();
+    auto out_shape = output_tensors[0].GetTensorTypeAndShapeInfo().GetShape();
+
+    std::cout << "Output shape: [";
+    for (size_t i = 0; i < out_shape.size(); ++i) {
+        std::cout << out_shape[i];
+        if (i != out_shape.size() - 1) std::cout << ", ";
+    }
+    std::cout << "]" << std::endl;
+
     return 0;
 }