CbmRichMCbmDenoiseCnn.h

Go to the documentation of this file.

/* Copyright (C) 2024 UGiessen, Giessen
   SPDX-License-Identifier: GPL-3.0-only
   Authors: Martin Beyer [committer] */
 
#if HAVE_ONNXRUNTIME
 
#pragma once
 
#include <TObject.h>
#include <TSystem.h>
 
#include <array>
#include <string>
#include <vector>
 
#include <onnxruntime/core/session/onnxruntime_cxx_api.h>
 
class CbmEvent;
class CbmRichDigiMapManager;
class TClonesArray;
 
class CbmRichMCbmDenoiseCnn : public TObject {
 public:
  CbmRichMCbmDenoiseCnn() = default;
 
  ~CbmRichMCbmDenoiseCnn() = default;
 
  CbmRichMCbmDenoiseCnn(const CbmRichMCbmDenoiseCnn&) = delete;
 
  CbmRichMCbmDenoiseCnn operator=(const CbmRichMCbmDenoiseCnn&) = delete;
 
  void Init();
 
  void Process(CbmEvent* event, const TClonesArray* richHits);
 
  void ProcessTimeWindow(const std::vector<int>& timeWindowHitIndices, const double& seedHitTime,
                         const TClonesArray* richHits);
 
  int AddressToGridIndex(const int address);
 
  std::vector<float> Inference(const std::vector<int>& gridIndices);
 
  void SetClassificationThreshold(float threshold) { fClassificationThreshold = threshold; }
 
 private:
  const double fTimeWindowLength{25.};
  const int fMinHitsInTimeWindow{7};
 
  float fClassificationThreshold{0.5};
 
  const std::string fOnnxFilePath =
    std::string(gSystem->Getenv("VMCWORKDIR")) + "/parameters/rich/mRich/rich_v21c_mcbm_UNet.onnx";
 
  // TODO: read names/shapes from onnx file
  const std::vector<const char*> fInputNames{"input"};
  const std::vector<const char*> fOutputNames{"output"};
 
  const std::array<int64_t, 4> fInputShape{1, 1, 72, 32};
  const std::array<int64_t, 4> fOutputShape{1, 1, 72, 32};
 
  std::array<float, 2304> fInput{};
 
  std::unique_ptr<Ort::Env> fOrtEnv{nullptr};
  std::unique_ptr<Ort::SessionOptions> fOrtSessionOptions{nullptr};
  std::unique_ptr<Ort::Session> fOrtSession{nullptr};
  std::unique_ptr<Ort::MemoryInfo> fOrtAllocatorInfo{nullptr};
  std::unique_ptr<Ort::RunOptions> fOrtRunOptions{nullptr};
 
  std::unique_ptr<Ort::Value> fOrtInput{nullptr};
 
  CbmRichDigiMapManager* fCbmRichDigiMapManager{nullptr};
 
  ClassDef(CbmRichMCbmDenoiseCnn, 1);
};
 
#endif