Spaces:
Sleeping
Sleeping
File size: 3,468 Bytes
0332bda |
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 |
import numpy as np
import os
from plyfile import PlyElement, PlyData
import open3d as o3d
def get_f_dc(colors):
return RGB2SH(colors)[:, :, np.newaxis]
def get_f_rest(points, max_sh_degree=3):
f_rest_shape = (points.shape[0], (max_sh_degree + 1) ** 2 - 1, 3)
return np.zeros(f_rest_shape)
def get_opacity(points):
return inverse_sigmoid(0.5 * np.ones((points.shape[0], 1)))
def get_scales(points):
scales = np.ones((points.shape[0], 3)) * 0.0015
scales[:, 2] = 1e-6
return np.log(scales)
def get_rotation(normals):
if normals is not None and np.any(normals):
return normal2rotation(normals)
else:
return np.zeros((normals.shape[0], 4))
def RGB2SH(rgb):
return (rgb - 0.5) / 0.28209479177387814
def inverse_sigmoid(x):
return np.log(x / (1 - x))
def normal2rotation(n):
n = n / np.linalg.norm(n, axis=1, keepdims=True)
w0 = np.tile([[1, 0, 0]], (n.shape[0], 1))
R0 = w0 - np.sum(w0 * n, axis=1, keepdims=True) * n
R0 *= np.sign(R0[:, :1])
R0 /= np.linalg.norm(R0, axis=1, keepdims=True)
R1 = np.cross(n, R0)
R1 *= np.sign(R1[:, 1:2]) * np.sign(n[:, 2:])
R = np.stack([R0, R1, n], axis=-1)
q = rotmat2quaternion(R)
return q
def rotmat2quaternion(R, normalize=False):
tr = R[:, 0, 0] + R[:, 1, 1] + R[:, 2, 2] + 1e-6
r = np.sqrt(1 + tr) / 2
q = np.stack([
r,
(R[:, 2, 1] - R[:, 1, 2]) / (4 * r),
(R[:, 0, 2] - R[:, 2, 0]) / (4 * r),
(R[:, 1, 0] - R[:, 0, 1]) / (4 * r)
], axis=-1)
if normalize:
q /= np.linalg.norm(q, axis=-1, keepdims=True)
return q
def point2gs(path, pcd, scale=None, max_sh_degree=1):
# Ensure the directory exists
os.makedirs(os.path.dirname(path), exist_ok=True)
# Get point cloud data
xyz = np.asarray(pcd.points)
normals = np.asarray(pcd.normals) if pcd.has_normals() else np.zeros_like(xyz)
colors = np.asarray(pcd.colors) if pcd.has_colors() else np.ones_like(xyz)
# Generate additional attributes
f_dc = get_f_dc(colors).reshape(xyz.shape[0], -1)
f_rest = get_f_rest(xyz, max_sh_degree).reshape(xyz.shape[0], -1)
opacities = get_opacity(xyz)
if scale is not None:
scale = np.log(scale)
else:
scale = get_scales(xyz)
rotation = get_rotation(normals)
# Construct list of attributes
attribute_names = ['x', 'y', 'z', 'nx', 'ny', 'nz']
attribute_names.extend([f'f_dc_{i}' for i in range(f_dc.shape[-1])])
attribute_names.extend([f'f_rest_{i}' for i in range(f_rest.shape[-1])])
attribute_names.append('opacity')
attribute_names.extend([f'scale_{i}' for i in range(scale.shape[1])])
attribute_names.extend([f'rot_{i}' for i in range(rotation.shape[1])])
# Create dtype for numpy structured array
dtype_full = [(attribute, 'f4') for attribute in attribute_names]
# Combine all attributes
attributes = np.concatenate((
xyz, normals,
f_dc,
f_rest,
opacities, scale, rotation
), axis=1)
# Ensure attributes match the dtype
assert attributes.shape[1] == len(dtype_full), f"Mismatch in attribute count. Expected {len(dtype_full)}, got {attributes.shape[1]}"
# Create structured array
elements = np.empty(xyz.shape[0], dtype=dtype_full)
elements[:] = list(map(tuple, attributes))
# Create PlyElement and save
el = PlyElement.describe(elements, 'vertex')
PlyData([el]).write(path)
|