From 674851578734b82d3d121876af34f06f37cae75c Mon Sep 17 00:00:00 2001
From: Martin Valigursky <59932779+mvaligursky@users.noreply.github.com>
Date: Fri, 20 Oct 2023 16:12:26 +0100
Subject: [PATCH] use centers from the texture in the shader (#257)

Co-authored-by: Martin Valigursky <mvaligursky@snapchat.com>
---
 src/splat/splat-material.ts | 17 ++++++++++++-----
 src/splat/splat.ts          | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/src/splat/splat-material.ts b/src/splat/splat-material.ts
index 7652344..2e21d0c 100644
--- a/src/splat/splat-material.ts
+++ b/src/splat/splat-material.ts
@@ -6,13 +6,11 @@ import {
     GraphicsDevice,
     Material,
     SEMANTIC_POSITION,
-    SEMANTIC_ATTR11,
     SEMANTIC_ATTR13
 } from "playcanvas";
 
 const splatVS = `
     attribute vec3 vertex_position;
-    attribute vec3 splat_center;
 
     uniform mat4 matrix_model;
     uniform mat4 matrix_view;
@@ -50,6 +48,7 @@ const splatVS = `
     uniform sampler2D splatColor;
     uniform highp sampler2D splatScale;
     uniform highp sampler2D splatRotation;
+    uniform highp sampler2D splatCenter;
 
     #ifdef WEBGPU
 
@@ -78,6 +77,10 @@ const splatVS = `
             return texelFetch(splatRotation, dataUV, 0).xyz;
         }
 
+        vec3 getCenter() {
+            return texelFetch(splatCenter, dataUV, 0).xyz;
+        }
+
     #else
 
         // TODO: use texture2DLodEXT on WebGL
@@ -108,6 +111,10 @@ const splatVS = `
             return texture(splatRotation, dataUV).xyz;
         }
 
+        vec3 getCenter() {
+            return texture(splatCenter, dataUV).xyz;
+        }
+
     #endif
 
     void computeCov3d(in vec3 rot, in vec3 scale, out vec3 covA, out vec3 covB)
@@ -144,7 +151,8 @@ const splatVS = `
     {
         evalDataUV();
 
-        vec4 splat_cam = matrix_view * matrix_model * vec4(splat_center, 1.0);
+        vec3 center = getCenter();
+        vec4 splat_cam = matrix_view * matrix_model * vec4(center, 1.0);
         vec4 splat_proj = matrix_projection * splat_cam;
 
         // cull behind camera
@@ -199,7 +207,7 @@ const splatVS = `
 
         #ifdef DEBUG_RENDER
 
-            vec3 local = quatToMat3(rotation) * (vertex_position * scale * 2.0) + splat_center;
+            vec3 local = quatToMat3(rotation) * (vertex_position * scale * 2.0) + center;
             gl_Position = matrix_viewProjection * matrix_model * vec4(local, 1.0);
 
             color = getColor();
@@ -243,7 +251,6 @@ const createSplatMaterial = (device: GraphicsDevice, debugRender = false) => {
 
     result.shader = createShaderFromCode(device, vs, fs, `splatShader-${debugRender}`, {
         vertex_position: SEMANTIC_POSITION,
-        splat_center: SEMANTIC_ATTR11,
         vertex_id: SEMANTIC_ATTR13
     });
 
diff --git a/src/splat/splat.ts b/src/splat/splat.ts
index a1a910c..7e2771e 100644
--- a/src/splat/splat.ts
+++ b/src/splat/splat.ts
@@ -265,6 +265,36 @@ class Splat {
         return texture;
     }
 
+    createCenterTexture(splatData: SplatData, size: Vec2, format: object) {
+
+        // texture format based vars
+        const { numComponents, isHalf } = format;
+
+        const x = splatData.getProp('x');
+        const y = splatData.getProp('y');
+        const z = splatData.getProp('z');
+
+        const texture = this.createTexture('splatCenter', format.format, size);
+        const data = texture.lock();
+
+        for (let i = 0; i < splatData.numSplats; i++) {
+
+            if (isHalf) {
+                data[i * numComponents + 0] = float2Half(x[i]);
+                data[i * numComponents + 1] = float2Half(y[i]);
+                data[i * numComponents + 2] = float2Half(z[i]);
+            } else {
+                data[i * numComponents + 0] = x[i];
+                data[i * numComponents + 1] = y[i];
+                data[i * numComponents + 2] = z[i];
+            }
+        }
+
+        texture.unlock();
+        return texture;
+    }
+
+
     create(splatData: SplatData, options: any) {
         const x = splatData.getProp('x');
         const y = splatData.getProp('y');
@@ -280,12 +310,12 @@ class Splat {
         const colorTexture = this.createColorTexture(splatData, textureSize);
         const scaleTexture = this.createScaleTexture(splatData, textureSize, this.getTextureFormat(false));
         const rotationTexture = this.createRotationTexture(splatData, textureSize, this.getTextureFormat(false));
+        const centerTexture = this.createCenterTexture(splatData, textureSize, this.getTextureFormat(false));
 
         // position.xyz, rotation.xyz
         const floatData = new Float32Array(splatData.numSplats * stride);
         const uint32Data = new Uint32Array(floatData.buffer);
 
-        // const quat = new Quat();
         const isWebGPU = this.device.isWebGPU;
 
         for (let i = 0; i < splatData.numSplats; ++i) {
@@ -307,6 +337,7 @@ class Splat {
         this.material.setParameter('splatColor', colorTexture);
         this.material.setParameter('splatScale', scaleTexture);
         this.material.setParameter('splatRotation', rotationTexture);
+        this.material.setParameter('splatCenter', centerTexture);
         this.material.setParameter('tex_params', new Float32Array([textureSize.x, textureSize.y, 1 / textureSize.x, 1 / textureSize.y]));
 
         // create instance data