NearestNeighborSearch.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
#pragma once
 
#include <optional>
#include <vector>
 
#include "open3d/core/Tensor.h"
#include "open3d/core/nns/FixedRadiusIndex.h"
#include "open3d/core/nns/KnnIndex.h"
#include "open3d/core/nns/NanoFlannIndex.h"
 
namespace open3d {
namespace core {
namespace nns {
 
class NearestNeighborSearch {
public:
    // NearestNeighborSearch(const Tensor &dataset_points)
    //     : dataset_points_(dataset_points){};
    NearestNeighborSearch(const Tensor &dataset_points,
                          const Dtype &index_dtype = core::Int32)
        : dataset_points_(dataset_points), index_dtype_(index_dtype) {}
    ~NearestNeighborSearch();
    NearestNeighborSearch(const NearestNeighborSearch &) = delete;
    NearestNeighborSearch &operator=(const NearestNeighborSearch &) = delete;
 
public:
    bool KnnIndex();
 
    bool MultiRadiusIndex();
 
    bool FixedRadiusIndex(std::optional<double> radius = {});
 
    bool HybridIndex(std::optional<double> radius = {});
 
    std::pair<Tensor, Tensor> KnnSearch(const Tensor &query_points, int knn);
 
    std::tuple<Tensor, Tensor, Tensor> FixedRadiusSearch(
            const Tensor &query_points, double radius, bool sort = true);
 
    std::tuple<Tensor, Tensor, Tensor> MultiRadiusSearch(
            const Tensor &query_points, const Tensor &radii);
 
    std::tuple<Tensor, Tensor, Tensor> HybridSearch(const Tensor &query_points,
                                                    const double radius,
                                                    const int max_knn) const;
 
private:
    bool SetIndex();
 
    void AssertCPU(const Tensor &t) const;
 
protected:
    std::unique_ptr<NanoFlannIndex> nanoflann_index_;
    std::unique_ptr<nns::FixedRadiusIndex> fixed_radius_index_;
    std::unique_ptr<nns::KnnIndex> knn_index_;
    const Tensor dataset_points_;
    const Dtype index_dtype_;
};
}  // namespace nns
}  // namespace core
}  // namespace open3d