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