Spaces:
Running
Running
import numpy as np | |
import torch | |
import torch.nn.functional as F | |
import argparse | |
import tqdm | |
import json | |
import cv2 as cv | |
import os, glob | |
import math | |
from render_utils.lib.utils.graphics_utils import focal2fov, getProjectionMatrix | |
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer | |
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, | |
override_color = None, | |
compute_cov3D_python = False | |
): | |
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. | |
scales = None | |
rotations = None | |
cov3D_precomp = None | |
scales = gaussian_vals['scales'] | |
rotations = gaussian_vals['rotations'] | |
# 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. | |
shs = None | |
# colors_precomp = None | |
# if override_color is None: | |
# shs = gaussian_vals['shs'] | |
# else: | |
# colors_precomp = override_color | |
if 'colors' in gaussian_vals: | |
colors_precomp = gaussian_vals['colors'] | |
else: | |
colors_precomp = None | |
# 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) | |
# Rasterize visible Gaussians to image, obtain their radii (on screen). | |
rendered_image, radii = 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, | |
"viewspace_points": screenspace_points, | |
"visibility_filter": radii > 0, | |
"radii": radii | |
} | |
def blend_color(head_facial_color, body_facial_color, blend_weight): | |
blend_weight = blend_weight.reshape([len(blend_weight)] + [1]*(len(head_facial_color.shape)-1)) | |
result = head_facial_color * blend_weight + body_facial_color * (1-blend_weight) | |
return result | |
def paste_back_with_linear_interp(pasteback_scale, pasteback_center, src, tgt_size): | |
pasteback_topleft = [pasteback_center[0] - src.shape[1]/2/pasteback_scale, | |
pasteback_center[1] - src.shape[0]/2/pasteback_scale] | |
h, w = src.shape[0], src.shape[1] | |
grayscale = False | |
if len(src.shape) == 2: | |
src = src.reshape([h, w, 1]) | |
grayscale = True | |
src = torch.from_numpy(src) | |
src = src.permute(2, 0, 1).unsqueeze(0) | |
grid = torch.meshgrid(torch.arange(0, tgt_size[0]), torch.arange(0, tgt_size[1]), indexing='xy') | |
grid = torch.stack(grid, dim = -1).float().to(src.device).unsqueeze(0) | |
grid[..., 0] = (grid[..., 0] - pasteback_topleft[0]) * pasteback_scale | |
grid[..., 1] = (grid[..., 1] - pasteback_topleft[1]) * pasteback_scale | |
grid[..., 0] = grid[..., 0] / (src.shape[-1] / 2.0) - 1.0 | |
grid[..., 1] = grid[..., 1] / (src.shape[-2] / 2.0) - 1.0 | |
out = F.grid_sample(src, grid, align_corners = True) | |
out = out[0].detach().permute(1, 2, 0).cpu().numpy() | |
if grayscale: | |
out = out[:, :, 0] | |
return out | |
def soften_blending_mask(blending_mask, valid_mask): | |
blending_mask = np.clip(blending_mask*2.0, 0.0, 1.0) | |
blending_mask = cv.erode(blending_mask, np.ones((5, 5))) * valid_mask | |
blending_mask_bk = np.copy(blending_mask) | |
blending_mask = cv.blur(blending_mask*valid_mask, (25, 25)) | |
valid_mask = cv.blur(valid_mask, (25, 25)) | |
blending_mask = blending_mask / (valid_mask + 1e-6) * blending_mask_bk | |
return blending_mask |