File size: 1,954 Bytes
909940e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
import os
import torch
from torch.utils.cpp_extension import load
from torch.autograd import Function
from torch.autograd.function import once_differentiable
build_path = os.path.join(os.path.split(os.path.abspath(__file__))[0], 'build')
os.makedirs(build_path, exist_ok=True)
file_path = os.path.split(os.path.abspath(__file__))[0]
GSWrapper = load(
name="gscuda",
# sources=["gs_cuda/gswrapper.cpp", "gs_cuda/gs.cu"],
sources=[os.path.join(file_path, "gswrapper.cpp"),
os.path.join(file_path, "gs.cu")],
build_directory=build_path,
verbose=True)
class GSCUDA(Function):
@staticmethod
def forward(ctx, sigmas, coords, colors, rendered_img):
ctx.save_for_backward(sigmas, coords, colors)
h, w, c = rendered_img.shape
s = sigmas.shape[0]
GSWrapper.gs_render(sigmas, coords, colors, rendered_img, s, h, w, c)
return rendered_img
@staticmethod
@once_differentiable
def backward(ctx, grad_output):
sigmas, coords, colors = ctx.saved_tensors
h, w, c = grad_output.shape
s = sigmas.shape[0]
grads_sigmas = torch.zeros_like(sigmas)
grads_coords = torch.zeros_like(coords)
grads_colors = torch.zeros_like(colors)
GSWrapper.gs_render_backward(sigmas, coords, colors, grad_output.contiguous(), grads_sigmas, grads_coords, grads_colors, s, h, w, c)
return (grads_sigmas, grads_coords, grads_colors, None)
def gaussiansplatting_render(sigmas, coords, colors, image_size):
sigmas = sigmas.contiguous() # (gs num, 3)
coords = coords.contiguous() # (gs num, 2)
colors = colors.contiguous() # (gs num, c)
h, w = image_size[:2]
c = colors.shape[-1]
rendered_img = torch.zeros(h, w, c).to(colors.device).to(torch.float32)
return GSCUDA.apply(sigmas, coords, colors, rendered_img)
|