55 lines
2.6 KiB
C++
55 lines
2.6 KiB
C++
#ifndef ANSONNXCL_H
|
|
#define ANSONNXCL_H
|
|
#pragma once
|
|
#include "ANSEngineCommon.h"
|
|
#include "engine.h"
|
|
|
|
namespace ANSCENTER {
|
|
class ANSENGINE_API ANSONNXCL :public ANSODBase {
|
|
public:
|
|
virtual bool Initialize(std::string licenseKey, ModelConfig modelConfig, const std::string& modelZipFilePath, const std::string& modelZipPassword, std::string& labelMap) override;
|
|
virtual bool LoadModel(const std::string& modelZipFilePath, const std::string& modelZipPassword)override;
|
|
virtual bool LoadModelFromFolder(std::string licenseKey, ModelConfig modelConfig, std::string modelName, std::string className, const std::string& modelFolder, std::string& labelMap)override;
|
|
virtual bool OptimizeModel(bool fp16, std::string& optimizedModelFolder);
|
|
std::vector<Object> RunInference(const cv::Mat& input);
|
|
std::vector<Object> RunInference(const cv::Mat& input, const std::string& camera_id);
|
|
bool Destroy();
|
|
~ANSONNXCL();
|
|
private:
|
|
std::string _modelFilePath;
|
|
bool _modelLoadValid;
|
|
bool _fp16{ false };
|
|
size_t vectorProduct(const std::vector<int64_t>& vector);
|
|
bool preprocessImageToTensor(const cv::Mat& image, cv::Mat& outImage,
|
|
const cv::Size& targetShape,
|
|
const cv::Scalar& color = cv::Scalar(0, 0, 0),
|
|
bool scaleUp = true,
|
|
const std::string& strategy = "resize");
|
|
void warmupModel();
|
|
bool Init(const std::string& modelPath, const cv::Size& targetInputShape, bool useGPU = true);
|
|
bool preprocess(const cv::Mat& image, float*& blob, std::vector<int64_t>& inputTensorShape);
|
|
std::vector<Object> postprocess(const std::vector<Ort::Value>& outputTensors, const std::string& camera_id);
|
|
std::vector<Object> classify(const cv::Mat& image, const std::string& camera_id);
|
|
|
|
private:
|
|
Ort::Env env_{ nullptr };
|
|
Ort::SessionOptions sessionOptions_{ nullptr };
|
|
Ort::Session session_{ nullptr };
|
|
|
|
bool isDynamicInputShape_{};
|
|
cv::Size inputImageShape_{};
|
|
|
|
std::vector<Ort::AllocatedStringPtr> inputNodeNameAllocatedStrings_{};
|
|
std::vector<const char*> inputNames_{};
|
|
std::vector<Ort::AllocatedStringPtr> outputNodeNameAllocatedStrings_{};
|
|
std::vector<const char*> outputNames_{};
|
|
|
|
size_t numInputNodes_{}, numOutputNodes_{};
|
|
int numClasses_{ 0 };
|
|
float m_imgWidth = 0;
|
|
float m_imgHeight = 0;
|
|
};
|
|
}
|
|
#endif
|
|
#pragma once
|