#include "wheel-detector.h"

TensorrtWheelDetector::TensorrtWheelDetector(const std::string &model_file, const std::string &class_file){
    cudaSetDevice(0);
    
    yolov8_ = new YOLOv8_seg(model_file);
    yolov8_->make_pipe(false);
    imgsz_=cv::Size{640, 480};
    seg_h_        = 120;
    seg_w_        = 160;
    seg_channels_ = 32;
}

TensorrtWheelDetector::~TensorrtWheelDetector(){
    if(yolov8_!=nullptr){
        delete yolov8_;
        yolov8_=nullptr;
    }

}

bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs){
    if(yolov8_==nullptr){
        return false;
    }
    
    yolov8_->copy_from_Mat(img, imgsz_);
    yolov8_->infer();
    float score_thres=0.9;
    float iou_thres=0.65;
    int topk=10;
    yolov8_->postprocess(objs, score_thres, iou_thres, topk, seg_channels_, seg_h_, seg_w_);
    return true;

}

bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res){
    if(detect(img,objs))
    {
        const std::vector<std::string> classes={"none","wheel"};
        const std::vector<std::vector<unsigned int>> colors = {{0, 114, 189},   {0, 255, 0}};
        const std::vector<std::vector<unsigned int>> mask_colors = {{255, 56, 56},  {255, 0, 0}};
        yolov8_->draw_objects(img, res, objs, classes, colors, mask_colors);
        return true;
    }else{
        return false;
    }
}

std::vector<cv::Point> TensorrtWheelDetector::getPointsFromObj(Object obj){
    std::vector<cv::Point> ret;
    int x=int(obj.rect.x+0.5);
    int y=int(obj.rect.y+0.5);
    int width=int(obj.rect.width);
    int height=int(obj.rect.height);
    
    for(int i=0;i<height;++i){
        for(int j=0;j<width;++j){
            if(obj.boxMask.at<uchar>(i,j)!=0){
                ret.push_back(cv::Point((x+j),y+i));
            }
        }
    }
    return ret;
}