Spaces:
Running on Zero
Running on Zero
| # conda install -y pytorch pytorch-cuda=12.1 -c pytorch -c nvidia | |
| # pip install Pillow plyfile opencv-python numpy git+https://github.com/graphdeco-inria/diff-gaussian-rasterization | |
| import math | |
| import os | |
| import cv2 | |
| import matplotlib | |
| import numpy as np | |
| import torch | |
| from diff_gaussian_rasterization import ( | |
| GaussianRasterizationSettings, | |
| GaussianRasterizer, | |
| ) | |
| from einops import rearrange | |
| from plyfile import PlyData, PlyElement | |
| from torch import nn | |
| from collections import OrderedDict | |
| import imageio | |
| from typing import Union | |
| from pdb import set_trace as stx | |
| import kiui | |
| from diffusionGS.utils.mesh_utils import decimate_mesh, clean_mesh | |
| import trimesh | |
| def gaussian_3d_coeff(xyzs, covs): | |
| # xyzs: [N, 3] | |
| # covs: [N, 6] | |
| x, y, z = xyzs[:, 0], xyzs[:, 1], xyzs[:, 2] | |
| a, b, c, d, e, f = covs[:, 0], covs[:, 1], covs[:, 2], covs[:, 3], covs[:, 4], covs[:, 5] | |
| # eps must be small enough !!! | |
| inv_det = 1 / (a * d * f + 2 * e * c * b - e**2 * a - c**2 * d - b**2 * f + 1e-24) | |
| inv_a = (d * f - e**2) * inv_det | |
| inv_b = (e * c - b * f) * inv_det | |
| inv_c = (e * b - c * d) * inv_det | |
| inv_d = (a * f - c**2) * inv_det | |
| inv_e = (b * c - e * a) * inv_det | |
| inv_f = (a * d - b**2) * inv_det | |
| power = -0.5 * (x**2 * inv_a + y**2 * inv_d + z**2 * inv_f) - x * y * inv_b - x * z * inv_c - y * z * inv_e | |
| power[power > 0] = -1e10 # abnormal values... make weights 0 | |
| return torch.exp(power) | |
| def get_turntable_cameras( | |
| hfov=50, | |
| num_views=8, | |
| w=384, | |
| h=384, | |
| radius=2.7, | |
| elevation=0, | |
| up_vector=np.array([0, 0, 1]), | |
| ): | |
| fx = w / (2 * np.tan(np.deg2rad(hfov) / 2.0)) | |
| fy = fx | |
| cx, cy = w / 2.0, h / 2.0 | |
| fxfycxcy = ( | |
| np.array([fx, fy, cx, cy]).reshape(1, 4).repeat(num_views, axis=0) | |
| ) # [num_views, 4] | |
| azimuths = np.linspace(0, 360, num_views, endpoint=False) | |
| elevations = np.ones_like(azimuths) * elevation | |
| c2ws = [] | |
| for elev, azim in zip(elevations, azimuths): | |
| elev, azim = np.deg2rad(elev), np.deg2rad(azim) | |
| z = radius * np.sin(elev) | |
| base = radius * np.cos(elev) | |
| x = base * np.cos(azim) | |
| y = base * np.sin(azim) | |
| cam_pos = np.array([x, y, z]) | |
| forward = -cam_pos / np.linalg.norm(cam_pos) | |
| right = np.cross(forward, up_vector) | |
| right = right / np.linalg.norm(right) | |
| up = np.cross(right, forward) | |
| up = up / np.linalg.norm(up) | |
| R = np.stack((right, -up, forward), axis=1) | |
| c2w = np.eye(4) | |
| c2w[:3, :4] = np.concatenate((R, cam_pos[:, None]), axis=1) | |
| c2ws.append(c2w) | |
| c2ws = np.stack(c2ws, axis=0) # [num_views, 4, 4] | |
| return w, h, num_views, fxfycxcy, c2ws | |
| def imageseq2video(images, filename, fps=24): | |
| # if images is uint8, convert to float32 | |
| if images.dtype == np.uint8: | |
| images = images.astype(np.float32) / 255.0 | |
| imageio.mimwrite(filename, images, fps=fps, quality=8) | |
| # copied from: utils.general_utils | |
| def strip_lowerdiag(L): | |
| uncertainty = torch.zeros((L.shape[0], 6), dtype=torch.float, device=L.device) | |
| uncertainty[:, 0] = L[:, 0, 0] | |
| uncertainty[:, 1] = L[:, 0, 1] | |
| uncertainty[:, 2] = L[:, 0, 2] | |
| uncertainty[:, 3] = L[:, 1, 1] | |
| uncertainty[:, 4] = L[:, 1, 2] | |
| uncertainty[:, 5] = L[:, 2, 2] | |
| return uncertainty | |
| def strip_symmetric(sym): | |
| return strip_lowerdiag(sym) | |
| def build_rotation(r): | |
| norm = torch.sqrt( | |
| r[:, 0] * r[:, 0] + r[:, 1] * r[:, 1] + r[:, 2] * r[:, 2] + r[:, 3] * r[:, 3] | |
| ) | |
| q = r / norm[:, None] | |
| R = torch.zeros((q.size(0), 3, 3), device=r.device) | |
| r = q[:, 0] | |
| x = q[:, 1] | |
| y = q[:, 2] | |
| z = q[:, 3] | |
| R[:, 0, 0] = 1 - 2 * (y * y + z * z) | |
| R[:, 0, 1] = 2 * (x * y - r * z) | |
| R[:, 0, 2] = 2 * (x * z + r * y) | |
| R[:, 1, 0] = 2 * (x * y + r * z) | |
| R[:, 1, 1] = 1 - 2 * (x * x + z * z) | |
| R[:, 1, 2] = 2 * (y * z - r * x) | |
| R[:, 2, 0] = 2 * (x * z - r * y) | |
| R[:, 2, 1] = 2 * (y * z + r * x) | |
| R[:, 2, 2] = 1 - 2 * (x * x + y * y) | |
| return R | |
| def build_scaling_rotation(s, r): | |
| L = torch.zeros((s.shape[0], 3, 3), dtype=torch.float, device=s.device) | |
| R = build_rotation(r) | |
| L[:, 0, 0] = s[:, 0] | |
| L[:, 1, 1] = s[:, 1] | |
| L[:, 2, 2] = s[:, 2] | |
| L = R @ L | |
| return L | |
| # copied from: utils.sh_utils | |
| C0 = 0.28209479177387814 | |
| C1 = 0.4886025119029199 | |
| C2 = [ | |
| 1.0925484305920792, | |
| -1.0925484305920792, | |
| 0.31539156525252005, | |
| -1.0925484305920792, | |
| 0.5462742152960396, | |
| ] | |
| C3 = [ | |
| -0.5900435899266435, | |
| 2.890611442640554, | |
| -0.4570457994644658, | |
| 0.3731763325901154, | |
| -0.4570457994644658, | |
| 1.445305721320277, | |
| -0.5900435899266435, | |
| ] | |
| C4 = [ | |
| 2.5033429417967046, | |
| -1.7701307697799304, | |
| 0.9461746957575601, | |
| -0.6690465435572892, | |
| 0.10578554691520431, | |
| -0.6690465435572892, | |
| 0.47308734787878004, | |
| -1.7701307697799304, | |
| 0.6258357354491761, | |
| ] | |
| def eval_sh(deg, sh, dirs): | |
| """ | |
| Evaluate spherical harmonics at unit directions | |
| using hardcoded SH polynomials. | |
| Works with torch/np/jnp. | |
| ... Can be 0 or more batch dimensions. | |
| Args: | |
| deg: int SH deg. Currently, 0-3 supported | |
| sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2] | |
| dirs: jnp.ndarray unit directions [..., 3] | |
| Returns: | |
| [..., C] | |
| """ | |
| assert deg <= 4 and deg >= 0 | |
| coeff = (deg + 1) ** 2 | |
| assert sh.shape[-1] >= coeff | |
| result = C0 * sh[..., 0] | |
| if deg > 0: | |
| x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3] | |
| result = ( | |
| result - C1 * y * sh[..., 1] + C1 * z * sh[..., 2] - C1 * x * sh[..., 3] | |
| ) | |
| if deg > 1: | |
| xx, yy, zz = x * x, y * y, z * z | |
| xy, yz, xz = x * y, y * z, x * z | |
| result = ( | |
| result | |
| + C2[0] * xy * sh[..., 4] | |
| + C2[1] * yz * sh[..., 5] | |
| + C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] | |
| + C2[3] * xz * sh[..., 7] | |
| + C2[4] * (xx - yy) * sh[..., 8] | |
| ) | |
| if deg > 2: | |
| result = ( | |
| result | |
| + C3[0] * y * (3 * xx - yy) * sh[..., 9] | |
| + C3[1] * xy * z * sh[..., 10] | |
| + C3[2] * y * (4 * zz - xx - yy) * sh[..., 11] | |
| + C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] | |
| + C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] | |
| + C3[5] * z * (xx - yy) * sh[..., 14] | |
| + C3[6] * x * (xx - 3 * yy) * sh[..., 15] | |
| ) | |
| if deg > 3: | |
| result = ( | |
| result | |
| + C4[0] * xy * (xx - yy) * sh[..., 16] | |
| + C4[1] * yz * (3 * xx - yy) * sh[..., 17] | |
| + C4[2] * xy * (7 * zz - 1) * sh[..., 18] | |
| + C4[3] * yz * (7 * zz - 3) * sh[..., 19] | |
| + C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] | |
| + C4[5] * xz * (7 * zz - 3) * sh[..., 21] | |
| + C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] | |
| + C4[7] * xz * (xx - 3 * yy) * sh[..., 23] | |
| + C4[8] | |
| * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) | |
| * sh[..., 24] | |
| ) | |
| return result | |
| def RGB2SH(rgb): | |
| return (rgb - 0.5) / C0 | |
| def SH2RGB(sh): | |
| return sh * C0 + 0.5 | |
| def create_video(image_folder, output_video_file, framerate=30): | |
| # Get all image file paths to a list. | |
| images = [img for img in os.listdir(image_folder) if img.endswith(".png")] | |
| images.sort() | |
| # Read the first image to know the height and width | |
| frame = cv2.imread(os.path.join(image_folder, images[0])) | |
| height, width, layers = frame.shape | |
| video = cv2.VideoWriter( | |
| output_video_file, cv2.VideoWriter_fourcc(*"mp4v"), framerate, (width, height) | |
| ) | |
| # iterate over each image and add it to the video sequence | |
| for image in images: | |
| video.write(cv2.imread(os.path.join(image_folder, image))) | |
| cv2.destroyAllWindows() | |
| video.release() | |
| class Camera(nn.Module): | |
| def __init__(self, C2W, fxfycxcy, h, w): | |
| """ | |
| C2W: 4x4 camera-to-world matrix; opencv convention | |
| fxfycxcy: 4 | |
| """ | |
| super().__init__() | |
| self.C2W = C2W.clone().float() | |
| self.W2C = self.C2W.inverse() | |
| self.h = h | |
| self.w = w | |
| self.znear = 0.01 | |
| self.zfar = 100.0 | |
| fx, fy, cx, cy = fxfycxcy[0], fxfycxcy[1], fxfycxcy[2], fxfycxcy[3] | |
| self.tanfovX = w / (2 * fx) | |
| self.tanfovY = h / (2 * fy) | |
| def getProjectionMatrix(W, H, fx, fy, cx, cy, znear, zfar): | |
| P = torch.zeros(4, 4, device=fx.device) | |
| P[0, 0] = 2 * fx / W | |
| P[1, 1] = 2 * fy / H | |
| P[0, 2] = 2 * (cx / W) - 1 | |
| P[1, 2] = 2 * (cy / H) - 1 | |
| P[2, 2] = -(zfar + znear) / (zfar - znear) | |
| P[3, 2] = 1.0 | |
| P[2, 3] = -(2 * zfar * znear) / (zfar - znear) | |
| return P | |
| self.world_view_transform = self.W2C.transpose(0, 1) | |
| self.projection_matrix = getProjectionMatrix( | |
| self.w, self.h, fx, fy, cx, cy, self.znear, self.zfar | |
| ).transpose(0, 1) | |
| self.full_proj_transform = ( | |
| self.world_view_transform.unsqueeze(0).bmm( | |
| self.projection_matrix.unsqueeze(0) | |
| ) | |
| ).squeeze(0) | |
| self.camera_center = self.C2W[:3, 3] | |
| # modified from scene/gaussian_model.py | |
| # 从 Gaussian Splatting 中拿来的代码 | |
| class GaussianModel: | |
| # 将各类函数 setup | |
| def setup_functions(self): | |
| def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation): | |
| L = build_scaling_rotation(scaling_modifier * scaling, rotation) | |
| actual_covariance = L @ L.transpose(1, 2) | |
| symm = strip_symmetric(actual_covariance) | |
| return symm | |
| self.scaling_activation = torch.exp | |
| self.inv_scaling_activation = torch.log | |
| self.rotation_activation = torch.nn.functional.normalize | |
| self.opacity_activation = torch.sigmoid | |
| self.covariance_activation = build_covariance_from_scaling_rotation | |
| def __init__(self, sh_degree: int, scaling_modifier=None): | |
| self.sh_degree = sh_degree | |
| self._xyz = torch.empty(0) # 创建一个空的张量 | |
| self._features_dc = torch.empty(0) | |
| if self.sh_degree > 0: | |
| self._features_rest = torch.empty(0) | |
| else: | |
| self._features_rest = None | |
| self._scaling = torch.empty(0) | |
| self._rotation = torch.empty(0) | |
| self._opacity = torch.empty(0) | |
| self.setup_functions() # 将各类函数 setup | |
| self.scaling_modifier = scaling_modifier | |
| # 重新初始化 | |
| def empty(self): | |
| self.__init__(self.sh_degree, self.scaling_modifier) | |
| # 将各个属性设置好 | |
| def set_data(self, xyz, features, scaling, rotation, opacity): | |
| """ | |
| xyz : torch.tensor of shape (N, 3) | |
| features : torch.tensor of shape (N, (self.sh_degree + 1) ** 2, 3) | |
| scaling : torch.tensor of shape (N, 3) | |
| rotation : torch.tensor of shape (N, 4) | |
| opacity : torch.tensor of shape (N, 1) | |
| """ | |
| self._xyz = xyz | |
| self._features_dc = features[:, :1, :].contiguous() | |
| if self.sh_degree > 0: | |
| self._features_rest = features[:, 1:, :].contiguous() | |
| else: | |
| self._features_rest = None | |
| self._scaling = scaling | |
| self._rotation = rotation | |
| self._opacity = opacity | |
| return self | |
| # 放到 cuda 上边 | |
| def to(self, device): | |
| self._xyz = self._xyz.to(device) | |
| self._features_dc = self._features_dc.to(device) | |
| if self.sh_degree > 0: | |
| self._features_rest = self._features_rest.to(device) | |
| self._scaling = self._scaling.to(device) | |
| self._rotation = self._rotation.to(device) | |
| self._opacity = self._opacity.to(device) | |
| return self | |
| # 用一个 mask 去把一些数据选出来，为什么要 filter 掉一些数据呢？ | |
| # 因为要 crop 出中间的一个 Gaussian Splatting | |
| ''' | |
| 主要有以下几个地方需要 filter: | |
| [1] 点云的位置超出了取值范围 | |
| [2] 近平面和远平面之外的点云 | |
| [3] adaptive density control 时需要 prune 的点云 | |
| ''' | |
| def filter(self, valid_mask): | |
| self._xyz = self._xyz[valid_mask] | |
| self._features_dc = self._features_dc[valid_mask] | |
| if self.sh_degree > 0: | |
| self._features_rest = self._features_rest[valid_mask] | |
| self._scaling = self._scaling[valid_mask] | |
| self._rotation = self._rotation[valid_mask] | |
| self._opacity = self._opacity[valid_mask] | |
| return self | |
| # xyz 的取值范围都限定在 [-1, 1] 之间 | |
| def crop(self, crop_bbx=[-1, 1, -1, 1, -1, 1]): | |
| x_min, x_max, y_min, y_max, z_min, z_max = crop_bbx | |
| xyz = self._xyz | |
| invalid_mask = ( | |
| (xyz[:, 0] < x_min) | |
| | (xyz[:, 0] > x_max) | |
| | (xyz[:, 1] < y_min) | |
| | (xyz[:, 1] > y_max) | |
| | (xyz[:, 2] < z_min) | |
| | (xyz[:, 2] > z_max) | |
| ) | |
| valid_mask = ~invalid_mask | |
| return self.filter(valid_mask) | |
| def prune(self, opacity_thres=0.05): | |
| opacity = self.get_opacity.squeeze(1) | |
| valid_mask = opacity > opacity_thres | |
| return self.filter(valid_mask) | |
| def prune_by_nearfar(self, cam_origins, nearfar_percent=(0.01, 0.99)): | |
| # cam_origins: [num_cams, 3] | |
| # nearfar_percent: [near, far] | |
| assert len(nearfar_percent) == 2 | |
| assert nearfar_percent[0] < nearfar_percent[1] | |
| assert nearfar_percent[0] >= 0 and nearfar_percent[1] <= 1 | |
| device = self._xyz.device | |
| # compute distance of all points to all cameras | |
| # [num_points, num_cams] | |
| ''' | |
| torch.cdist 用来计算 self._xyz 中的每个点到 cam_origins 中的每个点之间的距离, dists 的 shape 为 [num_points, num_cams] | |
| [None] 用来增加矩阵的维度 | |
| ''' | |
| dists = torch.cdist(self._xyz[None], cam_origins[None].to(device))[0] | |
| # [2, num_cams] | |
| ''' | |
| torch.quantile 用来计算分位数的 | |
| 第一行是近平面距离 - 0.01 分位数 | |
| 第二行是远平面距离 - 0.99 分位数 | |
| ''' | |
| dists_percentile = torch.quantile( | |
| dists, torch.tensor(nearfar_percent).to(device), dim=0 | |
| ) | |
| # prune all points that are outside the percentile range | |
| # [num_points, num_cams] | |
| # goal: prune points that are too close or too far from any camera | |
| reject_mask = (dists < dists_percentile[0:1, :]) | ( | |
| dists > dists_percentile[1:2, :] | |
| ) | |
| reject_mask = reject_mask.any(dim=1) | |
| valid_mask = ~reject_mask | |
| return self.filter(valid_mask) | |
| def apply_all_filters( | |
| self, | |
| opacity_thres=0.05, | |
| crop_bbx=[-1, 1, -1, 1, -1, 1], | |
| cam_origins=None, | |
| nearfar_percent=(0.005, 1.0), | |
| ): | |
| self.prune(opacity_thres) | |
| if crop_bbx is not None: | |
| self.crop(crop_bbx) # cropbox 是针对中心矩阵的 | |
| if cam_origins is not None: | |
| self.prune_by_nearfar(cam_origins, nearfar_percent) # nearfar 是针对相机原点的 | |
| return self | |
| # 将中间的 bbox 缩小 | |
| def shrink_bbx(self, drop_ratio=0.05): | |
| xyz = self._xyz | |
| xyz_min, xyz_max = torch.quantile( | |
| xyz, | |
| torch.tensor([drop_ratio, 1 - drop_ratio]).float().to(xyz.device), | |
| dim=0, | |
| ) # [2, N] | |
| xyz_min = xyz_min.detach().cpu().numpy() | |
| xyz_max = xyz_max.detach().cpu().numpy() | |
| crop_bbx = [ | |
| xyz_min[0], | |
| xyz_max[0], | |
| xyz_min[1], | |
| xyz_max[1], | |
| xyz_min[2], | |
| xyz_max[2], | |
| ] | |
| print(f"Shrinking bbx to {crop_bbx}") | |
| return self.crop(crop_bbx) | |
| def report_stats(self): | |
| print( | |
| f"xyz: {self._xyz.shape}, {self._xyz.min().item()}, {self._xyz.max().item()}" | |
| ) | |
| print( | |
| f"features_dc: {self._features_dc.shape}, {self._features_dc.min().item()}, {self._features_dc.max().item()}" | |
| ) | |
| if self.sh_degree > 0: | |
| print( | |
| f"features_rest: {self._features_rest.shape}, {self._features_rest.min().item()}, {self._features_rest.max().item()}" | |
| ) | |
| print( | |
| f"scaling: {self._scaling.shape}, {self._scaling.min().item()}, {self._scaling.max().item()}" | |
| ) | |
| print( | |
| f"rotation: {self._rotation.shape}, {self._rotation.min().item()}, {self._rotation.max().item()}" | |
| ) | |
| print( | |
| f"opacity: {self._opacity.shape}, {self._opacity.min().item()}, {self._opacity.max().item()}" | |
| ) | |
| print( | |
| f"after activation, xyz: {self.get_xyz.shape}, {self.get_xyz.min().item()}, {self.get_xyz.max().item()}" | |
| ) | |
| print( | |
| f"after activation, features: {self.get_features.shape}, {self.get_features.min().item()}, {self.get_features.max().item()}" | |
| ) | |
| print( | |
| f"after activation, scaling: {self.get_scaling.shape}, {self.get_scaling.min().item()}, {self.get_scaling.max().item()}" | |
| ) | |
| print( | |
| f"after activation, rotation: {self.get_rotation.shape}, {self.get_rotation.min().item()}, {self.get_rotation.max().item()}" | |
| ) | |
| print( | |
| f"after activation, opacity: {self.get_opacity.shape}, {self.get_opacity.min().item()}, {self.get_opacity.max().item()}" | |
| ) | |
| print( | |
| f"after activation, covariance: {self.get_covariance().shape}, {self.get_covariance().min().item()}, {self.get_covariance().max().item()}" | |
| ) | |
| ''' | |
| 两种装饰器的用法和区别： | |
| (1) @property, 把类的方法变成属性来调用, 例如: gs.get_xyz, 不需要加括号 | |
| (2) @staticmethod, 静态方法, 不需要实例化, 直接通过类名调用这个类内的方法 | |
| ''' | |
| def get_scaling(self): | |
| if self.scaling_modifier is not None: | |
| return self.scaling_activation(self._scaling) * self.scaling_modifier | |
| else: | |
| return self.scaling_activation(self._scaling) | |
| def get_rotation(self): | |
| return self.rotation_activation(self._rotation) | |
| def get_xyz(self): | |
| return self._xyz | |
| def get_features(self): | |
| if self.sh_degree > 0: | |
| features_dc = self._features_dc | |
| features_rest = self._features_rest | |
| return torch.cat((features_dc, features_rest), dim=1) | |
| else: | |
| return self._features_dc | |
| def get_opacity(self): | |
| return self.opacity_activation(self._opacity) | |
| def get_covariance(self, scaling_modifier=1): | |
| return self.covariance_activation( | |
| self.get_scaling, scaling_modifier, self._rotation | |
| ) | |
| # 设置好各个属性变量的数值类型 | |
| def construct_dtypes(self, use_fp16=False, enable_gs_viewer=True): | |
| if not use_fp16: | |
| l = [ | |
| ("x", "f4"), | |
| ("y", "f4"), | |
| ("z", "f4"), | |
| ("red", "u1"), | |
| ("green", "u1"), | |
| ("blue", "u1"), | |
| ] | |
| # All channels except the 3 DC | |
| for i in range(self._features_dc.shape[1] * self._features_dc.shape[2]): | |
| l.append((f"f_dc_{i}", "f4")) | |
| if enable_gs_viewer: | |
| assert self.sh_degree <= 3, "GS viewer only supports SH up to degree 3" | |
| sh_degree = 3 | |
| for i in range(((sh_degree + 1) ** 2 - 1) * 3): | |
| l.append((f"f_rest_{i}", "f4")) | |
| else: | |
| if self.sh_degree > 0: | |
| for i in range( | |
| self._features_rest.shape[1] * self._features_rest.shape[2] | |
| ): | |
| l.append((f"f_rest_{i}", "f4")) | |
| l.append(("opacity", "f4")) | |
| for i in range(self._scaling.shape[1]): | |
| l.append((f"scale_{i}", "f4")) | |
| for i in range(self._rotation.shape[1]): | |
| l.append((f"rot_{i}", "f4")) | |
| else: | |
| l = [ | |
| ("x", "f2"), | |
| ("y", "f2"), | |
| ("z", "f2"), | |
| ("red", "u1"), | |
| ("green", "u1"), | |
| ("blue", "u1"), | |
| ] | |
| # All channels except the 3 DC | |
| for i in range(self._features_dc.shape[1] * self._features_dc.shape[2]): | |
| l.append((f"f_dc_{i}", "f2")) | |
| if self.sh_degree > 0: | |
| for i in range( | |
| self._features_rest.shape[1] * self._features_rest.shape[2] | |
| ): | |
| l.append((f"f_rest_{i}", "f2")) | |
| l.append(("opacity", "f2")) | |
| for i in range(self._scaling.shape[1]): | |
| l.append((f"scale_{i}", "f2")) | |
| for i in range(self._rotation.shape[1]): | |
| l.append((f"rot_{i}", "f2")) | |
| return l | |
| # dump 点云文件 | |
| def save_ply( | |
| self, | |
| path, | |
| use_fp16=False, | |
| enable_gs_viewer=True, | |
| color_code=False, | |
| filter_mask=None, | |
| ): | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| xyz = self._xyz.detach().cpu().numpy() | |
| f_dc = ( | |
| self._features_dc.detach() | |
| .transpose(1, 2) | |
| .flatten(start_dim=1) | |
| .contiguous() | |
| .cpu() | |
| .numpy() | |
| ) | |
| if not color_code: | |
| rgb = (SH2RGB(f_dc) * 255.0).clip(0.0, 255.0).astype(np.uint8) | |
| else: | |
| # use an color map to color code the index of points | |
| index = np.linspace(0, 1, xyz.shape[0]) | |
| rgb = matplotlib.colormaps["viridis"](index)[..., :3] | |
| rgb = (rgb * 255.0).clip(0.0, 255.0).astype(np.uint8) | |
| opacities = self._opacity.detach().cpu().numpy() | |
| if self.scaling_modifier is not None: | |
| scale = self.inv_scaling_activation(self.get_scaling).detach().cpu().numpy() | |
| else: | |
| scale = self._scaling.detach().cpu().numpy() | |
| rotation = self._rotation.detach().cpu().numpy() | |
| dtype_full = self.construct_dtypes(use_fp16, enable_gs_viewer) | |
| elements = np.empty(xyz.shape[0], dtype=dtype_full) | |
| f_rest = None | |
| if self.sh_degree > 0: | |
| f_rest = ( | |
| self._features_rest.detach() | |
| .transpose(1, 2) | |
| .flatten(start_dim=1) | |
| .contiguous() | |
| .cpu() | |
| .numpy() | |
| ) # (3, (self.sh_degree + 1) ** 2 - 1) | |
| if enable_gs_viewer: | |
| sh_degree = 3 | |
| if f_rest is None: | |
| f_rest = np.zeros( | |
| (xyz.shape[0], 3 * ((sh_degree + 1) ** 2 - 1)), dtype=np.float32 | |
| ) | |
| elif f_rest.shape[1] < 3 * ((sh_degree + 1) ** 2 - 1): | |
| f_rest_pad = np.zeros( | |
| (xyz.shape[0], 3 * ((sh_degree + 1) ** 2 - 1)), dtype=np.float32 | |
| ) | |
| f_rest_pad[:, : f_rest.shape[1]] = f_rest | |
| f_rest = f_rest_pad | |
| if f_rest is not None: | |
| attributes = np.concatenate( | |
| (xyz, rgb, f_dc, f_rest, opacities, scale, rotation), axis=1 | |
| ) | |
| else: | |
| attributes = np.concatenate( | |
| (xyz, rgb, f_dc, opacities, scale, rotation), axis=1 | |
| ) | |
| if filter_mask is not None: | |
| attributes = attributes[filter_mask] | |
| elements = elements[filter_mask] | |
| elements[:] = list(map(tuple, attributes)) | |
| el = PlyElement.describe(elements, "vertex") | |
| PlyData([el]).write(path) | |
| # load 点云文件 | |
| def load_ply(self, path): | |
| 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, | |
| ) | |
| opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis] | |
| features_dc = np.zeros((xyz.shape[0], 3, 1)) | |
| features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"]) | |
| features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"]) | |
| features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"]) | |
| if self.sh_degree > 0: | |
| extra_f_names = [ | |
| p.name | |
| for p in plydata.elements[0].properties | |
| if p.name.startswith("f_rest_") | |
| ] | |
| extra_f_names = sorted(extra_f_names, key=lambda x: int(x.split("_")[-1])) | |
| assert len(extra_f_names) == 3 * (self.sh_degree + 1) ** 2 - 3 | |
| features_extra = np.zeros((xyz.shape[0], len(extra_f_names))) | |
| for idx, attr_name in enumerate(extra_f_names): | |
| features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name]) | |
| # Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC) | |
| features_extra = features_extra.reshape( | |
| (features_extra.shape[0], 3, (self.sh_degree + 1) ** 2 - 1) | |
| ) | |
| scale_names = [ | |
| p.name | |
| for p in plydata.elements[0].properties | |
| if p.name.startswith("scale_") | |
| ] | |
| scale_names = sorted(scale_names, key=lambda x: int(x.split("_")[-1])) | |
| 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") | |
| ] | |
| rot_names = sorted(rot_names, key=lambda x: int(x.split("_")[-1])) | |
| 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]) | |
| self._xyz = torch.from_numpy(xyz.astype(np.float32)) | |
| self._features_dc = ( | |
| torch.from_numpy(features_dc.astype(np.float32)) | |
| .transpose(1, 2) | |
| .contiguous() | |
| ) | |
| if self.sh_degree > 0: | |
| self._features_rest = ( | |
| torch.from_numpy(features_extra.astype(np.float32)) | |
| .transpose(1, 2) | |
| .contiguous() | |
| ) | |
| self._opacity = torch.from_numpy( | |
| np.copy(opacities).astype(np.float32) | |
| ).contiguous() | |
| self._scaling = torch.from_numpy(scales.astype(np.float32)).contiguous() | |
| self._rotation = torch.from_numpy(rots.astype(np.float32)).contiguous() | |
| def extract_fields(self, resolution=128, num_blocks=16, relax_ratio=1.5): | |
| from scipy.spatial import cKDTree | |
| if num_blocks < 1 or resolution < num_blocks or resolution % num_blocks: | |
| raise ValueError("Field resolution must be a positive multiple of num_blocks.") | |
| # resolution: resolution of field | |
| block_size = 2 / num_blocks | |
| split_size = resolution // num_blocks | |
| ## recenter | |
| xyzs = self.get_xyz | |
| opacities = self.get_opacity | |
| covs = self.get_covariance() | |
| # normalize to ~ [-1, 1] | |
| stds = self.get_scaling | |
| mn, mx = xyzs.amin(0), xyzs.amax(0) | |
| self.mesh_center = (mn + mx) / 2 | |
| self.mesh_scale = 1.8 / (mx - mn).amax().item() | |
| xyzs = (xyzs - self.mesh_center) * self.mesh_scale | |
| stds = stds * self.mesh_scale | |
| covs = self.covariance_activation(stds, 1, self._rotation) | |
| # tile | |
| device = opacities.device | |
| occ = torch.zeros([resolution] * 3, dtype=torch.float32, device=device) | |
| X = torch.linspace(-1, 1, resolution).split(split_size) | |
| Y = torch.linspace(-1, 1, resolution).split(split_size) | |
| Z = torch.linspace(-1, 1, resolution).split(split_size) | |
| # Find each tile's Gaussian centers on the CPU. Scanning the full cloud | |
| # and synchronizing mask.any() for every tile is especially slow for | |
| # sparse clouds and shared GPUs (64**3 tiles in the default mesh). | |
| centers = xyzs.detach().cpu().numpy() | |
| tree = cKDTree(centers) | |
| # loop blocks (assume max size of gaussian is small than relax_ratio * block_size !!!) | |
| for xi, xs in enumerate(X): | |
| for yi, ys in enumerate(Y): | |
| for zi, zs in enumerate(Z): | |
| lower = np.array([xs[0], ys[0], zs[0]], dtype=np.float32) - block_size * relax_ratio | |
| upper = np.array([xs[-1], ys[-1], zs[-1]], dtype=np.float32) + block_size * relax_ratio | |
| midpoint = (lower.astype(np.float64) + upper) / 2 | |
| radius = np.max((upper.astype(np.float64) - lower) / 2) | |
| candidates = tree.query_ball_point(midpoint, np.nextafter(radius, np.inf), p=np.inf) | |
| if not candidates: | |
| continue | |
| candidates = np.asarray(candidates, dtype=np.int64) | |
| in_tile = ((centers[candidates] > lower) & (centers[candidates] < upper)).all(-1) | |
| candidates = np.sort(candidates[in_tile]) | |
| if len(candidates) == 0: | |
| continue | |
| # Preserve the original accumulation order and strict tile | |
| # bounds so the density field stays numerically equivalent. | |
| mask = torch.as_tensor(candidates, device=device) | |
| xx, yy, zz = torch.meshgrid(xs, ys, zs) | |
| # sample points [M, 3] | |
| pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1).to(device) | |
| mask_xyzs = xyzs[mask] # [L, 3] | |
| mask_covs = covs[mask] # [L, 6] | |
| mask_opas = opacities[mask].view(1, -1) # [L, 1] --> [1, L] | |
| # query per point-gaussian pair. | |
| g_pts = pts.unsqueeze(1).repeat(1, mask_covs.shape[0], 1) - mask_xyzs.unsqueeze(0) # [M, L, 3] | |
| g_covs = mask_covs.unsqueeze(0).repeat(pts.shape[0], 1, 1) # [M, L, 6] | |
| # batch on gaussian to avoid OOM | |
| batch_g = 1024 | |
| val = 0 | |
| for start in range(0, g_covs.shape[1], batch_g): | |
| end = min(start + batch_g, g_covs.shape[1]) | |
| w = gaussian_3d_coeff(g_pts[:, start:end].reshape(-1, 3), g_covs[:, start:end].reshape(-1, 6)).reshape(pts.shape[0], -1) # [M, l] | |
| val += (mask_opas[:, start:end] * w).sum(-1) | |
| # kiui.lo(val, mask_opas, w) | |
| occ[xi * split_size: xi * split_size + len(xs), | |
| yi * split_size: yi * split_size + len(ys), | |
| zi * split_size: zi * split_size + len(zs)] = val.reshape(len(xs), len(ys), len(zs)) | |
| kiui.lo(occ, verbose=1) | |
| return occ | |
| def extract_mesh(self, density_thresh=0.005, resolution=256, decimate_target=1e5): | |
| occ = self.extract_fields(resolution, num_blocks=64).detach().cpu().numpy() | |
| import mcubes | |
| # breakpoint() | |
| vertices, triangles = mcubes.marching_cubes(occ, density_thresh) | |
| vertices = vertices / (resolution - 1.0) * 2 - 1 | |
| vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.015) | |
| if decimate_target > 0 and triangles.shape[0] > decimate_target: | |
| vertices, triangles = decimate_mesh(vertices, triangles, decimate_target) | |
| # Field queries use centered/scaled coordinates; export in the same | |
| # world coordinates as the original Gaussian cloud. | |
| vertices = vertices / self.mesh_scale + self.mesh_center.detach().cpu().numpy() | |
| # print( | |
| # f"[INFO] marching cubes result: {v.shape} ({v.min().item()}-{v.max().item()}), {f.shape}" | |
| # ) | |
| mesh = trimesh.Trimesh(vertices=vertices, faces=triangles) | |
| #mesh = Mesh(v=v, f=f, device='cuda') | |
| return mesh | |
| # 调用 Rasterizer 来进行 render | |
| def render_opencv_cam( | |
| pc: GaussianModel, | |
| height: int, | |
| width: int, | |
| C2W: torch.Tensor, | |
| fxfycxcy: torch.Tensor, | |
| bg_color=(1.0, 1.0, 1.0), | |
| scaling_modifier=1.0, | |
| ): | |
| """ | |
| Render the scene. | |
| Background tensor (bg_color) must be on GPU! | |
| """ | |
| # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means | |
| screenspace_points = torch.empty_like( | |
| pc.get_xyz, dtype=pc.get_xyz.dtype, requires_grad=True, device="cuda" | |
| ) | |
| # try: | |
| # screenspace_points.retain_grad() | |
| # except: | |
| # pass | |
| viewpoint_camera = Camera(C2W=C2W, fxfycxcy=fxfycxcy, h=height, w=width) | |
| bg_color = torch.tensor(list(bg_color), dtype=torch.float32, device=C2W.device) | |
| # Set up rasterization configuration | |
| raster_settings = GaussianRasterizationSettings( | |
| image_height=int(viewpoint_camera.h), | |
| image_width=int(viewpoint_camera.w), | |
| tanfovx=viewpoint_camera.tanfovX, | |
| tanfovy=viewpoint_camera.tanfovY, | |
| bg=bg_color, | |
| scale_modifier=scaling_modifier, | |
| viewmatrix=viewpoint_camera.world_view_transform, | |
| projmatrix=viewpoint_camera.full_proj_transform, | |
| sh_degree=pc.sh_degree, | |
| campos=viewpoint_camera.camera_center, | |
| prefiltered=False, | |
| debug=False, | |
| ) | |
| rasterizer = GaussianRasterizer(raster_settings=raster_settings) | |
| means3D = pc.get_xyz | |
| means2D = screenspace_points | |
| opacity = pc.get_opacity | |
| scales = pc.get_scaling | |
| rotations = pc.get_rotation | |
| shs = pc.get_features | |
| # Rasterize visible Gaussians to image, obtain their radii (on screen). | |
| rendered_image, radii = rasterizer( | |
| means3D=means3D, | |
| means2D=means2D, | |
| shs=shs, | |
| colors_precomp=None, | |
| opacities=opacity, | |
| scales=scales, | |
| rotations=rotations, | |
| cov3D_precomp=None, | |
| ) | |
| # 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, | |
| } | |
| class DeferredGaussianRender(torch.autograd.Function): | |
| def forward( | |
| ctx, | |
| xyz, | |
| features, | |
| scaling, | |
| rotation, | |
| opacity, | |
| height, | |
| width, | |
| C2W, | |
| fxfycxcy, | |
| scaling_modifier=None, | |
| use_gssplat=False, | |
| ): | |
| """ | |
| xyz: [b, n_gaussians, 3] | |
| features: [b, n_gaussians, (sh_degree+1)^2, 3] | |
| scaling: [b, n_gaussians, 3] | |
| rotation: [b, n_gaussians, 4] | |
| opacity: [b, n_gaussians, 1] | |
| height: int | |
| width: int | |
| C2W: [b, v, 4, 4] | |
| fxfycxcy: [b, v, 4] | |
| output: [b, v, 3, height, width] | |
| """ | |
| ctx.scaling_modifier = scaling_modifier | |
| # Infer sh_degree from features, features 里面存的是 Gaussian 的 attributes | |
| sh_degree = int(math.sqrt(features.shape[-2])) - 1 | |
| # Create a temp class to hold the data and for rendering | |
| gaussians_model = GaussianModel(sh_degree, scaling_modifier) | |
| with torch.no_grad(): | |
| b, v = C2W.size(0), C2W.size(1) | |
| renders = [] | |
| for i in range(b): | |
| # 设置高斯的属性 | |
| pc = gaussians_model.set_data( | |
| xyz[i], features[i], scaling[i], rotation[i], opacity[i] | |
| ) | |
| # 二次循环渲染出每个 view 下的图像 | |
| for j in range(v): | |
| renders.append( | |
| render_opencv_cam(pc, height, width, C2W[i, j], fxfycxcy[i, j])[ | |
| "render" | |
| ] | |
| ) | |
| renders = torch.stack(renders, dim=0) | |
| renders = renders.reshape(b, v, 3, height, width) | |
| renders = renders.requires_grad_() | |
| # Save_for_backward only supports tensors | |
| ctx.save_for_backward(xyz, features, scaling, rotation, opacity, C2W, fxfycxcy) | |
| ctx.rendering_size = (height, width) | |
| ctx.sh_degree = sh_degree | |
| ctx.use_gssplat = use_gssplat | |
| # Release the temp class; do not save it. | |
| del gaussians_model | |
| return renders | |
| def backward(ctx, grad_output): | |
| # Restore params | |
| xyz, features, scaling, rotation, opacity, C2W, fxfycxcy = ctx.saved_tensors | |
| height, width = ctx.rendering_size | |
| sh_degree = ctx.sh_degree | |
| use_gssplat = ctx.use_gssplat | |
| # **The order of this dict should not be changed** | |
| input_dict = OrderedDict( | |
| [ | |
| ("xyz", xyz), | |
| ("features", features), | |
| ("scaling", scaling), | |
| ("rotation", rotation), | |
| ("opacity", opacity), | |
| ] | |
| ) | |
| input_dict = {k: v.detach().requires_grad_() for k, v in input_dict.items()} | |
| # Create a temp class to hold the data and for rendering | |
| gaussians_model = GaussianModel(sh_degree, ctx.scaling_modifier) | |
| with torch.enable_grad(): | |
| b, v = C2W.size(0), C2W.size(1) | |
| for i in range(b): | |
| for j in range(v): | |
| # The backward will remove the diff graph, thus each time we need a copy | |
| pc = gaussians_model.set_data( | |
| **{k: v[i] for k, v in input_dict.items()} | |
| ) | |
| # Forward | |
| render = render_opencv_cam( | |
| pc, height, width, C2W[i, j], fxfycxcy[i, j] | |
| )["render"] | |
| # Backward, suppose that only values in input_dict will get gradients. | |
| render.backward(grad_output[i, j]) | |
| del gaussians_model | |
| return *[var.grad for var in input_dict.values()], None, None, None, None, None, None | |
| # Function for the class | |
| deferred_gaussian_render = DeferredGaussianRender.apply | |
| class DeferredGaussianRender_scene(torch.autograd.Function): | |
| def forward( | |
| ctx, | |
| xyz, | |
| features, | |
| scaling, | |
| rotation, | |
| opacity, | |
| height, | |
| width, | |
| C2W, | |
| fxfycxcy, | |
| scaling_modifier=None, | |
| ): | |
| """ | |
| xyz: [b, n_gaussians, 3] | |
| features: [b, n_gaussians, (sh_degree+1)^2, 3] | |
| scaling: [b, n_gaussians, 3] | |
| rotation: [b, n_gaussians, 4] | |
| opacity: [b, n_gaussians, 1] | |
| height: int | |
| width: int | |
| C2W: [b, v, 4, 4] | |
| fxfycxcy: [b, v, 4] | |
| output: [b, v, 3, height, width] | |
| """ | |
| ctx.scaling_modifier = scaling_modifier | |
| # Infer sh_degree from features, features 里面存的是 Gaussian 的 attributes | |
| sh_degree = int(math.sqrt(features.shape[-2])) - 1 | |
| # Create a temp class to hold the data and for rendering | |
| gaussians_model = GaussianModel(sh_degree, scaling_modifier) | |
| with torch.no_grad(): | |
| b, v = C2W.size(0), C2W.size(1) | |
| renders = torch.zeros( | |
| b, v, 3, height, width, dtype=torch.float32, device=xyz.device | |
| ) | |
| depth = torch.zeros( | |
| b, v, 1, height, width, dtype=torch.float32, device=xyz.device | |
| ) | |
| alpha = torch.zeros( | |
| b, v, 1, height, width, dtype=torch.float32, device=xyz.device | |
| ) | |
| for i in range(b): | |
| # 设置高斯的属性 | |
| pc = gaussians_model.set_data( | |
| xyz[i], features[i], scaling[i], rotation[i], opacity[i] | |
| ) | |
| # 二次循环渲染出每个 view 下的图像 | |
| for j in range(v): | |
| buffers = render_opencv_cam( | |
| pc, height, width, C2W[i, j], fxfycxcy[i, j] | |
| ) | |
| renders[i, j] = buffers["render"] | |
| if "depth" in buffers and buffers["depth"] is not None: | |
| depth[i, j] = buffers["depth"] | |
| if "alpha" in buffers and buffers["alpha"] is not None: | |
| alpha[i, j] = buffers["alpha"] | |
| renders = renders.requires_grad_() | |
| # depth = depth.requires_grad_() | |
| # alpha = alpha.requires_grad_() | |
| # Save_for_backward only supports tensors | |
| ctx.save_for_backward(xyz, features, scaling, rotation, opacity, C2W, fxfycxcy) | |
| ctx.rendering_size = (height, width) | |
| ctx.sh_degree = sh_degree | |
| # Release the temp class; do not save it. | |
| del gaussians_model | |
| return renders, depth, alpha | |
| def backward(ctx, grad_render, grad_depth, grad_alpha): | |
| # Restore params | |
| xyz, features, scaling, rotation, opacity, C2W, fxfycxcy = ctx.saved_tensors | |
| height, width = ctx.rendering_size | |
| sh_degree = ctx.sh_degree | |
| # **The order of this dict should not be changed** | |
| input_dict = OrderedDict( | |
| [ | |
| ("xyz", xyz), | |
| ("features", features), | |
| ("scaling", scaling), | |
| ("rotation", rotation), | |
| ("opacity", opacity), | |
| ] | |
| ) | |
| input_dict = {k: v.detach().requires_grad_() for k, v in input_dict.items()} | |
| # Create a temp class to hold the data and for rendering | |
| gaussians_model = GaussianModel(sh_degree, ctx.scaling_modifier) | |
| with torch.enable_grad(): | |
| b, v = C2W.size(0), C2W.size(1) | |
| for i in range(b): | |
| for j in range(v): | |
| # The backward will remove the diff graph, thus each time we need a copy | |
| pc = gaussians_model.set_data( | |
| **{k: v[i] for k, v in input_dict.items()} | |
| ) | |
| # Forward | |
| buffers = render_opencv_cam( | |
| pc, height, width, C2W[i, j], fxfycxcy[i, j] | |
| ) | |
| render = buffers["render"] | |
| # if "depth" in buffers and buffers["depth"] is not None: | |
| # depth = buffers["depth"] | |
| # if "alpha" in buffers and buffers["alpha"] is not None: | |
| # alpha = buffers["alpha"] | |
| # Backward, suppose that only values in input_dict will get gradients. | |
| render.backward(grad_render[i, j]) | |
| # depth.backward(grad_depth[i, j]) | |
| # alpha.backward(grad_alpha[i, j]) | |
| del gaussians_model | |
| return *[var.grad for var in input_dict.values()], None, None, None, None, None | |
| # Function for the class | |
| deferred_gaussian_render_scene = DeferredGaussianRender_scene.apply | |
| def render_turntable(pc: GaussianModel, rendering_resolution=384, num_views=8): # 渲染一圈图像 - 8 个 view | |
| w, h, v, fxfycxcy, c2w = get_turntable_cameras( | |
| h=rendering_resolution, w=rendering_resolution, num_views=num_views | |
| ) # 这里的 v - viewpoint 被重置了，从 get_turntable_cameras 中抽取 | |
| device = pc._xyz.device | |
| fxfycxcy = torch.from_numpy(fxfycxcy).float().to(device) # [v, 4] | |
| c2w = torch.from_numpy(c2w).float().to(device) # [v, 4, 4] | |
| renderings = torch.zeros(v, 3, h, w, dtype=torch.float32, device=device) | |
| for j in range(v): | |
| renderings[j] = render_opencv_cam(pc, h, w, c2w[j], fxfycxcy[j])["render"] | |
| torch.cuda.empty_cache() # free up memory on GPU | |
| renderings = renderings.detach().cpu().numpy() | |
| renderings = (renderings * 255).clip(0, 255).astype(np.uint8) | |
| renderings = rearrange(renderings, "v c h w -> h (v w) c") | |
| return renderings | |
| class GaussianRenderer(nn.Module): | |
| def __init__(self, sh_degree, scaling_modifier=None): | |
| super().__init__() | |
| self.sh_degree = sh_degree | |
| self.scaling_modifier = scaling_modifier | |
| def forward( | |
| self, | |
| xyz, | |
| features, | |
| scaling, | |
| rotation, | |
| opacity, | |
| height, | |
| width, | |
| C2W, | |
| fxfycxcy, | |
| deferred=True, | |
| ): | |
| """ | |
| xyz: [b, n_gaussians, 3] | |
| features: [b, n_gaussians, (sh_degree+1)^2, 3] | |
| scaling: [b, n_gaussians, 3] | |
| rotation: [b, n_gaussians, 4] | |
| opacity: [b, n_gaussians, 1] | |
| height: int | |
| width: int | |
| C2W: [b, v, 4, 4] | |
| fxfycxcy: [b, v, 4] | |
| output: [b, v, 3, height, width] | |
| """ | |
| if deferred: | |
| renderings = deferred_gaussian_render( | |
| xyz, | |
| features, | |
| scaling, | |
| rotation, | |
| opacity, | |
| height, | |
| width, | |
| C2W, | |
| fxfycxcy, | |
| self.scaling_modifier, | |
| ) | |
| b, v = C2W.size(0), C2W.size(1) | |
| gaussians_models = [] | |
| for i in range(b): | |
| pc = GaussianModel(self.sh_degree, self.scaling_modifier) | |
| pc = pc.set_data( | |
| xyz[i], features[i], scaling[i], rotation[i], opacity[i] | |
| ) | |
| gaussians_models.append(pc) | |
| else: | |
| b, v = C2W.size(0), C2W.size(1) | |
| renderings = torch.zeros( | |
| b, v, 3, height, width, dtype=torch.float32, device=xyz.device | |
| ) | |
| gaussians_models = [] | |
| for i in range(b): | |
| pc = GaussianModel(self.sh_degree, self.scaling_modifier) | |
| pc = pc.set_data( | |
| xyz[i], features[i], scaling[i], rotation[i], opacity[i] | |
| ) | |
| gaussians_models.append(pc) | |
| for j in range(v): | |
| renderings[i, j] = render_opencv_cam( | |
| pc, height, width, C2W[i, j], fxfycxcy[i, j] | |
| )["render"] | |
| return renderings, gaussians_models | |
| def render_generic(pc: GaussianModel, c2ws, fxfycxcy, h=512, w=512): | |
| # cw2s: [v, 4, 4] | |
| # fxfycxcy: [v, 4] | |
| v = c2ws.shape[0] | |
| device = pc._xyz.device | |
| fxfycxcy = fxfycxcy.float().to(device) # [v, 4] | |
| renderings = torch.zeros(v, 3, h, w, dtype=torch.float32, device=device) | |
| for j in range(v): | |
| renderings[j] = render_opencv_cam(pc, h, w, c2ws[j], fxfycxcy[j])["render"] | |
| torch.cuda.empty_cache() # free up memory on GPU | |
| renderings = renderings.detach().cpu().numpy() | |
| renderings = (renderings * 255).clip(0, 255).astype(np.uint8) | |
| renderings = rearrange(renderings, "v c h w -> v h w c") | |
| return renderings | |
| if __name__ == "__main__": | |
| import json | |
| from PIL import Image | |
| from tqdm import tqdm | |
| out_dir = "/mnt/localssd/debug-3dgs" | |
| os.makedirs(out_dir, exist_ok=True) | |
| os.system( | |
| f"wget https://phidias.s3.us-west-2.amazonaws.com/kaiz/neural-capture/eval-3dgs-lowres/AWS_test_set/results/1.fashion_boots_rubber_boots__short__Feb_21__2023_at_5_19_25_PM_yf/point_cloud/iteration_30000_fg/point_cloud.ply -O {out_dir}/point_cloud.ply" | |
| ) | |
| os.system( | |
| f"wget https://neural-capture.s3.us-west-2.amazonaws.com/data/AWS_test_set/preprocessed/1.fashion_boots_rubber_boots__short__Feb_21__2023_at_5_19_25_PM_yf/opencv_cameras_traj_norm.json -O {out_dir}/opencv_cameras_traj_norm.json" | |
| ) | |
| device = "cuda:0" | |
| pc = GaussianModel(sh_degree=3) | |
| pc.load_ply(f"{out_dir}/point_cloud.ply") | |
| pc = pc.to(device) | |
| # pc.save_ply(f"{out_dir}/point_cloud_shrink.ply") | |
| # pc.load_ply(f"{out_dir}/point_cloud_shrink.ply") | |
| # pc = pc.to(device) | |
| # pc.prune(opacity_thres=0.05) | |
| # pc.save_ply(f"{out_dir}/point_cloud_shrink_prune.ply") | |
| # pc = pc.to(device) | |
| # pc.shrink_bbx(drop_ratio=0.01) | |
| # pc.save_ply(f"{out_dir}/point_cloud_shrink_prune.ply") | |
| # pc = pc.to(device) | |
| pc.report_stats() | |
| with open(f"{out_dir}/opencv_cameras_traj_norm.json", "r") as f: | |
| cam_traj = json.load(f) | |
| for i, cam in tqdm(enumerate(cam_traj["frames"]), desc="Rendering progress"): | |
| w2c = np.array(cam["w2c"]) | |
| c2w = np.linalg.inv(w2c) | |
| c2w = torch.from_numpy(c2w.astype(np.float32)).to(device) | |
| fx = cam["fx"] | |
| fy = cam["fy"] | |
| cx = cam["cx"] | |
| cy = cam["cy"] | |
| cx = cx - 5 | |
| cy = cy + 4 | |
| fxfycxcy = torch.tensor([fx, fy, cx, cy], dtype=torch.float32, device=device) | |
| h = cam["h"] | |
| w = cam["w"] | |
| im = render_opencv_cam(pc, h, w, c2w, fxfycxcy, bg_color=[0.0, 0.0, 0.0])[ | |
| "render" | |
| ] | |
| im = im.detach().cpu().numpy().transpose(1, 2, 0) | |
| im = (im * 255).astype(np.uint8) | |
| Image.fromarray(im).save(f"{out_dir}/render_{i:08d}.png") | |
| create_video(out_dir, f"{out_dir}/render.mp4", framerate=30) | |
| print(f"Saved {out_dir}/render.mp4") | |