ExtinctionValues.hpp

#pragma once
 
#include "../trees/TreeAltitudeAlgorithms.hpp"
#include "../trees/WeightedMorphologicalTree.hpp"
#include "../trees/WeightedTreeView.hpp"
#include "../trees/detail/TreeKindValidation.hpp"
#include "../utils/Image.hpp"
#include "../utils/Common.hpp"
#include "../contours/ContoursComputedIncrementally.hpp"
 
#include <algorithm>
#include <concepts>
#include <limits>
#include <memory>
#include <stack>
#include <stdexcept>
#include <string>
#include <vector>
 
namespace mmcfilters {
 
template <std::floating_point Real = float>
struct RegionalExtremaNode {
    NodeId leaf;
 
    NodeId cutoffNode;
 
    Real extinction;
 
    RegionalExtremaNode(NodeId leaf, NodeId cutoffNode, Real extinction)
        : leaf(leaf), cutoffNode(cutoffNode), extinction(extinction) {}
};
 
template<AltitudeValue T, std::floating_point Real = float>
class ExtinctionValues {
protected:
    using AltitudeView = WeightedTreeView<T>;
 
    std::vector<RegionalExtremaNode<Real>> regionalExtremaNodes;
    AltitudeView view_;
    const WeightedMorphologicalTree<T>* weighted_ = nullptr;
    const MorphologicalTree& tree;
    std::size_t treeMutationVersion_ = 0;
 
    AltitudeView view() const {
        return weighted_ != nullptr ? weighted_->asView() : view_;
    }
 
    void requireStableTree(const char* context) const {
        tree.requireMutationVersion(treeMutationVersion_, context);
    }
 
    static void requireAttributePointer(const Real* attr, const char* context) {
        if (attr == nullptr) {
            throw std::invalid_argument(std::string(context) + " requires a non-null attribute buffer.");
        }
    }
 
    static const Real* requireAttributeBuffer(const MorphologicalTree& tree, const std::vector<Real>& attr, const char* context) {
        if (attr.size() != static_cast<std::size_t>(tree.getNumInternalNodeSlots())) {
            throw std::invalid_argument(std::string(context) + " attribute size must match the internal node slot count.");
        }
        return attr.data();
    }
 
    static T altitudeOf(const AltitudeView& view, NodeId nodeId) {
        return view.getAltitude(nodeId);
    }
 
    std::vector<NodeId> collectComponentTreeExtrema() const {
        return this->tree.getLeaves();
    }
 
    static Real dominantExtremumSentinel() noexcept {
        return std::numeric_limits<Real>::max();
    }
 
    static void requireNonNegativeExtremaToKeep(int extremaToKeep, const char* context) {
        if (extremaToKeep < 0) {
            throw std::invalid_argument(std::string(context) + " requires a non-negative extremaToKeep value.");
        }
    }
 
    void initialize(const Real* attr) {
        requireAttributePointer(attr, "ExtinctionValues");
        std::vector<NodeId> leaves = collectComponentTreeExtrema();
        regionalExtremaNodes.reserve(leaves.size());
        std::vector<uint8_t> visited(this->tree.getNumInternalNodeSlots(), false);
        for (NodeId leafNodeId : leaves) {
            Real extinction = dominantExtremumSentinel();
            NodeId cutoffNodeId = leafNodeId;
            NodeId parentNodeId = this->tree.getNodeParent(cutoffNodeId);
            bool flag = true;
            while (flag && !this->tree.isRoot(cutoffNodeId)) {
                if (this->tree.getNumChildren(parentNodeId) > 1) {
                    for (NodeId sonNodeId : this->tree.getChildren(parentNodeId)) {
                        if (flag) {
                            if (visited[sonNodeId] && sonNodeId != cutoffNodeId && attr[sonNodeId] == attr[cutoffNodeId]) {
                                flag = false;
                            } else if (sonNodeId != cutoffNodeId && attr[sonNodeId] > attr[cutoffNodeId]) {
                                flag = false;
                            }
                            visited[sonNodeId] = true;
                        }
                    }
                }
                if (flag) {
                    cutoffNodeId = parentNodeId;
                    parentNodeId = this->tree.getNodeParent(cutoffNodeId);
                }
            }
            if (!this->tree.isRoot(cutoffNodeId)) {
                extinction = attr[cutoffNodeId];
            }
            regionalExtremaNodes.emplace_back(leafNodeId, cutoffNodeId, extinction);
        }
 
        std::sort(regionalExtremaNodes.begin(), regionalExtremaNodes.end(), [](const auto& a, const auto& b) {
            if (a.extinction != b.extinction) {
                return a.extinction > b.extinction;
            }
            if (a.cutoffNode != b.cutoffNode) {
                return a.cutoffNode < b.cutoffNode;
            }
            return a.leaf < b.leaf;
        });
    }
 
public:
    ExtinctionValues(const AltitudeView& view, const std::shared_ptr<Real[]>& attr)
        : ExtinctionValues(view, attr.get()) {}
 
