OnnxTaskFactory.cpp

Go to the documentation of this file.

#include <onnxruntime_cxx_api.h>

#include <AlexandriaKernel/memory_tools.h>
#include <NdArray/NdArray.h>

#include "SEImplementation/Common/OnnxCommon.h"

#include "SEImplementation/Plugin/Onnx/OnnxPlugin.h"
#include "SEImplementation/Plugin/Onnx/OnnxSourceTask.h"
#include "SEImplementation/Plugin/Onnx/OnnxProperty.h"
#include "SEImplementation/Plugin/Onnx/OnnxConfig.h"

#include "SEImplementation/Plugin/Onnx/OnnxTaskFactory.h"

namespace SourceXtractor {

static std::string generatePropertyName(const OnnxModel& model) {
  std::stringstream prop_name;

  std::string domain = model.getDomain();
  if (!domain.empty()) {
    prop_name << domain << '.';
  }

  std::string graph_name = model.getGraphName();
  if (!graph_name.empty()) {
    prop_name << graph_name << '.';
  }

  prop_name << model.getOutputName();

  return prop_name.str();
}

OnnxTaskFactory::OnnxTaskFactory() {}

std::shared_ptr<Task> OnnxTaskFactory::createTask(const PropertyId& property_id) const {
  if (property_id == PropertyId::create<OnnxProperty>()) {
    return std::make_shared<OnnxSourceTask>(m_model_infos);
  }
  return nullptr;
}

void OnnxTaskFactory::reportConfigDependencies(Euclid::Configuration::ConfigManager& manager) const {
  manager.registerConfiguration<OnnxConfig>();
}

void OnnxTaskFactory::configure(Euclid::Configuration::ConfigManager& manager) {
  const auto& onnx_config = manager.getConfiguration<OnnxConfig>();
  const auto& models = onnx_config.getModels();

  for (auto model_path : models) {
    auto model = std::make_shared<OnnxModel>(model_path);

    if (model->getInputType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
      throw Elements::Exception() << "Only ONNX models with float input are supported";
    }

    if (model->getInputShape().size() != 4) {
      throw Elements::Exception() << "Expected 4 axes for the input layer, got " << model->getInputShape().size();
    }

    auto prop_name = generatePropertyName(*model);
    onnx_logger.info() << "Output name will be " << prop_name;

    m_model_infos.emplace_back(OnnxSourceTask::OnnxModelInfo {model, prop_name});

  }
}

template<typename T>
static void registerColumnConverter(OutputRegistry& registry, const OnnxSourceTask::OnnxModelInfo& model_info) {
  auto key = model_info.prop_name;

  registry.registerColumnConverter<OnnxProperty, Euclid::NdArray::NdArray<T>>(
      model_info.prop_name, [key](const OnnxProperty& prop) {
      return prop.getData<T>(key);
    }, "", model_info.model->getModelPath()
  );
}

void OnnxTaskFactory::registerPropertyInstances(OutputRegistry& registry) {
  for (const auto& model_info : m_model_infos) {
    switch (model_info.model->getOutputType()) {
      case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT:
        registerColumnConverter<float>(registry, model_info);
        break;
      case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32:
        registerColumnConverter<int32_t>(registry, model_info);
        break;
      default:
        throw Elements::Exception() << "Unsupported output type: " << model_info.model->getOutputType();
    }
  }
}

} // end of namespace SourceXtractor