56 lines
1.7 KiB
C++
Executable File
56 lines
1.7 KiB
C++
Executable File
#ifndef AUTO_BUFF__YOLO11_BUFF_HPP
|
||
#define AUTO_BUFF__YOLO11_BUFF_HPP
|
||
#include <yaml-cpp/yaml.h>
|
||
|
||
#include <filesystem>
|
||
#include <opencv2/opencv.hpp>
|
||
#include <openvino/openvino.hpp>
|
||
|
||
#include "tools/logger.hpp"
|
||
|
||
namespace auto_buff
|
||
{
|
||
const std::vector<std::string> class_names = {"buff", "r"};
|
||
|
||
class YOLO11_BUFF
|
||
{
|
||
public:
|
||
struct Object
|
||
{
|
||
cv::Rect_<float> rect;
|
||
int label;
|
||
float prob;
|
||
std::vector<cv::Point2f> kpt;
|
||
};
|
||
|
||
YOLO11_BUFF(const std::string & config);
|
||
|
||
// 使用NMS，用来获取多个框
|
||
std::vector<Object> get_multicandidateboxes(cv::Mat & image);
|
||
|
||
// 寻找置信度最高的框
|
||
std::vector<Object> get_onecandidatebox(cv::Mat & image);
|
||
|
||
private:
|
||
ov::Core core; // 创建OpenVINO Runtime Core对象
|
||
std::shared_ptr<ov::Model> model;
|
||
ov::CompiledModel compiled_model;
|
||
ov::InferRequest infer_request;
|
||
ov::Tensor input_tensor;
|
||
const int NUM_POINTS = 6;
|
||
|
||
// 转换图像数据: 先转换元素类型, (可选)然后归一化到[0, 1], (可选)然后交换RB通道
|
||
void convert(
|
||
const cv::Mat & input, cv::Mat & output, const bool normalize, const bool exchangeRB) const;
|
||
|
||
// 对网络的输入为图片数据的节点进行赋值，实现图片数据输入网络,return 缩放因子, 该缩放是为了将input_image塞进input_tensor
|
||
float fill_tensor_data_image(ov::Tensor & input_tensor, const cv::Mat & input_image) const;
|
||
|
||
// 打印模型信息, 这个函数修改自$${OPENVINO_COMMON}/utils/src/args_helper.cpp的同名函数
|
||
void printInputAndOutputsInfo(const ov::Model & network);
|
||
|
||
// 将image保存为"../result/$${programName}.jpg"
|
||
void save(const std::string & programName, const cv::Mat & image);
|
||
};
|
||
} // namespace auto_buff
|
||
#endif |