OnnxRT 模块移除参数，改为传参，并更改其关联的模块

cv-rmvl · Jan 26, 2024 · cd911f3 · cd911f3
1 parent bc3ca73
commit cd911f3
Show file tree

Hide file tree

Showing 11 changed files with 124 additions and 108 deletions.
diff --git a/extra/combo/src/armor/armor_cal.cpp b/extra/combo/src/armor/armor_cal.cpp
@@ -167,6 +167,5 @@ Mat Armor::getNumberROI(Mat src, combo::ptr p_combo)
     cvtColor(roi, roi, COLOR_BGR2GRAY);
     resize(roi, roi, Size(armor_param.ROI_SIZE, armor_param.ROI_SIZE));
     threshold(roi, roi, 0, 255, THRESH_OTSU);
-    cvtColor(roi, roi, COLOR_GRAY2BGR);
     return roi;
 }
diff --git a/extra/detector/param/armor_detector.para b/extra/detector/param/armor_detector.para
@@ -1,4 +1,7 @@
 int GRAY_THRESHOLD_RED = 80       # 二值化 R-B 阈值
 int GRAY_THRESHOLD_BLUE = 80      # 二值化 B-R 阈值
 double MIN_CONTOUR_AREA = 25      # 最小轮廓面积
-float MAX_TRACKER_DELTA_DIS = 255 # 识别为相同装甲板序列时，装甲板中心在两帧之间允许的最大距离
+float MAX_TRACKER_DELTA_DIS = 255 # 识别为相同装甲板序列时，装甲板中心在两帧之间允许的最大距离
+
+float MODEL_MEAN = 0.449 # 分类网络归一化后的颜色均值
+float MODEL_STD = 0.226  # 分类网络归一化后的颜色标准差
diff --git a/extra/detector/param/gyro_detector.para b/extra/detector/param/gyro_detector.para
@@ -12,3 +12,6 @@ float MIN_X_DISTANCE_RATIO = 2  # 两个装甲板允许合并至一个集合时
 float MAX_Y_DISTANCE_RATIO = 3  # 两个装甲板允许合并至一个集合时，中心点的垂直距离 dy 与装甲板高度的最大比值
 
 float MIN_CENTER_DIS = 30 # 序列组之间的最小间距，小于该值则需要弃用其中一个 group
