Open3D (C++ API)  0.17.0
FixedRadiusSearchOpKernel.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// Copyright (c) 2018-2023 www.open3d.org
5// SPDX-License-Identifier: MIT
6// ----------------------------------------------------------------------------
7
8#pragma once
9
10#include "../TensorFlowHelper.h"
13#include "tensorflow/core/framework/op.h"
14#include "tensorflow/core/framework/op_kernel.h"
15#include "tensorflow/core/lib/core/errors.h"
16
18// namespace for code that is common for all kernels
19namespace fixed_radius_search_opkernel {
20
21// class for the allocator object
22template <class T, class TIndex>
23class OutputAllocator {
24public:
25 OutputAllocator(tensorflow::OpKernelContext* context) : context(context) {}
26
27 void AllocIndices(TIndex** ptr, size_t num) {
28 using namespace tensorflow;
29 *ptr = nullptr;
30 Tensor* tensor = 0;
31 TensorShape shape({int64_t(num)});
32 OP_REQUIRES_OK(context, context->allocate_output(0, shape, &tensor));
33
34 auto flat_tensor = tensor->flat<TIndex>();
35 *ptr = (TIndex*)flat_tensor.data();
36 }
37
38 void AllocDistances(T** ptr, size_t num) {
39 using namespace tensorflow;
40 *ptr = nullptr;
41 Tensor* tensor = 0;
42 TensorShape shape({int64_t(num)});
43 OP_REQUIRES_OK(context, context->allocate_output(2, shape, &tensor));
44 auto flat_tensor = tensor->flat<T>();
45 *ptr = flat_tensor.data();
46 }
47
48private:
49 tensorflow::OpKernelContext* context;
50};
51
52class FixedRadiusSearchOpKernel : public tensorflow::OpKernel {
53public:
54 explicit FixedRadiusSearchOpKernel(
55 tensorflow::OpKernelConstruction* construction)
56 : OpKernel(construction) {
57 using namespace tensorflow;
58 using namespace open3d::core::nns;
59
60 std::string metric_str;
61 OP_REQUIRES_OK(construction,
62 construction->GetAttr("metric", &metric_str));
63 if (metric_str == "L1")
64 metric = L1;
65 else if (metric_str == "L2")
66 metric = L2;
67 else
68 metric = Linf;
69
70 OP_REQUIRES_OK(construction,
71 construction->GetAttr("ignore_query_point",
72 &ignore_query_point));
73
74 OP_REQUIRES_OK(construction, construction->GetAttr("return_distances",
75 &return_distances));
76 }
77
78 void Compute(tensorflow::OpKernelContext* context) override {
79 using namespace tensorflow;
80 static_assert(sizeof(int64) == sizeof(int64_t),
81 "int64 type is not compatible");
82
83 const Tensor& points = context->input(0);
84 const Tensor& queries = context->input(1);
85
86 const Tensor& radius = context->input(2);
87 OP_REQUIRES(context, TensorShapeUtils::IsScalar(radius.shape()),
88 errors::InvalidArgument("radius must be scalar, got shape ",
89 radius.shape().DebugString()));
90
91 const Tensor& points_row_splits = context->input(3);
92 const Tensor& queries_row_splits = context->input(4);
93
94 const Tensor& hash_table_splits = context->input(5);
95 const Tensor& hash_table_index = context->input(6);
96 const Tensor& hash_table_cell_splits = context->input(7);
97
98 {
99 using namespace open3d::ml::op_util;
100
101 Dim num_points("num_points");
102 Dim num_queries("num_queries");
103 Dim batch_size("batch_size");
104 Dim num_cells("num_cells");
105 CHECK_SHAPE(context, points, num_points, 3);
106 CHECK_SHAPE(context, hash_table_index, num_points);
107 CHECK_SHAPE(context, queries, num_queries, 3);
108 CHECK_SHAPE(context, points_row_splits, batch_size + 1);
109 CHECK_SHAPE(context, queries_row_splits, batch_size + 1);
110 CHECK_SHAPE(context, hash_table_splits, batch_size + 1);
111 CHECK_SHAPE(context, hash_table_cell_splits, num_cells + 1);
112 }
113 Tensor* query_neighbors_row_splits = 0;
114 TensorShape query_neighbors_row_splits_shape(
115 {queries.shape().dim_size(0) + 1});
116 OP_REQUIRES_OK(context, context->allocate_output(
117 1, query_neighbors_row_splits_shape,
118 &query_neighbors_row_splits));
119
120 Kernel(context, points, queries, radius, points_row_splits,
121 queries_row_splits, hash_table_splits, hash_table_index,
122 hash_table_cell_splits, *query_neighbors_row_splits);
123 }
124
125 virtual void Kernel(tensorflow::OpKernelContext* context,
126 const tensorflow::Tensor& points,
127 const tensorflow::Tensor& queries,
128 const tensorflow::Tensor& radius,
129 const tensorflow::Tensor& points_row_splits,
130 const tensorflow::Tensor& queries_row_splits,
131 const tensorflow::Tensor& hash_table_splits,
132 const tensorflow::Tensor& hash_table_index,
133 const tensorflow::Tensor& hash_table_cell_splits,
134 tensorflow::Tensor& query_neighbors_row_splits) = 0;
135
136protected:
138 bool ignore_query_point;
139 bool return_distances;
140};
141} // namespace fixed_radius_search_opkernel
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:186
ImGuiContext * context
Definition: Window.cpp:76
Definition: Tensor.h:32
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:50
int points
Definition: FilePCD.cpp:54
Definition: FixedRadiusIndex.cpp:16
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:19
@ Linf
Definition: NeighborSearchCommon.h:19
@ L1
Definition: NeighborSearchCommon.h:19
@ L2
Definition: NeighborSearchCommon.h:19
Definition: ShapeChecking.h:16