NDEVRCorrespondenceEstimation.h

#pragma once
#include <NDEVR/RTree.h>
#include <NDEVR/Buffer.h>
#include <NDEVR/LineSegment.h>
#include <NDEVR/Distance.h>
#include <pcl/registration/correspondence_types.h>
namespace NDEVR
{
    using namespace pcl;
    using namespace pcl::registration;
    template <typename PointSource,
        typename PointTarget,
        typename NormalT,
        typename Scalar = float>
    class NDEVRCorrespondenceEstimation
    {
    public:
        using Ptr = shared_ptr<NDEVRCorrespondenceEstimation<PointSource,
            PointTarget,
            NormalT,
            Scalar>>;
        using ConstPtr = shared_ptr<const NDEVRCorrespondenceEstimation<PointSource,
            PointTarget,
            NormalT,
            Scalar>>;
 
        
        using PointCloudSource = pcl::PointCloud<PointSource>;
        using PointCloudSourcePtr = typename PointCloudSource::Ptr;
        using PointCloudSourceConstPtr = typename PointCloudSource::ConstPtr;
 
        using PointCloudTarget = pcl::PointCloud<PointTarget>;
        using PointCloudTargetPtr = typename PointCloudTarget::Ptr;
        using PointCloudTargetConstPtr = typename PointCloudTarget::ConstPtr;
 
        using PointCloudNormals = pcl::PointCloud<NormalT>;
        using NormalsPtr = typename PointCloudNormals::Ptr;
        using NormalsConstPtr = typename PointCloudNormals::ConstPtr;
        NDEVRCorrespondenceEstimation()
            : source_normals_()
            , source_normals_transformed_()
            , target_normals_()
        {}

        virtual ~NDEVRCorrespondenceEstimation() = default;

        void setTree(const RTree<3, fltp04>* tree, const Buffer<Vertex<3, fltp04>>* params, const Buffer<Vertex<3, fltp04>>* tree_sorted_params)
        {
           m_tree = tree;
           m_params = params;
           m_tree_sorted_params = tree_sorted_params;
        }

        inline void setSourceNormals(const NormalsConstPtr& normals)
        {
            source_normals_ = normals;
        }

        inline NormalsConstPtr getSourceNormals() const
        {
            return (source_normals_);
        }

        inline void setTargetNormals(const NormalsConstPtr& normals)
        {
            target_normals_ = normals;
        }

        inline NormalsConstPtr getTargetNormals() const
        {
            return (target_normals_);
        }

        bool requiresSourceNormals() const
        {
            return (true);
        }

        void setInputTarget(PointCloudSourcePtr input)
        {
            input_ = input;
        }

        void setSourceNormals(pcl::PCLPointCloud2::ConstPtr cloud2)
        {
            NormalsPtr cloud(new PointCloudNormals);
            fromPCLPointCloud2(*cloud2, *cloud);
            setSourceNormals(cloud);
        }

        bool requiresTargetNormals() const
        {
            return (true);
        }

        void setTargetNormals(pcl::PCLPointCloud2::ConstPtr cloud2)
        {
            NormalsPtr cloud(new PointCloudNormals);
            fromPCLPointCloud2(*cloud2, *cloud);
            setTargetNormals(cloud);
        }

        Vertex<3, fltp04> vertex(uint04 index) const
        {
            const auto& pt = input_->at(index);
            return Vertex<3, fltp04>(pt.x, pt.y, pt.z);
        }

        LineSegment<3, fltp04> lineSegment(uint04 index) const
        {
            LineSegment<3, fltp04> ls;
            const auto& pt = input_->at(index);
            ls[A] = Vertex<3, fltp04>(pt.x, pt.y, pt.z);
            Vertex<3, fltp04> nm((*target_normals_)[index].normal_x
                , (*target_normals_)[index].normal_y
                , (*target_normals_)[index].normal_z);
            ls[B] = ls[A] + m_normal_expansion_factor * nm;
            ls[A] = ls[A] - m_normal_expansion_factor * nm;
            return ls;
        }

        bool radiusSearch(uint04 tgt_idx, uint04& index, fltp04& k_sqr_distances) const
        {
            Vertex<3, fltp04> vert = vertex(tgt_idx);
            index = m_tree->closestElementPresorted(vert, *m_tree_sorted_params, k_sqr_distances, m_epsilon);
            if (IsInvalid(index))
                return false;
            index = m_tree->indices()[index];
            return true;
        }

        bool radiusNormalSearch(uint04 tgt_idx, uint04& index, fltp04& k_sqr_distances) const
        {
            LineSegment<3, fltp04> seg = lineSegment(tgt_idx);
            index = m_tree->closestElementPresorted(seg, *m_tree_sorted_params, k_sqr_distances, m_epsilon);
            if (IsInvalid(index))
                return false;
            index = m_tree->indices()[index];
            return true;
        }

