infer.cpp 744 B

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