GraphInferenceToolBase.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
#ifndef MUONINFERENCETOOLS_GRAPHINFERENCETOOL_H
#define MUONINFERENCETOOLS_GRAPHINFERENCETOOL_H
 
#include "MuonInferenceInterfaces/IGraphInferenceTool.h"
#include "MuonInferenceInterfaces/NodeFeatureList.h"
#include "MuonInferenceInterfaces/GraphData.h"
 
#include "AthOnnxInterfaces/IOnnxRuntimeSessionTool.h"
 
#include "MuonSpacePoint/SpacePointContainer.h"
 
#include "AthenaBaseComps/AthAlgTool.h"
#include "StoreGate/ReadHandleKey.h"
#include <onnxruntime_cxx_api.h> // is this somewhere else?
#include "nlohmann/json.hpp"
 
namespace MuonML{
    class GraphInferenceToolBase : public extends<AthAlgTool, IGraphInferenceTool> {
        public:
            using base_class::base_class;
            StatusCode buildGraph(const EventContext& ctx,
                                    GraphRawData& graphData) const;
 
            StatusCode runInference(GraphRawData& graphData) const;
        
 
        protected:
            StatusCode setupModel();
    
            Ort::Session& model() const;
            SG::ReadHandleKey<MuonR4::SpacePointContainer> m_readKey{this, "ReadSpacePoints", "MuonSpacePoints"};
        private:
            NodeFeatureList m_graphFeatures{};
            ToolHandle<AthOnnx::IOnnxRuntimeSessionTool> m_onnxSessionTool{this, "ModelSession", "" };
 
    }; 
 
}
 
#endif