File OnnxEdgeClassifier.hpp

namespace Acts


This file is foreseen for the Geometry module to replace Extent

class OnnxEdgeClassifier : public Acts::EdgeClassificationBase
#include <Acts/Plugins/ExaTrkX/OnnxEdgeClassifier.hpp>

Public Functions

OnnxEdgeClassifier(const Config &cfg)
inline Config config() const
virtual std::tuple<std::any, std::any, std::any> operator()(std::any nodes, std::any edges, const Logger &logger) override

Perform edge classification.

  • nodes – Node tensor with shape (n_nodes, n_node_features)

  • edges – Edge-index tensor with shape (2, n_edges)

  • logger – Logger instance


(node_tensor, edge_tensor, score_tensor)

Private Members

Config m_cfg
std::unique_ptr<Ort::Env> m_env
std::string m_inputNameEdges
std::string m_inputNameNodes
std::unique_ptr<Ort::Session> m_model
std::string m_outputNameScores
struct Config
#include <Acts/Plugins/ExaTrkX/OnnxEdgeClassifier.hpp>

Public Members

float cut = 0.21
std::string modelPath
namespace Ort