Expose modifyCovariance and unify gsplat shader customization API (playcanvas#8064)

* Expose modifyCovariance and unify gsplat shader customization API

* lint

---------

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>

Author: mvaligursky

examples/src/examples/gaussian-splatting/multi-splat.shader.glsl.vert

@@ -1,15 +1,16 @@
 uniform float uTime;
 
-vec3 modifyPosition(vec3 center) {
+void modifyCenter(inout vec3 center) {
     // modify center
     float heightIntensity = center.y * 0.2;
     center.x += sin(uTime * 5.0 + center.y) * 0.3 * heightIntensity;
+}
 
-    // output y-coordinate
-    return center;
+void modifyCovariance(vec3 originalCenter, vec3 modifiedCenter, inout vec3 covA, inout vec3 covB) {
+    // no modification
 }
 
-vec4 modifyColor(vec3 center, vec4 clr) {
+void modifyColor(vec3 center, inout vec4 clr) {
     float sineValue = abs(sin(uTime * 5.0 + center.y));
 
     #ifdef CUTOUT
@@ -23,7 +24,5 @@ vec4 modifyColor(vec3 center, vec4 clr) {
         float blend = smoothstep(0.9, 1.0, sineValue);
         clr.xyz = mix(clr.xyz, gold, blend);
     #endif
-
-    return clr;
 }
 

examples/src/examples/gaussian-splatting/multi-splat.shader.wgsl.vert

@@ -1,31 +1,28 @@
 uniform uTime: f32;
 
-fn modifyPosition(center: vec3f) -> vec3f {
+fn modifyCenter(center: ptr<function, vec3f>) {
     // modify center
-    var result = center;
-    let heightIntensity = center.y * 0.2;
-    result.x += sin(uniform.uTime * 5.0 + center.y) * 0.3 * heightIntensity;
+    let heightIntensity = (*center).y * 0.2;
+    (*center).x += sin(uniform.uTime * 5.0 + (*center).y) * 0.3 * heightIntensity;
+}
 
-    // output y-coordinate
-    return result;
+fn modifyCovariance(originalCenter: vec3f, modifiedCenter: vec3f, covA: ptr<function, vec3f>, covB: ptr<function, vec3f>) {
+    // no modification
 }
 
-fn modifyColor(center: vec3f, clr: vec4f) -> vec4f {
-    var result = clr;
+fn modifyColor(center: vec3f, clr: ptr<function, vec4f>) {
     let sineValue = abs(sin(uniform.uTime * 5.0 + center.y));
 
     #ifdef CUTOUT
         // in cutout mode, remove pixels along the wave
         if (sineValue < 0.5) {
-            result.a = 0.0;
+            (*clr).a = 0.0;
         }
     #else
         // in non-cutout mode, add a golden tint to the wave
         let gold = vec3f(1.0, 0.85, 0.0);
         let blend = smoothstep(0.9, 1.0, sineValue);
-        result = vec4f(mix(clr.xyz, gold, blend), clr.a);
+        (*clr) = vec4f(mix((*clr).xyz, gold, blend), (*clr).a);
     #endif
-
-    return result;
 }
 

examples/src/examples/gaussian-splatting/picking.shader.glsl.vert

@@ -1,16 +1,20 @@
 uniform float fade;
 
-// animate position based on the fade value
-vec3 modifyPosition(vec3 pos) {
+// animate center position based on the fade value
+void modifyCenter(inout vec3 pos) {
     // Use a sine wave to create a smooth scale down and back up animation
     float angle = fade * 3.14159265;
     float shrinkFactor = sin(angle) * 0.3;
     float scale = 1.0 - shrinkFactor;
-    return pos * scale;
+    pos *= scale;
+}
+
+void modifyCovariance(vec3 originalCenter, vec3 modifiedCenter, inout vec3 covA, inout vec3 covB) {
+    // no modification
 }
 
 // animate color based on the fade value
-vec4 modifyColor(vec3 center, vec4 clr) {
+void modifyColor(vec3 center, inout vec4 clr) {
 
     // Check if the color is approximately grayscale
     float r = clr.r;
@@ -29,7 +33,5 @@ vec4 modifyColor(vec3 center, vec4 clr) {
         // cross fade blue to original orange color based on fade value
         clr.rgb = mix(clr.bgr * 0.5, clr.rgb, fade);
     }
-
-    return clr;
 }
 

examples/src/examples/gaussian-splatting/picking.shader.wgsl.vert

@@ -1,22 +1,25 @@
 uniform fade: f32;
 
-// animate position based on the fade value
-fn modifyPosition(pos: vec3f) -> vec3f {
+// animate center position based on the fade value
+fn modifyCenter(pos: ptr<function, vec3f>) {
     // Use a sine wave to create a smooth scale down and back up animation
     let angle = uniform.fade * 3.14159265;
     let shrinkFactor = sin(angle) * 0.3;
     let scale = 1.0 - shrinkFactor;
-    return pos * scale;
+    *pos = *pos * scale;
+}
+
+fn modifyCovariance(originalCenter: vec3f, modifiedCenter: vec3f, covA: ptr<function, vec3f>, covB: ptr<function, vec3f>) {
+    // no modification
 }
 
 // animate color based on the fade value
-fn modifyColor(center: vec3f, clr: vec4f) -> vec4f {
-    var result = clr;
+fn modifyColor(center: vec3f, clr: ptr<function, vec4f>) {
 
     // Check if the color is approximately grayscale
-    let r = clr.r;
-    let g = clr.g;
-    let b = clr.b;
+    let r = (*clr).r;
+    let g = (*clr).g;
+    let b = (*clr).b;
 
     let grayscaleThreshold = 0.01;
     let isGrayscale = abs(r - g) < grayscaleThreshold &&
@@ -25,12 +28,10 @@ fn modifyColor(center: vec3f, clr: vec4f) -> vec4f {
 
     if (isGrayscale) {
         // If the color is grayscale, make it very bright (the PC logo)
-        result = vec4f(clr.rgb * 10.0, clr.a);
+        *clr = vec4f((*clr).rgb * 10.0, (*clr).a);
     } else {
         // cross fade blue to original orange color based on fade value
-        result = vec4f(mix(clr.bgr * 0.5, clr.rgb, uniform.fade), clr.a);
+        *clr = vec4f(mix((*clr).bgr * 0.5, (*clr).rgb, uniform.fade), (*clr).a);
     }
-
-    return result;
 }
 

src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplat.js

@@ -1,4 +1,6 @@
 export default /* glsl */`
+#include "gsplatHelpersVS"
+#include "gsplatCustomizeVS"
 #include "gsplatCommonVS"
 
 varying mediump vec2 gaussianUV;
@@ -19,8 +21,6 @@ mediump vec4 discardVec = vec4(0.0, 0.0, 2.0, 1.0);
     uniform float colorRampIntensity;
 #endif
 
-#include "gsplatCustomizeVS"
-
 void main(void) {
     // read gaussian details
     SplatSource source;
@@ -34,9 +34,13 @@ void main(void) {
     #endif
 
     vec3 modelCenter = readCenter(source);
-    modelCenter = modifyPosition(modelCenter);
 
     SplatCenter center;
+    center.modelCenterOriginal = modelCenter;
+    
+    modifyCenter(modelCenter);
+    center.modelCenterModified = modelCenter;
+
     if (!initCenter(modelCenter, center)) {
         gl_Position = discardVec;
         return;
@@ -71,7 +75,7 @@ void main(void) {
         clr.xyz += evalSH(sh, dir) * scale;
     #endif
 
-    clr = modifyColor(modelCenter, clr);
+    modifyColor(modelCenter, clr);
 
     clipCorner(corner, clr.w);
 

src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatCorner.js

@@ -74,6 +74,7 @@ bool initCornerCov(SplatSource source, SplatCenter center, out SplatCorner corne
 bool initCorner(SplatSource source, SplatCenter center, out SplatCorner corner) {
     vec3 covA, covB;
     readCovariance(source, covA, covB);
+    modifyCovariance(center.modelCenterOriginal, center.modelCenterModified, covA, covB);
     return initCornerCov(source, center, corner, covA, covB);
 }
 `;

src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatCustomize.js

@@ -1,9 +1,22 @@
 export default /* glsl */`
-vec3 modifyPosition(vec3 center) {
-    return center;
+void modifyCenter(inout vec3 center) {
+    // Modify the splat center position
+    // Example: center.y += 1.0; // offset all splats up by 1 unit
 }
 
-vec4 modifyColor(vec3 center, vec4 color) {
-    return color;
+void modifyCovariance(vec3 originalCenter, vec3 modifiedCenter, inout vec3 covA, inout vec3 covB) {
+    // Modify the splat size/covariance
+    // Example to scale all splats by 2x:
+    // gsplatApplyUniformScale(covA, covB, 2.0);
+    //
+    // Example to clamp size to a range:
+    // float size = gsplatExtractSize(covA, covB);
+    // float newSize = clamp(size, 0.01, 0.5);
+    // gsplatApplyUniformScale(covA, covB, newSize / size);
+}
+
+void modifyColor(vec3 center, inout vec4 color) {
+    // Modify the splat color
+    // Example: color.rgb *= 0.5; // darken all splats
 }
 `;

src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatHelpers.js

@@ -0,0 +1,14 @@
+export default /* glsl */`
+// Helper function to extract the RMS size from covariance
+float gsplatExtractSize(vec3 covA, vec3 covB) {
+    float tr = covA.x + covB.x + covB.z;
+    return sqrt(max(tr, 0.0) / 3.0);
+}
+
+// Helper function to apply uniform scale to covariance
+void gsplatApplyUniformScale(inout vec3 covA, inout vec3 covB, float scale) {
+    float s2 = scale * scale;
+    covA *= s2;
+    covB *= s2;
+}
+`;

src/scene/shader-lib/glsl/chunks/gsplat/vert/gsplatStructs.js

@@ -14,6 +14,8 @@ struct SplatCenter {
     vec4 proj;          // center in clip space
     mat4 modelView;     // model-view matrix
     float projMat00;    // element [0][0] of the projection matrix
+    vec3 modelCenterOriginal;   // original model center before modification
+    vec3 modelCenterModified;   // model center after modification
 };
 
 // stores the offset from center for the current gaussian

src/scene/shader-lib/glsl/collections/shader-chunks-glsl.js

@@ -52,6 +52,7 @@ import gsplatCompressedDataVS from '../chunks/gsplat/vert/gsplatCompressedData.j
 import gsplatCompressedSHVS from '../chunks/gsplat/vert/gsplatCompressedSH.js';
 import gsplatCustomizeVS from '../chunks/gsplat/vert/gsplatCustomize.js';
 import gsplatEvalSHVS from '../chunks/gsplat/vert/gsplatEvalSH.js';
+import gsplatHelpersVS from '../chunks/gsplat/vert/gsplatHelpers.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';
@@ -237,6 +238,7 @@ const shaderChunksGLSL = {
     gsplatCompressedSHVS,
     gsplatCustomizeVS,
     gsplatEvalSHVS,
+    gsplatHelpersVS,
     gsplatQuatToMat3VS,
     gsplatSogsColorVS,
     gsplatSogsDataVS,