# modified from : 2dgs/gaussian_renderer/__init__.py import numpy as np from pdb import set_trace as st import torch import torch.nn as nn import torch.nn.functional as F # from diff_gaussian_rasterization import ( # GaussianRasterizationSettings, # GaussianRasterizer, # ) from torch.profiler import profile, record_function, ProfilerActivity from diff_surfel_rasterization import GaussianRasterizationSettings, GaussianRasterizer from utils.point_utils import depth_to_normal, depth_to_normal_2 import kiui class GaussianRenderer2DGS: def __init__(self, output_size, out_chans, rendering_kwargs, **kwargs): # self.opt = opt self.bg_color = torch.tensor([1, 1, 1], dtype=torch.float32, device="cuda") # self.bg_color = torch.tensor([0,0,1], dtype=torch.float32, device="cuda") self.output_size = output_size self.out_chans = out_chans self.rendering_kwargs = rendering_kwargs # intrinsics # self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy)) # self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32) # self.proj_matrix[0, 0] = 1 / self.tan_half_fov # self.proj_matrix[1, 1] = 1 / self.tan_half_fov # self.proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear) # self.proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear) # self.proj_matrix[2, 3] = 1 def render(self, gaussians, cam_view, cam_view_proj, cam_pos, tanfov, bg_color=None, scale_modifier=1, output_size=None): # gaussians: [B, N, 14-1] # cam_view, cam_view_proj: [B, V, 4, 4] # cam_pos: [B, V, 3] if output_size is None: output_size = self.output_size device = gaussians.device B, V = cam_view.shape[:2] assert gaussians.shape[2] == 13 # scale with 2dof gaussians = gaussians.contiguous().float() # gs rendering in fp32 # loop of loop... images = [] alphas = [] depths = [] # surf_normals = [] rend_normals = [] dists = [] if bg_color is None: bg_color = self.bg_color for b in range(B): # pos, opacity, scale, rotation, shs means3D = gaussians[b, :, 0:3].contiguous().float() opacity = gaussians[b, :, 3:4].contiguous().float() scales = gaussians[b, :, 4:6].contiguous().float() rotations = gaussians[b, :, 6:10].contiguous().float() rgbs = gaussians[b, :, 10:13].contiguous().float() # [N, 3] for v in range(V): # render novel views view_matrix = cam_view[b, v].float() # world_view_transform view_proj_matrix = cam_view_proj[b, v].float() campos = cam_pos[b, v].float() # with profile(activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU,], record_shapes=True) as prof: # with record_function("rendering"): raster_settings = GaussianRasterizationSettings( image_height=output_size, image_width=output_size, tanfovx=tanfov, tanfovy=tanfov, bg=bg_color, scale_modifier=scale_modifier, viewmatrix=view_matrix, projmatrix=view_proj_matrix, sh_degree=0, campos=campos, prefiltered=False, debug=False, ) rasterizer = GaussianRasterizer(raster_settings=raster_settings) # Rasterize visible Gaussians to image, obtain their radii (on screen). # rendered_image, radii, rendered_depth, rendered_alpha = rasterizer( rendered_image, radii, allmap = rasterizer( means3D=means3D, means2D=torch.zeros_like(means3D, dtype=torch.float32, device=device), shs=None, colors_precomp=rgbs, opacities=opacity, scales=scales, rotations=rotations, cov3D_precomp=None, # cov3D_precomp = cov3D_precomp ) # print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) # with profile(activities=[ProfilerActivity.CUDA, ProfilerActivity.CPU,], record_shapes=True) as prof: # ! additional regularizations render_alpha = allmap[1:2] # get normal map # transform normal from view space to world space # with record_function("render_normal"): render_normal = allmap[2:5] # render_normal = (render_normal.permute(1,2,0) @ (viewpoint_camera.world_view_transform[:3,:3].T)).permute(2,0,1) render_normal = (render_normal.permute(1,2,0) @ (view_matrix[:3,:3].T)).permute(2,0,1) # with record_function("render_depth"): # get median depth map render_depth_median = allmap[5:6] render_depth_median = torch.nan_to_num(render_depth_median, 0, 0) # get expected depth map render_depth_expected = allmap[0:1] render_depth_expected = (render_depth_expected / render_alpha) render_depth_expected = torch.nan_to_num(render_depth_expected, 0, 0) # get depth distortion map render_dist = allmap[6:7] # psedo surface attributes # surf depth is either median or expected by setting depth_ratio to 1 or 0 # for bounded scene, use median depth, i.e., depth_ratio = 1; # for unbounded scene, use expected depth, i.e., depth_ration = 0, to reduce disk anliasing. # ! hard coded depth_ratio = 1 for objaverse surf_depth = render_depth_median # with record_function("surf_normal"): # depth_ratio = 1 # # surf_depth = render_depth_expected * (1-depth_ratio) + (depth_ratio) * render_depth_median # # assume the depth points form the 'surface' and generate psudo surface normal for regularizations. # # surf_normal = depth_to_normal(viewpoint_camera, surf_depth) # surf_normal = depth_to_normal_2(world_view_transform=view_matrix, tanfov=tanfov, W=self.output_size, H=self.output_size, depth=surf_depth) # surf_normal = surf_normal.permute(2,0,1) # # remember to multiply with accum_alpha since render_normal is unnormalized. # surf_normal = surf_normal * (render_alpha).detach() # ! images rendered_image = rendered_image.clamp(0, 1) # images.append(rendered_image) # alphas.append(rendered_alpha) # depths.append(rendered_depth) images.append(rendered_image) alphas.append(render_alpha) depths.append(surf_depth) # surf_normals.append(surf_normal) rend_normals.append(render_normal) dists.append(render_dist) # print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20)) # st() pass images = torch.stack(images, dim=0).view(B, V, 3, output_size, output_size) alphas = torch.stack(alphas, dim=0).view(B, V, 1, output_size, output_size) depths = torch.stack(depths, dim=0).view(B, V, 1, output_size, output_size) # approximated surface normal? No, direct depth supervision here. # surf_normals = torch.stack(surf_normals, dim=0).view(B, V, 3, self.output_size, self.output_size) # disk normal rend_normals = torch.stack(rend_normals, dim=0).view(B, V, 3, output_size, output_size) dists = torch.stack(dists, dim=0).view(B, V, 1, output_size, output_size) # images = torch.stack(images, dim=0).view(B*V, 3, self.output_size, self.output_size) # alphas = torch.stack(alphas, dim=0).view(B*V, 1, self.output_size, self.output_size) # depths = torch.stack(depths, dim=0).view(B*V, 1, self.output_size, self.output_size) return { "image": images, # [B, V, 3, H, W] "alpha": alphas, # [B, V, 1, H, W] "depth": depths, # "surf_normal": surf_normals, "rend_normal": rend_normals, "dist": dists } # TODO, save/load 2dgs Gaussians def save_2dgs_ply(self, path, gaussians, compatible=True): # gaussians: [B, N, 13] mkdir_p(os.path.dirname(path)) assert gaussians.shape[0] == 1, 'only support batch size 1' from plyfile import PlyData, PlyElement means3D = gaussians[0, :, 0:3].contiguous().float() opacity = gaussians[0, :, 3:4].contiguous().float() scales = gaussians[0, :, 4:6].contiguous().float() rotations = gaussians[0, :, 6:10].contiguous().float() shs = gaussians[0, :, 10:].unsqueeze(1).contiguous().float() # [N, 1, 3] # invert activation to make it compatible with the original ply format if compatible: opacity = kiui.op.inverse_sigmoid(opacity) scales = torch.log(scales + 1e-8) shs = (shs - 0.5) / 0.28209479177387814 xyzs = means3D.detach().cpu().numpy() f_dc = shs.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() opacities = opacity.detach().cpu().numpy() scales = scales.detach().cpu().numpy() rotations = rotations.detach().cpu().numpy() # xyz = self._xyz.detach().cpu().numpy() # normals = np.zeros_like(xyz) # f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() # f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() # opacities = self._opacity.detach().cpu().numpy() # scale = self._scaling.detach().cpu().numpy() # rotation = self._rotation.detach().cpu().numpy() # dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] l = ['x', 'y', 'z'] # All channels except the 3 DC for i in range(f_dc.shape[1]): l.append('f_dc_{}'.format(i)) # save normals also for i in range(f_dc.shape[1]): l.append('f_dc_{}'.format(i)) l.append('opacity') for i in range(scales.shape[1]): l.append('scale_{}'.format(i)) for i in range(rotations.shape[1]): l.append('rot_{}'.format(i)) dtype_full = [(attribute, 'f4') for attribute in l] elements = np.empty(xyz.shape[0], dtype=dtype_full) # attributes = np.concatenate((xyzs, f_dc, opacities, scales, rotations), axis=1) # attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1) attributes = np.concatenate((xyz, normals, f_dc, opacities, scale, rotation), axis=1) elements[:] = list(map(tuple, attributes)) el = PlyElement.describe(elements, 'vertex') PlyData([el]).write(path) # def save_ply(self, gaussians, path, compatible=True): # # gaussians: [B, N, 14] # # compatible: save pre-activated gaussians as in the original paper # assert gaussians.shape[0] == 1, 'only support batch size 1' # from plyfile import PlyData, PlyElement # means3D = gaussians[0, :, 0:3].contiguous().float() # opacity = gaussians[0, :, 3:4].contiguous().float() # scales = gaussians[0, :, 4:7].contiguous().float() # rotations = gaussians[0, :, 7:11].contiguous().float() # shs = gaussians[0, :, 11:].unsqueeze(1).contiguous().float() # [N, 1, 3] # # prune by opacity # mask = opacity.squeeze(-1) >= 0.005 # means3D = means3D[mask] # opacity = opacity[mask] # scales = scales[mask] # rotations = rotations[mask] # shs = shs[mask] # # invert activation to make it compatible with the original ply format # if compatible: # opacity = kiui.op.inverse_sigmoid(opacity) # scales = torch.log(scales + 1e-8) # shs = (shs - 0.5) / 0.28209479177387814 # xyzs = means3D.detach().cpu().numpy() # f_dc = shs.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy() # opacities = opacity.detach().cpu().numpy() # scales = scales.detach().cpu().numpy() # rotations = rotations.detach().cpu().numpy() # l = ['x', 'y', 'z'] # # All channels except the 3 DC # for i in range(f_dc.shape[1]): # l.append('f_dc_{}'.format(i)) # l.append('opacity') # for i in range(scales.shape[1]): # l.append('scale_{}'.format(i)) # for i in range(rotations.shape[1]): # l.append('rot_{}'.format(i)) # dtype_full = [(attribute, 'f4') for attribute in l] # elements = np.empty(xyzs.shape[0], dtype=dtype_full) # attributes = np.concatenate((xyzs, f_dc, opacities, scales, rotations), axis=1) # elements[:] = list(map(tuple, attributes)) # el = PlyElement.describe(elements, 'vertex') # PlyData([el]).write(path) def load_2dgs_ply(self, path, compatible=True): from plyfile import PlyData, PlyElement plydata = PlyData.read(path) xyz = np.stack((np.asarray(plydata.elements[0]["x"]), np.asarray(plydata.elements[0]["y"]), np.asarray(plydata.elements[0]["z"])), axis=1) print("Number of points at loading : ", xyz.shape[0]) opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] shs = np.zeros((xyz.shape[0], 3)) shs[:, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) shs[:, 1] = np.asarray(plydata.elements[0]["f_dc_1"]) shs[:, 2] = np.asarray(plydata.elements[0]["f_dc_2"]) scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")] scales = np.zeros((xyz.shape[0], len(scale_names))) for idx, attr_name in enumerate(scale_names): scales[:, idx] = np.asarray(plydata.elements[0][attr_name]) rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot_")] rots = np.zeros((xyz.shape[0], len(rot_names))) for idx, attr_name in enumerate(rot_names): rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) normal_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot_")] rots = np.zeros((xyz.shape[0], len(rot_names))) for idx, attr_name in enumerate(rot_names): rots[:, idx] = np.asarray(plydata.elements[0][attr_name]) gaussians = np.concatenate([xyz, opacities, scales, rots, shs], axis=1) gaussians = torch.from_numpy(gaussians).float() # cpu if compatible: gaussians[..., 3:4] = torch.sigmoid(gaussians[..., 3:4]) gaussians[..., 4:7] = torch.exp(gaussians[..., 4:7]) gaussians[..., 11:] = 0.28209479177387814 * gaussians[..., 11:] + 0.5 return gaussians