72 lines
2.4 KiB
C++
72 lines
2.4 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(std::vector<HDF5DataFileStatistics> &v) {
|
|
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();
|
|
}
|
|
writer.GetStatistics(v);
|
|
}
|
|
|
|
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();
|
|
}
|
|
|
|
StreamWriterStatistics StreamWriter::Run() {
|
|
StreamWriterStatistics ret;
|
|
StartDataCollection();
|
|
try {
|
|
CollectImages(ret.data_file_stats);
|
|
} 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();
|
|
|
|
ret.image_puller_stats = image_puller.GetStatistics();
|
|
logger.Info("Write task done. Images = {} Throughput = {:.0f} MB/s Frame rate = {:.0f} Hz",
|
|
ret.image_puller_stats.processed_images, ret.image_puller_stats.performance_MBs, ret.image_puller_stats.performance_Hz);
|
|
return ret;
|
|
}
|