Improvements before MAX IV test
This commit is contained in:
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user