segments.h

Go to the documentation of this file.

#pragma once
 
#include <map>
#include <vector>
 
#include <pcl/PointIndices.h>
#include <pcl/point_cloud.h>
 
#include <SimoxUtility/shapes/AxisAlignedBoundingBox.h>
 
#include <VisionX/interface/core/DataTypes.h>
 
#include "AABB.h"
 
namespace visionx::tools::detail
{
    // Version for pcl::PointIndices
    void addSegmentIndex(std::map<uint32_t, pcl::PointIndices>& segmentIndices,
                         uint32_t label,
                         std::size_t index);
    // Version for pcl::PointIndices::Ptr
    void addSegmentIndex(std::map<uint32_t, pcl::PointIndices::Ptr>& segmentIndices,
                         uint32_t label,
                         std::size_t index);
 
    template <class LabeledPointT, class MapValueT>
    void
    addSegmentIndices(const pcl::PointCloud<LabeledPointT>& labeledCloud,
                      std::map<uint32_t, MapValueT>& segmentIndices,
                      bool excludeZero = false)
    {
        for (std::size_t i = 0; i < labeledCloud.points.size(); i++)
        {
            const uint32_t currentLabel = labeledCloud.points[i].label;
            if (!(excludeZero && currentLabel == 0))
            {
                addSegmentIndex(segmentIndices, currentLabel, i);
            }
        }
    }
} // namespace visionx::tools::detail
 
namespace visionx::tools
{
    // Segment Indices
 
    /**
     * @brief Add indices of each segment's points to `segmentIndices`.
     * @param segmentIndices The map to fill (label -> point indicies).
     * @param labeledCloud A labeled point cloud.
     * @param excludeZero If true, points with label == 0 are ignored.
     */
    template <class LabeledPointT>
    void
    addSegmentIndices(std::map<uint32_t, pcl::PointIndices>& segmentIndices,
                      const pcl::PointCloud<LabeledPointT>& labeledCloud,
                      bool excludeZero = false)
    {
        detail::addSegmentIndices(labeledCloud, segmentIndices, excludeZero);
    }
 
    template <class LabeledPointT>
    void
    addSegmentIndices(std::map<uint32_t, pcl::PointIndices::Ptr>& segmentIndices,
                      const pcl::PointCloud<LabeledPointT>& labeledCloud,
                      bool excludeZero = false)
    {
        detail::addSegmentIndices(labeledCloud, segmentIndices, excludeZero);
    }
 
    /**
     * @brief Get a map from segment labels to indices of segments' points.
     * @param labeledCloud A labeled point cloud.
     * @param excludeZero If true, points with label == 0 are ignored.
     * @return Map from segment labels to segment indices.
     */
    template <class LabeledPointT>
    std::map<uint32_t, pcl::PointIndices>
    getSegmentIndices(const pcl::PointCloud<LabeledPointT>& labeledCloud, bool excludeZero = false)
    {
        std::map<uint32_t, pcl::PointIndices> segmentIndices;
        addSegmentIndices(segmentIndices, labeledCloud, excludeZero);
        return segmentIndices;
    }
 
    /// Version of `addSegmentIndices()` containing `pcl::PointIndices::Ptr`.
    /// @see addSegmentIndices()
    template <class LabeledPointT>
    std::map<uint32_t, pcl::PointIndices::Ptr>
    getSegmentIndicesPtr(const pcl::PointCloud<LabeledPointT>& labeledCloud,
                         bool excludeZero = false)
    {
        std::map<uint32_t, pcl::PointIndices::Ptr> segmentIndices;
        addSegmentIndices(segmentIndices, labeledCloud, excludeZero);
        return segmentIndices;
    }
 
    /// Return the set of labels in `pointCloud`.
    template <typename LabeledPointT>
    std::set<uint32_t>
    getLabelSet(const pcl::PointCloud<LabeledPointT>& pointCloud)
    {
        std::set<uint32_t> labels;
        for (const auto& point : pointCloud)
        {
            labels.insert(point.label);
        }
        return labels;
    }
 
    /// Return the labels in `segmentIndices` as `std::set`.
    template <class ValueT>
    std::set<uint32_t>
    getLabelSet(const std::map<uint32_t, ValueT>& segmentIndices)
    {
        std::set<uint32_t> labels;
        for (const auto& [label, _] : segmentIndices)
        {
            labels.insert(label);
        }
        return labels;
    }
 
    /// Return the labels in `segmentIndices` as `std::vector`.
    template <class ValueT>
    std::vector<uint32_t>
    getUniqueLabels(const std::map<uint32_t, ValueT>& segmentIndices)
    {
        std::vector<uint32_t> labels;
        labels.reserve(segmentIndices.size());
        for (const auto& [label, _] : segmentIndices)
        {
            labels.push_back(label);
        }
        return labels;
    }
 
    template <class PointT>
    void
    relabel(pcl::PointCloud<PointT>& pointCloud,
            const std::map<uint32_t, pcl::PointIndices>& segmentIndices)
    {
        for (const auto& [label, indices] : segmentIndices)
        {
            for (int index : indices.indices)
            {
                pointCloud.at(index).label = label;
            }
        }
    }
 
    template <class PointT>
    void
    relabel(pcl::PointCloud<PointT>& pointCloud,
            const std::map<uint32_t, pcl::PointIndices::Ptr>& segmentIndices)
    {
        for (const auto& [label, indices] : segmentIndices)
        {
            for (int index : indices->indices)
            {
                pointCloud.at(index).label = label;
            }
        }
    }
 
    template <class PointT, class IndicesT = pcl::PointIndices>
    std::vector<simox::AxisAlignedBoundingBox>
    getSegmentAABBs(const pcl::PointCloud<PointT>& pointCloud,
                    const std::vector<IndicesT>& segmentIndices)
    {
        std::vector<simox::AxisAlignedBoundingBox> aabbs;
        aabbs.reserve(segmentIndices.size());
        for (const auto& indices : segmentIndices)
        {
            aabbs.push_back(getAABB(pointCloud, indices));
        }
        return aabbs;
    }
 
    template <class PointT, class IndicesT = pcl::PointIndices>
    std::map<uint32_t, simox::AxisAlignedBoundingBox>
    getSegmentAABBs(const pcl::PointCloud<PointT>& pointCloud,
                    const std::map<uint32_t, IndicesT>& segmentIndices)
    {
        std::map<uint32_t, simox::AxisAlignedBoundingBox> aabbs;
        for (const auto& [label, indices] : segmentIndices)
        {
            aabbs[label] = getAABB(pointCloud, indices);
        }
        return aabbs;
    }
 
} // namespace visionx::tools