File size: 5,205 Bytes
8c48cce | 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 | import math
import numpy as np
from plyfile import PlyData, PlyElement
import math
import os
import warp as wp
def load_ply(filepath):
"""
Load a Gaussian splat PLY file.
Returns dict with: positions, scales, rotations, opacities, shs
"""
plydata = PlyData.read(filepath)
vertex = plydata['vertex']
num_points = len(vertex)
# Load positions
positions = np.stack([
vertex['x'], vertex['y'], vertex['z']
], axis=-1).astype(np.float32)
# Load scales (stored in log space)
scales = np.stack([
np.exp(vertex['scale_0']),
np.exp(vertex['scale_1']),
np.exp(vertex['scale_2'])
], axis=-1).astype(np.float32)
# Load opacities
opacities = vertex['opacity'].astype(np.float32).reshape(-1, 1)
# Load rotations (quaternion)
rotations = np.stack([
vertex['rot_0'], vertex['rot_1'], vertex['rot_2'], vertex['rot_3']
], axis=-1).astype(np.float32)
# Load SH coefficients
# DC term
sh_dc = np.stack([
vertex['f_dc_0'], vertex['f_dc_1'], vertex['f_dc_2']
], axis=-1).astype(np.float32)
# Rest of SH coefficients
sh_rest = []
for i in range(45):
sh_rest.append(vertex[f'f_rest_{i}'])
sh_rest = np.stack(sh_rest, axis=-1).astype(np.float32) # (N, 45)
sh_rest = sh_rest.reshape(num_points, 15, 3) # (N, 15, 3)
# Combine into (N*16, 3) format expected by renderer
shs = np.zeros((num_points * 16, 3), dtype=np.float32)
for i in range(num_points):
shs[i * 16] = sh_dc[i]
for j in range(15):
shs[i * 16 + j + 1] = sh_rest[i, j]
return {
'positions': positions,
'scales': scales,
'rotations': rotations,
'opacities': opacities,
'shs': shs,
'num_points': num_points
}
# Function to save point cloud to PLY file
def save_ply(params, filepath, num_points, colors=None):
# Get numpy arrays
positions = params['positions'].numpy()
scales = params['scales'].numpy()
rotations = params['rotations'].numpy()
opacities = params['opacities'].numpy()
shs = params['shs'].numpy()
# Handle colors - either provided or computed from SH coefficients
if colors is not None:
# Use provided colors
if hasattr(colors, 'numpy'):
colors_np = colors.numpy()
else:
colors_np = colors
else:
# Compute colors from SH coefficients (DC term only for simplicity)
# SH DC coefficients are stored in the first coefficient (index 0)
colors_np = np.zeros((num_points, 3), dtype=np.float32)
for i in range(num_points):
# Get DC term from SH coefficients
sh_dc = shs[i * 16] # First SH coefficient contains DC term
# Convert from SH to RGB (simplified - just use DC term)
colors_np[i] = np.clip(sh_dc + 0.5, 0.0, 1.0) # Add 0.5 offset and clamp
# Create vertex data
vertex_data = []
for i in range(num_points):
# Basic properties
vertex = (
positions[i][0], positions[i][1], positions[i][2],
np.log(scales[i][0]), np.log(scales[i][1]), np.log(scales[i][2]), # Log-space encoding
(opacities[i])
)
# Add rotation quaternion elements
quat = rotations[i]
rot_elements = (quat[0], quat[1], quat[2], quat[3]) # x, y, z, w
vertex += rot_elements
# Add RGB colors (convert to 0-255 range)
color_255 = (
int(np.clip(colors_np[i][0] * 255, 0, 255)),
int(np.clip(colors_np[i][1] * 255, 0, 255)),
int(np.clip(colors_np[i][2] * 255, 0, 255))
)
vertex += color_255
# Add SH coefficients
sh_dc = tuple(shs[i * 16][j] for j in range(3))
vertex += sh_dc
# Add remaining SH coefficients
sh_rest = []
for j in range(1, 16):
for c in range(3):
sh_rest.append(shs[i * 16 + j][c])
vertex += tuple(sh_rest)
vertex_data.append(vertex)
# Define the structure of the PLY file
vertex_type = [
('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
('scale_0', 'f4'), ('scale_1', 'f4'), ('scale_2', 'f4'),
('opacity', 'f4')
]
# Add rotation quaternion elements
vertex_type.extend([('rot_0', 'f4'), ('rot_1', 'f4'), ('rot_2', 'f4'), ('rot_3', 'f4')])
# Add RGB color fields
vertex_type.extend([('red', 'u1'), ('green', 'u1'), ('blue', 'u1')])
# Add SH coefficients
vertex_type.extend([('f_dc_0', 'f4'), ('f_dc_1', 'f4'), ('f_dc_2', 'f4')])
# Add remaining SH coefficients
for i in range(45): # 15 coeffs * 3 channels
vertex_type.append((f'f_rest_{i}', 'f4'))
vertex_array = np.array(vertex_data, dtype=vertex_type)
el = PlyElement.describe(vertex_array, 'vertex')
# Create directory if it doesn't exist
os.makedirs(os.path.dirname(filepath), exist_ok=True)
# Save the PLY file
PlyData([el], text=False).write(filepath)
print(f"Point cloud saved to {filepath}")
|