#pragma once #include #include #include #include "yolov8_utils.h" #include //#include //if use OrtTensorRTProviderOptionsV2 //#include class Yolov8Onnx { public: Yolov8Onnx() :_OrtMemoryInfo(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeCPUOutput)) {}; ~Yolov8Onnx() { if (_OrtSession != nullptr) delete _OrtSession; };// delete _OrtMemoryInfo; public: /** \brief Read onnx-model * \param[in] modelPath:onnx-model path * \param[in] isCuda:if true,use Ort-GPU,else run it on cpu. * \param[in] cudaID:if isCuda==true,run Ort-GPU on cudaID. * \param[in] warmUp:if isCuda==true,warm up GPU-model. */ bool ReadModel(const std::string& modelPath, bool isCuda = false, int cudaID = 0, bool warmUp = true); /** \brief detect. * \param[in] srcImg:a 3-channels image. * \param[out] output:detection results of input image. */ bool OnnxDetect(cv::Mat& srcImg, std::vector& output); /** \brief detect,batch size= _batchSize * \param[in] srcImg:A batch of images. * \param[out] output:detection results of input images. */ bool OnnxBatchDetect(std::vector& srcImg, std::vector>& output); private: template T VectorProduct(const std::vector& v) { return std::accumulate(v.begin(), v.end(), 1, std::multiplies()); }; int Preprocessing(const std::vector& srcImgs, std::vector& outSrcImgs, std::vector& params); const int _netWidth = 640; //ONNX-net-input-width const int _netHeight = 640; //ONNX-net-input-height int _batchSize = 1; //if multi-batch,set this bool _isDynamicShape = false;//onnx support dynamic shape float _classThreshold = 0.55; float _nmsThreshold = 0.45; float _maskThreshold = 0.5; //ONNXRUNTIME Ort::Env _OrtEnv = Ort::Env(OrtLoggingLevel::ORT_LOGGING_LEVEL_ERROR, "Yolov8"); Ort::SessionOptions _OrtSessionOptions = Ort::SessionOptions(); Ort::Session* _OrtSession = nullptr; Ort::MemoryInfo _OrtMemoryInfo; #if ORT_API_VERSION < ORT_OLD_VISON char* _inputName, * _output_name0; #else std::shared_ptr _inputName, _output_name0; #endif std::vector _inputNodeNames; //输入节点名 std::vector _outputNodeNames;//输出节点名 size_t _inputNodesNum = 0; //输入节点数 size_t _outputNodesNum = 0; //输出节点数 ONNXTensorElementDataType _inputNodeDataType; //数据类型 ONNXTensorElementDataType _outputNodeDataType; std::vector _inputTensorShape; //输入张量shape std::vector _outputTensorShape; public: std::vector _className = { "earAxis", "auxHook", "mainHook" }; };