mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 18:44:21 +08:00
146332cfb6
* in progress. camera_indices broken * faster dataloader. starting on pixel sampler logic * imagesampler and pixelsampler * lint and black * continuing merge * pytest working * fixed docs * docs * minor changes
30 lines
732 B
C++
30 lines
732 B
C++
/*
|
|
This file uses pybind to make CUDA calls.
|
|
*/
|
|
|
|
#include <torch/extension.h>
|
|
|
|
#include <vector>
|
|
|
|
// CUDA forward declarations
|
|
|
|
std::vector<torch::Tensor> sample_uniformly_along_ray_bundle(
|
|
torch::Tensor origins,
|
|
torch::Tensor directions,
|
|
torch::Tensor nears,
|
|
torch::Tensor fars,
|
|
torch::Tensor offsets,
|
|
int max_num_samples);
|
|
|
|
// C++ interface
|
|
|
|
#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
|
|
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
|
#define CHECK_INPUT(x) \
|
|
CHECK_CUDA(x); \
|
|
CHECK_CONTIGUOUS(x)
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
m.def("sample_uniformly_along_ray_bundle", &sample_uniformly_along_ray_bundle);
|
|
} |