pengc02's picture
upload_part1
7648567
raw
history blame
5.37 kB
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
@torch.no_grad()
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