ml_onnx_model.h

Go to the documentation of this file.

//=============================================================================================================
 
#ifndef ML_ONNX_MODEL_H
#define ML_ONNX_MODEL_H
 
//=============================================================================================================
// INCLUDES
//=============================================================================================================
 
#include "ml_global.h"
#include "ml_model.h"
 
//=============================================================================================================
// QT INCLUDES
//=============================================================================================================
 
#include <QString>
 
//=============================================================================================================
// STL INCLUDES
//=============================================================================================================
 
#include <memory>
#include <string>
#include <vector>
 
//=============================================================================================================
// FORWARD DECLARE ORT TYPES
//=============================================================================================================
 
#ifdef MNE_USE_ONNXRUNTIME
namespace Ort { class Env; class Session; class MemoryInfo; class RunOptions; }
#endif
 
//=============================================================================================================
// DEFINE NAMESPACE MLLIB
//=============================================================================================================
 
namespace MLLIB{
 
//=============================================================================================================
class MLSHARED_EXPORT MlOnnxModel : public MlModel
{
public:
    //=========================================================================================================
    MlOnnxModel();
 
    //=========================================================================================================
    ~MlOnnxModel() override;
 
    //=========================================================================================================
    MlTensor predict(const MlTensor& input) const override;
    bool save(const QString& path) const override;
    bool load(const QString& path) override;
    QString modelType() const override;
    MlTaskType taskType() const override;
 
    //=========================================================================================================
    bool isLoaded() const;
 
private:
#ifdef MNE_USE_ONNXRUNTIME
    static Ort::Env& ortEnv();
#endif
 
    QString     m_modelPath;                                
    MlTaskType  m_taskType = MlTaskType::Classification;    
 
#ifdef MNE_USE_ONNXRUNTIME
    std::unique_ptr<Ort::Session>    m_session;              
    std::unique_ptr<Ort::MemoryInfo> m_memoryInfo;           
    std::vector<std::string>         m_inputNames;           
    std::vector<std::string>         m_outputNames;          
    std::vector<std::vector<int64_t>> m_inputShapes;         
#endif
};
 
} // namespace MLLIB
 
#endif // ML_ONNX_MODEL_H