diff --git a/src/framework/components/gsplat/component.js b/src/framework/components/gsplat/component.js index 67e3276bc98..ff79ef2161d 100644 --- a/src/framework/components/gsplat/component.js +++ b/src/framework/components/gsplat/component.js @@ -61,6 +61,9 @@ class GSplatComponent extends Component { */ _materialTmp = null; + /** @private */ + _highQualitySH = false; + /** * @type {BoundingBox|null} * @private @@ -206,6 +209,37 @@ class GSplatComponent extends Component { return this._instance?.material ?? this._materialTmp ?? null; } + /** + * Sets whether to use the high quality or the approximate (but fast) spherical-harmonic calculation when rendering SOGS data. + * + * The low quality approximation evaluates the scene's spherical harmonic contributions + * along the camera's Z-axis instead of using each gaussian's view vector. This results + * in gaussians being accurate at the center of the screen and becoming less accurate + * as they appear further from the center. This is a good trade-off for performance + * when rendering large SOGS datasets, especially on mobile devices. + * + * Defaults to false. + * + * @type {boolean} + */ + set highQualitySH(value) { + if (value !== this._highQualitySH) { + this._highQualitySH = value; + if (this._instance) { + this._instance.setHighQualitySH(value); + } + } + } + + /** + * Gets whether the high quality (true) or the fast approximate (false) spherical-harmonic calculation is used when rendering SOGS data. + * + * @type {boolean} + */ + get highQualitySH() { + return this._highQualitySH; + } + /** * Sets whether gsplat will cast shadows for lights that have shadow casting enabled. Defaults * to false. @@ -469,7 +503,10 @@ class GSplatComponent extends Component { // create new instance const asset = this._assetReference.asset; if (asset) { - this.instance = new GSplatInstance(asset.resource, this._materialTmp); + this.instance = new GSplatInstance(asset.resource, { + material: this._materialTmp, + highQualitySH: this._highQualitySH + }); this._materialTmp = null; this.customAabb = this.instance.resource.aabb.clone(); } diff --git a/src/framework/components/gsplat/system.js b/src/framework/components/gsplat/system.js index 434d3c69dee..4e0fe162dc3 100644 --- a/src/framework/components/gsplat/system.js +++ b/src/framework/components/gsplat/system.js @@ -16,9 +16,10 @@ const _schema = [ // order matters here const _properties = [ 'castShadows', + 'material', + 'highQualitySH', 'asset', - 'layers', - 'material' + 'layers' ]; /** diff --git a/src/scene/gsplat/gsplat-instance.js b/src/scene/gsplat/gsplat-instance.js index b39b75dd047..552fc15bf7f 100644 --- a/src/scene/gsplat/gsplat-instance.js +++ b/src/scene/gsplat/gsplat-instance.js @@ -2,7 +2,9 @@ import { Mat4 } from '../../core/math/mat4.js'; import { Vec3 } from '../../core/math/vec3.js'; import { CULLFACE_NONE, SEMANTIC_ATTR13, SEMANTIC_POSITION, PIXELFORMAT_R32U } from '../../platform/graphics/constants.js'; import { MeshInstance } from '../mesh-instance.js'; +import { GSplatResolveSH } from './gsplat-resolve-sh.js'; import { GSplatSorter } from './gsplat-sorter.js'; +import { GSplatSogsData } from './gsplat-sogs-data.js'; import { ShaderMaterial } from '../materials/shader-material.js'; import { BLEND_NONE, BLEND_PREMULTIPLIED } from '../constants.js'; @@ -34,13 +36,16 @@ class GSplatInstance { options = {}; - /** @type {GSplatSorter | null} */ + /** @type {GSplatSorter|null} */ sorter = null; lastCameraPosition = new Vec3(); lastCameraDirection = new Vec3(); + /** @type {GSplatResolveSH|null} */ + resolveSH = null; + /** * List of cameras this instance is visible for. Updated every frame by the renderer. * @@ -51,9 +56,11 @@ class GSplatInstance { /** * @param {GSplatResourceBase} resource - The splat instance. - * @param {ShaderMaterial|null} material - The material instance. + * @param {object} [options] - Options for the instance. + * @param {ShaderMaterial|null} [options.material] - The material instance. + * @param {boolean} [options.highQualitySH] - Whether to use the high quality or the approximate spherical harmonic calculation. Only applies to SOGS data. */ - constructor(resource, material) { + constructor(resource, options = {}) { this.resource = resource; // create the order texture @@ -63,9 +70,9 @@ class GSplatInstance { resource.evalTextureSize(resource.numSplats) ); - if (material) { + if (options.material) { // material is provided - this._material = material; + this._material = options.material; // patch splat order this._material.setParameter('splatOrder', this.orderTexture); @@ -111,9 +118,13 @@ class GSplatInstance { // update splat count on the material this.material.setParameter('numSplats', count); }); + + // configure sogs sh resolve + this.setHighQualitySH(options.highQualitySH ?? false); } destroy() { + this.resolveSH?.destroy(); this.material?.destroy(); this.meshInstance?.destroy(); this.sorter?.destroy(); @@ -212,10 +223,30 @@ class GSplatInstance { const camera = this.cameras[0]; this.sort(camera._node); + // resolve spherical harmonics + this.resolveSH?.render(camera._node, this.meshInstance.node.getWorldTransform()); + // we get new list of cameras each frame this.cameras.length = 0; } } + + setHighQualitySH(value) { + const { resource } = this; + const { gsplatData } = resource; + + if (gsplatData instanceof GSplatSogsData && + gsplatData.shBands > 0 && + value === !!this.resolveSH) { + + if (this.resolveSH) { + this.resolveSH.destroy(); + this.resolveSH = null; + } else { + this.resolveSH = new GSplatResolveSH(resource.device, this); + } + } + } } export { GSplatInstance }; diff --git a/src/scene/gsplat/gsplat-resolve-sh.js b/src/scene/gsplat/gsplat-resolve-sh.js new file mode 100644 index 00000000000..769db1381b1 --- /dev/null +++ b/src/scene/gsplat/gsplat-resolve-sh.js @@ -0,0 +1,305 @@ +import { Vec3 } from '../../core/math/vec3.js'; +import { Mat4 } from '../../core/math/mat4.js'; +import { BlendState } from '../../platform/graphics/blend-state.js'; +import { + CULLFACE_NONE, + PIXELFORMAT_RGBA8, + SEMANTIC_POSITION +} from '../../platform/graphics/constants.js'; +import { DepthState } from '../../platform/graphics/depth-state.js'; +import { RenderTarget } from '../../platform/graphics/render-target.js'; +import { ShaderUtils } from '../shader-lib/shader-utils.js'; +import { ShaderChunks } from '../shader-lib/shader-chunks.js'; +import { Vec2 } from '../../core/math/vec2.js'; +import { RenderPass } from '../../platform/graphics/render-pass.js'; +import { QuadRender } from '../graphics/quad-render.js'; + +const vertexGLSL = /* glsl */` + attribute vec2 vertex_position; + void main(void) { + gl_Position = vec4(vertex_position, 0.0, 1.0); + } +`; + +const fragmentGLSL = /* glsl */` + #include "gsplatEvalSHVS" + + // takes a normalized 3-component value, convert to (11, 11, 10) bit range and + // then package into RGBA8 + vec4 packRgb(vec3 v) { + uvec3 vb = uvec3(clamp(v, vec3(0.0), vec3(1.0)) * vec3(2047.0, 2047.0, 1023.0)); + uint bits = (vb.x << 21) | (vb.y << 10) | vb.z; + return vec4((uvec4(bits) >> uvec4(24, 16, 8, 0)) & uvec4(0xff)) / vec4(255.0); + } + + uniform mediump vec3 dir; + uniform mediump sampler2D centroids; + uniform mediump float shN_mins; + uniform mediump float shN_maxs; + + void main(void) { + ivec2 uv = ivec2(gl_FragCoord.xy) * ivec2(SH_COEFFS, 1); + mediump vec3 coefficients[SH_COEFFS]; + + // read coefficients + for (int i = 0; i < SH_COEFFS; i++) { + vec3 s = texelFetch(centroids, ivec2(uv.x + i, uv.y), 0).xyz; + coefficients[i] = mix(vec3(shN_mins), vec3(shN_maxs), s); + } + + // evaluate + gl_FragColor = packRgb(evalSH(coefficients, dir) * 0.25 + 0.5); + } +`; + +const vertexWGSL = /* wgsl */` + attribute vertex_position: vec2f; + @vertex + fn vertexMain(input: VertexInput) -> VertexOutput { + var output: VertexOutput; + output.position = vec4f(vertex_position, 0.0, 1.0); + return output; + } +`; + +const fragmentWGSL = /* wgsl */` + #include "gsplatEvalSHVS" + + // takes a normalized 3-component value, convert to (11, 11, 10) bit range and + // then package into RGBA8 + fn packRgb(v: vec3f) -> vec4f { + let vb = vec3u(clamp(v, vec3f(0.0), vec3f(1.0)) * vec3f(2047.0, 2047.0, 1023.0)); + let bits = dot(vb, vec3u(1 << 21, 1 << 10, 1)); + return vec4f((vec4u(bits) >> vec4u(24, 16, 8, 0)) & vec4u(0xff)) / vec4f(255.0); + } + + uniform dir: vec3f; + uniform shN_mins: f32; + uniform shN_maxs: f32; + + var centroids: texture_2d; + + @fragment + fn fragmentMain(input: FragmentInput) -> FragmentOutput { + var output: FragmentOutput; + + var uv = vec2i(input.position.xy) * vec2i(SH_COEFFS, 1); + var coefficients: array; + + // read coefficients + for (var i: i32 = 0; i < SH_COEFFS; i++) { + let s: vec3f = textureLoad(centroids, vec2i(uv.x + i, uv.y), 0).xyz; + coefficients[i] = mix(vec3f(uniform.shN_mins), vec3f(uniform.shN_maxs), s); + } + + // evaluate + output.color = packRgb(evalSH(&coefficients, uniform.dir) * 0.25 + 0.5); + + return output; + } +`; + +const gsplatSogsColorGLSL = /* glsl */` + uniform mediump sampler2D sh0; + uniform highp sampler2D sh_labels; + uniform mediump sampler2D sh_result; + + uniform vec4 sh0_mins; + uniform vec4 sh0_maxs; + + float SH_C0 = 0.28209479177387814; + + // unpack 11, 11, 10 normalized value from rgba8 texture sample + vec3 unpackRgb(vec4 v) { + uvec4 uv = uvec4(v * 255.0); + uint bits = (uv.x << 24) | (uv.y << 16) | (uv.z << 8) | uv.w; + uvec3 vb = (uvec3(bits) >> uvec3(21, 10, 0)) & uvec3(0x7ffu, 0x7ffu, 0x3ffu); + return vec3(vb) / vec3(2047.0, 2047.0, 1023.0); + } + + vec4 readColor(in SplatSource source) { + // sample base color + vec4 baseSample = mix(sh0_mins, sh0_maxs, texelFetch(sh0, source.uv, 0)); + + // resolve base color + vec4 base = vec4(vec3(0.5) + baseSample.xyz * SH_C0, 1.0 / (1.0 + exp(-baseSample.w))); + + // extract spherical harmonics palette index + ivec2 labelSample = ivec2(texelFetch(sh_labels, source.uv, 0).xy * 255.0); + int n = labelSample.x + labelSample.y * 256; + + vec4 shSample = texelFetch(sh_result, ivec2(n % 64, n / 64), 0); + vec3 sh = (unpackRgb(shSample) - vec3(0.5)) * 4.0; + + return vec4(base.xyz + sh, base.w); + } +`; + +const gsplatSogsColorWGSL = /* wgsl */` + var sh0: texture_2d; + var sh_labels: texture_2d; + var sh_result: texture_2d; + + uniform sh0_mins: vec4f; + uniform sh0_maxs: vec4f; + + const SH_C0: f32 = 0.28209479177387814; + + // unpack 11, 11, 10 normalized value from rgba8 texture sample + fn unpackRgb(v: vec4f) -> vec3f { + let bits = dot(vec4u(v * 255.0), vec4u(1u << 24, 1u << 16, 1u << 8, 1u)); + let vb = (vec3u(bits) >> vec3u(21, 10, 0)) & vec3u(0x7ffu, 0x7ffu, 0x3ffu); + return vec3f(vb) / vec3f(2047.0, 2047.0, 1023.0); + } + + fn readColor(source: ptr) -> vec4f { + // sample base color + let baseSample: vec4f = mix(uniform.sh0_mins, uniform.sh0_maxs, textureLoad(sh0, source.uv, 0)); + let base = vec4f(vec3f(0.5) + baseSample.xyz * SH_C0, 1.0 / (1.0 + exp(-baseSample.w))); + + // extract spherical harmonics palette index + let labelSample: vec2i = vec2i(textureLoad(sh_labels, source.uv, 0).xy * 255.0); + let n = labelSample.x + labelSample.y * 256; + + let shSample: vec4f = textureLoad(sh_result, vec2i(n % 64, n / 64), 0); + let sh: vec3f = (unpackRgb(shSample) - vec3f(0.5)) * 4.0; + + return vec4f(base.xyz + sh, base.w); + } +`; + +const resolve = (scope, values) => { + for (const key in values) { + scope.resolve(key).setValue(values[key]); + } +}; + +class CustomRenderPass extends RenderPass { + /** + * @type {() => void | null} + */ + executeCallback = null; + + execute() { + this.executeCallback?.(); + } +} + +const invModelMat = new Mat4(); +const dir = new Vec3(); + +class GSplatResolveSH { + prevDir = new Vec3(); + + updateMode = 'enable'; // 'enable', 'disable', 'always' + + constructor(device, gsplatInstance) { + this.device = device; + this.gsplatInstance = gsplatInstance; + + const { resource } = gsplatInstance; + + const includes = new Map(ShaderChunks.get(device, device.isWebGPU ? 'wgsl' : 'glsl')); + + this.shader = ShaderUtils.createShader(device, { + uniqueName: 'gsplatResolveSH', + vertexGLSL, + fragmentGLSL, + vertexWGSL, + fragmentWGSL, + vertexIncludes: includes, + fragmentIncludes: includes, + fragmentDefines: new Map([ + ['SH_BANDS', resource.gsplatData.shBands.toString()] + ]), + attributes: { + vertex_position: SEMANTIC_POSITION + } + }); + + this.texture = resource.createTexture('centroids', PIXELFORMAT_RGBA8, new Vec2(64, 1024)); + this.renderTarget = new RenderTarget({ + colorBuffer: this.texture, + depth: false + }); + + this.renderPass = new CustomRenderPass(device); + this.renderPass.init(this.renderTarget, {}); + this.renderPass.colorOps.clear = true; + this.quadRender = new QuadRender(this.shader); + + const { material } = gsplatInstance; + material.setDefine('SH_BANDS', '0'); + + const { shaderChunks } = material; + shaderChunks.glsl.set('gsplatSogsColorVS', gsplatSogsColorGLSL); + shaderChunks.wgsl.set('gsplatSogsColorVS', gsplatSogsColorWGSL); + + material.update(); + + device.scope.resolve('sh_result').setValue(this.texture); + } + + destroy() { + const { gsplatInstance } = this; + + const { material } = gsplatInstance; + material.setDefine('SH_BANDS', gsplatInstance.resource.gsplatData.shBands.toString()); + + const { shaderChunks } = material; + shaderChunks.glsl.delete('gsplatSogsColorVS'); + shaderChunks.wgsl.delete('gsplatSogsColorVS'); + + material.update(); + + this.quadRender.destroy(); + this.renderPass.destroy(); + this.renderTarget.destroy(); + this.texture.destroy(); + this.shader.destroy(); + } + + render(camera, modelMat) { + const { prevDir, updateMode } = this; + + // disabled + if (updateMode === 'disable') { + return; + } + + // calculate camera Z in model space + invModelMat.invert(modelMat); + invModelMat.transformVector(camera.forward, dir); + + // if direction hasn't changed early out + dir.normalize(); + if (updateMode === 'enable' && dir.equalsApprox(prevDir, 1e-3)) { + return; + } + prevDir.copy(dir); + + const execute = () => { + const { device } = this; + const { sh_centroids, meta } = this.gsplatInstance.resource.gsplatData; + + resolve(device.scope, { + dir: dir.toArray(), + centroids: sh_centroids, + shN_mins: meta.shN.mins, + shN_maxs: meta.shN.maxs + }); + + device.setCullMode(CULLFACE_NONE); + device.setDepthState(DepthState.NODEPTH); + device.setStencilState(null, null); + device.setBlendState(BlendState.NOBLEND); + + this.quadRender.render(); + }; + + this.renderPass.executeCallback = execute; + this.renderPass.render(); + } +} + +export { GSplatResolveSH }; diff --git a/src/scene/gsplat/gsplat-resource-base.js b/src/scene/gsplat/gsplat-resource-base.js index 8d34c338042..b701a7f26cf 100644 --- a/src/scene/gsplat/gsplat-resource-base.js +++ b/src/scene/gsplat/gsplat-resource-base.js @@ -1,4 +1,5 @@ import { Debug } from '../../core/debug.js'; +import { Vec2 } from '../../core/math/vec2.js'; import { BoundingBox } from '../../core/shape/bounding-box.js'; import { ADDRESS_CLAMP_TO_EDGE, BUFFER_STATIC, FILTER_NEAREST, SEMANTIC_ATTR13, TYPE_UINT32 } from '../../platform/graphics/constants.js'; import { Texture } from '../../platform/graphics/texture.js'; @@ -7,7 +8,6 @@ import { VertexBuffer } from '../../platform/graphics/vertex-buffer.js'; import { Mesh } from '../mesh.js'; /** - * @import { Vec2 } from '../../core/math/vec2.js'; * @import { GraphicsDevice } from '../../platform/graphics/graphics-device.js' * @import { GSplatData } from './gsplat-data.js'; * @import { GSplatCompressedData } from './gsplat-compressed-data.js'; @@ -115,7 +115,14 @@ class GSplatResourceBase { configureMaterial(material) { } + /** + * Evaluates the size of the texture based on the number of splats. + * + * @param {number} count - Number of gaussians. + * @returns {Vec2} Returns a Vec2 object representing the size of the texture. + */ evalTextureSize(count) { + return Vec2.ZERO; } /** diff --git a/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplat.js b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplat.js index dc8197c04e2..16f926a415c 100644 --- a/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplat.js +++ b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplat.js @@ -49,7 +49,14 @@ void main(void) { #if SH_BANDS > 0 // calculate the model-space view direction vec3 dir = normalize(center.view * mat3(center.modelView)); - clr.xyz += evalSH(source, dir); + + // read sh coefficients + vec3 sh[SH_COEFFS]; + float scale; + readSHData(source, sh, scale); + + // evaluate + clr.xyz += evalSH(sh, dir) * scale; #endif clipCorner(corner, clr.w); diff --git a/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatCommon.js b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatCommon.js index 3171b6a0dec..966ca35db3e 100644 --- a/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatCommon.js +++ b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatCommon.js @@ -16,26 +16,6 @@ struct SplatCenter { float projMat00; // element [0][0] of the projection matrix }; -mat3 quatToMat3(vec4 R) { - vec4 R2 = R + R; - float X = R2.x * R.w; - vec4 Y = R2.y * R; - vec4 Z = R2.z * R; - float W = R2.w * R.w; - - return mat3( - 1.0 - Z.z - W, - Y.z + X, - Y.w - Z.x, - Y.z - X, - 1.0 - Y.y - W, - Z.w + Y.x, - Y.w + Z.x, - Z.w - Y.x, - 1.0 - Y.y - Z.z - ); -} - // stores the offset from center for the current gaussian struct SplatCorner { vec2 offset; // corner offset from center in clip space @@ -43,17 +23,13 @@ struct SplatCorner { #if GSPLAT_AA float aaFactor; // for scenes generated with antialiasing #endif + + vec2 v; + float dlen; }; -#if SH_BANDS == 1 - #define SH_COEFFS 3 -#elif SH_BANDS == 2 - #define SH_COEFFS 8 -#elif SH_BANDS == 3 - #define SH_COEFFS 15 -#else - #define SH_COEFFS 0 -#endif +#include "gsplatEvalSHVS" +#include "gsplatQuatToMat3VS" #if GSPLAT_COMPRESSED_DATA == true #include "gsplatCompressedDataVS" @@ -85,77 +61,4 @@ void clipCorner(inout SplatCorner corner, float alpha) { corner.offset *= clip; corner.uv *= clip; } - -// spherical Harmonics - -#if SH_BANDS > 0 - - #define SH_C1 0.4886025119029199f - - #if SH_BANDS > 1 - #define SH_C2_0 1.0925484305920792f - #define SH_C2_1 -1.0925484305920792f - #define SH_C2_2 0.31539156525252005f - #define SH_C2_3 -1.0925484305920792f - #define SH_C2_4 0.5462742152960396f - #endif - - #if SH_BANDS > 2 - #define SH_C3_0 -0.5900435899266435f - #define SH_C3_1 2.890611442640554f - #define SH_C3_2 -0.4570457994644658f - #define SH_C3_3 0.3731763325901154f - #define SH_C3_4 -0.4570457994644658f - #define SH_C3_5 1.445305721320277f - #define SH_C3_6 -0.5900435899266435f - #endif - - // see https://github.com/graphdeco-inria/gaussian-splatting/blob/main/utils/sh_utils.py - vec3 evalSH(in SplatSource source, in vec3 dir) { - - vec3 sh[SH_COEFFS]; - - // read sh coefficients - float scale; - readSHData(source, sh, scale); - - float x = dir.x; - float y = dir.y; - float z = dir.z; - - // 1st degree - vec3 result = SH_C1 * (-sh[0] * y + sh[1] * z - sh[2] * x); - - #if SH_BANDS > 1 - // 2nd degree - float xx = x * x; - float yy = y * y; - float zz = z * z; - float xy = x * y; - float yz = y * z; - float xz = x * z; - - result += - sh[3] * (SH_C2_0 * xy) + - sh[4] * (SH_C2_1 * yz) + - sh[5] * (SH_C2_2 * (2.0 * zz - xx - yy)) + - sh[6] * (SH_C2_3 * xz) + - sh[7] * (SH_C2_4 * (xx - yy)); - #endif - - #if SH_BANDS > 2 - // 3rd degree - result += - sh[8] * (SH_C3_0 * y * (3.0 * xx - yy)) + - sh[9] * (SH_C3_1 * xy * z) + - sh[10] * (SH_C3_2 * y * (4.0 * zz - xx - yy)) + - sh[11] * (SH_C3_3 * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) + - sh[12] * (SH_C3_4 * x * (4.0 * zz - xx - yy)) + - sh[13] * (SH_C3_5 * z * (xx - yy)) + - sh[14] * (SH_C3_6 * x * (xx - 3.0 * yy)); - #endif - - return result * scale; - } -#endif `; diff --git a/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatEvalSH.js b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatEvalSH.js new file mode 100644 index 00000000000..782154e35a0 --- /dev/null +++ b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatEvalSH.js @@ -0,0 +1,75 @@ +export default /* glsl */` + #if SH_BANDS == 1 + #define SH_COEFFS 3 + #elif SH_BANDS == 2 + #define SH_COEFFS 8 + #elif SH_BANDS == 3 + #define SH_COEFFS 15 + #else + #define SH_COEFFS 0 + #endif + + #if SH_BANDS > 0 + + const float SH_C1 = 0.4886025119029199f; + + #if SH_BANDS > 1 + const float SH_C2_0 = 1.0925484305920792f; + const float SH_C2_1 = -1.0925484305920792f; + const float SH_C2_2 = 0.31539156525252005f; + const float SH_C2_3 = -1.0925484305920792f; + const float SH_C2_4 = 0.5462742152960396f; + #endif + + #if SH_BANDS > 2 + const float SH_C3_0 = -0.5900435899266435f; + const float SH_C3_1 = 2.890611442640554f; + const float SH_C3_2 = -0.4570457994644658f; + const float SH_C3_3 = 0.3731763325901154f; + const float SH_C3_4 = -0.4570457994644658f; + const float SH_C3_5 = 1.445305721320277f; + const float SH_C3_6 = -0.5900435899266435f; + #endif + + // see https://github.com/graphdeco-inria/gaussian-splatting/blob/main/utils/sh_utils.py + vec3 evalSH(in vec3 sh[SH_COEFFS], in vec3 dir) { + float x = dir.x; + float y = dir.y; + float z = dir.z; + + // 1st degree + vec3 result = SH_C1 * (-sh[0] * y + sh[1] * z - sh[2] * x); + + #if SH_BANDS > 1 + // 2nd degree + float xx = x * x; + float yy = y * y; + float zz = z * z; + float xy = x * y; + float yz = y * z; + float xz = x * z; + + result += + sh[3] * (SH_C2_0 * xy) + + sh[4] * (SH_C2_1 * yz) + + sh[5] * (SH_C2_2 * (2.0 * zz - xx - yy)) + + sh[6] * (SH_C2_3 * xz) + + sh[7] * (SH_C2_4 * (xx - yy)); + #endif + + #if SH_BANDS > 2 + // 3rd degree + result += + sh[8] * (SH_C3_0 * y * (3.0 * xx - yy)) + + sh[9] * (SH_C3_1 * xy * z) + + sh[10] * (SH_C3_2 * y * (4.0 * zz - xx - yy)) + + sh[11] * (SH_C3_3 * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) + + sh[12] * (SH_C3_4 * x * (4.0 * zz - xx - yy)) + + sh[13] * (SH_C3_5 * z * (xx - yy)) + + sh[14] * (SH_C3_6 * x * (xx - 3.0 * yy)); + #endif + + return result; + } + #endif +`; diff --git a/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatQuatToMat3.js b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatQuatToMat3.js new file mode 100644 index 00000000000..e29cb79796b --- /dev/null +++ b/src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatQuatToMat3.js @@ -0,0 +1,21 @@ +export default /* glsl */` +mat3 quatToMat3(vec4 R) { + vec4 R2 = R + R; + float X = R2.x * R.w; + vec4 Y = R2.y * R; + vec4 Z = R2.z * R; + float W = R2.w * R.w; + + return mat3( + 1.0 - Z.z - W, + Y.z + X, + Y.w - Z.x, + Y.z - X, + 1.0 - Y.y - W, + Z.w + Y.x, + Y.w + Z.x, + Z.w - Y.x, + 1.0 - Y.y - Z.z + ); +} +`; diff --git a/src/scene/shader-lib/glsl/collections/shader-chunks-glsl.js b/src/scene/shader-lib/glsl/collections/shader-chunks-glsl.js index cb7d15a0ab5..38169fbad35 100644 --- a/src/scene/shader-lib/glsl/collections/shader-chunks-glsl.js +++ b/src/scene/shader-lib/glsl/collections/shader-chunks-glsl.js @@ -50,6 +50,8 @@ import gsplatColorVS from '../chunks/gsplat/vert/gsplatColor.js'; import gsplatCommonVS from '../chunks/gsplat/vert/gsplatCommon.js'; import gsplatCompressedDataVS from '../chunks/gsplat/vert/gsplatCompressedData.js'; import gsplatCompressedSHVS from '../chunks/gsplat/vert/gsplatCompressedSH.js'; +import gsplatEvalSHVS from '../chunks/gsplat/vert/gsplatEvalSH.js'; +import gsplatQuatToMat3VS from '../chunks/gsplat/vert/gsplatQuatToMat3.js'; import gsplatSogsColorVS from '../chunks/gsplat/vert/gsplatSogsColor.js'; import gsplatSogsDataVS from '../chunks/gsplat/vert/gsplatSogsData.js'; import gsplatSogsSHVS from '../chunks/gsplat/vert/gsplatSogsSH.js'; @@ -228,6 +230,8 @@ const shaderChunksGLSL = { gsplatCommonVS, gsplatCompressedDataVS, gsplatCompressedSHVS, + gsplatEvalSHVS, + gsplatQuatToMat3VS, gsplatSogsColorVS, gsplatSogsDataVS, gsplatSogsSHVS, diff --git a/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplat.js b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplat.js index 133681931ba..0cc55e13b9c 100644 --- a/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplat.js +++ b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplat.js @@ -53,7 +53,14 @@ fn vertexMain(input: VertexInput) -> VertexOutput { // calculate the model-space view direction let modelView3x3 = mat3x3f(center.modelView[0].xyz, center.modelView[1].xyz, center.modelView[2].xyz); let dir = normalize(modelView3x3 * center.view); - clr = vec4f(clr.xyz + evalSH(&source, dir), clr.a); + + // read sh coefficients + var sh: array; + var scale: f32; + readSHData(&source, &sh, &scale); + + // evaluate + clr = vec4f(clr.xyz + evalSH(&sh, dir) * scale, clr.a); #endif clipCorner(&corner, clr.w); diff --git a/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatCommon.js b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatCommon.js index 6b99e0fac68..d53dfce1f22 100644 --- a/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatCommon.js +++ b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatCommon.js @@ -25,29 +25,8 @@ struct SplatCorner { #endif } -fn quatToMat3(R: vec4) -> mat3x3 { - let R2: vec4 = R + R; - let X: f32 = R2.x * R.w; - let Y: vec4 = R2.y * R; - let Z: vec4 = R2.z * R; - let W: f32 = R2.w * R.w; - - return mat3x3( - 1.0 - Z.z - W, Y.z + X, Y.w - Z.x, - Y.z - X, 1.0 - Y.y - W, Z.w + Y.x, - Y.w + Z.x, Z.w - Y.x, 1.0 - Y.y - Z.z - ); -} - -#if SH_BANDS == 1 - const SH_COEFFS: i32 = 3; -#elif SH_BANDS == 2 - const SH_COEFFS: i32 = 8; -#elif SH_BANDS == 3 - const SH_COEFFS: i32 = 15; -#else - const SH_COEFFS: i32 = 0; -#endif +#include "gsplatEvalSHVS" +#include "gsplatQuatToMat3VS" #if GSPLAT_COMPRESSED_DATA #include "gsplatCompressedDataVS" @@ -80,76 +59,4 @@ fn clipCorner(corner: ptr, alpha: f32) { corner.uv = corner.uv * clip; } -// spherical Harmonics - -#if SH_BANDS > 0 - const SH_C1: f32 = 0.4886025119029199; - - #if SH_BANDS > 1 - const SH_C2_0: f32 = 1.0925484305920792; - const SH_C2_1: f32 = -1.0925484305920792; - const SH_C2_2: f32 = 0.31539156525252005; - const SH_C2_3: f32 = -1.0925484305920792; - const SH_C2_4: f32 = 0.5462742152960396; - #endif - - #if SH_BANDS > 2 - const SH_C3_0: f32 = -0.5900435899266435; - const SH_C3_1: f32 = 2.890611442640554; - const SH_C3_2: f32 = -0.4570457994644658; - const SH_C3_3: f32 = 0.3731763325901154; - const SH_C3_4: f32 = -0.4570457994644658; - const SH_C3_5: f32 = 1.445305721320277; - const SH_C3_6: f32 = -0.5900435899266435; - #endif - - // see https://github.com/graphdeco-inria/gaussian-splatting/blob/main/utils/sh_utils.py - fn evalSH(source: ptr, dir: vec3f) -> vec3f { - - var sh: array; - - var scale: f32; - readSHData(source, &sh, &scale); - - let x = dir.x; - let y = dir.y; - let z = dir.z; - - // 1st degree - var result = SH_C1 * (-sh[0] * y + sh[1] * z - sh[2] * x); - - #if SH_BANDS > 1 - // 2nd degree - let xx = x * x; - let yy = y * y; - let zz = z * z; - let xy = x * y; - let yz = y * z; - let xz = x * z; - - result = result + ( - sh[3] * (SH_C2_0 * xy) + - sh[4] * (SH_C2_1 * yz) + - sh[5] * (SH_C2_2 * (2.0 * zz - xx - yy)) + - sh[6] * (SH_C2_3 * xz) + - sh[7] * (SH_C2_4 * (xx - yy)) - ); - #endif - - #if SH_BANDS > 2 - // 3rd degree - result = result + ( - sh[8] * (SH_C3_0 * y * (3.0 * xx - yy)) + - sh[9] * (SH_C3_1 * xy * z) + - sh[10] * (SH_C3_2 * y * (4.0 * zz - xx - yy)) + - sh[11] * (SH_C3_3 * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) + - sh[12] * (SH_C3_4 * x * (4.0 * zz - xx - yy)) + - sh[13] * (SH_C3_5 * z * (xx - yy)) + - sh[14] * (SH_C3_6 * x * (xx - 3.0 * yy)) - ); - #endif - - return result * scale; - } -#endif `; diff --git a/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatEvalSH.js b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatEvalSH.js new file mode 100644 index 00000000000..18c1e14ecc3 --- /dev/null +++ b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatEvalSH.js @@ -0,0 +1,77 @@ +export default /* wgsl */` + #if SH_BANDS == 1 + const SH_COEFFS: i32 = 3; + #elif SH_BANDS == 2 + const SH_COEFFS: i32 = 8; + #elif SH_BANDS == 3 + const SH_COEFFS: i32 = 15; + #else + const SH_COEFFS: i32 = 0; + #endif + + #if SH_BANDS > 0 + + const SH_C1: f32 = 0.4886025119029199; + + #if SH_BANDS > 1 + const SH_C2_0: f32 = 1.0925484305920792; + const SH_C2_1: f32 = -1.0925484305920792; + const SH_C2_2: f32 = 0.31539156525252005; + const SH_C2_3: f32 = -1.0925484305920792; + const SH_C2_4: f32 = 0.5462742152960396; + #endif + + #if SH_BANDS > 2 + const SH_C3_0: f32 = -0.5900435899266435; + const SH_C3_1: f32 = 2.890611442640554; + const SH_C3_2: f32 = -0.4570457994644658; + const SH_C3_3: f32 = 0.3731763325901154; + const SH_C3_4: f32 = -0.4570457994644658; + const SH_C3_5: f32 = 1.445305721320277; + const SH_C3_6: f32 = -0.5900435899266435; + #endif + + // see https://github.com/graphdeco-inria/gaussian-splatting/blob/main/utils/sh_utils.py + fn evalSH(sh: ptr>, dir: vec3f) -> vec3f { + let x = dir.x; + let y = dir.y; + let z = dir.z; + + // 1st degree + var result = SH_C1 * (-sh[0] * y + sh[1] * z - sh[2] * x); + + #if SH_BANDS > 1 + // 2nd degree + let xx = x * x; + let yy = y * y; + let zz = z * z; + let xy = x * y; + let yz = y * z; + let xz = x * z; + + result = result + ( + sh[3] * (SH_C2_0 * xy) + + sh[4] * (SH_C2_1 * yz) + + sh[5] * (SH_C2_2 * (2.0 * zz - xx - yy)) + + sh[6] * (SH_C2_3 * xz) + + sh[7] * (SH_C2_4 * (xx - yy)) + ); + #endif + + #if SH_BANDS > 2 + // 3rd degree + result = result + ( + sh[8] * (SH_C3_0 * y * (3.0 * xx - yy)) + + sh[9] * (SH_C3_1 * xy * z) + + sh[10] * (SH_C3_2 * y * (4.0 * zz - xx - yy)) + + sh[11] * (SH_C3_3 * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)) + + sh[12] * (SH_C3_4 * x * (4.0 * zz - xx - yy)) + + sh[13] * (SH_C3_5 * z * (xx - yy)) + + sh[14] * (SH_C3_6 * x * (xx - 3.0 * yy)) + ); + #endif + + return result; + } + #endif +`; diff --git a/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatQuatToMat3.js b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatQuatToMat3.js new file mode 100644 index 00000000000..665638fcb53 --- /dev/null +++ b/src/scene/shader-lib/wgsl/chunks/gsplat/vert/gsplatQuatToMat3.js @@ -0,0 +1,15 @@ +export default /* wgsl */` +fn quatToMat3(R: vec4) -> mat3x3 { + let R2: vec4 = R + R; + let X: f32 = R2.x * R.w; + let Y: vec4 = R2.y * R; + let Z: vec4 = R2.z * R; + let W: f32 = R2.w * R.w; + + return mat3x3( + 1.0 - Z.z - W, Y.z + X, Y.w - Z.x, + Y.z - X, 1.0 - Y.y - W, Z.w + Y.x, + Y.w + Z.x, Z.w - Y.x, 1.0 - Y.y - Z.z + ); +} +`; diff --git a/src/scene/shader-lib/wgsl/collections/shader-chunks-wgsl.js b/src/scene/shader-lib/wgsl/collections/shader-chunks-wgsl.js index 90d6c90dfa6..55b671826b2 100644 --- a/src/scene/shader-lib/wgsl/collections/shader-chunks-wgsl.js +++ b/src/scene/shader-lib/wgsl/collections/shader-chunks-wgsl.js @@ -48,6 +48,8 @@ import gsplatColorVS from '../chunks/gsplat/vert/gsplatColor.js'; import gsplatCommonVS from '../chunks/gsplat/vert/gsplatCommon.js'; import gsplatCompressedDataVS from '../chunks/gsplat/vert/gsplatCompressedData.js'; import gsplatCompressedSHVS from '../chunks/gsplat/vert/gsplatCompressedSH.js'; +import gsplatEvalSHVS from '../chunks/gsplat/vert/gsplatEvalSH.js'; +import gsplatQuatToMat3VS from '../chunks/gsplat/vert/gsplatQuatToMat3.js'; import gsplatSogsColorVS from '../chunks/gsplat/vert/gsplatSogsColor.js'; import gsplatSogsDataVS from '../chunks/gsplat/vert/gsplatSogsData.js'; import gsplatSogsSHVS from '../chunks/gsplat/vert/gsplatSogsSH.js'; @@ -225,6 +227,8 @@ const shaderChunksWGSL = { gsplatCommonVS, gsplatCompressedDataVS, gsplatCompressedSHVS, + gsplatEvalSHVS, + gsplatQuatToMat3VS, gsplatSogsColorVS, gsplatSogsDataVS, gsplatSogsSHVS,