68 lines
2.0 KiB
C++
68 lines
2.0 KiB
C++
#ifndef RKYOLOV5S_H
|
|
#define RKYOLOV5S_H
|
|
|
|
#include "rknn_api.h"
|
|
#include "opencv2/core/core.hpp"
|
|
#include <map> // 用于存储跟踪目标
|
|
#include <mutex> // 确保线程安全
|
|
#include <string> // 使用 std::string
|
|
#include <vector> // 使用 std::vector
|
|
#include "postprocess.h"
|
|
|
|
// 前置声明
|
|
static void dump_tensor_attr(rknn_tensor_attr *attr);
|
|
static unsigned char *load_data(FILE *fp, size_t ofst, size_t sz);
|
|
static unsigned char *load_model(const char *filename, int *model_size);
|
|
static int saveFloat(const char *file_name, float *output, int element_size);
|
|
|
|
// 用于跟踪单个目标的结构体
|
|
struct TrackedPerson
|
|
{
|
|
int id; // 唯一ID
|
|
cv::Rect box; // 当前位置
|
|
double entry_time; // 进入入侵区域的时间戳 (秒)
|
|
bool is_in_zone; // 是否在区域内
|
|
bool alarm_triggered; // 是否已触发报警
|
|
int frames_unseen; // 消失的帧数
|
|
};
|
|
|
|
class rkYolov5s
|
|
{
|
|
private:
|
|
int ret;
|
|
std::mutex mtx;
|
|
std::string model_path;
|
|
unsigned char *model_data;
|
|
|
|
rknn_context ctx;
|
|
rknn_input_output_num io_num;
|
|
rknn_tensor_attr *input_attrs;
|
|
rknn_tensor_attr *output_attrs;
|
|
rknn_input inputs[1];
|
|
|
|
int channel, width, height;
|
|
int img_width, img_height;
|
|
|
|
float nms_threshold, box_conf_threshold;
|
|
|
|
// 入侵检测和跟踪相关成员
|
|
cv::Rect intrusion_zone; // 入侵区域
|
|
std::map<int, TrackedPerson> tracked_persons; // 存储所有被跟踪的人
|
|
int next_track_id; // 用于分配新的唯一ID
|
|
double intrusion_time_threshold; // 入侵时间阈值 (秒)
|
|
|
|
// 跟踪逻辑的私有方法
|
|
void update_tracker(detect_result_group_t &detect_result_group);
|
|
|
|
public:
|
|
rkYolov5s(const std::string &model_path);
|
|
int init(rknn_context *ctx_in, bool isChild);
|
|
rknn_context *get_pctx();
|
|
cv::Mat infer(cv::Mat &ori_img);
|
|
~rkYolov5s();
|
|
|
|
// 用于从外部设置入侵区域的公共方法
|
|
void set_intrusion_zone(const cv::Rect& zone);
|
|
};
|
|
|
|
#endif // RKYOLOV5S_H
|