From 6e25c3d879eac7ab79be143227106e6ba2dcec54 Mon Sep 17 00:00:00 2001 From: hybridherbst Date: Mon, 14 Sep 2026 22:16:41 +0200 Subject: [PATCH] Expose decoded spherical harmonics per splat --- docs/docs/ext-splats.md | 2 +- docs/docs/packed-splats.md | 18 +-- docs/docs/splat-mesh.md | 6 +- package.json | 2 +- src/ExtSplats.ts | 103 +++++++++++--- src/PackedSplats.ts | 73 +++++++--- src/SplatData.ts | 32 +++++ src/SplatMesh.ts | 45 ++++--- src/SplatPager.ts | 60 +++++++-- src/index.ts | 5 + src/utils.ts | 180 ++++++++++++++++++++++--- test/spherical-harmonics.test.ts | 221 +++++++++++++++++++++++++++++++ test/wasm-loader.mjs | 35 +++++ 13 files changed, 674 insertions(+), 108 deletions(-) create mode 100644 src/SplatData.ts create mode 100644 test/spherical-harmonics.test.ts create mode 100644 test/wasm-loader.mjs diff --git a/docs/docs/ext-splats.md b/docs/docs/ext-splats.md index 4e08d1ae..9d2e7d58 100644 --- a/docs/docs/ext-splats.md +++ b/docs/docs/ext-splats.md @@ -81,7 +81,7 @@ Each `ExtSplat` occupies 8 × `uint32` total = 32 bytes. import { utils } from "@sparkjsdev/spark"; utils.encodeExtSplat(extSplats.extArrays, index, x, y, z, sx, sy, sz, qx, qy, qz, qw, opacity, r, g, b); -const { center, scales, quaternion, color, opacity } = utils.decodeExtSplat(extSplats.extArrays, index); +const { center, scales, quaternion, color, opacity, sphericalHarmonics } = utils.decodeExtSplat(extSplats.extArrays, index, extSplats.extra); ``` You can also use instance helpers: `setSplat`, `pushSplat`, `getSplat`, and `forEachSplat`. diff --git a/docs/docs/packed-splats.md b/docs/docs/packed-splats.md index 1c93034f..97ea3d95 100644 --- a/docs/docs/packed-splats.md +++ b/docs/docs/packed-splats.md @@ -62,7 +62,7 @@ Utility functions are provided in Javascript to pack/unpack these encodings: ```javascript // Set via packedSplats interface -packedSplats.setSplat(index, center, scales, quaternion, opacity, color); +packedSplats.setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics); // Set underlying Uint32 array directly import { utils } from "@sparkjsdev/spark"; @@ -72,15 +72,15 @@ utils.setPackedSplat(packedSplats.packedArray, index, x, y, z, scaleX, scaleY, . utils.setPackedSplatQuat(packedSplats.packedArray, index, quatX, quatY, quatZ, quatW); // Unpack all splat components from the Uint32 array -const { center, scales, quaternion, color, opacity } = utils.unpackSplat(packedSplats.packedArray, index); +const { center, scales, quaternion, color, opacity, sphericalHarmonics } = utils.unpackSplat(packedSplats.packedArray, index, packedSplats.splatEncoding, packedSplats.extra); // Unpack all splats with callback -packedSplats.forEachSplat((index, center, scales, quaternion, opacity, color) => { +packedSplats.forEachSplat((index, center, scales, quaternion, opacity, color, sphericalHarmonics) => { // Use unpacked splat data. Changing the inputs directly has no effect. // Update just the scales component utils.setPackedSplatScales(packedSplat.packedArray, index, 0.005, 0.01, 0.015); // Update the entire splat - packedSplat.setSplat(index, center, scales, quaternion, opacity, color); + packedSplat.setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics); }); ``` @@ -175,19 +175,19 @@ Ensures that `this.packedArray` can fit `numSplats` splats. If it's too small, r Typically you don't need to call this, because calling `this.setSplat(index, ...)` and `this.pushSplat(...)` will automatically call `ensureSplats()` so we have enough splats. -### `getSplat(index): { center, scales, quaternion, opacity, color }` +### `getSplat(index): { center, scales, quaternion, opacity, color, sphericalHarmonics }` -Unpack the 16-byte splat data at `index` into the THREE.js components `center: THREE.Vector3`, `scales: THREE.Vector3`, `quaternion: THREE.Quaternion`, `opacity: number 0..1`, `color: THREE.Color 0..1`. +Unpack the splat data at `index`, including decoded `sphericalHarmonics.sh1`, `.sh2`, and `.sh3` bands when present. -### `setSplat(index, center, scales, quaternion, opacity, color)` +### `setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics?)` Set all PackedSplat components at `index` with the provided splat attributes (can be the same objects returned by `getSplat`). Ensures there is capacity for at least `index+1` splats. -### `pushSplat(center, scales, quaternion, opacity, color)` +### `pushSplat(center, scales, quaternion, opacity, color, sphericalHarmonics?)` Effectively calls `this.setSplat(this.numSplats++, center, ...)`, useful on construction where you just want to iterate and create a collection of splats. -### `forEachSplat(callback: (index, center, scales, quaternion, opacity, color) => void)` +### `forEachSplat(callback: (index, center, scales, quaternion, opacity, color, sphericalHarmonics) => void)` Iterate over splats index `0..=(this.numSplats-1)`, unpack each splat and invoke the callback function with the splat attributes. diff --git a/docs/docs/splat-mesh.md b/docs/docs/splat-mesh.md index bbeb926d..0cc202bb 100644 --- a/docs/docs/splat-mesh.md +++ b/docs/docs/splat-mesh.md @@ -141,13 +141,13 @@ Additional properties on a `SplatMesh` instance: Call this when you are finished with the `SplatMesh` and want to free any buffers it holds (via `packedSplats`). -## `pushSplat(center, scales, quaternion, opacity, color)` +## `pushSplat(center, scales, quaternion, opacity, color, sphericalHarmonics?)` Creates a new splat with the provided parameters (all values in "float" space, i.e. 0-1 for opacity and color) and adds it to the end of the `packedSplats`, increasing `numSplats` by 1. If necessary, reallocates the buffer with an exponential doubling strategy to fit the new data, so it's fairly efficient to `pushSplat(...)` each splat you want to create in a loop. -## `forEachSplat(callback: (index, center, scales, quaternion, opacity, color) => void)` +## `forEachSplat(callback: (index, center, scales, quaternion, opacity, color, sphericalHarmonics) => void)` -This method iterates over all splats in this instance's `packedSplats`, invoking the provided callback with `index: number` in `0..=(this.numSplats-1)`, `center: THREE.Vector3`, `scales: THREE.Vector3`, `quaternion: THREE.Quaternion`, `opacity: number` (0..1), and `color: THREE.Color` (rgb values in 0..1). Note that the objects passed in as `center` etc. are the same for every callback invocation: they are reused for efficiency. *Changing these values has no effect* as they are decoded/unpacked copies of the underlying data. To update the `packedSplats`, call `.packedSplats.setSplat(index, center, scales, quaternion, opacity, color)`. +This method iterates over all splats in this instance's `packedSplats`, invoking the provided callback with `index: number` in `0..=(this.numSplats-1)`, `center: THREE.Vector3`, `scales: THREE.Vector3`, `quaternion: THREE.Quaternion`, `opacity: number` (0..1), `color: THREE.Color` (rgb values in 0..1), and decoded `sphericalHarmonics`. Note that the objects passed in as `center` etc. are the same for every callback invocation: these objects are reused for efficiency. *Changing these values has no effect* as they are decoded/unpacked copies of the underlying data. To update the `packedSplats`, call `.packedSplats.setSplat(index, center, scales, quaternion, opacity, color, sphericalHarmonics)`. ## `getBoundingBox(centers_only=true)` diff --git a/package.json b/package.json index bb1357fa..cc27f4d8 100644 --- a/package.json +++ b/package.json @@ -38,7 +38,7 @@ "site:deploy": "npm run site:build && node scripts/deploy-site.js", "site:serve": "node scripts/serve-site.js site", "start": "npm run dev", - "test": "node --no-warnings --loader ts-node/esm --test test/**/*.test.ts" + "test": "node --no-warnings --loader ts-node/esm --loader ./test/wasm-loader.mjs --test test/**/*.test.ts" }, "repository": "sparkjs-dev/spark", "files": ["dist"], diff --git a/src/ExtSplats.ts b/src/ExtSplats.ts index 0184732f..2a673e3a 100644 --- a/src/ExtSplats.ts +++ b/src/ExtSplats.ts @@ -1,5 +1,10 @@ import * as THREE from "three"; import type { RgbaArray } from "./RgbaArray"; +import type { + ForEachSplatCallback, + SplatSphericalHarmonics, + UnpackedSplat, +} from "./SplatData"; import { SplatLoader } from "./SplatLoader"; import type { SplatSource } from "./SplatMesh"; import { workerPool } from "./SplatWorker"; @@ -24,7 +29,14 @@ import { unindent, unindentLines, } from "./dyno"; -import { decodeExtSplat, encodeExtSplat, getTextureSize } from "./utils"; +import { + decodeExtSplat, + encodeExt3Rgb, + encodeExtRgb, + encodeExtSh1Rgb, + encodeExtSplat, + getTextureSize, +} from "./utils"; export type ExtSplatsOptions = { // URL to fetch a Gaussian splat file from (supports .ply, .splat, .ksplat, @@ -472,17 +484,11 @@ export class ExtSplats implements SplatSource { // Unpack the 16-byte Gsplat data at index into the Three.js components // center: THREE.Vector3, scales: THREE.Vector3, quaternion: THREE.Quaternion, // opacity: number 0..1, color: THREE.Color 0..1. - getSplat(index: number): { - center: THREE.Vector3; - scales: THREE.Vector3; - quaternion: THREE.Quaternion; - opacity: number; - color: THREE.Color; - } { + getSplat(index: number): UnpackedSplat { if (index >= this.numSplats) { throw new Error("Invalid index"); } - return decodeExtSplat(this.extArrays, index); + return decodeExtSplat(this.extArrays, index, this.extra); } // Set all ExtSplat components at index with the provided Gsplat attributes @@ -495,6 +501,7 @@ export class ExtSplats implements SplatSource { quaternion: THREE.Quaternion, opacity: number, color: THREE.Color, + sphericalHarmonics?: SplatSphericalHarmonics, ) { const extArrays = this.ensureSplats(index + 1); encodeExtSplat( @@ -515,9 +522,71 @@ export class ExtSplats implements SplatSource { color.g, color.b, ); + this.setSplatSphericalHarmonics(index, sphericalHarmonics); this.numSplats = Math.max(this.numSplats, index + 1); } + private ensureSplatsSh(key: string, numSplats: number): Uint32Array { + const current = this.extra[key] as Uint32Array | undefined; + const currentSplats = current?.length ? current.length / 4 : 0; + if (current && currentSplats >= numSplats) { + return current; + } + const targetSplats = getTextureSize( + Math.max(numSplats, 2 * currentSplats), + ).maxSplats; + const result = new Uint32Array(targetSplats * 4); + if (current) { + result.set(current); + } + this.extra[key] = result; + return result; + } + + private setSplatSphericalHarmonics( + index: number, + sphericalHarmonics?: SplatSphericalHarmonics, + ) { + if (!sphericalHarmonics) { + return; + } + if (sphericalHarmonics.sh1) { + encodeExtSh1Rgb( + this.ensureSplatsSh("sh1", index + 1), + index, + sphericalHarmonics.sh1, + ); + } + if (sphericalHarmonics.sh2) { + const sh1 = this.ensureSplatsSh("sh1", index + 1); + const sh2 = this.ensureSplatsSh("sh2", index + 1); + const base = index * 4; + sh1[base + 3] = encodeExtRgb( + sphericalHarmonics.sh2[0], + sphericalHarmonics.sh2[1], + sphericalHarmonics.sh2[2], + ); + for (let coefficient = 1; coefficient < 5; ++coefficient) { + const offset = coefficient * 3; + sh2[base + coefficient - 1] = encodeExtRgb( + sphericalHarmonics.sh2[offset], + sphericalHarmonics.sh2[offset + 1], + sphericalHarmonics.sh2[offset + 2], + ); + } + } + if (sphericalHarmonics.sh3) { + this.ensureSplatsSh("sh1", index + 1); + this.ensureSplatsSh("sh2", index + 1); + encodeExt3Rgb( + this.ensureSplatsSh("sh3a", index + 1), + this.ensureSplatsSh("sh3b", index + 1), + index, + sphericalHarmonics.sh3, + ); + } + } + // Effectively calls this.setSplat(this.numSplats++, center, ...), useful on // construction where you just want to iterate and create a collection of Gsplats. pushSplat( @@ -526,6 +595,7 @@ export class ExtSplats implements SplatSource { quaternion: THREE.Quaternion, opacity: number, color: THREE.Color, + sphericalHarmonics?: SplatSphericalHarmonics, ) { const extArrays = this.ensureSplats(this.numSplats + 1); encodeExtSplat( @@ -546,26 +616,18 @@ export class ExtSplats implements SplatSource { color.g, color.b, ); + this.setSplatSphericalHarmonics(this.numSplats, sphericalHarmonics); ++this.numSplats; } // Iterate over Gsplats index 0..=(this.numSplats-1), unpack each Gsplat // and invoke the callback function with the Gsplat attributes. - forEachSplat( - callback: ( - index: number, - center: THREE.Vector3, - scales: THREE.Vector3, - quaternion: THREE.Quaternion, - opacity: number, - color: THREE.Color, - ) => void, - ) { + forEachSplat(callback: ForEachSplatCallback) { if (!this.numSplats) { return; } for (let i = 0; i < this.numSplats; ++i) { - const unpacked = decodeExtSplat(this.extArrays, i); + const unpacked = decodeExtSplat(this.extArrays, i, this.extra); callback( i, unpacked.center, @@ -573,6 +635,7 @@ export class ExtSplats implements SplatSource { unpacked.quaternion, unpacked.opacity, unpacked.color, + unpacked.sphericalHarmonics, ); } } diff --git a/src/PackedSplats.ts b/src/PackedSplats.ts index cd363c17..6ca83926 100644 --- a/src/PackedSplats.ts +++ b/src/PackedSplats.ts @@ -2,6 +2,11 @@ import * as THREE from "three"; import { FullScreenQuad } from "three/addons/postprocessing/Pass.js"; import type { RgbaArray } from "./RgbaArray"; +import type { + ForEachSplatCallback, + SplatSphericalHarmonics, + UnpackedSplat, +} from "./SplatData"; import type { GsplatGenerator } from "./SplatGenerator"; import { SplatLoader } from "./SplatLoader"; import type { SplatSource } from "./SplatMesh"; @@ -44,7 +49,14 @@ import { splitGsplat, } from "./dyno/splats"; import { getShaders } from "./shaders"; -import { getTextureSize, setPackedSplat, unpackSplat } from "./utils"; +import { + encodeSh1Rgb, + encodeSh2Rgb, + encodeSh3Rgb, + getTextureSize, + setPackedSplat, + unpackSplat, +} from "./utils"; // Initialize a PackedSplats collection from source data via // url, fileBytes, or packedArray. Creates an empty array if none are set, @@ -527,17 +539,11 @@ export class PackedSplats implements SplatSource { // Unpack the 16-byte Gsplat data at index into the Three.js components // center: THREE.Vector3, scales: THREE.Vector3, quaternion: THREE.Quaternion, // opacity: number 0..1, color: THREE.Color 0..1. - getSplat(index: number): { - center: THREE.Vector3; - scales: THREE.Vector3; - quaternion: THREE.Quaternion; - opacity: number; - color: THREE.Color; - } { + getSplat(index: number): UnpackedSplat { if (!this.packedArray || index >= this.numSplats) { throw new Error("Invalid index"); } - return unpackSplat(this.packedArray, index, this.splatEncoding); + return unpackSplat(this.packedArray, index, this.splatEncoding, this.extra); } // Set all PackedSplat components at index with the provided Gsplat attributes @@ -550,6 +556,7 @@ export class PackedSplats implements SplatSource { quaternion: THREE.Quaternion, opacity: number, color: THREE.Color, + sphericalHarmonics?: SplatSphericalHarmonics, ) { const packedSplats = this.ensureSplats(index + 1); setPackedSplat( @@ -570,9 +577,36 @@ export class PackedSplats implements SplatSource { color.g, color.b, ); + this.setSplatSphericalHarmonics(index, sphericalHarmonics); this.numSplats = Math.max(this.numSplats, index + 1); } + private setSplatSphericalHarmonics( + index: number, + sphericalHarmonics?: SplatSphericalHarmonics, + ) { + if (!sphericalHarmonics) { + return; + } + if (sphericalHarmonics.sh1) { + const sh1 = this.ensureSplatsSh(1, index + 1); + sh1.fill(0, index * 2, index * 2 + 2); + encodeSh1Rgb(sh1, index, sphericalHarmonics.sh1, this.splatEncoding); + } + if (sphericalHarmonics.sh2) { + this.ensureSplatsSh(1, index + 1); + const sh2 = this.ensureSplatsSh(2, index + 1); + encodeSh2Rgb(sh2, index, sphericalHarmonics.sh2, this.splatEncoding); + } + if (sphericalHarmonics.sh3) { + this.ensureSplatsSh(1, index + 1); + this.ensureSplatsSh(2, index + 1); + const sh3 = this.ensureSplatsSh(3, index + 1); + sh3.fill(0, index * 4, index * 4 + 4); + encodeSh3Rgb(sh3, index, sphericalHarmonics.sh3, this.splatEncoding); + } + } + // Effectively calls this.setSplat(this.numSplats++, center, ...), useful on // construction where you just want to iterate and create a collection of Gsplats. pushSplat( @@ -581,6 +615,7 @@ export class PackedSplats implements SplatSource { quaternion: THREE.Quaternion, opacity: number, color: THREE.Color, + sphericalHarmonics?: SplatSphericalHarmonics, ) { const packedSplats = this.ensureSplats(this.numSplats + 1); setPackedSplat( @@ -601,26 +636,23 @@ export class PackedSplats implements SplatSource { color.g, color.b, ); + this.setSplatSphericalHarmonics(this.numSplats, sphericalHarmonics); ++this.numSplats; } // Iterate over Gsplats index 0..=(this.numSplats-1), unpack each Gsplat // and invoke the callback function with the Gsplat attributes. - forEachSplat( - callback: ( - index: number, - center: THREE.Vector3, - scales: THREE.Vector3, - quaternion: THREE.Quaternion, - opacity: number, - color: THREE.Color, - ) => void, - ) { + forEachSplat(callback: ForEachSplatCallback) { if (!this.packedArray || !this.numSplats) { return; } for (let i = 0; i < this.numSplats; ++i) { - const unpacked = unpackSplat(this.packedArray, i, this.splatEncoding); + const unpacked = unpackSplat( + this.packedArray, + i, + this.splatEncoding, + this.extra, + ); callback( i, unpacked.center, @@ -628,6 +660,7 @@ export class PackedSplats implements SplatSource { unpacked.quaternion, unpacked.opacity, unpacked.color, + unpacked.sphericalHarmonics, ); } } diff --git a/src/SplatData.ts b/src/SplatData.ts new file mode 100644 index 00000000..33ae76c9 --- /dev/null +++ b/src/SplatData.ts @@ -0,0 +1,32 @@ +import type * as THREE from "three"; + +/** + * Decoded RGB spherical-harmonic coefficients for a splat, grouped by band. + * Arrays contain 9, 15, and 21 floats for SH1, SH2, and SH3 respectively. + */ +export type SplatSphericalHarmonics = { + sh1?: Float32Array; + sh2?: Float32Array; + sh3?: Float32Array; +}; + +/** Decoded attributes returned by splat unpacking helpers. */ +export type UnpackedSplat = { + center: THREE.Vector3; + scales: THREE.Vector3; + quaternion: THREE.Quaternion; + opacity: number; + color: THREE.Color; + sphericalHarmonics: SplatSphericalHarmonics; +}; + +/** Callback invoked by `forEachSplat`. Values are reused between invocations. */ +export type ForEachSplatCallback = ( + index: number, + center: THREE.Vector3, + scales: THREE.Vector3, + quaternion: THREE.Quaternion, + opacity: number, + color: THREE.Color, + sphericalHarmonics: SplatSphericalHarmonics, +) => void; diff --git a/src/SplatMesh.ts b/src/SplatMesh.ts index 976f4455..e57962fb 100644 --- a/src/SplatMesh.ts +++ b/src/SplatMesh.ts @@ -1,4 +1,8 @@ import * as THREE from "three"; +import type { + ForEachSplatCallback, + SplatSphericalHarmonics, +} from "./SplatData"; import { get_raycast_buffer, @@ -182,16 +186,7 @@ export interface SplatSource { typeof Gsplat >; - forEachSplat( - callback: ( - index: number, - center: THREE.Vector3, - scales: THREE.Vector3, - quaternion: THREE.Quaternion, - opacity: number, - color: THREE.Color, - ) => void, - ): void; + forEachSplat(callback: ForEachSplatCallback): void; } export class EmptySplatSource implements SplatSource { @@ -530,11 +525,26 @@ export class SplatMesh extends SplatGenerator { quaternion: THREE.Quaternion, opacity: number, color: THREE.Color, + sphericalHarmonics?: SplatSphericalHarmonics, ) { if (this.packedSplats) { - this.packedSplats.pushSplat(center, scales, quaternion, opacity, color); + this.packedSplats.pushSplat( + center, + scales, + quaternion, + opacity, + color, + sphericalHarmonics, + ); } else if (this.extSplats) { - this.extSplats.pushSplat(center, scales, quaternion, opacity, color); + this.extSplats.pushSplat( + center, + scales, + quaternion, + opacity, + color, + sphericalHarmonics, + ); } } @@ -547,16 +557,7 @@ export class SplatMesh extends SplatGenerator { // no effect as they are decoded/unpacked copies of the underlying data. To update // the packedSplats, call .packedSplats.setSplat(index, center, scales, // quaternion, opacity, color). - forEachSplat( - callback: ( - index: number, - center: THREE.Vector3, - scales: THREE.Vector3, - quaternion: THREE.Quaternion, - opacity: number, - color: THREE.Color, - ) => void, - ) { + forEachSplat(callback: ForEachSplatCallback) { this.splats?.forEachSplat(callback); } diff --git a/src/SplatPager.ts b/src/SplatPager.ts index baa2869e..302c13e0 100644 --- a/src/SplatPager.ts +++ b/src/SplatPager.ts @@ -1,20 +1,23 @@ import * as THREE from "three"; import { decode_rad_header } from "spark-rs"; -import { LN_SCALE_MAX, LN_SCALE_MIN, dyno } from "."; import { evaluateExtSH } from "./ExtSplats"; import { evaluatePackedSH } from "./PackedSplats"; +import type { ForEachSplatCallback } from "./SplatData"; import { getSplatFileType, getSplatFileTypeFromPath } from "./SplatLoader"; import type { SplatSource } from "./SplatMesh"; import { workerPool } from "./SplatWorker"; import { DEFAULT_SPLAT_ENCODING, type ExtResult, + LN_SCALE_MAX, + LN_SCALE_MIN, type PackedResult, type RadMeta, type SplatEncoding, SplatFileType, } from "./defines"; +import * as dyno from "./dyno"; import { type DynoUsampler2DArray, pagedSplatTexCoord } from "./dyno"; import { decodeExtSplat, @@ -444,16 +447,7 @@ export class PagedSplats implements SplatSource { // Iterate over Gsplats index 0..=(this.numSplats-1), unpack each Gsplat // and invoke the callback function with the Gsplat attributes. - forEachSplat( - callback: ( - index: number, - center: THREE.Vector3, - scales: THREE.Vector3, - quaternion: THREE.Quaternion, - opacity: number, - color: THREE.Color, - ) => void, - ) { + forEachSplat(callback: ForEachSplatCallback) { if (!this.pager || !this.numSplats) { return; } @@ -467,12 +461,49 @@ export class PagedSplats implements SplatSource { packedSplatArray, extPackedSplatArray, ]; - + const shExtra = extSplats + ? { + sh1: + this.numSh >= 1 + ? this.pager.shTextures[0].value.image.data + : undefined, + sh2: + this.numSh >= 2 + ? this.pager.shTextures[1].value.image.data + : undefined, + sh3a: + this.numSh >= 3 + ? this.pager.shTextures[2].value.image.data + : undefined, + sh3b: + this.numSh >= 3 + ? this.pager.shTextures[3].value.image.data + : undefined, + } + : { + sh1: + this.numSh >= 1 + ? this.pager.shTextures[0].value.image.data + : undefined, + sh2: + this.numSh >= 2 + ? this.pager.shTextures[1].value.image.data + : undefined, + sh3: + this.numSh >= 3 + ? this.pager.shTextures[2].value.image.data + : undefined, + }; for (let i = 0; i < this.numSplats; ++i) { const splatIndex = indices[i]; const unpacked = extSplats - ? decodeExtSplat(extArrays, splatIndex) - : unpackSplat(packedSplatArray, splatIndex, this.splatEncoding); + ? decodeExtSplat(extArrays, splatIndex, shExtra) + : unpackSplat( + packedSplatArray, + splatIndex, + this.splatEncoding, + shExtra, + ); callback( i, unpacked.center, @@ -480,6 +511,7 @@ export class PagedSplats implements SplatSource { unpacked.quaternion, unpacked.opacity, unpacked.color, + unpacked.sphericalHarmonics, ); } } diff --git a/src/index.ts b/src/index.ts index 2ce39cc9..5645b341 100644 --- a/src/index.ts +++ b/src/index.ts @@ -7,6 +7,11 @@ export { SplatAccumulator, type GeneratorMapping } from "./SplatAccumulator"; export * as dyno from "./dyno"; export { RgbaArray, readRgbaArray } from "./RgbaArray"; +export type { + ForEachSplatCallback, + SplatSphericalHarmonics, + UnpackedSplat, +} from "./SplatData"; export { SplatLoader, diff --git a/src/utils.ts b/src/utils.ts index 9c0afad8..27862e7b 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -3,6 +3,7 @@ import * as THREE from "three"; // Miscellaneous utility functions for Spark +import type { SplatSphericalHarmonics, UnpackedSplat } from "./SplatData.js"; import { LN_SCALE_MAX, LN_SCALE_MIN, @@ -430,13 +431,8 @@ export function encodeExtSplat( export function decodeExtSplat( extArrays: [Uint32Array, Uint32Array], index: number, -): { - center: THREE.Vector3; - scales: THREE.Vector3; - quaternion: THREE.Quaternion; - color: THREE.Color; - opacity: number; -} { + extra?: Record, +): UnpackedSplat { // Returns a static object which is reused each time const result = packedFields; const i4 = index * 4; @@ -452,6 +448,7 @@ export function decodeExtSplat( result.scales.y = Math.exp(fromHalf(extB[i4 + 2] & 0xffff)); result.scales.z = Math.exp(fromHalf(extB[i4 + 2] >>> 16)); decodeQuatOctXy1010R12(extB[i4 + 3], result.quaternion); + decodeExtSphericalHarmonics(extra ?? {}, index, result.sphericalHarmonics); return result; } @@ -723,6 +720,7 @@ const packedFields = { quaternion: packedQuaternion, color: packedColor, opacity: 0.0, + sphericalHarmonics: {}, }; // Unpack all components of a PackedSplat from the packedSplats Uint32Array into @@ -736,14 +734,12 @@ export function unpackSplat( lnScaleMin?: number; lnScaleMax?: number; lodOpacity?: boolean; + sh1Max?: number; + sh2Max?: number; + sh3Max?: number; }, -): { - center: THREE.Vector3; - scales: THREE.Vector3; - quaternion: THREE.Quaternion; - color: THREE.Color; - opacity: number; -} { + extra?: Record, +): UnpackedSplat { // Returns a static object which is reused each time const result = packedFields; @@ -789,6 +785,12 @@ export function unpackSplat( // decodeQuatXyz888(uQuat, result.quaternion); // decodeQuatEulerXyz888(uQuat, result.quaternion); + decodePackedSphericalHarmonics( + extra ?? {}, + index, + encoding, + result.sphericalHarmonics, + ); return result; } @@ -1483,6 +1485,89 @@ export function encodeSh3Rgb( } } +function decodeSignedBits( + words: Uint32Array, + wordBase: number, + bitStart: number, + numBits: number, +): number { + const wordOffset = bitStart >>> 5; + const bitOffset = bitStart & 31; + let value = words[wordBase + wordOffset] >>> bitOffset; + if (bitOffset + numBits > 32) { + value |= words[wordBase + wordOffset + 1] << (32 - bitOffset); + } + const shift = 32 - numBits; + return (value << shift) >> shift; +} + +function ensureShBand( + target: SplatSphericalHarmonics, + band: keyof SplatSphericalHarmonics, + length: number, +): Float32Array { + const current = target[band]; + if (current?.length === length) { + return current; + } + const result = new Float32Array(length); + target[band] = result; + return result; +} + +/** Decode one splat's compact SH1..SH3 data into reusable float arrays. */ +function decodePackedSphericalHarmonics( + extra: Record, + index: number, + encoding?: { + sh1Max?: number; + sh2Max?: number; + sh3Max?: number; + }, + target: SplatSphericalHarmonics = {}, +): SplatSphericalHarmonics { + const sh1 = extra.sh1; + if (!(sh1 instanceof Uint32Array)) { + target.sh1 = undefined; + target.sh2 = undefined; + target.sh3 = undefined; + return target; + } + + const decodedSh1 = ensureShBand(target, "sh1", 9); + for (let i = 0; i < decodedSh1.length; ++i) { + decodedSh1[i] = + decodeSignedBits(sh1, index * 2, i * 7, 7) * + ((encoding?.sh1Max ?? 1) / 63); + } + + const sh2 = extra.sh2; + if (!(sh2 instanceof Uint32Array)) { + target.sh2 = undefined; + target.sh3 = undefined; + return target; + } + const decodedSh2 = ensureShBand(target, "sh2", 15); + for (let i = 0; i < decodedSh2.length; ++i) { + decodedSh2[i] = + decodeSignedBits(sh2, index * 4, i * 8, 8) * + ((encoding?.sh2Max ?? 1) / 127); + } + + const sh3 = extra.sh3; + if (!(sh3 instanceof Uint32Array)) { + target.sh3 = undefined; + return target; + } + const decodedSh3 = ensureShBand(target, "sh3", 21); + for (let i = 0; i < decodedSh3.length; ++i) { + decodedSh3[i] = + decodeSignedBits(sh3, index * 4, i * 6, 6) * + ((encoding?.sh3Max ?? 1) / 31); + } + return target; +} + export function encodeExtRgb(r: number, g: number, b: number): number { const ar = Math.abs(r); const ag = Math.abs(g); @@ -1513,6 +1598,65 @@ export function decodeExtRgb(encoded: number): THREE.Color { return color; } +function decodeExtRgbInto( + encoded: number, + target: Float32Array, + offset: number, +) { + const color = decodeExtRgb(encoded); + target[offset] = color.r; + target[offset + 1] = color.g; + target[offset + 2] = color.b; +} + +/** Decode one splat's extended SH1..SH3 data into reusable float arrays. */ +function decodeExtSphericalHarmonics( + extra: Record, + index: number, + target: SplatSphericalHarmonics = {}, +): SplatSphericalHarmonics { + const sh1 = extra.sh1; + if (!(sh1 instanceof Uint32Array)) { + target.sh1 = undefined; + target.sh2 = undefined; + target.sh3 = undefined; + return target; + } + + const base = index * 4; + const decodedSh1 = ensureShBand(target, "sh1", 9); + for (let i = 0; i < 3; ++i) { + decodeExtRgbInto(sh1[base + i], decodedSh1, i * 3); + } + + const sh2 = extra.sh2; + if (!(sh2 instanceof Uint32Array)) { + target.sh2 = undefined; + target.sh3 = undefined; + return target; + } + const decodedSh2 = ensureShBand(target, "sh2", 15); + decodeExtRgbInto(sh1[base + 3], decodedSh2, 0); + for (let i = 0; i < 4; ++i) { + decodeExtRgbInto(sh2[base + i], decodedSh2, (i + 1) * 3); + } + + const sh3a = extra.sh3a; + const sh3b = extra.sh3b; + if (!(sh3a instanceof Uint32Array) || !(sh3b instanceof Uint32Array)) { + target.sh3 = undefined; + return target; + } + const decodedSh3 = ensureShBand(target, "sh3", 21); + for (let i = 0; i < 4; ++i) { + decodeExtRgbInto(sh3a[base + i], decodedSh3, i * 3); + } + for (let i = 0; i < 3; ++i) { + decodeExtRgbInto(sh3b[base + i], decodedSh3, (i + 4) * 3); + } + return target; +} + export function encodeExtSh1Rgb( sh1Array: Uint32Array, index: number, @@ -1539,11 +1683,11 @@ export function encodeExtSh12Rgb( } sh1Array[i4 + 3] = encodeExtRgb(sh2Rgb[0], sh2Rgb[1], sh2Rgb[2]); for (let k = 1; k < 5; ++k) { - const k5 = k * 5; + const k3 = k * 3; sh2Array[i4 + (k - 1)] = encodeExtRgb( - sh2Rgb[k5], - sh2Rgb[k5 + 1], - sh2Rgb[k5 + 2], + sh2Rgb[k3], + sh2Rgb[k3 + 1], + sh2Rgb[k3 + 2], ); } } diff --git a/test/spherical-harmonics.test.ts b/test/spherical-harmonics.test.ts new file mode 100644 index 00000000..62dca881 --- /dev/null +++ b/test/spherical-harmonics.test.ts @@ -0,0 +1,221 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import * as THREE from "three"; +import { ExtSplats } from "../src/ExtSplats.js"; +import { PackedSplats } from "../src/PackedSplats.js"; +import { + decodeExtSplat, + encodeExt3Rgb, + encodeExtSh12Rgb, + encodeSh1Rgb, + encodeSh2Rgb, + encodeSh3Rgb, + unpackSplat, +} from "../src/utils.js"; + +function coefficients(length: number, scale = 1) { + return Float32Array.from( + { length }, + (_, index) => (((index * 7) % 13) / 6 - 1) * scale, + ); +} + +function extendedCoefficients(length: number) { + const componentScale = [1, -0.5, 0.25]; + return Float32Array.from({ length }, (_, index) => { + const coefficient = Math.floor(index / 3); + return 2 ** ((coefficient % 4) - 1) * componentScale[index % 3]; + }); +} + +function assertArraysClose( + actual: Float32Array | undefined, + expected: Float32Array, + tolerance: number, +) { + assert.ok(actual); + assert.equal(actual.length, expected.length); + for (let i = 0; i < expected.length; ++i) { + assert.ok( + Math.abs(actual[i] - expected[i]) <= tolerance, + `coefficient ${i}: expected ${expected[i]}, got ${actual[i]}`, + ); + } +} + +test("decodes packed SH1 through SH3", () => { + const encoding = { sh1Max: 2, sh2Max: 3, sh3Max: 4 }; + const sh1 = coefficients(9, encoding.sh1Max); + const sh2 = coefficients(15, encoding.sh2Max); + const sh3 = coefficients(21, encoding.sh3Max); + const extra = { + sh1: new Uint32Array(2), + sh2: new Uint32Array(4), + sh3: new Uint32Array(4), + }; + encodeSh1Rgb(extra.sh1, 0, sh1, encoding); + encodeSh2Rgb(extra.sh2, 0, sh2, encoding); + encodeSh3Rgb(extra.sh3, 0, sh3, encoding); + + const decoded = unpackSplat( + new Uint32Array(4), + 0, + encoding, + extra, + ).sphericalHarmonics; + assertArraysClose(decoded.sh1, sh1, encoding.sh1Max / 63 + 1e-6); + assertArraysClose(decoded.sh2, sh2, encoding.sh2Max / 127 + 1e-6); + assertArraysClose(decoded.sh3, sh3, encoding.sh3Max / 31 + 1e-6); +}); + +test("decodes extended SH1 through SH3", () => { + const sh1 = extendedCoefficients(9); + const sh2 = extendedCoefficients(15); + const sh3 = extendedCoefficients(21); + const extra = { + sh1: new Uint32Array(4), + sh2: new Uint32Array(4), + sh3a: new Uint32Array(4), + sh3b: new Uint32Array(4), + }; + encodeExtSh12Rgb(extra.sh1, extra.sh2, 0, sh1, sh2); + encodeExt3Rgb(extra.sh3a, extra.sh3b, 0, sh3); + + const decoded = decodeExtSplat( + [new Uint32Array(4), new Uint32Array(4)], + 0, + extra, + ).sphericalHarmonics; + assertArraysClose(decoded.sh1, sh1, 0.02); + assertArraysClose(decoded.sh2, sh2, 0.02); + assertArraysClose(decoded.sh3, sh3, 0.02); +}); + +test("unpackSplat omits unavailable SH bands", () => { + const decoded = unpackSplat( + new Uint32Array(4), + 0, + undefined, + {}, + ).sphericalHarmonics; + assert.deepEqual(decoded, { + sh1: undefined, + sh2: undefined, + sh3: undefined, + }); +}); + +const center = new THREE.Vector3(1, 2, 3); +const scales = new THREE.Vector3(0.1, 0.2, 0.3); +const quaternion = new THREE.Quaternion(); +const color = new THREE.Color(0.25, 0.5, 0.75); + +test("PackedSplats writes callback SH changes back through setSplat", () => { + const encoding = { sh1Max: 2, sh2Max: 3, sh3Max: 4 }; + const splats = new PackedSplats({ maxSplats: 1, splatEncoding: encoding }); + splats.pushSplat(center, scales, quaternion, 0.8, color, { + sh1: coefficients(9, encoding.sh1Max), + sh2: coefficients(15, encoding.sh2Max), + sh3: coefficients(21, encoding.sh3Max), + }); + + splats.forEachSplat( + (index, nextCenter, nextScales, nextQuaternion, opacity, nextColor, sh) => { + sh.sh1?.fill(0); + sh.sh2?.fill(0); + sh.sh3?.fill(0); + splats.setSplat( + index, + nextCenter, + nextScales, + nextQuaternion, + opacity, + nextColor, + sh, + ); + }, + ); + + const sh = splats.getSplat(0).sphericalHarmonics; + assert.deepEqual(sh.sh1, new Float32Array(9)); + assert.deepEqual(sh.sh2, new Float32Array(15)); + assert.deepEqual(sh.sh3, new Float32Array(21)); +}); + +test("PackedSplats preserves SH bands omitted from setSplat", () => { + const splats = new PackedSplats({ maxSplats: 1 }); + const originalSh2 = coefficients(15); + splats.pushSplat(center, scales, quaternion, 0.8, color, { + sh1: coefficients(9), + sh2: originalSh2, + }); + const unpacked = splats.getSplat(0); + splats.setSplat( + 0, + unpacked.center, + unpacked.scales, + unpacked.quaternion, + unpacked.opacity, + unpacked.color, + { sh1: new Float32Array(9) }, + ); + + assertArraysClose( + splats.getSplat(0).sphericalHarmonics.sh2, + originalSh2, + 1 / 127 + 1e-6, + ); +}); + +test("ExtSplats writes callback SH changes back through setSplat", () => { + const splats = new ExtSplats({ maxSplats: 1 }); + splats.pushSplat(center, scales, quaternion, 0.8, color, { + sh1: extendedCoefficients(9), + sh2: extendedCoefficients(15), + sh3: extendedCoefficients(21), + }); + + splats.forEachSplat( + (index, nextCenter, nextScales, nextQuaternion, opacity, nextColor, sh) => { + sh.sh1?.fill(0); + sh.sh2?.fill(0); + sh.sh3?.fill(0); + splats.setSplat( + index, + nextCenter, + nextScales, + nextQuaternion, + opacity, + nextColor, + sh, + ); + }, + ); + + const sh = splats.getSplat(0).sphericalHarmonics; + assert.deepEqual(sh.sh1, new Float32Array(9)); + assert.deepEqual(sh.sh2, new Float32Array(15)); + assert.deepEqual(sh.sh3, new Float32Array(21)); +}); + +test("pushSplat creates missing lower SH bands", () => { + const packed = new PackedSplats({ maxSplats: 1 }); + packed.pushSplat(center, scales, quaternion, 0.8, color, { + sh3: coefficients(21), + }); + assert.equal(packed.getNumSh(), 3); + assert.deepEqual( + packed.getSplat(0).sphericalHarmonics.sh1, + new Float32Array(9), + ); + + const extended = new ExtSplats({ maxSplats: 1 }); + extended.pushSplat(center, scales, quaternion, 0.8, color, { + sh3: extendedCoefficients(21), + }); + assert.equal(extended.getNumSh(), 3); + assert.deepEqual( + extended.getSplat(0).sphericalHarmonics.sh2, + new Float32Array(15), + ); +}); diff --git a/test/wasm-loader.mjs b/test/wasm-loader.mjs new file mode 100644 index 00000000..edee572b --- /dev/null +++ b/test/wasm-loader.mjs @@ -0,0 +1,35 @@ +export async function resolve(specifier, context, nextResolve) { + if (specifier.includes("?worker")) { + return { + shortCircuit: true, + url: "data:text/javascript,export default class TestWorker {}", + }; + } + return nextResolve(specifier, context); +} + +export async function load(url, context, nextLoad) { + if (url.endsWith(".wasm?arraybuffer&base64")) { + const wasmUrl = url.slice(0, url.indexOf("?")); + return { + format: "module", + shortCircuit: true, + source: ` + import { readFileSync } from "node:fs"; + const bytes = readFileSync(new URL(${JSON.stringify(wasmUrl)})); + export default Uint8Array.from(bytes).buffer; + `, + }; + } + if (url.endsWith(".glsl")) { + return { + format: "module", + shortCircuit: true, + source: ` + import { readFileSync } from "node:fs"; + export default readFileSync(new URL(${JSON.stringify(url)}), "utf8"); + `, + }; + } + return nextLoad(url, context); +}