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;
 
    //=========================================================================================================
    QString metadata(const QString& key) 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