如何在lite.ai.toolkit(1.5k+ stars)中增加您的模型?

目录

0. 近期一些感想

1. 这篇文章讲什么

2. Lite.Ai.ToolKit代码结构介绍

3. 添加模型的步骤

4. 总结


0. 近期一些感想

业余写了个C++ AI模型工具箱,叫 lite.ai.toolkit;写这个repo的最开始的出发点,并不是有多远大的愿景或规划,只是为了更好地理解我所接触到的算法,更好地理解它们的细节。因为在我看来,有些东西,只有放在C++的层面来看,你才能真正接触到它最真实的细节。比如目标检测,后处理涉及anchor的生成和坐标反算,给定我一个20x20x?的输出,如果我能通过C++来准确地计算出anchor,并将预测的offsets在内存中索引位置和计算anchor索引对应上,准确无误地反算出最后的 bounding box,那么,我就认为自己算是理解了这段逻辑的细节了。所以我看一篇paper,了解一个新的算法,通常会非常倾向于用C++把它的推理过一遍,lite.ai.toolkit 就是一个我用来保存这些代码的地方。只是意料之外的是,后来捏的模型慢慢变多了,开始有一些些人关注,我也在想要不把这些代码都捏成一个库吧,所以就有了现在的 lite.ai.toolkit,目前为止收获了 1.5k+ stars,但说实在的,个人真的是诚惶诚恐,因为我很清楚,这事情本身的难度并不大, lite.ai.toolkit 的代码质量也一般,只是我恰好这么做了。所以还是很感谢社区小伙伴们的支持~ lite.ai.toolkit 一直还没达到我心里面认可的1.0版本,比如说它至少要有100+模型,要对关键的前后处理有一些针对性的性能优化,要充分可拆解可二次开发(比如想要在mediapipe的基础上做二次开发,那可太难了),要有基本的不同平台的预编译库等等。不过,只要我还继续在做算法或工程,这个小工具应该就会继续维护下去,来日方长,咱们也且行且看 ~

1. 这篇文章讲什么

那么这篇文章讲些什么呢?有一些小伙伴在用 lite.ai.toolkit,最近慢慢地开始有越来越多的小伙伴提出一些issue说,想在lite.ai.toolkit中增加一些新模型,这其实就是如何对 lite.ai.toolkit 做二次开发的问题。最近的 issues 基本是长这样的。

但无奈的是,本人的业余时间实在是有限的,很难及时地去满足这些增加新模型的需求。这篇文章主要讲一下,如何把你自己的模型添加到lite.ai.toolkit。lite.ai.toolkit集成了一些比较新的基础模型,比如人脸检测、人脸识别、抠图、人脸属性分析、图像分类、人脸关键点识别、图像着色、目标检测等等,可以直接用到具体的场景中。但是,毕竟lite.ai.toolkit的模型还是有限的,具体的场景下,可能有你经过优化的模型,比如你自己训了一个目标检测器,可能效果更好。那么,如何把你的模型加入到lite.ai.toolkit中呢?这样既能用到lite.ai.toolkit一些已有的算法能力,也能兼容您的具体场景。我也在lite.ai.toolkit的repo中,同步开了一个issue,这个issue主要是讲这个问题。大家有疑惑的可以提在这个issue,我会尽可能回答~

👉CONTRIBUTING (如何添加你的模型?) · Issue #191 · DefTruth/lite.ai.toolkit · GitHub这个issue主要讲一下,如何把你自己的模型添加到lite.ai.toolkit。lite.ai.toolkit集成了一些比较新的基础模型,比如人脸检测、人脸识别、抠图、人脸属性分析、图像分类、人脸关键点识别、图像着色、目标检测等等,可以直接用到具体的场景中。但是,毕竟lite.ai.toolkit的模型还是有限的,具体的场景下,可能有你经过优化的模型,比如你自己训了一个目标检测器,可能效果更好。那么,如何把你的模型加入到lite.ai.toolkit中呢?这样既能用到...https://github.com/DefTruth/lite.ai.toolkit/issues/191

2. Lite.Ai.ToolKit代码结构介绍

