Open3D (C++ API)  0.18.0+252c867
BuildSpatialHashTableOpKernel.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"
12 #include "tensorflow/core/framework/op.h"
13 #include "tensorflow/core/framework/op_kernel.h"
14 #include "tensorflow/core/lib/core/errors.h"
15 
16 class BuildSpatialHashTableOpKernel : public tensorflow::OpKernel {
17 public:
19  tensorflow::OpKernelConstruction* construction)
20  : OpKernel(construction) {
21  using namespace tensorflow;
22 
23  OP_REQUIRES_OK(construction,
24  construction->GetAttr("max_hash_table_size",
26  }
27 
28  void Compute(tensorflow::OpKernelContext* context) override {
29  using namespace tensorflow;
30  using namespace open3d::ml::op_util;
31 
32  const Tensor& points = context->input(0);
33  const Tensor& radius = context->input(1);
34  OP_REQUIRES(context, TensorShapeUtils::IsScalar(radius.shape()),
35  errors::InvalidArgument("radius must be scalar, got shape ",
36  radius.shape().DebugString()));
37 
38  const Tensor& points_row_splits = context->input(2);
39 
40  const Tensor& hash_table_size_factor_tensor = context->input(3);
41  OP_REQUIRES(
42  context,
43  TensorShapeUtils::IsScalar(
44  hash_table_size_factor_tensor.shape()),
45  errors::InvalidArgument(
46  "hash_table_size_factor must be scalar, got shape ",
47  hash_table_size_factor_tensor.shape().DebugString()));
48  const double hash_table_size_factor =
49  hash_table_size_factor_tensor.scalar<double>()();
50 
51  Dim num_points("num_points");
52  Dim batch_size("batch_size");
53  CHECK_SHAPE(context, points, num_points, 3);
54  CHECK_SHAPE(context, points_row_splits, batch_size + 1);
55 
56  std::vector<uint32_t> hash_table_splits(batch_size.value() + 1, 0);
57  for (int i = 0; i < batch_size.value(); ++i) {
58  int64_t num_points_i = points_row_splits.flat<int64>()(i + 1) -
59  points_row_splits.flat<int64>()(i);
60 
61  int64_t hash_table_size = std::min<int64_t>(
62  std::max<int64_t>(hash_table_size_factor * num_points_i, 1),
64  hash_table_splits[i + 1] = hash_table_splits[i] + hash_table_size;
65  }
66 
67  Tensor* hash_table_index = 0;
68  TensorShape hash_table_index_shape({num_points.value()});
69  OP_REQUIRES_OK(context,
70  context->allocate_output(0, hash_table_index_shape,
71  &hash_table_index));
72 
73  Tensor* hash_table_cell_splits = 0;
74  TensorShape hash_table_cell_splits_shape(
75  {hash_table_splits.back() + 1});
76  OP_REQUIRES_OK(context,
77  context->allocate_output(1, hash_table_cell_splits_shape,
78  &hash_table_cell_splits));
79 
80  Tensor* out_hash_table_splits = 0;
81  TensorShape out_hash_table_splits_shape({batch_size.value() + 1});
82  OP_REQUIRES_OK(context,
83  context->allocate_output(2, out_hash_table_splits_shape,
84  &out_hash_table_splits));
85  for (size_t i = 0; i < hash_table_splits.size(); ++i) {
86  out_hash_table_splits->flat<uint32_t>()(i) = hash_table_splits[i];
87  }
88 
89  Kernel(context, points, radius, points_row_splits, hash_table_splits,
90  *hash_table_index, *hash_table_cell_splits);
91  }
92 
93  virtual void Kernel(tensorflow::OpKernelContext* context,
94  const tensorflow::Tensor& points,
95  const tensorflow::Tensor& radius,
96  const tensorflow::Tensor& points_row_splits,
97  const std::vector<uint32_t>& hash_table_splits,
98  tensorflow::Tensor& hash_table_index,
99  tensorflow::Tensor& hash_table_cell_splits) = 0;
100 
101 protected:
103 };
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:186
ImGuiContext * context
Definition: Window.cpp:76
Definition: BuildSpatialHashTableOpKernel.h:16
int max_hash_table_size
Definition: BuildSpatialHashTableOpKernel.h:102
virtual void Kernel(tensorflow::OpKernelContext *context, const tensorflow::Tensor &points, const tensorflow::Tensor &radius, const tensorflow::Tensor &points_row_splits, const std::vector< uint32_t > &hash_table_splits, tensorflow::Tensor &hash_table_index, tensorflow::Tensor &hash_table_cell_splits)=0
void Compute(tensorflow::OpKernelContext *context) override
Definition: BuildSpatialHashTableOpKernel.h:28
BuildSpatialHashTableOpKernel(tensorflow::OpKernelConstruction *construction)
Definition: BuildSpatialHashTableOpKernel.h:18
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:50
int64_t & value()
Definition: ShapeChecking.h:70
int points
Definition: FilePCD.cpp:54
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 k4a_image_t image_handle uint8_t image_handle image_handle image_handle image_handle uint32_t
Definition: K4aPlugin.cpp:548
Definition: ShapeChecking.h:16