From 738d7c399eedcdba10a4494146295464a634fbc3 Mon Sep 17 00:00:00 2001 From: Noeri Huisman <8823461+mrxz@users.noreply.github.com> Date: Fri, 18 Jul 2025 17:29:32 +0200 Subject: [PATCH] Expand splat represention to 32 bytes using two array texture targets --- package-lock.json | 7 +- package.json | 6 +- src/PackedSplats.ts | 99 ++++++++++++++++--------- src/SparkRenderer.ts | 10 ++- src/SplatLoader.ts | 22 ++++-- src/SplatMesh.ts | 2 +- src/antisplat.ts | 10 ++- src/dyno/output.ts | 9 ++- src/dyno/splats.ts | 10 ++- src/ksplat.ts | 7 +- src/pcsogs.ts | 9 ++- src/shaders/computeUvec4x2.glsl | 38 ++++++++++ src/shaders/splatDefines.glsl | 48 ++++++------ src/shaders/splatFragment.glsl | 2 +- src/shaders/splatVertex.glsl | 6 +- src/spz.ts | 2 +- src/utils.ts | 126 +++++++++++++++----------------- src/worker.ts | 18 +++-- 18 files changed, 266 insertions(+), 165 deletions(-) create mode 100644 src/shaders/computeUvec4x2.glsl diff --git a/package-lock.json b/package-lock.json index c312171..0394832 100644 --- a/package-lock.json +++ b/package-lock.json @@ -19,7 +19,7 @@ "onchange": "7.1.0", "spark-internal-rs": "file:rust/spark-internal-rs/pkg", "stats.js": "^0.17.0", - "three": "^0.172.0", + "three": "git+https://github.com/mrxz/three.js.git#array-texture-mrt-build", "ts-node": "10.9.2", "typescript": "^5.7.3", "vite": "^6.0.11", @@ -2623,9 +2623,8 @@ } }, "node_modules/three": { - "version": "0.172.0", - "resolved": "https://registry.npmjs.org/three/-/three-0.172.0.tgz", - "integrity": "sha512-6HMgMlzU97MsV7D/tY8Va38b83kz8YJX+BefKjspMNAv0Vx6dxMogHOrnRl/sbMIs3BPUKijPqDqJ/+UwJbIow==", + "version": "0.178.0", + "resolved": "git+ssh://git@github.com/mrxz/three.js.git#bc2d35e0b7c33c4fd9d3f7ef58a6765f9527adbf", "dev": true, "license": "MIT" }, diff --git a/package.json b/package.json index 743d283..27536ed 100644 --- a/package.json +++ b/package.json @@ -48,14 +48,14 @@ "lefthook": "1.11.12", "lil-gui": "^0.20.0", "onchange": "7.1.0", + "spark-internal-rs": "file:rust/spark-internal-rs/pkg", "stats.js": "^0.17.0", - "three": "^0.172.0", + "three": "git+https://github.com/mrxz/three.js.git#array-texture-mrt-build", "ts-node": "10.9.2", "typescript": "^5.7.3", "vite": "^6.0.11", "vite-plugin-dts": "^4.5.4", - "vite-plugin-glsl": "^1.3.1", - "spark-internal-rs": "file:rust/spark-internal-rs/pkg" + "vite-plugin-glsl": "^1.3.1" }, "dependencies": { "fflate": "^0.8.2" diff --git a/src/PackedSplats.ts b/src/PackedSplats.ts index 4cc8387..b2c8380 100644 --- a/src/PackedSplats.ts +++ b/src/PackedSplats.ts @@ -12,7 +12,7 @@ import { outputPackedSplat, } from "./dyno"; import { TPackedSplats, definePackedSplats } from "./dyno/splats"; -import computeUvec4Template from "./shaders/computeUvec4.glsl"; +import computeUvec4x2Template from "./shaders/computeUvec4x2.glsl"; import { getTextureSize, setPackedSplat, unpackSplat } from "./utils"; // Initialize a PackedSplats collection from source data via @@ -38,7 +38,7 @@ export type PackedSplatsOptions = { maxSplats?: number; // Use provided packed data array, where each 4 consecutive uint32 values // encode one "packed" Gsplat. (default: undefined) - packedArray?: Uint32Array; + packedArray?: [Uint32Array, Uint32Array]; // Override number of splats in packed array to use only a subset. // (default: length of packed array / 4) numSplats?: number; @@ -59,7 +59,7 @@ export type PackedSplatsOptions = { export class PackedSplats { maxSplats = 0; numSplats = 0; - packedArray: Uint32Array | null = null; + packedArray: [Uint32Array, Uint32Array] | null = null; extra: Record; initialized: Promise; @@ -68,7 +68,7 @@ export class PackedSplats { // Either target or source will be non-null, depending on whether the PackedSplats // is being used as a data source or generated to. target: THREE.WebGLArrayRenderTarget | null = null; - source: THREE.DataArrayTexture | null = null; + source: THREE.DataArrayTexture[] | null = null; // Set to true if source packedArray is updated to have it upload to GPU needsUpdate = true; @@ -106,7 +106,7 @@ export class PackedSplats { // Calculate number of horizontal texture rows that could fit in array. // A properly initialized packedArray should already take into account the // width and height of the texture and be rounded up with padding. - this.maxSplats = Math.floor(this.packedArray.length / 4); + this.maxSplats = Math.floor(this.packedArray[0].length / 4); this.maxSplats = Math.floor(this.maxSplats / SPLAT_TEX_WIDTH) * SPLAT_TEX_WIDTH; this.numSplats = Math.min( @@ -152,7 +152,8 @@ export class PackedSplats { this.target = null; } if (this.source) { - this.source.dispose(); + this.source[0].dispose(); + this.source[1].dispose(); this.source = null; } } @@ -163,22 +164,26 @@ export class PackedSplats { // 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. - ensureSplats(numSplats: number): Uint32Array { + ensureSplats(numSplats: number): [Uint32Array, Uint32Array] { const targetSize = numSplats <= this.maxSplats ? this.maxSplats : // Grow exponentially to avoid frequent reallocations Math.max(numSplats, 2 * this.maxSplats); - const currentSize = !this.packedArray ? 0 : this.packedArray.length / 4; + const currentSize = !this.packedArray ? 0 : this.packedArray[0].length / 4; if (!this.packedArray || targetSize > currentSize) { this.maxSplats = getTextureSize(targetSize).maxSplats; - const newArray = new Uint32Array(this.maxSplats * 4); + const newArrays: [Uint32Array, Uint32Array] = [ + new Uint32Array(this.maxSplats * 4), + new Uint32Array(this.maxSplats * 4), + ]; if (this.packedArray) { // Copy over existing data - newArray.set(this.packedArray); + newArrays[0].set(this.packedArray[0]); + newArrays[1].set(this.packedArray[1]); } - this.packedArray = newArray; + this.packedArray = newArrays; } return this.packedArray; } @@ -188,7 +193,7 @@ export class PackedSplats { let wordsPerSplat: number; let key: string; if (level === 0) { - return this.ensureSplats(numSplats); + return this.ensureSplats(numSplats)[0]; // FIXME } if (level === 1) { // 3 x 3 uint7 = 63 bits = 2 uint32 @@ -359,6 +364,7 @@ export class PackedSplats { this.target.texture.type = THREE.UnsignedIntType; this.target.texture.internalFormat = "RGBA32UI"; this.target.scissorTest = true; + this.target.textures = [this.target.texture, this.target.texture.clone()]; return true; } @@ -372,7 +378,7 @@ export class PackedSplats { let maxSplats = 0; const mapping = splatCounts.map((numSplats) => { const base = maxSplats; - // Generation happens in horizonal row chunks, so round up to full width + // Generation happens in horizontal row chunks, so round up to full width const rounded = Math.ceil(numSplats / SPLAT_TEX_WIDTH) * SPLAT_TEX_WIDTH; maxSplats += rounded; return { base, count: numSplats }; @@ -382,10 +388,10 @@ export class PackedSplats { // Returns a THREE.DataArrayTexture representing the PackedSplats content as // a Uint32x4 data array texture (2048 x 2048 x depth in size) - getTexture(): THREE.DataArrayTexture { + getTexture(): THREE.DataArrayTexture[] { if (this.target) { // Return the render target's texture - return this.target.texture; + return this.target.textures; } if (this.source || this.packedArray) { // Update source texture if needed and return @@ -397,7 +403,7 @@ export class PackedSplats { } // Check if source texture needs to be created/updated - private maybeUpdateSource(): THREE.DataArrayTexture { + private maybeUpdateSource(): THREE.DataArrayTexture[] { if (!this.packedArray) { throw new Error("No packed splats"); } @@ -406,32 +412,55 @@ export class PackedSplats { this.needsUpdate = false; if (this.source) { - const { width, height, depth } = this.source.image; + const { width, height, depth } = this.source[0].image; if (this.maxSplats !== width * height * depth) { // The existing source texture isn't the right size, so dispose it - this.source.dispose(); + this.source[0].dispose(); + this.source[1].dispose(); this.source = null; } } if (!this.source) { // Allocate a new source texture of the right size const { width, height, depth } = getTextureSize(this.maxSplats); - this.source = new THREE.DataArrayTexture( - this.packedArray, + const sourceTexture = new THREE.DataArrayTexture( + this.packedArray[0], width, height, depth, ); - this.source.format = THREE.RGBAIntegerFormat; - this.source.type = THREE.UnsignedIntType; - this.source.internalFormat = "RGBA32UI"; - this.source.needsUpdate = true; - } else if (this.packedArray.buffer !== this.source.image.data.buffer) { - // The source texture is the right size, update the data - this.source.image.data = new Uint8Array(this.packedArray.buffer); + sourceTexture.format = THREE.RGBAIntegerFormat; + sourceTexture.type = THREE.UnsignedIntType; + sourceTexture.internalFormat = "RGBA32UI"; + sourceTexture.needsUpdate = true; + + const sourceTexture2 = new THREE.DataArrayTexture( + this.packedArray[1], + width, + height, + depth, + ); + sourceTexture2.format = THREE.RGBAIntegerFormat; + sourceTexture2.type = THREE.UnsignedIntType; + sourceTexture2.internalFormat = "RGBA32UI"; + sourceTexture2.needsUpdate = true; + + this.source = [sourceTexture, sourceTexture2]; + } else { + if (this.packedArray[0].buffer !== this.source[0].image.data.buffer) { + this.source[0].image.data = new Uint8Array( + this.packedArray[0].buffer, + ); + } + if (this.packedArray[1].buffer !== this.source[1].image.data.buffer) { + this.source[1].image.data = new Uint8Array( + this.packedArray[1].buffer, + ); + } } // Indicate to Three.js that the source texture needs to be uploaded to the GPU - this.source.needsUpdate = true; + this.source[0].needsUpdate = true; + this.source[1].needsUpdate = true; } return this.source; } @@ -440,7 +469,7 @@ export class PackedSplats { // Can be used where you need an uninitialized THREE.DataArrayTexture like // a uniform you will update with the result of this.getTexture() later. - static getEmpty(): THREE.DataArrayTexture { + static getEmpty(): THREE.DataArrayTexture[] { if (!PackedSplats.emptySource) { const { width, height, depth, maxSplats } = getTextureSize(1); const emptyArray = new Uint32Array(maxSplats * 4); @@ -455,7 +484,7 @@ export class PackedSplats { PackedSplats.emptySource.internalFormat = "RGBA32UI"; PackedSplats.emptySource.needsUpdate = true; } - return PackedSplats.emptySource; + return [PackedSplats.emptySource, PackedSplats.emptySource]; } // Get a program and THREE.RawShaderMaterial for a given GsplatGenerator, @@ -479,7 +508,7 @@ export class PackedSplats { ); if (!PackedSplats.programTemplate) { PackedSplats.programTemplate = new DynoProgramTemplate( - computeUvec4Template, + computeUvec4x2Template, ); } // Create a program from the template and graph @@ -615,6 +644,7 @@ export class DynoPackedSplats extends DynoUniform< "packedSplats", { texture: THREE.DataArrayTexture; + texture2: THREE.DataArrayTexture; numSplats: number; } > { @@ -626,12 +656,15 @@ export class DynoPackedSplats extends DynoUniform< type: TPackedSplats, globals: () => [definePackedSplats], value: { - texture: PackedSplats.getEmpty(), + texture: PackedSplats.getEmpty()[0], + texture2: PackedSplats.getEmpty()[1], numSplats: 0, }, update: (value) => { value.texture = - this.packedSplats?.getTexture() ?? PackedSplats.getEmpty(); + this.packedSplats?.getTexture()?.[0] ?? PackedSplats.getEmpty()[0]; + value.texture2 = + this.packedSplats?.getTexture()?.[1] ?? PackedSplats.getEmpty()[1]; value.numSplats = this.packedSplats?.numSplats ?? 0; return value; }, diff --git a/src/SparkRenderer.ts b/src/SparkRenderer.ts index 18da9a0..1fe774a 100644 --- a/src/SparkRenderer.ts +++ b/src/SparkRenderer.ts @@ -265,6 +265,7 @@ export class SparkRenderer extends THREE.Mesh { uniforms, transparent: true, blending: THREE.NormalBlending, + premultipliedAlpha: true, depthTest: true, depthWrite: false, side: THREE.DoubleSide, @@ -382,7 +383,8 @@ export class SparkRenderer extends THREE.Mesh { // Splat texture mid plane distance, or 0.0 to disable splatTexMid: { value: 0.0 }, // Gsplat collection to render - packedSplats: { type: "t", value: PackedSplats.getEmpty() }, + packedSplats: { type: "t", value: PackedSplats.getEmpty()[0] }, + packedSplats2: { type: "t", value: PackedSplats.getEmpty()[1] }, // Time in seconds for time-based effects time: { value: 0 }, // Delta time in seconds since last frame @@ -577,12 +579,14 @@ export class SparkRenderer extends THREE.Mesh { if (this.viewpoint.display) { const { accumulator, geometry } = this.viewpoint.display; this.uniforms.numSplats.value = accumulator.splats.numSplats; - this.uniforms.packedSplats.value = accumulator.splats.getTexture(); + this.uniforms.packedSplats.value = accumulator.splats.getTexture()[0]; + this.uniforms.packedSplats2.value = accumulator.splats.getTexture()[1]; this.geometry = geometry; } else { // No Gsplats to display for this viewpoint yet this.uniforms.numSplats.value = 0; - this.uniforms.packedSplats.value = PackedSplats.getEmpty(); + this.uniforms.packedSplats.value = PackedSplats.getEmpty()[0]; + this.uniforms.packedSplats2.value = PackedSplats.getEmpty()[1]; this.geometry = EMPTY_GEOMETRY; } } diff --git a/src/SplatLoader.ts b/src/SplatLoader.ts index d77bde7..5448a8c 100644 --- a/src/SplatLoader.ts +++ b/src/SplatLoader.ts @@ -402,7 +402,7 @@ export async function unpackSplats({ fileType?: SplatFileType; pathOrUrl?: string; }): Promise<{ - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra?: Record; }> { @@ -422,13 +422,19 @@ export async function unpackSplats({ await ply.parseHeader(); const numSplats = ply.numSplats; const maxSplats = getTextureSize(numSplats).maxSplats; - const args = { fileBytes, packedArray: new Uint32Array(maxSplats * 4) }; + const args = { + fileBytes, + packedArray: [ + new Uint32Array(maxSplats * 4), + new Uint32Array(maxSplats * 4), + ], + }; return await withWorker(async (worker) => { const { packedArray, numSplats, extra } = (await worker.call( "unpackPly", args, )) as { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }; @@ -443,7 +449,7 @@ export async function unpackSplats({ fileBytes, }, )) as { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }; @@ -457,7 +463,7 @@ export async function unpackSplats({ { fileBytes, }, - )) as { packedArray: Uint32Array; numSplats: number }; + )) as { packedArray: [Uint32Array, Uint32Array]; numSplats: number }; return { packedArray, numSplats }; }); } @@ -467,7 +473,7 @@ export async function unpackSplats({ "decodeKsplat", { fileBytes }, )) as { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }; @@ -480,7 +486,7 @@ export async function unpackSplats({ "decodePcSogs", { fileBytes, extraFiles }, )) as { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }; @@ -493,7 +499,7 @@ export async function unpackSplats({ "decodePcSogsZip", { fileBytes }, )) as { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }; diff --git a/src/SplatMesh.ts b/src/SplatMesh.ts index e34579f..a1023a9 100644 --- a/src/SplatMesh.ts +++ b/src/SplatMesh.ts @@ -537,7 +537,7 @@ export class SplatMesh extends SplatGenerator { near, far, this.packedSplats.numSplats, - this.packedSplats.packedArray, + this.packedSplats.packedArray[0], // FIXME RAYCAST_ELLIPSOID, ); diff --git a/src/antisplat.ts b/src/antisplat.ts index d24017d..49cd429 100644 --- a/src/antisplat.ts +++ b/src/antisplat.ts @@ -66,18 +66,22 @@ export function decodeAntiSplat( } export function unpackAntiSplat(fileBytes: Uint8Array): { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; } { let numSplats = 0; let maxSplats = 0; - let packedArray = new Uint32Array(0); + const emptyArray = new Uint32Array(0); + let packedArray: [Uint32Array, Uint32Array] = [emptyArray, emptyArray]; decodeAntiSplat( fileBytes, (cbNumSplats) => { numSplats = cbNumSplats; maxSplats = computeMaxSplats(numSplats); - packedArray = new Uint32Array(maxSplats * 4); + packedArray = [ + new Uint32Array(maxSplats * 4), + new Uint32Array(maxSplats * 4), + ]; }, ( index, diff --git a/src/dyno/output.ts b/src/dyno/output.ts index 494295b..e995ba6 100644 --- a/src/dyno/output.ts +++ b/src/dyno/output.ts @@ -30,13 +30,18 @@ export class OutputPackedSplat if (gsplat) { return unindentLines(` if (isGsplatActive(${gsplat}.flags)) { - ${output} = packSplat(${gsplat}.center, ${gsplat}.scales, ${gsplat}.quaternion, ${gsplat}.rgba); + uvec4[2] packed = packSplat(${gsplat}.center, ${gsplat}.scales, ${gsplat}.quaternion, ${gsplat}.rgba); + ${output} = packed[0]; + ${output}2 = packed[1]; } else { ${output} = uvec4(0u, 0u, 0u, 0u); + ${output}2 = uvec4(0u, 0u, 0u, 0u); } `); } - return [`${output} = uvec4(0u, 0u, 0u, 0u);`]; + return [ + `${output} = uvec4(0u, 0u, 0u, 0u); ${output}2 = uvec4(0u, 0u, 0u, 0u);`, + ]; }, }); } diff --git a/src/dyno/splats.ts b/src/dyno/splats.ts index 51abe39..a0ee9ff 100644 --- a/src/dyno/splats.ts +++ b/src/dyno/splats.ts @@ -117,6 +117,7 @@ export const defineGsplat = unindent(` export const definePackedSplats = unindent(` struct PackedSplats { usampler2DArray texture; + usampler2DArray texture2; int numSplats; }; `); @@ -137,9 +138,12 @@ export class NumPackedSplats extends UnaryOp< } const defineReadPackedSplat = unindent(` - bool readPackedSplat(usampler2DArray texture, int numSplats, int index, out Gsplat gsplat) { + bool readPackedSplat(usampler2DArray texture, usampler2DArray texture2, int numSplats, int index, out Gsplat gsplat) { if ((index >= 0) && (index < numSplats)) { - uvec4 packed = texelFetch(texture, splatTexCoord(index), 0); + uvec4[2] packed = uvec4[2]( + texelFetch(texture, splatTexCoord(index), 0), + texelFetch(texture2, splatTexCoord(index), 0) + ); unpackSplat(packed, gsplat.center, gsplat.scales, gsplat.quaternion, gsplat.rgba); return true; } else { @@ -173,7 +177,7 @@ export class ReadPackedSplat let statements: string[]; if (packedSplats && index) { statements = unindentLines(` - if (readPackedSplat(${packedSplats}.texture, ${packedSplats}.numSplats, ${index}, ${gsplat})) { + if (readPackedSplat(${packedSplats}.texture, ${packedSplats}.texture2, ${packedSplats}.numSplats, ${index}, ${gsplat})) { bool zeroSize = all(equal(${gsplat}.scales, vec3(0.0, 0.0, 0.0))); ${gsplat}.flags = zeroSize ? 0u : GSPLAT_FLAG_ACTIVE; } else { diff --git a/src/ksplat.ts b/src/ksplat.ts index 33a892c..0133019 100644 --- a/src/ksplat.ts +++ b/src/ksplat.ts @@ -353,7 +353,7 @@ export function decodeKsplat( } export function unpackKsplat(fileBytes: Uint8Array): { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; } { @@ -387,7 +387,10 @@ export function unpackKsplat(fileBytes: Uint8Array): { const numSplats = splatCount; const maxSplats = computeMaxSplats(numSplats); - const packedArray = new Uint32Array(maxSplats * 4); + const packedArray: [Uint32Array, Uint32Array] = [ + new Uint32Array(maxSplats * 4), + new Uint32Array(maxSplats * 4), + ]; const extra: Record = {}; let sectionBase = HEADER_BYTES + maxSectionCount * SECTION_BYTES; diff --git a/src/pcsogs.ts b/src/pcsogs.ts index 7eb090d..b037814 100644 --- a/src/pcsogs.ts +++ b/src/pcsogs.ts @@ -15,7 +15,7 @@ export async function unpackPcSogs( json: PcSogsJson, extraFiles: Record, ): Promise<{ - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }> { @@ -25,7 +25,10 @@ export async function unpackPcSogs( const numSplats = json.means.shape[0]; const maxSplats = computeMaxSplats(numSplats); - const packedArray = new Uint32Array(maxSplats * 4); + const packedArray: [Uint32Array, Uint32Array] = [ + new Uint32Array(maxSplats * 4), + new Uint32Array(maxSplats * 4), + ]; const extra: Record = {}; const meansPromise = Promise.all([ @@ -249,7 +252,7 @@ async function decodeImageRgba(fileBytes: ArrayBuffer) { } export async function unpackPcSogsZip(fileBytes: Uint8Array): Promise<{ - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }> { diff --git a/src/shaders/computeUvec4x2.glsl b/src/shaders/computeUvec4x2.glsl new file mode 100644 index 0000000..4d33f0b --- /dev/null +++ b/src/shaders/computeUvec4x2.glsl @@ -0,0 +1,38 @@ +precision highp float; +precision highp int; +precision highp sampler2D; +precision highp usampler2D; +precision highp isampler2D; +precision highp sampler2DArray; +precision highp usampler2DArray; +precision highp isampler2DArray; +precision highp sampler3D; +precision highp usampler3D; +precision highp isampler3D; + +#include + +uniform uint targetLayer; +uniform int targetBase; +uniform int targetCount; + +layout(location = 0) out uvec4 target; +layout(location = 1) out uvec4 target2; + +{{ GLOBALS }} + +void produceSplat(int index) { + {{ STATEMENTS }} +} + +void main() { + int targetIndex = int(targetLayer << SPLAT_TEX_LAYER_BITS) + int(uint(gl_FragCoord.y) << SPLAT_TEX_WIDTH_BITS) + int(gl_FragCoord.x); + int index = targetIndex - targetBase; + + if ((index >= 0) && (index < targetCount)) { + produceSplat(index); + } else { + target = uvec4(0u, 0u, 0u, 0u); + target2 = uvec4(0u, 0u, 0u, 0u); + } +} diff --git a/src/shaders/splatDefines.glsl b/src/shaders/splatDefines.glsl index 4a4fc10..ebe9b1c 100644 --- a/src/shaders/splatDefines.glsl +++ b/src/shaders/splatDefines.glsl @@ -203,14 +203,7 @@ vec4 decodeQuatOctXy88R8(uint encoded) { // } // Pack a Gsplat into a uvec4 -uvec4 packSplat(vec3 center, vec3 scales, vec4 quaternion, vec4 rgba) { - uvec4 uRgba = uvec4(round(clamp(rgba * 255.0, 0.0, 255.0))); - - uint uQuat = encodeQuatOctXy88R8(quaternion); - // uint uQuat = encodeQuatXyz888(quaternion); - // uint uQuat = encodeQuatEulerXyz888(quaternion); - uvec3 uQuat3 = uvec3(uQuat & 0xffu, (uQuat >> 8u) & 0xffu, (uQuat >> 16u) & 0xffu); - +uvec4[2] packSplat(vec3 center, vec3 scales, vec4 quaternion, vec4 rgba) { // Encode scales in three uint8s, where 0=>0.0 and 1..=255 stores log scale uvec3 uScales = uvec3( (scales.x == 0.0) ? 0u : uint(round(clamp((log(scales.x) - LN_SCALE_MIN) / LN_RESCALE, 0.0, 254.0))) + 1u, @@ -218,25 +211,30 @@ uvec4 packSplat(vec3 center, vec3 scales, vec4 quaternion, vec4 rgba) { (scales.z == 0.0) ? 0u : uint(round(clamp((log(scales.z) - LN_SCALE_MIN) / LN_RESCALE, 0.0, 254.0))) + 1u ); - // Pack it all into 4 x uint32 - uint word0 = uRgba.r | (uRgba.g << 8u) | (uRgba.b << 16u) | (uRgba.a << 24u); - uint word1 = packHalf2x16(center.xy); - uint word2 = packHalf2x16(vec2(center.z, 0.0)) | (uQuat3.x << 16u) | (uQuat3.y << 24u); - uint word3 = uScales.x | (uScales.y << 8u) | (uScales.z << 16u) | (uQuat3.z << 24u); - return uvec4(word0, word1, word2, word3); + uvec4 uRgba = uvec4(round(clamp(rgba * 255.0, 0.0, 65535.0))); + + uint uQuat = encodeQuatOctXy88R8(quaternion); + + return uvec4[2]( + uvec4( + floatBitsToUint(center), + uScales.x | (uScales.y << 8u) | (uScales.z << 16u) + ), + uvec4( + uRgba.r | (uRgba.g << 16u), + uRgba.b | (uRgba.a << 16u), + uQuat, + 0u + ) + ); } // Unpack a Gsplat from a uvec4 -void unpackSplat(uvec4 packed, out vec3 center, out vec3 scales, out vec4 quaternion, out vec4 rgba) { - uint word0 = packed.x, word1 = packed.y, word2 = packed.z, word3 = packed.w; +void unpackSplat(uvec4[2] packed, out vec3 center, out vec3 scales, out vec4 quaternion, out vec4 rgba) { + uint word0 = packed[0].x, word1 = packed[0].y, word2 = packed[0].z, word3 = packed[0].w, + word4 = packed[1].x, word5 = packed[1].y, word6 = packed[1].z, word7 = packed[1].w; - uvec4 uRgba = uvec4(word0 & 0xffu, (word0 >> 8u) & 0xffu, (word0 >> 16u) & 0xffu, (word0 >> 24u) & 0xffu); - rgba = vec4(uRgba) / 255.0; - - center = vec4( - unpackHalf2x16(word1), - unpackHalf2x16(word2 & 0xffffu) - ).xyz; + center = uintBitsToFloat(packed[0].xyz); uvec3 uScales = uvec3(word3 & 0xffu, (word3 >> 8u) & 0xffu, (word3 >> 16u) & 0xffu); scales = vec3( @@ -245,8 +243,10 @@ void unpackSplat(uvec4 packed, out vec3 center, out vec3 scales, out vec4 quater (uScales.z == 0u) ? 0.0 : exp(LN_SCALE_MIN + float(uScales.z - 1u) * LN_RESCALE) ); + uvec4 uRgba = uvec4(word4 & 0xffffu, (word4 >> 16u) & 0xffffu, word5 & 0xffffu, (word5 >> 16u) & 0xffffu); + rgba = vec4(uRgba) / 255.0; - uint uQuat = ((word2 >> 16u) & 0xFFFFu) | ((word3 >> 8u) & 0xFF0000u); + uint uQuat = word6 & 0xFFFFFFu; quaternion = decodeQuatOctXy88R8(uQuat); // quaternion = decodeQuatXyz888(uQuat); // quaternion = decodeQuatEulerXyz888(uQuat); diff --git a/src/shaders/splatFragment.glsl b/src/shaders/splatFragment.glsl index c49de39..7c2a625 100644 --- a/src/shaders/splatFragment.glsl +++ b/src/shaders/splatFragment.glsl @@ -66,5 +66,5 @@ void main() { if (encodeLinear) { rgba.rgb = srgbToLinear(rgba.rgb); } - fragColor = rgba; + fragColor = vec4(rgba.rgb * rgba.a, rgba.a); } diff --git a/src/shaders/splatVertex.glsl b/src/shaders/splatVertex.glsl index aea0475..4125f59 100644 --- a/src/shaders/splatVertex.glsl +++ b/src/shaders/splatVertex.glsl @@ -28,6 +28,7 @@ uniform float clipXY; uniform float focalAdjustment; uniform usampler2DArray packedSplats; +uniform usampler2DArray packedSplats2; void main() { // Default to outside the frustum so it's discarded if we return early @@ -46,7 +47,10 @@ void main() { (splatIndex >> SPLAT_TEX_WIDTH_BITS) & SPLAT_TEX_HEIGHT_MASK, splatIndex >> SPLAT_TEX_LAYER_BITS ); - uvec4 packed = texelFetch(packedSplats, texCoord, 0); + uvec4[2] packed = uvec4[2]( + texelFetch(packedSplats, texCoord, 0), + texelFetch(packedSplats2, texCoord, 0) + ); vec3 center, scales; vec4 quaternion, rgba; diff --git a/src/spz.ts b/src/spz.ts index 298987b..dc15bb1 100644 --- a/src/spz.ts +++ b/src/spz.ts @@ -6,7 +6,7 @@ import { getSplatFileType, getSplatFileTypeFromPath, } from "./SplatLoader"; -import { GunzipReader, fromHalf, unpackSplat } from "./utils"; +import { GunzipReader, fromHalf } from "./utils"; import { decodeAntiSplat } from "./antisplat"; import { decodeKsplat } from "./ksplat"; diff --git a/src/utils.ts b/src/utils.ts index efeeb8b..c988a4e 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -337,7 +337,7 @@ export class FreeList { // logarithmic uint8, rotation as three uint8s representing rotation axis and angle, // and RGBA as 4xuint8. export function setPackedSplat( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, x: number, y: number, @@ -354,10 +354,10 @@ export function setPackedSplat( g: number, b: number, ) { - const uR = floatToUint8(r); - const uG = floatToUint8(g); - const uB = floatToUint8(b); - const uA = floatToUint8(opacity); + const uR = Math.max(0, Math.min(65535, Math.round(r * 255))); + const uG = Math.max(0, Math.min(65535, Math.round(g * 255))); + const uB = Math.max(0, Math.min(65535, Math.round(b * 255))); + const uA = Math.max(0, Math.min(65535, Math.round(opacity * 255))); // Alternate internal encodings commented out below. const uQuat = encodeQuatOctXy88R8( @@ -365,9 +365,6 @@ export function setPackedSplat( ); // const uQuat = encodeQuatXyz888(new THREE.Quaternion(quatX, quatY, quatZ, quatW)); // const uQuat = encodeQuatEulerXyz888(new THREE.Quaternion(quatX, quatY, quatZ, quatW)); - const uQuatX = uQuat & 0xff; - const uQuatY = (uQuat >>> 8) & 0xff; - const uQuatZ = (uQuat >>> 16) & 0xff; // Allow scales below LN_SCALE_MIN to be encoded as 0, which signifies a 2DGS const uScaleX = @@ -401,40 +398,36 @@ export function setPackedSplat( ), ); - const uCenterX = toHalf(x); - const uCenterY = toHalf(y); - const uCenterZ = toHalf(z); - // Encode the splat as 4 consecutive Uint32 elements const i4 = index * 4; - packedSplats[i4] = uR | (uG << 8) | (uB << 16) | (uA << 24); - packedSplats[i4 + 1] = uCenterX | (uCenterY << 16); - packedSplats[i4 + 2] = uCenterZ | (uQuatX << 16) | (uQuatY << 24); - packedSplats[i4 + 3] = - uScaleX | (uScaleY << 8) | (uScaleZ << 16) | (uQuatZ << 24); + packedSplats[0][i4] = floatBitsToUint(x); + packedSplats[0][i4 + 1] = floatBitsToUint(y); + packedSplats[0][i4 + 2] = floatBitsToUint(z); + packedSplats[0][i4 + 3] = uScaleX | (uScaleY << 8) | (uScaleZ << 16); + packedSplats[1][i4] = uR | (uG << 16); + packedSplats[1][i4 + 1] = uB | (uA << 16); + packedSplats[1][i4 + 2] = uQuat; + packedSplats[1][i4 + 3] = 0; } // Encode the center coordinates x,y,z in the packedSplats Uint32Array, // leaving all other fields as is. export function setPackedSplatCenter( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, x: number, y: number, z: number, ) { - const uCenterX = toHalf(x); - const uCenterY = toHalf(y); - const uCenterZ = toHalf(z); - const i4 = index * 4; - packedSplats[i4 + 1] = uCenterX | (uCenterY << 16); - packedSplats[i4 + 2] = uCenterZ | (packedSplats[i4 + 2] & 0xffff0000); + packedSplats[0][i4] = floatBitsToUint(x); + packedSplats[0][i4 + 1] = floatBitsToUint(y); + packedSplats[0][i4 + 2] = floatBitsToUint(z); } // Encode the scales x,y,z in the packedSplats Uint32Array, leaving all other fields as is. export function setPackedSplatScales( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, scaleX: number, scaleY: number, @@ -473,17 +466,13 @@ export function setPackedSplatScales( ); const i4 = index * 4; - packedSplats[i4 + 3] = - uScaleX | - (uScaleY << 8) | - (uScaleZ << 16) | - (packedSplats[i4 + 3] & 0xff000000); + packedSplats[0][i4 + 3] = uScaleX | (uScaleY << 8) | (uScaleZ << 16); } // Encode the rotation quatX, quatY, quatZ, quatW in the packedSplats Uint32Array, // leaving all other fields as is. export function setPackedSplatQuat( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, quatX: number, quatY: number, @@ -495,60 +484,57 @@ export function setPackedSplatQuat( ); // const uQuat = encodeQuatXyz888(new THREE.Quaternion(quatX, quatY, quatZ, quatW)); // const uQuat = encodeQuatEulerXyz888(new THREE.Quaternion(quatX, quatY, quatZ, quatW)); - const uQuatX = uQuat & 0xff; - const uQuatY = (uQuat >>> 8) & 0xff; - const uQuatZ = (uQuat >>> 16) & 0xff; const i4 = index * 4; - packedSplats[i4 + 2] = - (packedSplats[i4 + 2] & 0x0000ffff) | (uQuatX << 16) | (uQuatY << 24); - packedSplats[i4 + 3] = (packedSplats[i4 + 3] & 0x00ffffff) | (uQuatZ << 24); + packedSplats[1][i4 + 2] = uQuat; } // Encode the RGBA color in the packedSplats Uint32Array, leaving other fields alone. export function setPackedSplatRgba( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, r: number, g: number, b: number, a: number, ) { - const uR = floatToUint8(r); - const uG = floatToUint8(g); - const uB = floatToUint8(b); - const uA = floatToUint8(a); + const uR = Math.max(0, Math.min(65535, Math.round(r * 255))); + const uG = Math.max(0, Math.min(65535, Math.round(g * 255))); + const uB = Math.max(0, Math.min(65535, Math.round(b * 255))); + const uA = Math.max(0, Math.min(65535, Math.round(a * 255))); + const i4 = index * 4; - packedSplats[i4] = uR | (uG << 8) | (uB << 16) | (uA << 24); + packedSplats[1][i4] = uR | (uG << 16); + packedSplats[1][i4 + 1] = uB | (uA << 16); } // Encode the RGB color in the packedSplats Uint32Array, leaving other fields alone. export function setPackedSplatRgb( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, r: number, g: number, b: number, ) { - const uR = floatToUint8(r); - const uG = floatToUint8(g); - const uB = floatToUint8(b); + const uR = Math.max(0, Math.min(65535, Math.round(r * 255))); + const uG = Math.max(0, Math.min(65535, Math.round(g * 255))); + const uB = Math.max(0, Math.min(65535, Math.round(b * 255))); const i4 = index * 4; - packedSplats[i4] = - uR | (uG << 8) | (uB << 16) | (packedSplats[i4] & 0xff000000); + packedSplats[1][i4] = uR | (uG << 16); + packedSplats[1][i4 + 1] = uB | (packedSplats[1][i4 + 1] & 0xffff0000); } // Encode the opacity in the packedSplats Uint32Array, leaving other fields alone. export function setPackedSplatOpacity( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, opacity: number, ) { - const uA = floatToUint8(opacity); + const uA = Math.max(0, Math.min(65535, Math.round(opacity * 255))); const i4 = index * 4; - packedSplats[i4] = (packedSplats[i4] & 0x00ffffff) | (uA << 24); + packedSplats[1][i4 + 1] = (packedSplats[1][i4 + 1] & 0x0000ffff) | (uA << 16); } const packedCenter = new THREE.Vector3(); @@ -566,7 +552,7 @@ const packedFields = { // Unpack all components of a PackedSplat from the packedSplats Uint32Array into // THREE.js vector objects. The returned objects will be reused each call. export function unpackSplat( - packedSplats: Uint32Array, + packedSplats: [Uint32Array, Uint32Array], index: number, ): { center: THREE.Vector3; @@ -579,21 +565,19 @@ export function unpackSplat( const result = packedFields; const i4 = index * 4; - const word0 = packedSplats[i4]; - const word1 = packedSplats[i4 + 1]; - const word2 = packedSplats[i4 + 2]; - const word3 = packedSplats[i4 + 3]; + const word0 = packedSplats[0][i4]; + const word1 = packedSplats[0][i4 + 1]; + const word2 = packedSplats[0][i4 + 2]; + const word3 = packedSplats[0][i4 + 3]; + const word4 = packedSplats[1][i4]; + const word5 = packedSplats[1][i4 + 1]; + const word6 = packedSplats[1][i4 + 2]; + const word7 = packedSplats[1][i4 + 3]; - result.color.set( - (word0 & 0xff) / 255, - ((word0 >>> 8) & 0xff) / 255, - ((word0 >>> 16) & 0xff) / 255, - ); - result.opacity = ((word0 >>> 24) & 0xff) / 255; result.center.set( - fromHalf(word1 & 0xffff), - fromHalf((word1 >>> 16) & 0xffff), - fromHalf(word2 & 0xffff), + uintBitsToFloat(word0), + uintBitsToFloat(word1), + uintBitsToFloat(word2), ); const uScalesX = word3 & 0xff; @@ -606,8 +590,14 @@ export function unpackSplat( result.scales.z = uScalesZ === 0 ? 0.0 : Math.exp(LN_SCALE_MIN + (uScalesZ - 1) * LN_RESCALE); - const uQuat = ((word2 >>> 16) & 0xffff) | ((word3 >>> 8) & 0xff0000); - decodeQuatOctXy88R8(uQuat, result.quaternion); + result.color.set( + (word4 & 0xffff) / 255, + ((word4 >>> 16) & 0xffff) / 255, + (word5 & 0xffff) / 255, + ); + result.opacity = ((word5 >>> 16) & 0xffff) / 255; + + decodeQuatOctXy88R8(word6 & 0xffffff, result.quaternion); // decodeQuatXyz888(uQuat, result.quaternion); // decodeQuatEulerXyz888(uQuat, result.quaternion); diff --git a/src/worker.ts b/src/worker.ts index 98c95a5..f85ac05 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -11,7 +11,9 @@ import { encodeSh1Rgb, encodeSh2Rgb, encodeSh3Rgb, + floatBitsToUint, getArrayBuffers, + getTextureSize, setPackedSplat, setPackedSplatCenter, setPackedSplatOpacity, @@ -37,7 +39,7 @@ async function onMessage(event: MessageEvent) { switch (name) { case "unpackPly": { const { packedArray, fileBytes } = args as { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; fileBytes: Uint8Array; }; const decoded = await unpackPly({ packedArray, fileBytes }); @@ -183,8 +185,11 @@ async function onMessage(event: MessageEvent) { async function unpackPly({ packedArray, fileBytes, -}: { packedArray: Uint32Array; fileBytes: Uint8Array }): Promise<{ - packedArray: Uint32Array; +}: { + packedArray: [Uint32Array, Uint32Array]; + fileBytes: Uint8Array; +}): Promise<{ + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; }> { @@ -257,14 +262,17 @@ async function unpackPly({ } function unpackSpz(fileBytes: Uint8Array): { - packedArray: Uint32Array; + packedArray: [Uint32Array, Uint32Array]; numSplats: number; extra: Record; } { const spz = new SpzReader({ fileBytes }); const numSplats = spz.numSplats; const maxSplats = computeMaxSplats(numSplats); - const packedArray = new Uint32Array(maxSplats * 4); + const packedArray: [Uint32Array, Uint32Array] = [ + new Uint32Array(maxSplats * 4), + new Uint32Array(maxSplats * 4), + ]; const extra: Record = {}; spz.parseSplats(