Open3D (C++ API)  0.12.0
VoxelizeOpKernel.h
Go to the documentation of this file.
1 // ----------------------------------------------------------------------------
2 // - Open3D: www.open3d.org -
3 // ----------------------------------------------------------------------------
4 // The MIT License (MIT)
5 //
6 // Copyright (c) 2020 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 
30 #include "torch/script.h"
31 
32 template <class T>
33 void VoxelizeCPU(const torch::Tensor& points,
34  const torch::Tensor& voxel_size,
35  const torch::Tensor& points_range_min,
36  const torch::Tensor& points_range_max,
37  const int64_t max_points_per_voxel,
38  const int64_t max_voxels,
39  torch::Tensor& voxel_coords,
40  torch::Tensor& voxel_point_indices,
41  torch::Tensor& voxel_point_row_splits);
42 
43 #ifdef BUILD_CUDA_MODULE
44 template <class T>
45 void VoxelizeCUDA(const torch::Tensor& points,
46  const torch::Tensor& voxel_size,
47  const torch::Tensor& points_range_min,
48  const torch::Tensor& points_range_max,
49  const int64_t max_points_per_voxel,
50  const int64_t max_voxels,
51  torch::Tensor& voxel_coords,
52  torch::Tensor& voxel_point_indices,
53  torch::Tensor& voxel_point_row_splits);
54 #endif
55 
57 public:
58  VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
59  : device_type(device_type), device_idx(device_idx) {}
60 
61  void AllocVoxelCoords(int32_t** ptr, int64_t rows, int64_t cols) {
62  voxel_coords = torch::empty({rows, cols},
63  torch::dtype(ToTorchDtype<int32_t>())
64  .device(device_type, device_idx));
65  *ptr = voxel_coords.data_ptr<int32_t>();
66  }
67 
68  void AllocVoxelPointIndices(int64_t** ptr, int64_t num) {
69  voxel_point_indices =
70  torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
71  .device(device_type, device_idx));
72  *ptr = voxel_point_indices.data_ptr<int64_t>();
73  }
74 
75  void AllocVoxelPointRowSplits(int64_t** ptr, int64_t num) {
76  voxel_point_row_splits =
77  torch::empty({num}, torch::dtype(ToTorchDtype<int64_t>())
78  .device(device_type, device_idx));
79  *ptr = voxel_point_row_splits.data_ptr<int64_t>();
80  }
81 
82  const torch::Tensor& VoxelCoords() const { return voxel_coords; }
83  const torch::Tensor& VoxelPointIndices() const {
84  return voxel_point_indices;
85  }
86  const torch::Tensor& VoxelPointRowSplits() const {
87  return voxel_point_row_splits;
88  }
89 
90 private:
91  torch::Tensor voxel_coords;
92  torch::Tensor voxel_point_indices;
93  torch::Tensor voxel_point_row_splits;
94  torch::DeviceType device_type;
95  int device_idx;
96 };
Definition: VoxelizeOpKernel.h:56
VoxelizeOutputAllocator(torch::DeviceType device_type, int device_idx)
Definition: VoxelizeOpKernel.h:58
TorchDtype_t ToTorchDtype< int64_t >()
Definition: TorchHelper.h:94
const torch::Tensor & VoxelCoords() const
Definition: VoxelizeOpKernel.h:82
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
void AllocVoxelPointRowSplits(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:75
void AllocVoxelPointIndices(int64_t **ptr, int64_t num)
Definition: VoxelizeOpKernel.h:68
const torch::Tensor & VoxelPointIndices() const
Definition: VoxelizeOpKernel.h:83
int points
Definition: FilePCD.cpp:73
void VoxelizeCPU(const torch::Tensor &points, const torch::Tensor &voxel_size, const torch::Tensor &points_range_min, const torch::Tensor &points_range_max, const int64_t max_points_per_voxel, const int64_t max_voxels, torch::Tensor &voxel_coords, torch::Tensor &voxel_point_indices, torch::Tensor &voxel_point_row_splits)
Definition: VoxelizeOpKernel.cpp:37
void AllocVoxelCoords(int32_t **ptr, int64_t rows, int64_t cols)
Definition: VoxelizeOpKernel.h:61
TorchDtype_t ToTorchDtype< int32_t >()
Definition: TorchHelper.h:90
const torch::Tensor & VoxelPointRowSplits() const
Definition: VoxelizeOpKernel.h:86