tea_cv/inference.h

#ifndef INFERENCE_H
#define INFERENCE_H

// Cpp native
#include <fstream>
#include <vector>
#include <string>
#include <random>

// OpenCV / DNN / Inference
#include <opencv2/imgproc.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>

struct Detection
{
    int class_id{0};
    std::string className{};
    float confidence{0.0};
    cv::Scalar color{};
    cv::Rect box{};
	std::vector<cv::Point> kpts{};
};

class Inference
{
public:
    Inference(const std::string &onnxModelPath, const cv::Size2f &modelInputShape, const std::string &classesTxtFile, const bool &runWithCuda = true);
    std::vector<Detection> runInference(const cv::Mat &input);
	
	void setModelConfidenseThreshold(float t) { modelConfidenseThreshold = t; };
	void setModelScoreThreshold(float t) { modelScoreThreshold = t; };
	void setModelNMSThreshold(float t) { modelNMSThreshold = t; }
private:
    void loadClassesFromFile();
    void loadOnnxNetwork();
    cv::Mat formatToSquare(const cv::Mat &source);

    std::string modelPath{};
    std::string classesPath{};
    bool cudaEnabled{};

    std::vector<std::string> classes{};
    cv::Size2f modelShape{};

    float modelConfidenseThreshold {0.25};
    float modelScoreThreshold      {0.45f};
    float modelNMSThreshold        {0.50};

    bool letterBoxForSquare = true;

    cv::dnn::Net net;
};

#endif // INFERENCE_H