Spaces:
Running
on
Zero
Running
on
Zero
| import os | |
| from plyfile import PlyData, PlyElement | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import numpy as np | |
| import math | |
| import copy | |
| from lam.models.rendering.utils.typing import * | |
| from lam.models.rendering.utils.utils import trunc_exp, MLP | |
| from einops import rearrange, repeat | |
| inverse_sigmoid = lambda x: np.log(x / (1 - x)) | |
| class GaussianModel: | |
| def __init__(self, xyz=None, opacity=None, rotation=None, scaling=None, shs=None, offset=None, ply_path=None, sh2rgb=False, albedo=None, lights=None) -> None: | |
| self.xyz: Tensor = xyz | |
| self.opacity: Tensor = opacity | |
| self.rotation: Tensor = rotation | |
| self.scaling: Tensor = scaling | |
| self.shs: Tensor = shs | |
| self.albedo: Tensor = albedo | |
| self.offset: Tensor = offset | |
| self.lights: Tensor = lights | |
| if ply_path is not None: | |
| self.load_ply(ply_path, sh2rgb=sh2rgb) | |
| def update_lights(self, lights): | |
| self.lights = lights | |
| def update_albedo(self, albedo): | |
| self.albedo = albedo | |
| def update_shs(self, shs): | |
| self.shs = shs | |
| def to_cuda(self): | |
| self.xyz = self.xyz.cuda() | |
| self.opacity = self.opacity.cuda() | |
| self.rotation = self.rotation.cuda() | |
| self.scaling = self.scaling.cuda() | |
| self.shs = self.shs.cuda() | |
| self.offset = self.offset.cuda() | |
| self.albedo = self.albedo.cuda() | |
| def construct_list_of_attributes(self): | |
| l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] | |
| if len(self.shs.shape) == 2: | |
| features_dc = self.shs[:, :3].unsqueeze(1) | |
| features_rest = self.shs[:, 3:].unsqueeze(1) | |
| else: | |
| features_dc = self.shs[:, :1] | |
| features_rest = self.shs[:, 1:] | |
| for i in range(features_dc.shape[1]*features_dc.shape[2]): | |
| l.append('f_dc_{}'.format(i)) | |
| for i in range(features_rest.shape[1]*features_rest.shape[2]): | |
| l.append('f_rest_{}'.format(i)) | |
| l.append('opacity') | |
| for i in range(self.scaling.shape[1]): | |
| l.append('scale_{}'.format(i)) | |
| for i in range(self.rotation.shape[1]): | |
| l.append('rot_{}'.format(i)) | |
| return l | |
| def save_ply(self, path, rgb2sh=False, offset2xyz=False, albedo2rgb=False): | |
| if offset2xyz: | |
| xyz = self.offset.detach().cpu().float().numpy() | |
| else: | |
| xyz = self.xyz.detach().cpu().float().numpy() | |
| if albedo2rgb: | |
| self.shs = self.albedo | |
| normals = np.zeros_like(xyz) | |
| if len(self.shs.shape) == 2: | |
| features_dc = self.shs[:, :3].unsqueeze(1).float() | |
| features_rest = self.shs[:, 3:].unsqueeze(1).float() | |
| else: | |
| features_dc = self.shs[:, :1].float() | |
| features_rest = self.shs[:, 1:].float() | |
| f_dc = features_dc.detach().flatten(start_dim=1).contiguous().cpu().numpy() | |
| f_rest = features_rest.detach().flatten(start_dim=1).contiguous().cpu().numpy() | |
| if rgb2sh: | |
| from lam.models.rendering.utils.sh_utils import RGB2SH | |
| f_dc = RGB2SH(f_dc) | |
| opacities = inverse_sigmoid(torch.clamp(self.opacity, 1e-3, 1 - 1e-3).detach().cpu().float().numpy()) | |
| scale = np.log(self.scaling.detach().cpu().float().numpy()) | |
| rotation = self.rotation.detach().cpu().float().numpy() | |
| dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] | |
| elements = np.empty(xyz.shape[0], dtype=dtype_full) | |
| attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1) | |
| elements[:] = list(map(tuple, attributes)) | |
| el = PlyElement.describe(elements, 'vertex') | |
| PlyData([el]).write(path) | |
| def save_ply_nodeact(self, path, rgb2sh=False, albedo2rgb=False): | |
| if albedo2rgb: | |
| self.shs = self.albedo | |
| xyz = self.xyz.detach().cpu().float().numpy() | |
| normals = np.zeros_like(xyz) | |
| if len(self.shs.shape) == 2: | |
| features_dc = self.shs[:, :3].unsqueeze(1).float() | |
| features_rest = self.shs[:, 3:].unsqueeze(1).float() | |
| else: | |
| features_dc = self.shs[:, :1].float() | |
| features_rest = self.shs[:, 1:].float() | |
| f_dc = features_dc.detach().flatten(start_dim=1).contiguous().cpu().numpy() | |
| f_rest = features_rest.detach().flatten(start_dim=1).contiguous().cpu().numpy() | |
| if rgb2sh: | |
| from lam.models.rendering.utils.sh_utils import RGB2SH | |
| f_dc = RGB2SH(f_dc) | |
| opacities = self.opacity.detach().cpu().float().numpy() | |
| scale = self.scaling.detach().cpu().float().numpy() | |
| rotation = self.rotation.detach().cpu().float().numpy() | |
| dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()] | |
| elements = np.empty(xyz.shape[0], dtype=dtype_full) | |
| attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1) | |
| elements[:] = list(map(tuple, attributes)) | |
| el = PlyElement.describe(elements, 'vertex') | |
| PlyData([el]).write(path) | |
| def load_ply(self, path, sh2rgb=False): | |
| 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"]) | |
| 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])) | |
| 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 = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cpu").requires_grad_(False)) | |
| self.features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cpu").transpose(1, 2).contiguous().requires_grad_(False)) | |
| if sh2rgb: | |
| from lam.models.rendering.utils.sh_utils import SH2RGB | |
| self.features_dc = SH2RGB(self.features_dc) | |
| self.features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cpu").transpose(1, 2).contiguous().requires_grad_(False)) | |
| self.shs = torch.cat([self.features_dc, self.features_rest], dim=1) | |
| self.opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cpu").requires_grad_(False)) | |
| self.scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cpu").requires_grad_(False)) | |
| self.rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cpu").requires_grad_(False)) | |
| self.offset = nn.Parameter(torch.zeros_like(self.xyz).requires_grad_(False)) | |
| self.albedo = nn.Parameter(torch.zeros_like(self.shs).requires_grad_(False)) | |
| self.lights = nn.Parameter(torch.zeros_like(self.shs).requires_grad_(False)) | |
| if sh2rgb: | |
| self.opacity = nn.functional.sigmoid(self.opacity) | |
| self.scaling = trunc_exp(self.scaling) | |
| self.active_sh_degree = self.sh_degree | |