#pragma once

#include "inference.h"

struct Object {
    cv::Rect_<float> rect;
    int              label = 0;
    float            prob  = 0.0;
    cv::Mat          boxMask;
};
    
class TensorrtWheelDetector{
public:
    TensorrtWheelDetector(const std::string &model_file, const std::string &class_file);
    ~TensorrtWheelDetector();

    bool detect(cv::Mat& img,std::vector<Object>& objs);
    bool detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res);

    static std::vector<cv::Point> getPointsFromObj(const Object &obj);

private:
    Inference*  yolov8_;
    cv::Size     imgsz_;
    int      seg_h_        = 120;
    int      seg_w_        = 160;
    int      seg_channels_ = 32;
};