Spaces:
Running
Running
File size: 7,261 Bytes
7648567 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 |
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
|