Spaces:
Running
Running
File size: 4,184 Bytes
ec9a6bc |
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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use
# under the terms of the LICENSE.md file.
#
# For inquiries contact george.drettakis@inria.fr
#
import torch
import math
from diff_gaussian_rasterization_depth_alpha import GaussianRasterizationSettings, GaussianRasterizer
from AnimatableGaussians.utils.graphics_utils import focal2fov, getProjectionMatrix
from AnimatableGaussians.utils.sh_utils import eval_sh
def render3(
gaussian_vals: dict,
bg_color: torch.Tensor,
extr: torch.Tensor,
intr: torch.Tensor,
img_w: int,
img_h: int,
scaling_modifier = 1.0,
):
means3D = gaussian_vals['positions']
# Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
screenspace_points = torch.zeros_like(means3D, dtype = means3D.dtype, requires_grad = True, device = "cuda") + 0
try:
screenspace_points.retain_grad()
except:
pass
means2D = screenspace_points
opacity = gaussian_vals['opacity']
# If precomputed 3d covariance is provided, use it. If not, then it will be computed from
# scaling / rotation by the rasterizer.
cov3D_precomp = None
scales = gaussian_vals['scales']
rotations = gaussian_vals['rotations']
# Set up rasterization configuration
FoVx = focal2fov(intr[0, 0].item(), img_w)
FoVy = focal2fov(intr[1, 1].item(), img_h)
tanfovx = math.tan(FoVx * 0.5)
tanfovy = math.tan(FoVy * 0.5)
world_view_transform = extr.transpose(1, 0).cuda()
projection_matrix = getProjectionMatrix(znear = 0.1, zfar = 100, fovX = FoVx, fovY = FoVy, K = intr, img_w = img_w, img_h = img_h).transpose(0, 1).cuda()
full_proj_transform = (world_view_transform.unsqueeze(0).bmm(projection_matrix.unsqueeze(0))).squeeze(0)
camera_center = torch.linalg.inv(extr)[:3, 3]
raster_settings = GaussianRasterizationSettings(
image_height = img_h,
image_width = img_w,
tanfovx = tanfovx,
tanfovy = tanfovy,
bg = bg_color,
scale_modifier = scaling_modifier,
viewmatrix = world_view_transform,
projmatrix = full_proj_transform,
sh_degree = gaussian_vals['max_sh_degree'],
campos = camera_center,
prefiltered = False,
debug = False
)
rasterizer = GaussianRasterizer(raster_settings = raster_settings)
# If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
# from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
assert not ('colors' in gaussian_vals and 'shs' in gaussian_vals), "Cannot use both color and SH!"
if 'colors' in gaussian_vals:
colors_precomp = gaussian_vals['colors']
else:
colors_precomp = None
if 'shs' in gaussian_vals:
shs_view = gaussian_vals['shs']
dir_pp = (means3D - camera_center.repeat(means3D.shape[0], 1))
dir_pp_normalized = dir_pp / dir_pp.norm(dim = 1, keepdim = True)
sh2rgb = eval_sh(gaussian_vals['max_sh_degree'], shs_view, dir_pp_normalized)
colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)
shs = None
# Rasterize visible Gaussians to image, obtain their radii (on screen).
rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
means3D = means3D,
means2D = means2D,
shs = shs,
colors_precomp = colors_precomp,
opacities = opacity,
scales = scales,
rotations = rotations,
cov3D_precomp = cov3D_precomp)
# Those Gaussians that were frustum culled or had a radius of 0 were not visible.
# They will be excluded from value updates used in the splitting criteria.
return {
"render": rendered_image,
"depth": rendered_depth,
"mask": rendered_alpha,
"viewspace_points": screenspace_points,
"visibility_filter": radii > 0,
"radii": radii
}
|