souging/TRELLIS_TextTo3D
0
1#2# Copyright (C) 2023, Inria3# GRAPHDECO research group, https://team.inria.fr/graphdeco4# All rights reserved.5#6# This software is free for non-commercial, research and evaluation use 7# under the terms of the LICENSE.md file.8#9# For inquiries contact george.drettakis@inria.fr10#11 12import torch13import math14from easydict import EasyDict as edict15import numpy as np16from ..representations.gaussian import Gaussian17from .sh_utils import eval_sh18import torch.nn.functional as F19from easydict import EasyDict as edict20 21 22def intrinsics_to_projection(23 intrinsics: torch.Tensor,24 near: float,25 far: float,26 ) -> torch.Tensor:27 """28 OpenCV intrinsics to OpenGL perspective matrix29 30 Args:31 intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix32 near (float): near plane to clip33 far (float): far plane to clip34 Returns:35 (torch.Tensor): [4, 4] OpenGL perspective matrix36 """37 fx, fy = intrinsics[0, 0], intrinsics[1, 1]38 cx, cy = intrinsics[0, 2], intrinsics[1, 2]39 ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device)40 ret[0, 0] = 2 * fx41 ret[1, 1] = 2 * fy42 ret[0, 2] = 2 * cx - 143 ret[1, 2] = - 2 * cy + 144 ret[2, 2] = far / (far - near)45 ret[2, 3] = near * far / (near - far)46 ret[3, 2] = 1.47 return ret48 49 50def render(viewpoint_camera, pc : Gaussian, pipe, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None):51 """52 Render the scene. 53 54 Background tensor (bg_color) must be on GPU!55 """56 # lazy import57 if 'GaussianRasterizer' not in globals():58 from diff_gaussian_rasterization import GaussianRasterizer, GaussianRasterizationSettings59 60 # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means61 screenspace_points = torch.zeros_like(pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda") + 062 try:63 screenspace_points.retain_grad()64 except:65 pass66 # Set up rasterization configuration67 tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)68 tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)69 70 kernel_size = pipe.kernel_size71 subpixel_offset = torch.zeros((int(viewpoint_camera.image_height), int(viewpoint_camera.image_width), 2), dtype=torch.float32, device="cuda")72 73 raster_settings = GaussianRasterizationSettings(74 image_height=int(viewpoint_camera.image_height),75 image_width=int(viewpoint_camera.image_width),76 tanfovx=tanfovx,77 tanfovy=tanfovy,78 kernel_size=kernel_size,79 subpixel_offset=subpixel_offset,80 bg=bg_color,81 scale_modifier=scaling_modifier,82 viewmatrix=viewpoint_camera.world_view_transform,83 projmatrix=viewpoint_camera.full_proj_transform,84 sh_degree=pc.active_sh_degree,85 campos=viewpoint_camera.camera_center,86 prefiltered=False,87 debug=pipe.debug88 )89 90 rasterizer = GaussianRasterizer(raster_settings=raster_settings)91 92 means3D = pc.get_xyz93 means2D = screenspace_points94 opacity = pc.get_opacity95 96 # If precomputed 3d covariance is provided, use it. If not, then it will be computed from97 # scaling / rotation by the rasterizer.98 scales = None99 rotations = None100 cov3D_precomp = None101 if pipe.compute_cov3D_python:102 cov3D_precomp = pc.get_covariance(scaling_modifier)103 else:104 scales = pc.get_scaling105 rotations = pc.get_rotation106 107 # If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors108 # from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.109 shs = None110 colors_precomp = None111 if override_color is None:112 if pipe.convert_SHs_python:113 shs_view = pc.get_features.transpose(1, 2).view(-1, 3, (pc.max_sh_degree+1)**2)114 dir_pp = (pc.get_xyz - viewpoint_camera.camera_center.repeat(pc.get_features.shape[0], 1))115 dir_pp_normalized = dir_pp/dir_pp.norm(dim=1, keepdim=True)116 sh2rgb = eval_sh(pc.active_sh_degree, shs_view, dir_pp_normalized)117 colors_precomp = torch.clamp_min(sh2rgb + 0.5, 0.0)118 else:119 shs = pc.get_features120 else:121 colors_precomp = override_color122 123 # Rasterize visible Gaussians to image, obtain their radii (on screen). 124 rendered_image, radii = rasterizer(125 means3D = means3D,126 means2D = means2D,127 shs = shs,128 colors_precomp = colors_precomp,129 opacities = opacity,130 scales = scales,131 rotations = rotations,132 cov3D_precomp = cov3D_precomp133 )134 135 # Those Gaussians that were frustum culled or had a radius of 0 were not visible.136 # They will be excluded from value updates used in the splitting criteria.137 return edict({"render": rendered_image,138 "viewspace_points": screenspace_points,139 "visibility_filter" : radii > 0,140 "radii": radii})141 142 143class GaussianRenderer:144 """145 Renderer for the Voxel representation.146 147 Args:148 rendering_options (dict): Rendering options.149 """150 151 def __init__(self, rendering_options={}) -> None:152 self.pipe = edict({153 "kernel_size": 0.1,154 "convert_SHs_python": False,155 "compute_cov3D_python": False,156 "scale_modifier": 1.0,157 "debug": False158 })159 self.rendering_options = edict({160 "resolution": None,161 "near": None,162 "far": None,163 "ssaa": 1,164 "bg_color": 'random',165 })166 self.rendering_options.update(rendering_options)167 self.bg_color = None168 169 def render(170 self,171 gausssian: Gaussian,172 extrinsics: torch.Tensor,173 intrinsics: torch.Tensor,174 colors_overwrite: torch.Tensor = None175 ) -> edict:176 """177 Render the gausssian.178 179 Args:180 gaussian : gaussianmodule181 extrinsics (torch.Tensor): (4, 4) camera extrinsics182 intrinsics (torch.Tensor): (3, 3) camera intrinsics183 colors_overwrite (torch.Tensor): (N, 3) override color184 185 Returns:186 edict containing:187 color (torch.Tensor): (3, H, W) rendered color image188 """189 resolution = self.rendering_options["resolution"]190 near = self.rendering_options["near"]191 far = self.rendering_options["far"]192 ssaa = self.rendering_options["ssaa"]193 194 if self.rendering_options["bg_color"] == 'random':195 self.bg_color = torch.zeros(3, dtype=torch.float32, device="cuda")196 if np.random.rand() < 0.5:197 self.bg_color += 1198 else:199 self.bg_color = torch.tensor(self.rendering_options["bg_color"], dtype=torch.float32, device="cuda")200 201 view = extrinsics202 perspective = intrinsics_to_projection(intrinsics, near, far)203 camera = torch.inverse(view)[:3, 3]204 focalx = intrinsics[0, 0]205 focaly = intrinsics[1, 1]206 fovx = 2 * torch.atan(0.5 / focalx)207 fovy = 2 * torch.atan(0.5 / focaly)208 209 camera_dict = edict({210 "image_height": resolution * ssaa,211 "image_width": resolution * ssaa,212 "FoVx": fovx,213 "FoVy": fovy,214 "znear": near,215 "zfar": far,216 "world_view_transform": view.T.contiguous(),217 "projection_matrix": perspective.T.contiguous(),218 "full_proj_transform": (perspective @ view).T.contiguous(),219 "camera_center": camera220 })221 222 # Render223 render_ret = render(camera_dict, gausssian, self.pipe, self.bg_color, override_color=colors_overwrite, scaling_modifier=self.pipe.scale_modifier)224 225 if ssaa > 1:226 render_ret.render = F.interpolate(render_ret.render[None], size=(resolution, resolution), mode='bilinear', align_corners=False, antialias=True).squeeze()227 228 ret = edict({229 'color': render_ret['render']230 })231 return ret232