algorithms/onnx/ONNXInference.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 - 2024 Sylvester Joosten, Dmitry Romanov, Wouter Deconinck, Dmitry Kalinkin
0003
0004 #pragma once
0005
0006 #include <algorithms/algorithm.h>
0007 #include <cstdint>
0008 #include <onnxruntime_cxx_api.h>
0009 #include <string>
0010 #include <string_view>
0011 #include <vector>
0012 #include <edm4eic/TensorCollection.h>
0013
0014 #include "algorithms/interfaces/WithPodConfig.h"
0015 #include "algorithms/onnx/ONNXInferenceConfig.h"
0016
0017 namespace eicrecon {
0018
0019 using ONNXInferenceAlgorithm =
0020     algorithms::Algorithm<algorithms::Input<std::vector<edm4eic::TensorCollection>>,
0021                           algorithms::Output<std::vector<edm4eic::TensorCollection>>>;
0022
0023 class ONNXInference : public ONNXInferenceAlgorithm,
0024                               public WithPodConfig<ONNXInferenceConfig> {
0025
0026 public:
0027   ONNXInference(std::string_view name)
0028       : ONNXInferenceAlgorithm{name,
0029                                        {"inputTensors"},
0030                                        {"outputTensors"},
0031                                        ""} {
0032   }
0033
0034   void init() final;
0035   void process(const Input&, const Output&) const final;
0036
0037 private:
0038   mutable Ort::Env m_env{nullptr};
0039   mutable Ort::Session m_session{nullptr};
0040
0041   std::vector<std::string> m_input_names;
0042   std::vector<const char*> m_input_names_char;
0043   std::vector<std::vector<std::int64_t>> m_input_shapes;
0044
0045   std::vector<std::string> m_output_names;
0046   std::vector<const char*> m_output_names_char;
0047   std::vector<std::vector<std::int64_t>> m_output_shapes;
0048 };
0049
0050 } // namespace eicrecon