lite.ai.toolkit采用了尽可能解耦的方式来管理代码,基于不同的推理引擎的实现是相互独立的:

  • 代码结构思路:ONNXRuntime版本的YOLOX和MNN版本的YOLOX是相互独立的,他们的代码分别管理在不同的目录下,你可以只编译ONNXRuntime版本的实现。

  • 坏处和好处:虽然说这样做的坏处是难以避免出现重复的代码,但在AI原子能力层面做到了完全的独立,比如,如果你只想用onnxruntime版本的YOLOX,那么你完全可以把其他的模型cpp和h文件全部删掉,只保留yolox.cpp和yolox.h。另一个好处就是,可以在编译期决定你想要用什么模型,不想要用的模型完全不需要进行编译,这样编出来的库,其实也不会很大,这也是我后面准备添加的一个功能,这是后话,先不展开了。

回到二次开发这个话题,为方便lite.ai.toolkit的用户添加自己的模型,这里简单介绍下lite.ai.toolkit的代码布局。

  • lite文件夹

该文件夹的根目录下包含了所有的主要代码,其结构如下:

# --------------------------- 这部分是管理整理工程的代码 被下游的各个模块依赖 ----------------------------------
├── backend.h          # 宏处理 决定基础的推理引擎 目前必须是ONNXRuntime
├── config.h           # 宏处理
├── config.h.in        # cmake编译时宏处理
├── lite.ai.defs.h     # 宏处理
├── lite.ai.headers.h  # 引入基础依赖库
├── lite.h             # 引入项目的基础对外的模块
├── pipeline           # 暂时没用
├── pipeline.h         # 暂时没用
├── types.h            # 基础类型,很重要,会被下游模块复用,比如实际上ort::types是types的alias
├── utils.cpp          # 基础功能函数实现,很重要,会被下游模块复用
└── utils.h            # 基础功能函数头文件
├── models.h           # 模型整体的命名空间管理,很重要,所有被实现的模型需要在这里被导出

# --------------------------- 以下的各个部分是相互独立的 但是需要依赖上面这个整体的部分 --------------------------
...
├── mnn
│   ├── core     #  MNN基础父类和特定功能的实现,必须要阅读
│   │   ├── mnn_config.h
│   │   ├── mnn_core.h       # MNN模型命名空间管理,实现一个类前,现在这里添加签名
│   │   ├── mnn_defs.h
│   │   ├── mnn_handler.cpp  # 基础父类实现,必须阅读
│   │   ├── mnn_handler.h
│   │   ├── mnn_types.h
│   │   ├── mnn_utils.cpp
│   │   └── mnn_utils.h
│   └── cv      #  各个模型的具体实现,会引用core中实现的父类和功能函数
│       ├── mnn_age_googlenet.cpp
│       ├── mnn_age_googlenet.h
│       ├── mnn_cava_combined_face.cpp
...
├── ncnn
│   ├── core     #  NCNN基础父类和特定功能的实现,必须要阅读
│   │   ├── ncnn_config.h
│   │   ├── ncnn_core.h        # NCNN模型命名空间管理,实现一个类前,现在这里添加签名
│   │   ├── ncnn_custom.cpp
│   │   ├── ncnn_custom.h
│   │   ├── ncnn_defs.h
│   │   ├── ncnn_handler.cpp   # 基础父类实现,必须阅读
│   │   ├── ncnn_handler.h
│   │   ├── ncnn_types.h
│   │   ├── ncnn_utils.cpp
│   │   └── ncnn_utils.h
│   └── cv      #  各个模型的具体实现,会引用core中实现的父类和功能函数
│       ├── ncnn_age_googlenet.cpp
│       ├── ncnn_age_googlenet.h
│       ├── ncnn_cava_combined_face.cpp
│       ├── ncnn_cava_combined_face.h
│       ├── ncnn_cava_ghost_arcface.cpp
...
├── ort
│   ├── core     #  ONNXRuntime基础父类和特定功能的实现,必须要阅读
│   │   ├── ort_config.h
│   │   ├── ort_core.h        # ONNXRuntime模型命名空间管理,实现一个类前,现在这里添加签名
│   │   ├── ort_defs.h
│   │   ├── ort_handler.cpp   # 基础父类实现,必须阅读
│   │   ├── ort_handler.h
│   │   ├── ort_types.h
│   │   ├── ort_utils.cpp
│   │   └── ort_utils.h
│   └── cv      #  各个模型的具体实现,会引用core中实现的父类和功能函数
│       ├── age_googlenet.cpp
│       ├── age_googlenet.h
│       ├── cava_combined_face.cpp
│       ├── cava_combined_face.h
│       ├── cava_ghost_arcface.cpp
│       ├── cava_ghost_arcface.h
...
├── tnn
│   ├── core     #  TNN基础父类和特定功能的实现,必须要阅读
│   │   ├── tnn_config.h
│   │   ├── tnn_core.h        # TNN模型命名空间管理,实现一个类前,现在这里添加签名
│   │   ├── tnn_defs.h
│   │   ├── tnn_handler.cpp   # 基础父类实现,必须阅读
│   │   ├── tnn_handler.h
│   │   ├── tnn_types.h
│   │   ├── tnn_utils.cpp
│   │   └── tnn_utils.h
│   └── cv    #  各个模型的具体实现,会引用core中实现的父类和功能函数
│       ├── tnn_age_googlenet.cpp
│       ├── tnn_age_googlenet.h
│       ├── tnn_cava_combined_face.cpp
│       ├── tnn_cava_combined_face.h
│       ├── tnn_cava_ghost_arcface.cpp
│       ├── tnn_cava_ghost_arcface.h
│       ├── tnn_center_loss_face.cpp

