infer.cpp 785 B

123456789101112131415161718192021222324252627
  1. #include "infer/infer.hpp"
  2. #include "infer/trt/yolo/yolo.hpp"
  3. #include "infer/trt/depth/yolo.hpp"
  4. #include <iostream>
  5. #include <memory>
  6. 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)
  7. {
  8. std::shared_ptr<Infer> infer;
  9. switch (model_type)
  10. {
  11. case ModelType::YOLOV5:
  12. case ModelType::YOLOV8:
  13. case ModelType::YOLO11:
  14. case ModelType::YOLO11POSE:
  15. infer = yolo::load_yolo(model_path, model_type, names, gpu_id, confidence_threshold, nms_threshold);
  16. break;
  17. case ModelType::DEPTH_ANYTHING:
  18. infer = depth::load_depth(model_path, gpu_id);
  19. break;
  20. default:
  21. break;
  22. }
  23. return infer;
  24. }