ONNXPredictorBase.h

Go to the documentation of this file.

// Copyright (c) 2002-present, OpenMS Team -- EKU Tuebingen, ETH Zurich, and FU Berlin
// SPDX-License-Identifier: BSD-3-Clause
//
// --------------------------------------------------------------------------
// $Maintainer: Timo Sachsenberg $
// $Authors: Satyam Yadav $
// --------------------------------------------------------------------------
 
#pragma once
 
#include <OpenMS/config.h>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <string>
#include <memory>
#include <vector>
 
// Forward declaration to avoid exposing ONNX headers globally
namespace Ort {
    struct Session;
    struct SessionOptions;
    struct MemoryInfo;
}
 
namespace OpenMS
{
    class OPENMS_DLLAPI ONNXPredictorBase
    {
    public:
        explicit ONNXPredictorBase(const std::string& model_path, int threads = 4);
 
        virtual ~ONNXPredictorBase();
 
        ONNXPredictorBase(const ONNXPredictorBase&) = delete;
        ONNXPredictorBase& operator=(const ONNXPredictorBase&) = delete;
 
        ONNXPredictorBase(ONNXPredictorBase&&) noexcept;
        ONNXPredictorBase& operator=(ONNXPredictorBase&&) noexcept;
 
        Ort::Session& session();
        const Ort::Session& session() const;
 
        Ort::MemoryInfo& memoryInfo();
        const Ort::MemoryInfo& memoryInfo() const;
 
        std::vector<std::string> getInputNames() const;
 
        std::vector<std::string> getOutputNames() const;
 
        std::vector<int64_t> getInputShape(size_t input_index) const;
 
        size_t getInputCount() const;
 
        size_t getOutputCount() const;
 
        void forEachBatch(const std::vector<std::vector<size_t>>& groups, size_t batch_size,
                          const std::function<void(const size_t* indices, size_t n)>& run) const;
 
    private:
        int threads_ = 1;
        std::unique_ptr<Ort::SessionOptions> session_options_;
        std::unique_ptr<Ort::Session> session_;
        std::unique_ptr<Ort::MemoryInfo> memory_info_;
    };
} // namespace OpenMS