3. 添加模型的步骤

以下以添加YOLOX的ONNXRuntime C++版本为例,讲解如何添加一个新模型。

  • 第一步: 添加YOLOX类签名。在 lite/ort/core/ort_core.h 中添加YoloX签名,如果是其他推理引擎则还需要加具体的推理引擎作为前缀,如MNNYoloX.

// lite/ort/core/ort_core.h 中
namespace ortcv
{
  // ... 
  class LITE_EXPORTS YoloX;                      // [56] * reference: https://github.com/Megvii-BaseDetection/YOLOX
}
// lite/mnn/core/mnn_core.h 中
namespace mnncv
{
// ... 
class LITE_EXPORTS MNNYoloX;                      // [3] * reference: https://github.com/Megvii-BaseDetection/YOLOX
}
  • 第二步: 在 lite/ort/cv 中新建 yolox.h 和 yolox.cpp,注意最好文件名和你在xxx_core.h中的保持一致,方便管理,对于非ONNXRuntime的版本,还应该加上推理引擎作为前缀,如 mnn_yolox.h 和 mnn_yolox.cpp.

├── ort
│   ├── core     #  ONNXRuntime基础父类和特定功能的实现,必须要阅读
│   │   ├── ort_config.h
│...
│   │   └── ort_utils.h
│   └── cv      #  各个模型的具体实现,会引用core中实现的父类和功能函数
│       ├── yolox.cpp
│       ├── yolox.h
  • 第三步: 编写 YoloX 类的头文件,由于是静态维度推理,并且是单输入多(单)输出模型,所以可以继承 BasicOrtHandler(请自行阅读lite/ort/core/ort_handler.cpp具体实现),注意BasicOrtHandler有个transform的虚函数是要重写的。另外对于最终public的接口,也请保持detect的命名规范以确保API语义的统一性。请务必保持detect、detect_video等的命名规范,detect是图片级别的检测接口,detect_video是视频级别的检测接口。对于types命名空间,实际上对是全局types的引用,因此请在 lite/types.h中查看是否有合适的类型,如果没有则需要在 lite::types 命名空间中添加,然后再在 yolox.h 中使用,请尽可能保持自定义类型的简洁性。

#ifndef LITE_AI_ORT_CV_YOLOX_H
#define LITE_AI_ORT_CV_YOLOX_H

#include "lite/ort/core/ort_core.h"

namespace ortcv
{
  class LITE_EXPORTS YoloX : public BasicOrtHandler
  {
  public:
    explicit YoloX(const std::string &_onnx_path, unsigned int _num_threads = 1) :
        BasicOrtHandler(_onnx_path, _num_threads)
    {};

    ~YoloX() override = default;

  private:
    // nested classes
    typedef struct GridAndStride
    {
      int grid0;
      int grid1;
      int stride;
    } YoloXAnchor;

