zone_alarm/src/yolo_detector.h

#pragma once
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>
#include "rknn_api.h"

struct Detection {
    cv::Rect bbox;
    float confidence;
    int class_id;
    std::string class_name;
};

class YoloDetector {
public:
    YoloDetector() = default;
    ~YoloDetector();
    
    // core_mask: 0=自动; 1=Core0; 2=Core1; 4=Core2; 7=0/1/2
    bool load(const std::string& model_path, int input_size, 
              float conf_thresh, float nms_thresh, int core_mask = 0);
    
    std::vector<Detection> detect(const float* rgb_data, int img_w, int img_h,
                                   float scale, int pad_x, int pad_y);
    
    bool isLoaded() const { return loaded_; }
    int inputSize() const { return input_size_; }

private:
    rknn_context ctx_ = 0;
    bool loaded_ = false;
    int input_size_ = 640;
    float conf_thresh_ = 0.35f;
    float nms_thresh_ = 0.45f;
    uint32_t n_input_ = 0;
    uint32_t n_output_ = 0;
    int frame_count_ = 0;
    std::vector<rknn_tensor_attr> input_attrs_;
    std::vector<rknn_tensor_attr> output_attrs_;
    
    static void nms(std::vector<Detection>& dets, float threshold);
};