Skip to content

Commit ae9a5fd

Browse files
quadjrquadjr
committed
Improve the render speed
1 parent 7642b06 commit ae9a5fd

File tree

5 files changed

+147
-84
lines changed

5 files changed

+147
-84
lines changed

‎LICENSE

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
MIT License
22

3-
Copyright (c) 2023 Kevin Kwok, Junya Kuwada
3+
Copyright (c) 2023 Kevin Kwok, Mark Kellogg, Junya Kuwada
44

55
Permission is hereby granted, free of charge, to any person obtaining a copy
66
of this software and associated documentation files (the "Software"), to deal

‎README.md

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ About the splat file, please refer [antimatter15](https://github.com/antimatter1
2323
<html>
2424
<head>
2525
<script src="https://aframe.io/releases/1.4.2/aframe.min.js"></script>
26-
<script src="https://unpkg.com/[email protected].7/dist/aframe-gaussian-splatting-component.min.js"></script>
26+
<script src="https://unpkg.com/[email protected].8/dist/aframe-gaussian-splatting-component.min.js"></script>
2727
</head>
2828
<body>
2929
<a-scene renderer="antialias: false" stats>

‎index.html

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
<html>
33
<head>
44
<script src="https://aframe.io/releases/1.4.2/aframe.min.js"></script>
5-
<script src="https://unpkg.com/[email protected].7/dist/aframe-gaussian-splatting-component.min.js"></script>
5+
<script src="https://unpkg.com/[email protected].8/dist/aframe-gaussian-splatting-component.min.js"></script>
66
</head>
77
<body>
88
<a-scene renderer="antialias: false" stats>

‎index.js

Lines changed: 143 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -14,20 +14,140 @@ AFRAME.registerComponent("gaussian_splatting", {
1414

1515
const focal = (size.y / 2.0) / Math.tan(this.el.sceneEl.camera.el.components.camera.data.fov / 2.0 * Math.PI / 180.0);
1616

17-
const geometry = new THREE.PlaneGeometry( 4, 4);
17+
let u_buffer = new Uint8Array(buffer);
18+
if (
19+
u_buffer[0] == 112 &&
20+
u_buffer[1] == 108 &&
21+
u_buffer[2] == 121 &&
22+
u_buffer[3] == 10
23+
) {
24+
buffer = this.processPlyBuffer(buffer);
25+
u_buffer = new Uint8Array(buffer);
26+
}
27+
28+
const rowLength = 3 * 4 + 3 * 4 + 4 + 4;
29+
let vertexCount = Math.floor(buffer.byteLength / rowLength);
30+
let f_buffer = new Float32Array(buffer);
31+
32+
if(vertexCount > 4096*4096/3){
33+
console.log("vertexCount limited to 4096*4096/3", vertexCount);
34+
vertexCount = Math.floor(4096*4096/3);
35+
}
36+
37+
let matrices = new Float32Array(vertexCount * 16);
38+
const paddedCenterCovariances = new Float32Array(4096 * 4096 * 4);
39+
const paddedColors = new Float32Array(4096 * 4096 * 4);
40+
for (let i = 0; i < vertexCount; i++) {
41+
let quat = new THREE.Quaternion(
42+
(u_buffer[32 * i + 28 + 1] - 128) / 128.0,
43+
(u_buffer[32 * i + 28 + 2] - 128) / 128.0,
44+
-(u_buffer[32 * i + 28 + 3] - 128) / 128.0,
45+
(u_buffer[32 * i + 28 + 0] - 128) / 128.0,
46+
);
47+
let center = new THREE.Vector3(
48+
f_buffer[8 * i + 0],
49+
f_buffer[8 * i + 1],
50+
-f_buffer[8 * i + 2]
51+
);
52+
let scale = new THREE.Vector3(
53+
f_buffer[8 * i + 3 + 0],
54+
f_buffer[8 * i + 3 + 1],
55+
f_buffer[8 * i + 3 + 2]
56+
);
57+
58+
let mtx = new THREE.Matrix4();
59+
mtx.makeRotationFromQuaternion(quat);
60+
mtx.transpose();
61+
mtx.scale(scale);
62+
let mtx_t = mtx.clone()
63+
mtx.transpose();
64+
mtx.premultiply(mtx_t);
65+
mtx.setPosition(center);
66+
67+
let destOffset = i * 12;
68+
paddedCenterCovariances[destOffset + 0] = center.x;
69+
paddedCenterCovariances[destOffset + 1] = center.y;
70+
paddedCenterCovariances[destOffset + 2] = center.z;
71+
paddedCenterCovariances[destOffset + 3] = mtx.elements[0];
72+
paddedCenterCovariances[destOffset + 4] = mtx.elements[1];
73+
paddedCenterCovariances[destOffset + 5] = mtx.elements[2];
74+
paddedCenterCovariances[destOffset + 6] = mtx.elements[5];
75+
paddedCenterCovariances[destOffset + 7] = mtx.elements[6];
76+
paddedCenterCovariances[destOffset + 8] = mtx.elements[10];
77+
78+
// RGBA
79+
destOffset = i * 4;
80+
paddedColors[destOffset + 0] = u_buffer[32 * i + 24 + 0] / 255;
81+
paddedColors[destOffset + 1] = u_buffer[32 * i + 24 + 1] / 255;
82+
paddedColors[destOffset + 2] = u_buffer[32 * i + 24 + 2] / 255;
83+
paddedColors[destOffset + 3] = u_buffer[32 * i + 24 + 3] / 255;
84+
85+
for(let j = 0; j < 16; j++){
86+
matrices[i * 16 + j] = mtx.elements[j];
87+
}
88+
}
89+
90+
const centerCovarianceTexture = new THREE.DataTexture(paddedCenterCovariances, 4096, 4096, THREE.RGBAFormat, THREE.FloatType);
91+
centerCovarianceTexture.needsUpdate = true;
92+
const colorTexture = new THREE.DataTexture(paddedColors, 4096, 4096, THREE.RGBAFormat, THREE.FloatType);
93+
colorTexture.needsUpdate = true;
94+
95+
const camera_mtx = this.el.sceneEl.camera.el.object3D.matrixWorld.elements;
96+
let view = new Float32Array([camera_mtx[2], camera_mtx[6], camera_mtx[10]]);
97+
let splatIndexArray = this.sortSplats(matrices, view);
98+
const splatIndexes = new THREE.InstancedBufferAttribute(splatIndexArray, 1, false);
99+
splatIndexes.setUsage(THREE.DynamicDrawUsage);
100+
101+
const baseGeometry = new THREE.BufferGeometry();
102+
const positionsArray = new Float32Array(6 * 3);
103+
const positions = new THREE.BufferAttribute(positionsArray, 3);
104+
baseGeometry.setAttribute('position', positions);
105+
positions.setXYZ(2, -2.0, 2.0, 0.0);
106+
positions.setXYZ(1, 2.0, 2.0, 0.0);
107+
positions.setXYZ(0, -2.0, -2.0, 0.0);
108+
positions.setXYZ(5, -2.0, -2.0, 0.0);
109+
positions.setXYZ(4, 2.0, 2.0, 0.0);
110+
positions.setXYZ(3, 2.0, -2.0, 0.0);
111+
positions.needsUpdate = true;
112+
113+
const geometry = new THREE.InstancedBufferGeometry().copy(baseGeometry);
114+
geometry.setAttribute('splatIndex', splatIndexes);
115+
geometry.instanceCount = vertexCount;
116+
18117
const material = new THREE.ShaderMaterial( {
19118
uniforms : {
20119
viewport: {value: new Float32Array([size.x, size.y])},
21-
focal: {value: focal}
120+
focal: {value: focal},
121+
centerCovarianceTexture: {value: centerCovarianceTexture},
122+
colorTexture: {value: colorTexture}
22123
},
23124
vertexShader: `
24125
out vec4 vColor;
25126
out vec2 vPosition;
26127
uniform vec2 viewport;
27128
uniform float focal;
28129
130+
attribute uint splatIndex;
131+
uniform sampler2D centerCovarianceTexture;
132+
uniform sampler2D colorTexture;
133+
134+
ivec2 getDataUV(in int stride, in int offset) {
135+
uint covarianceD = splatIndex * uint(stride) + uint(offset);
136+
return ivec2(covarianceD%uint(4096),covarianceD/uint(4096));
137+
}
138+
29139
void main () {
30-
vec4 center = vec4(instanceMatrix[3][0], instanceMatrix[3][1], instanceMatrix[3][2], 1);
140+
vec4 sampledCenterCovarianceA = texelFetch(centerCovarianceTexture, getDataUV(3, 0), 0);
141+
vec4 sampledCenterCovarianceB = texelFetch(centerCovarianceTexture, getDataUV(3, 1), 0);
142+
vec4 sampledCenterCovarianceC = texelFetch(centerCovarianceTexture, getDataUV(3, 2), 0);
143+
144+
vec4 center = vec4(sampledCenterCovarianceA.xyz, 1);
145+
vec3 cov3D_M11_M12_M13 = vec3(sampledCenterCovarianceA.w, sampledCenterCovarianceB.xy);
146+
vec3 cov3D_M22_M23_M33 = vec3(sampledCenterCovarianceB.zw, sampledCenterCovarianceC.x);
147+
148+
ivec2 colorUV = ivec2(splatIndex%uint(4096),splatIndex/uint(4096));
149+
vec4 sampledColor = texelFetch(colorTexture, colorUV, 0);
150+
31151
// Adjust View Pose
32152
mat4 adjViewMatrix = inverse(viewMatrix);
33153
adjViewMatrix[0][1] *= -1.0;
@@ -50,6 +170,12 @@ AFRAME.registerComponent("gaussian_splatting", {
50170
return;
51171
}
52172
173+
mat3 Vrk = mat3(
174+
cov3D_M11_M12_M13.x, cov3D_M11_M12_M13.y, cov3D_M11_M12_M13.z,
175+
cov3D_M11_M12_M13.y, cov3D_M22_M23_M33.x, cov3D_M22_M23_M33.y,
176+
cov3D_M11_M12_M13.z, cov3D_M22_M23_M33.y, cov3D_M22_M23_M33.z
177+
);
178+
53179
mat3 J = mat3(
54180
focal / camspace.z, 0., -(focal * camspace.x) / (camspace.z * camspace.z),
55181
0., -focal / camspace.z, (focal * camspace.y) / (camspace.z * camspace.z),
@@ -58,7 +184,7 @@ AFRAME.registerComponent("gaussian_splatting", {
58184
59185
mat3 W = transpose(mat3(modelView));
60186
mat3 T = W * J;
61-
mat3 cov = transpose(T) * mat3(instanceMatrix) * T;
187+
mat3 cov = transpose(T) * Vrk * T;
62188
63189
vec2 vCenter = vec2(pos2d) / pos2d.w;
64190
@@ -74,7 +200,7 @@ AFRAME.registerComponent("gaussian_splatting", {
74200
vec2 v1 = min(sqrt(2.0 * lambda1), 1024.0) * diagonalVector;
75201
vec2 v2 = min(sqrt(2.0 * lambda2), 1024.0) * vec2(diagonalVector.y, -diagonalVector.x);
76202
77-
vColor = vec4(instanceMatrix[0][3], instanceMatrix[1][3], instanceMatrix[2][3], instanceMatrix[3][3]);
203+
vColor = sampledColor;
78204
vPosition = position.xy;
79205
80206
gl_Position = vec4(
@@ -97,7 +223,8 @@ AFRAME.registerComponent("gaussian_splatting", {
97223
blending : THREE.CustomBlending,
98224
blendSrcAlpha : THREE.OneFactor,
99225
depthTest : true,
100-
depthWrite: false
226+
depthWrite: false,
227+
transparent: false
101228
} );
102229

103230
window.addEventListener('resize', () => {
@@ -109,67 +236,9 @@ AFRAME.registerComponent("gaussian_splatting", {
109236
material.uniforms.focal.value = focal;
110237
});
111238

112-
let u_buffer = new Uint8Array(buffer);
113-
if (
114-
u_buffer[0] == 112 &&
115-
u_buffer[1] == 108 &&
116-
u_buffer[2] == 121 &&
117-
u_buffer[3] == 10
118-
) {
119-
buffer = this.processPlyBuffer(buffer);
120-
u_buffer = new Uint8Array(buffer);
121-
}
122-
123-
const rowLength = 3 * 4 + 3 * 4 + 4 + 4;
124-
let vertexCount = Math.floor(buffer.byteLength / rowLength);
125-
let f_buffer = new Float32Array(buffer);
126-
127-
let matrices = new Float32Array(vertexCount * 16);
128-
for (let i = 0; i < vertexCount; i++) {
129-
let quat = new THREE.Quaternion(
130-
(u_buffer[32 * i + 28 + 1] - 128) / 128.0,
131-
(u_buffer[32 * i + 28 + 2] - 128) / 128.0,
132-
-(u_buffer[32 * i + 28 + 3] - 128) / 128.0,
133-
(u_buffer[32 * i + 28 + 0] - 128) / 128.0,
134-
);
135-
let center = new THREE.Vector3(
136-
f_buffer[8 * i + 0],
137-
f_buffer[8 * i + 1],
138-
-f_buffer[8 * i + 2]
139-
);
140-
let scale = new THREE.Vector3(
141-
f_buffer[8 * i + 3 + 0],
142-
f_buffer[8 * i + 3 + 1],
143-
f_buffer[8 * i + 3 + 2]
144-
);
145-
146-
let mtx = new THREE.Matrix4();
147-
mtx.makeRotationFromQuaternion(quat);
148-
mtx.transpose();
149-
mtx.scale(scale);
150-
let mtx_t = mtx.clone()
151-
mtx.transpose();
152-
mtx.premultiply(mtx_t);
153-
mtx.setPosition(center);
154-
155-
// RGBA
156-
mtx.elements[3] = u_buffer[32 * i + 24 + 0] / 255;
157-
mtx.elements[7] = u_buffer[32 * i + 24 + 1] / 255;
158-
mtx.elements[11] = u_buffer[32 * i + 24 + 2] / 255;
159-
mtx.elements[15] = u_buffer[32 * i + 24 + 3] / 255;
160-
161-
for(let j = 0; j < 16; j++){
162-
matrices[i * 16 + j] = mtx.elements[j];
163-
}
164-
}
165-
166-
const camera_mtx = this.el.sceneEl.camera.el.object3D.matrixWorld.elements;
167-
let view = new Float32Array([camera_mtx[2], camera_mtx[6], camera_mtx[10]]);
168-
this.iMesh = new THREE.InstancedMesh(geometry, material, vertexCount);
169-
this.iMesh.frustumCulled = false;
170-
this.iMesh.instanceMatrix.array = this.sortSplats(matrices, view);
171-
this.iMesh.instanceMatrix.needsUpdate = true;
172-
this.el.object3D.add(this.iMesh);
239+
let mesh = new THREE.Mesh(geometry, material, vertexCount);
240+
mesh.frustumCulled = false;
241+
this.el.object3D.add(mesh);
173242

174243
this.worker = new Worker(
175244
URL.createObjectURL(
@@ -185,8 +254,10 @@ AFRAME.registerComponent("gaussian_splatting", {
185254
}, [matrices.buffer]);
186255

187256
this.worker.onmessage = (e) => {
188-
this.iMesh.instanceMatrix.array = new Float32Array(e.data.sortedMatrices);
189-
this.iMesh.instanceMatrix.needsUpdate = true;
257+
let indexes = new Uint32Array(e.data.sortedIndexes);
258+
mesh.geometry.attributes.splatIndex.set(indexes);
259+
mesh.geometry.attributes.splatIndex.needsUpdate = true;
260+
mesh.geometry.instanceCount = indexes.length;
190261
this.sortReady = true;
191262
};
192263
this.sortReady = true;
@@ -213,8 +284,8 @@ AFRAME.registerComponent("gaussian_splatting", {
213284
}
214285
if(e.data.view){
215286
const view = new Float32Array(e.data.view);
216-
const sortedMatrices = sortFunction(matrices, view);
217-
self.postMessage({sortedMatrices}, [sortedMatrices.buffer]);
287+
const sortedIndexes = sortFunction(matrices, view);
288+
self.postMessage({sortedIndexes}, [sortedIndexes.buffer]);
218289
}
219290
};
220291
},
@@ -247,15 +318,7 @@ AFRAME.registerComponent("gaussian_splatting", {
247318
let depthIndex = new Uint32Array(vertexCount);
248319
for (let i = 0; i < vertexCount; i++) depthIndex[starts0[sizeList[i]]++] = i;
249320

250-
let sortedMatrices = new Float32Array(vertexCount * 16);
251-
for (let j = 0; j < vertexCount; j++) {
252-
let i = depthIndex[j];
253-
for(let k = 0; k < 16; k++){
254-
sortedMatrices[j * 16 + k] = matrices[i * 16 + k];
255-
}
256-
}
257-
258-
return sortedMatrices;
321+
return depthIndex;
259322
},
260323
processPlyBuffer: function (inputBuffer) {
261324
const ubuf = new Uint8Array(inputBuffer);

‎package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "aframe-gaussian-splatting-component",
3-
"version": "0.0.7",
3+
"version": "0.0.8",
44
"description": "This component is an A-Frame implementation of real-time rendering for '3D Gaussian Splatting for Real-Time Radiance Field Rendering'",
55
"main": "index.js",
66
"scripts": {

0 commit comments

Comments
 (0)