HierarchicalWatershedSaliency.hpp

#pragma once
 
#include "HierarchySaliencyMap.hpp"
#include "../ValuedMorphologicalTree.hpp"
#include "../TreeAltitudeAlgorithms.hpp"
#include "../ValuedMorphologicalTreeView.hpp"
 
#include <algorithm>
#include <cmath>
#include <concepts>
#include <numeric>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
namespace mmcfilters {
 
class HierarchicalWatershedSaliency {
  private:
    class DisjointSet {
        std::vector<NodeId> parent_;
        std::vector<int> size_;
 
      public:
        explicit DisjointSet(int count) : parent_(static_cast<std::size_t>(count)), size_(static_cast<std::size_t>(count), 1) {
            std::iota(parent_.begin(), parent_.end(), NodeId{0});
        }
 
        [[nodiscard]] NodeId find(NodeId id) {
            NodeId root = id;
            while (parent_[static_cast<std::size_t>(root)] != root) {
                root = parent_[static_cast<std::size_t>(root)];
            }
            while (parent_[static_cast<std::size_t>(id)] != id) {
                const NodeId next = parent_[static_cast<std::size_t>(id)];
                parent_[static_cast<std::size_t>(id)] = root;
                id = next;
            }
            return root;
        }
 
        [[nodiscard]] NodeId unite(NodeId lhs, NodeId rhs) {
            lhs = find(lhs);
            rhs = find(rhs);
            if (lhs == rhs) {
                return lhs;
            }
            if (size_[static_cast<std::size_t>(lhs)] < size_[static_cast<std::size_t>(rhs)]) {
                std::swap(lhs, rhs);
            }
            parent_[static_cast<std::size_t>(rhs)] = lhs;
            size_[static_cast<std::size_t>(lhs)] += size_[static_cast<std::size_t>(rhs)];
            return lhs;
        }
    };
 
    template <AltitudeValue T> struct OrderedGraphEdge {
        NodeId source = InvalidNode;
        NodeId target = InvalidNode;
        NodeId lca = InvalidNode;
        T altitude{};
        bool finestRegionEdge = false;
        std::size_t order = 0;
    };
 
    template <AltitudeValue T, std::floating_point Real> struct PersistenceEdge {
        OrderedGraphEdge<T> graphEdge;
        Real persistence = Real{0};
    };
 
    template <AltitudeValue T> static bool edgePrecedes(const OrderedGraphEdge<T>& lhs, const OrderedGraphEdge<T>& rhs, NodeAltitudeOrder nodeAltitudeOrder) {
        if (lhs.finestRegionEdge != rhs.finestRegionEdge) {
            return lhs.finestRegionEdge;
        }
        if (lhs.altitude != rhs.altitude) {
            if (nodeAltitudeOrder == NodeAltitudeOrder::Increasing) {
                return rhs.altitude < lhs.altitude;
            }
            return lhs.altitude < rhs.altitude;
        }
        if (lhs.source != rhs.source) {
            return lhs.source < rhs.source;
        }
        return lhs.target < rhs.target;
    }
 
    template <std::floating_point Real>
    static void validateLeafExtinctions(const MorphologicalTree& tree, std::span<const Real> leafExtinction, const char* context) {
        if (leafExtinction.size() != static_cast<std::size_t>(tree.numInternalNodeSlots())) {
            throw std::invalid_argument(std::string(context) + " requires one extinction slot per internal NodeId.");
        }
        for (NodeId leaf : tree.leaves()) {
            const Real value = leafExtinction[static_cast<std::size_t>(leaf)];
            if (!std::isfinite(value) || value < Real{0}) {
                throw std::invalid_argument(std::string(context) + " requires finite non-negative leaf extinction values.");
            }
        }
    }
 
