66 lines
2.4 KiB
C++
66 lines
2.4 KiB
C++
// 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"
|
|
#include "../common/PixelMask.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 PixelMask& mask,
|
|
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 PixelMask& mask, const int16_t* image, Quarter q);
|
|
std::vector<float> Prepare(const DiffractionExperiment& experiment, const PixelMask& mask, const int32_t* image, Quarter q);
|
|
std::vector<float> Prepare(const DiffractionExperiment& experiment, const PixelMask& mask, const int8_t* image, Quarter q);
|
|
|
|
std::optional<float> Inference(const DiffractionExperiment& experiment, const PixelMask& mask, const void* image, int nquads);
|
|
std::optional<float> Inference(const DiffractionExperiment& experiment, const PixelMask& mask, const void* image, Quarter q);
|
|
};
|
|
|
|
|
|
#endif //JUNGFRAUJOCH_NEURALNETRESPREDICTOR_H
|