diff --git a/examples/01_scene/09_gaussian_splats.py b/examples/01_scene/09_gaussian_splats.py index 86f061181..7d164afb1 100644 --- a/examples/01_scene/09_gaussian_splats.py +++ b/examples/01_scene/09_gaussian_splats.py @@ -45,6 +45,9 @@ class SplatFile(TypedDict): """(N, 1). Range [0, 1].""" covariances: npt.NDArray[np.floating] """(N, 3, 3).""" + sh_coeffs: npt.NDArray[np.floating] | None + """(N, K, 3) spherical harmonics coefficients (including the DC term), or + None if the file only stores per-Gaussian colors.""" def load_splat_file(splat_path: Path, center: bool = False) -> SplatFile: @@ -88,6 +91,8 @@ def load_splat_file(splat_path: Path, center: bool = False) -> SplatFile: opacities=splat_uint8[:, 27:28] / 255.0, # Covariances should have shape (N, 3, 3). covariances=covariances, + # .splat files don't store spherical harmonics. + sh_coeffs=None, ) @@ -102,9 +107,26 @@ def load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile: positions = np.stack([v["x"], v["y"], v["z"]], axis=-1) scales = np.exp(np.stack([v["scale_0"], v["scale_1"], v["scale_2"]], axis=-1)) wxyzs = np.stack([v["rot_0"], v["rot_1"], v["rot_2"], v["rot_3"]], axis=1) - colors = 0.5 + SH_C0 * np.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1) + dc_terms = np.stack([v["f_dc_0"], v["f_dc_1"], v["f_dc_2"]], axis=1) + colors = 0.5 + SH_C0 * dc_terms opacities = 1.0 / (1.0 + np.exp(-v["opacity"][:, None])) + # Read higher-order spherical harmonics coefficients, if present. 3DGS + # checkpoints store them channel-major (all R terms, then G, then B); we + # want (N, K, 3) with the DC term first. + field_names = {prop.name for prop in v.properties} + num_rest = sum(1 for name in field_names if name.startswith("f_rest_")) + sh_coeffs = None + if num_rest > 0: + assert num_rest % 3 == 0 + rest_per_channel = num_rest // 3 + rest_terms = np.stack( + [v[f"f_rest_{i}"] for i in range(num_rest)], axis=1 + ).reshape((-1, 3, rest_per_channel)) + sh_coeffs = np.concatenate( + [dc_terms[:, :, None], rest_terms], axis=2 + ).transpose(0, 2, 1) + Rs = tf.SO3(wxyzs).as_matrix() covariances = np.einsum( "nij,njk,nlk->nil", Rs, np.eye(3)[None, :, :] * scales[:, None, :] ** 2, Rs @@ -121,6 +143,7 @@ def load_ply_file(ply_file_path: Path, center: bool = False) -> SplatFile: rgbs=colors, opacities=opacities, covariances=covariances, + sh_coeffs=sh_coeffs, ) @@ -147,6 +170,7 @@ def main( rgbs=splat_data["rgbs"], opacities=splat_data["opacities"], covariances=splat_data["covariances"], + sh_coeffs=splat_data["sh_coeffs"], ) remove_button = server.gui.add_button(f"Remove splat object {i}") diff --git a/src/viser/_messages.py b/src/viser/_messages.py index 609bf4971..a3d20fb9a 100644 --- a/src/viser/_messages.py +++ b/src/viser/_messages.py @@ -2146,6 +2146,17 @@ class GaussianSplatsProps: - rgba (int32) Where cov1-6 are the upper-triangular terms of covariance matrices.""" + sh_degree: int + """Spherical harmonics degree for view-dependent colors. 0 means no + spherical harmonics (colors are read from `buffer`).""" + sh_buffer: Optional[npt.NDArray[np.uint32]] + """Optional spherical harmonics coefficients, used when `sh_degree > 0`. + + Each Gaussian gets `4 * ceil(3 * (sh_degree + 1)^2 / 8)` uint32 elements: + all `3 * (sh_degree + 1)^2` coefficients (including the DC term) as + float16, in coefficient-major order (c0.rgb, c1.rgb, ...), zero-padded to + a multiple of 8 float16s so each Gaussian spans a whole number of RGBA32UI + texels.""" scale: Union[float, Tuple[float, float, float]] = 1.0 """Scale of the Gaussian splats. A single float for uniform scaling or a tuple of (x, y, z) for per-axis scaling.""" diff --git a/src/viser/_scene_api.py b/src/viser/_scene_api.py index 4b1dbbfba..2acdb3821 100644 --- a/src/viser/_scene_api.py +++ b/src/viser/_scene_api.py @@ -2169,6 +2169,7 @@ def add_gaussian_splats( rgbs: np.ndarray, opacities: np.ndarray, *, + sh_coeffs: np.ndarray | None = None, scale: float | tuple[float, float, float] = 1.0, wxyz: Tuple[float, float, float, float] | np.ndarray = (1.0, 0.0, 0.0, 0.0), position: Tuple[float, float, float] | np.ndarray = (0.0, 0.0, 0.0), @@ -2185,6 +2186,13 @@ def add_gaussian_splats( covariances: Second moment for each Gaussian. (N, 3, 3). rgbs: Color for each Gaussian. (N, 3). opacities: Opacity for each Gaussian. (N, 1). + sh_coeffs: Optional spherical harmonics coefficients for + view-dependent colors, in the 3DGS (inria) convention. (N, K, + 3), where K is 4, 9, or 16 for SH degrees 1, 2, and 3. The + first coefficient is the DC term (`f_dc` in standard 3DGS + checkpoints); the remainder are the higher-order terms + (`f_rest`), lowest order first. When provided, colors are + computed from the harmonics and `rgbs` is ignored. scale: Scale of the Gaussian splats. A single float for uniform scaling or a tuple of (x, y, z) for per-axis scaling. wxyz: R_parent_local transformation. @@ -2200,6 +2208,29 @@ def add_gaussian_splats( assert opacities.shape == (num_gaussians, 1) assert covariances.shape == (num_gaussians, 3, 3) + sh_degree = 0 + sh_buffer = None + if sh_coeffs is not None: + degree_from_coeff_count = {4: 1, 9: 2, 16: 3} + assert ( + sh_coeffs.ndim == 3 + and sh_coeffs.shape[0] == num_gaussians + and sh_coeffs.shape[1] in degree_from_coeff_count + and sh_coeffs.shape[2] == 3 + ), ( + "sh_coeffs must have shape (N, K, 3) with K in (4, 9, 16)," + f" got {sh_coeffs.shape}" + ) + sh_degree = degree_from_coeff_count[sh_coeffs.shape[1]] + + # Pack coefficients as float16, zero-padded so each Gaussian spans + # a whole number of RGBA32UI texels (8 float16s each) client-side. + num_floats = sh_coeffs.shape[1] * 3 + num_floats_padded = -(-num_floats // 8) * 8 + sh_f16 = np.zeros((num_gaussians, num_floats_padded), dtype=np.float16) + sh_f16[:, :num_floats] = sh_coeffs.reshape(num_gaussians, num_floats) + sh_buffer = sh_f16.view(np.uint32) + # Get upper-triangular terms of covariance matrix. cov_triu = covariances.reshape((-1, 9))[:, np.array([0, 1, 2, 4, 5, 8])] buffer = np.concatenate( @@ -2225,6 +2256,8 @@ def add_gaussian_splats( props=_messages.GaussianSplatsProps( buffer=buffer, scale=scale, + sh_degree=sh_degree, + sh_buffer=sh_buffer, ), ) node_handle = GaussianSplatHandle._make( diff --git a/src/viser/_scene_handles.py b/src/viser/_scene_handles.py index 792a1b845..010d93538 100644 --- a/src/viser/_scene_handles.py +++ b/src/viser/_scene_handles.py @@ -999,6 +999,12 @@ class GaussianSplatHandle( - [3]: reserved for renderer - [4:7]: covariance upper-triangular (6x float16) - [7]: RGBA (4x uint8) + + When `sh_degree > 0`, view-dependent colors are computed from the + float16 spherical harmonics coefficients in `sh_buffer` instead of the + RGB values in `buffer`. The sub-property setters below (`centers`, + `rgbs`, ...) update `buffer` only; if they change the number of + Gaussians, a stale `sh_buffer` is ignored by the renderer. """ def _ensure_buffer_size(self, num_gaussians: int) -> None: diff --git a/src/viser/client/src/SceneTree.tsx b/src/viser/client/src/SceneTree.tsx index 3df11ec02..9e7c49727 100644 --- a/src/viser/client/src/SceneTree.tsx +++ b/src/viser/client/src/SceneTree.tsx @@ -580,6 +580,8 @@ function createObjectFactory( diff --git a/src/viser/client/src/Splatting/GaussianSplats.tsx b/src/viser/client/src/Splatting/GaussianSplats.tsx index 3c3e557ef..a0bb4e2f3 100644 --- a/src/viser/client/src/Splatting/GaussianSplats.tsx +++ b/src/viser/client/src/Splatting/GaussianSplats.tsx @@ -40,9 +40,11 @@ import { v4 as uuidv4 } from "uuid"; import { GaussianSplatsContext, + SH_TEXELS_PER_GAUSSIAN, createGaussianMeshProps, useGaussianSplatStore, type GaussianMeshProps, + type SplatGroupData, } from "./GaussianSplatsHelpers"; import { ViewerContext } from "../ViewerContext"; @@ -70,10 +72,15 @@ export const SplatObject = React.forwardRef< THREE.Group, { buffer: Uint32Array; + shBuffer?: Uint32Array | null; + shDegree?: number; sceneNodeName?: string; children?: React.ReactNode; } ->(function SplatObject({ buffer, sceneNodeName, children }, ref) { +>(function SplatObject( + { buffer, shBuffer = null, shDegree = 0, sceneNodeName, children }, + ref, +) { const splatContext = React.useContext(GaussianSplatsContext)!; const { setBuffer, removeBuffer } = splatContext.gaussianSplatState.actions; const nodeRefFromId = splatContext.gaussianSplatState.store( @@ -96,8 +103,8 @@ export const SplatObject = React.forwardRef< // Update buffer when it changes. React.useEffect(() => { - setBuffer(name, buffer); - }, [name, buffer, setBuffer]); + setBuffer(name, { buffer, shBuffer, shDegree }); + }, [name, buffer, shBuffer, shDegree, setBuffer]); return ( (null); const isFirstRenderRef = React.useRef(true); const initializedBufferTextureRef = React.useRef(false); @@ -187,11 +196,14 @@ function SplatRendererImpl() { !prevMergedRef.current || merged.gaussianBuffer !== prevMergedRef.current.gaussianBuffer; - // Check if number of Gaussians or groups changed (requires texture resize). + // Check if number of Gaussians or groups changed (requires texture + // resize). A spherical harmonics degree change also resizes the SH + // texture, so it takes the same path. const sizeChanged = prevMergedRef.current && (merged.numGaussians !== prevMergedRef.current.numGaussians || - merged.numGroups !== prevMergedRef.current.numGroups); + merged.numGroups !== prevMergedRef.current.numGroups || + merged.shDegree !== prevMergedRef.current.shDegree); // Initialize resources on first render. if (isFirstRenderRef.current) { @@ -200,6 +212,8 @@ function SplatRendererImpl() { merged.gaussianBuffer, merged.numGroups, maxTextureSize, + merged.shBuffer, + merged.shDegree, ); // Show splats immediately with identity sort order. This makes splats @@ -248,10 +262,13 @@ function SplatRendererImpl() { merged.gaussianBuffer, merged.numGroups, maxTextureSize, + merged.shBuffer, + merged.shDegree, ); // Dispose old resources. oldProps.textureBuffer.dispose(); + oldProps.shTextureBuffer.dispose(); oldProps.geometry.dispose(); oldProps.material.dispose(); oldProps.textureT_camera_groups.dispose(); @@ -278,6 +295,18 @@ function SplatRendererImpl() { textureData.set(merged.gaussianBuffer); meshPropsRef.current.textureBuffer.needsUpdate = true; + if ( + merged.shBuffer !== null && + merged.shBuffer !== prevMergedRef.current?.shBuffer + ) { + // Same degree (else sizeChanged), so the texture layouts match. + const shTextureData = meshPropsRef.current.shTextureBuffer.image + .data as Uint32Array; + shTextureData.fill(0); + shTextureData.set(merged.shBuffer); + meshPropsRef.current.shTextureBuffer.needsUpdate = true; + } + // Update worker with new buffer. postToWorker({ updateBuffer: merged.gaussianBuffer, @@ -299,6 +328,7 @@ function SplatRendererImpl() { return () => { if (meshPropsRef.current) { meshPropsRef.current.textureBuffer.dispose(); + meshPropsRef.current.shTextureBuffer.dispose(); meshPropsRef.current.geometry.dispose(); meshPropsRef.current.material.dispose(); meshPropsRef.current.textureT_camera_groups.dispose(); @@ -589,21 +619,22 @@ function SplatRendererImpl() { /**Consolidate groups of Gaussians into a single buffer, to make it possible * for them to be sorted globally.*/ function mergeGaussianGroups(groupBufferFromName: { - [name: string]: Uint32Array; + [name: string]: SplatGroupData; }) { // Create geometry. Each Gaussian will be rendered as a quad. let totalBufferLength = 0; - for (const buffer of Object.values(groupBufferFromName)) { - totalBufferLength += buffer.length; + for (const group of Object.values(groupBufferFromName)) { + totalBufferLength += group.buffer.length; } const numGaussians = totalBufferLength / 8; const gaussianBuffer = new Uint32Array(totalBufferLength); const groupIndices = new Uint32Array(numGaussians); let offset = 0; - for (const [groupIndex, groupBuffer] of Object.values( + for (const [groupIndex, group] of Object.values( groupBufferFromName, ).entries()) { + const groupBuffer = group.buffer; groupIndices.fill( groupIndex, offset / 8, @@ -623,6 +654,92 @@ function mergeGaussianGroups(groupBufferFromName: { offset += groupBuffer.length; } + const { shBuffer, shDegree } = mergeShBuffers( + groupBufferFromName, + numGaussians, + ); + const numGroups = Object.keys(groupBufferFromName).length; - return { numGaussians, gaussianBuffer, numGroups, groupIndices }; + return { + numGaussians, + gaussianBuffer, + numGroups, + groupIndices, + shBuffer, + shDegree, + }; +} + +/**Consolidate spherical harmonics coefficients, with the same Gaussian + * ordering as the merged splat buffer. + * + * The shader uses a single global SH degree, so when groups have different + * degrees we promote everything to the maximum: lower-degree groups keep + * their coefficients with zeros for the missing higher-order terms, and + * groups without spherical harmonics get DC-only coefficients synthesized + * from their RGBA colors. Returns a null buffer when no group has spherical + * harmonics, which disables the feature entirely.*/ +function mergeShBuffers( + groupBufferFromName: { [name: string]: SplatGroupData }, + numGaussians: number, +): { shBuffer: Uint32Array | null; shDegree: number } { + let shDegree = 0; + for (const group of Object.values(groupBufferFromName)) { + if (group.shBuffer !== null && group.shDegree > 0) { + shDegree = Math.max(shDegree, Math.min(group.shDegree, 3)); + } + } + if (shDegree === 0) { + return { shBuffer: null, shDegree: 0 }; + } + + const uint32PerGaussian = SH_TEXELS_PER_GAUSSIAN[shDegree] * 4; + const shBuffer = new Uint32Array(numGaussians * uint32PerGaussian); + const SH_C0 = 0.28209479177387814; + + let gaussianOffset = 0; + for (const group of Object.values(groupBufferFromName)) { + const numGroupGaussians = group.buffer.length / 8; + const srcPerGaussianExpected = + group.shDegree > 0 ? SH_TEXELS_PER_GAUSSIAN[group.shDegree] * 4 : 0; + if ( + group.shBuffer !== null && + group.shDegree > 0 && + // Guard against a stale coefficient buffer, e.g. after the main buffer + // was resized without updating the harmonics. + group.shBuffer.length === numGroupGaussians * srcPerGaussianExpected + ) { + const srcPerGaussian = srcPerGaussianExpected; + if (srcPerGaussian === uint32PerGaussian) { + shBuffer.set(group.shBuffer, gaussianOffset * uint32PerGaussian); + } else { + // Lower degree than the merged scene; copy per-Gaussian. The source + // rows end with zero padding (see GaussianSplatsProps.sh_buffer), + // which lands on higher-order coefficients and is harmless. + for (let i = 0; i < numGroupGaussians; i++) { + for (let j = 0; j < srcPerGaussian; j++) { + shBuffer[(gaussianOffset + i) * uint32PerGaussian + j] = + group.shBuffer[i * srcPerGaussian + j]; + } + } + } + } else { + // No spherical harmonics: synthesize a DC-only coefficient that + // reproduces the RGBA color, since `color = C0 * dc + 0.5`. + for (let i = 0; i < numGroupGaussians; i++) { + const rgba = group.buffer[i * 8 + 7]; + const r = ((rgba & 0xff) / 255.0 - 0.5) / SH_C0; + const g = (((rgba >> 8) & 0xff) / 255.0 - 0.5) / SH_C0; + const b = (((rgba >> 16) & 0xff) / 255.0 - 0.5) / SH_C0; + const base = (gaussianOffset + i) * uint32PerGaussian; + shBuffer[base] = + THREE.DataUtils.toHalfFloat(r) | + (THREE.DataUtils.toHalfFloat(g) << 16); + shBuffer[base + 1] = THREE.DataUtils.toHalfFloat(b); + } + } + gaussianOffset += numGroupGaussians; + } + + return { shBuffer, shDegree }; } diff --git a/src/viser/client/src/Splatting/GaussianSplatsHelpers.ts b/src/viser/client/src/Splatting/GaussianSplatsHelpers.ts index f439cd1d7..4fcdeedf7 100644 --- a/src/viser/client/src/Splatting/GaussianSplatsHelpers.ts +++ b/src/viser/client/src/Splatting/GaussianSplatsHelpers.ts @@ -5,6 +5,11 @@ import { Object3D } from "three"; import { useThree } from "@react-three/fiber"; import { shaderMaterial } from "@react-three/drei"; +/** Number of RGBA32UI texels each Gaussian occupies in the spherical + * harmonics texture, indexed by SH degree. Each texel holds 8 float16 + * coefficients; a degree-d Gaussian has 3 * (d + 1)^2 coefficients. */ +export const SH_TEXELS_PER_GAUSSIAN = [1, 2, 4, 6]; + const GaussianSplatMaterial = /* @__PURE__ */ shaderMaterial( { numGaussians: 0, @@ -15,6 +20,9 @@ const GaussianSplatMaterial = /* @__PURE__ */ shaderMaterial( depthWrite: false, transparent: true, textureBuffer: null as THREE.DataTexture | null, + shTextureBuffer: null as THREE.DataTexture | null, + shTexelsPerGaussian: 0, + shDegree: 0, textureT_camera_groups: null as THREE.DataTexture | null, transitionInState: 0.0, projectionMatrixCustom: new THREE.Matrix4(), @@ -32,6 +40,13 @@ const GaussianSplatMaterial = /* @__PURE__ */ shaderMaterial( // copy quadjr for this. uniform usampler2D textureBuffer; + // Spherical harmonics coefficients, as float16s packed into RGBA32UI + // texels. Each Gaussian gets shTexelsPerGaussian texels; 0 disables + // spherical harmonics (colors come from textureBuffer's RGBA). + uniform usampler2D shTextureBuffer; + uniform uint shTexelsPerGaussian; + uniform uint shDegree; + // We could also use a uniform to store transforms, but this would be more // limiting in terms of the # of groups we can have. uniform sampler2D textureT_camera_groups; @@ -68,6 +83,99 @@ const GaussianSplatMaterial = /* @__PURE__ */ shaderMaterial( return transpose(transform); } + // Evaluate view-dependent color from spherical harmonics coefficients, in + // the 3DGS (inria) convention. Uses the same recurrence relations as + // gsplat's spherical_harmonics(); coefficients are laid out + // coefficient-major (c0.rgb, c1.rgb, ...) as float16 pairs in RGBA32UI + // texels. + vec3 evalSphericalHarmonics(vec3 center, mat4 T_camera_group) { + // Unpack this Gaussian's coefficients; up to 48 for degree 3. Elements + // beyond shTexelsPerGaussian * 8 are left untouched and must not be read. + float coeffs[48]; + ivec2 texSize = textureSize(shTextureBuffer, 0); + int texStart = int(sortedIndex * shTexelsPerGaussian); + for (int i = 0; i < 6; i++) { + if (i >= int(shTexelsPerGaussian)) break; + int texIndex = texStart + i; + uvec4 texel = texelFetch( + shTextureBuffer, + ivec2(texIndex % texSize.x, texIndex / texSize.x), + 0); + vec2 v01 = unpackHalf2x16(texel.x); + vec2 v23 = unpackHalf2x16(texel.y); + vec2 v45 = unpackHalf2x16(texel.z); + vec2 v67 = unpackHalf2x16(texel.w); + coeffs[i * 8 + 0] = v01.x; + coeffs[i * 8 + 1] = v01.y; + coeffs[i * 8 + 2] = v23.x; + coeffs[i * 8 + 3] = v23.y; + coeffs[i * 8 + 4] = v45.x; + coeffs[i * 8 + 5] = v45.y; + coeffs[i * 8 + 6] = v67.x; + coeffs[i * 8 + 7] = v67.y; + } + + // View direction in the group frame, which the harmonics are defined in. + // The camera position in the group frame is -R^T t, from + // T_camera_group = [R | t]. + vec3 t_group_camera = -(transpose(mat3(T_camera_group)) * T_camera_group[3].xyz); + vec3 dir = normalize(center - t_group_camera); + float x = dir.x; + float y = dir.y; + float z = dir.z; + + // Degree 0. + vec3 rgb = 0.2820947917738781 * vec3(coeffs[0], coeffs[1], coeffs[2]); + + // Degree 1. + rgb += 0.48860251190291987 * ( + -y * vec3(coeffs[3], coeffs[4], coeffs[5]) + + z * vec3(coeffs[6], coeffs[7], coeffs[8]) + - x * vec3(coeffs[9], coeffs[10], coeffs[11])); + + if (shDegree >= 2u) { + float xx = x * x; + float yy = y * y; + float zz = z * z; + float fTmp0B = -1.092548430592079 * z; + float fC1 = xx - yy; + float fS1 = 2.0 * x * y; + float pSH6 = 0.9461746957575601 * zz - 0.3153915652525201; + float pSH7 = fTmp0B * x; + float pSH5 = fTmp0B * y; + float pSH8 = 0.5462742152960395 * fC1; + float pSH4 = 0.5462742152960395 * fS1; + rgb += pSH4 * vec3(coeffs[12], coeffs[13], coeffs[14]) + + pSH5 * vec3(coeffs[15], coeffs[16], coeffs[17]) + + pSH6 * vec3(coeffs[18], coeffs[19], coeffs[20]) + + pSH7 * vec3(coeffs[21], coeffs[22], coeffs[23]) + + pSH8 * vec3(coeffs[24], coeffs[25], coeffs[26]); + + if (shDegree >= 3u) { + float fTmp0C = -2.285228997322329 * zz + 0.4570457994644658; + float fTmp1B = 1.445305721320277 * z; + float fC2 = x * fC1 - y * fS1; + float fS2 = x * fS1 + y * fC1; + float pSH12 = z * (1.865881662950577 * zz - 1.119528997770346); + float pSH13 = fTmp0C * x; + float pSH11 = fTmp0C * y; + float pSH14 = fTmp1B * fC1; + float pSH10 = fTmp1B * fS1; + float pSH15 = -0.5900435899266435 * fC2; + float pSH9 = -0.5900435899266435 * fS2; + rgb += pSH9 * vec3(coeffs[27], coeffs[28], coeffs[29]) + + pSH10 * vec3(coeffs[30], coeffs[31], coeffs[32]) + + pSH11 * vec3(coeffs[33], coeffs[34], coeffs[35]) + + pSH12 * vec3(coeffs[36], coeffs[37], coeffs[38]) + + pSH13 * vec3(coeffs[39], coeffs[40], coeffs[41]) + + pSH14 * vec3(coeffs[42], coeffs[43], coeffs[44]) + + pSH15 * vec3(coeffs[45], coeffs[46], coeffs[47]); + } + } + + return max(rgb + 0.5, vec3(0.0)); + } + void main () { // Get position + scale from float buffer. ivec2 texSize = textureSize(textureBuffer, 0); @@ -147,6 +255,9 @@ const GaussianSplatMaterial = /* @__PURE__ */ shaderMaterial( float((rgbaUint32 >> uint(16)) & uint(0xFF)) / 255.0, float(rgbaUint32 >> uint(24)) / 255.0 ); + if (shDegree > 0u) { + vRgba.rgb = evalSphericalHarmonics(center, T_camera_group); + } // Throw the Gaussian off the screen if it's too close, too far, or too small. float weightedDeterminant = vRgba.a * (diag1 * diag2 - offDiag * offDiag); @@ -192,6 +303,8 @@ export function createGaussianMeshProps( gaussianBuffer: Uint32Array, numGroups: number, maxTextureSize: number, + shBuffer: Uint32Array | null = null, + shDegree: number = 0, ) { const numGaussians = gaussianBuffer.length / 8; @@ -243,9 +356,43 @@ export function createGaussianMeshProps( textureT_camera_groups.internalFormat = "RGBA32F"; textureT_camera_groups.needsUpdate = true; + // Optional texture for spherical harmonics coefficients. When absent, a + // 1x1 placeholder keeps the usampler2D uniform valid; the shader never + // samples it since shDegree is 0. + const shTexelsPerGaussian = + shDegree > 0 ? SH_TEXELS_PER_GAUSSIAN[shDegree] : 0; + let shTextureBuffer: THREE.DataTexture; + if (shBuffer !== null && shDegree > 0) { + const numShTexels = numGaussians * shTexelsPerGaussian; + const shTextureWidth = Math.min(numShTexels, maxTextureSize); + const shTextureHeight = Math.ceil(numShTexels / shTextureWidth); + const shBufferPadded = new Uint32Array(shTextureWidth * shTextureHeight * 4); + shBufferPadded.set(shBuffer); + shTextureBuffer = new THREE.DataTexture( + shBufferPadded, + shTextureWidth, + shTextureHeight, + THREE.RGBAIntegerFormat, + THREE.UnsignedIntType, + ); + } else { + shTextureBuffer = new THREE.DataTexture( + new Uint32Array(4), + 1, + 1, + THREE.RGBAIntegerFormat, + THREE.UnsignedIntType, + ); + } + shTextureBuffer.internalFormat = "RGBA32UI"; + shTextureBuffer.needsUpdate = true; + const material = new GaussianSplatMaterial(); material.fog = true; material.textureBuffer = textureBuffer; + material.shTextureBuffer = shTextureBuffer; + material.shTexelsPerGaussian = shTexelsPerGaussian; + material.shDegree = shDegree; material.textureT_camera_groups = textureT_camera_groups; material.numGaussians = numGaussians; @@ -255,6 +402,7 @@ export function createGaussianMeshProps( textureBuffer, textureWidth, textureHeight, + shTextureBuffer, sortedIndexAttribute, textureT_camera_groups, rowMajorT_camera_groups, @@ -267,14 +415,30 @@ export function createGaussianMeshProps( export function useGaussianMeshProps( gaussianBuffer: Uint32Array, numGroups: number, + shBuffer: Uint32Array | null = null, + shDegree: number = 0, ) { const maxTextureSize = useThree((state) => state.gl).capabilities .maxTextureSize; - return createGaussianMeshProps(gaussianBuffer, numGroups, maxTextureSize); + return createGaussianMeshProps( + gaussianBuffer, + numGroups, + maxTextureSize, + shBuffer, + shDegree, + ); +} +/**Per-group Gaussian data: the main splat buffer, plus optional spherical + * harmonics coefficients for view-dependent colors.*/ +export interface SplatGroupData { + buffer: Uint32Array; + shBuffer: Uint32Array | null; + shDegree: number; } + /**Global splat state.*/ interface SplatState { - groupBufferFromId: { [id: string]: Uint32Array }; + groupBufferFromId: { [id: string]: SplatGroupData }; nodeRefFromId: React.MutableRefObject<{ [name: string]: undefined | Object3D; }>; @@ -284,7 +448,7 @@ interface SplatState { } interface SplatActions { - setBuffer: (id: string, buffer: Uint32Array) => void; + setBuffer: (id: string, data: SplatGroupData) => void; removeBuffer: (id: string) => void; } @@ -302,9 +466,9 @@ export function useGaussianSplatStore() { }); const actions: SplatActions = { - setBuffer: (id, buffer) => { + setBuffer: (id, data) => { store.set((state) => ({ - groupBufferFromId: { ...state.groupBufferFromId, [id]: buffer }, + groupBufferFromId: { ...state.groupBufferFromId, [id]: data }, })); }, removeBuffer: (id) => { diff --git a/src/viser/client/src/WebsocketMessages.ts b/src/viser/client/src/WebsocketMessages.ts index 973839cd8..87047410a 100644 --- a/src/viser/client/src/WebsocketMessages.ts +++ b/src/viser/client/src/WebsocketMessages.ts @@ -501,7 +501,12 @@ export interface CubicBezierSplineMessage { export interface GaussianSplatsMessage { type: "GaussianSplatsMessage"; name: string; - props: { buffer: Uint32Array; scale: number | [number, number, number] }; + props: { + buffer: Uint32Array; + sh_degree: number; + sh_buffer: Uint32Array | null; + scale: number | [number, number, number]; + }; } /** Remove a particular node from the scene. * @@ -3093,6 +3098,14 @@ export const SceneNodePropsSchema: { kind: "default", tsType: "Uint32Array", }, + sh_degree: { + kind: "default", + tsType: "number", + }, + sh_buffer: { + kind: "default", + tsType: "(Uint32Array | null)", + }, scale: { kind: "default", tsType: "(number | [number, number, number])", diff --git a/tests/e2e/test_gaussian_splat_spherical_harmonics.py b/tests/e2e/test_gaussian_splat_spherical_harmonics.py new file mode 100644 index 000000000..e8d182413 --- /dev/null +++ b/tests/e2e/test_gaussian_splat_spherical_harmonics.py @@ -0,0 +1,145 @@ +"""E2E tests for Gaussian splat spherical harmonics rendering. + +Verifies that splats with spherical harmonics coefficients are view-dependent: +a degree-1 coefficient along +X makes the splats red-dominant when viewed from ++X and red-suppressed when viewed from -X. Also verifies that plain RGB splats +keep their color when composited with a spherical harmonics group (the +renderer promotes them to DC-only harmonics internally). +""" + +from __future__ import annotations + +from io import BytesIO + +import numpy as np +import pytest +from PIL import Image +from playwright.sync_api import Page + +import viser + +from .utils import wait_for_scene_node + +SH_C0 = 0.28209479177387814 +SH_C1 = 0.4886025119029199 + + +def _canvas_mean_color(page: Page) -> np.ndarray: + """Mean RGB of non-background canvas pixels (background is near-white).""" + canvas = page.locator("canvas").first + screenshot = canvas.screenshot() + img = np.array(Image.open(BytesIO(screenshot)).convert("RGB")).astype(np.float64) + non_background = img.min(axis=2) < 220 + assert non_background.sum() > 500, "Expected the splats to cover some pixels" + return img[non_background].mean(axis=0) + + +def _look_from(client: viser.ClientHandle, position: tuple) -> None: + client.camera.position = position + client.camera.look_at = (0.0, 0.0, 0.0) + + +@pytest.mark.parametrize("sh_degree", [1, 3]) +def test_gaussian_splat_sh_view_dependent_color( + viser_server: viser.ViserServer, + viser_page: Page, + sh_degree: int, +) -> None: + """Colors computed from spherical harmonics must change with view direction.""" + num_gaussians = 100 + + # Gray DC term, and a red coefficient on a basis function that is odd in + # x, so the red channel flips as the camera moves from +X to -X while + # green/blue stay at 0.5. + # + # Degree 1 uses coefficient 3, whose basis is `-SH_C1 * x`. Degree 3 uses + # coefficient 15 (the last one, exercising the texture fetch of all six + # RGBA32UI texels), whose basis is `-0.59 * x * (x^2 - 3 y^2)`; along the + # X axis both evaluate to -/+ their constant for view direction +/-X. + num_coeffs = (sh_degree + 1) ** 2 + sh_coeffs = np.zeros((num_gaussians, num_coeffs, 3), dtype=np.float32) + if sh_degree == 1: + sh_coeffs[:, 3, 0] = 0.5 / SH_C1 + else: + sh_coeffs[:, 15, 0] = 0.5 / 0.5900435899266435 + + viser_server.scene.add_gaussian_splats( + "/sh_splat", + centers=np.zeros((num_gaussians, 3), dtype=np.float32), + rgbs=np.full((num_gaussians, 3), 0.5, dtype=np.float32), + opacities=np.ones((num_gaussians, 1), dtype=np.float32), + covariances=np.tile(np.eye(3, dtype=np.float32), (num_gaussians, 1, 1)), + sh_coeffs=sh_coeffs, + ) + + wait_for_scene_node(viser_page, "/sh_splat") + # Give time for WASM init + sort + rendering. + viser_page.wait_for_timeout(3000) + + clients = viser_server.get_clients() + assert len(clients) == 1 + client = next(iter(clients.values())) + + _look_from(client, (8.0, 0.0, 0.0)) + viser_page.wait_for_timeout(1500) + color_from_pos_x = _canvas_mean_color(viser_page) + + _look_from(client, (-8.0, 0.0, 0.0)) + viser_page.wait_for_timeout(1500) + color_from_neg_x = _canvas_mean_color(viser_page) + + # Red must dominate from +X and vanish from -X; green/blue barely move. + assert color_from_pos_x[0] - color_from_neg_x[0] > 100, ( + f"Expected strong view-dependent red channel, got " + f"{color_from_pos_x=} vs {color_from_neg_x=}" + ) + for channel in (1, 2): + assert abs(color_from_pos_x[channel] - color_from_neg_x[channel]) < 40, ( + f"Channel {channel} should be view-independent, got " + f"{color_from_pos_x=} vs {color_from_neg_x=}" + ) + + +def test_gaussian_splat_sh_mixed_with_rgb_group( + viser_server: viser.ViserServer, + viser_page: Page, +) -> None: + """Plain RGB splats keep their color when a SH group is also present.""" + num_gaussians = 100 + + # Green splats without spherical harmonics... + viser_server.scene.add_gaussian_splats( + "/rgb_splat", + centers=np.zeros((num_gaussians, 3), dtype=np.float32), + rgbs=np.tile(np.array([[0.0, 1.0, 0.0]], dtype=np.float32), (num_gaussians, 1)), + opacities=np.ones((num_gaussians, 1), dtype=np.float32), + covariances=np.tile(np.eye(3, dtype=np.float32), (num_gaussians, 1, 1)), + ) + # ...composited with a degree-1 SH group far off to the side, which forces + # the whole scene through the spherical harmonics shader path. + sh_coeffs = np.zeros((num_gaussians, 4, 3), dtype=np.float32) + sh_coeffs[:, 0, :] = 0.5 / SH_C0 # White DC term. + viser_server.scene.add_gaussian_splats( + "/sh_splat", + centers=np.full((num_gaussians, 3), 50.0, dtype=np.float32), + rgbs=np.full((num_gaussians, 3), 0.5, dtype=np.float32), + opacities=np.ones((num_gaussians, 1), dtype=np.float32), + covariances=np.tile(np.eye(3, dtype=np.float32), (num_gaussians, 1, 1)), + sh_coeffs=sh_coeffs, + ) + + wait_for_scene_node(viser_page, "/rgb_splat") + wait_for_scene_node(viser_page, "/sh_splat") + viser_page.wait_for_timeout(3000) + + canvas = viser_page.locator("canvas").first + img = np.array( + Image.open(BytesIO(canvas.screenshot())).convert("RGB") + ).astype(np.float64) + green_mask = ( + (img[:, :, 1] > 150) & (img[:, :, 0] < 100) & (img[:, :, 2] < 100) + ) + assert green_mask.sum() > 1000, ( + f"RGB-only splats should still render green when mixed with a " + f"spherical harmonics group; found {green_mask.sum()} green pixels" + )