27 #include "tensorflow/core/framework/op.h" 28 #include "tensorflow/core/framework/op_kernel.h" 29 #include "tensorflow/core/lib/core/errors.h" 35 class OutputAllocator {
37 OutputAllocator(tensorflow::OpKernelContext* context) : context(context) {}
39 void AllocVoxelCoords(
int32_t** ptr, int64_t rows, int64_t cols) {
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();
49 void AllocVoxelPointIndices(int64_t** ptr, int64_t num) {
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();
59 void AllocVoxelPointRowSplits(int64_t** ptr, int64_t num) {
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();
70 tensorflow::OpKernelContext* context;
74 class VoxelizeOpKernel :
public tensorflow::OpKernel {
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));
85 void Compute(tensorflow::OpKernelContext* context)
override {
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);
94 Dim num_points(
"num_points");
101 context, ndim.value() > 0 && ndim.value() < 9,
102 errors::InvalidArgument(
103 "the number of dimensions must be in [1,..,8]"));
106 Kernel(context, points, voxel_size, points_range_min, points_range_max);
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;
117 tensorflow::int64 max_points_per_voxel;
118 tensorflow::int64 max_voxels;
#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