CasfComponentTrees.hpp

#pragma once
 
#include "DualMinMaxTreeIncrementalFilter.hpp"
#include "../MorphologicalTreeFactory.hpp"
 
#include <memory>
#include <stdexcept>
#include <utility>
#include <vector>
 
namespace mmcfilters::adjust {
 
enum class CasfComponentTreesAttribute {
    AREA,
    BOUNDING_BOX_WIDTH,
    BOUNDING_BOX_HEIGHT,
    BOUNDING_BOX_DIAGONAL
};
 
template<AltitudeValue T>
class CasfComponentTrees {
private:
    using tree_t = WeightedMorphologicalTree<T>;
    using image_ptr_t = ImagePtr<T>;
    using attribute_computer_t = DynamicTreeAttributeComputer<T>;
 
    AdjacencyRelation 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::BOUNDING_BOX_WIDTH:
                return std::make_unique<DynamicBoundingBoxAttributeComputer<T>>(BoundingBoxMeasure::WIDTH);
            case CasfComponentTreesAttribute::BOUNDING_BOX_HEIGHT:
                return std::make_unique<DynamicBoundingBoxAttributeComputer<T>>(BoundingBoxMeasure::HEIGHT);
            case CasfComponentTreesAttribute::BOUNDING_BOX_DIAGONAL:
                return std::make_unique<DynamicBoundingBoxAttributeComputer<T>>(BoundingBoxMeasure::DIAGONAL_LENGTH);
        }
        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.getRoot();
        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.getChildren(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->getNumCols() : 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->getNumCols() <= 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().getNumInternalNodeSlots(), maxTree_.topology().getNumInternalNodeSlots()));
        pruneCandidateQueue_.reserve(maxNodes);
        selectedPruneCandidates_.reserve(maxNodes);
    }
 
    [[nodiscard]] image_ptr_t filter(const std::vector<double>& thresholds) {
        for (double threshold : thresholds) {
            applyFilterStep(threshold);
        }
        return minTree_.reconstructionImage();
    }
 
    [[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