+
+float MODEL_MEAN = 0.449 # 分类网络归一化后的颜色均值
+float MODEL_STD = 0.226  # 分类网络归一化后的颜色标准差
diff --git a/extra/detector/src/armor_detector/find.cpp b/extra/detector/src/armor_detector/find.cpp
@@ -40,7 +40,7 @@ void ArmorDetector::find(Mat &src, vector<feature::ptr> &features, vector<combo:
             for (const auto &armor : armors)
             {
                 Mat roi = Armor::getNumberROI(src, armor);
-                auto type = _ort->inference({roi})[0];
+                auto type = _ort->inference({roi}, {armor_detector_param.MODEL_MEAN}, {armor_detector_param.MODEL_STD})[0];
                 armor->setType(_robot_t[type]);
                 rois.emplace_back(roi);
             }

diff --git a/extra/detector/src/gyro_detector/find.cpp b/extra/detector/src/gyro_detector/find.cpp
@@ -42,7 +42,7 @@ void GyroDetector::find(Mat &src, vector<feature::ptr> &features, vector<combo::
             for (const auto &armor : armors)
             {
                 Mat roi = Armor::getNumberROI(src, armor);
-                auto type = _ort->inference({roi})[0];
+                auto type = _ort->inference({roi}, {gyro_detector_param.MODEL_MEAN}, {gyro_detector_param.MODEL_STD})[0];
                 armor->setType(_robot_t[type]);
                 rois.emplace_back(roi);
             }

diff --git a/modules/ml/CMakeLists.txt b/modules/ml/CMakeLists.txt
@@ -10,11 +10,6 @@ rmvl_add_module(
 if(WITH_ONNXRUNTIME)
   find_package(Ort)
 endif()
-
-rmvl_generate_para(
-  ort
-  MODULE ml
-)
 
 set(BUILD_rmvl_ort_INIT ${WITH_ONNXRUNTIME})
 
@@ -25,8 +20,6 @@ rmvl_add_module(
   EXTERNAL ${Ort_LIBS}
 )
 
-rmvl_generate_module_para(ml)
-
 # test
 if(BUILD_TESTS)
 

diff --git a/modules/ml/include/rmvl/ml/ort.h b/modules/ml/include/rmvl/ml/ort.h
@@ -24,81 +24,30 @@ namespace rm
 //! ONNX-Runtime (Ort) 部署库 \cite onnx-rt
 class OnnxRT
 {
-    using session_ptr = std::unique_ptr<Ort::Session>;
-
-    Ort::Env _env;                               //!< 环境配置
-    Ort::SessionOptions _session_options;        //!< Session 配置
-    Ort::MemoryInfo _memory_info;                //!< Tensor 内存分配信息
-    Ort::AllocatorWithDefaultOptions _allocator; //!< 默认配置的内存分配器
-    session_ptr _p_session;
-
-    std::vector<std::vector<float>> _input_arrays; //!< 输入数组
-    std::vector<const char *> _input_names;        //!< 输入名
-    std::vector<const char *> _output_names;       //!< 输出名
-
 public:
     /**
-     * @brief Construct the OnnxRT object
+     * @brief 创建 OnnxRT 对象
      *
      * @param[in] model_path 模型路径，如果该路径不存在，则程序将因错误而退出
      */
-    OnnxRT(const std::string &model_path);
-    ~OnnxRT() = default;
+    OnnxRT(std::string_view model_path);
+    ~OnnxRT();
 
     void printModelInfo();
 
     /**
      * @brief 预处理，推理和后处理
      *
      * @param[in] images 所有的输入图像
+     * @param[in] means 网络模型各通道的均值
+     * @param[in] stds 网络模型各通道的标准差
      * @return 与概率最高的值对应的索引向量
      */
-    std::vector<size_t> inference(const std::vector<cv::Mat> &images);
+    std::vector<size_t> inference(const std::vector<cv::Mat> &images, const std::vector<float> &means, const std::vector<float> &stds);
 
 private:
-    /**
-     * @brief 初始化 Ort 引擎
-     *
-     * @param[in] model_path 模型路径
-     */
-    void setupEngine(const std::string &model_path);
-
-    /**
-     * @brief 分配内存，将图像平展为 NCHW 格式的一维数组，同时将数组归一化
-     * @note 参数 input_array 需要被初始化过，其长度需与 input_image 的大小一致
-     *
-     * @param[in] input_image 输入图像
-     * @param[out] input_array 从输入图像输入数组
-     */
-    void imageToVector(cv::Mat &input_image, std::vector<float> &input_array);
-
-    /**
-     * @brief 预处理
-     *
-     * @param[in] images 所有的输入图像
-     * @return 用于网络输入的 Tensors
-     */
-    std::vector<Ort::Value> preProcess(const std::vector<cv::Mat> &images);
-
-    /**
-     * @brief 后处理
-     *
-     * @param[in] output_tensors 网络输出的 Tensors
-     * @return 具有最高可信度的 index 或值
-     */
-    std::vector<size_t> postProcess(const std::vector<Ort::Value> &output_tensors);
-
-    /**
-     * @brief 推理并返回输出 Tensors
-     *
-     * @param[in] input_tensors 输入 Tensors
-     * @return 输出 Tensors
-     */
-    inline std::vector<Ort::Value> doInference(const std::vector<Ort::Value> &input_tensors)
-    {
-        return _p_session->Run(Ort::RunOptions{nullptr}, _input_names.data(), input_tensors.data(),
-                               input_tensors.size(), _output_names.data(), _output_names.size());
-    }
+    class Impl;
+    Impl *_pimpl;
 };
 
 //! @} ml_ort

diff --git a/modules/ml/param/ort.para b/modules/ml/param/ort.para
diff --git a/modules/ml/src/ort/ort.cpp b/modules/ml/src/ort/ort.cpp
@@ -14,12 +14,27 @@
 
 #include <opencv2/imgproc.hpp>
 
+#include "ort_impl.h"
 #include "rmvl/core/util.hpp"
-#include "rmvl/ml/ort.h"
 
-#include "rmvlpara/ml/ort.h"
+rm::OnnxRT::OnnxRT(std::string_view model_path) : _pimpl(new Impl(model_path)) {}
+rm::OnnxRT::~OnnxRT() { delete _pimpl; }
+void rm::OnnxRT::printModelInfo() { _pimpl->printModelInfo(); }
+std::vector<size_t> rm::OnnxRT::inference(const std::vector<cv::Mat> &images, const std::vector<float> &means,
+                                          const std::vector<float> &stds)
+{
+    try
+    {
+        return _pimpl->inference(images, means, stds);
+    }
+    catch (const rm::Exception &e)
+    {
+        delete _pimpl;
+        throw e;
+    }
+}
 
-rm::OnnxRT::OnnxRT(const std::string &model_path)
+rm::OnnxRT::Impl::Impl(std::string_view model_path)
     : _env(OrtLoggingLevel::ORT_LOGGING_LEVEL_WARNING, "OnnxDeployment"),
       _memory_info(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault))
 {
@@ -28,7 +43,7 @@ rm::OnnxRT::OnnxRT(const std::string &model_path)
     setupEngine(model_path);
 }
 
-void rm::OnnxRT::setupEngine(const std::string &model_path)
+void rm::OnnxRT::Impl::setupEngine(std::string_view model_path) noexcept
 {
 #ifdef WITH_ORT_CUDA
     OrtSessionOptionsAppendExecutionProvider_CUDA(session_options, 0);
@@ -39,7 +54,7 @@ void rm::OnnxRT::setupEngine(const std::string &model_path)
 #endif // WITH_ORT_TensorRT
 
     _session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
-    _p_session = std::make_unique<Ort::Session>(_env, model_path.c_str(), _session_options);
+    _p_session = std::make_unique<Ort::Session>(_env, model_path.data(), _session_options);
 
     // define the names of the I/O nodes
     for (size_t i = 0; i < _p_session->GetInputCount(); i++)
@@ -55,18 +70,22 @@ void rm::OnnxRT::setupEngine(const std::string &model_path)
     }
 }
 
-std::vector<size_t> rm::OnnxRT::inference(const std::vector<cv::Mat> &images)
+std::vector<size_t> rm::OnnxRT::Impl::inference(const std::vector<cv::Mat> &images,
+                                                const std::vector<float> &means,
+                                                const std::vector<float> &stds)
 {
-    std::vector<Ort::Value> input_tensors = preProcess(images);
+    std::vector<Ort::Value> input_tensors = preProcess(images, means, stds);
     std::vector<Ort::Value> output_tensors = doInference(input_tensors);
     return postProcess(output_tensors);
 }
 
-std::vector<Ort::Value> rm::OnnxRT::preProcess(const std::vector<cv::Mat> &images)
+std::vector<Ort::Value> rm::OnnxRT::Impl::preProcess(const std::vector<cv::Mat> &images,
+                                                     const std::vector<float> &means,
+                                                     const std::vector<float> &stds)
 {
     size_t input_count = _p_session->GetInputCount();
     if (input_count != images.size())
-        CV_Error(RMVL_StsBadArg, "Size of the \"images\" are not equal to the model input_count.");
+        RMVL_Error(RMVL_StsBadArg, "Size of the \"images\" are not equal to the model input_count.");
     // get the correct data of each input layer
     std::vector<Ort::Value> input_tensors;
     for (size_t i = 0; i < input_count; i++)
@@ -82,15 +101,15 @@ std::vector<Ort::Value> rm::OnnxRT::preProcess(const std::vector<cv::Mat> &image
         cv::Mat input_image;
         resize(images[i], input_image, cv::Size(input_shape[2], input_shape[3]));
         // allocate memory and normalization
-        imageToVector(input_image, _input_arrays[i]);
+        imageToVector(input_image, _input_arrays[i], means, stds);
         input_tensors.emplace_back(Ort::Value::CreateTensor<float>(_memory_info,
                                                                    _input_arrays[i].data(), _input_arrays[i].size(),
                                                                    input_shape.data(), input_shape.size()));
     }
     return input_tensors;
 }
 
-std::vector<size_t> rm::OnnxRT::postProcess(const std::vector<Ort::Value> &output_tensors)
+std::vector<size_t> rm::OnnxRT::Impl::postProcess(const std::vector<Ort::Value> &output_tensors) noexcept
 {
     // 所有输出对应的置信度最高的索引
     std::vector<size_t> output_indexs;
@@ -100,16 +119,16 @@ std::vector<size_t> rm::OnnxRT::postProcess(const std::vector<Ort::Value> &outpu
         const float *output = output_tensor.GetTensorData<float>();
         std::vector<size_t> indexs(output_tensor.GetTensorTypeAndShapeInfo().GetElementCount());
         iota(indexs.begin(), indexs.end(), 0);
-        auto it = max_element(indexs.begin(), indexs.end(),
-                              [&output](size_t lhs, size_t rhs) {
-                                  return output[lhs] < output[rhs];
-                              });
+        auto it = max_element(indexs.begin(), indexs.end(), [&output](size_t lhs, size_t rhs) {
+            return output[lhs] < output[rhs];
+        });
         output_indexs.emplace_back(*it);
     }
     return output_indexs;
 }
 
-void rm::OnnxRT::imageToVector(cv::Mat &input_image, std::vector<float> &input_array)
+void rm::OnnxRT::Impl::imageToVector(cv::Mat &input_image, std::vector<float> &input_array,
+                                     const std::vector<float> &means, const std::vector<float> &stds)
 {
     // CHW
     int C = input_image.channels();
@@ -120,22 +139,8 @@ void rm::OnnxRT::imageToVector(cv::Mat &input_image, std::vector<float> &input_a
     size_t pixels = C * H * W;
     if (pixels != input_array.size())
         RMVL_Error(RMVL_StsBadArg, "The size of the arguments: \"input_image\" and \"input_array\" are not equal");
-
-    std::vector<double> means(C);
-    std::vector<double> stds(C);
-    if (C == 1)
-    {
-        means[0] = para::ort_param.MONO_MEANS;
-        stds[0] = para::ort_param.MONO_STDS;
-    }
-    else
-    {
-        for (int i = 0; i < C; ++i)
-        {
-            means[i] = para::ort_param.RGB_MEANS[i];
-            stds[i] = para::ort_param.RGB_STDS[i];
-        }
-    }
+    if (static_cast<int>(means.size()) != C || static_cast<int>(stds.size()) != C)
+        RMVL_Error_(RMVL_StsBadArg, "Bad size of the input argument: the size of \"means\" or \"stds\" must be %d", C);
     // 转 Tensor 的 NCHW 格式，做归一化和标准化
     float *p_input_array = input_array.data();
     for (int c = 0; c < C; c++)

diff --git a/modules/ml/src/ort/ort_impl.h b/modules/ml/src/ort/ort_impl.h
@@ -0,0 +1,68 @@
+/**
+ * @file ort_impl.h
+ * @author zhaoxi ([email protected])
+ * @brief the deployment library of the ONNXruntime (Ort)
+ * @version 1.0
+ * @date 2024-01-25
+ *
+ * @copyright Copyright 2023 (c), zhaoxi
+ *
+ */
+
+#pragma once
+
+#include <opencv2/core/mat.hpp>
+
+#include <onnxruntime_cxx_api.h>
+
+#include "rmvl/ml/ort.h"
+
+namespace rm
+{
+
+//! ONNX-Runtime (Ort) 部署库 \cite onnx-rt
+class OnnxRT::Impl
+{
+    using session_ptr = std::unique_ptr<Ort::Session>;
+
+    Ort::Env _env;                               //!< 环境配置
+    Ort::SessionOptions _session_options;        //!< Session 配置
+    Ort::MemoryInfo _memory_info;                //!< Tensor 内存分配信息
+    Ort::AllocatorWithDefaultOptions _allocator; //!< 默认配置的内存分配器
+    session_ptr _p_session;
+
+    std::vector<std::vector<float>> _input_arrays; //!< 输入数组
+    std::vector<const char *> _input_names;        //!< 输入名
+    std::vector<const char *> _output_names;       //!< 输出名
+
+public:
+    Impl(std::string_view model_path);
+    ~Impl() = default;
+
+    void printModelInfo() noexcept;
+
+    std::vector<size_t> inference(const std::vector<cv::Mat> &images, const std::vector<float> &means, const std::vector<float> &stds);
+
+private:
+    //! 初始化 Ort 引擎
+    void setupEngine(std::string_view model_path) noexcept;
+
+    //! 分配内存，将图像平展为 NCHW 格式的一维数组，同时将数组归一化
+    void imageToVector(cv::Mat &input_image, std::vector<float> &input_array,
+                       const std::vector<float> &means, const std::vector<float> &stds);
+
+    //! 预处理
+    std::vector<Ort::Value> preProcess(const std::vector<cv::Mat> &images, const std::vector<float> &means, const std::vector<float> &stds);
+
+    //! 后处理
+    std::vector<size_t> postProcess(const std::vector<Ort::Value> &output_tensors) noexcept;
+
+    //! 推理并返回输出 Tensors
+    inline std::vector<Ort::Value> doInference(const std::vector<Ort::Value> &input_tensors) noexcept
+    {
+        return _p_session->Run(Ort::RunOptions{nullptr}, _input_names.data(), input_tensors.data(),
+                               input_tensors.size(), _output_names.data(), _output_names.size());
+    }
+};
+
+} // namespace rm
diff --git a/modules/ml/src/ort/ort_print.cpp b/modules/ml/src/ort/ort_print.cpp
@@ -9,12 +9,12 @@
  *
  */
 
-#include "rmvl/ml/ort.h"
+#include "ort_impl.h"
 
 using namespace std;
 using namespace Ort;
 
-void rm::OnnxRT::printModelInfo()
+void rm::OnnxRT::Impl::printModelInfo() noexcept
 {
     printf("-------------- Input Layer --------------\n");
     int input_node = _p_session->GetInputCount();