CasfComponentTrees.hpp

#pragma once
 
#include "DualMinMaxTreeIncrementalFilter.hpp"
#include "../MorphologicalTreeFactory.hpp"
#include "../../utils/Image.hpp"
 
#include <memory>
#include <stdexcept>
#include <utility>
#include <vector>
 
namespace mmcfilters::adjust {
 
enum class CasfComponentTreesAttribute { Area, BoundingBoxWidth, BoundingBoxHeight, BoundingBoxDiagonal };
 
template <AltitudeValue T> class CasfComponentTrees {
  private:
    using tree_t = ValuedMorphologicalTree<T>;
    using image_ptr_t = ImagePtr<T>;
    using attribute_computer_t = DynamicTreeAttributeComputer<T>;
 
    RegularGridAdjacency2D adjacency_;
    tree_t minTree_;
    tree_t maxTree_;
    std::unique_ptr<attribute_computer_t> minAttributeComputer_;
    std::unique_ptr<attribute_computer_t> maxAttributeComputer_;
    std::unique_ptr<DualMinMaxTreeIncrementalFilter<T>> adjust_;
    std::vector<double> minAttributeBuffer_;
    std::vector<double> maxAttributeBuffer_;
    std::vector<NodeId> pruneCandidateQueue_;
    std::vector<NodeId> selectedPruneCandidates_;
    CasfComponentTreesAttribute attribute_ = CasfComponentTreesAttribute::Area;
 
    static std::unique_ptr<attribute_computer_t> makeAttributeComputer(CasfComponentTreesAttribute attribute) {
        switch (attribute) {
        case CasfComponentTreesAttribute::Area:
            return std::make_unique<DynamicAreaAttributeComputer<T>>();
        case CasfComponentTreesAttribute::BoundingBoxWidth:
            return std::make_unique<DynamicBoundingBoxAttributeComputer<T>>(BoundingBoxMeasure::Width);
        case CasfComponentTreesAttribute::BoundingBoxHeight:
            return std::make_unique<DynamicBoundingBoxAttributeComputer<T>>(BoundingBoxMeasure::Height);
        case CasfComponentTreesAttribute::BoundingBoxDiagonal:
            return std::make_unique<DynamicBoundingBoxAttributeComputer<T>>(BoundingBoxMeasure::DiagonalLength);
        }
        throw std::runtime_error("Unknown CASF component-tree attribute.");
    }
 
    const std::vector<NodeId>& selectPruneCandidates(const tree_t& tree, const std::vector<double>& attribute, double threshold) {
        pruneCandidateQueue_.clear();
        selectedPruneCandidates_.clear();
 
        const MorphologicalTree& topology = tree.topology();
        const NodeId root = topology.root();
        if (root == InvalidNode || !topology.isAlive(root)) {
            return selectedPruneCandidates_;
        }
 
        pruneCandidateQueue_.push_back(root);
        size_t head = 0;
 
        while (head < pruneCandidateQueue_.size()) {
            const NodeId nodeId = pruneCandidateQueue_[head++];
            if (!topology.isAlive(nodeId)) {
                continue;
            }
 
            const double nodeAttribute = attribute[static_cast<size_t>(nodeId)];
            if (!topology.isRoot(nodeId) && nodeAttribute <= threshold) {
                selectedPruneCandidates_.push_back(nodeId);
                continue;
            }
 
            for (NodeId childId : topology.children(nodeId)) {
                if (topology.isAlive(childId)) {
                    pruneCandidateQueue_.push_back(childId);
                }
            }
        }
 
        return selectedPruneCandidates_;
    }
 
    void applyFilterStep(double threshold) {
        const std::vector<NodeId> maxCandidates = selectPruneCandidates(maxTree_, maxAttributeBuffer_, threshold);
        adjust_->pruneMaxTreeAndUpdateMinTree(maxCandidates);
 
        const std::vector<NodeId> minCandidates = selectPruneCandidates(minTree_, minAttributeBuffer_, threshold);
        adjust_->pruneMinTreeAndUpdateMaxTree(minCandidates);
    }
 
  public:
    CasfComponentTrees(image_ptr_t image, CasfComponentTreesAttribute attribute = CasfComponentTreesAttribute::Area, double radius = 1.5)
        : adjacency_(image ? image->getNumRows() : 0, image ? image->getNumColumns() : 0, radius),
          minTree_(MorphologicalTreeFactory::createMinTree(image, radius)), maxTree_(MorphologicalTreeFactory::createMaxTree(image, radius)),
          minAttributeComputer_(makeAttributeComputer(attribute)), maxAttributeComputer_(makeAttributeComputer(attribute)), attribute_(attribute) {
        if (!image || image->getNumRows() <= 0 || image->getNumColumns() <= 0 || image->getSize() <= 0) {
            throw std::invalid_argument("CasfComponentTrees requires a non-empty image.");
        }
 
        adjust_ = std::make_unique<DualMinMaxTreeIncrementalFilter<T>>(&minTree_, &maxTree_, adjacency_);
        adjust_->setAttributeComputer(*minAttributeComputer_, *maxAttributeComputer_, minAttributeBuffer_, maxAttributeBuffer_);
        minAttributeComputer_->computeAttribute(minTree_, minAttributeBuffer_);
        maxAttributeComputer_->computeAttribute(maxTree_, maxAttributeBuffer_);
 
        const size_t maxNodes = static_cast<size_t>(std::max(minTree_.topology().numInternalNodeSlots(), maxTree_.topology().numInternalNodeSlots()));
        pruneCandidateQueue_.reserve(maxNodes);
        selectedPruneCandidates_.reserve(maxNodes);
    }
 
    [[nodiscard]] image_ptr_t filter(const std::vector<double>& thresholds) {
        for (double threshold : thresholds) {
            applyFilterStep(threshold);
        }
        return minTree_.reconstructFromNodeAltitudes();
    }
 
    [[nodiscard]] const tree_t& minTree() const noexcept { return minTree_; }
 
    [[nodiscard]] const tree_t& maxTree() const noexcept { return maxTree_; }
 
    [[nodiscard]] CasfComponentTreesAttribute attribute() const noexcept { return attribute_; }
 
    [[nodiscard]] std::pair<std::vector<NodeId>, std::vector<T>> exportMinTree() const { return minTree_.exportHigraHierarchy(); }
 
    [[nodiscard]] std::pair<std::vector<NodeId>, std::vector<T>> exportMaxTree() const { return maxTree_.exportHigraHierarchy(); }
};
 
} // namespace mmcfilters::adjust