SIM_VolumetricConnectedComponentBuilder.h

Go to the documentation of this file.

#ifndef __SIM__VOLUMETRICCONNECTEDCOMPONENTBUILDER_H__
#define __SIM__VOLUMETRICCONNECTEDCOMPONENTBUILDER_H__

#include <SIM/SIM_FieldUtils.h>

#include <UT/UT_ArrayMap.h>
#include <UT/UT_ParallelUtil.h>
#include <SYS/SYS_Math.h>

#include <utility>

// Helper class that links together regions of an index field that 
// consist of non-negative values.

template <typename FieldType = SIM_RawIndexField>
class SIM_API SIM_VolumetricConnectedComponentBuilder
{
private:
    
    using IntegerMap = UT_ArrayMap<exint, exint>;
    using RootPairs = std::pair<exint, exint>;
    using ScalarType = typename FieldType::ScalarType;

    static constexpr exint UNVISITED_CELL = -2;
    static constexpr exint VISITED_CELL = -1;
    static constexpr bool USING_VOXEL_ARRAY = std::is_same_v<UT_VoxelArray<ScalarType>, FieldType>;

public:
    static constexpr exint INACTIVE_REGION = -1;

    // The provided labelled cell field should have active cells labelled as zero or greater.
    // If FieldType is SIM_RawField or SIM_RawIndexField, then connected_regions will be aligned
    // with label_cells. If FieldType is a UT_VoxelArray, then connected_regions will be merely
    // of the same dimension.
    // Faces are considered "open" wherever value of faceWeights exceeds min_weight.
    SIM_VolumetricConnectedComponentBuilder(SIM_RawIndexField &connected_regions,
                                            const FieldType &label_cells,
                                            const SIM_RawField **faceWeights,
                                            float min_weight = 0)
        : m_connected_regions(connected_regions)
        , m_face_weights(faceWeights)
        , m_min_weight(min_weight)
        {
            if constexpr (USING_VOXEL_ARRAY)
            {
                m_labelled_cells = &label_cells;
                m_connected_regions.init(SIM_SAMPLE_CENTER, UT_Vector3(0.0f), UT_Vector3(1.0f),
                                         label_cells.getXRes(), label_cells.getYRes(), label_cells.getZRes());
            }
            else
            {
                m_labelled_cells = label_cells.field();
                m_connected_regions.match(label_cells);
            }

            m_connected_regions.makeConstant(INACTIVE_REGION);
        }

    template<typename IsLabelActiveFunctor>
    exint buildConnectedComponents(const IsLabelActiveFunctor &isCellLabelActiveFunctor);

private:

    SYS_FORCE_INLINE
    exint
    flattenIndex(const UT_Vector3I& index, const UT_Vector3I& size) const
    {
        return (index[2] * size[1] + index[1]) * size[0] + index[0];
    }

    template<typename IsLabelActiveFunctor>
    void buildLocalRoots(SIM_RawIndexField &visited_cells, const IsLabelActiveFunctor &isCellLabelActiveFunctor);

    THREADED_METHOD1_CONST(SIM_VolumetricConnectedComponentBuilder, m_labelled_cells->numTiles() > 20,
                            linkLocalRoots,
                            UT_Array<UT_Array<RootPairs>> &, parallel_root_pair_lists)

    void linkLocalRootsPartial(UT_Array<UT_Array<RootPairs>> &parallel_root_pair_lists,
                                const UT_JobInfo &info) const;

    THREADED_METHOD2(SIM_VolumetricConnectedComponentBuilder, m_connected_regions.shouldMultiThread(),
                        assignRegionLabels,
                        const IntegerMap &, root_map,
                        const IntegerMap &, region_index_map)

    void assignRegionLabelsPartial(const IntegerMap &root_map,
                                    const IntegerMap &region_index_map,
                                    const UT_JobInfo &info);

    const UT_VoxelArray<ScalarType> *m_labelled_cells;
    SIM_RawIndexField &m_connected_regions;

