Open3D (C++ API)  0.16.1
BallQueryOpKernel.h
Go to the documentation of this file.
1// ----------------------------------------------------------------------------
2// - Open3D: www.open3d.org -
3// ----------------------------------------------------------------------------
4// The MIT License (MIT)
5//
6// Copyright (c) 2018-2021 www.open3d.org
7//
8// Permission is hereby granted, free of charge, to any person obtaining a copy
9// of this software and associated documentation files (the "Software"), to deal
10// in the Software without restriction, including without limitation the rights
11// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12// copies of the Software, and to permit persons to whom the Software is
13// furnished to do so, subject to the following conditions:
14//
15// The above copyright notice and this permission notice shall be included in
16// all copies or substantial portions of the Software.
17//
18// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
23// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
24// IN THE SOFTWARE.
25// ----------------------------------------------------------------------------
26
27#pragma once
28
29#include "../TensorFlowHelper.h"
30#include "tensorflow/core/framework/op.h"
31#include "tensorflow/core/framework/op_kernel.h"
32#include "tensorflow/core/lib/core/errors.h"
33
34class BallQueryOpKernel : public tensorflow::OpKernel {
35public:
36 explicit BallQueryOpKernel(tensorflow::OpKernelConstruction* construction)
37 : OpKernel(construction) {
38 using namespace tensorflow;
39
40 OP_REQUIRES_OK(construction,
41 construction->GetAttr("nsample", &nsample));
42 OP_REQUIRES_OK(construction, construction->GetAttr("radius", &radius));
43 OP_REQUIRES(
44 construction, nsample > 0,
45 errors::InvalidArgument("BallQuery expects positive nsample"));
46 }
47
48 void Compute(tensorflow::OpKernelContext* context) override {
49 using namespace tensorflow;
50
51 const Tensor& inp_tensor = context->input(0);
52 OP_REQUIRES(
53 context,
54 inp_tensor.dims() == 3 && inp_tensor.shape().dim_size(2) == 3,
55 errors::InvalidArgument("BallQuery expects "
56 "(batch_size,num_points,3) inp shape"));
57 int batch_size = inp_tensor.shape().dim_size(0);
58 int pts_size = inp_tensor.shape().dim_size(1);
59 auto inp_flat = inp_tensor.flat<float>();
60 const float* inp = &(inp_flat(0));
61
62 const Tensor& center_tensor = context->input(1);
63 OP_REQUIRES(context,
64 center_tensor.dims() == 3 &&
65 center_tensor.shape().dim_size(2) == 3,
66 errors::InvalidArgument(
67 "BallQuery expects "
68 "(batch_size,num_points,3) center shape"));
69 int ball_size = center_tensor.shape().dim_size(1);
70 auto center_flat = center_tensor.flat<float>();
71 const float* center = &(center_flat(0));
72
73 Tensor* out_tensor;
74 OP_REQUIRES_OK(context,
75 context->allocate_output(
76 0, TensorShape{batch_size, ball_size, nsample},
77 &out_tensor));
78 auto out_flat = out_tensor->flat<int>();
79 int* out = &(out_flat(0));
80
81 Kernel(context, batch_size, pts_size, ball_size, radius, nsample,
82 center, inp, out);
83 }
84
85 virtual void Kernel(tensorflow::OpKernelContext* context,
86 int b,
87 int n,
88 int m,
89 float radius,
90 int nsample,
91 const float* new_xyz,
92 const float* xyz,
93 int* idx) = 0;
94
95protected:
97 float radius;
98};
ImGuiContext * context
Definition: Window.cpp:95
Definition: BallQueryOpKernel.h:34
int nsample
Definition: BallQueryOpKernel.h:96
BallQueryOpKernel(tensorflow::OpKernelConstruction *construction)
Definition: BallQueryOpKernel.h:36
void Compute(tensorflow::OpKernelContext *context) override
Definition: BallQueryOpKernel.h:48
virtual void Kernel(tensorflow::OpKernelContext *context, int b, int n, int m, float radius, int nsample, const float *new_xyz, const float *xyz, int *idx)=0
float radius
Definition: BallQueryOpKernel.h:97