Vzense-Tof-3D/detect/onnx/wheel-detector.cpp

#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);
    }
    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;
}