| 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]) |
| |
| 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 |
|
|