    const SIM_RawField **m_face_weights;

    float m_min_weight;
};

SIM_EXTERN_TEMPLATE(SIM_VolumetricConnectedComponentBuilder<SIM_RawField>);
SIM_EXTERN_TEMPLATE(SIM_VolumetricConnectedComponentBuilder<SIM_RawIndexField>);
SIM_EXTERN_TEMPLATE(SIM_VolumetricConnectedComponentBuilder<UT_VoxelArray<fpreal32>>);
SIM_EXTERN_TEMPLATE(SIM_VolumetricConnectedComponentBuilder<UT_VoxelArray<fpreal64>>);

template<typename FieldType>
template<typename IsLabelActiveFunctor>
exint SIM_VolumetricConnectedComponentBuilder<FieldType>::buildConnectedComponents(const IsLabelActiveFunctor &isCellLabelActiveFunctor)
{
    // We want to replicate a flood fill throughout the entire contiguous
    // volume without actually flooding in a BFS sense.

    // First we perform a flood locally within each tile and point every voxel
    // in a connected region to a single representative cell that acts as a local root.
    {
        SIM_RawIndexField visited_cells;
        visited_cells.match(m_connected_regions);
        visited_cells.makeConstant(UNVISITED_CELL);

        buildLocalRoots(visited_cells, isCellLabelActiveFunctor);
    }

    const int thread_count = UT_Thread::getNumProcessors();
    UT_Array<UT_Array<RootPairs>> parallel_root_pair_lists;
    parallel_root_pair_lists.setSize(thread_count);

    // Once we have a local root for each connected component in each tile, we can
    // link the roots together across tile boundaries. Because there are far fewer tiles than voxels,
    // this process is extremely fast.

    linkLocalRoots(parallel_root_pair_lists);

    // Compile the pairs of roots into a single list
    UT_Array<RootPairs> root_pair_list;
    exint list_size = 0;
    for (int thread = 0; thread < thread_count; ++thread)
        list_size += parallel_root_pair_lists[thread].size();

    root_pair_list.bumpCapacity(list_size);

    for (int thread = 0; thread < thread_count; ++thread)
        root_pair_list.concat(parallel_root_pair_lists[thread]);

    // Build a hash map with the local root values as the keys and a list of all the roots they point to 
    // as the value. Since this is done at the tile level, there won't be many entries and it doesn't need
    // to be parallelized.

    UT_Map<exint, UT_Array<exint>> root_to_adjacent_root_map;
    root_to_adjacent_root_map.reserve(root_pair_list.size());

    // Sort the root pair list so we can pull from the hash table once per local root and fill its list.
    UTparallelSort(root_pair_list.begin(), root_pair_list.end(), [](const RootPairs &pair0, const RootPairs &pair1)
    {
        if (pair0.first < pair1.first) return true;
            return false;
    });

    const exint rpl_size = root_pair_list.size();
    for (int i = 0; i < rpl_size; )
    {
        exint root_key = root_pair_list[i].first;

        auto &local_root_list = root_to_adjacent_root_map[root_key];

        // Iterate over entries in the list that correspond to the same root key
        // and append the adjacent roots to the list
        while (i < rpl_size && root_pair_list[i].first == root_key)
        {
            exint adjacent_root_key = root_pair_list[i].second;

            // If the list of pairs was built correctly there shouldn't be any repeated pairs.
            UT_ASSERT(local_root_list.find(adjacent_root_key) == -1);

            local_root_list.append(adjacent_root_key);
            ++i;
        }
    }

#if UT_ASSERT_LEVEL > 0
    // Debug check: make sure that every link has a reciprocating link back
    for (const auto& local_root_map : root_to_adjacent_root_map)
    {
        exint local_root_key = local_root_map.first;
        const auto &local_root_list = local_root_map.second;

        for (const exint adjacent_root_key : local_root_list)
        {
            // Make sure the link exists
            UT_ASSERT(root_to_adjacent_root_map.contains(adjacent_root_key));

            const auto &adjacent_root_list = root_to_adjacent_root_map[adjacent_root_key];

            // Make sure that the root we are linking to does indeed link back.
            UT_ASSERT(adjacent_root_list.find(local_root_key) >= 0);
        }
    }
#endif

    // Now that the root connections are compiled, we can walk through the implied graph of root connections
    // per region. Each root walks through the graph and stores the largest root they come across. After walking
    // through the graph, the final root is stored in a separate hash map.

    IntegerMap final_root_map;
    final_root_map.reserve(root_pair_list.size());

    IntegerMap visited_roots_map;
    visited_roots_map.reserve(root_pair_list.size());

    {
        UT_Array<exint> to_visit_list;
        to_visit_list.bumpCapacity(root_pair_list.size());

        // This list stores all the nodes in the graph within
        // a connected component so we only have to walk through
        // the graph once to map all the local nodes to a global root
        // within the connected component.
        UT_Array<exint> roots_to_finish_list;
        roots_to_finish_list.bumpCapacity(root_pair_list.size());

        for (const auto &local_root_map : root_to_adjacent_root_map)
        {
            exint final_root_index = local_root_map.first;

            // If this current root has been visited already, we don't need to walk through the graph.
            if (visited_roots_map.contains(final_root_index))
            {
                UT_ASSERT(visited_roots_map[final_root_index] == VISITED_CELL);
                continue;
            }

            to_visit_list.clear();
            to_visit_list.append(final_root_index);

            visited_roots_map[final_root_index] = VISITED_CELL;

            roots_to_finish_list.clear();
            roots_to_finish_list.append(final_root_index);

            // Walk through the graph without retracing any steps. Record the largest
            // root index along the way.
            while (!to_visit_list.isEmpty())
            {
                exint local_root_index = to_visit_list.last();
                to_visit_list.removeLast();

                // Store the largest index as the global root in the connected component chain.
                final_root_index = SYSmax(final_root_index, local_root_index);

                // Queue up the new list. Only insert unvisited roots.
                const auto &adjacent_root_list = root_to_adjacent_root_map[local_root_index];
                for (const exint adjacent_root_index : adjacent_root_list)
                {
                    // We should only add a local root to the list if it has not yet been added to the list
                    // and has not been visited previously.
                    if (!visited_roots_map.contains(adjacent_root_index))
                    {
                        visited_roots_map[adjacent_root_index] = VISITED_CELL;
                        roots_to_finish_list.append(adjacent_root_index);
                        to_visit_list.append(adjacent_root_index);
                    }
                    else
                    {
                        UT_ASSERT(visited_roots_map[adjacent_root_index] == VISITED_CELL);
                    }
                }
            }

            // Update all nodes in connected component
            for (const exint root : roots_to_finish_list)
                final_root_map[root] = final_root_index;
        }
    }

    // With each local root mapping to a global root per region, we need to give those global
    // roots a unique indentifier.
    IntegerMap region_index_map;

    exint region_index = 0;
        
    for (const auto &root : final_root_map)
    {
        // If the global root for a region has not yet been created,
        // make a new entry in the hash and assign it a unique index.
        if (!region_index_map.contains(root.second))
            region_index_map[root.second] = region_index++;
    }

    // Assign region number to all valid voxels.
    assignRegionLabels(final_root_map, region_index_map);

    m_connected_regions.fieldNC()->collapseAllTiles();

    return region_index;
}

