diff --git a/src/PackedSplats.ts b/src/PackedSplats.ts index f9ba443..e975411 100644 --- a/src/PackedSplats.ts +++ b/src/PackedSplats.ts @@ -33,6 +33,7 @@ export type SplatEncoding = { sh2Max?: number; sh3Min?: number; sh3Max?: number; + maxSh?: number; }; export const DEFAULT_SPLAT_ENCODING: SplatEncoding = { @@ -46,6 +47,7 @@ export const DEFAULT_SPLAT_ENCODING: SplatEncoding = { sh2Max: 1, sh3Min: -1, sh3Max: 1, + maxSh: 3, }; // Initialize a PackedSplats collection from source data via diff --git a/src/SplatMesh.ts b/src/SplatMesh.ts index 2940b16..195bcbe 100644 --- a/src/SplatMesh.ts +++ b/src/SplatMesh.ts @@ -38,7 +38,7 @@ import { unindent, unindentLines, } from "./dyno"; -import { getTextureSize } from "./utils"; +import { getShArrayStride, getTextureSize } from "./utils"; export type SplatMeshOptions = { // URL to fetch a Gaussian splat file from(supports .ply, .splat, .ksplat, @@ -380,9 +380,8 @@ export class SplatMesh extends SplatGenerator { if (this.maxSh >= 1) { // Inject lighting from SH1..SH3 - const { sh1Texture, sh2Texture, sh3Texture } = - this.ensureShTextures(); - if (sh1Texture) { + const { shTexture, shDegrees } = this.ensureShTexture(); + if (shTexture && shDegrees) { //Calculate view direction in object space const viewCenterInObject = viewToObject.translate; const { center } = splitGsplat(gsplat).outputs; @@ -399,22 +398,16 @@ export class SplatMesh extends SplatGenerator { } // Evaluate Spherical Harmonics - const sh1Snorm = evaluateSH1(gsplat, sh1Texture, viewDir); - let rgb = rescaleSh(sh1Snorm, this.packedSplats.dynoSh1MinMax); - if (this.maxSh >= 2 && sh2Texture) { - const sh2Snorm = evaluateSH2(gsplat, sh2Texture, viewDir); - rgb = add( - rgb, - rescaleSh(sh2Snorm, this.packedSplats.dynoSh2MinMax), - ); - } - if (this.maxSh >= 3 && sh3Texture) { - const sh3Snorm = evaluateSH3(gsplat, sh3Texture, viewDir); - rgb = add( - rgb, - rescaleSh(sh3Snorm, this.packedSplats.dynoSh3MinMax), - ); - } + const rgb = evaluateSH( + gsplat, + shTexture, + viewDir, + shDegrees, + this.maxSh, + this.packedSplats.dynoSh1MinMax, + this.packedSplats.dynoSh2MinMax, + this.packedSplats.dynoSh3MinMax, + ); // Flash off for 0.3 / 1.0 sec for debugging // const fractTime = fract(SplatMesh.dynoTime); @@ -520,7 +513,7 @@ export class SplatMesh extends SplatGenerator { const viewToObjectMatrix = worldToObject.multiply(viewToWorld); if ( viewToObject.updateFromMatrix(viewToObjectMatrix) && - (this.enableViewToObject || this.packedSplats.extra.sh1) + (this.enableViewToObject || this.packedSplats.extra.sh) ) { // Only trigger update if we have view-dependent spherical harmonics updated = true; @@ -642,215 +635,112 @@ export class SplatMesh extends SplatGenerator { } } - private ensureShTextures(): { - sh1Texture?: DynoUsampler2DArray<"sh1", THREE.DataArrayTexture>; - sh2Texture?: DynoUsampler2DArray<"sh2", THREE.DataArrayTexture>; - sh3Texture?: DynoUsampler2DArray<"sh3", THREE.DataArrayTexture>; + private ensureShTexture(): { + shTexture?: DynoUsampler2DArray<"sh", THREE.DataArrayTexture>; + shDegrees?: 0 | 1 | 2 | 3; } { // Ensure we have textures for SH1..SH3 if we have data - if (!this.packedSplats.extra.sh1) { + if (!this.packedSplats.extra.sh) { return {}; } - let sh1Texture = this.packedSplats.extra.sh1Texture as - | DynoUsampler2DArray<"sh1", THREE.DataArrayTexture> - | undefined; - if (!sh1Texture) { - let sh1 = this.packedSplats.extra.sh1 as Uint32Array; - const { width, height, depth, maxSplats } = getTextureSize( - sh1.length / 2, - ); - if (sh1.length < maxSplats * 2) { - const newSh1 = new Uint32Array(maxSplats * 2); - newSh1.set(sh1); - this.packedSplats.extra.sh1 = newSh1; - sh1 = newSh1; - } - - const texture = new THREE.DataArrayTexture(sh1, width, height, depth); - texture.format = THREE.RGIntegerFormat; - texture.type = THREE.UnsignedIntType; - texture.internalFormat = "RG32UI"; - texture.needsUpdate = true; - - sh1Texture = new DynoUsampler2DArray({ - value: texture, - key: "sh1", - }); - this.packedSplats.extra.sh1Texture = sh1Texture; - } - - if (!this.packedSplats.extra.sh2) { - return { sh1Texture }; - } - - let sh2Texture = this.packedSplats.extra.sh2Texture as - | DynoUsampler2DArray<"sh2", THREE.DataArrayTexture> + const shDegrees = (this.packedSplats.extra.shDegrees as 0 | 1 | 2 | 3) ?? 0; + let shTexture = this.packedSplats.extra.shTexture as + | DynoUsampler2DArray<"sh", THREE.DataArrayTexture> | undefined; - if (!sh2Texture) { - let sh2 = this.packedSplats.extra.sh2 as Uint32Array; + if (!shTexture && shDegrees) { + let sh = this.packedSplats.extra.sh as Uint8Array; const { width, height, depth, maxSplats } = getTextureSize( - sh2.length / 4, + this.numSplats * shDegrees, // 1st order = 1 RGB pixel, 2nd = 2 RGB pixels, 3rd = 3 RGBA pixels ); - if (sh2.length < maxSplats * 4) { - const newSh2 = new Uint32Array(maxSplats * 4); - newSh2.set(sh2); - this.packedSplats.extra.sh2 = newSh2; - sh2 = newSh2; + const coefficientsWithPadding = getShArrayStride(shDegrees); + if (sh.length < maxSplats * coefficientsWithPadding) { + const newSh = new Uint8Array(maxSplats * coefficientsWithPadding); + newSh.set(sh); + this.packedSplats.extra.sh = newSh; + sh = newSh; } - const texture = new THREE.DataArrayTexture(sh2, width, height, depth); - texture.format = THREE.RGBAIntegerFormat; - texture.type = THREE.UnsignedIntType; - texture.internalFormat = "RGBA32UI"; - texture.needsUpdate = true; - - sh2Texture = new DynoUsampler2DArray({ - value: texture, - key: "sh2", - }); - this.packedSplats.extra.sh2Texture = sh2Texture; - } - - if (!this.packedSplats.extra.sh3) { - return { sh1Texture, sh2Texture }; - } - - let sh3Texture = this.packedSplats.extra.sh3Texture as - | DynoUsampler2DArray<"sh3", THREE.DataArrayTexture> - | undefined; - if (!sh3Texture) { - let sh3 = this.packedSplats.extra.sh3 as Uint32Array; - const { width, height, depth, maxSplats } = getTextureSize( - sh3.length / 4, + const texture = new THREE.DataArrayTexture( + new Uint32Array(sh.buffer), + width, + height, + depth, ); - if (sh3.length < maxSplats * 4) { - const newSh3 = new Uint32Array(maxSplats * 4); - newSh3.set(sh3); - this.packedSplats.extra.sh3 = newSh3; - sh3 = newSh3; - } - - const texture = new THREE.DataArrayTexture(sh3, width, height, depth); - texture.format = THREE.RGBAIntegerFormat; + texture.format = ( + shDegrees === 3 ? THREE.RGBAIntegerFormat : "RGB_INTEGER" + ) as THREE.AnyPixelFormat; texture.type = THREE.UnsignedIntType; - texture.internalFormat = "RGBA32UI"; + texture.internalFormat = shDegrees === 3 ? "RGBA32UI" : "RGB32UI"; texture.needsUpdate = true; - sh3Texture = new DynoUsampler2DArray({ + shTexture = new DynoUsampler2DArray({ value: texture, - key: "sh3", + key: "sh", }); - this.packedSplats.extra.sh3Texture = sh3Texture; + this.packedSplats.extra.shTexture = shTexture; } - return { sh1Texture, sh2Texture, sh3Texture }; + return { shTexture, shDegrees }; } } -const defineEvaluateSH1 = unindent(` - vec3 evaluateSH1(Gsplat gsplat, usampler2DArray sh1, vec3 viewDir) { - // Extract sint7 values packed into 2 x uint32 - uvec2 packed = texelFetch(sh1, splatTexCoord(gsplat.index), 0).rg; - vec3 sh1_0 = vec3(ivec3( - int(packed.x << 25u) >> 25, - int(packed.x << 18u) >> 25, - int(packed.x << 11u) >> 25 - )) / 63.0; - vec3 sh1_1 = vec3(ivec3( - int(packed.x << 4u) >> 25, - int((packed.x >> 3u) | (packed.y << 29u)) >> 25, - int(packed.y << 22u) >> 25 - )) / 63.0; - vec3 sh1_2 = vec3(ivec3( - int(packed.y << 15u) >> 25, - int(packed.y << 8u) >> 25, - int(packed.y << 1u) >> 25 - )) / 63.0; - - return sh1_0 * (-0.4886025 * viewDir.y) - + sh1_1 * (0.4886025 * viewDir.z) - + sh1_2 * (-0.4886025 * viewDir.x); +const defineUnpackSint8 = unindent(/*glsl*/ ` + vec4 unpackSint8(uint packed) { + return vec4(ivec4( + int(packed << 24u) >> 24, + int(packed << 16u) >> 24, + int(packed << 8u) >> 24, + int(packed) >> 24 + )) / 127.0; } `); -const defineEvaluateSH2 = unindent(` - vec3 evaluateSH2(Gsplat gsplat, usampler2DArray sh2, vec3 viewDir) { - // Extract sint8 values packed into 4 x uint32 - uvec4 packed = texelFetch(sh2, splatTexCoord(gsplat.index), 0); - vec3 sh2_0 = vec3(ivec3( - int(packed.x << 24u) >> 24, - int(packed.x << 16u) >> 24, - int(packed.x << 8u) >> 24 - )) / 127.0; - vec3 sh2_1 = vec3(ivec3( - int(packed.x) >> 24, - int(packed.y << 24u) >> 24, - int(packed.y << 16u) >> 24 - )) / 127.0; - vec3 sh2_2 = vec3(ivec3( - int(packed.y << 8u) >> 24, - int(packed.y) >> 24, - int(packed.z << 24u) >> 24 - )) / 127.0; - vec3 sh2_3 = vec3(ivec3( - int(packed.z << 16u) >> 24, - int(packed.z << 8u) >> 24, - int(packed.z) >> 24 - )) / 127.0; - vec3 sh2_4 = vec3(ivec3( - int(packed.w << 24u) >> 24, - int(packed.w << 16u) >> 24, - int(packed.w << 8u) >> 24 - )) / 127.0; +const defineEvaluateSH = unindent(/* glsl */ ` + vec3 evaluateSH(Gsplat gsplat, usampler2DArray sh, vec3 viewDir, vec2 sh1MinMax, vec2 sh2MinMax, vec2 sh3MinMax) { + // Extract sint8 values packed into 3, 8 and 12 x uint32 + uvec4 packedA = texelFetch(sh, splatTexCoord(gsplat.index * SH_TEXEL_STRIDE), 0); + vec4 a1 = unpackSint8(packedA.x); + vec4 a2 = unpackSint8(packedA.y); + vec4 a3 = unpackSint8(packedA.z); + vec4 a4 = unpackSint8(packedA.w); +#if SH_DEGREES > 1 + uvec4 packedB = texelFetch(sh, splatTexCoord(gsplat.index * SH_TEXEL_STRIDE + 1), 0); + vec4 b1 = unpackSint8(packedB.x); + vec4 b2 = unpackSint8(packedB.y); + vec4 b3 = unpackSint8(packedB.z); + vec4 b4 = unpackSint8(packedB.w); +#endif +#if SH_DEGREES > 2 + uvec4 packedC = texelFetch(sh, splatTexCoord(gsplat.index * SH_TEXEL_STRIDE + 2), 0); + vec4 c1 = unpackSint8(packedC.x); + vec4 c2 = unpackSint8(packedC.y); + vec4 c3 = unpackSint8(packedC.z); + vec4 c4 = unpackSint8(packedC.w); +#endif + +#if SH_DEGREES <= 2 + // RGB + vec3 sh1_0 = a1.xyz; + vec3 sh1_1 = vec3(a1.w, a2.xy); + vec3 sh1_2 = vec3(a2.zw, a3.x); +#else + // RGBA + vec3 sh1_0 = a1.xyz; + vec3 sh1_1 = vec3(a1.w, a2.xy); + vec3 sh1_2 = vec3(a2.zw, a3.x); +#endif + + vec3 sh1 = sh1_0 * (-0.4886025 * viewDir.y) + + sh1_1 * (0.4886025 * viewDir.z) + + sh1_2 * (-0.4886025 * viewDir.x); - return sh2_0 * (1.0925484 * viewDir.x * viewDir.y) - + sh2_1 * (-1.0925484 * viewDir.y * viewDir.z) - + sh2_2 * (0.3153915 * (2.0 * viewDir.z * viewDir.z - viewDir.x * viewDir.x - viewDir.y * viewDir.y)) - + sh2_3 * (-1.0925484 * viewDir.x * viewDir.z) - + sh2_4 * (0.5462742 * (viewDir.x * viewDir.x - viewDir.y * viewDir.y)); - } -`); + // rescale + sh1 = (sh1MinMax.x + sh1MinMax.y) / 2.0 + sh1 * (sh1MinMax.y - sh1MinMax.x) / 2.0; -const defineEvaluateSH3 = unindent(` - vec3 evaluateSH3(Gsplat gsplat, usampler2DArray sh3, vec3 viewDir) { - // Extract sint6 values packed into 4 x uint32 - uvec4 packed = texelFetch(sh3, splatTexCoord(gsplat.index), 0); - vec3 sh3_0 = vec3(ivec3( - int(packed.x << 26u) >> 26, - int(packed.x << 20u) >> 26, - int(packed.x << 14u) >> 26 - )) / 31.0; - vec3 sh3_1 = vec3(ivec3( - int(packed.x << 8u) >> 26, - int(packed.x << 2u) >> 26, - int((packed.x >> 4u) | (packed.y << 28u)) >> 26 - )) / 31.0; - vec3 sh3_2 = vec3(ivec3( - int(packed.y << 22u) >> 26, - int(packed.y << 16u) >> 26, - int(packed.y << 10u) >> 26 - )) / 31.0; - vec3 sh3_3 = vec3(ivec3( - int(packed.y << 4u) >> 26, - int((packed.y >> 2u) | (packed.z << 30u)) >> 26, - int(packed.z << 24u) >> 26 - )) / 31.0; - vec3 sh3_4 = vec3(ivec3( - int(packed.z << 18u) >> 26, - int(packed.z << 12u) >> 26, - int(packed.z << 6u) >> 26 - )) / 31.0; - vec3 sh3_5 = vec3(ivec3( - int(packed.z) >> 26, - int(packed.w << 26u) >> 26, - int(packed.w << 20u) >> 26 - )) / 31.0; - vec3 sh3_6 = vec3(ivec3( - int(packed.w << 14u) >> 26, - int(packed.w << 8u) >> 26, - int(packed.w << 2u) >> 26 - )) / 31.0; +#if SH_DEGREES == 1 || MAX_SH == 1 + return sh1; +#else float xx = viewDir.x * viewDir.x; float yy = viewDir.y * viewDir.y; @@ -859,74 +749,90 @@ const defineEvaluateSH3 = unindent(` float yz = viewDir.y * viewDir.z; float zx = viewDir.z * viewDir.x; - return sh3_0 * (-0.5900436 * viewDir.y * (3.0 * xx - yy)) +#if SH_DEGREES <= 2 + // RGB + vec3 sh2_0 = a3.yzw; + vec3 sh2_1 = b1.xyz; + vec3 sh2_2 = vec3(b1.w, b2.xy); + vec3 sh2_3 = vec3(b2.zw, b3.x); + vec3 sh2_4 = b3.yzw; +#else + // RGBA + vec3 sh2_0 = vec3(a3.yzw); + vec3 sh2_1 = vec3(a4.xyz); + vec3 sh2_2 = vec3(a4.w, b1.xy); + vec3 sh2_3 = vec3(b1.zw, b2.x); + vec3 sh2_4 = vec3(b2.yzw); +#endif + vec3 sh2 = sh2_0 * (1.0925484 * xy) + + sh2_1 * (-1.0925484 * yz) + + sh2_2 * (0.3153915 * (2.0 * zz - xx - yy)) + + sh2_3 * (-1.0925484 * zx) + + sh2_4 * (0.5462742 * (xx - yy)); + + // rescale + sh2 = (sh2MinMax.x + sh2MinMax.y) / 2.0 + sh2 * (sh2MinMax.y - sh2MinMax.x) / 2.0; + +#if SH_DEGREES == 2 || MAX_SH == 2 + return sh1 + sh2; +#else + vec3 sh3_0 = vec3(b3.xyz); + vec3 sh3_1 = vec3(b3.w, b4.xy); + vec3 sh3_2 = vec3(b4.zw, c1.x); + vec3 sh3_3 = vec3(c1.yzw); + vec3 sh3_4 = vec3(c2.xyz); + vec3 sh3_5 = vec3(c2.w, c3.xy); + vec3 sh3_6 = vec3(c3.zw, c4.x); + vec3 sh3 = sh3_0 * (-0.5900436 * viewDir.y * (3.0 * xx - yy)) + sh3_1 * (2.8906114 * xy * viewDir.z) + + sh3_2 * (-0.4570458 * viewDir.y * (4.0 * zz - xx - yy)) + sh3_3 * (0.3731763 * viewDir.z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) + sh3_4 * (-0.4570458 * viewDir.x * (4.0 * zz - xx - yy)) + sh3_5 * (1.4453057 * viewDir.z * (xx - yy)) + sh3_6 * (-0.5900436 * viewDir.x * (xx - 3.0 * yy)); + + // rescale + sh3 = (sh3MinMax.x + sh3MinMax.y) / 2.0 + sh3 * (sh3MinMax.y - sh3MinMax.x) / 2.0; + + return sh1 + sh2 + sh3; +#endif +#endif } `); -export function evaluateSH1( +export function evaluateSH( gsplat: DynoVal, - sh1: DynoUsampler2DArray<"sh1", THREE.DataArrayTexture>, + sh: DynoVal<"usampler2DArray">, viewDir: DynoVal<"vec3">, + shDegrees: 1 | 2 | 3, + maxSh: number, + sh1MinMax: DynoVal<"vec2">, + sh2MinMax: DynoVal<"vec2">, + sh3MinMax: DynoVal<"vec2">, ): DynoVal<"vec3"> { return dyno({ - inTypes: { gsplat: Gsplat, sh1: "usampler2DArray", viewDir: "vec3" }, + inTypes: { + gsplat: Gsplat, + sh: "usampler2DArray", + viewDir: "vec3", + sh1MinMax: "vec2", + sh2MinMax: "vec2", + sh3MinMax: "vec2", + }, outTypes: { rgb: "vec3" }, - inputs: { gsplat, sh1, viewDir }, - globals: () => [defineGsplat, defineEvaluateSH1], - statements: ({ inputs, outputs }) => { - const statements = unindentLines(` - if (isGsplatActive(${inputs.gsplat}.flags)) { - ${outputs.rgb} = evaluateSH1(${inputs.gsplat}, ${inputs.sh1}, ${inputs.viewDir}); - } else { - ${outputs.rgb} = vec3(0.0); - } + inputs: { gsplat, sh, viewDir, sh1MinMax, sh2MinMax, sh3MinMax }, + globals: () => { + const defines = unindent(` + #define MAX_SH ${maxSh} + #define SH_DEGREES ${shDegrees} + #define SH_TEXEL_STRIDE ${[0, 1, 2, 3][shDegrees]} `); - return statements; + return [defines, defineGsplat, defineUnpackSint8, defineEvaluateSH]; }, - }).outputs.rgb; -} - -export function evaluateSH2( - gsplat: DynoVal, - sh2: DynoVal<"usampler2DArray">, - viewDir: DynoVal<"vec3">, -): DynoVal<"vec3"> { - return dyno({ - inTypes: { gsplat: Gsplat, sh2: "usampler2DArray", viewDir: "vec3" }, - outTypes: { rgb: "vec3" }, - inputs: { gsplat, sh2, viewDir }, - globals: () => [defineGsplat, defineEvaluateSH2], - statements: ({ inputs, outputs }) => - unindentLines(` - if (isGsplatActive(${inputs.gsplat}.flags)) { - ${outputs.rgb} = evaluateSH2(${inputs.gsplat}, ${inputs.sh2}, ${inputs.viewDir}); - } else { - ${outputs.rgb} = vec3(0.0); - } - `), - }).outputs.rgb; -} - -export function evaluateSH3( - gsplat: DynoVal, - sh3: DynoVal<"usampler2DArray">, - viewDir: DynoVal<"vec3">, -): DynoVal<"vec3"> { - return dyno({ - inTypes: { gsplat: Gsplat, sh3: "usampler2DArray", viewDir: "vec3" }, - outTypes: { rgb: "vec3" }, - inputs: { gsplat, sh3, viewDir }, - globals: () => [defineGsplat, defineEvaluateSH3], statements: ({ inputs, outputs }) => unindentLines(` if (isGsplatActive(${inputs.gsplat}.flags)) { - ${outputs.rgb} = evaluateSH3(${inputs.gsplat}, ${inputs.sh3}, ${inputs.viewDir}); + ${outputs.rgb} = evaluateSH(${inputs.gsplat}, ${inputs.sh}, ${inputs.viewDir}, ${inputs.sh1MinMax}, ${inputs.sh2MinMax}, ${inputs.sh3MinMax}); } else { ${outputs.rgb} = vec3(0.0); } diff --git a/src/ksplat.ts b/src/ksplat.ts index 17b1be3..127dd05 100644 --- a/src/ksplat.ts +++ b/src/ksplat.ts @@ -1,10 +1,9 @@ import type { SplatEncoding } from "./PackedSplats"; import { computeMaxSplats, - encodeSh1Rgb, - encodeSh2Rgb, - encodeSh3Rgb, + encodeShRgb, fromHalf, + getShArrayStride, setPackedSplat, } from "./utils"; @@ -416,7 +415,11 @@ export function unpackKsplat( const bucketsMetaDataSizeBytes = partiallyFilledBucketCount * 4; const bucketsStorageSizeBytes = bucketStorageSizeBytes * bucketCount + bucketsMetaDataSizeBytes; - const sphericalHarmonicsDegree = section.getUint16(40, true); + const sphericalHarmonicsDegree = section.getUint16(40, true) as + | 0 + | 1 + | 2 + | 3; const shComponents = KSPLAT_SH_DEGREE_TO_COMPONENTS[sphericalHarmonicsDegree]; @@ -601,33 +604,38 @@ export function unpackKsplat( ); if (sphericalHarmonicsDegree >= 1) { + if (!extra.sh) { + extra.sh = new Uint8Array( + numSplats * getShArrayStride(sphericalHarmonicsDegree), + ); + extra.shDegrees = sphericalHarmonicsDegree; + } + if (sh1) { - if (!extra.sh1) { - extra.sh1 = new Uint32Array(numSplats * 2); - } for (const [i, key] of sh1Index.entries()) { sh1[i] = getSh(splatOffset, key); } - encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1, splatEncoding); } if (sh2) { - if (!extra.sh2) { - extra.sh2 = new Uint32Array(numSplats * 4); - } for (const [i, key] of sh2Index.entries()) { sh2[i] = getSh(splatOffset, key); } - encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2, splatEncoding); } if (sh3) { - if (!extra.sh3) { - extra.sh3 = new Uint32Array(numSplats * 4); - } for (const [i, key] of sh3Index.entries()) { sh3[i] = getSh(splatOffset, key); } - encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3, splatEncoding); } + + encodeShRgb( + extra.sh as Uint8Array, + sphericalHarmonicsDegree, + i, + sh1, + sh2, + sh3, + splatEncoding, + ); } } sectionBase += storageSizeBytes; diff --git a/src/pcsogs.ts b/src/pcsogs.ts index ac1b3eb..c97bd20 100644 --- a/src/pcsogs.ts +++ b/src/pcsogs.ts @@ -3,9 +3,9 @@ import type { SplatEncoding } from "./PackedSplats"; import { type PcSogsJson, tryPcSogsZip } from "./SplatLoader"; import { computeMaxSplats, - encodeSh1Rgb, - encodeSh2Rgb, - encodeSh3Rgb, + encodeShRgb, + getShArrayStride, + getShDegrees, setPackedSplatCenter, setPackedSplatQuat, setPackedSplatRgba, @@ -130,9 +130,11 @@ export async function unpackPcSogs( const useSH2 = json.shN.shape[1] >= 27 - 3; const useSH1 = json.shN.shape[1] >= 12 - 3; - if (useSH1) extra.sh1 = new Uint32Array(numSplats * 2); - if (useSH2) extra.sh2 = new Uint32Array(numSplats * 4); - if (useSH3) extra.sh3 = new Uint32Array(numSplats * 4); + const shDegrees = getShDegrees(useSH1, useSH2, useSH3, splatEncoding); + if (shDegrees > 0) { + extra.sh = new Uint8Array(numSplats * getShArrayStride(shDegrees)); + extra.shDegrees = shDegrees; + } const sh1 = new Float32Array(9); const sh2 = new Float32Array(15); @@ -181,12 +183,15 @@ export async function unpackPcSogs( } } - if (useSH1) - encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1, splatEncoding); - if (useSH2) - encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2, splatEncoding); - if (useSH3) - encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3, splatEncoding); + encodeShRgb( + extra.sh as Uint8Array, + shDegrees, + i, + sh1, + sh2, + sh3, + splatEncoding, + ); } }); promises.push(shNPromise); diff --git a/src/utils.ts b/src/utils.ts index f73e317..13a06fc 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -1130,137 +1130,88 @@ export function decodeQuatEulerXyz888( return out; } -// Pack four signed 8-bit values into a single uint32. -function packSint8Bytes( - b0: number, - b1: number, - b2: number, - b3: number, -): number { - const clampedB0 = Math.max(-127, Math.min(127, b0 * 127)); - const clampedB1 = Math.max(-127, Math.min(127, b1 * 127)); - const clampedB2 = Math.max(-127, Math.min(127, b2 * 127)); - const clampedB3 = Math.max(-127, Math.min(127, b3 * 127)); - return ( - (clampedB0 & 0xff) | - ((clampedB1 & 0xff) << 8) | - ((clampedB2 & 0xff) << 16) | - ((clampedB3 & 0xff) << 24) - ); -} - -// Encode an array of 9 signed RGB SH1 coefficients (clamped to [-1,1]) into -// a pair of uint32 values, where each coefficient is stored as a sint7 -export function encodeSh1Rgb( - sh1Array: Uint32Array, - index: number, - sh1Rgb: Float32Array, +export function getShDegrees( + hasSh1?: Float32Array | boolean, + hasSh2?: Float32Array | boolean, + hasSh3?: Float32Array | boolean, encoding?: { - sh1Min?: number; - sh1Max?: number; + maxSh?: number; }, -) { - const sh1Min = encoding?.sh1Min ?? -1; - const sh1Max = encoding?.sh1Max ?? 1; - const sh1Mid = 0.5 * (sh1Min + sh1Max); - const sh1Scale = 126 / (sh1Max - sh1Min); - - // Pack sint7 values into 2 x uint32 - const base = index * 2; - for (let i = 0; i < 9; ++i) { - const s = (sh1Rgb[i] - sh1Mid) * sh1Scale; - const value = Math.round(Math.max(-63, Math.min(63, s))) & 0x7f; - const bitStart = i * 7; - const bitEnd = bitStart + 7; - - const wordStart = Math.floor(bitStart / 32); - const bitOffset = bitStart - wordStart * 32; - const firstWord = (value << bitOffset) & 0xffffffff; - sh1Array[base + wordStart] |= firstWord; - - if (bitEnd > wordStart * 32 + 32) { - const secondWord = (value >>> (32 - bitOffset)) & 0xffffffff; - sh1Array[base + wordStart + 1] |= secondWord; - } - } +): 0 | 1 | 2 | 3 { + return Math.min( + encoding?.maxSh ?? 3, + hasSh3 ? 3 : hasSh2 ? 2 : hasSh1 ? 1 : 0, + ) as 0 | 1 | 2 | 3; +} + +export function getShArrayStride(shDegrees: 0 | 1 | 2 | 3): number { + return [0, 12, 24, 48][shDegrees]; } -// Encode an array of 15 signed RGB SH2 coefficients (clamped to [-1,1]) into -// an array of 4 uint32 values, where each coefficient is stored as a sint8. -export function encodeSh2Rgb( - sh2Array: Uint32Array, +export function encodeShRgb( + shArray: Uint8Array, + shDegrees: 0 | 1 | 2 | 3, index: number, - sh2Rgb: Float32Array, + sh1Rgb?: Float32Array, + sh2Rgb?: Float32Array, + sh3Rgb?: Float32Array, encoding?: { + sh1Min?: number; + sh1Max?: number; sh2Min?: number; sh2Max?: number; + sh3Min?: number; + sh3Max?: number; + maxSh?: number; }, ) { - const sh2Min = encoding?.sh2Min ?? -1; - const sh2Max = encoding?.sh2Max ?? 1; - const sh2Mid = 0.5 * (sh2Min + sh2Max); - const sh2Scale = 2 / (sh2Max - sh2Min); - - // Pack sint8 values into 4 x uint32 - sh2Array[index * 4 + 0] = packSint8Bytes( - (sh2Rgb[0] - sh2Mid) * sh2Scale, - (sh2Rgb[1] - sh2Mid) * sh2Scale, - (sh2Rgb[2] - sh2Mid) * sh2Scale, - (sh2Rgb[3] - sh2Mid) * sh2Scale, - ); - sh2Array[index * 4 + 1] = packSint8Bytes( - (sh2Rgb[4] - sh2Mid) * sh2Scale, - (sh2Rgb[5] - sh2Mid) * sh2Scale, - (sh2Rgb[6] - sh2Mid) * sh2Scale, - (sh2Rgb[7] - sh2Mid) * sh2Scale, - ); - sh2Array[index * 4 + 2] = packSint8Bytes( - (sh2Rgb[8] - sh2Mid) * sh2Scale, - (sh2Rgb[9] - sh2Mid) * sh2Scale, - (sh2Rgb[10] - sh2Mid) * sh2Scale, - (sh2Rgb[11] - sh2Mid) * sh2Scale, - ); - sh2Array[index * 4 + 3] = packSint8Bytes( - (sh2Rgb[12] - sh2Mid) * sh2Scale, - (sh2Rgb[13] - sh2Mid) * sh2Scale, - (sh2Rgb[14] - sh2Mid) * sh2Scale, - 0, - ); + const coefficientsWithPadding = getShArrayStride(shDegrees); + if (sh1Rgb && shDegrees >= 1) { + encodeShCoefficientsRgb( + shArray, + index * coefficientsWithPadding, + sh1Rgb, + encoding?.sh1Min ?? -1, + encoding?.sh1Max ?? 1, + ); + } + if (sh2Rgb && shDegrees >= 2) { + encodeShCoefficientsRgb( + shArray, + index * coefficientsWithPadding + 9, + sh2Rgb, + encoding?.sh2Min ?? -1, + encoding?.sh2Max ?? 1, + ); + } + if (sh3Rgb && shDegrees >= 3) { + encodeShCoefficientsRgb( + shArray, + index * coefficientsWithPadding + 24, + sh3Rgb, + encoding?.sh3Min ?? -1, + encoding?.sh3Max ?? 1, + ); + } } -// Encode an array of 21 signed RGB SH3 coefficients (clamped to [-1,1]) into -// an array of 4 uint32 values, where each coefficient is stored as a sint6. -export function encodeSh3Rgb( - sh3Array: Uint32Array, +// Encode an array of signed RGB SH coefficients (clamped to [-1,1]) into +// an uint8 array, where each coefficient is stored as a sint8 +function encodeShCoefficientsRgb( + shArray: Uint8Array, index: number, - sh3Rgb: Float32Array, - encoding?: { - sh3Min?: number; - sh3Max?: number; - }, + shRgb: Float32Array, + shMin: number, + shMax: number, ) { - const sh3Min = encoding?.sh3Min ?? -1; - const sh3Max = encoding?.sh3Max ?? 1; - const sh3Mid = 0.5 * (sh3Min + sh3Max); - const sh3Scale = 62 / (sh3Max - sh3Min); - - // Pack sint6 values into 4 x uint32 - const base = index * 4; - for (let i = 0; i < 21; ++i) { - const s = (sh3Rgb[i] - sh3Mid) * sh3Scale; - const value = Math.round(Math.max(-31, Math.min(31, s))) & 0x3f; - const bitStart = i * 6; - const bitEnd = bitStart + 6; - - const wordStart = Math.floor(bitStart / 32); - const bitOffset = bitStart - wordStart * 32; - const firstWord = (value << bitOffset) & 0xffffffff; - sh3Array[base + wordStart] |= firstWord; - - if (bitEnd > wordStart * 32 + 32) { - const secondWord = (value >>> (32 - bitOffset)) & 0xffffffff; - sh3Array[base + wordStart + 1] |= secondWord; - } + const sh1Mid = 0.5 * (shMin + shMax); + const sh1Scale = 2 / (shMax - shMin); + + for (let i = 0; i < shRgb.length; ++i) { + shArray[index + i] = Math.max( + -127, + Math.min(127, (shRgb[i] - sh1Mid) * sh1Scale * 127), + ); } } diff --git a/src/worker.ts b/src/worker.ts index a7ff2fa..1f05331 100644 --- a/src/worker.ts +++ b/src/worker.ts @@ -9,10 +9,10 @@ import { PlyReader } from "./ply"; import { SpzReader, transcodeSpz } from "./spz"; import { computeMaxSplats, - encodeSh1Rgb, - encodeSh2Rgb, - encodeSh3Rgb, + encodeShRgb, getArrayBuffers, + getShArrayStride, + getShDegrees, setPackedSplat, setPackedSplatCenter, setPackedSplatOpacity, @@ -355,24 +355,25 @@ async function unpackPly({ ); }, (index, sh1, sh2, sh3) => { - if (sh1) { - if (!extra.sh1) { - extra.sh1 = new Uint32Array(numSplats * 2); - } - encodeSh1Rgb(extra.sh1 as Uint32Array, index, sh1, splatEncoding); - } - if (sh2) { - if (!extra.sh2) { - extra.sh2 = new Uint32Array(numSplats * 4); - } - encodeSh2Rgb(extra.sh2 as Uint32Array, index, sh2, splatEncoding); + const shDegrees = getShDegrees(sh1, sh2, sh3, splatEncoding); + if (shDegrees === 0) { + return; } - if (sh3) { - if (!extra.sh3) { - extra.sh3 = new Uint32Array(numSplats * 4); - } - encodeSh3Rgb(extra.sh3 as Uint32Array, index, sh3, splatEncoding); + + if (!extra.sh) { + extra.sh = new Uint8Array(numSplats * getShArrayStride(shDegrees)); + extra.shDegrees = shDegrees; } + + encodeShRgb( + extra.sh as Uint8Array, + shDegrees, + index, + sh1, + sh2, + sh3, + splatEncoding, + ); }, ); @@ -417,24 +418,25 @@ function unpackSpz( setPackedSplatQuat(packedArray, index, quatX, quatY, quatZ, quatW); }, (index, sh1, sh2, sh3) => { - if (sh1) { - if (!extra.sh1) { - extra.sh1 = new Uint32Array(numSplats * 2); - } - encodeSh1Rgb(extra.sh1 as Uint32Array, index, sh1, splatEncoding); - } - if (sh2) { - if (!extra.sh2) { - extra.sh2 = new Uint32Array(numSplats * 4); - } - encodeSh2Rgb(extra.sh2 as Uint32Array, index, sh2, splatEncoding); + const shDegrees = getShDegrees(sh1, sh2, sh3, splatEncoding); + if (shDegrees === 0) { + return; } - if (sh3) { - if (!extra.sh3) { - extra.sh3 = new Uint32Array(numSplats * 4); - } - encodeSh3Rgb(extra.sh3 as Uint32Array, index, sh3, splatEncoding); + + if (!extra.sh) { + extra.sh = new Uint8Array(numSplats * getShArrayStride(shDegrees)); + extra.shDegrees = shDegrees; } + + encodeShRgb( + extra.sh as Uint8Array, + shDegrees, + index, + sh1, + sh2, + sh3, + splatEncoding, + ); }, ); return { packedArray, numSplats, extra };