Gaussian Splatting Shadows

packages/dev/core/src/Materials/GaussianSplatting/gaussianSplattingMaterial.ts

@@ -4,19 +4,25 @@ import type { Mesh } from "../../Meshes/mesh";
 import type { Effect, IEffectCreationOptions } from "../../Materials/effect";
 import type { Scene } from "../../scene";
 import type { Matrix } from "../../Maths/math.vector";
-import type { GaussianSplattingMesh } from "core/Meshes";
+import type { GaussianSplattingMesh } from "../../Meshes";
 import { SerializationHelper } from "../../Misc/decorators.serialization";
 import { VertexBuffer } from "../../Buffers/buffer";
 import { MaterialDefines } from "../../Materials/materialDefines";
 import { PushMaterial } from "../../Materials/pushMaterial";
 import { RegisterClass } from "../../Misc/typeStore";
 import { AddClipPlaneUniforms, BindClipPlane } from "../clipPlaneMaterialHelper";
-import { Camera } from "core/Cameras/camera";
+import { Camera } from "../../Cameras/camera";
+import { ShadowDepthWrapper } from "../../Materials/shadowDepthWrapper";
+import { ShaderMaterial } from "../../Materials/shaderMaterial";
 
 import "../../Shaders/gaussianSplatting.fragment";
 import "../../Shaders/gaussianSplatting.vertex";
 import "../../ShadersWGSL/gaussianSplatting.fragment";
 import "../../ShadersWGSL/gaussianSplatting.vertex";
