v1.0.0-rc.36

image_analysis/NeuralNetResPredictor.cpp

@@ -0,0 +1,115 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include <cmath>
+#include "NeuralNetResPredictor.h"
+
+#include "../common/JFJochException.h"
+
+NeuralNetResPredictor::NeuralNetResPredictor(const std::string& model_path)
+        : model_input(512*512),
+          enable(!model_path.empty())
+#ifdef JFJOCH_USE_TORCH
+        , device(torch::kCUDA)
+#endif
+{
+#ifdef JFJOCH_USE_TORCH
+    if (enable) {
+        module = torch::jit::load(model_path);
+        module.to(device);
+    }
+#else
+    enable = false;
+#endif
+}
+
+template<class T>
+void NeuralNetResPredictor::PrepareInternal(const DiffractionExperiment& experiment, const T* image) {
+    size_t pool_factor = GetMaxPoolFactor(experiment);
+    size_t xpixel = experiment.GetXPixelsNum();
+    size_t ypixel = experiment.GetYPixelsNum();
+    size_t start_x = std::lround(experiment.GetBeamX_pxl());
+    size_t start_y = std::lround(experiment.GetBeamY_pxl());
+
+    for (size_t y = 0; y < 512; y++) {
+        size_t y0 = y * pool_factor + start_y;
+        size_t min_yp = std::min(y0 + pool_factor, ypixel);
+        for (size_t x = 0; x < 512; x++) {
+            float val = 0.0;
+            size_t x0 = x * pool_factor + start_x;
+            size_t min_xp = std::min(x0 + pool_factor, ypixel);
+
+            for (size_t yp = y0; yp < min_yp; yp++) {
+                for (size_t xp = x0; xp < min_xp; xp++) {
+                    int16_t pxl = image[yp * xpixel + xp];
+                    if (val < pxl)
+                        val = pxl;
+                }
+            }
+            float max_pool = floorf(sqrtf(val));
+            model_input[512 * y + x] = max_pool;
+        }
+    }
+}
+
+void NeuralNetResPredictor::Prepare(const DiffractionExperiment& experiment, const int16_t *image) {
+    PrepareInternal(experiment, image);
+}
+
+void NeuralNetResPredictor::Prepare(const DiffractionExperiment& experiment, const int32_t *image) {
+    PrepareInternal(experiment, image);
+}
+
+void NeuralNetResPredictor::Prepare(const DiffractionExperiment& experiment, const int8_t *image) {
+    PrepareInternal(experiment, image);
+}
+
+size_t NeuralNetResPredictor::GetMaxPoolFactor(const DiffractionExperiment& experiment) const {
+    float max_direction = std::max(experiment.GetXPixelsNum(), experiment.GetYPixelsNum()) / 2.0;
+    size_t pool_factor = std::lround(max_direction / 512.0f);
+    if (pool_factor <= 0)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector size is too small");
+
+    if (pool_factor > 8)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector size is too large");
+    return pool_factor;
+}
+
+std::optional<float> NeuralNetResPredictor::Inference(const DiffractionExperiment& experiment, const void *image) {
+    if (!enable)
+        return {};
+#ifdef JFJOCH_USE_TORCH
+    switch (experiment.GetByteDepthImage()) {
+        case 1:
+            Prepare(experiment, (int8_t *) image);
+            break;
+        case 2:
+            Prepare(experiment, (int16_t *) image);
+            break;
+        case 4:
+            Prepare(experiment, (int32_t *) image);
+            break;
+        default:
+            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Bit depth not supported");
+    }
+
+    auto options = torch::TensorOptions().dtype(at::kFloat);
+    auto model_input_tensor = torch::from_blob(model_input.data(), {1,1,512,512}, options).to(device);
+
+    std::vector<torch::jit::IValue> pixels;
+    pixels.emplace_back(std::move(model_input_tensor));
+
+    auto output = module.forward(pixels).toTensor();
+    auto tensor_output = output[0].item<float>();
+    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);
+    return resolution;
+#else
+    return {};
+#endif
+}
+
+const std::vector<float> &NeuralNetResPredictor::GetModelInput() const {
+    return model_input;
+}