Open3D (C++ API)  0.18.0
Loading...
Searching...
No Matches
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
16class BuildSpatialHashTableOpKernel : public tensorflow::OpKernel {
17public:
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
101protected:
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
Definition ShapeChecking.h:16