src/plugins/Onnx.cpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 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 https://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/EventData/NeuralCalibrator.hpp"
0010 #include "ActsExamples/TrackFindingML/AmbiguityResolutionMLAlgorithm.hpp"
0011 #include "ActsExamples/TrackFindingML/SeedFilterMLAlgorithm.hpp"
0012 #include "ActsPython/Utilities/Helpers.hpp"
0013 #include "ActsPython/Utilities/Macros.hpp"
0014
0015 #include <pybind11/pybind11.h>
0016 #include <pybind11/stl.h>
0017
0018 namespace py = pybind11;
0019
0020 using namespace Acts;
0021 using namespace ActsExamples;
0022 using namespace ActsPython;
0023
0024 PYBIND11_MODULE(ActsExamplesPythonBindingsOnnx, onnx) {
0025   onnx.def(
0026       "makeNeuralCalibrator",
0027       [](const char *modelPath, std::size_t nComp,
0028          std::vector<std::size_t> volumeIds)
0029           -> std::shared_ptr<MeasurementCalibrator> {
0030         return std::make_shared<NeuralCalibrator>(modelPath, nComp, volumeIds);
0031       },
0032       py::arg("modelPath"), py::arg("nComp") = 1,
0033       py::arg("volumeIds") = std::vector<std::size_t>({7, 8, 9}));
0034
0035   ACTS_PYTHON_DECLARE_ALGORITHM(
0036       AmbiguityResolutionMLAlgorithm, onnx, "AmbiguityResolutionMLAlgorithm",
0037       inputTracks, inputDuplicateNN, outputTracks, nMeasurementsMin);
0038
0039   ACTS_PYTHON_DECLARE_ALGORITHM(SeedFilterMLAlgorithm, onnx,
0040                                 "SeedFilterMLAlgorithm", inputTrackParameters,
0041                                 inputSimSeeds, inputSeedFilterNN,
0042                                 outputTrackParameters, outputSimSeeds,
0043                                 epsilonDBScan, minPointsDBScan, minSeedScore);
0044 }