PCLTreeWrapper.h

#pragma once
#include "../Headers/DLLInfo.h"
#include "../Headers/PCLFunctions.h"
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/search/search.h>
#include <NDEVR/RTree.h>
#include <NDEVR/Matrix.h>
#include <NDEVR/Buffer.h>
#include <NDEVR/Translator.h>
namespace NDEVR
{
    template<typename PointT>
    class PCLRTreeWrapper : public pcl::search::Search<PointT>
    {
    public:
        using PointCloud = pcl::PointCloud<PointT>;
        using PointCloudPtr = typename PointCloud::Ptr;
        using PointCloudConstPtr = typename PointCloud::ConstPtr;
      
        typedef std::shared_ptr<std::vector<int>> IndicesPtr;
        typedef std::shared_ptr<const std::vector<int>> IndicesConstPtr;
 
        typedef std::shared_ptr<PCLRTreeWrapper<PointT>> Ptr;
        typedef std::shared_ptr<const PCLRTreeWrapper<PointT> > ConstPtr;

        PCLRTreeWrapper(const Buffer<Vertex<3, fltp04>>& vertices, const RTree<3, fltp04>& tree, uint04& total_search_ptr, uint04 search_count, LogPtr log = LogPtr())
            : pcl::search::Search<PointT>()
            , m_tree(tree)
            , m_values(vertices)
            , m_search_count(search_count)
            , m_current_search_count(total_search_ptr)
            , m_log(_t("Tree Operation"), log)
        {}
 
        PCLRTreeWrapper(const Buffer<Vertex<3, fltp04>>& vertices, const RTree<3, fltp04>& tree, uint04& search_count, LogPtr log = LogPtr())
            : pcl::search::Search<PointT>()
            , m_tree(tree)
            , m_values(vertices)
            , m_search_count(m_values.size())
            , m_current_search_count(search_count)
            , m_log(_t("Tree Operation"), log)
        {}

        virtual ~PCLRTreeWrapper()
        {
        }
 
 
        bool setInputCloud(const PointCloudConstPtr& cloud, const IndicesConstPtr& indices = IndicesConstPtr()) final override
        {
            return pcl::search::Search<PointT>::setInputCloud(cloud, indices);
        }

 
        int nearestKSearch(const PointT& point, int k,
            std::vector<int>& k_indices,
            std::vector<float>& k_sqr_distances) const final override
        {
            return radiusSearch(point, m_max_tree_span,
                k_indices, k_sqr_distances,
                cast<uint04>(k));
        }

        int radiusSearch(const PointT& point, double radius,
            std::vector<int>& k_indices,
            std::vector<float>& k_sqr_distances,
            unsigned int k = 0) const final override
        {
            Vertex<3, fltp04> vertex(point.x, point.y, point.z);
            fltp04 max_distance_squared = cast<fltp04>(radius * radius);
            uint04 ku;
            if (k == 0)
                ku = getMin(m_max_k, m_values.size());
            else
                ku = getMin(m_max_k, cast<uint04>(k));
            if (ku == 1)
            {
                uint04 index = m_tree.closestElementPresorted(vertex, m_values, max_distance_squared, 0.0f);
                k_indices.resize(1);
                k_sqr_distances.resize(1);
                if (IsInvalid(index))
                {
                    k_indices[0] = 0;
                    k_sqr_distances[0] = Constant<fltp04>::Max;
                }
                else
                {
                    k_indices[0] = m_tree.indices()[index];
                    k_sqr_distances[0] = max_distance_squared;
                }
            }
            else
            {
                
                //vertex = m_input_transform * vertex;
                MaxHeap<fltp04, uint04> heap(cast<uint04>(ku + 1));
                m_tree.presortedClosestElements(vertex, ku, heap, m_values, max_distance_squared);
                k_indices.resize(heap.size());
                k_sqr_distances.resize(heap.size());
                const uint04 size = heap.size();
                auto sorted_values = heap.sortedValues();
                for (uint04 i = 0; i < sorted_values.size(); i++)
                {
                    k_indices[i] = cast<int>(m_tree.indices()[sorted_values[i].first]);
                    k_sqr_distances[i] = sorted_values[i].second;
                }
            }
            if (m_current_search_count % 5000 == 0)
                m_log.setProgress(cast<fltp04>(m_current_search_count) / cast<fltp04>(m_search_count));
            m_current_search_count++;
            return cast<int>(k_indices.size());
        }
        void setMaxSearchDistance(fltp04 max)
        {
            m_max_tree_span = max;
        }
        void setMaxK(uint04 max)
        {
            m_max_k = max;
        }
    protected:
        const RTree<3, fltp04>& m_tree;
        const Buffer<Vertex<3, fltp04>>& m_values;
        uint04 m_search_count;
        uint04& m_current_search_count;
        fltp04 m_max_tree_span = Constant<fltp04>::Max;
        uint04 m_max_k = Constant<uint04>::Max;
        //Matrix<fltp08> m_input_transform = Matrix<fltp08>(1.0);
        mutable ProgressInfo m_log;
    };
}