File model.h

File List > model.h

Go to the documentation of this file

#pragma once
#ifdef DLL_EXPORT
#define DLL_API __declspec(dllexport)
#else
#define DLL_API __declspec(dllimport)
#endif

#include "config.h"
#include "type.h"
#include "processer.h"
#include "model_result.h"
#include "onnxruntime_cxx_api.h"

#include <Windows.h>
#include <opencv2/opencv.hpp>

using namespace cv;

namespace inspection
{
    class Model
    {
    private:
        TaskType _task;
        string _weightPath;
        int _imageWidth, _imageHeight;
        string _device;
        int _batchSize;
        DomainConfig _domainConfig;
        int _numOfClass;
        string _outputDir;

        // for onnx
        Ort::Env _env;
        Ort::SessionOptions _sessionOption;
        Ort::Session* _session;
        Ort::AllocatorWithDefaultOptions _allocator;
        Ort::RunOptions _runOptions;

        std::vector<float> _inputImage;
        std::vector<int64_t> _inputShape, _outputShape;
        std::vector<float> _results;
        size_t _numInputNodes, _numOutputNodes;
        std::vector<const char*> _inputNodeNames, _outputNodeNames;

    public:
        Model(ModelConfig mc);
        ModelResult Run(Mat image);
        float* RunOnnx(Mat image);
        void SetEnvForOnnxruntime();
        void ReleaseEnv();
        DomainConfig GetDomainConfig();
    };
}