leonarski_f
d315506633
* Enhancements for EIGER * Writer is more flexible and capable of handling DECTRIS data
36 lines
1.2 KiB
C++
36 lines
1.2 KiB
C++
// Copyright (2019-2024) Paul Scherrer Institute
|
|
|
|
#ifndef JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
|
|
#define JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
|
|
|
|
#ifdef JFJOCH_USE_TORCH
|
|
#include <torch/script.h>
|
|
#endif
|
|
|
|
#include "../common/DiffractionExperiment.h"
|
|
|
|
// Based on model described in:
|
|
// Mendez, D., Holton, J. M., Lyubimov, A. Y., Hollatz, S., Mathews, I. I., Cichosz, A., Martirosyan, V.,
|
|
// Zeng, T., Stofer, R., Liu, R., Song, J., McPhillips, S., Soltis, M. & Cohen, A. E. (2024).
|
|
// Acta Cryst. D80, 26-43.
|
|
|
|
class NeuralNetResPredictor {
|
|
std::vector<float> model_input;
|
|
#ifdef JFJOCH_USE_TORCH
|
|
torch::Device device;
|
|
torch::jit::script::Module module;
|
|
#endif
|
|
template<class T>
|
|
void PrepareInternal(const DiffractionExperiment& experiment, const T* image);
|
|
public:
|
|
explicit NeuralNetResPredictor(const std::string& model_path);
|
|
void Prepare(const DiffractionExperiment& experiment, const int16_t* image);
|
|
void Prepare(const DiffractionExperiment& experiment, const int32_t* image);
|
|
float Inference(const DiffractionExperiment& experiment, const void* image);
|
|
size_t GetMaxPoolFactor(const DiffractionExperiment& experiment) const;
|
|
const std::vector<float> &GetModelInput() const;
|
|
};
|
|
|
|
|
|
#endif //JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
|