    typedef struct
    {
      float r;
      int dw;
      int dh;
      int new_unpad_w;
      int new_unpad_h;
      bool flag;
    } YoloXScaleParams;

  private:
    const float mean_vals[3] = {255.f * 0.485f, 255.f * 0.456, 255.f * 0.406f};
    const float scale_vals[3] = {1 / (255.f * 0.229f), 1 / (255.f * 0.224f), 1 / (255.f * 0.225f)};

    const char *class_names[80] = {
        "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
        "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
        "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
        "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
        "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
        "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
        "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard",
        "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
        "scissors", "teddy bear", "hair drier", "toothbrush"
    };
    enum NMS
    {
      HARD = 0, BLEND = 1, OFFSET = 2
    };
    static constexpr const unsigned int max_nms = 30000;

  private:
    // 需要被重写的方法
    Ort::Value transform(const cv::Mat &mat_rs) override; // without resize

    void resize_unscale(const cv::Mat &mat,
                        cv::Mat &mat_rs,
                        int target_height,
                        int target_width,
                        YoloXScaleParams &scale_params);

    void generate_anchors(const int target_height,
                          const int target_width,
                          std::vector<int> &strides,
                          std::vector<YoloXAnchor> &anchors);

    void generate_bboxes(const YoloXScaleParams &scale_params,
                         std::vector<types::Boxf> &bbox_collection,
                         std::vector<Ort::Value> &output_tensors,
                         float score_threshold, int img_height,
                         int img_width); // rescale & exclude

    void nms(std::vector<types::Boxf> &input, std::vector<types::Boxf> &output,
             float iou_threshold, unsigned int topk, unsigned int nms_type);

  public:
    // 请保持detect、detect_video等的命名规范,detect是图片级别的检测接口,detect_video是视频级别的检测接口
    void detect(const cv::Mat &mat, std::vector<types::Boxf> &detected_boxes,
                float score_threshold = 0.25f, float iou_threshold = 0.45f,
                unsigned int topk = 100, unsigned int nms_type = NMS::OFFSET);

  };
}
#endif //LITE_AI_ORT_CV_YOLOX_H
  • 第四步: 在 yolox.cpp 中实现 YoloX 类的所有方法,并可能会引用全局的 lite/utils.h 进行复用,这基本是唯一的全局依赖了。

#include "yolox.h"
#include "lite/ort/core/ort_utils.h" // 引入onnxruntime特定的自定义功能函数,依赖于推理引擎
#include "lite/utils.h" // 引入全局定义的功能函数,不依赖推理引擎,如NMS

using ortcv::YoloX;

Ort::Value YoloX::transform(const cv::Mat &mat_rs)
{
  cv::Mat canvas;
  cv::cvtColor(mat_rs, canvas, cv::COLOR_BGR2RGB);
  // resize without padding, (Done): add padding as the official Python implementation.
  // cv::resize(canva, canva, cv::Size(input_node_dims.at(3),
  //                                  input_node_dims.at(2)));
  // (1,3,640,640) 1xCXHXW
  ortcv::utils::transform::normalize_inplace(canvas, mean_vals, scale_vals); // float32
  // Note !!!: Comment out this line if you use the newest YOLOX model.
  // There is no normalization for the newest official C++ implementation
  // using ncnn. Reference:
  // [1] https://github.com/Megvii-BaseDetection/YOLOX/blob/main/demo/ncnn/cpp/yolox.cpp
  // ortcv::utils::transform::normalize_inplace(canva, mean_vals, scale_vals); // float32
  return ortcv::utils::transform::create_tensor(
      canvas, input_node_dims, memory_info_handler,
      input_values_handler, ortcv::utils::transform::CHW);
}

