diff --git a/extra/combo/src/armor/armor_cal.cpp b/extra/combo/src/armor/armor_cal.cpp index e5397400..bfde25b2 100755 --- 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 index 7714c2e5..d2a00fc3 100644 --- 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 # 识别为相同装甲板序列时,装甲板中心在两帧之间允许的最大距离 \ No newline at end of file +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 index 5950b0b0..8040612f 100644 --- 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 index a7fcbd88..fbd33b0d 100755 --- 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 &features, vectorinference({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 index acd8a679..fc70f66e 100755 --- 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 &features, vectorinference({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 index 3b86e635..52705e97 100644 --- 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 index 819b69a9..90222b25 100755 --- a/modules/ml/include/rmvl/ml/ort.h +++ b/modules/ml/include/rmvl/ml/ort.h @@ -24,26 +24,14 @@ namespace rm //! ONNX-Runtime (Ort) 部署库 \cite onnx-rt class OnnxRT { - using session_ptr = std::unique_ptr; - - Ort::Env _env; //!< 环境配置 - Ort::SessionOptions _session_options; //!< Session 配置 - Ort::MemoryInfo _memory_info; //!< Tensor 内存分配信息 - Ort::AllocatorWithDefaultOptions _allocator; //!< 默认配置的内存分配器 - session_ptr _p_session; - - std::vector> _input_arrays; //!< 输入数组 - std::vector _input_names; //!< 输入名 - std::vector _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(); @@ -51,54 +39,15 @@ class OnnxRT * @brief 预处理,推理和后处理 * * @param[in] images 所有的输入图像 + * @param[in] means 网络模型各通道的均值 + * @param[in] stds 网络模型各通道的标准差 * @return 与概率最高的值对应的索引向量 */ - std::vector inference(const std::vector &images); + std::vector inference(const std::vector &images, const std::vector &means, const std::vector &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 &input_array); - - /** - * @brief 预处理 - * - * @param[in] images 所有的输入图像 - * @return 用于网络输入的 Tensors - */ - std::vector preProcess(const std::vector &images); - - /** - * @brief 后处理 - * - * @param[in] output_tensors 网络输出的 Tensors - * @return 具有最高可信度的 index 或值 - */ - std::vector postProcess(const std::vector &output_tensors); - - /** - * @brief 推理并返回输出 Tensors - * - * @param[in] input_tensors 输入 Tensors - * @return 输出 Tensors - */ - inline std::vector doInference(const std::vector &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 deleted file mode 100644 index e836a0c9..00000000 --- a/modules/ml/param/ort.para +++ /dev/null @@ -1,4 +0,0 @@ -Vec3d RGB_MEANS = {0.485, 0.456, 0.406} # 归一化后的 RGB 颜色均值,顺序为 (R, G, B) -Vec3d RGB_STDS = {0.229, 0.224, 0.225} # 归一化后的 RGB 颜色标准差,顺序为 (R, G, B) -double MONO_MEANS = 0.449 # 归一化后的 单通道颜色均值 -double MONO_STDS = 0.226 # 归一化后的 单通道颜色标准差 \ No newline at end of file diff --git a/modules/ml/src/ort/ort.cpp b/modules/ml/src/ort/ort.cpp index 8a46823e..de04020a 100755 --- a/modules/ml/src/ort/ort.cpp +++ b/modules/ml/src/ort/ort.cpp @@ -14,12 +14,27 @@ #include +#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 rm::OnnxRT::inference(const std::vector &images, const std::vector &means, + const std::vector &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(_env, model_path.c_str(), _session_options); + _p_session = std::make_unique(_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 rm::OnnxRT::inference(const std::vector &images) +std::vector rm::OnnxRT::Impl::inference(const std::vector &images, + const std::vector &means, + const std::vector &stds) { - std::vector input_tensors = preProcess(images); + std::vector input_tensors = preProcess(images, means, stds); std::vector output_tensors = doInference(input_tensors); return postProcess(output_tensors); } -std::vector rm::OnnxRT::preProcess(const std::vector &images) +std::vector rm::OnnxRT::Impl::preProcess(const std::vector &images, + const std::vector &means, + const std::vector &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 input_tensors; for (size_t i = 0; i < input_count; i++) @@ -82,7 +101,7 @@ std::vector rm::OnnxRT::preProcess(const std::vector &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(_memory_info, _input_arrays[i].data(), _input_arrays[i].size(), input_shape.data(), input_shape.size())); @@ -90,7 +109,7 @@ std::vector rm::OnnxRT::preProcess(const std::vector &image return input_tensors; } -std::vector rm::OnnxRT::postProcess(const std::vector &output_tensors) +std::vector rm::OnnxRT::Impl::postProcess(const std::vector &output_tensors) noexcept { // 所有输出对应的置信度最高的索引 std::vector output_indexs; @@ -100,16 +119,16 @@ std::vector rm::OnnxRT::postProcess(const std::vector &outpu const float *output = output_tensor.GetTensorData(); std::vector 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 &input_array) +void rm::OnnxRT::Impl::imageToVector(cv::Mat &input_image, std::vector &input_array, + const std::vector &means, const std::vector &stds) { // CHW int C = input_image.channels(); @@ -120,22 +139,8 @@ void rm::OnnxRT::imageToVector(cv::Mat &input_image, std::vector &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 means(C); - std::vector 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(means.size()) != C || static_cast(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 new file mode 100644 index 00000000..3ef4fc84 --- /dev/null +++ b/modules/ml/src/ort/ort_impl.h @@ -0,0 +1,68 @@ +/** + * @file ort_impl.h + * @author zhaoxi (535394140@qq.com) + * @brief the deployment library of the ONNXruntime (Ort) + * @version 1.0 + * @date 2024-01-25 + * + * @copyright Copyright 2023 (c), zhaoxi + * + */ + +#pragma once + +#include + +#include + +#include "rmvl/ml/ort.h" + +namespace rm +{ + +//! ONNX-Runtime (Ort) 部署库 \cite onnx-rt +class OnnxRT::Impl +{ + using session_ptr = std::unique_ptr; + + Ort::Env _env; //!< 环境配置 + Ort::SessionOptions _session_options; //!< Session 配置 + Ort::MemoryInfo _memory_info; //!< Tensor 内存分配信息 + Ort::AllocatorWithDefaultOptions _allocator; //!< 默认配置的内存分配器 + session_ptr _p_session; + + std::vector> _input_arrays; //!< 输入数组 + std::vector _input_names; //!< 输入名 + std::vector _output_names; //!< 输出名 + +public: + Impl(std::string_view model_path); + ~Impl() = default; + + void printModelInfo() noexcept; + + std::vector inference(const std::vector &images, const std::vector &means, const std::vector &stds); + +private: + //! 初始化 Ort 引擎 + void setupEngine(std::string_view model_path) noexcept; + + //! 分配内存,将图像平展为 NCHW 格式的一维数组,同时将数组归一化 + void imageToVector(cv::Mat &input_image, std::vector &input_array, + const std::vector &means, const std::vector &stds); + + //! 预处理 + std::vector preProcess(const std::vector &images, const std::vector &means, const std::vector &stds); + + //! 后处理 + std::vector postProcess(const std::vector &output_tensors) noexcept; + + //! 推理并返回输出 Tensors + inline std::vector doInference(const std::vector &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 index 09ab294d..2587aaa4 100755 --- 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();