Skip to content

Commit

Permalink
Splatting uses colors stored in a texture instead of vertex buffer (#250
Browse files Browse the repository at this point in the history 
)

* 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 <[email protected]>
  • Loading branch information
@mvaligursky
mvaligursky and Martin Valigursky authored Oct 19, 2023
1 parent e1e9e67 commit bdab09e
Showing 1 changed file with 130 additions and 37 deletions.
167 changes: 130 additions & 37 deletions src/splat-resource.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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)
{
Expand Down Expand Up @@ -159,7 +209,8 @@ void main(void)
0.0, 0.0) * splat_proj.w;
texCoord = vertex_position * 2.0;
color = splat_color;
color = getColor();
}
`;

Expand All @@ -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;
Expand All @@ -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();
}
`;

Expand Down Expand Up @@ -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
Expand All @@ -253,17 +304,17 @@ 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);
};

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(
Expand Down Expand Up @@ -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();
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
Expand All @@ -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
Expand All @@ -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);

Expand Down

0 comments on commit bdab09e

Please sign in to comment.