void YoloX::resize_unscale(const cv::Mat &mat, cv::Mat &mat_rs,
                           int target_height, int target_width,
                           YoloXScaleParams &scale_params)
{
  if (mat.empty()) return;
  int img_height = static_cast<int>(mat.rows);
  int img_width = static_cast<int>(mat.cols);

  mat_rs = cv::Mat(target_height, target_width, CV_8UC3,
                   cv::Scalar(114, 114, 114));
  // scale ratio (new / old) new_shape(h,w)
  float w_r = (float) target_width / (float) img_width;
  float h_r = (float) target_height / (float) img_height;
  float r = std::min(w_r, h_r);
  // compute padding
  int new_unpad_w = static_cast<int>((float) img_width * r); // floor
  int new_unpad_h = static_cast<int>((float) img_height * r); // floor
  int pad_w = target_width - new_unpad_w; // >=0
  int pad_h = target_height - new_unpad_h; // >=0

  int dw = pad_w / 2;
  int dh = pad_h / 2;

  // resize with unscaling
  cv::Mat new_unpad_mat = mat.clone();
  cv::resize(new_unpad_mat, new_unpad_mat, cv::Size(new_unpad_w, new_unpad_h));
  new_unpad_mat.copyTo(mat_rs(cv::Rect(dw, dh, new_unpad_w, new_unpad_h)));

  // record scale params.
  scale_params.r = r;
  scale_params.dw = dw;
  scale_params.dh = dh;
  scale_params.new_unpad_w = new_unpad_w;
  scale_params.new_unpad_h = new_unpad_h;
  scale_params.flag = true;
}

void YoloX::detect(const cv::Mat &mat, std::vector<types::Boxf> &detected_boxes,
                   float score_threshold, float iou_threshold,
                   unsigned int topk, unsigned int nms_type)
{
  if (mat.empty()) return;
  const int input_height = input_node_dims.at(2);
  const int input_width = input_node_dims.at(3);
  int img_height = static_cast<int>(mat.rows);
  int img_width = static_cast<int>(mat.cols);

  // resize & unscale
  cv::Mat mat_rs;
  YoloXScaleParams scale_params;
  this->resize_unscale(mat, mat_rs, input_height, input_width, scale_params);

  // 1. make input tensor
  Ort::Value input_tensor = this->transform(mat_rs);
  // 2. inference scores & boxes.
  auto output_tensors = ort_session->Run(
      Ort::RunOptions{nullptr}, input_node_names.data(),
      &input_tensor, 1, output_node_names.data(), num_outputs
  );
  // 3. rescale & exclude.
  std::vector<types::Boxf> bbox_collection;
  this->generate_bboxes(scale_params, bbox_collection, output_tensors, score_threshold, img_height, img_width);
  // 4. hard|blend|offset nms with topk.
  this->nms(bbox_collection, detected_boxes, iou_threshold, topk, nms_type);
}

void YoloX::generate_anchors(const int target_height,
                             const int target_width,
                             std::vector<int> &strides,
                             std::vector<YoloXAnchor> &anchors)
{
  for (auto stride : strides)
  {
    int num_grid_w = target_width / stride;
    int num_grid_h = target_height / stride;
    for (int g1 = 0; g1 < num_grid_h; ++g1)
    {
      for (int g0 = 0; g0 < num_grid_w; ++g0)
      {
#ifdef LITE_WIN32
        YoloXAnchor anchor;
        anchor.grid0 = g0;
        anchor.grid1 = g1;
        anchor.stride = stride;
        anchors.push_back(anchor);
#else
        anchors.push_back((YoloXAnchor) {g0, g1, stride});
#endif
      }
    }
  }
}

