Open3D (C++ API)  0.12.0
VoxelizeOpKernel.h
Go to the documentation of this file.
1 // The MIT License (MIT)
2 //
3 // Copyright (c) 2020 www.open3d.org
4 //
5 // Permission is hereby granted, free of charge, to any person obtaining a copy
6 // of this software and associated documentation files (the "Software"), to deal
7 // in the Software without restriction, including without limitation the rights
8 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9 // copies of the Software, and to permit persons to whom the Software is
10 // furnished to do so, subject to the following conditions:
11 //
12 // The above copyright notice and this permission notice shall be included in
13 // all copies or substantial portions of the Software.
14 //
15 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
20 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
21 // IN THE SOFTWARE.
22 // ----------------------------------------------------------------------------
23 #pragma once
24 
25 //#include "open3d/ml/impl/misc/VoxelPooling.h"
27 #include "tensorflow/core/framework/op.h"
28 #include "tensorflow/core/framework/op_kernel.h"
29 #include "tensorflow/core/lib/core/errors.h"
30 
32 // namespace for code that is common for all kernels
33 namespace voxelize_opkernel {
34 
35 class OutputAllocator {
36 public:
37  OutputAllocator(tensorflow::OpKernelContext* context) : context(context) {}
38 
39  void AllocVoxelCoords(int32_t** ptr, int64_t rows, int64_t cols) {
40  using namespace tensorflow;
41  *ptr = nullptr;
42  Tensor* tensor = 0;
43  TensorShape shape({rows, cols});
44  OP_REQUIRES_OK(context, context->allocate_output(0, shape, &tensor));
45  auto flat_tensor = tensor->flat<int32_t>();
46  *ptr = flat_tensor.data();
47  }
48 
49  void AllocVoxelPointIndices(int64_t** ptr, int64_t num) {
50  using namespace tensorflow;
51  *ptr = nullptr;
52  Tensor* tensor = 0;
53  TensorShape shape({num});
54  OP_REQUIRES_OK(context, context->allocate_output(1, shape, &tensor));
55  auto flat_tensor = tensor->flat<int64>();
56  *ptr = (int64_t*)flat_tensor.data();
57  }
58 
59  void AllocVoxelPointRowSplits(int64_t** ptr, int64_t num) {
60  using namespace tensorflow;
61  *ptr = nullptr;
62  Tensor* tensor = 0;
63  TensorShape shape({num});
64  OP_REQUIRES_OK(context, context->allocate_output(2, shape, &tensor));
65  auto flat_tensor = tensor->flat<int64>();
66  *ptr = (int64_t*)flat_tensor.data();
67  }
68 
69 private:
70  tensorflow::OpKernelContext* context;
71 };
72 
73 // Base class with common code for the OpKernel implementations
74 class VoxelizeOpKernel : public tensorflow::OpKernel {
75 public:
76  explicit VoxelizeOpKernel(tensorflow::OpKernelConstruction* construction)
77  : OpKernel(construction) {
78  OP_REQUIRES_OK(construction,
79  construction->GetAttr("max_points_per_voxel",
80  &max_points_per_voxel));
81  OP_REQUIRES_OK(construction,
82  construction->GetAttr("max_voxels", &max_voxels));
83  }
84 
85  void Compute(tensorflow::OpKernelContext* context) override {
86  using namespace tensorflow;
87  const Tensor& points = context->input(0);
88  const Tensor& voxel_size = context->input(1);
89  const Tensor& points_range_min = context->input(2);
90  const Tensor& points_range_max = context->input(3);
91 
92  {
93  using namespace open3d::ml::op_util;
94  Dim num_points("num_points");
95  Dim ndim("ndim");
96  CHECK_SHAPE(context, points, num_points, ndim);
97  CHECK_SHAPE(context, voxel_size, ndim);
98  CHECK_SHAPE(context, points_range_min, ndim);
99  CHECK_SHAPE(context, points_range_max, ndim);
100  OP_REQUIRES(
101  context, ndim.value() > 0 && ndim.value() < 9,
102  errors::InvalidArgument(
103  "the number of dimensions must be in [1,..,8]"));
104  }
105 
106  Kernel(context, points, voxel_size, points_range_min, points_range_max);
107  }
108 
109  // Function with the device specific code
110  virtual void Kernel(tensorflow::OpKernelContext* context,
111  const tensorflow::Tensor& points,
112  const tensorflow::Tensor& voxel_size,
113  const tensorflow::Tensor& points_range_min,
114  const tensorflow::Tensor& points_range_max) = 0;
115 
116 protected:
117  tensorflow::int64 max_points_per_voxel;
118  tensorflow::int64 max_voxels;
119 };
120 
121 } // namespace voxelize_opkernel
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:204
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 int32_t
Definition: K4aPlugin.cpp:398
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:69
int points
Definition: FilePCD.cpp:73
Definition: ShapeChecking.h:35