    template <AltitudeValue T>
    static std::vector<OrderedGraphEdge<T>> collectOrderedGraphEdges(const ValuedMorphologicalTreeView<T>& valuedTree, const RegularGridAdjacency2D& adjacency) {
        const MorphologicalTree& tree = valuedTree.topology();
        std::vector<OrderedGraphEdge<T>> edges;
        std::size_t order = 0;
        for (NodeId source = 0; source < tree.numPixels(); ++source) {
            const NodeId sourceSmallestNode = detail::TreeKernelAccess::smallestNodeMap(tree, source);
            for (int targetValue : detail::CommittedGridAccess::forwardNeighbors(adjacency, source)) {
                const NodeId target = static_cast<NodeId>(targetValue);
                const NodeId targetSmallestNode = detail::TreeKernelAccess::smallestNodeMap(tree, target);
                const bool finestRegionEdge = sourceSmallestNode == targetSmallestNode;
                const NodeId lca = finestRegionEdge ? sourceSmallestNode : detail::TreeKernelAccess::lowestCommonAncestor(tree, sourceSmallestNode, targetSmallestNode);
                edges.push_back(OrderedGraphEdge<T>{source, target, lca, valuedTree.nodeAltitudes()[static_cast<std::size_t>(lca)], finestRegionEdge, order++});
            }
        }
        const NodeAltitudeOrder nodeAltitudeOrder = tree.nodeAltitudeOrder();
        std::stable_sort(edges.begin(), edges.end(), [nodeAltitudeOrder](const auto& lhs, const auto& rhs) { return edgePrecedes(lhs, rhs, nodeAltitudeOrder); });
        return edges;
    }
 
    template <AltitudeValue T>
    static std::vector<OrderedGraphEdge<T>> selectMinimumSpanningTree(std::vector<OrderedGraphEdge<T>> orderedEdges, int numVertices) {
        DisjointSet components(numVertices);
        std::vector<OrderedGraphEdge<T>> mst;
        mst.reserve(static_cast<std::size_t>(std::max(0, numVertices - 1)));
        for (const OrderedGraphEdge<T>& edge : orderedEdges) {
            if (components.find(edge.source) == components.find(edge.target)) {
                continue;
            }
            static_cast<void>(components.unite(edge.source, edge.target));
            mst.push_back(edge);
            if (mst.size() == static_cast<std::size_t>(numVertices - 1)) {
                break;
            }
        }
        return mst;
    }
 
    template <AltitudeValue T, std::floating_point Real>
    static std::vector<PersistenceEdge<T, Real>> assignPersistence(const MorphologicalTree& tree, const std::vector<OrderedGraphEdge<T>>& mst,
                                                                   std::span<const Real> leafExtinction) {
        const int numVertices = tree.numPixels();
        std::vector<Real> componentExtinction(static_cast<std::size_t>(numVertices), Real{0});
        for (NodeId leaf : tree.leaves()) {
            const auto properParts = detail::TreeKernelAccess::properParts(tree, leaf);
            const auto it = properParts.begin();
            componentExtinction[static_cast<std::size_t>(*it)] = leafExtinction[static_cast<std::size_t>(leaf)];
        }
 
        DisjointSet components(numVertices);
        std::vector<PersistenceEdge<T, Real>> persistenceEdges;
        persistenceEdges.reserve(mst.size());
        for (const OrderedGraphEdge<T>& edge : mst) {
            const NodeId lhsRoot = components.find(edge.source);
            const NodeId rhsRoot = components.find(edge.target);
            const Real lhsExtinction = componentExtinction[static_cast<std::size_t>(lhsRoot)];
            const Real rhsExtinction = componentExtinction[static_cast<std::size_t>(rhsRoot)];
            persistenceEdges.push_back(PersistenceEdge<T, Real>{edge, std::min(lhsExtinction, rhsExtinction)});
            const NodeId mergedRoot = components.unite(lhsRoot, rhsRoot);
            componentExtinction[static_cast<std::size_t>(mergedRoot)] = std::max(lhsExtinction, rhsExtinction);
        }
        return persistenceEdges;
    }
 
