From bdab09ebc6785befdc46fa1a96923067b740c3d1 Mon Sep 17 00:00:00 2001 From: Martin Valigursky <59932779+mvaligursky@users.noreply.github.com> Date: Thu, 19 Oct 2023 13:59:49 +0100 Subject: [PATCH] Splatting uses colors stored in a texture instead of vertex buffer (#250) * Splatting uses colors stored in a texture instead of vertex buffer * spaces * update to avoid rounding errors (based on morphing code) * removed colors from vertex --------- Co-authored-by: Martin Valigursky --- src/splat-resource.ts | 167 ++++++++++++++++++++++++++++++++---------- 1 file changed, 130 insertions(+), 37 deletions(-) diff --git a/src/splat-resource.ts b/src/splat-resource.ts index d745855..15a0e2c 100644 --- a/src/splat-resource.ts +++ b/src/splat-resource.ts @@ -19,15 +19,20 @@ import { SEMANTIC_ATTR11, SEMANTIC_ATTR12, SEMANTIC_ATTR13, - SEMANTIC_COLOR, SEMANTIC_POSITION, Texture, - TYPE_UINT8, + math, + Vec2, Vec3, VertexFormat, VertexBuffer, TYPE_FLOAT32, - BLEND_NORMAL + BLEND_NORMAL, + FILTER_NEAREST, + ADDRESS_CLAMP_TO_EDGE, + PIXELFORMAT_RGBA8, + SEMANTIC_ATTR14, + TYPE_UINT32 } from 'playcanvas'; import { PlyElement } from './ply-reader'; @@ -37,7 +42,7 @@ import { SortWorker } from './sort-worker'; const debugRender = false; const debugRenderBounds = false; -const quatToMat3 = ` +const sharedShader = ` mat3 quatToMat3(vec3 R) { float x = R.x; @@ -59,12 +64,57 @@ mat3 quatToMat3(vec3 R) 1.0 - 2.0 * (y * y + z * z) ); } + +#ifdef WEBGPU + attribute uint vertex_id; +#else + attribute float vertex_id; +#endif + +uniform vec4 tex_params; +uniform sampler2D splatColor; + +#ifdef WEBGPU + ivec2 getTextureCoords() { + + // turn vertex_id into int grid coordinates + ivec2 textureSize = ivec2(tex_params.xy); + vec2 invTextureSize = tex_params.zw; + + int gridV = int(float(vertex_id) * invTextureSize.x); + int gridU = int(vertex_id - gridV * textureSize.x); + return ivec2(gridU, gridV); + } + +#else + vec2 getTextureCoords() { + vec2 textureSize = tex_params.xy; + vec2 invTextureSize = tex_params.zw; + + // turn vertex_id into int grid coordinates + float gridV = floor(vertex_id * invTextureSize.x); + float gridU = vertex_id - (gridV * textureSize.x); + + // convert grid coordinates to uv coordinates with half pixel offset + return vec2(gridU, gridV) * invTextureSize + (0.5 * invTextureSize); + } +#endif + +vec4 getColor() { + #ifdef WEBGPU + ivec2 textureUV = getTextureCoords(); + return texelFetch(splatColor, ivec2(textureUV), 0); + #else + vec2 textureUV = getTextureCoords(); + return texture2D(splatColor, textureUV); + #endif +} + `; -const splatVS = /* glsl_ */ ` +const splatVS = ` attribute vec2 vertex_position; attribute vec3 splat_center; -attribute vec4 splat_color; attribute vec3 splat_rotation; attribute vec3 splat_scale; @@ -77,7 +127,7 @@ uniform vec2 viewport; varying vec2 texCoord; varying vec4 color; -${quatToMat3} +${sharedShader} void computeCov3d(in vec3 rot, in vec3 scale, out vec3 covA, out vec3 covB) { @@ -159,7 +209,8 @@ void main(void) 0.0, 0.0) * splat_proj.w; texCoord = vertex_position * 2.0; - color = splat_color; + + color = getColor(); } `; @@ -179,7 +230,6 @@ void main(void) const splatDebugVS = /* glsl_ */ ` attribute vec3 vertex_position; attribute vec3 splat_center; -attribute vec4 splat_color; attribute vec3 splat_rotation; attribute vec3 splat_scale; @@ -188,13 +238,14 @@ uniform mat4 matrix_viewProjection; varying vec4 color; -${quatToMat3} +${sharedShader} void main(void) { vec3 local = quatToMat3(splat_rotation) * (vertex_position * splat_scale * 2.0) + splat_center; gl_Position = matrix_viewProjection * matrix_model * vec4(local, 1.0); - color = splat_color; + + color = getColor(); } `; @@ -224,9 +275,9 @@ const getSplatMat = (result: Mat4, data: Float32Array) => { const px = data[0]; const py = data[1]; const pz = data[2]; - const x = data[4]; - const y = data[5]; - const z = data[6]; + const x = data[3]; + const y = data[4]; + const z = data[5]; const w = Math.sqrt(1 - (x * x + y * y + z * z)); // build rotation matrix @@ -253,7 +304,7 @@ const getSplatMat = (result: Mat4, data: Float32Array) => { const getSplatAabb = (result: BoundingBox, data: Float32Array) => { getSplatMat(mat4, data); aabb.center.set(0, 0, 0); - aabb.halfExtents.set(data[7] * 2, data[8] * 2, data[9] * 2); + aabb.halfExtents.set(data[6] * 2, data[7] * 2, data[8] * 2); result.setFromTransformedAabb(aabb, mat4); }; @@ -261,9 +312,9 @@ const renderDebugSplat = (app: AppBase, worldMat: Mat4, data: Float32Array) => { getSplatMat(mat4, data); mat4.mul2(worldMat, mat4); - const sx = data[7]; - const sy = data[8]; - const sz = data[9]; + const sx = data[6]; + const sy = data[7]; + const sz = data[8]; for (let i = 0; i < 8; ++i) { vec3.set( @@ -310,9 +361,9 @@ class SplatResource extends ContainerResource { this.quadMaterial.shader = createShaderFromCode(this.device, vs, fs, 'splatShader', { vertex_position: SEMANTIC_POSITION, splat_center: SEMANTIC_ATTR11, - splat_color: SEMANTIC_COLOR, splat_rotation: SEMANTIC_ATTR12, - splat_scale: SEMANTIC_ATTR13 + splat_scale: SEMANTIC_ATTR13, + vertex_id: SEMANTIC_ATTR14 }); this.quadMaterial.update(); @@ -334,6 +385,27 @@ class SplatResource extends ContainerResource { } + evalTextureSize(count: number) : Vec2 { + const width = Math.ceil(Math.sqrt(count)); + const height = Math.ceil(count / width); + return new Vec2(width, height); + } + + createTexture(name: string, format: number, size: Vec2) { + return new Texture(this.device, { + width: size.x, + height: size.y, + format: format, + cubemap: false, + mipmaps: false, + minFilter: FILTER_NEAREST, + magFilter: FILTER_NEAREST, + addressU: ADDRESS_CLAMP_TO_EDGE, + addressV: ADDRESS_CLAMP_TO_EDGE, + name: name + }); + } + instantiateModelEntity(/* options: any */): Entity { return null; } @@ -373,11 +445,16 @@ class SplatResource extends ContainerResource { const stride = 10; + const textureSize = this.evalTextureSize(vertexElement.count); + const colorTexture = this.createTexture('splatColor', PIXELFORMAT_RGBA8, textureSize); + const colorData = colorTexture.lock(); + // position.xyz, color, rotation.xyz, scale.xyz const floatData = new Float32Array(vertexElement.count * stride); - const uint8Data = new Uint8ClampedArray(floatData.buffer); + const uint32Data = new Uint32Array(floatData.buffer); const quat = new Quat(); + const isWebGPU = this.device.isWebGPU; for (let i = 0; i < vertexElement.count; ++i) { const j = i; @@ -396,9 +473,13 @@ class SplatResource extends ContainerResource { // vertex colors if (f_dc_0 && f_dc_1 && f_dc_2) { const SH_C0 = 0.28209479177387814; - uint8Data[i * stride * 4 + 12] = (0.5 + SH_C0 * f_dc_0[j]) * 255; - uint8Data[i * stride * 4 + 13] = (0.5 + SH_C0 * f_dc_1[j]) * 255; - uint8Data[i * stride * 4 + 14] = (0.5 + SH_C0 * f_dc_2[j]) * 255; + const r = math.clamp((0.5 + SH_C0 * f_dc_0[j]) * 255, 0, 255); + const g = math.clamp((0.5 + SH_C0 * f_dc_1[j]) * 255, 0, 255); + const b = math.clamp((0.5 + SH_C0 * f_dc_2[j]) * 255, 0, 255); + + colorData[i * 4 + 0] = r; + colorData[i * 4 + 1] = g; + colorData[i * 4 + 2] = b; } // opacity @@ -411,36 +492,48 @@ class SplatResource extends ContainerResource { const t = Math.exp(v); return t / (1 + t); }; - uint8Data[i * stride * 4 + 15] = sigmoid(opacity[j]) * 255; + const a = sigmoid(opacity[j]) * 255; + colorData[i * 4 + 3] = a; } else { - uint8Data[i * stride * 4 + 15] = 255; + colorData[i * 4 + 3] = 255; } quat.set(rot_0[j], rot_1[j], rot_2[j], rot_3[j]).normalize(); // rotation if (quat.w < 0) { - floatData[i * stride + 4] = -quat.x; - floatData[i * stride + 5] = -quat.y; - floatData[i * stride + 6] = -quat.z; + floatData[i * stride + 3] = -quat.x; + floatData[i * stride + 4] = -quat.y; + floatData[i * stride + 5] = -quat.z; } else { - floatData[i * stride + 4] = quat.x; - floatData[i * stride + 5] = quat.y; - floatData[i * stride + 6] = quat.z; + floatData[i * stride + 3] = quat.x; + floatData[i * stride + 4] = quat.y; + floatData[i * stride + 5] = quat.z; } // scale - floatData[i * stride + 7] = Math.exp(scale_0[j]); - floatData[i * stride + 8] = Math.exp(scale_1[j]); - floatData[i * stride + 9] = Math.exp(scale_2[j]); + floatData[i * stride + 6] = Math.exp(scale_0[j]); + floatData[i * stride + 7] = Math.exp(scale_1[j]); + floatData[i * stride + 8] = Math.exp(scale_2[j]); + + // index + if (isWebGPU) { + uint32Data[i * stride + 9] = i; + } else { + floatData[i * stride + 9] = i + 0.2; + } } + colorTexture.unlock(); + this.quadMaterial.setParameter('splatColor', colorTexture); + this.quadMaterial.setParameter('tex_params', new Float32Array([textureSize.x, textureSize.y, 1 / textureSize.x, 1 / textureSize.y])); + // create instance data const vertexFormat = new VertexFormat(this.device, [ { semantic: SEMANTIC_ATTR11, components: 3, type: TYPE_FLOAT32 }, - { semantic: SEMANTIC_COLOR, components: 4, type: TYPE_UINT8, normalize: true }, { semantic: SEMANTIC_ATTR12, components: 3, type: TYPE_FLOAT32 }, - { semantic: SEMANTIC_ATTR13, components: 3, type: TYPE_FLOAT32 } + { semantic: SEMANTIC_ATTR13, components: 3, type: TYPE_FLOAT32 }, + { semantic: SEMANTIC_ATTR14, components: 1, type: isWebGPU ? TYPE_UINT32 : TYPE_FLOAT32 } ]); const vertexBuffer = new VertexBuffer(this.device, vertexFormat, vertexElement.count, BUFFER_DYNAMIC, floatData.buffer);