void YoloX::generate_bboxes(const YoloXScaleParams &scale_params,
                            std::vector<types::Boxf> &bbox_collection,
                            std::vector<Ort::Value> &output_tensors,
                            float score_threshold, int img_height,
                            int img_width)
{
  Ort::Value &pred = output_tensors.at(0); // (1,n,85=5+80=cxcy+cwch+obj_conf+cls_conf)
  auto pred_dims = output_node_dims.at(0); // (1,n,85)
  const unsigned int num_anchors = pred_dims.at(1); // n = ?
  const unsigned int num_classes = pred_dims.at(2) - 5;
  const float input_height = static_cast<float>(input_node_dims.at(2)); // e.g 640
  const float input_width = static_cast<float>(input_node_dims.at(3)); // e.g 640

  std::vector<YoloXAnchor> anchors;
  std::vector<int> strides = {8, 16, 32}; // might have stride=64
  this->generate_anchors(input_height, input_width, strides, anchors);

  float r_ = scale_params.r;
  int dw_ = scale_params.dw;
  int dh_ = scale_params.dh;

  bbox_collection.clear();
  unsigned int count = 0;
  for (unsigned int i = 0; i < num_anchors; ++i)
  {
    float obj_conf = pred.At<float>({0, i, 4});
    if (obj_conf < score_threshold) continue; // filter first.

    float cls_conf = pred.At<float>({0, i, 5});
    unsigned int label = 0;
    for (unsigned int j = 0; j < num_classes; ++j)
    {
      float tmp_conf = pred.At<float>({0, i, j + 5});
      if (tmp_conf > cls_conf)
      {
        cls_conf = tmp_conf;
        label = j;
      }
    } // argmax
    float conf = obj_conf * cls_conf; // cls_conf (0.,1.)
    if (conf < score_threshold) continue; // filter

    const int grid0 = anchors.at(i).grid0;
    const int grid1 = anchors.at(i).grid1;
    const int stride = anchors.at(i).stride;

    float dx = pred.At<float>({0, i, 0});
    float dy = pred.At<float>({0, i, 1});
    float dw = pred.At<float>({0, i, 2});
    float dh = pred.At<float>({0, i, 3});

    float cx = (dx + (float) grid0) * (float) stride;
    float cy = (dy + (float) grid1) * (float) stride;
    float w = std::exp(dw) * (float) stride;
    float h = std::exp(dh) * (float) stride;
    float x1 = ((cx - w / 2.f) - (float) dw_) / r_;
    float y1 = ((cy - h / 2.f) - (float) dh_) / r_;
    float x2 = ((cx + w / 2.f) - (float) dw_) / r_;
    float y2 = ((cy + h / 2.f) - (float) dh_) / r_;

    types::Boxf box;
    box.x1 = std::max(0.f, x1);
    box.y1 = std::max(0.f, y1);
    box.x2 = std::min(x2, (float) img_width - 1.0f);
    box.y2 = std::min(y2, (float) img_height - 1.0f);
    box.score = conf;
    box.label = label;
    box.label_text = class_names[label];
    box.flag = true;
    bbox_collection.push_back(box);

    count += 1; // limit boxes for nms.
    if (count > max_nms)
      break;
  }
#if LITEORT_DEBUG
  std::cout << "detected num_anchors: " << num_anchors << "\n";
  std::cout << "generate_bboxes num: " << bbox_collection.size() << "\n";
#endif
}

void YoloX::nms(std::vector<types::Boxf> &input, std::vector<types::Boxf> &output,
                float iou_threshold, unsigned int topk, unsigned int nms_type)
{
  if (nms_type == NMS::BLEND) lite::utils::blending_nms(input, output, iou_threshold, topk);
  else if (nms_type == NMS::OFFSET) lite::utils::offset_nms(input, output, iou_threshold, topk);
  else lite::utils::hard_nms(input, output, iou_threshold, topk);
}
  • 第五步: 在 lite/models.h 中添加类型别名,进行命名空间管理,其他的推理引擎版本的命名空间管理类似。(这个步骤后面会优化,比如会增加一些类似DECLARE_ORT_MODEL(...)的宏来方便用户使用,不过目前就先按照手动添加的方式来吧)

// ENABLE_ONNXRUNTIME
#ifdef ENABLE_ONNXRUNTIME
// ...
#include "lite/ort/cv/yolox.h"
#endif

// 默认版本
namespace lite 
{
  namespace cv 
  {
#ifdef BACKEND_ONNXRUNTIME
    typedef ortcv::YoloX _YoloX;
#endif
  }
    // 2. general object detection
  namespace detection
  {
#ifdef BACKEND_ONNXRUNTIME
    typedef _YoloX YoloX;
#endif
  }
}
// 还有个onnxruntime的命名空间也要添加
namespace lite
{
  namespace onnxruntime 
  {
    namespace cv
    {
      typedef ortcv::YoloX _ONNXYoloX;
    }
    // 2. general object detection
    namespace detection
    {
      typedef _ONNXYoloX YoloX;
    }
  }
}
  • 第六步: 编写测试工程 examples/lite/cv/test_lite_yolox.cpp

#include "lite/lite.h"

