intrusion-detective/services/intrusion-detection/include/rkYolov5s.hpp

#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