    template <AltitudeValue T, std::floating_point Real>
    static ValuedMorphologicalTree<Real> buildPersistenceDendrogram(std::vector<PersistenceEdge<T, Real>> persistenceEdges, int rows, int columns,
                                                                      const RegularGridAdjacency2D& adjacency) {
        const int numVertices = rows * columns;
        std::stable_sort(persistenceEdges.begin(), persistenceEdges.end(), [](const auto& lhs, const auto& rhs) {
            if (lhs.persistence != rhs.persistence) {
                return lhs.persistence < rhs.persistence;
            }
            return lhs.graphEdge.order < rhs.graphEdge.order;
        });
 
        const int numNodes = numVertices == 0 ? 0 : 2 * numVertices - 1;
        std::vector<NodeId> parent(static_cast<std::size_t>(numNodes), InvalidNode);
        std::vector<NodeId> smallestNodeMap(static_cast<std::size_t>(numVertices), InvalidNode);
        std::vector<Real> altitude(static_cast<std::size_t>(numNodes), Real{0});
        for (PixelId pixel = 0; pixel < numVertices; ++pixel) {
            parent[static_cast<std::size_t>(pixel)] = pixel;
            smallestNodeMap[static_cast<std::size_t>(pixel)] = pixel;
        }
 
        DisjointSet components(numVertices);
        std::vector<NodeId> componentNode(static_cast<std::size_t>(numVertices));
        std::iota(componentNode.begin(), componentNode.end(), NodeId{0});
        NodeId nextNode = numVertices;
        for (const PersistenceEdge<T, Real>& edge : persistenceEdges) {
            const NodeId lhsRoot = components.find(edge.graphEdge.source);
            const NodeId rhsRoot = components.find(edge.graphEdge.target);
            const NodeId lhsNode = componentNode[static_cast<std::size_t>(lhsRoot)];
            const NodeId rhsNode = componentNode[static_cast<std::size_t>(rhsRoot)];
            parent[static_cast<std::size_t>(lhsNode)] = nextNode;
            parent[static_cast<std::size_t>(rhsNode)] = nextNode;
            parent[static_cast<std::size_t>(nextNode)] = nextNode;
            altitude[static_cast<std::size_t>(nextNode)] = edge.persistence;
            const NodeId mergedRoot = components.unite(lhsRoot, rhsRoot);
            componentNode[static_cast<std::size_t>(mergedRoot)] = nextNode;
            ++nextNode;
        }
 
        const NodeId root = numVertices == 1 ? NodeId{0} : nextNode - 1;
        MorphologicalTreeSemantics semantics{
            MorphologicalTreeKind::Generic, NodeAltitudeOrder::Unconstrained, SharedAdjacencyContext{adjacency}};
        return detail::materializeNativeHierarchy<Real>(makeNativeHierarchyOwned(
            std::move(parent), std::move(smallestNodeMap), std::move(altitude), root, GridDomain2D{rows, columns}, std::move(semantics)));
    }
 
  public:
    template <AltitudeValue T, std::floating_point Real>
    [[nodiscard]] static EdgeSaliencyMap<Real> compute(const ValuedMorphologicalTreeView<T>& valuedTree, std::span<const Real> leafExtinction,
                                                       const RegularGridAdjacency2D& adjacency) {
        constexpr const char* context = "HierarchicalWatershedSaliency::compute";
        valuedTree.requireTopologyUnchanged(context);
        const MorphologicalTree& tree = valuedTree.topology();
        TreeAltitudeAlgorithms::requireGlobalMonotoneAltitudeOrder(tree, context);
        validateLeafExtinctions(tree, leafExtinction, context);
        HierarchySaliencyMapValidation::validateHierarchyConnectivity(tree, adjacency, context);
 
        const int numVertices = tree.numPixels();
        auto graphEdges = collectOrderedGraphEdges(valuedTree, adjacency);
        auto mst = selectMinimumSpanningTree(std::move(graphEdges), numVertices);
        auto persistenceEdges = assignPersistence(tree, mst, leafExtinction);
        auto dendrogram = buildPersistenceDendrogram(std::move(persistenceEdges), tree.numRows(), tree.numColumns(), adjacency);
        return HierarchySaliencyMap::computeSaliencyEdgeMap(dendrogram.topology(), adjacency, dendrogram.nodeAltitudeSpan(),
                                                            HierarchyValuationPolicy::AllowLevelCollapse, HierarchyLevelConvention::EdgeSaliencyValue,
                                                            HierarchyConnectivityPolicy::AssumeConnected);
    }
 
    template <AltitudeValue T, std::floating_point Real>
    [[nodiscard]] static EdgeSaliencyMap<int> computeRanked(const ValuedMorphologicalTreeView<T>& valuedTree, std::span<const Real> leafExtinction,
                                                            const RegularGridAdjacency2D& adjacency) {
        return HierarchySaliencyMap::rankEdgeSaliencyMap(compute(valuedTree, leafExtinction, adjacency));
    }
};
 
} // namespace mmcfilters