#pragma once #include #include #include #include "yolov8_utils.h" #include //#include //if use OrtTensorRTProviderOptionsV2 //#include class RTDETROnnx { public: RTDETROnnx() : _OrtMemoryInfo(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtDeviceAllocator, OrtMemType::OrtMemTypeCPUOutput)) {}; ~RTDETROnnx() { 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.25; 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 = { "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush" }; };