template<typename FieldType>
template<typename IsCellLabelActiveFunctor>
void
SIM_VolumetricConnectedComponentBuilder<FieldType>::buildLocalRoots(SIM_RawIndexField &visited_cells,
                                                            const IsCellLabelActiveFunctor &isCellLabelActiveFunctor)
{
    using namespace SIM::FieldUtils;

    UT_Interrupt *boss = UTgetInterrupt();

    UT_Vector3I voxel_res = m_labelled_cells->getVoxelRes();
    
    const exint tiles = m_labelled_cells->numTiles();

    UTparallelForEachNumber(tiles, [&](const UT_BlockedRange<exint> &range)
    {
        UT_VoxelArrayIterator<typename FieldType::ScalarType> vit;
        vit.setConstArray(m_labelled_cells);

        UT_VoxelTileIterator<typename FieldType::ScalarType> vitt;

        // TODO: change from for each number to something that uses a larger batch.
        // Allocating this list for every time could be slow
        UT_Array<UT_Vector3I> to_visit_cell_list;

        // TILESIZE is a Houdini defined value for the number of
        // voxels in a dimension of a voxel tile.
        to_visit_cell_list.bumpCapacity(TILESIZE * TILESIZE * TILESIZE);

        for (exint i = range.begin(); i != range.end(); ++i)
        {
            vit.myTileStart = i;
            vit.myTileEnd = i + 1;
            vit.rewind();

            if (boss->opInterrupt())
                return;

            if (!vit.isTileConstant() || isCellLabelActiveFunctor(vit.getValue()))
            {
                UT_Vector3I tile_start, tile_end;
                vit.getTileVoxels(tile_start, tile_end);

                vitt.setTile(vit);
                for (vitt.rewind(); !vitt.atEnd(); vitt.advance())
                {
                    UT_Vector3I cell(vitt.x(), vitt.y(), vitt.z());

                    UT_ASSERT(cell[0] >= tile_start[0] && cell[0] < tile_end[0] &&
                                cell[1] >= tile_start[1] && cell[1] < tile_end[1] &&
                                cell[2] >= tile_start[2] && cell[2] < tile_end[2]);

                    if (isCellLabelActiveFunctor(vitt.getValue()) && getFieldValue(visited_cells, cell) == UNVISITED_CELL)
                    {
                        const exint flat_root = flattenIndex(cell, voxel_res);

                        setFieldValue(visited_cells, cell, VISITED_CELL);
                        setFieldValue(m_connected_regions, cell, flat_root);

                        to_visit_cell_list.clear();
                        to_visit_cell_list.append(cell);

                        // Begin the flood fill within the tile.
                        while (!to_visit_cell_list.isEmpty())
                        {
                            UT_Vector3I local_cell = to_visit_cell_list.last();
                            to_visit_cell_list.removeLast();

                            // Load up adjacent cells that have not yet been labelled.
                            for (int axis : {0,1,2})
                                for (int direction : {0,1})
                                {
                                    UT_Vector3I adjacent_cell = cellToCellMap(local_cell, axis, direction);

                                    // Check that the adjacent cell is within the same tile.
                                    if (adjacent_cell[axis] < tile_start[axis] || adjacent_cell[axis] >= tile_end[axis])
                                        continue;

                                    // We don't want to touch a cell that has already been visited.
                                    // We ignore any cells that do not have a valid label.
                                    if (getFieldValue(visited_cells, adjacent_cell) == UNVISITED_CELL &&
                                        isCellLabelActiveFunctor((*m_labelled_cells)(adjacent_cell[0],
                                                                                     adjacent_cell[1],
                                                                                     adjacent_cell[2])))
                                    {
                                        UT_Vector3I face = cellToFaceMap(local_cell, axis, direction);
                                        
                                        // We only want to connect cells across non-zero face weights.
                                        if (getFieldValue(*m_face_weights[axis], face) > m_min_weight)
                                        {
                                            setFieldValue(visited_cells, adjacent_cell, VISITED_CELL);
                                            setFieldValue(m_connected_regions, adjacent_cell, flat_root);

                                            to_visit_cell_list.append(adjacent_cell);
                                        }
                                    }
                                }
                        }
                    }
                }
            }
        }
    });

    // Compress down the connected region tiles
    m_connected_regions.fieldNC()->collapseAllTiles();
}

#endif