Open3D (C++ API)  0.18.0+252c867
ContinuousConvOpKernel.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 
11 #include "tensorflow/core/framework/op.h"
12 #include "tensorflow/core/framework/op_kernel.h"
13 #include "tensorflow/core/lib/core/errors.h"
14 
15 template <class TIndex>
16 class ContinuousConvOpKernel : public tensorflow::OpKernel {
17 public:
19  tensorflow::OpKernelConstruction* construction)
20  : OpKernel(construction) {
21  using namespace tensorflow;
22  using namespace open3d::ml::impl;
23  OP_REQUIRES_OK(construction,
24  construction->GetAttr("align_corners", &align_corners));
25  OP_REQUIRES_OK(construction,
26  construction->GetAttr("normalize", &normalize));
27 
28  std::string interpolation_str;
29  OP_REQUIRES_OK(construction, construction->GetAttr("interpolation",
30  &interpolation_str));
31 
32  if (interpolation_str == "linear")
33  interpolation = InterpolationMode::LINEAR;
34  else if (interpolation_str == "linear_border")
35  interpolation = InterpolationMode::LINEAR_BORDER;
36  else
38 
39  std::string mapping_str;
40  OP_REQUIRES_OK(construction, construction->GetAttr("coordinate_mapping",
41  &mapping_str));
42 
43  if (mapping_str == "ball_to_cube_radial")
44  coordinate_mapping = CoordinateMapping::BALL_TO_CUBE_RADIAL;
45  else if (mapping_str == "ball_to_cube_volume_preserving")
47  CoordinateMapping::BALL_TO_CUBE_VOLUME_PRESERVING;
48  else
49  coordinate_mapping = CoordinateMapping::IDENTITY;
50 
51  OP_REQUIRES_OK(construction, construction->GetAttr("max_temp_mem_MB",
52  &max_temp_mem_MB));
53  }
54 
55  void Compute(tensorflow::OpKernelContext* context) override {
56  using namespace tensorflow;
57  static_assert(sizeof(int64) == sizeof(int64_t),
58  "int64 type is not compatible");
59  const Tensor& filter = context->input(0);
60 
61  const Tensor& out_positions = context->input(1);
62  OP_REQUIRES(context,
63  out_positions.shape().dim_size(0) <=
64  std::numeric_limits<TIndex>::max(),
65  errors::InvalidArgument("Too many output points"));
66 
67  const Tensor& extents = context->input(2);
68  OP_REQUIRES(context, extents.shape().dims() == 2,
69  errors::InvalidArgument("extents must be a rank 2 tensor"));
70  OP_REQUIRES(context,
71  extents.shape().dim_size(0) ==
72  out_positions.shape().dim_size(0) ||
73  extents.shape().dim_size(0) == 1,
74  errors::InvalidArgument("number of extents must match the "
75  "number of out_positions or must "
76  "be 1"));
77  OP_REQUIRES(context,
78  extents.shape().dim_size(1) == 3 ||
79  extents.shape().dim_size(1) == 1,
80  errors::InvalidArgument(
81  "number of components for extents must be 3 or 1"));
82 
83  const Tensor& offset = context->input(3);
84  OP_REQUIRES(context, offset.shape().dims() == 1,
85  errors::InvalidArgument("offset must be a rank 1 tensor"));
86  OP_REQUIRES(context, offset.shape().dim_size(0) == 3,
87  errors::InvalidArgument("offset length must be 3"));
88 
89  const Tensor& inp_positions = context->input(4);
90  OP_REQUIRES(context,
91  inp_positions.shape().dim_size(0) <=
92  std::numeric_limits<TIndex>::max(),
93  errors::InvalidArgument("Too many input points"));
94 
95  const Tensor& inp_features = context->input(5);
96 
97  const Tensor& inp_importance = context->input(6);
98 
99  const Tensor& neighbors_index = context->input(7);
100 
101  const Tensor& neighbors_importance = context->input(8);
102 
103  const Tensor& neighbors_row_splits = context->input(9);
104 
105  OP_REQUIRES(
106  context,
107  inp_positions.shape().dim_size(0) ==
108  inp_features.shape().dim_size(0),
109  errors::InvalidArgument("first dim of inp_positions does not "
110  "match the first dim of inp_features"));
111 
112  OP_REQUIRES(context,
113  inp_positions.shape().dim_size(0) ==
114  inp_importance.shape().dim_size(0) ||
115  inp_importance.shape().dim_size(0) == 0,
116  errors::InvalidArgument("first dim of inp_positions does "
117  "not match the first dim of "
118  "inp_importance"));
119 
120  OP_REQUIRES(context,
121  neighbors_importance.shape().dim_size(0) ==
122  neighbors_index.shape().dim_size(0) ||
123  neighbors_importance.shape().dim_size(0) == 0,
124  errors::InvalidArgument("first dim of neighbors_importance "
125  "does not match the first dim of "
126  "neighbors_index"));
127 
128  OP_REQUIRES(
129  context,
130  filter.shape().dim_size(3) == inp_features.shape().dim_size(1),
131  errors::InvalidArgument("number of input channels in filter "
132  "and inp_features does not match"));
133 
134  TensorShape out_features_shape({out_positions.shape().dim_size(0),
135  filter.shape().dim_size(4)});
136  Tensor* out_features = nullptr;
137  OP_REQUIRES_OK(context, context->allocate_output(0, out_features_shape,
138  &out_features));
139 
140  std::vector<int> filter_dims({
141  int(filter.shape().dim_size(0)),
142  int(filter.shape().dim_size(1)),
143  int(filter.shape().dim_size(2)),
144  int(filter.shape().dim_size(3)),
145  int(filter.shape().dim_size(4)),
146  });
147 
148  bool individual_extents = extents.shape().dim_size(0) ==
149  out_positions.shape().dim_size(0) &&
150  extents.shape().dim_size(0) > 1;
151 
152  bool isotropic_extents = extents.shape().dim_size(1) == 1;
153 
154  bool point_importances = inp_importance.shape().dim_size(0) != 0;
155 
156  bool has_neighbors_importances =
157  neighbors_importance.shape().dim_size(0) != 0;
158 
159  Kernel(context, filter, out_positions, extents, offset, inp_positions,
160  inp_features, inp_importance, neighbors_index,
161  neighbors_importance, neighbors_row_splits, filter_dims,
162  individual_extents, isotropic_extents, point_importances,
163  has_neighbors_importances, *out_features);
164  }
165 
166  virtual void Kernel(tensorflow::OpKernelContext* context,
167  const tensorflow::Tensor& filter,
168  const tensorflow::Tensor& out_positions,
169  const tensorflow::Tensor& extents,
170  const tensorflow::Tensor& offset,
171  const tensorflow::Tensor& inp_positions,
172  const tensorflow::Tensor& inp_features,
173  const tensorflow::Tensor& inp_importance,
174  const tensorflow::Tensor& neighbors_index,
175  const tensorflow::Tensor& neighbors_importance,
176  const tensorflow::Tensor& neighbors_row_splits,
177  const std::vector<int>& filter_dims,
178  const bool individual_extents,
179  const bool isotropic_extents,
180  const bool point_importances,
181  const bool has_neighbors_importances,
182  tensorflow::Tensor& out_features) = 0;
183 
184 public:
186  bool normalize;
190 };
ImGuiContext * context
Definition: Window.cpp:76
Definition: ContinuousConvOpKernel.h:16
bool normalize
Definition: ContinuousConvOpKernel.h:186
open3d::ml::impl::InterpolationMode interpolation
Definition: ContinuousConvOpKernel.h:187
void Compute(tensorflow::OpKernelContext *context) override
Definition: ContinuousConvOpKernel.h:55
bool align_corners
Definition: ContinuousConvOpKernel.h:185
virtual void Kernel(tensorflow::OpKernelContext *context, const tensorflow::Tensor &filter, const tensorflow::Tensor &out_positions, const tensorflow::Tensor &extents, const tensorflow::Tensor &offset, const tensorflow::Tensor &inp_positions, const tensorflow::Tensor &inp_features, const tensorflow::Tensor &inp_importance, const tensorflow::Tensor &neighbors_index, const tensorflow::Tensor &neighbors_importance, const tensorflow::Tensor &neighbors_row_splits, const std::vector< int > &filter_dims, const bool individual_extents, const bool isotropic_extents, const bool point_importances, const bool has_neighbors_importances, tensorflow::Tensor &out_features)=0
ContinuousConvOpKernel(tensorflow::OpKernelConstruction *construction)
Definition: ContinuousConvOpKernel.h:18
int max_temp_mem_MB
Definition: ContinuousConvOpKernel.h:189
open3d::ml::impl::CoordinateMapping coordinate_mapping
Definition: ContinuousConvOpKernel.h:188
int offset
Definition: FilePCD.cpp:45
const char const char value recording_handle imu_sample recording_handle uint8_t size_t data_size k4a_record_configuration_t config target_format k4a_capture_t capture_handle k4a_imu_sample_t imu_sample playback_handle k4a_logging_message_cb_t void min_level device_handle k4a_imu_sample_t timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c int
Definition: K4aPlugin.cpp:474
Definition: ContinuousConv.h:16
InterpolationMode
Definition: ContinuousConvTypes.h:18
@ NEAREST_NEIGHBOR
Definition: VoxelPooling.h:21
CoordinateMapping
Definition: ContinuousConvTypes.h:26