Spaces:
Sleeping
Sleeping
| 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) | |