Plugins/Onnx/MLTrackClassifier.hpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2020 CERN for the benefit of the Acts project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include "Acts/Plugins/Onnx/OnnxRuntimeBase.hpp"
0012
0013 #include <vector>
0014
0015 namespace Acts {
0016
0017 // Specialized class that labels tracks as good/duplicate/fake using a
0018 // deep neural network.
0019 class MLTrackClassifier : public OnnxRuntimeBase {
0020   using OnnxRuntimeBase::OnnxRuntimeBase;
0021
0022  public:
0023   /// @brief The labels for track quality
0024   enum class TrackLabels { eGood, eDuplicate, eFake };
0025
0026   /// @brief Predict the track label
0027   ///
0028   /// @param inputFeatures The vector of input features for the trajectory to be
0029   /// classified
0030   /// @param decisionThreshProb The probability threshold used to predict the
0031   /// track label
0032   ///
0033   /// @return The predicted track label of the trajectory
0034   TrackLabels predictTrackLabel(std::vector<float>& inputFeatures,
0035                                 double decisionThreshProb) const;
0036
0037   /// @brief Check if the predicted track label is 'duplicate'
0038   ///
0039   /// @param inputFeatures The vector of input features for the trajectory to be
0040   /// classified
0041   /// @param decisionThreshProb The probability threshold used to predict the
0042   /// track label
0043   ///
0044   /// @return If the predicted track label is 'duplicate'
0045   bool isDuplicate(std::vector<float>& inputFeatures,
0046                    double decisionThreshProb) const;
0047 };
0048
0049 }  // namespace Acts