+import "../../Shaders/gaussianSplattingDepth.fragment";
+import "../../Shaders/gaussianSplattingDepth.vertex";
+import "../../ShadersWGSL/gaussianSplattingDepth.fragment";
+import "../../ShadersWGSL/gaussianSplattingDepth.vertex";
 import {
     BindFogParameters,
     BindLogDepth,
@@ -67,6 +73,7 @@ export class GaussianSplattingMaterial extends PushMaterial {
         super(name, scene);
 
         this.backFaceCulling = false;
+        this.shadowDepthWrapper = GaussianSplattingMaterial._MakeGaussianSplattingShadowDepthWrapper(scene!, this.shaderLanguage);
     }
 
     /**
@@ -126,6 +133,23 @@ export class GaussianSplattingMaterial extends PushMaterial {
         return true;
     }
 
+    protected static _Attribs = [VertexBuffer.PositionKind, "splatIndex"];
+    protected static _Samplers = ["covariancesATexture", "covariancesBTexture", "centersTexture", "colorsTexture", "shTexture0", "shTexture1", "shTexture2"];
+    protected static _UniformBuffers = ["Scene", "Mesh"];
+    protected static _Uniforms = [
+        "world",
+        "view",
+        "projection",
+        "vFogInfos",
+        "vFogColor",
+        "logarithmicDepthConstant",
+        "invViewport",
+        "dataTextureSize",
+        "focal",
+        "eyePosition",
+        "kernelSize",
+        "viewDirectionFactor",
+    ];
     /**
      * Checks whether the material is ready to be rendered for a given mesh.
      * @param mesh The mesh to render
@@ -197,44 +221,25 @@ export class GaussianSplattingMaterial extends PushMaterial {
             scene.resetCachedMaterial();
 
             //Attributes
-            const attribs = [VertexBuffer.PositionKind, "splatIndex"];
-
-            PrepareAttributesForInstances(attribs, defines);
-
-            const uniforms = [
-                "world",
-                "view",
-                "projection",
-                "vFogInfos",
-                "vFogColor",
-                "logarithmicDepthConstant",
-                "invViewport",
-                "dataTextureSize",
-                "focal",
-                "eyePosition",
-                "kernelSize",
-                "viewDirectionFactor",
-            ];
-            const samplers = ["covariancesATexture", "covariancesBTexture", "centersTexture", "colorsTexture", "shTexture0", "shTexture1", "shTexture2"];
-            const uniformBuffers = ["Scene", "Mesh"];
+            PrepareAttributesForInstances(GaussianSplattingMaterial._Attribs, defines);
 
             PrepareUniformsAndSamplersList(<IEffectCreationOptions>{
-                uniformsNames: uniforms,
-                uniformBuffersNames: uniformBuffers,
-                samplers: samplers,
+                uniformsNames: GaussianSplattingMaterial._Uniforms,
+                uniformBuffersNames: GaussianSplattingMaterial._UniformBuffers,
+                samplers: GaussianSplattingMaterial._Samplers,
                 defines: defines,
             });
 
-            AddClipPlaneUniforms(uniforms);
+            AddClipPlaneUniforms(GaussianSplattingMaterial._Uniforms);
 
             const join = defines.toString();
             const effect = scene.getEngine().createEffect(
                 "gaussianSplatting",
                 <IEffectCreationOptions>{
-                    attributes: attribs,
-                    uniformsNames: uniforms,
-                    uniformBuffersNames: uniformBuffers,
-                    samplers: samplers,
+                    attributes: GaussianSplattingMaterial._Attribs,
+                    uniformsNames: GaussianSplattingMaterial._Uniforms,
+                    uniformBuffersNames: GaussianSplattingMaterial._UniformBuffers,
+                    samplers: GaussianSplattingMaterial._Samplers,
                     defines: join,
                     onCompiled: this.onCompiled,
                     onError: this.onError,
@@ -374,6 +379,61 @@ export class GaussianSplattingMaterial extends PushMaterial {
         this._afterBind(mesh, this._activeEffect, subMesh);
     }
 
+    protected static _MakeGaussianSplattingShadowDepthWrapper(scene: Scene, shaderLanguage: ShaderLanguage): ShadowDepthWrapper {
+        const shaderMaterial = new ShaderMaterial(
+            "gaussianSplattingDepth",
+            scene,
+            {
+                vertex: "gaussianSplattingDepth",
+                fragment: "gaussianSplattingDepth",
+            },
+            {
+                attributes: GaussianSplattingMaterial._Attribs,
+                uniforms: GaussianSplattingMaterial._Uniforms,
+                samplers: GaussianSplattingMaterial._Samplers,
+                uniformBuffers: GaussianSplattingMaterial._UniformBuffers,
+                shaderLanguage: shaderLanguage,
+                defines: ["#define GS_DISABLE_COLOR"],
+            }
+        );
+
+        const shadowDepthWrapper = new ShadowDepthWrapper(shaderMaterial, scene, {
+            standalone: true,
+        });
+
+        shaderMaterial.onBindObservable.add((mesh: AbstractMesh) => {
+            const effect = shaderMaterial.getEffect()!;
+            const gsMaterial = mesh.material as GaussianSplattingMaterial;
+            const gsMesh = mesh as GaussianSplattingMesh;
+
+            mesh.getMeshUniformBuffer().bindToEffect(effect, "Mesh");
+            shaderMaterial.bindView(effect);
+            shaderMaterial.bindViewProjection(effect);
+
+            const shadowmapWidth = scene.getEngine().getRenderWidth();
+            const shadowmapHeight = scene.getEngine().getRenderHeight();
+            effect.setFloat2("invViewport", 1 / shadowmapWidth, 1 / shadowmapHeight);
+
+            const projection = scene.getProjectionMatrix();
+            const t = projection.m[5];
+            const focal = (shadowmapWidth * t) / 2.0;
+
+            effect.setFloat2("focal", focal, focal);
+            effect.setFloat("kernelSize", gsMaterial && gsMaterial.kernelSize ? gsMaterial.kernelSize : GaussianSplattingMaterial.KernelSize);
+
+            if (gsMesh.covariancesATexture) {
+                const textureSize = gsMesh.covariancesATexture.getSize();
+                effect.setFloat2("dataTextureSize", textureSize.width, textureSize.height);
+
+                effect.setTexture("covariancesATexture", gsMesh.covariancesATexture);
+                effect.setTexture("covariancesBTexture", gsMesh.covariancesBTexture);
+                effect.setTexture("centersTexture", gsMesh.centersTexture);
+            }
+        });
+
+        return shadowDepthWrapper;
+    }
+
     /**
      * Clones the material.
      * @param name The cloned name.

packages/dev/core/src/Shaders/ShadersInclude/gaussianSplatting.fx

@@ -22,7 +22,9 @@ ivec2 getDataUVint(float index, vec2 textureSize) {
 
 struct Splat {
     vec4 center;
+#ifndef GS_DISABLE_COLOR
     vec4 color;
+#endif
     vec4 covA;
     vec4 covB;
 #if SH_DEGREE > 0
@@ -41,7 +43,9 @@ Splat readSplat(float splatIndex)
     Splat splat;
     vec2 splatUV = getDataUV(splatIndex, dataTextureSize);
     splat.center = texture2D(centersTexture, splatUV);
+#ifndef GS_DISABLE_COLOR
     splat.color = texture2D(colorsTexture, splatUV);
+#endif
     splat.covA = texture2D(covariancesATexture, splatUV) * splat.center.w;
     splat.covB = texture2D(covariancesBTexture, splatUV) * splat.center.w;
 #if SH_DEGREE > 0

packages/dev/core/src/Shaders/gaussianSplattingDepth.fragment.fx

@@ -0,0 +1,6 @@
+precision highp float;
+varying vec2 vPosition;
+void main(void) {
+    float A = -dot(vPosition, vPosition);
+    if (A < -1.) discard;
+}

packages/dev/core/src/Shaders/gaussianSplattingDepth.vertex.fx

@@ -0,0 +1,23 @@
+#include<__decl__gaussianSplattingVertex>
+attribute float splatIndex;
+
+uniform vec2 invViewport;
+uniform vec2 dataTextureSize;
+uniform vec2 focal;
+uniform float kernelSize;
+
+uniform sampler2D covariancesATexture;
+uniform sampler2D covariancesBTexture;
+uniform sampler2D centersTexture;
+
+varying vec2 vPosition;
+#include<gaussianSplatting>
+
+void main(void) {
+    Splat splat = readSplat(splatIndex);
+    vec3 covA = splat.covA.xyz;
+    vec3 covB = vec3(splat.covA.w, splat.covB.xy);
+    vec4 worldPosGS = world * vec4(splat.center.xyz, 1.0);
+    vPosition = position.xy;
+    gl_Position = gaussianSplatting(position.xy, worldPosGS.xyz, vec2(1.,1.), covA, covB, world, view, projection);
+}

packages/dev/core/src/ShadersWGSL/ShadersInclude/gaussianSplatting.fx

@@ -6,7 +6,9 @@ fn getDataUV(index: f32, dataTextureSize: vec2f) -> vec2<f32> {
 
 struct Splat {
     center: vec4f,
+#ifndef GS_DISABLE_COLOR
     color: vec4f,
+#endif
     covA: vec4f,
     covB: vec4f,
 #if SH_DEGREE > 0
@@ -25,7 +27,9 @@ fn readSplat(splatIndex: f32, dataTextureSize: vec2f) -> Splat {
     let splatUV = getDataUV(splatIndex, dataTextureSize);
     let splatUVi32 = vec2<i32>(i32(splatUV.x), i32(splatUV.y));
     splat.center = textureLoad(centersTexture, splatUVi32, 0);
+#ifndef GS_DISABLE_COLOR
     splat.color = textureLoad(colorsTexture, splatUVi32, 0);
+#endif
     splat.covA = textureLoad(covariancesATexture, splatUVi32, 0) * splat.center.w;
     splat.covB = textureLoad(covariancesBTexture, splatUVi32, 0) * splat.center.w;
 #if SH_DEGREE > 0

packages/dev/core/src/ShadersWGSL/gaussianSplattingDepth.fragment.fx

@@ -0,0 +1,15 @@
+#include<gaussianSplattingFragmentDeclaration>
+varying vPosition: vec2f;
+
+// move discard logic to a function to avoid early return issues and parsing/compilation errors with last '}
+fn checkDiscard(inPosition: vec2f) -> vec4f {
+    var A : f32 = -dot(inPosition, inPosition);
+    if (A < -1.) {
+        discard;
+    }
+    return vec4f(0.0);
+}
+@fragment
+fn main(input: FragmentInputs) -> FragmentOutputs {
+    fragmentOutputs.color = checkDiscard(fragmentInputs.vPosition);
+}

packages/dev/core/src/ShadersWGSL/gaussianSplattingDepth.vertex.fx

@@ -0,0 +1,27 @@
+#include<sceneUboDeclaration>
+#include<meshUboDeclaration>
+attribute splatIndex: f32;
+attribute position: vec2f;
+
+uniform invViewport: vec2f;
+uniform dataTextureSize: vec2f;
+uniform focal: vec2f;
+uniform kernelSize: f32;
+
+var covariancesATexture: texture_2d<f32>;
+var covariancesBTexture: texture_2d<f32>;
+var centersTexture: texture_2d<f32>;
+
+varying vPosition: vec2f;
+
+#include<gaussianSplatting>
+
+@vertex
+fn main(input : VertexInputs) -> FragmentInputs {
+    var splat: Splat = readSplat(input.splatIndex, uniforms.dataTextureSize);
+    var covA: vec3f = splat.covA.xyz;
+    var covB: vec3f = vec3f(splat.covA.w, splat.covB.xy);
+    let worldPos: vec4f = mesh.world * vec4f(splat.center.xyz, 1.0);
+    vertexOutputs.vPosition = input.position;
+    vertexOutputs.position = gaussianSplatting(input.position, worldPos.xyz, vec2f(1.0, 1.0), covA, covB, mesh.world, scene.view, scene.projection, uniforms.focal, uniforms.invViewport, uniforms.kernelSize);
+}

packages/tools/tests/test/visualization/config.json

@@ -1,6 +1,11 @@
 {
     "root": "https://cdn.babylonjs.com",
     "tests": [
+        {
+            "title": "Gaussian Splatting Shadows",
+            "playgroundId": "#OE54M5#16",
+            "referenceImage": "gsplat-shadows.png"
+        },
         {
             "title": "NPE - Angle align",
             "playgroundId": "#H5RP91",