OnnxProperty.h

Go to the documentation of this file.

 
#ifndef _SEIMPLEMENTATION_PLUGIN_ONNXATTRIBUTE_H_
#define _SEIMPLEMENTATION_PLUGIN_ONNXATTRIBUTE_H_
 
#include "SEFramework/Property/Property.h"
#include <memory>
#include <NdArray/NdArray.h>
 
namespace SourceXtractor {

class OnnxProperty : public Property {
public:
  struct NdWrapperBase {
    virtual ~NdWrapperBase() = default;
  };

  template<typename T>
  struct NdWrapper : public NdWrapperBase {
    Euclid::NdArray::NdArray<T> m_ndarray;
 
    template <typename ...Args>
    explicit NdWrapper(Args&&... args) : m_ndarray(std::forward<Args>(args)...) {}
  };

  virtual ~OnnxProperty() = default;

  explicit OnnxProperty(std::map<std::string, std::unique_ptr<NdWrapperBase>>&& output) : m_output(std::move(output)) {}

  template<typename T>
  const Euclid::NdArray::NdArray<T>& getData(const std::string& key) const {
    return dynamic_cast<NdWrapper<T> *>(m_output.at(key).get())->m_ndarray;
  }
 
private:
  std::map<std::string, std::unique_ptr<NdWrapperBase>> m_output;
};
 
} // end of namespace SourceXtractor
 
#endif //_SEIMPLEMENTATION_PLUGIN_ONNXATTRIBUTE_H_