Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import torch | |
| from plyfile import PlyData, PlyElement | |
| from typing import NamedTuple | |
| import smplx | |
| import tqdm | |
| import cv2 as cv | |
| import os | |
| from scipy.spatial.transform import Rotation as R | |
| class GaussianAttributes(NamedTuple): | |
| xyz: np.ndarray | |
| opacities: np.ndarray | |
| features_dc: np.ndarray | |
| features_extra: np.ndarray | |
| scales: np.ndarray | |
| rot: np.ndarray | |
| def load_gaussians_from_ply(path): | |
| max_sh_degree = 3 | |
| 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"]) | |
| 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 * (max_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, (max_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]) | |
| return GaussianAttributes(xyz, opacities, features_dc, features_extra, scales, rots) | |
| def construct_list_of_attributes(_features_dc, _features_rest, _scaling, _rotation): | |
| l = ['x', 'y', 'z', 'nx', 'ny', 'nz'] | |
| # All channels except the 3 DC | |
| 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(_scaling.shape[1]): | |
| l.append('scale_{}'.format(i)) | |
| for i in range(_rotation.shape[1]): | |
| l.append('rot_{}'.format(i)) | |
| return l | |
| def select_gaussians(gaussian_attrs, select_mask_or_idx): | |
| return GaussianAttributes( | |
| xyz=gaussian_attrs.xyz[select_mask_or_idx], | |
| opacities=gaussian_attrs.opacities[select_mask_or_idx], | |
| features_dc=gaussian_attrs.features_dc[select_mask_or_idx], | |
| features_extra=gaussian_attrs.features_extra[select_mask_or_idx], | |
| scales=gaussian_attrs.scales[select_mask_or_idx], | |
| rot=gaussian_attrs.rot[select_mask_or_idx] | |
| ) | |
| def combine_gaussians(gaussian_attrs_list): | |
| return GaussianAttributes( | |
| xyz=np.concatenate([gau.xyz for gau in gaussian_attrs_list], axis=0), | |
| opacities=np.concatenate([gau.opacities for gau in gaussian_attrs_list], axis=0), | |
| features_dc=np.concatenate([gau.features_dc for gau in gaussian_attrs_list], axis=0), | |
| features_extra=np.concatenate([gau.features_extra for gau in gaussian_attrs_list], axis=0), | |
| scales=np.concatenate([gau.scales for gau in gaussian_attrs_list], axis=0), | |
| rot=np.concatenate([gau.rot for gau in gaussian_attrs_list], axis=0), | |
| ) | |
| def apply_transformation_to_gaussians(gaussian_attrs, spatial_transformation, color_transformation=None): | |
| xyzs = np.copy(gaussian_attrs.xyz) | |
| xyzs = np.matmul(xyzs, spatial_transformation[:3, :3].transpose()) + spatial_transformation[:3, 3].reshape([1, 3]) | |
| gaussian_rotmats = R.from_quat(gaussian_attrs.rot[:, (1, 2, 3, 0)]).as_matrix() | |
| new_rots = [] | |
| for rotmat in gaussian_rotmats: | |
| rotmat = np.matmul(spatial_transformation[:3, :3], rotmat) | |
| rotq = R.from_matrix(rotmat).as_quat() | |
| rotq = np.array([rotq[3], rotq[0], rotq[1], rotq[2]]) | |
| new_rots.append(rotq) | |
| new_rots = np.stack(new_rots, axis=0) | |
| if color_transformation is not None: | |
| if color_transformation.shape[0] == 3 and color_transformation.shape[1] == 3: | |
| new_clrs = np.matmul(gaussian_attrs.features_dc[:, :, 0], color_transformation)[:, :, np.newaxis] | |
| elif color_transformation.shape[0] == 4 and color_transformation.shape[1] == 4: | |
| clrs = gaussian_attrs.features_dc[:, :, 0] | |
| clrs = np.concatenate([clrs, np.ones_like(clrs[:, :1])], axis=1) | |
| new_clrs = np.matmul(clrs, color_transformation) | |
| new_clrs = new_clrs[:, :3, np.newaxis] | |
| else: | |
| new_clrs = gaussian_attrs.features_dc | |
| return GaussianAttributes( | |
| xyz=xyzs, | |
| opacities=gaussian_attrs.opacities, | |
| features_dc=new_clrs, | |
| features_extra=gaussian_attrs.features_extra, | |
| scales=gaussian_attrs.scales, | |
| rot=new_rots, | |
| ) | |
| def update_gaussian_attributes( | |
| orig_gaussian, | |
| new_xyz=None, new_rgb=None, new_rot=None, new_opacity=None, new_scale=None): | |
| return GaussianAttributes( | |
| xyz=orig_gaussian.xyz if new_xyz is None else new_xyz, | |
| opacities=orig_gaussian.opacities if new_opacity is None else new_opacity, | |
| features_dc=orig_gaussian.features_dc if new_rgb is None else new_rgb, | |
| features_extra=orig_gaussian.features_extra, | |
| scales=orig_gaussian.scales if new_scale is None else new_scale, | |
| rot=orig_gaussian.rot if new_rot is None else new_rot, | |
| ) | |
| def save_gaussians_as_ply(path, gaussian_attrs: GaussianAttributes): | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| xyz = gaussian_attrs.xyz | |
| normals = np.zeros_like(xyz) | |
| features_dc = gaussian_attrs.features_dc | |
| features_rest = gaussian_attrs.features_extra | |
| opacities = gaussian_attrs.opacities | |
| scale = gaussian_attrs.scales | |
| rotation = gaussian_attrs.rot | |
| dtype_full = [(attribute, 'f4') for attribute in construct_list_of_attributes(features_dc, features_rest, scale, rotation)] | |
| elements = np.empty(xyz.shape[0], dtype = dtype_full) | |
| attributes = np.concatenate((xyz, normals, features_dc.reshape(features_dc.shape[0], -1), | |
| features_rest.reshape(features_rest.shape[0], -1), | |
| opacities, scale, rotation), axis=1) | |
| elements[:] = list(map(tuple, attributes)) | |
| el = PlyElement.describe(elements, 'vertex') | |
| PlyData([el]).write(path) | |
| return | |