59 lines
1.3 KiB
C++
59 lines
1.3 KiB
C++
#ifndef RKYOLOV8_HPP
|
|
#define RKYOLOV8_HPP
|
|
|
|
#include <memory>
|
|
#include <opencv2/opencv.hpp>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "rknn/rknn_api.h"
|
|
#include "rknn_test/postprocess.h"
|
|
|
|
// 引入我们刚才修复的 DMA 分配器
|
|
#include "rknn_test/dma_allocator.hpp"
|
|
|
|
namespace rknn_test {
|
|
|
|
class rkYolov8 {
|
|
public:
|
|
rkYolov8(const std::string& model_path, const std::string& label_path, int class_num);
|
|
~rkYolov8();
|
|
|
|
int init(rknn_context* ctx_in, bool is_slave);
|
|
rknn_context* get_pctx();
|
|
detect_result_group_t infer(const cv::Mat& ori_img);
|
|
|
|
private:
|
|
unsigned char* load_model(const char* filename, int* model_size);
|
|
void post_process_v8_dfl(rknn_output* outputs, float scale, int pad_w, int pad_h,
|
|
detect_result_group_t* group);
|
|
|
|
private:
|
|
std::string model_path;
|
|
std::string m_label_path;
|
|
int m_class_num;
|
|
|
|
rknn_context ctx;
|
|
bool is_slave = false;
|
|
unsigned char* model_data = nullptr;
|
|
|
|
rknn_input_output_num io_num;
|
|
rknn_tensor_attr* input_attrs = nullptr;
|
|
rknn_tensor_attr* output_attrs = nullptr;
|
|
|
|
rknn_input inputs[1];
|
|
|
|
int width = 0;
|
|
int height = 0;
|
|
int channel = 0;
|
|
|
|
float conf_threshold;
|
|
float nms_threshold;
|
|
|
|
// 【修改点】使用 DmaBuffer 智能指针管理 NPU 输入内存
|
|
std::unique_ptr<DmaBuffer> input_dma_buf_;
|
|
};
|
|
|
|
} // namespace rknn_test
|
|
|
|
#endif // RKYOLOV8_HPP
|