#ifndef INFER_HPP__
#define INFER_HPP__

#include "common/data.hpp"
#include "common/meta.hpp"
#include "common/image.hpp"

#include <variant>

enum class ModelType : int{
    YOLOV5  = 0,
    YOLOV8  = 1,
    YOLO11 = 2,
    YOLO11POSE = 3,
    DEPTH_ANYTHING = 4
};

using Result = std::variant<data::BoxArray, cv::Mat>;

class Infer{
public:
    virtual Result forward(const tensor::Image &image, int slice_width, int slice_height, float overlap_width_ratio, float overlap_height_ratio, void *stream = nullptr) = 0;
    virtual Result forward(const tensor::Image &image, void *stream = nullptr) = 0;

protected:
    std::mutex mutex_;
};

std::shared_ptr<Infer> load(const std::string& model_path, ModelType model_type, const std::vector<std::string>& names, int gpu_id=0, float confidence_threshold=0.5f, float nms_threshold=0.45f);

#endif // INFER_HPP__