297 lines
10 KiB
C++
297 lines
10 KiB
C++
// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include <cmath>
|
|
#include <utility>
|
|
#include "NeuralNetInferenceClient.h"
|
|
|
|
#include "../common/JFJochException.h"
|
|
#include "httplib/httplib.h"
|
|
|
|
void NeuralNetInferenceClient::AddLogger(Logger *in_logger) {
|
|
logger = in_logger;
|
|
}
|
|
|
|
size_t NeuralNetInferenceClient::GetHostCount() {
|
|
std::unique_lock ul(m);
|
|
return hosts.size();
|
|
}
|
|
|
|
void NeuralNetInferenceClient::AddHost(std::string hostname, uint16_t port) {
|
|
std::unique_lock ul(m);
|
|
hosts.emplace_back(PredictorAddr{.hostname = hostname, .port = port, .busy = false});
|
|
enable = true;
|
|
}
|
|
|
|
void NeuralNetInferenceClient::AddHost(std::string addr) {
|
|
if (addr.empty()) {
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Host address cannot be empty");
|
|
}
|
|
|
|
size_t pos = addr.find(':');
|
|
std::string hostname;
|
|
uint16_t port = 8000; // Default port
|
|
|
|
if (pos != std::string::npos) {
|
|
hostname = addr.substr(0, pos);
|
|
if (hostname.empty()) {
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Hostname cannot be empty");
|
|
}
|
|
|
|
std::string portStr = addr.substr(pos + 1);
|
|
try {
|
|
int portValue = std::stoi(portStr);
|
|
if (portValue < 0 || portValue > 65535) {
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"Port must be in range 0-65535, got: " + portStr);
|
|
}
|
|
port = static_cast<uint16_t>(portValue);
|
|
} catch (const std::invalid_argument&) {
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"Invalid port number format: " + portStr);
|
|
} catch (const std::out_of_range&) {
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"Port number out of range: " + portStr);
|
|
}
|
|
} else {
|
|
hostname = addr;
|
|
if (hostname.empty()) {
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Hostname cannot be empty");
|
|
}
|
|
}
|
|
|
|
AddHost(hostname, port);
|
|
}
|
|
|
|
|
|
template<class T>
|
|
std::vector<float> NeuralNetInferenceClient::PrepareInternal(const DiffractionExperiment& experiment,
|
|
const PixelMask& mask,
|
|
const T* image,
|
|
Quarter q) {
|
|
std::vector<float> ret(512*512);
|
|
|
|
int64_t pool_factor = GetMaxPoolFactor(experiment);
|
|
int64_t xpixel = experiment.GetXPixelsNum();
|
|
int64_t ypixel = experiment.GetYPixelsNum();
|
|
int64_t start_x = std::lround(experiment.GetBeamX_pxl());
|
|
int64_t start_y = std::lround(experiment.GetBeamY_pxl());
|
|
|
|
for (int64_t y = 0; y < 512; y++) {
|
|
int64_t y0;
|
|
|
|
if ((q == Quarter::BottomLeft) || (q == Quarter::BottomRight))
|
|
y0 = start_y + y * pool_factor;
|
|
else
|
|
y0 = start_y - (y + 1) * pool_factor + 1;
|
|
|
|
int64_t max_yp = std::min(y0 + pool_factor, ypixel);
|
|
if (y0 < 0)
|
|
y0 = 0;
|
|
|
|
for (int64_t x = 0; x < 512; x++) {
|
|
int64_t x0;
|
|
|
|
if ((q == Quarter::TopRight) || (q == Quarter::BottomRight))
|
|
x0 = start_x + x * pool_factor;
|
|
else
|
|
x0 = start_x - (x + 1) * pool_factor + 1;
|
|
|
|
int64_t max_xp = std::min(x0 + pool_factor, xpixel);
|
|
if (x0 < 0)
|
|
x0 = 0;
|
|
|
|
int64_t val = 0.0;
|
|
for (int64_t yp = y0; yp < max_yp; yp++) {
|
|
for (int64_t xp = x0; xp < max_xp; xp++) {
|
|
int64_t pxl = image[yp * xpixel + xp];
|
|
if (mask.GetMask().at(yp * xpixel + xp) != 0)
|
|
pxl = INT64_MAX;
|
|
else if (pxl > INT16_MAX)
|
|
pxl = INT16_MAX;
|
|
|
|
if (pxl > val)
|
|
val = pxl;
|
|
}
|
|
}
|
|
|
|
if (val == INT64_MAX)
|
|
ret[512 * y + x] = 0;
|
|
else
|
|
ret[512 * y + x] = floor(sqrt(static_cast<double>(val)));
|
|
}
|
|
}
|
|
|
|
return ret;
|
|
}
|
|
|
|
std::vector<float> NeuralNetInferenceClient::Prepare(const DiffractionExperiment& experiment,
|
|
const PixelMask& mask,
|
|
const int16_t *image, Quarter q) {
|
|
return PrepareInternal(experiment, mask, image, q);
|
|
}
|
|
|
|
std::vector<float> NeuralNetInferenceClient::Prepare(const DiffractionExperiment& experiment,
|
|
const PixelMask& mask,
|
|
const int32_t *image, Quarter q) {
|
|
return PrepareInternal(experiment, mask, image, q);
|
|
}
|
|
|
|
std::vector<float> NeuralNetInferenceClient::Prepare(const DiffractionExperiment& experiment,
|
|
const PixelMask& mask,
|
|
const int8_t *image, Quarter q) {
|
|
return PrepareInternal(experiment, mask, image, q);
|
|
}
|
|
|
|
size_t NeuralNetInferenceClient::GetMaxPoolFactor(const DiffractionExperiment& experiment) const {
|
|
float max_direction = std::max(experiment.GetXPixelsNum(), experiment.GetYPixelsNum()) / 2.0;
|
|
|
|
size_t pool_factor = std::ceil(max_direction / 512.0f - 0.25);
|
|
|
|
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<size_t> NeuralNetInferenceClient::GetFreeNode() {
|
|
//Assumes m is locked!
|
|
|
|
for (int i = 0; i < hosts.size(); i++) {
|
|
if (!hosts[i].busy) {
|
|
hosts[i].busy = true;
|
|
return i;
|
|
}
|
|
}
|
|
return {};
|
|
}
|
|
std::optional<size_t> NeuralNetInferenceClient::AcquireNode() {
|
|
std::unique_lock ul(m);
|
|
|
|
if (hosts.empty())
|
|
return {};
|
|
|
|
auto ret = GetFreeNode();
|
|
if (ret)
|
|
return ret;
|
|
|
|
// No free node available, wait for one with a timeout
|
|
c.wait_for(ul, std::chrono::milliseconds(200),
|
|
[this]() { return GetFreeNode().has_value(); });
|
|
|
|
return GetFreeNode();
|
|
}
|
|
|
|
void NeuralNetInferenceClient::ReturnNode(size_t node) {
|
|
std::unique_lock ul(m);
|
|
hosts[node].busy = false;
|
|
c.notify_one();
|
|
}
|
|
|
|
std::optional<float> NeuralNetInferenceClient::Run(const std::vector<float> &input) {
|
|
auto h = AcquireNode();
|
|
|
|
if (!h.has_value())
|
|
return {};
|
|
|
|
if (logger)
|
|
logger->Debug("Using host {:d}", h.value());
|
|
|
|
httplib::Result resp;
|
|
try {
|
|
httplib::Client cli(hosts[h.value()].hostname, hosts[h.value()].port);
|
|
cli.set_read_timeout(1);
|
|
cli.set_write_timeout(1);
|
|
resp = cli.Post("/infer", (char *) input.data(), input.size() * sizeof(float), "application/octet-stream");
|
|
} catch (...) {
|
|
ReturnNode(h.value());
|
|
}
|
|
ReturnNode(h.value());
|
|
|
|
if (resp && resp->status == httplib::StatusCode::OK_200) {
|
|
auto j = nlohmann::json::parse(resp->body);
|
|
if (j.contains("result") && j["result"].is_number_float())
|
|
return j["result"].get<float>();
|
|
}
|
|
return {};
|
|
|
|
}
|
|
|
|
std::optional<float>
|
|
NeuralNetInferenceClient::Inference(const DiffractionExperiment &experiment,
|
|
const PixelMask& mask,
|
|
const void *image,
|
|
int nquads) {
|
|
if (!enable)
|
|
return {};
|
|
|
|
std::optional<float> quad[4];
|
|
|
|
if (nquads >= 1)
|
|
quad[0] = Inference(experiment, mask, image, Quarter::BottomRight);
|
|
if (nquads >= 2)
|
|
quad[1] = Inference(experiment, mask, image, Quarter::BottomRight);
|
|
if (nquads >= 3)
|
|
quad[2] = Inference(experiment, mask, image, Quarter::BottomRight);
|
|
if (nquads >= 4)
|
|
quad[3] = Inference(experiment, mask, image, Quarter::BottomLeft);
|
|
|
|
int count = 0;
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < 4; i++) {
|
|
if (quad[i]) {
|
|
count++;
|
|
sum += quad[i].value();
|
|
}
|
|
}
|
|
if (count == 0)
|
|
return {};
|
|
return sum / count;
|
|
}
|
|
|
|
std::optional<float> NeuralNetInferenceClient::Inference(const DiffractionExperiment& experiment,
|
|
const PixelMask& mask,
|
|
const void *image,
|
|
Quarter q) {
|
|
if (!enable)
|
|
return {};
|
|
|
|
std::vector<float> v;
|
|
|
|
switch (experiment.GetByteDepthImage()) {
|
|
case 1:
|
|
if (experiment.IsPixelSigned())
|
|
v = PrepareInternal(experiment, mask, (int8_t *) image, q);
|
|
else
|
|
v = PrepareInternal(experiment, mask, (uint8_t *) image, q);
|
|
break;
|
|
case 2:
|
|
if (experiment.IsPixelSigned())
|
|
v = PrepareInternal(experiment, mask, (int16_t *) image, q);
|
|
else
|
|
v = PrepareInternal(experiment, mask, (uint16_t *) image, q);
|
|
break;
|
|
case 4:
|
|
if (experiment.IsPixelSigned())
|
|
v = PrepareInternal(experiment, mask, (int32_t *) image, q);
|
|
else
|
|
v = PrepareInternal(experiment, mask, (uint32_t *) image, q);
|
|
break;
|
|
default:
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Bit depth not supported");
|
|
}
|
|
auto resp = Run(v);
|
|
|
|
if (!resp.has_value())
|
|
return {};
|
|
|
|
float two_theta = atanf(
|
|
((2.0f * GetMaxPoolFactor(experiment) * experiment.GetPixelSize_mm()
|
|
/ experiment.GetDetectorDistance_mm()) * resp.value()));
|
|
float stheta = sinf(two_theta * 0.5f);
|
|
float resolution = experiment.GetWavelength_A() / (2.0f * stheta);
|
|
return resolution;
|
|
}
|