static void test_default()
{
  std::string onnx_path = "../../../hub/onnx/cv/yolox_s.onnx";
  std::string test_img_path = "../../../examples/lite/resources/test_lite_yolox_1.jpg";
  std::string save_img_path = "../../../logs/test_lite_yolox_1.jpg";

  // 1. Test Default Engine ONNXRuntime
  lite::cv::detection::YoloX *yolox = new lite::cv::detection::YoloX(onnx_path); // default

  std::vector<lite::types::Boxf> detected_boxes;
  cv::Mat img_bgr = cv::imread(test_img_path);
  yolox->detect(img_bgr, detected_boxes);

  lite::utils::draw_boxes_inplace(img_bgr, detected_boxes);

  cv::imwrite(save_img_path, img_bgr);

  std::cout << "Default Version Detected Boxes Num: " << detected_boxes.size() << std::endl;

  delete yolox;

}

static void test_onnxruntime()
{
#ifdef ENABLE_ONNXRUNTIME
  std::string onnx_path = "../../../hub/onnx/cv/yolox_s.onnx";
  std::string test_img_path = "../../../examples/lite/resources/test_lite_yolox_2.jpg";
  std::string save_img_path = "../../../logs/test_lite_yolox_2.jpg";

  // 2. Test Specific Engine ONNXRuntime
  lite::onnxruntime::cv::detection::YoloX *yolox =
      new lite::onnxruntime::cv::detection::YoloX(onnx_path);

  std::vector<lite::types::Boxf> detected_boxes;
  cv::Mat img_bgr = cv::imread(test_img_path);
  yolox->detect(img_bgr, detected_boxes);

  lite::utils::draw_boxes_inplace(img_bgr, detected_boxes);

  cv::imwrite(save_img_path, img_bgr);

  std::cout << "ONNXRuntime Version Detected Boxes Num: " << detected_boxes.size() << std::endl;

  delete yolox;
#endif
}

static void test_mnn()
{
#ifdef ENABLE_MNN
  // ...
#endif
}

static void test_ncnn()
{
#ifdef ENABLE_NCNN
  // ...
#endif
}

static void test_tnn()
{
#ifdef ENABLE_TNN
  // ...
#endif
}

static void test_lite()
{
  test_default();
  test_onnxruntime();
  test_mnn();
  test_ncnn();
  test_tnn();
}

int main(__unused int argc, __unused char *argv[])
{
  test_lite();
  return 0;
}
  • 第七步: 在 examples/CMakeLists.txt中增加编译可执行文件的选项。

# ...
add_lite_executable(lite_yolox cv)

注意这步是需要一定的命名规范的,为了能够使用add_lite_executable函数,你必须把测试用例的cpp命名为以下格式:

test_lite_xxx.cpp # 添加到CMakeLists时,使用add_lite_executable(lite_xxx, cv)
  • 第八步: 重新编译工程,并测试示例(Mac/Linux,Windows还需要手动拷贝编译好的lite.ai.toolkit.dll以及其他依赖库到 build/lite.ai.toolkit/bin)

sh ./build.sh && cd build/lite.ai.toolkit/bin && ./lite_yolox

提示:对于多输入多输出模型,不能继承BasicOrtHandler,需要单独的实现,请参考 rvm.h 和 rvm.cpp 的做法。

4. 总结

OK,往 lite.ai.toolkit 工具箱中添加一个新模型的基本步骤就是这样。无论你是否熟悉推理部署,理解上面的步骤应该是不难的。按照上面的步骤,您就可以使用 lite.ai.toolkit 原有的一些能力,并在需要的时候,增加自己需要的新模型。 另外,我也在lite.ai.toolkit的repo中,同步开了一个issue,这个issue主要是讲这个问题。大家有疑惑的可以提在这个issue,我会尽可能回答~

https://github.com/DefTruth/lite.ai.toolkit/issues/191https://github.com/DefTruth/lite.ai.toolkit/issues/191

还是那句话,受到一些关注,真是诚惶诚恐,这个小工具应该就会继续维护下去,对我来说也是个不错的学习的过程,所以来日方长,咱们也且行且看 ~ (预告:盲猜今年内能出个1.0版本~)

平时会写点破破烂烂的技术文章,感兴趣的可以关注我的技术专栏:

C++/推理引擎/渲染引擎笔记 - 知乎C++/推理引擎/渲染引擎笔记https://www.zhihu.com/column/c_1360887484541452288

  • 1
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值