#include "wheel-detector.h"

TensorrtWheelDetector::TensorrtWheelDetector(const std::string &model_file, const std::string &class_file){
    yolov8_ = new Inference(model_file, cv::Size{640, 640}, class_file, false);
}

TensorrtWheelDetector::~TensorrtWheelDetector(){
}

bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs){
    std::vector<Detection> rets = yolov8_->runInference(img);

    for (auto &ret: rets) {
        Object obj;
       obj.rect = ret.box;
        obj.label = ret.class_id;
        obj.prob = ret.confidence;
        objs.push_back(obj);
    }
    return true;
}

bool TensorrtWheelDetector::detect(cv::Mat& img,std::vector<Object>& objs,cv::Mat& res){
    std::vector<Detection> rets = yolov8_->runInference(img);

    for (auto &ret: rets) {
        Object obj;
        obj.rect = ret.box;
        obj.label = ret.class_id;
        obj.prob = ret.confidence;
        obj.boxMask = img;
        objs.push_back(obj);

        cv::Rect box = ret.box;
        cv::Scalar color = ret.color;

        // Detection box
        cv::rectangle(res, box, color, 2);

        // Detection box text
        std::string classString = ret.className + ' ' + std::to_string(ret.confidence).substr(0, 4);
        cv::Size textSize = cv::getTextSize(classString, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);
        cv::Rect textBox(box.x, box.y - 40, textSize.width + 10, textSize.height + 20);

        cv::rectangle(res, textBox, color, cv::FILLED);
        cv::putText(res, classString, cv::Point(box.x + 5, box.y - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);
    }
    return true;
}

std::vector<cv::Point> TensorrtWheelDetector::getPointsFromObj(const 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){
            ret.emplace_back(x+j,y+i);
        }
    }

    return ret;
}