        MaxHeap<fltp04, uint04> radiusKSearch(uint04 tgt_idx, fltp04 max_distance_squared, uint04 k = 0) const
        {
            Vertex<3, fltp04> vert = vertex(tgt_idx);
            //vertex = m_input_transform * vertex;
            MaxHeap<fltp04, uint04> heap(cast<uint04>(k + 1));
            m_tree->presortedClosestElements(vert, k, heap, *m_tree_sorted_params, max_distance_squared, m_epsilon);
            const uint04 size = heap.size();
            //auto sorted_values = heap.sortedValues();
            for (uint04 i = 0; i < size; i++)
                heap.values()[i].first = m_tree->indices()[heap.values()[i].first];
            return heap;
        }

        MaxHeap<fltp04, uint04> radiusNormalKSearch(uint04 tgt_idx, fltp04 max_distance_squared, uint04 k = 0) const
        {
            LineSegment<3, fltp04> seg = lineSegment(tgt_idx);
            //vertex = m_input_transform * vertex;
            MaxHeap<fltp04, uint04> heap(cast<uint04>(k + 1));
            m_tree->presortedClosestElements(seg, k, heap, *m_tree_sorted_params, max_distance_squared, m_epsilon);
            const uint04 size = heap.size();
            //auto sorted_values = heap.sortedValues();
            for (uint04 i = 0; i < size; i++)
                heap.values()[i].first = m_tree->indices()[heap.values()[i].first];
            return heap;
        }

        fltp04 calcDistance(uint04 a, uint04 b) const
        {
            fltp04 d = distance<fltp04>(vertex(a), (*m_params)[b]);
            fltp04 cos_angle = (*source_normals_)[b].normal_x *
                (*target_normals_)[a].normal_x +
                (*source_normals_)[b].normal_y *
                (*target_normals_)[a].normal_y +
                (*source_normals_)[b].normal_z *
                (*target_normals_)[a].normal_z;
            return d * (2.0f - cos_angle * cos_angle);
        }

        void determineCorrespondencesMinDistance(uint04 offset, pcl::Correspondences& correspondences, fltp04 max_distance)
        {
            // Iterate over the input set of source indices
            int count = input_->size();
            correspondences.clear();
            //#pragma omp parallel for schedule(dynamic) num_threads(8)
            for (int idx_i = offset; idx_i < count; idx_i += m_skip_factor)
            {
                uint04 idx;
                fltp04 dist = max_distance;
                if (radiusSearch(idx_i, idx, dist))
                {
                    if (dist <= max_distance)
                    {
                        pcl::Correspondence corr;
                        corr.index_query = idx_i;
                        corr.index_match = idx;
                        corr.distance = dist;
                        correspondences.push_back(corr);
                    }
                }
            }
        }

        void determineCorrespondencesBackProjection(uint04 offset, pcl::Correspondences& correspondences, fltp04 max_distance)
        {
            // Iterate over the input set of source indices
            int count = input_->size();
            correspondences.clear();
            //#pragma omp parallel for schedule(dynamic) num_threads(8)
            for (int idx_i = offset; idx_i < count; idx_i += m_skip_factor)
            {
                
                if (k_ == 1)
                {
                    uint04 idx;
                    fltp04 dist = max_distance;
                    if (radiusNormalSearch(idx_i, idx, dist))
                    {
                        fltp04 distance = calcDistance(idx_i, idx);
                        if (distance <= max_distance)
                        {
                            pcl::Correspondence corr;
                            corr.index_query = idx_i;
                            corr.index_match = idx;
                            corr.distance = distance;
                            correspondences.push_back(corr);
                        }
                    }
                }
                else
                {
                    MinHeap<fltp04, uint04> heap = radiusNormalKSearch(idx_i, max_distance, k_);
                    const uint04 size = heap.size();
                    if (size == 0)
                        continue;
                    uint04 min_index = heap.values()[0].first;
                    fltp04 min_dist = calcDistance(idx_i, min_index);
                    for (uint04 i = 1; i < size; i++)
                    {
                        float dist = calcDistance(idx_i, heap.values()[i].first);
                        if (dist < min_dist)
                        {
                            min_dist = dist;
                            min_index = heap.values()[i].first;
                        }
                    }
                    if (min_dist <= max_distance)
                    {
                        pcl::Correspondence corr;
                        corr.index_query = idx_i;
                        corr.index_match = min_index;
                        corr.distance = min_dist;
                        correspondences.push_back(corr);
                    }
                }
            }
        }

        inline void setKSearch(unsigned int k)
        {
            k_ = k;
        }

        inline unsigned int getKSearch() const
        {
            return (k_);
        }
    private:
        const RTree<3, fltp04>* m_tree; 
        const Buffer<Vertex<3, fltp04>>* m_params; 
        const Buffer<Vertex<3, fltp04>>* m_tree_sorted_params; 
        fltp04 m_normal_expansion_factor = 0.1f; 
        NormalsConstPtr source_normals_;

        NormalsPtr source_normals_transformed_;
        PointCloudSourcePtr input_; 

        PointCloudSourcePtr source_cloud_transformed_;

        NormalsConstPtr target_normals_;
        fltp04 m_epsilon = 0.0f; 

        uint04 k_{ 10 };
        int m_skip_factor = 5; 
    };
}