ComponentTreeBuilder.hpp

#pragma once
 
#include "MorphologicalTreeSemantics.hpp"
#include "NativeHierarchyOwned.hpp"
#include "../utils/Altitude.hpp"
#include "../utils/Common.hpp"
#include "../utils/Image.hpp"
#include "../utils/RegularGridAdjacency2D.hpp"
#include "../utils/CommittedGridAccess.hpp"
 
#include <algorithm>
#include <cassert>
#include <cstddef>
#include <cstdint>
#include <iostream>
#include <numeric>
#include <tuple>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace mmcfilters {
 
class ComponentTreeBuilder {
    RegularGridAdjacency2D adjacency_;
    bool isMaxTree_;
 
    template <typename PixelType> [[nodiscard]] std::vector<PixelId> sort(const ImagePtr<PixelType>& image) const {
        const int numPixels = image->getSize();
        std::vector<PixelId> orderedPixels(static_cast<std::size_t>(numPixels));
        const PixelType* values = image->rawData();
 
        if constexpr (!std::is_same_v<PixelType, std::uint8_t>) {
            if (PRINT_LOG) {
                std::cout << "Sorting image with comparison sort, size: " << numPixels << std::endl;
            }
            std::iota(orderedPixels.begin(), orderedPixels.end(), 0);
            if (isMaxTree_) {
                std::stable_sort(orderedPixels.begin(), orderedPixels.end(), [&](PixelId lhs, PixelId rhs) { return values[lhs] < values[rhs]; });
            } else {
                std::stable_sort(orderedPixels.begin(), orderedPixels.end(), [&](PixelId lhs, PixelId rhs) { return values[lhs] > values[rhs]; });
            }
        } else {
            if (PRINT_LOG) {
                std::cout << "Sorting uint8 image with counting sort, size: " << numPixels << std::endl;
            }
 
            int maxValue = static_cast<int>(values[0]);
            for (int index = 1; index < numPixels; ++index) {
                if (maxValue < values[index]) {
                    maxValue = values[index];
                }
            }
 
            std::vector<std::uint32_t> counter(static_cast<std::size_t>(maxValue) + 1, 0);
            if (isMaxTree_) {
                for (int index = 0; index < numPixels; ++index) {
                    ++counter[values[index]];
                }
                for (int level = 1; level <= maxValue; ++level) {
                    counter[level] += counter[level - 1];
                }
                for (int index = numPixels - 1; index >= 0; --index) {
                    orderedPixels[--counter[values[index]]] = index;
                }
            } else {
                for (int index = 0; index < numPixels; ++index) {
                    ++counter[maxValue - values[index]];
                }
                for (int level = 1; level <= maxValue; ++level) {
                    counter[level] += counter[level - 1];
                }
                for (int index = numPixels - 1; index >= 0; --index) {
                    orderedPixels[--counter[maxValue - values[index]]] = index;
                }
            }
        }
        return orderedPixels;
    }
 
    template <typename PixelType>
    [[nodiscard]] std::tuple<std::vector<PixelId>, std::vector<PixelId>, int> buildTree(const ImagePtr<PixelType>& image) const {
        std::vector<PixelId> orderedPixels = sort(image);
        const PixelType* values = image->rawData();
 
        const int numPixels = image->getSize();
        std::vector<PixelId> unionParent(static_cast<std::size_t>(numPixels), InvalidPixel);
        std::vector<PixelId> pixelParent(static_cast<std::size_t>(numPixels), InvalidPixel);
        auto findRoot = [&](PixelId pixel) {
            while (unionParent[pixel] != pixel) {
                unionParent[pixel] = unionParent[unionParent[pixel]];
                pixel = unionParent[pixel];
            }
            return pixel;
        };
 
        for (int index = numPixels - 1; index >= 0; --index) {
            const PixelId pixel = orderedPixels[static_cast<std::size_t>(index)];
            pixelParent[pixel] = pixel;
            unionParent[pixel] = pixel;
            for (PixelId neighbor : detail::CommittedGridAccess::neighbors(adjacency_, pixel)) {
                if (unionParent[neighbor] == InvalidPixel) {
                    continue;
                }
                const PixelId neighborRoot = findRoot(neighbor);
                if (pixel != neighborRoot) {
                    pixelParent[neighborRoot] = pixel;
                    unionParent[neighborRoot] = pixel;
                }
            }
        }
 
        int numNodes = 0;
        for (int index = 0; index < numPixels; ++index) {
            const PixelId pixel = orderedPixels[static_cast<std::size_t>(index)];
            const PixelId parent = pixelParent[pixel];
            if (values[pixelParent[parent]] == values[parent]) {
                pixelParent[pixel] = pixelParent[parent];
            }
            if (pixelParent[pixel] == pixel || values[pixelParent[pixel]] != values[pixel]) {
                ++numNodes;
            }
        }
 
        return {std::move(pixelParent), std::move(orderedPixels), numNodes};
    }
 
    template <AltitudeValue T>
    [[nodiscard]] NativeHierarchyOwned<T> computeIntrinsicAttributes(const ImagePtr<T>& image, const std::vector<PixelId>& pixelParent,
                                                                   const std::vector<PixelId>& orderedPixels, int numNodes) {
        const MorphologicalTreeKind kind = isMaxTree_ ? MorphologicalTreeKind::MaxTree : MorphologicalTreeKind::MinTree;
        const GridDomain2D gridDomain{image->getNumRows(), image->getNumColumns()};
        const auto* imageValues = image->rawData();
        std::vector<NodeId> nodeParent(static_cast<std::size_t>(numNodes));
        std::vector<T> nodeValuation(static_cast<std::size_t>(numNodes));
        std::vector<NodeId> smallestNode(static_cast<std::size_t>(image->getSize()));
 
        NodeId idNode = 0;
        for (const PixelId pixel : orderedPixels) {
            const PixelId parent = pixelParent[pixel];
            if (pixel == parent || imageValues[pixel] != imageValues[parent]) {
                assert(pixel != parent || idNode == 0);
                smallestNode[pixel] = idNode;
                nodeParent[idNode] = pixel == parent ? idNode : smallestNode[parent];
                nodeValuation[idNode] = imageValues[pixel];
                ++idNode;
            } else {
                smallestNode[pixel] = smallestNode[parent];
            }
        }
 
        assert(idNode == numNodes);
        return makeNativeHierarchyOwned<T>(std::move(nodeParent), std::move(smallestNode), std::move(nodeValuation), 0, gridDomain,
                                           makeMorphologicalTreeSemantics(kind, SharedAdjacencyContext{std::move(adjacency_)}));
    }
 
  public:
    explicit ComponentTreeBuilder(RegularGridAdjacency2D adjacency, bool isMaxTree) noexcept
        : adjacency_(std::move(adjacency)), isMaxTree_(isMaxTree) {}
 
    template <AltitudeValue T> [[nodiscard]] NativeHierarchyOwned<T> build(const ImagePtr<T>& image) && {
        assert(adjacency_.getNumRows() == image->getNumRows() && adjacency_.getNumColumns() == image->getNumColumns());
        auto [pixelParent, orderedPixels, numNodes] = buildTree(image);
        return computeIntrinsicAttributes(image, pixelParent, orderedPixels, numNodes);
    }
};
 
} // namespace mmcfilters