roibin_split_stage.h

Go to the documentation of this file.

#pragma once
 
#include "stage/stage.h"
#include "fzm_format.h"
#include "backend/types.h"
#include <array>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace fz {
 
struct RoiPeak {
    uint32_t z;   
    uint16_t x;   
    uint16_t y;   
};
static_assert(sizeof(RoiPeak) == 8, "RoiPeak must be 8 bytes to match the .roi format");
 
template <typename TData>
class ROIBinSplitStage : public Stage {
public:
    ROIBinSplitStage() = default;
    ~ROIBinSplitStage() override = default;
 
    // ── Configuration ───────────────────────────────────────────────────────
 
    void setPeaksFile(const std::string& path);
 
    void setPeaks(const std::vector<RoiPeak>& peaks);
 
    void setRoiHalfWidth(uint32_t hw) { half_width_ = hw; }
    uint32_t getRoiHalfWidth() const { return half_width_; }
 
    void setBinFactor(uint32_t b) {
        if (b == 0) throw std::runtime_error("ROIBinSplit: bin_factor must be >= 1");
        bin_ = b;
    }
    uint32_t getBinFactor() const { return bin_; }
 
    void setDims(const std::array<size_t, 3>& dims) override {
        if (dims_from_header_ || dims[0] == 0) return;
        dims_ = dims;
    }
    void setDims(size_t x, size_t y = 1, size_t z = 1) {
        setDims(std::array<size_t, 3>{x, y, z});
    }
 
    size_t getNumPeaks() const { return peaks_.size(); }
    size_t getBoxArea()  const { return size_t(2 * half_width_ + 1) * (2 * half_width_ + 1); }
    size_t getBgNx()     const { return (dims_[0] + bin_ - 1) / bin_; }
    size_t getBgNy()     const { return (dims_[1] + bin_ - 1) / bin_; }
    size_t getBgCount()  const { return getBgNx() * getBgNy() * dims_[2]; }
    size_t getRoiCount() const { return peaks_.size() * getBoxArea(); }
 
    double getRoiOverlapFraction() const { return overlap_frac_; }
 
    // ── Stage control ───────────────────────────────────────────────────────
    void setInverse(bool inv) override { is_inverse_ = inv; }
    bool isInverse() const override    { return is_inverse_; }
 
    bool isGraphCompatible() const override { return true; }
 
    // ── Port model ──────────────────────────────────────────────────────────
    size_t getNumInputs()  const override { return is_inverse_ ? 3 : 1; }
    size_t getNumOutputs() const override { return is_inverse_ ? 1 : 3; }
 
    std::vector<std::string> getOutputNames() const override {
        return is_inverse_ ? std::vector<std::string>{"output"}
                           : std::vector<std::string>{"roi", "bg", "peaks"};
    }
 
    std::string getName() const override { return "ROIBinSplit"; }
 
    uint16_t getStageTypeId() const override {
        return static_cast<uint16_t>(StageType::ROIBIN_SPLIT);
    }
 
    uint8_t getOutputDataType(size_t output_index) const override {
        if (is_inverse_) return elemType();
        switch (output_index) {
            case 0: return elemType();                             // roi
            case 1: return elemType();                             // bg
            default: return static_cast<uint8_t>(DataType::UINT8);  // peaks
        }
    }
    uint8_t getInputDataType(size_t input_index) const override {
        if (!is_inverse_) return elemType();
        switch (input_index) {
            case 0: return elemType();
            case 1: return elemType();
            default: return static_cast<uint8_t>(DataType::UINT8);
        }
    }
 
    // ── Execution ───────────────────────────────────────────────────────────
    void execute(
        fz::stream_t stream,
        MemoryPool* pool,
        const std::vector<void*>& inputs,
        const std::vector<void*>& outputs,
        const std::vector<size_t>& sizes
    ) override;
 
    void onFinalize(size_t estimated_inlen, MemoryPool* pool) override;
 
    size_t estimateDeviceFootprintBytes(size_t /*inlen*/) const override {
        return peaks_.size() * sizeof(RoiPeak);
    }
 
    // ── Size estimation ─────────────────────────────────────────────────────
    std::vector<size_t> estimateOutputSizes(
        const std::vector<size_t>& input_sizes
    ) const override {
        if (is_inverse_) {
            return {dims_[0] * dims_[1] * dims_[2] * sizeof(TData)};
        }
        (void)input_sizes;
        return {getRoiCount() * sizeof(TData),
                getBgCount()  * sizeof(TData),
                peaks_.size() * sizeof(RoiPeak)};
    }
 
