Improvements before MAX IV test

This commit is contained in:
2024-04-25 20:11:58 +02:00
parent 2c8e1fd83d
commit ea70b27e85
80 changed files with 1835 additions and 1781 deletions
+17 -10
View File
@@ -6,14 +6,19 @@
#include "../common/JFJochException.h"
NeuralNetResPredictor::NeuralNetResPredictor(const std::string& model_path)
: model_input(512*512)
: model_input(512*512),
enable(!model_path.empty())
#ifdef JFJOCH_USE_TORCH
, device(torch::kCUDA)
, device(torch::kCUDA)
#endif
{
{
#ifdef JFJOCH_USE_TORCH
module = torch::jit::load(model_path);
module.to(device);
if (enable) {
module = torch::jit::load(model_path);
module.to(device);
}
#else
enable = false;
#endif
}
@@ -65,13 +70,15 @@ size_t NeuralNetResPredictor::GetMaxPoolFactor(const DiffractionExperiment& expe
return pool_factor;
}
float NeuralNetResPredictor::Inference(const DiffractionExperiment& experiment, const void *image) {
std::optional<float> NeuralNetResPredictor::Inference(const DiffractionExperiment& experiment, const void *image) {
if (!enable)
return {};
#ifdef JFJOCH_USE_TORCH
if (experiment.GetPixelDepth() == 2)
Prepare(experiment, (int16_t *) image);
else
Prepare(experiment, (int32_t *) image);
#ifdef JFJOCH_USE_TORCH
auto options = torch::TensorOptions().dtype(at::kFloat);
auto model_input_tensor = torch::from_blob(model_input.data(), {1,1,512,512}, options).to(device);
@@ -83,10 +90,10 @@ float NeuralNetResPredictor::Inference(const DiffractionExperiment& experiment,
float two_theta = atanf(((2.0f * experiment.GetPixelSize_mm() / experiment.GetDetectorDistance_mm()) * tensor_output));
float stheta = sinf(two_theta * 0.5f);
float resolution = experiment.GetWavelength_A() / (2.0f * stheta);
#else
float resolution = 50.0;
#endif
return resolution;
#else
return {};
#endif
}
const std::vector<float> &NeuralNetResPredictor::GetModelInput() const {