    ExtinctionValues(const AltitudeView& view, const std::vector<Real>& attr)
        : ExtinctionValues(view, requireAttributeBuffer(view.topology(), attr, "ExtinctionValues")) {}
 
    ExtinctionValues(const AltitudeView& view, const Real* attr)
        : view_(view),
          tree(view_.topology()),
          treeMutationVersion_(tree.getMutationVersion()) {
        view_.requireTopologyUnchanged("ExtinctionValues");
        detail::validateComponentTreeKind(this->tree, "ExtinctionValues");
        initialize(attr);
    }
 
    ExtinctionValues(const WeightedMorphologicalTree<T>& weighted, const std::shared_ptr<Real[]>& attr)
        : ExtinctionValues(weighted.asView(), attr.get()) {
        weighted_ = &weighted;
    }
 
    ExtinctionValues(const WeightedMorphologicalTree<T>& weighted, const std::vector<Real>& attr)
        : ExtinctionValues(weighted.asView(), attr) {
        weighted_ = &weighted;
    }
 
    ExtinctionValues(const WeightedMorphologicalTree<T>& weighted, const Real* attr)
        : ExtinctionValues(weighted.asView(), attr) {
        weighted_ = &weighted;
    }
 
    [[nodiscard]] ImagePtr<Real> saliencyMap(int extremaToKeep, bool unweighted = true) {
        requireStableTree("ExtinctionValues::saliencyMap");
        requireNonNegativeExtremaToKeep(extremaToKeep, "ExtinctionValues::saliencyMap");
        std::vector<uint8_t> keep(tree.getNumInternalNodeSlots(), false);
        std::vector<Real> extinctionByNode(tree.getNumInternalNodeSlots(), Real{0});
        std::vector<NodeId> keptNodes;
        const int leafToKeep = std::min(extremaToKeep, static_cast<int>(regionalExtremaNodes.size()));
        for (int i = 0; i < leafToKeep; ++i) {
            const NodeId cutoffNode = this->regionalExtremaNodes[i].cutoffNode;
            if (!keep[cutoffNode]) {
                keptNodes.push_back(cutoffNode);
            }
            keep[cutoffNode] = true;
            extinctionByNode[cutoffNode] = unweighted ? static_cast<Real>(leafToKeep - i) : this->regionalExtremaNodes[i].extinction;
        }
 
        ImagePtr<Real> imgOutputPtr = Image<Real>::create(tree.getNumRowsOfImage(), tree.getNumColsOfImage(), Real{0});
        auto saliencyOutput = imgOutputPtr->rawData();
 
        auto contours = ContoursComputedIncrementally::extractCompactContours(tree);
        for (NodeId node : keptNodes) {
            for (int p : contours.getContour(node)) {
                saliencyOutput[p] = extinctionByNode[node];
            }
        }
 
        return imgOutputPtr;
    }
 
    [[nodiscard]] ImagePtr<T> filtering(int extremaToKeep) {
        requireStableTree("ExtinctionValues::filtering");
        requireNonNegativeExtremaToKeep(extremaToKeep, "ExtinctionValues::filtering");
        const AltitudeView altitudeView = view();
        std::vector<uint8_t> criterion(tree.getNumInternalNodeSlots(), false);
        const int leafToKeep = std::min(extremaToKeep, static_cast<int>(regionalExtremaNodes.size()));
        for (int i = 0; i < leafToKeep; i++) {
            criterion[regionalExtremaNodes[i].leaf] = true;
        }
        for (NodeId nodeId : tree.getPostOrderNodes()) {
            if (!tree.isRoot(nodeId) && criterion[nodeId]) {
                criterion[tree.getNodeParent(nodeId)] = true;
            }
        }
 
        ImagePtr<T> imgOutputPtr = Image<T>::create(tree.getNumRowsOfImage(), tree.getNumColsOfImage(), T{});
        auto imgOutput = imgOutputPtr->rawData();
        std::stack<NodeId> stack;
        stack.push(tree.getRoot());
        while (!stack.empty()) {
            const NodeId nodeId = stack.top();
            stack.pop();
            const T level = altitudeOf(altitudeView, nodeId);
            for (int pixel : tree.getProperParts(nodeId)) {
                imgOutput[pixel] = level;
            }
            for (NodeId childNodeId : tree.getChildren(nodeId)) {
                if (criterion[childNodeId]) {
                    stack.push(childNodeId);
                } else {
                    for (NodeId subtreeNodeId : tree.getNodeSubtree(childNodeId)) {
                        for (int pixel : tree.getProperParts(subtreeNodeId)) {
                            imgOutput[pixel] = level;
                        }
                    }
                }
            }
        }
        return imgOutputPtr;
    }
 
    std::vector<RegionalExtremaNode<Real>>& getExtinctionValues() {
        requireStableTree("ExtinctionValues::getExtinctionValues");
        return regionalExtremaNodes;
    }
};
 
 
} // namespace mmcfilters