    std::unordered_map<std::string, size_t>
    getActualOutputSizesByName() const override {
        if (is_inverse_)
            return {{"output", dims_[0] * dims_[1] * dims_[2] * sizeof(TData)}};
        return {{"roi",   getRoiCount() * sizeof(TData)},
                {"bg",    getBgCount()  * sizeof(TData)},
                {"peaks", peaks_.size() * sizeof(RoiPeak)}};
    }
 
    size_t getActualOutputSize(int index) const override {
        if (is_inverse_) return (index == 0)
            ? dims_[0] * dims_[1] * dims_[2] * sizeof(TData) : 0;
        switch (index) {
            case 0:  return getRoiCount() * sizeof(TData);
            case 1:  return getBgCount()  * sizeof(TData);
            case 2:  return peaks_.size() * sizeof(RoiPeak);
            default: return 0;
        }
    }
 
    std::vector<std::string> getRunNotes() const override;
 
    // ── Serialization ───────────────────────────────────────────────────────
    size_t serializeHeader(size_t, uint8_t* buf, size_t max_size) const override {
        if (max_size < kHeaderSize) return 0;
        size_t off = 0;
        auto put32 = [&](uint32_t v) { std::memcpy(buf + off, &v, 4); off += 4; };
        auto put16 = [&](uint16_t v) { std::memcpy(buf + off, &v, 2); off += 2; };
        put32(static_cast<uint32_t>(dims_[0]));
        put32(static_cast<uint32_t>(dims_[1]));
        put32(static_cast<uint32_t>(dims_[2]));
        put32(static_cast<uint32_t>(peaks_.size()));
        put16(static_cast<uint16_t>(half_width_));
        put16(static_cast<uint16_t>(bin_));
        buf[off++] = elemType();
        buf[off++] = 0;  // reserved
        return off;
    }
 
    void deserializeHeader(const uint8_t* buf, size_t size) override {
        if (size < kHeaderSize) return;
        size_t off = 0;
        auto get32 = [&]() { uint32_t v; std::memcpy(&v, buf + off, 4); off += 4; return v; };
        auto get16 = [&]() { uint16_t v; std::memcpy(&v, buf + off, 2); off += 2; return v; };
        dims_[0] = get32(); dims_[1] = get32(); dims_[2] = get32();
        dims_from_header_ = true;
        const uint32_t npeaks = get32();
        half_width_ = get16();
        bin_        = get16();
        // The peak *values* arrive on the `peaks` input port, not in this header;
        // only the count is needed here so the sizes line up before execute().
        peaks_.assign(npeaks, RoiPeak{0, 0, 0});
    }
 
    size_t getMaxHeaderSize(size_t) const override { return kHeaderSize; }
 
    void saveState() override {
        saved_dims_ = dims_; saved_hw_ = half_width_;
        saved_bin_ = bin_;   saved_npeaks_ = peaks_.size();
        state_saved_ = true;
    }
    void restoreState() override {
        if (!state_saved_) return;
        dims_ = saved_dims_; half_width_ = saved_hw_; bin_ = saved_bin_;
        if (peaks_.size() != saved_npeaks_) peaks_.resize(saved_npeaks_);
    }
 
private:
    // 4 x uint32 (dims, npeaks) + 2 x uint16 (hw, bin) + dtype + reserved.
    static constexpr size_t kHeaderSize = 22;
 
    static constexpr uint8_t elemType() {
        return static_cast<uint8_t>(sizeof(TData) == 4 ? DataType::FLOAT32
                                                       : DataType::FLOAT64);
    }
 
    void computeOverlapFraction();
 
    bool  is_inverse_ = false;
    bool  dims_from_header_ = false;  
    bool  state_saved_      = false;  
    std::array<size_t, 3> dims_ = {0, 0, 1};
    uint32_t half_width_ = 4;
    uint32_t bin_        = 1;
    double   overlap_frac_ = 0.0;
 
    std::vector<RoiPeak> peaks_;
    RoiPeak* d_peaks_ = nullptr;   
 
    std::array<size_t, 3> saved_dims_ = {0, 0, 1};
    uint32_t saved_hw_  = 4;
    uint32_t saved_bin_ = 1;
    size_t   saved_npeaks_ = 0;
};
 
extern template class ROIBinSplitStage<float>;
extern template class ROIBinSplitStage<double>;
 
} // namespace fz