80 lines
2.8 KiB
C++
80 lines
2.8 KiB
C++
// Copyright (2019-2022) Paul Scherrer Institute
|
|
// SPDX-License-Identifier: GPL-3.0-or-later
|
|
|
|
#include "../common/JFJochException.h"
|
|
#include "StreamWriter.h"
|
|
#include "HDF5Writer.h"
|
|
#include "HDF5NXmx.h"
|
|
#include "MakeDirectory.h"
|
|
|
|
StreamWriter::StreamWriter(ZMQContext &context, Logger &in_logger, const std::string &zmq_addr)
|
|
: image_puller(context), logger(in_logger) {
|
|
image_puller.Connect(zmq_addr);
|
|
logger.Info("Connected via ZMQ to {}", zmq_addr);
|
|
}
|
|
|
|
void StreamWriter::StartDataCollection() {
|
|
image_puller.WaitForImage();
|
|
while (image_puller.GetFrameType() != JFJochFrameDeserializer::Type::START) {
|
|
logger.Error("Expected START image");
|
|
image_puller.WaitForImage();
|
|
}
|
|
start_message = image_puller.GetStartMessage();
|
|
logger.Info("Starting writing for dataset {} of {} images", start_message.file_prefix,
|
|
start_message.number_of_images);
|
|
|
|
MakeDirectory(start_message.file_prefix);
|
|
}
|
|
|
|
void StreamWriter::CollectImages() {
|
|
HDF5Writer writer(start_message);
|
|
image_puller.WaitForImage();
|
|
while (image_puller.GetFrameType() == JFJochFrameDeserializer::Type::IMAGE) {
|
|
auto image_array = image_puller.GetDataMessage();
|
|
writer.Write(image_array);
|
|
image_puller.WaitForImage();
|
|
}
|
|
}
|
|
|
|
void StreamWriter::EndDataCollection() {
|
|
while (image_puller.GetFrameType() != JFJochFrameDeserializer::Type::END) {
|
|
logger.Error("Expected END image");
|
|
image_puller.WaitForImage();
|
|
}
|
|
EndMessage end_message = image_puller.GetEndMessage();
|
|
if (end_message.write_master_file)
|
|
HDF5Metadata::NXmx(start_message, end_message);
|
|
}
|
|
|
|
|
|
void StreamWriter::Abort() {
|
|
image_puller.Abort();
|
|
}
|
|
|
|
ZMQImagePullerStatistics StreamWriter::Run() {
|
|
StartDataCollection();
|
|
try {
|
|
CollectImages();
|
|
} catch (JFJochException &e) {
|
|
// Error during collecting images will skip to end data collection
|
|
// End data collection will consume all images till the end
|
|
logger.ErrorException(e);
|
|
}
|
|
EndDataCollection();
|
|
|
|
auto puller_stats = image_puller.GetStatistics();
|
|
logger.Info("Write task done. Images = {} Throughput = {:.0f} MB/s Frame rate = {:.0f} Hz",
|
|
puller_stats.processed_images, puller_stats.performance_MBs, puller_stats.performance_Hz);
|
|
return puller_stats;
|
|
}
|
|
|
|
void StreamWriter::BaseDirectory(const std::string &input) {
|
|
logger.Info("Setting base directory to " + input);
|
|
try {
|
|
std::filesystem::current_path(input);
|
|
} catch (const std::filesystem::filesystem_error& err) {
|
|
logger.Error("Cannot set base directory: " + std::string(err.what()));
|
|
throw JFJochException(JFJochExceptionCategory::FileWriteError,
|
|
"Cannot set base directory " + std::string(err.what()));
|
|
}
|
|
} |