Open3D (C++ API)  0.12.0
GeneralEWSharedImpl.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 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 #include <atomic>
28 
29 #include "open3d/core/Dispatch.h"
30 #include "open3d/core/Dtype.h"
31 #include "open3d/core/MemoryManager.h"
32 #include "open3d/core/SizeVector.h"
33 #include "open3d/core/Tensor.h"
34 #include "open3d/core/kernel/GeneralEW.h"
35 #include "open3d/core/kernel/GeneralEWMacros.h"
36 #include "open3d/core/kernel/GeneralIndexer.h"
37 #include "open3d/utility/Console.h"
38 
39 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
40 #define OPEN3D_ATOMIC_ADD(X, Y) atomicAdd(X, Y)
41 #else
42 #define OPEN3D_ATOMIC_ADD(X, Y) (*X).fetch_add(Y)
43 #endif
44 
45 #define DISPATCH_BYTESIZE_TO_VOXEL(BYTESIZE, ...) \
46  [&] { \
47  if (BYTESIZE == sizeof(ColoredVoxel32f)) { \
48  using voxel_t = ColoredVoxel32f; \
49  return __VA_ARGS__(); \
50  } else if (BYTESIZE == sizeof(ColoredVoxel16i)) { \
51  using voxel_t = ColoredVoxel16i; \
52  return __VA_ARGS__(); \
53  } else if (BYTESIZE == sizeof(Voxel32f)) { \
54  using voxel_t = Voxel32f; \
55  return __VA_ARGS__(); \
56  } else { \
57  utility::LogError("Unsupported voxel bytesize"); \
58  } \
59  }()
60 
61 namespace open3d {
62 namespace core {
63 namespace kernel {
64 
68 struct Voxel32f {
69  float tsdf;
70  float weight;
71 
72  static bool HasColor() { return false; }
73  OPEN3D_HOST_DEVICE float GetTSDF() { return tsdf; }
74  OPEN3D_HOST_DEVICE float GetWeight() { return static_cast<float>(weight); }
75  OPEN3D_HOST_DEVICE float GetR() { return 1.0; }
76  OPEN3D_HOST_DEVICE float GetG() { return 1.0; }
77  OPEN3D_HOST_DEVICE float GetB() { return 1.0; }
78 
79  OPEN3D_HOST_DEVICE void Integrate(float dsdf) {
80  tsdf = (weight * tsdf + dsdf) / (weight + 1);
81  weight += 1;
82  }
83  OPEN3D_HOST_DEVICE void Integrate(float dsdf,
84  float dr,
85  float dg,
86  float db) {
87  printf("[Voxel32f] should never reach here.\n");
88  }
89 };
90 
96 
97 struct ColoredVoxel16i {
98  static const uint16_t kMaxUint16 = 65535;
99  static constexpr float kColorFactor = 255.0f;
100 
101  float tsdf;
102  uint16_t weight;
103 
104  uint16_t r;
105  uint16_t g;
106  uint16_t b;
107 
108  static bool HasColor() { return true; }
109  OPEN3D_HOST_DEVICE float GetTSDF() { return tsdf; }
110  OPEN3D_HOST_DEVICE float GetWeight() { return static_cast<float>(weight); }
111  OPEN3D_HOST_DEVICE float GetR() {
112  return static_cast<float>(r / kColorFactor);
113  }
114  OPEN3D_HOST_DEVICE float GetG() {
115  return static_cast<float>(g / kColorFactor);
116  }
117  OPEN3D_HOST_DEVICE float GetB() {
118  return static_cast<float>(b / kColorFactor);
119  }
120  OPEN3D_HOST_DEVICE void Integrate(float dsdf) {
121  float inc_wsum = static_cast<float>(weight) + 1;
122  float inv_wsum = 1.0f / inc_wsum;
123  tsdf = (static_cast<float>(weight) * tsdf + dsdf) * inv_wsum;
124  weight = static_cast<uint16_t>(inc_wsum < static_cast<float>(kMaxUint16)
125  ? weight + 1
126  : kMaxUint16);
127  }
128  OPEN3D_HOST_DEVICE void Integrate(float dsdf,
129  float dr,
130  float dg,
131  float db) {
132  float inc_wsum = static_cast<float>(weight) + 1;
133  float inv_wsum = 1.0f / inc_wsum;
134  tsdf = (weight * tsdf + dsdf) * inv_wsum;
135  r = static_cast<uint16_t>(
136  round((weight * r + dr * kColorFactor) * inv_wsum));
137  g = static_cast<uint16_t>(
138  round((weight * g + dg * kColorFactor) * inv_wsum));
139  b = static_cast<uint16_t>(
140  round((weight * b + db * kColorFactor) * inv_wsum));
141  weight = static_cast<uint16_t>(inc_wsum < static_cast<float>(kMaxUint16)
142  ? weight + 1
143  : kMaxUint16);
144  }
145 };
146 
149 struct ColoredVoxel32f {
150  float tsdf;
151  float weight;
152 
153  float r;
154  float g;
155  float b;
156 
157  static bool HasColor() { return true; }
158  OPEN3D_HOST_DEVICE float GetTSDF() { return tsdf; }
159  OPEN3D_HOST_DEVICE float GetWeight() { return weight; }
160  OPEN3D_HOST_DEVICE float GetR() { return r; }
161  OPEN3D_HOST_DEVICE float GetG() { return g; }
162  OPEN3D_HOST_DEVICE float GetB() { return b; }
163  OPEN3D_HOST_DEVICE void Integrate(float dsdf) {
164  float inv_wsum = 1.0f / (weight + 1);
165  tsdf = (weight * tsdf + dsdf) * inv_wsum;
166  weight += 1;
167  }
168  OPEN3D_HOST_DEVICE void Integrate(float dsdf,
169  float dr,
170  float dg,
171  float db) {
172  float inv_wsum = 1.0f / (weight + 1);
173  tsdf = (weight * tsdf + dsdf) * inv_wsum;
174  r = (weight * r + dr) * inv_wsum;
175  g = (weight * g + dg) * inv_wsum;
176  b = (weight * b + db) * inv_wsum;
177 
178  weight += 1;
179  }
180 };
181 
182 // Get a voxel in a certain voxel block given the block id with its neighbors.
183 template <typename voxel_t>
184 inline OPEN3D_DEVICE voxel_t* DeviceGetVoxelAt(
185  int xo,
186  int yo,
187  int zo,
188  int curr_block_idx,
189  int resolution,
190  const NDArrayIndexer& nb_block_masks_indexer,
191  const NDArrayIndexer& nb_block_indices_indexer,
192  const NDArrayIndexer& blocks_indexer) {
193  int xn = (xo + resolution) % resolution;
194  int yn = (yo + resolution) % resolution;
195  int zn = (zo + resolution) % resolution;
196 
197  int64_t dxb = sign(xo - xn);
198  int64_t dyb = sign(yo - yn);
199  int64_t dzb = sign(zo - zn);
200 
201  int64_t nb_idx = (dxb + 1) + (dyb + 1) * 3 + (dzb + 1) * 9;
202 
203  bool block_mask_i = *static_cast<bool*>(
204  nb_block_masks_indexer.GetDataPtrFromCoord(curr_block_idx, nb_idx));
205  if (!block_mask_i) return nullptr;
206 
207  int64_t block_idx_i =
208  *static_cast<int64_t*>(nb_block_indices_indexer.GetDataPtrFromCoord(
209  curr_block_idx, nb_idx));
210 
211  return static_cast<voxel_t*>(
212  blocks_indexer.GetDataPtrFromCoord(xn, yn, zn, block_idx_i));
213 }
214 
215 // Get TSDF gradient as normal in a certain voxel block given the block id with
216 // its neighbors.
217 template <typename voxel_t>
218 inline OPEN3D_DEVICE void DeviceGetNormalAt(
219  int xo,
220  int yo,
221  int zo,
222  int curr_block_idx,
223  float* n,
224  int resolution,
225  float voxel_size,
226  const NDArrayIndexer& nb_block_masks_indexer,
227  const NDArrayIndexer& nb_block_indices_indexer,
228  const NDArrayIndexer& blocks_indexer) {
229  auto GetVoxelAt = [&] OPEN3D_DEVICE(int xo, int yo, int zo) {
230  return DeviceGetVoxelAt<voxel_t>(
231  xo, yo, zo, curr_block_idx, resolution, nb_block_masks_indexer,
232  nb_block_indices_indexer, blocks_indexer);
233  };
234  voxel_t* vxp = GetVoxelAt(xo + 1, yo, zo);
235  voxel_t* vxn = GetVoxelAt(xo - 1, yo, zo);
236  voxel_t* vyp = GetVoxelAt(xo, yo + 1, zo);
237  voxel_t* vyn = GetVoxelAt(xo, yo - 1, zo);
238  voxel_t* vzp = GetVoxelAt(xo, yo, zo + 1);
239  voxel_t* vzn = GetVoxelAt(xo, yo, zo - 1);
240  if (vxp && vxn) n[0] = (vxp->GetTSDF() - vxn->GetTSDF()) / (2 * voxel_size);
241  if (vyp && vyn) n[1] = (vyp->GetTSDF() - vyn->GetTSDF()) / (2 * voxel_size);
242  if (vzp && vzn) n[2] = (vzp->GetTSDF() - vzn->GetTSDF()) / (2 * voxel_size);
243 };
244 
245 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
246 void CUDAUnprojectKernel
247 #else
248 void CPUUnprojectKernel
249 #endif
250  (const std::unordered_map<std::string, Tensor>& srcs,
251  std::unordered_map<std::string, Tensor>& dsts) {
252  static std::vector<std::string> src_attrs = {
253  "depth", "intrinsics", "depth_scale", "depth_max", "stride",
254  };
255  for (auto& k : src_attrs) {
256  if (srcs.count(k) == 0) {
257  utility::LogError(
258  "[UnprojectKernel] expected Tensor {} in srcs, but "
259  "did not receive",
260  k);
261  }
262  }
263 
264  // Input
265  Tensor depth = srcs.at("depth");
266  Tensor intrinsics = srcs.at("intrinsics");
267  Tensor extrinsics = srcs.at("extrinsics");
268  float depth_scale = srcs.at("depth_scale").Item<float>();
269  float depth_max = srcs.at("depth_max").Item<float>();
270  int64_t stride = srcs.at("stride").Item<int64_t>();
271 
272  NDArrayIndexer depth_indexer(depth, 2);
273  TransformIndexer ti(intrinsics, extrinsics.Inverse(), 1.0f);
274 
275  // Output
276  int64_t rows_strided = depth_indexer.GetShape(0) / stride;
277  int64_t cols_strided = depth_indexer.GetShape(1) / stride;
278 
279  Tensor points({rows_strided * cols_strided, 3}, core::Dtype::Float32,
280  depth.GetDevice());
281  NDArrayIndexer point_indexer(points, 1);
282 
283  // Counter
284 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
285  Tensor count(std::vector<int>{0}, {}, core::Dtype::Int32,
286  depth.GetDevice());
287  int* count_ptr = static_cast<int*>(count.GetDataPtr());
288 #else
289  std::atomic<int> count_atomic(0);
290  std::atomic<int>* count_ptr = &count_atomic;
291 #endif
292 
293  int64_t n = rows_strided * cols_strided;
294 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
295  CUDALauncher::LaunchGeneralKernel(
296  n, [=] OPEN3D_DEVICE(int64_t workload_idx) {
297 #else
298  CPULauncher::LaunchGeneralKernel(n, [&](int64_t workload_idx) {
299 #endif
300  int64_t y = (workload_idx / cols_strided) * stride;
301  int64_t x = (workload_idx % cols_strided) * stride;
302 
303  float d = (*static_cast<uint16_t*>(
304  depth_indexer.GetDataPtrFromCoord(x, y))) /
305  depth_scale;
306  if (d > 0 && d < depth_max) {
307  int idx = OPEN3D_ATOMIC_ADD(count_ptr, 1);
308 
309  float x_c = 0, y_c = 0, z_c = 0;
310  ti.Unproject(static_cast<float>(x), static_cast<float>(y),
311  d, &x_c, &y_c, &z_c);
312 
313  float* vertex = static_cast<float*>(
314  point_indexer.GetDataPtrFromCoord(idx));
315  ti.RigidTransform(x_c, y_c, z_c, vertex + 0, vertex + 1,
316  vertex + 2);
317  }
318  });
319 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
320  int total_pts_count = count.Item<int>();
321 #else
322  int total_pts_count = (*count_ptr).load();
323 #endif
324  dsts.emplace("points", points.Slice(0, 0, total_pts_count));
325 }
326 
327 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
328 void CUDATSDFIntegrateKernel
329 #else
330 void CPUTSDFIntegrateKernel
331 #endif
332  (const std::unordered_map<std::string, Tensor>& srcs,
333  std::unordered_map<std::string, Tensor>& dsts) {
334  // Decode input tensors
335  static std::vector<std::string> src_attrs = {
336  "depth", "indices", "block_keys", "intrinsics",
337  "extrinsics", "resolution", "voxel_size", "sdf_trunc",
338  "depth_scale", "depth_max",
339  };
340  for (auto& k : src_attrs) {
341  if (srcs.count(k) == 0) {
342  utility::LogError(
343  "[TSDFIntegrateKernel] expected Tensor {} in srcs, but "
344  "did not receive",
345  k);
346  }
347  }
348 
349  Tensor depth = srcs.at("depth").To(core::Dtype::Float32);
350  Tensor indices = srcs.at("indices");
351  Tensor block_keys = srcs.at("block_keys");
352  Tensor block_values = dsts.at("block_values");
353 
354  // Transforms
355  Tensor intrinsics = srcs.at("intrinsics").To(core::Dtype::Float32);
356  Tensor extrinsics = srcs.at("extrinsics").To(core::Dtype::Float32);
357 
358  // Parameters
359  int64_t resolution = srcs.at("resolution").Item<int64_t>();
360  int64_t resolution3 = resolution * resolution * resolution;
361 
362  float voxel_size = srcs.at("voxel_size").Item<float>();
363  float sdf_trunc = srcs.at("sdf_trunc").Item<float>();
364  float depth_scale = srcs.at("depth_scale").Item<float>();
365  float depth_max = srcs.at("depth_max").Item<float>();
366 
367  // Shape / transform indexers, no data involved
368  NDArrayIndexer voxel_indexer({resolution, resolution, resolution});
369  TransformIndexer transform_indexer(intrinsics, extrinsics, voxel_size);
370 
371  // Real data indexer
372  NDArrayIndexer depth_indexer(depth, 2);
373  NDArrayIndexer block_keys_indexer(block_keys, 1);
374  NDArrayIndexer voxel_block_buffer_indexer(block_values, 4);
375 
376  // Optional color integration
377  Tensor color;
378  NDArrayIndexer color_indexer;
379  bool integrate_color = false;
380  if (srcs.count("color") != 0) {
381  color = srcs.at("color").To(core::Dtype::Float32);
382  color_indexer = NDArrayIndexer(color, 2);
383  integrate_color = true;
384  }
385 
386  // Plain arrays that does not require indexers
387  int64_t* indices_ptr = static_cast<int64_t*>(indices.GetDataPtr());
388 
389  int64_t n = indices.GetLength() * resolution3;
390 
391 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
392  CUDALauncher launcher;
393 #else
394  CPULauncher launcher;
395 #endif
396 
397  DISPATCH_BYTESIZE_TO_VOXEL(
398  voxel_block_buffer_indexer.ElementByteSize(), [&]() {
399  launcher.LaunchGeneralKernel(n, [=] OPEN3D_DEVICE(
400  int64_t workload_idx) {
401  // Natural index (0, N) -> (block_idx, voxel_idx)
402  int64_t block_idx = indices_ptr[workload_idx / resolution3];
403  int64_t voxel_idx = workload_idx % resolution3;
404 
406  // block_idx -> (x_block, y_block, z_block)
407  int* block_key_ptr = static_cast<int*>(
408  block_keys_indexer.GetDataPtrFromCoord(block_idx));
409  int64_t xb = static_cast<int64_t>(block_key_ptr[0]);
410  int64_t yb = static_cast<int64_t>(block_key_ptr[1]);
411  int64_t zb = static_cast<int64_t>(block_key_ptr[2]);
412 
413  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
414  int64_t xv, yv, zv;
415  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
416 
417  // coordinate in world (in voxel)
418  int64_t x = (xb * resolution + xv);
419  int64_t y = (yb * resolution + yv);
420  int64_t z = (zb * resolution + zv);
421 
422  // coordinate in camera (in voxel -> in meter)
423  float xc, yc, zc, u, v;
424  transform_indexer.RigidTransform(
425  static_cast<float>(x), static_cast<float>(y),
426  static_cast<float>(z), &xc, &yc, &zc);
427 
428  // coordinate in image (in pixel)
429  transform_indexer.Project(xc, yc, zc, &u, &v);
430  if (!depth_indexer.InBoundary(u, v)) {
431  return;
432  }
433 
434  // Associate image workload and compute SDF and TSDF.
435  float depth = *static_cast<const float*>(
436  depth_indexer.GetDataPtrFromCoord(
437  static_cast<int64_t>(u),
438  static_cast<int64_t>(v))) /
439  depth_scale;
440 
441  float sdf = (depth - zc);
442  if (depth <= 0 || depth > depth_max || zc <= 0 ||
443  sdf < -sdf_trunc) {
444  return;
445  }
446  sdf = sdf < sdf_trunc ? sdf : sdf_trunc;
447  sdf /= sdf_trunc;
448 
449  // Associate voxel workload and update TSDF/Weights
450  voxel_t* voxel_ptr = static_cast<voxel_t*>(
451  voxel_block_buffer_indexer.GetDataPtrFromCoord(
452  xv, yv, zv, block_idx));
453 
454  if (integrate_color) {
455  float* color_ptr = static_cast<float*>(
456  color_indexer.GetDataPtrFromCoord(
457  static_cast<int64_t>(u),
458  static_cast<int64_t>(v)));
459 
460  voxel_ptr->Integrate(sdf, color_ptr[0], color_ptr[1],
461  color_ptr[2]);
462  } else {
463  voxel_ptr->Integrate(sdf);
464  }
465  });
466  });
467 }
468 
469 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
470 void CUDAPointExtractionKernel
471 #else
472 void CPUPointExtractionKernel
473 #endif
474  (const std::unordered_map<std::string, Tensor>& srcs,
475  std::unordered_map<std::string, Tensor>& dsts) {
476  // Decode input tensors
477  static std::vector<std::string> src_attrs = {
478  "indices", "nb_indices", "nb_masks", "block_keys",
479  "block_values", "voxel_size", "resolution",
480  };
481  for (auto& k : src_attrs) {
482  if (srcs.count(k) == 0) {
483  utility::LogError(
484  "[TSDFSurfaceExtractionKernel] expected Tensor {} in "
485  "srcs, but did not receive",
486  k);
487  }
488  }
489 
490  Tensor indices = srcs.at("indices");
491  Tensor nb_indices = srcs.at("nb_indices");
492  Tensor nb_masks = srcs.at("nb_masks");
493  Tensor block_keys = srcs.at("block_keys");
494  Tensor block_values = srcs.at("block_values");
495 
496  // Parameters
497  int64_t resolution = srcs.at("resolution").Item<int64_t>();
498  int64_t resolution3 = resolution * resolution * resolution;
499 
500  float voxel_size = srcs.at("voxel_size").Item<float>();
501 
502  // Shape / transform indexers, no data involved
503  NDArrayIndexer voxel_indexer({resolution, resolution, resolution});
504 
505  // Real data indexer
506  NDArrayIndexer voxel_block_buffer_indexer(block_values, 4);
507  NDArrayIndexer block_keys_indexer(block_keys, 1);
508  NDArrayIndexer nb_block_masks_indexer(nb_masks, 2);
509  NDArrayIndexer nb_block_indices_indexer(nb_indices, 2);
510 
511  // Plain arrays that does not require indexers
512  int64_t* indices_ptr = static_cast<int64_t*>(indices.GetDataPtr());
513 
514  int64_t n_blocks = indices.GetLength();
515  int64_t n = n_blocks * resolution3;
516 
517  // Output
518 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
519  core::Tensor count(std::vector<int>{0}, {}, core::Dtype::Int32,
520  block_values.GetDevice());
521  int* count_ptr = static_cast<int*>(count.GetDataPtr());
522 #else
523  std::atomic<int> count_atomic(0);
524  std::atomic<int>* count_ptr = &count_atomic;
525 #endif
526 
527 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
528  CUDALauncher launcher;
529 #else
530  CPULauncher launcher;
531 #endif
532 
533  // This pass determines valid number of points.
534  DISPATCH_BYTESIZE_TO_VOXEL(
535  voxel_block_buffer_indexer.ElementByteSize(), [&]() {
536  launcher.LaunchGeneralKernel(n, [=] OPEN3D_DEVICE(
537  int64_t workload_idx) {
538  auto GetVoxelAt = [&] OPEN3D_DEVICE(
539  int xo, int yo, int zo,
540  int curr_block_idx) -> voxel_t* {
541  return DeviceGetVoxelAt<voxel_t>(
542  xo, yo, zo, curr_block_idx,
543  static_cast<int>(resolution),
544  nb_block_masks_indexer,
545  nb_block_indices_indexer,
546  voxel_block_buffer_indexer);
547  };
548 
549  // Natural index (0, N) -> (block_idx, voxel_idx)
550  int64_t workload_block_idx = workload_idx / resolution3;
551  int64_t block_idx = indices_ptr[workload_block_idx];
552  int64_t voxel_idx = workload_idx % resolution3;
553 
554  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
555  int64_t xv, yv, zv;
556  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
557 
558  voxel_t* voxel_ptr = static_cast<voxel_t*>(
559  voxel_block_buffer_indexer.GetDataPtrFromCoord(
560  xv, yv, zv, block_idx));
561  float tsdf_o = voxel_ptr->GetTSDF();
562  float weight_o = voxel_ptr->GetWeight();
563  if (weight_o <= kWeightThreshold) return;
564 
565  // Enumerate x-y-z directions
566  for (int i = 0; i < 3; ++i) {
567  voxel_t* ptr = GetVoxelAt(
568  static_cast<int>(xv) + (i == 0),
569  static_cast<int>(yv) + (i == 1),
570  static_cast<int>(zv) + (i == 2),
571  static_cast<int>(workload_block_idx));
572  if (ptr == nullptr) continue;
573 
574  float tsdf_i = ptr->GetTSDF();
575  float weight_i = ptr->GetWeight();
576 
577  if (weight_i > kWeightThreshold &&
578  tsdf_i * tsdf_o < 0) {
579  OPEN3D_ATOMIC_ADD(count_ptr, 1);
580  }
581  }
582  });
583  });
584 
585 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
586  int total_count = count.Item<int>();
587 #else
588  int total_count = (*count_ptr).load();
589 #endif
590  utility::LogInfo("Total point count = {}", total_count);
591 
592  core::Tensor points({total_count, 3}, core::Dtype::Float32,
593  block_values.GetDevice());
594  core::Tensor normals({total_count, 3}, core::Dtype::Float32,
595  block_values.GetDevice());
596  NDArrayIndexer point_indexer(points, 1);
597  NDArrayIndexer normal_indexer(normals, 1);
598 
599  // Reset count
600 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
601  count = core::Tensor(std::vector<int>{0}, {}, core::Dtype::Int32,
602  block_values.GetDevice());
603  count_ptr = static_cast<int*>(count.GetDataPtr());
604 #else
605  (*count_ptr) = 0;
606 #endif
607 
608  // This pass extracts exact surface points.
609  DISPATCH_BYTESIZE_TO_VOXEL(
610  voxel_block_buffer_indexer.ElementByteSize(), [&]() {
611  bool extract_color = false;
612  Tensor colors;
613  NDArrayIndexer color_indexer;
614  if (voxel_t::HasColor()) {
615  extract_color = true;
616  colors = Tensor({total_count, 3}, core::Dtype::Float32,
617  block_values.GetDevice());
618  color_indexer = NDArrayIndexer(colors, 1);
619  }
620 
621  launcher.LaunchGeneralKernel(n, [=] OPEN3D_DEVICE(
622  int64_t workload_idx) {
623  auto GetVoxelAt = [&] OPEN3D_DEVICE(
624  int xo, int yo, int zo,
625  int curr_block_idx) -> voxel_t* {
626  return DeviceGetVoxelAt<voxel_t>(
627  xo, yo, zo, curr_block_idx,
628  static_cast<int>(resolution),
629  nb_block_masks_indexer,
630  nb_block_indices_indexer,
631  voxel_block_buffer_indexer);
632  };
633  auto GetNormalAt = [&] OPEN3D_DEVICE(int xo, int yo, int zo,
634  int curr_block_idx,
635  float* n) {
636  return DeviceGetNormalAt<voxel_t>(
637  xo, yo, zo, curr_block_idx, n,
638  static_cast<int>(resolution), voxel_size,
639  nb_block_masks_indexer,
640  nb_block_indices_indexer,
641  voxel_block_buffer_indexer);
642  };
643 
644  // Natural index (0, N) -> (block_idx, voxel_idx)
645  int64_t workload_block_idx = workload_idx / resolution3;
646  int64_t block_idx = indices_ptr[workload_block_idx];
647  int64_t voxel_idx = workload_idx % resolution3;
648 
650  // block_idx -> (x_block, y_block, z_block)
651  int* block_key_ptr = static_cast<int*>(
652  block_keys_indexer.GetDataPtrFromCoord(block_idx));
653  int64_t xb = static_cast<int64_t>(block_key_ptr[0]);
654  int64_t yb = static_cast<int64_t>(block_key_ptr[1]);
655  int64_t zb = static_cast<int64_t>(block_key_ptr[2]);
656 
657  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
658  int64_t xv, yv, zv;
659  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
660 
661  voxel_t* voxel_ptr = static_cast<voxel_t*>(
662  voxel_block_buffer_indexer.GetDataPtrFromCoord(
663  xv, yv, zv, block_idx));
664  float tsdf_o = voxel_ptr->GetTSDF();
665  float weight_o = voxel_ptr->GetWeight();
666 
667  if (weight_o <= kWeightThreshold) return;
668 
669  int64_t x = xb * resolution + xv;
670  int64_t y = yb * resolution + yv;
671  int64_t z = zb * resolution + zv;
672 
673  float no[3] = {0}, ni[3] = {0};
674  GetNormalAt(static_cast<int>(xv), static_cast<int>(yv),
675  static_cast<int>(zv),
676  static_cast<int>(workload_block_idx), no);
677 
678  // Enumerate x-y-z axis
679  for (int i = 0; i < 3; ++i) {
680  voxel_t* ptr = GetVoxelAt(
681  static_cast<int>(xv) + (i == 0),
682  static_cast<int>(yv) + (i == 1),
683  static_cast<int>(zv) + (i == 2),
684  static_cast<int>(workload_block_idx));
685  if (ptr == nullptr) continue;
686 
687  float tsdf_i = ptr->GetTSDF();
688  float weight_i = ptr->GetWeight();
689 
690  if (weight_i > kWeightThreshold &&
691  tsdf_i * tsdf_o < 0) {
692  float ratio = (0 - tsdf_o) / (tsdf_i - tsdf_o);
693 
694  int idx = OPEN3D_ATOMIC_ADD(count_ptr, 1);
695 
696  float* point_ptr = static_cast<float*>(
697  point_indexer.GetDataPtrFromCoord(idx));
698  point_ptr[0] =
699  voxel_size * (x + ratio * int(i == 0));
700  point_ptr[1] =
701  voxel_size * (y + ratio * int(i == 1));
702  point_ptr[2] =
703  voxel_size * (z + ratio * int(i == 2));
704  GetNormalAt(static_cast<int>(xv) + (i == 0),
705  static_cast<int>(yv) + (i == 1),
706  static_cast<int>(zv) + (i == 2),
707  static_cast<int>(workload_block_idx),
708  ni);
709 
710  float* normal_ptr = static_cast<float*>(
711  normal_indexer.GetDataPtrFromCoord(idx));
712  float nx = (1 - ratio) * no[0] + ratio * ni[0];
713  float ny = (1 - ratio) * no[1] + ratio * ni[1];
714  float nz = (1 - ratio) * no[2] + ratio * ni[2];
715  float norm = static_cast<float>(
716  sqrt(nx * nx + ny * ny + nz * nz) + 1e-5);
717  normal_ptr[0] = nx / norm;
718  normal_ptr[1] = ny / norm;
719  normal_ptr[2] = nz / norm;
720 
721  if (extract_color) {
722  float* color_ptr = static_cast<float*>(
723  color_indexer.GetDataPtrFromCoord(idx));
724 
725  float r_o = voxel_ptr->GetR();
726  float g_o = voxel_ptr->GetG();
727  float b_o = voxel_ptr->GetB();
728 
729  float r_i = ptr->GetR();
730  float g_i = ptr->GetG();
731  float b_i = ptr->GetB();
732 
733  color_ptr[0] =
734  ((1 - ratio) * r_o + ratio * r_i) /
735  255.0f;
736  color_ptr[1] =
737  ((1 - ratio) * g_o + ratio * g_i) /
738  255.0f;
739  color_ptr[2] =
740  ((1 - ratio) * b_o + ratio * b_i) /
741  255.0f;
742  }
743  }
744  }
745  });
746  dsts.emplace("points", points);
747  dsts.emplace("normals", normals);
748 
749  if (extract_color) {
750  dsts.emplace("colors", colors);
751  }
752  });
753 }
754 
755 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
756 void CUDAMeshExtractionKernel
757 #else
758 void CPUMeshExtractionKernel
759 #endif
760  (const std::unordered_map<std::string, Tensor>& srcs,
761  std::unordered_map<std::string, Tensor>& dsts) {
762  // Decode input tensors
763  static std::vector<std::string> src_attrs = {
764  "indices", "inv_indices", "nb_indices", "nb_masks",
765  "block_keys", "block_values", "voxel_size", "resolution",
766  };
767  for (auto& k : src_attrs) {
768  if (srcs.count(k) == 0) {
769  utility::LogError(
770  "[CUDAMarchingCubesKernel] expected Tensor {} in "
771  "srcs, but "
772  "did not receive",
773  k);
774  }
775  }
776 
777  Tensor indices = srcs.at("indices");
778  Tensor inv_indices = srcs.at("inv_indices");
779  Tensor nb_indices = srcs.at("nb_indices");
780  Tensor nb_masks = srcs.at("nb_masks");
781  Tensor block_keys = srcs.at("block_keys");
782  Tensor block_values = srcs.at("block_values");
783 
784  // Parameters
785  int64_t resolution = srcs.at("resolution").Item<int64_t>();
786  int64_t resolution3 = resolution * resolution * resolution;
787 
788  float voxel_size = srcs.at("voxel_size").Item<float>();
789 
790  // Shape / transform indexers, no data involved
791  NDArrayIndexer voxel_indexer({resolution, resolution, resolution});
792 
793  // Output
794 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
795  CUDACachedMemoryManager::ReleaseCache();
796 #endif
797 
798  int n_blocks = static_cast<int>(indices.GetLength());
799  // Voxel-wise mesh info. 4 channels correspond to:
800  // 3 edges' corresponding vertex index + 1 table index.
801  core::Tensor mesh_structure;
802  try {
803  mesh_structure = core::Tensor::Zeros(
804  {n_blocks, resolution, resolution, resolution, 4},
805  core::Dtype::Int32, block_keys.GetDevice());
806  } catch (const std::runtime_error&) {
807  utility::LogError(
808  "[MeshExtractionKernel] Unable to allocate assistance mesh "
809  "structure for Marching "
810  "Cubes with {} active voxel blocks. Please consider using a "
811  "larger voxel size (currently {}) for TSDF "
812  "integration, or using tsdf_volume.cpu() to perform mesh "
813  "extraction on CPU.",
814  n_blocks, voxel_size);
815  }
816 
817  // Real data indexer
818  NDArrayIndexer voxel_block_buffer_indexer(block_values, 4);
819  NDArrayIndexer mesh_structure_indexer(mesh_structure, 4);
820  NDArrayIndexer nb_block_masks_indexer(nb_masks, 2);
821  NDArrayIndexer nb_block_indices_indexer(nb_indices, 2);
822 
823  // Plain arrays that does not require indexers
824  int64_t* indices_ptr = static_cast<int64_t*>(indices.GetDataPtr());
825  int64_t* inv_indices_ptr = static_cast<int64_t*>(inv_indices.GetDataPtr());
826 
827  int64_t n = n_blocks * resolution3;
828 
829 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
830  CUDALauncher launcher;
831 #else
832  CPULauncher launcher;
833 #endif
834 
835  // Pass 0: analyze mesh structure, set up one-on-one correspondences from
836  // edges to vertices.
837  DISPATCH_BYTESIZE_TO_VOXEL(
838  voxel_block_buffer_indexer.ElementByteSize(), [&]() {
839  launcher.LaunchGeneralKernel(n, [=] OPEN3D_DEVICE(
840  int64_t workload_idx) {
841  auto GetVoxelAt = [&] OPEN3D_DEVICE(
842  int xo, int yo, int zo,
843  int curr_block_idx) -> voxel_t* {
844  return DeviceGetVoxelAt<voxel_t>(
845  xo, yo, zo, curr_block_idx,
846  static_cast<int>(resolution),
847  nb_block_masks_indexer,
848  nb_block_indices_indexer,
849  voxel_block_buffer_indexer);
850  };
851 
852  // Natural index (0, N) -> (block_idx, voxel_idx)
853  int64_t workload_block_idx = workload_idx / resolution3;
854  int64_t voxel_idx = workload_idx % resolution3;
855 
856  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
857  int64_t xv, yv, zv;
858  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
859 
860  // Check per-vertex sign in the cube to determine cube type
861  int table_idx = 0;
862  for (int i = 0; i < 8; ++i) {
863  voxel_t* voxel_ptr_i = GetVoxelAt(
864  static_cast<int>(xv) + vtx_shifts[i][0],
865  static_cast<int>(yv) + vtx_shifts[i][1],
866  static_cast<int>(zv) + vtx_shifts[i][2],
867  static_cast<int>(workload_block_idx));
868  if (voxel_ptr_i == nullptr) return;
869 
870  float tsdf_i = voxel_ptr_i->GetTSDF();
871  float weight_i = voxel_ptr_i->GetWeight();
872  if (weight_i <= kWeightThreshold) return;
873 
874  table_idx |= ((tsdf_i < 0) ? (1 << i) : 0);
875  }
876 
877  int* mesh_struct_ptr = static_cast<int*>(
878  mesh_structure_indexer.GetDataPtrFromCoord(
879  xv, yv, zv, workload_block_idx));
880  mesh_struct_ptr[3] = table_idx;
881 
882  if (table_idx == 0 || table_idx == 255) return;
883 
884  // Check per-edge sign in the cube to determine cube type
885  int edges_with_vertices = edge_table[table_idx];
886  for (int i = 0; i < 12; ++i) {
887  if (edges_with_vertices & (1 << i)) {
888  int64_t xv_i = xv + edge_shifts[i][0];
889  int64_t yv_i = yv + edge_shifts[i][1];
890  int64_t zv_i = zv + edge_shifts[i][2];
891  int edge_i = edge_shifts[i][3];
892 
893  int dxb = static_cast<int>(xv_i / resolution);
894  int dyb = static_cast<int>(yv_i / resolution);
895  int dzb = static_cast<int>(zv_i / resolution);
896 
897  int nb_idx =
898  (dxb + 1) + (dyb + 1) * 3 + (dzb + 1) * 9;
899 
900  int64_t block_idx_i = *static_cast<int64_t*>(
901  nb_block_indices_indexer
902  .GetDataPtrFromCoord(
903  workload_block_idx,
904  nb_idx));
905  int* mesh_ptr_i = static_cast<int*>(
906  mesh_structure_indexer.GetDataPtrFromCoord(
907  xv_i - dxb * resolution,
908  yv_i - dyb * resolution,
909  zv_i - dzb * resolution,
910  inv_indices_ptr[block_idx_i]));
911 
912  // Non-atomic write, but we are safe
913  mesh_ptr_i[edge_i] = -1;
914  }
915  }
916  });
917  });
918 
919  // Pass 1: determine valid number of vertices.
920 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
921  core::Tensor vtx_count(std::vector<int>{0}, {}, core::Dtype::Int32,
922  block_values.GetDevice());
923  int* vtx_count_ptr = static_cast<int*>(vtx_count.GetDataPtr());
924 #else
925  std::atomic<int> vtx_count_atomic(0);
926  std::atomic<int>* vtx_count_ptr = &vtx_count_atomic;
927 #endif
928 
929 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
930  CUDALauncher::LaunchGeneralKernel(
931  n, [=] OPEN3D_DEVICE(int64_t workload_idx) {
932 #else
933  CPULauncher::LaunchGeneralKernel(n, [&](int64_t workload_idx) {
934 #endif
935  // Natural index (0, N) -> (block_idx, voxel_idx)
936  int64_t workload_block_idx = workload_idx / resolution3;
937  int64_t voxel_idx = workload_idx % resolution3;
938 
939  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
940  int64_t xv, yv, zv;
941  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
942 
943  // Obtain voxel's mesh struct ptr
944  int* mesh_struct_ptr = static_cast<int*>(
945  mesh_structure_indexer.GetDataPtrFromCoord(
946  xv, yv, zv, workload_block_idx));
947 
948  // Early quit -- no allocated vertex to compute
949  if (mesh_struct_ptr[0] != -1 && mesh_struct_ptr[1] != -1 &&
950  mesh_struct_ptr[2] != -1) {
951  return;
952  }
953 
954  // Enumerate 3 edges in the voxel
955  for (int e = 0; e < 3; ++e) {
956  int vertex_idx = mesh_struct_ptr[e];
957  if (vertex_idx != -1) continue;
958 
959  OPEN3D_ATOMIC_ADD(vtx_count_ptr, 1);
960  }
961  });
962 
963  // Reset count_ptr
964 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
965  int total_vtx_count = vtx_count.Item<int>();
966  vtx_count = core::Tensor(std::vector<int>{0}, {}, core::Dtype::Int32,
967  block_values.GetDevice());
968  vtx_count_ptr = static_cast<int*>(vtx_count.GetDataPtr());
969 #else
970  int total_vtx_count = (*vtx_count_ptr).load();
971  (*vtx_count_ptr) = 0;
972 #endif
973 
974  utility::LogInfo("Total vertex count = {}", total_vtx_count);
975  core::Tensor vertices({total_vtx_count, 3}, core::Dtype::Float32,
976  block_values.GetDevice());
977  core::Tensor normals({total_vtx_count, 3}, core::Dtype::Float32,
978  block_values.GetDevice());
979 
980  NDArrayIndexer block_keys_indexer(block_keys, 1);
981  NDArrayIndexer vertex_indexer(vertices, 1);
982  NDArrayIndexer normal_indexer(normals, 1);
983 
984  // Pass 2: extract vertices.
985  DISPATCH_BYTESIZE_TO_VOXEL(
986  voxel_block_buffer_indexer.ElementByteSize(), [&]() {
987  bool extract_color = false;
988  Tensor colors;
989  NDArrayIndexer color_indexer;
990  if (voxel_t::HasColor()) {
991  extract_color = true;
992  colors = Tensor({total_vtx_count, 3}, core::Dtype::Float32,
993  block_values.GetDevice());
994  color_indexer = NDArrayIndexer(colors, 1);
995  }
996  launcher.LaunchGeneralKernel(n, [=] OPEN3D_DEVICE(
997  int64_t workload_idx) {
998  auto GetVoxelAt = [&] OPEN3D_DEVICE(
999  int xo, int yo, int zo,
1000  int curr_block_idx) -> voxel_t* {
1001  return DeviceGetVoxelAt<voxel_t>(
1002  xo, yo, zo, curr_block_idx,
1003  static_cast<int>(resolution),
1004  nb_block_masks_indexer,
1005  nb_block_indices_indexer,
1006  voxel_block_buffer_indexer);
1007  };
1008 
1009  auto GetNormalAt = [&] OPEN3D_DEVICE(int xo, int yo, int zo,
1010  int curr_block_idx,
1011  float* n) {
1012  return DeviceGetNormalAt<voxel_t>(
1013  xo, yo, zo, curr_block_idx, n,
1014  static_cast<int>(resolution), voxel_size,
1015  nb_block_masks_indexer,
1016  nb_block_indices_indexer,
1017  voxel_block_buffer_indexer);
1018  };
1019 
1020  // Natural index (0, N) -> (block_idx, voxel_idx)
1021  int64_t workload_block_idx = workload_idx / resolution3;
1022  int64_t block_idx = indices_ptr[workload_block_idx];
1023  int64_t voxel_idx = workload_idx % resolution3;
1024 
1025  // block_idx -> (x_block, y_block, z_block)
1026  int* block_key_ptr = static_cast<int*>(
1027  block_keys_indexer.GetDataPtrFromCoord(block_idx));
1028  int64_t xb = static_cast<int64_t>(block_key_ptr[0]);
1029  int64_t yb = static_cast<int64_t>(block_key_ptr[1]);
1030  int64_t zb = static_cast<int64_t>(block_key_ptr[2]);
1031 
1032  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
1033  int64_t xv, yv, zv;
1034  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
1035 
1036  // global coordinate (in voxels)
1037  int64_t x = xb * resolution + xv;
1038  int64_t y = yb * resolution + yv;
1039  int64_t z = zb * resolution + zv;
1040 
1041  // Obtain voxel's mesh struct ptr
1042  int* mesh_struct_ptr = static_cast<int*>(
1043  mesh_structure_indexer.GetDataPtrFromCoord(
1044  xv, yv, zv, workload_block_idx));
1045 
1046  // Early quit -- no allocated vertex to compute
1047  if (mesh_struct_ptr[0] != -1 && mesh_struct_ptr[1] != -1 &&
1048  mesh_struct_ptr[2] != -1) {
1049  return;
1050  }
1051 
1052  // Obtain voxel ptr
1053  voxel_t* voxel_ptr = static_cast<voxel_t*>(
1054  voxel_block_buffer_indexer.GetDataPtrFromCoord(
1055  xv, yv, zv, block_idx));
1056  float tsdf_o = voxel_ptr->GetTSDF();
1057  float no[3] = {0}, ne[3] = {0};
1058  GetNormalAt(static_cast<int>(xv), static_cast<int>(yv),
1059  static_cast<int>(zv),
1060  static_cast<int>(workload_block_idx), no);
1061 
1062  // Enumerate 3 edges in the voxel
1063  for (int e = 0; e < 3; ++e) {
1064  int vertex_idx = mesh_struct_ptr[e];
1065  if (vertex_idx != -1) continue;
1066 
1067  voxel_t* voxel_ptr_e = GetVoxelAt(
1068  static_cast<int>(xv) + (e == 0),
1069  static_cast<int>(yv) + (e == 1),
1070  static_cast<int>(zv) + (e == 2),
1071  static_cast<int>(workload_block_idx));
1072  float tsdf_e = voxel_ptr_e->GetTSDF();
1073  float ratio = (0 - tsdf_o) / (tsdf_e - tsdf_o);
1074 
1075  int idx = OPEN3D_ATOMIC_ADD(vtx_count_ptr, 1);
1076  mesh_struct_ptr[e] = idx;
1077 
1078  float ratio_x = ratio * int(e == 0);
1079  float ratio_y = ratio * int(e == 1);
1080  float ratio_z = ratio * int(e == 2);
1081 
1082  float* vertex_ptr = static_cast<float*>(
1083  vertex_indexer.GetDataPtrFromCoord(idx));
1084  vertex_ptr[0] = voxel_size * (x + ratio_x);
1085  vertex_ptr[1] = voxel_size * (y + ratio_y);
1086  vertex_ptr[2] = voxel_size * (z + ratio_z);
1087 
1088  float* normal_ptr = static_cast<float*>(
1089  normal_indexer.GetDataPtrFromCoord(idx));
1090  GetNormalAt(static_cast<int>(xv) + (e == 0),
1091  static_cast<int>(yv) + (e == 1),
1092  static_cast<int>(zv) + (e == 2),
1093  static_cast<int>(workload_block_idx), ne);
1094  float nx = (1 - ratio) * no[0] + ratio * ne[0];
1095  float ny = (1 - ratio) * no[1] + ratio * ne[1];
1096  float nz = (1 - ratio) * no[2] + ratio * ne[2];
1097  float norm = static_cast<float>(
1098  sqrt(nx * nx + ny * ny + nz * nz) + 1e-5);
1099  normal_ptr[0] = nx / norm;
1100  normal_ptr[1] = ny / norm;
1101  normal_ptr[2] = nz / norm;
1102 
1103  if (extract_color) {
1104  float* color_ptr = static_cast<float*>(
1105  color_indexer.GetDataPtrFromCoord(idx));
1106  float r_o = voxel_ptr->GetR();
1107  float g_o = voxel_ptr->GetG();
1108  float b_o = voxel_ptr->GetB();
1109 
1110  float r_e = voxel_ptr_e->GetR();
1111  float g_e = voxel_ptr_e->GetG();
1112  float b_e = voxel_ptr_e->GetB();
1113  color_ptr[0] =
1114  ((1 - ratio) * r_o + ratio * r_e) / 255.0f;
1115  color_ptr[1] =
1116  ((1 - ratio) * g_o + ratio * g_e) / 255.0f;
1117  color_ptr[2] =
1118  ((1 - ratio) * b_o + ratio * b_e) / 255.0f;
1119  }
1120  }
1121  });
1122  dsts.emplace("vertices", vertices);
1123  dsts.emplace("normals", normals);
1124 
1125  if (extract_color) {
1126  dsts.emplace("colors", colors);
1127  }
1128  });
1129 
1130  // Pass 3: connect vertices and form triangles.
1131 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
1132  core::Tensor triangle_count(std::vector<int>{0}, {}, core::Dtype::Int32,
1133  block_values.GetDevice());
1134  int* tri_count_ptr = static_cast<int*>(triangle_count.GetDataPtr());
1135 #else
1136  std::atomic<int> tri_count_atomic(0);
1137  std::atomic<int>* tri_count_ptr = &tri_count_atomic;
1138 #endif
1139 
1140  core::Tensor triangles({total_vtx_count * 3, 3}, core::Dtype::Int64,
1141  block_values.GetDevice());
1142  NDArrayIndexer triangle_indexer(triangles, 1);
1143 
1144 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
1145  CUDALauncher::LaunchGeneralKernel(
1146  n, [=] OPEN3D_DEVICE(int64_t workload_idx) {
1147 #else
1148  CPULauncher::LaunchGeneralKernel(n, [&](int64_t workload_idx) {
1149 #endif
1150  // Natural index (0, N) -> (block_idx, voxel_idx)
1151  int64_t workload_block_idx = workload_idx / resolution3;
1152  int64_t voxel_idx = workload_idx % resolution3;
1153 
1154  // voxel_idx -> (x_voxel, y_voxel, z_voxel)
1155  int64_t xv, yv, zv;
1156  voxel_indexer.WorkloadToCoord(voxel_idx, &xv, &yv, &zv);
1157 
1158  // Obtain voxel's mesh struct ptr
1159  int* mesh_struct_ptr = static_cast<int*>(
1160  mesh_structure_indexer.GetDataPtrFromCoord(
1161  xv, yv, zv, workload_block_idx));
1162 
1163  int table_idx = mesh_struct_ptr[3];
1164  if (tri_count[table_idx] == 0) return;
1165 
1166  for (size_t tri = 0; tri < 16; tri += 3) {
1167  if (tri_table[table_idx][tri] == -1) return;
1168 
1169  int tri_idx = OPEN3D_ATOMIC_ADD(tri_count_ptr, 1);
1170 
1171  for (size_t vertex = 0; vertex < 3; ++vertex) {
1172  int edge = tri_table[table_idx][tri + vertex];
1173 
1174  int64_t xv_i = xv + edge_shifts[edge][0];
1175  int64_t yv_i = yv + edge_shifts[edge][1];
1176  int64_t zv_i = zv + edge_shifts[edge][2];
1177  int64_t edge_i = edge_shifts[edge][3];
1178 
1179  int dxb = static_cast<int>(xv_i / resolution);
1180  int dyb = static_cast<int>(yv_i / resolution);
1181  int dzb = static_cast<int>(zv_i / resolution);
1182 
1183  int nb_idx = (dxb + 1) + (dyb + 1) * 3 + (dzb + 1) * 9;
1184 
1185  int64_t block_idx_i = *static_cast<int64_t*>(
1186  nb_block_indices_indexer.GetDataPtrFromCoord(
1187  workload_block_idx, nb_idx));
1188  int* mesh_struct_ptr_i = static_cast<int*>(
1189  mesh_structure_indexer.GetDataPtrFromCoord(
1190  xv_i - dxb * resolution,
1191  yv_i - dyb * resolution,
1192  zv_i - dzb * resolution,
1193  inv_indices_ptr[block_idx_i]));
1194 
1195  int64_t* triangle_ptr = static_cast<int64_t*>(
1196  triangle_indexer.GetDataPtrFromCoord(tri_idx));
1197  triangle_ptr[2 - vertex] = mesh_struct_ptr_i[edge_i];
1198  }
1199  }
1200  });
1201 
1202 #if defined(BUILD_CUDA_MODULE) && defined(__CUDACC__)
1203  int total_tri_count = triangle_count.Item<int>();
1204 #else
1205  int total_tri_count = (*tri_count_ptr).load();
1206 #endif
1207  utility::LogInfo("Total triangle count = {}", total_tri_count);
1208  triangles = triangles.Slice(0, 0, total_tri_count);
1209  dsts.emplace("triangles", triangles);
1210 }
1211 
1212 } // namespace kernel
1213 } // namespace core
1214 } // namespace open3d
open3d::core::kernel::ColoredVoxel32f
Definition: GeneralEWSharedImpl.h:149
open3d::core::kernel::ColoredVoxel16i::GetTSDF
OPEN3D_HOST_DEVICE float GetTSDF()
Definition: GeneralEWSharedImpl.h:109
open3d::core::kernel::ColoredVoxel32f::GetR
OPEN3D_HOST_DEVICE float GetR()
Definition: GeneralEWSharedImpl.h:160
open3d::core::kernel::ColoredVoxel32f::g
float g
Definition: GeneralEWSharedImpl.h:154
open3d::core::kernel::CPUUnprojectKernel
void CPUUnprojectKernel(const std::unordered_map< std::string, Tensor > &srcs, std::unordered_map< std::string, Tensor > &dsts)
Definition: GeneralEWSharedImpl.h:250
open3d::core::kernel::ColoredVoxel32f::GetG
OPEN3D_HOST_DEVICE float GetG()
Definition: GeneralEWSharedImpl.h:161
open3d::core::kernel::CPUTSDFIntegrateKernel
void CPUTSDFIntegrateKernel(const std::unordered_map< std::string, Tensor > &srcs, std::unordered_map< std::string, Tensor > &dsts)
Definition: GeneralEWSharedImpl.h:332
open3d::core::kernel::Voxel32f
Definition: GeneralEWSharedImpl.h:68
open3d::core::kernel::ColoredVoxel32f::r
float r
Definition: GeneralEWSharedImpl.h:153
open3d::core::cuda::ReleaseCache
void ReleaseCache()
Definition: CUDAUtils.cpp:55
open3d::core::Tensor::Inverse
Tensor Inverse() const
Definition: Tensor.cpp:1293
open3d::core::kernel::CPULauncher
Definition: CPULauncher.h:42
open3d::core::Tensor::GetDataPtr
void * GetDataPtr()
Definition: Tensor.h:961
open3d::core::kernel::CPUMeshExtractionKernel
void CPUMeshExtractionKernel(const std::unordered_map< std::string, Tensor > &srcs, std::unordered_map< std::string, Tensor > &dsts)
Definition: GeneralEWSharedImpl.h:760
open3d::core::kernel::Voxel32f::GetTSDF
OPEN3D_HOST_DEVICE float GetTSDF()
Definition: GeneralEWSharedImpl.h:73
open3d::core::kernel::ColoredVoxel32f::Integrate
OPEN3D_HOST_DEVICE void Integrate(float dsdf, float dr, float dg, float db)
Definition: GeneralEWSharedImpl.h:168
open3d::utility::LogError
void LogError(const char *format, const Args &... args)
Definition: Console.h:176
open3d::core::kernel::DeviceGetVoxelAt
OPEN3D_DEVICE voxel_t * DeviceGetVoxelAt(int xo, int yo, int zo, int curr_block_idx, int resolution, const NDArrayIndexer &nb_block_masks_indexer, const NDArrayIndexer &nb_block_indices_indexer, const NDArrayIndexer &blocks_indexer)
Definition: GeneralEWSharedImpl.h:184
open3d::core::kernel::ColoredVoxel16i::GetR
OPEN3D_HOST_DEVICE float GetR()
Definition: GeneralEWSharedImpl.h:111
open3d::core::kernel::ColoredVoxel16i
Definition: GeneralEWSharedImpl.h:97
open3d::core::kernel::NDArrayIndexer
Definition: GeneralIndexer.h:133
SizeVector.h
open3d::core::kernel::ColoredVoxel16i::GetG
OPEN3D_HOST_DEVICE float GetG()
Definition: GeneralEWSharedImpl.h:114
open3d::core::kernel::ColoredVoxel32f::GetTSDF
OPEN3D_HOST_DEVICE float GetTSDF()
Definition: GeneralEWSharedImpl.h:158
open3d::core::kernel::ColoredVoxel16i::b
uint16_t b
Definition: GeneralEWSharedImpl.h:106
open3d::core::kernel::ColoredVoxel16i::r
uint16_t r
Definition: GeneralEWSharedImpl.h:104
open3d::core::kernel::Voxel32f::GetWeight
OPEN3D_HOST_DEVICE float GetWeight()
Definition: GeneralEWSharedImpl.h:74
GeneralEWMacros.h
open3d::core::Tensor::GetDevice
Device GetDevice() const
Definition: Tensor.cpp:955
GeneralEW.h
OPEN3D_DEVICE
#define OPEN3D_DEVICE
Definition: CUDAUtils.h:55
open3d::core::kernel::Voxel32f::GetB
OPEN3D_HOST_DEVICE float GetB()
Definition: GeneralEWSharedImpl.h:77
open3d::core::kernel::NDArrayIndexer::GetDataPtrFromCoord
OPEN3D_HOST_DEVICE void * GetDataPtrFromCoord(int64_t x) const
Definition: GeneralIndexer.h:278
GeneralIndexer.h
open3d::core::Dtype::Int32
static const Dtype Int32
Definition: Dtype.h:44
open3d::core::kernel::ColoredVoxel16i::g
uint16_t g
Definition: GeneralEWSharedImpl.h:105
open3d::core::kernel::Voxel32f::GetR
OPEN3D_HOST_DEVICE float GetR()
Definition: GeneralEWSharedImpl.h:75
open3d::core::kernel::Voxel32f::Integrate
OPEN3D_HOST_DEVICE void Integrate(float dsdf, float dr, float dg, float db)
Definition: GeneralEWSharedImpl.h:83
open3d::core::kernel::Voxel32f::GetG
OPEN3D_HOST_DEVICE float GetG()
Definition: GeneralEWSharedImpl.h:76
color
math::float4 color
Definition: LineSetBuffers.cpp:64
Console.h
open3d::core::kernel::ColoredVoxel32f::GetB
OPEN3D_HOST_DEVICE float GetB()
Definition: GeneralEWSharedImpl.h:162
OPEN3D_HOST_DEVICE
#define OPEN3D_HOST_DEVICE
Definition: CUDAUtils.h:54
DISPATCH_BYTESIZE_TO_VOXEL
#define DISPATCH_BYTESIZE_TO_VOXEL(BYTESIZE,...)
Definition: GeneralEWSharedImpl.h:45
open3d::core::kernel::ColoredVoxel32f::tsdf
float tsdf
Definition: GeneralEWSharedImpl.h:150
Dispatch.h
open3d::core::kernel::ColoredVoxel16i::Integrate
OPEN3D_HOST_DEVICE void Integrate(float dsdf, float dr, float dg, float db)
Definition: GeneralEWSharedImpl.h:128
open3d::core::kernel::ColoredVoxel32f::HasColor
static bool HasColor()
Definition: GeneralEWSharedImpl.h:157
open3d::core::Tensor::To
Tensor To(Dtype dtype, bool copy=false) const
Definition: Tensor.cpp:453
open3d::core::kernel::ColoredVoxel16i::GetWeight
OPEN3D_HOST_DEVICE float GetWeight()
Definition: GeneralEWSharedImpl.h:110
open3d::core::kernel::ColoredVoxel16i::Integrate
OPEN3D_HOST_DEVICE void Integrate(float dsdf)
Definition: GeneralEWSharedImpl.h:120
open3d::core::kernel::ColoredVoxel16i::HasColor
static bool HasColor()
Definition: GeneralEWSharedImpl.h:108
open3d::core::kernel::NDArrayIndexer::ElementByteSize
OPEN3D_HOST_DEVICE int64_t ElementByteSize()
Definition: GeneralIndexer.h:189
open3d::core::kernel::ColoredVoxel16i::GetB
OPEN3D_HOST_DEVICE float GetB()
Definition: GeneralEWSharedImpl.h:117
open3d::io::k4a_plugin::int
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 timeout_in_ms capture_handle capture_handle capture_handle image_handle temperature_c int
Definition: K4aPlugin.cpp:479
MemoryManager.h
open3d::core::kernel::ColoredVoxel32f::weight
float weight
Definition: GeneralEWSharedImpl.h:151
open3d::core::Tensor::Zeros
static Tensor Zeros(const SizeVector &shape, Dtype dtype, const Device &device=Device("CPU:0"))
Create a tensor fill with zeros.
Definition: Tensor.cpp:182
stride
size_t stride
Definition: TriangleMeshBuffers.cpp:183
count
int count
Definition: FilePCD.cpp:61
open3d::core::Dtype::Float32
static const Dtype Float32
Definition: Dtype.h:42
open3d::core::kernel::CPUPointExtractionKernel
void CPUPointExtractionKernel(const std::unordered_map< std::string, Tensor > &srcs, std::unordered_map< std::string, Tensor > &dsts)
Definition: GeneralEWSharedImpl.h:474
open3d::core::kernel::CPULauncher::LaunchGeneralKernel
static void LaunchGeneralKernel(int64_t n, func_t element_kernel)
General kernels with non-conventional indexers.
Definition: CPULauncher.h:176
open3d::core::kernel::Voxel32f::weight
float weight
Definition: GeneralEWSharedImpl.h:70
open3d::core::kernel::Voxel32f::tsdf
float tsdf
Definition: GeneralEWSharedImpl.h:69
points
int points
Definition: FilePCD.cpp:73
open3d::core::Dtype::Int64
static const Dtype Int64
Definition: Dtype.h:45
open3d
Definition: PinholeCameraIntrinsic.cpp:35
open3d::core::Tensor
Definition: Tensor.h:48
open3d::core::kernel::TransformIndexer
Helper class for converting coordinates/indices between 3D/3D, 3D/2D, 2D/3D.
Definition: GeneralIndexer.h:40
open3d::core::kernel::ColoredVoxel32f::GetWeight
OPEN3D_HOST_DEVICE float GetWeight()
Definition: GeneralEWSharedImpl.h:159
open3d::core::kernel::NDArrayIndexer::InBoundary
OPEN3D_HOST_DEVICE bool InBoundary(float x, float y) const
Definition: GeneralIndexer.h:257
open3d::core::kernel::ColoredVoxel16i::weight
uint16_t weight
Definition: GeneralEWSharedImpl.h:102
open3d::core::kernel::ColoredVoxel32f::b
float b
Definition: GeneralEWSharedImpl.h:155
OPEN3D_ATOMIC_ADD
#define OPEN3D_ATOMIC_ADD(X, Y)
Definition: GeneralEWSharedImpl.h:42
Tensor.h
open3d::core::Tensor::GetLength
int64_t GetLength() const
Definition: Tensor.h:943
open3d::core::kernel::DeviceGetNormalAt
OPEN3D_DEVICE void DeviceGetNormalAt(int xo, int yo, int zo, int curr_block_idx, float *n, int resolution, float voxel_size, const NDArrayIndexer &nb_block_masks_indexer, const NDArrayIndexer &nb_block_indices_indexer, const NDArrayIndexer &blocks_indexer)
Definition: GeneralEWSharedImpl.h:218
open3d::core::kernel::Voxel32f::HasColor
static bool HasColor()
Definition: GeneralEWSharedImpl.h:72
open3d::core::Tensor::Item
T Item() const
Definition: Tensor.h:446
Dtype.h
open3d::utility::LogInfo
void LogInfo(const char *format, const Args &... args)
Definition: Console.h:186
open3d::core::kernel::Voxel32f::Integrate
OPEN3D_HOST_DEVICE void Integrate(float dsdf)
Definition: GeneralEWSharedImpl.h:79
open3d::core::kernel::ColoredVoxel32f::Integrate
OPEN3D_HOST_DEVICE void Integrate(float dsdf)
Definition: GeneralEWSharedImpl.h:163
open3d::core::kernel::ColoredVoxel16i::tsdf
float tsdf
Definition: GeneralEWSharedImpl.h:101
Generated by   doxygen 1.8.13