v1.0.0-rc.38

image_analysis/NeuralNetInferenceClient.h

@@ -0,0 +1,61 @@
+// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#ifndef JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
+#define JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
+
+#include <vector>
+#include <mutex>
+#include <condition_variable>
+
+#include "../common/Logger.h"
+#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.
+
+enum class Quarter {TopLeft, TopRight, BottomLeft, BottomRight};
+
+class NeuralNetInferenceClient {
+    struct PredictorAddr {
+        std::string hostname;
+        uint16_t port;
+        bool busy = false;
+    };
+
+    bool enable = false;
+
+    std::vector<PredictorAddr> hosts;
+
+    std::mutex m;
+    std::condition_variable c;
+
+    Logger *logger = nullptr;
+
+    std::optional<float> Run(const std::vector<float> &input);
+
+    std::optional<size_t> GetFreeNode();
+    std::optional<size_t> AcquireNode();
+    void ReturnNode(size_t node);
+
+    template<class T>
+    std::vector<float> PrepareInternal(const DiffractionExperiment& experiment, const T* image, Quarter q);
+public:
+    void AddHost(std::string hostname, uint16_t port);
+    void AddHost(std::string addr);
+    size_t GetHostCount();
+    void AddLogger(Logger *logger);
+
+    size_t GetMaxPoolFactor(const DiffractionExperiment& experiment) const;
+    std::vector<float> Prepare(const DiffractionExperiment& experiment, const int16_t* image, Quarter q);
+    std::vector<float> Prepare(const DiffractionExperiment& experiment, const int32_t* image, Quarter q);
+    std::vector<float> Prepare(const DiffractionExperiment& experiment, const int8_t* image, Quarter q);
+
+    std::optional<float> Inference(const DiffractionExperiment& experiment, const void* image, int nquads);
+    std::optional<float> Inference(const DiffractionExperiment& experiment, const void* image, Quarter q);
+};
+
+
+#endif //JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H