#include "infer/infer.hpp"

#include "infer/trt/yolo/yolo.hpp"
#include "infer/trt/depth/yolo.hpp"
#include <iostream>
#include <memory>


std::shared_ptr<Infer> load(const std::string& model_path, ModelType model_type, const std::vector<std::string>& names, int gpu_id, float confidence_threshold, float nms_threshold)
{
    std::shared_ptr<Infer> infer;
    switch (model_type)
    {
    case ModelType::YOLOV5:
    case ModelType::YOLOV8:
    case ModelType::YOLO11:
    case ModelType::YOLO11POSE:
        infer = yolo::load_yolo(model_path, model_type, names, gpu_id, confidence_threshold, nms_threshold);
        break;
    case ModelType::DEPTH_ANYTHING:
        infer = depth::load_depth(model_path, gpu_id);
        break;
    default:
        break;
    }
    return infer;
}