Spaces:
Sleeping
Sleeping
| 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 |