mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2026-02-26 16:26:00 +08:00
3ff6c5ab85
Summary: Adds CHECK_CPU macros that checks if a tensor is on the CPU device throughout csrc directories up to `marching_cubes`. Directories updated include those in `gather_scatter`, `interp_face_attrs`, `iou_box3d`, `knn`, and `marching_cubes`. Note that this is the second part of a larger change, and to keep diffs better organized, subsequent diffs will update the remaining directories. Reviewed By: bottler Differential Revision: D77558550 fbshipit-source-id: 762a0fe88548dc8d0901b198a11c40d0c36e173f
64 lines
2.0 KiB
C++
64 lines
2.0 KiB
C++
/*
|
|
* Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
* All rights reserved.
|
|
*
|
|
* This source code is licensed under the BSD-style license found in the
|
|
* LICENSE file in the root directory of this source tree.
|
|
*/
|
|
|
|
#pragma once
|
|
#include <torch/extension.h>
|
|
#include <tuple>
|
|
#include <vector>
|
|
#include "utils/pytorch3d_cutils.h"
|
|
|
|
// Run Marching Cubes algorithm over a batch of volume scalar fields
|
|
// with a pre-defined threshold and return a mesh composed of vertices
|
|
// and faces for the mesh.
|
|
//
|
|
// Args:
|
|
// vol: FloatTensor of shape (D, H, W) giving a volume
|
|
// scalar grids.
|
|
// isolevel: isosurface value to use as the threshoold to determine whether
|
|
// the points are within a volume.
|
|
//
|
|
// Returns:
|
|
// vertices: (N_verts, 3) FloatTensor of vertices
|
|
// faces: (N_faces, 3) LongTensor of faces
|
|
// ids: (N_verts,) LongTensor used to identify each vertex and deduplication
|
|
// to avoid floating point precision issues.
|
|
// For Cuda, will be used to dedupe redundant vertices.
|
|
// For cpp implementation, this tensor is just a placeholder.
|
|
|
|
// CPU implementation
|
|
std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCpu(
|
|
const at::Tensor& vol,
|
|
const float isolevel);
|
|
|
|
// CUDA implementation
|
|
std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubesCuda(
|
|
const at::Tensor& vol,
|
|
const float isolevel);
|
|
|
|
// Implementation which is exposed
|
|
inline std::tuple<at::Tensor, at::Tensor, at::Tensor> MarchingCubes(
|
|
const at::Tensor& vol,
|
|
const float isolevel) {
|
|
if (vol.is_cuda()) {
|
|
#ifdef WITH_CUDA
|
|
CHECK_CUDA(vol);
|
|
const int D = vol.size(0);
|
|
const int H = vol.size(1);
|
|
const int W = vol.size(2);
|
|
if (D > 1024 || H > 1024 || W > 1024) {
|
|
AT_ERROR("Maximum volume size allowed 1K x 1K x 1K");
|
|
}
|
|
return MarchingCubesCuda(vol.contiguous(), isolevel);
|
|
#else
|
|
AT_ERROR("Not compiled with GPU support.");
|
|
#endif
|
|
}
|
|
CHECK_CPU(vol);
|
|
return MarchingCubesCpu(vol.contiguous(), isolevel);
|
|
}
|