aframe-gaussian-splatting-component/dist/aframe-gaussian-splatting-component.min.js

This component is an A-Frame implementation of real-time rendering for '3D Gaussian Splatting for Real-Time Radiance Field Rendering'

AFRAME.registerComponent("gaussian_splatting",{schema:{src:{type:"string",default:"train.splat"}},init:function(){this.el.sceneEl.renderer.setPixelRatio(1),this.el.sceneEl.renderer.xr.setFramebufferScaleFactor(.5),this.initGL(this.el.sceneEl.camera.el.components.camera.camera,this.el.object3D,this.el.sceneEl.renderer),this.loadData(this.data.src)},initGL:function(camera,object,renderer){this.camera=camera,this.object=object,this.renderer=renderer,this.textureReady=!1,this.object.frustumCulled=!1,this.centerAndScaleData=new Float32Array(67108864),this.covAndColorData=new Uint32Array(67108864),this.centerAndScaleTexture=new THREE.DataTexture(this.centerAndScaleData,4096,4096,THREE.RGBA,THREE.FloatType),this.centerAndScaleTexture.needsUpdate=!0,this.covAndColorTexture=new THREE.DataTexture(this.covAndColorData,4096,4096,THREE.RGBAIntegerFormat,THREE.UnsignedIntType),this.covAndColorTexture.internalFormat="RGBA32UI",this.covAndColorTexture.needsUpdate=!0;camera=new Uint32Array(16777216),object=new THREE.InstancedBufferAttribute(camera,1,!1),object.setUsage(THREE.DynamicDrawUsage),renderer=new THREE.BufferGeometry,camera=new Float32Array(18),camera=new THREE.BufferAttribute(camera,3),renderer.setAttribute("position",camera),camera.setXYZ(2,-2,2,0),camera.setXYZ(1,2,2,0),camera.setXYZ(0,-2,-2,0),camera.setXYZ(5,-2,-2,0),camera.setXYZ(4,2,2,0),camera.setXYZ(3,2,-2,0),camera.needsUpdate=!0,camera=(new THREE.InstancedBufferGeometry).copy(renderer);camera.setAttribute("splatIndex",object),camera.instanceCount=1;const material=new THREE.ShaderMaterial({uniforms:{viewport:{value:new Float32Array([1980,1080])},focal:{value:1e3},centerAndScaleTexture:{value:this.centerAndScaleTexture},covAndColorTexture:{value:this.covAndColorTexture},gsProjectionMatrix:{value:this.getProjectionMatrix()},gsModelViewMatrix:{value:this.getModelViewMatrix()}},vertexShader:`
				precision highp usampler2D;

				out vec4 vColor;
				out vec2 vPosition;
				uniform vec2 viewport;
				uniform float focal;
				uniform mat4 gsProjectionMatrix;
				uniform mat4 gsModelViewMatrix;

				attribute uint splatIndex;
				uniform sampler2D centerAndScaleTexture;
				uniform usampler2D covAndColorTexture;

				vec2 unpackInt16(in uint value) {
					int v = int(value);
					int v0 = v >> 16;
					int v1 = (v & 0xFFFF);
					if((v & 0x8000) != 0)
						v1 |= 0xFFFF0000;
					return vec2(float(v1), float(v0));
				}

				void main () {
					ivec2 texPos = ivec2(splatIndex%uint(4096),splatIndex/uint(4096));
					vec4 centerAndScaleData = texelFetch(centerAndScaleTexture, texPos, 0);

					vec4 center = vec4(centerAndScaleData.xyz, 1);
					vec4 camspace = gsModelViewMatrix * center;
					vec4 pos2d = gsProjectionMatrix * camspace;

					float bounds = 1.2 * pos2d.w;
					if (pos2d.z < -pos2d.w || pos2d.x < -bounds || pos2d.x > bounds
						|| pos2d.y < -bounds || pos2d.y > bounds) {
						gl_Position = vec4(0.0, 0.0, 2.0, 1.0);
						return;
					}

					uvec4 covAndColorData = texelFetch(covAndColorTexture, texPos, 0);
					vec2 cov3D_M11_M12 = unpackInt16(covAndColorData.x) * centerAndScaleData.w;
					vec2 cov3D_M13_M22 = unpackInt16(covAndColorData.y) * centerAndScaleData.w;
					vec2 cov3D_M23_M33 = unpackInt16(covAndColorData.z) * centerAndScaleData.w;
					mat3 Vrk = mat3(
						cov3D_M11_M12.x, cov3D_M11_M12.y, cov3D_M13_M22.x,
						cov3D_M11_M12.y, cov3D_M13_M22.y, cov3D_M23_M33.x,
						cov3D_M13_M22.x, cov3D_M23_M33.x, cov3D_M23_M33.y
					);

					mat3 J = mat3(
						focal / camspace.z, 0., -(focal * camspace.x) / (camspace.z * camspace.z), 
						0., -focal / camspace.z, (focal * camspace.y) / (camspace.z * camspace.z), 
						0., 0., 0.
					);

					mat3 W = transpose(mat3(gsModelViewMatrix));
					mat3 T = W * J;
					mat3 cov = transpose(T) * Vrk * T;

					vec2 vCenter = vec2(pos2d) / pos2d.w;

					float diagonal1 = cov[0][0] + 0.3;
					float offDiagonal = cov[0][1];
					float diagonal2 = cov[1][1] + 0.3;

					float mid = 0.5 * (diagonal1 + diagonal2);
					float radius = length(vec2((diagonal1 - diagonal2) / 2.0, offDiagonal));
					float lambda1 = mid + radius;
					float lambda2 = max(mid - radius, 0.1);
					vec2 diagonalVector = normalize(vec2(offDiagonal, lambda1 - diagonal1));
					vec2 v1 = min(sqrt(2.0 * lambda1), 1024.0) * diagonalVector;
					vec2 v2 = min(sqrt(2.0 * lambda2), 1024.0) * vec2(diagonalVector.y, -diagonalVector.x);

					uint colorUint = covAndColorData.w;
					vColor = vec4(
						float(colorUint & uint(0xFF)) / 255.0,
						float((colorUint >> uint(8)) & uint(0xFF)) / 255.0,
						float((colorUint >> uint(16)) & uint(0xFF)) / 255.0,
						float(colorUint >> uint(24)) / 255.0
					);
					vPosition = position.xy;

					gl_Position = vec4(
						vCenter 
							+ position.x * v2 / viewport * 2.0 
							+ position.y * v1 / viewport * 2.0, pos2d.z / pos2d.w, 1.0);
				}
				`,fragmentShader:`
				in vec4 vColor;
				in vec2 vPosition;

				void main () {
					float A = -dot(vPosition, vPosition);
					if (A < -4.0) discard;
					float B = exp(A) * vColor.a;
					gl_FragColor = vec4(vColor.rgb, B);
				}
			`,blending:THREE.CustomBlending,blendSrcAlpha:THREE.OneFactor,depthTest:!0,depthWrite:!1,transparent:!0});material.onBeforeRender=(renderer,scene,camera,geometry,object,group)=>{var projectionMatrix=this.getProjectionMatrix(camera),camera=(mesh.material.uniforms.gsProjectionMatrix.value=projectionMatrix,mesh.material.uniforms.gsModelViewMatrix.value=this.getModelViewMatrix(camera),new THREE.Vector4),renderer=(renderer.getCurrentViewport(camera),camera.w/2*Math.abs(projectionMatrix.elements[5]));material.uniforms.viewport.value[0]=camera.z,material.uniforms.viewport.value[1]=camera.w,material.uniforms.focal.value=renderer},(mesh=new THREE.Mesh(camera,material)).frustumCulled=!1,this.object.add(mesh),this.worker=new Worker(URL.createObjectURL(new Blob(["(",this.createWorker.toString(),")(self)"],{type:"application/javascript"}))),this.worker.onmessage=e=>{e=new Uint32Array(e.data.sortedIndexes);mesh.geometry.attributes.splatIndex.set(e),mesh.geometry.attributes.splatIndex.needsUpdate=!0,mesh.geometry.instanceCount=e.length,this.sortReady=!0},this.sortReady=!0},loadData:function(src){this.loadedVertexCount=0,this.rowLength=32,this.worker.postMessage({method:"clear"});fetch(src).then(async data=>{var reader=data.body.getReader();let bytesDownloaded=0,bytesProcesses=0;var data=data.headers.get("Content-Length"),totalDownloadBytes=data?parseInt(data):void 0,chunks=[],start=Date.now();let lastReportedProgress=0;for(var isPly=src.endsWith(".ply");;)try{var mbps,percent,{value,done}=await reader.read();if(done){console.log("Completed download.");break}bytesDownloaded+=value.length,null!=totalDownloadBytes?(mbps=bytesDownloaded/1024/1024/((Date.now()-start)/1e3),1<(percent=bytesDownloaded/totalDownloadBytes*100)-lastReportedProgress&&(console.log("download progress:",percent.toFixed(2)+"%",mbps.toFixed(2)+" Mbps"),lastReportedProgress=percent)):console.log("download progress:",bytesDownloaded,", unknown total"),chunks.push(value),!this.textureReady&&this.renderer.properties.get(this.centerAndScaleTexture)&&this.renderer.properties.get(this.covAndColorTexture)&&(this.textureReady=!0);var bytesRemains=bytesDownloaded-bytesProcesses;if(!isPly&&this.textureReady&&bytesRemains>this.rowLength){var extra_data,vertexCount=Math.floor(bytesRemains/this.rowLength),concatenatedChunksbuffer=new Uint8Array(bytesRemains);let offset=0;for(const chunk of chunks)concatenatedChunksbuffer.set(chunk,offset),offset+=chunk.length;chunks.length=0,bytesRemains>vertexCount*this.rowLength&&((extra_data=new Uint8Array(bytesRemains-vertexCount*this.rowLength)).set(concatenatedChunksbuffer.subarray(bytesRemains-extra_data.length,bytesRemains),0),chunks.push(extra_data));var buffer=new Uint8Array(vertexCount*this.rowLength);buffer.set(concatenatedChunksbuffer.subarray(0,buffer.byteLength),0),this.pushDataBuffer(buffer.buffer,vertexCount),bytesProcesses+=vertexCount*this.rowLength}}catch(error){console.error(error);break}if(0<bytesDownloaded-bytesProcesses){let concatenatedChunks=new Uint8Array(chunks.reduce((acc,chunk)=>acc+chunk.length,0)),offset=0;for(const chunk of chunks)concatenatedChunks.set(chunk,offset),offset+=chunk.length;isPly&&(concatenatedChunks=new Uint8Array(this.processPlyBuffer(concatenatedChunks.buffer))),this.pushDataBuffer(concatenatedChunks.buffer,Math.floor(concatenatedChunks.byteLength/this.rowLength))}})},pushDataBuffer:function(buffer,vertexCount){if(16777216<this.loadedVertexCount+vertexCount&&(console.log("vertexCount limited to 4096*4096",vertexCount),vertexCount=16777216-this.loadedVertexCount),!(vertexCount<=0)){var u_buffer=new Uint8Array(buffer),f_buffer=new Float32Array(buffer),matrices=new Float32Array(16*vertexCount),covAndColorData_uint8=new Uint8Array(this.covAndColorData.buffer),covAndColorData_int16=new Int16Array(this.covAndColorData.buffer);for(let i=0;i<vertexCount;i++){var quat=new THREE.Quaternion((u_buffer[32*i+28+1]-128)/128,(u_buffer[32*i+28+2]-128)/128,-(u_buffer[32*i+28+3]-128)/128,(u_buffer[32*i+28]-128)/128),center=new THREE.Vector3(f_buffer[8*i+0],f_buffer[8*i+1],-f_buffer[8*i+2]),scale=new THREE.Vector3(f_buffer[8*i+3],f_buffer[8*i+3+1],f_buffer[8*i+3+2]),mtx=new THREE.Matrix4,quat=(mtx.makeRotationFromQuaternion(quat),mtx.transpose(),mtx.scale(scale),mtx.clone()),cov_indexes=(mtx.transpose(),mtx.premultiply(quat),mtx.setPosition(center),[0,1,2,5,6,10]);let max_value=0;for(let j=0;j<cov_indexes.length;j++)Math.abs(mtx.elements[cov_indexes[j]])>max_value&&(max_value=Math.abs(mtx.elements[cov_indexes[j]]));let destOffset=4*this.loadedVertexCount+4*i;this.centerAndScaleData[destOffset+0]=center.x,this.centerAndScaleData[destOffset+1]=center.y,this.centerAndScaleData[destOffset+2]=center.z,this.centerAndScaleData[destOffset+3]=max_value/32767,destOffset=8*this.loadedVertexCount+4*i*2;for(let j=0;j<cov_indexes.length;j++)covAndColorData_int16[destOffset+j]=parseInt(32767*mtx.elements[cov_indexes[j]]/max_value);covAndColorData_uint8[(destOffset=16*this.loadedVertexCount+4*(4*i+3))+0]=u_buffer[32*i+24],covAndColorData_uint8[destOffset+1]=u_buffer[32*i+24+1],covAndColorData_uint8[destOffset+2]=u_buffer[32*i+24+2],covAndColorData_uint8[destOffset+3]=u_buffer[32*i+24+3],mtx.elements[15]=Math.max(scale.x,scale.y,scale.z)*u_buffer[32*i+24+3]/255;for(let j=0;j<16;j++)matrices[16*i+j]=mtx.elements[j]}for(var gl=this.renderer.getContext();0<vertexCount;){let width=0,height=0;var xoffset=this.loadedVertexCount%4096,yoffset=Math.floor(this.loadedVertexCount/4096),centerAndScaleTextureProperties=(height=this.loadedVertexCount%4096!=0?(width=Math.min(4096,xoffset+vertexCount)-xoffset,1):0<Math.floor(vertexCount/4096)?(width=4096,Math.floor(vertexCount/4096)):(width=vertexCount%4096,1),this.renderer.properties.get(this.centerAndScaleTexture)),centerAndScaleTextureProperties=(gl.bindTexture(gl.TEXTURE_2D,centerAndScaleTextureProperties.__webglTexture),gl.texSubImage2D(gl.TEXTURE_2D,0,xoffset,yoffset,width,height,gl.RGBA,gl.FLOAT,this.centerAndScaleData,4*this.loadedVertexCount),this.renderer.properties.get(this.covAndColorTexture));gl.bindTexture(gl.TEXTURE_2D,centerAndScaleTextureProperties.__webglTexture),gl.texSubImage2D(gl.TEXTURE_2D,0,xoffset,yoffset,width,height,gl.RGBA_INTEGER,gl.UNSIGNED_INT,this.covAndColorData,4*this.loadedVertexCount),this.loadedVertexCount+=width*height,vertexCount-=width*height}this.worker.postMessage({method:"push",matrices:matrices.buffer},[matrices.buffer])}},tick:function(time,timeDelta){var camera_mtx;this.sortReady&&(this.sortReady=!1,camera_mtx=this.getModelViewMatrix().elements,camera_mtx=new Float32Array([camera_mtx[2],camera_mtx[6],camera_mtx[10],camera_mtx[14]]),this.worker.postMessage({method:"sort",view:camera_mtx.buffer},[camera_mtx.buffer]))},getProjectionMatrix:function(camera){camera=(camera=camera||this.camera).projectionMatrix.clone();return camera.elements[4]*=-1,camera.elements[5]*=-1,camera.elements[6]*=-1,camera.elements[7]*=-1,camera},getModelViewMatrix:function(camera){var camera=(camera=camera||this.camera).matrixWorld.clone(),mtx=(camera.elements[1]*=-1,camera.elements[4]*=-1,camera.elements[6]*=-1,camera.elements[9]*=-1,camera.elements[13]*=-1,this.object.matrixWorld.clone());return mtx.invert(),mtx.elements[1]*=-1,mtx.elements[4]*=-1,mtx.elements[6]*=-1,mtx.elements[9]*=-1,mtx.elements[13]*=-1,mtx.multiply(camera),mtx.invert(),mtx},createWorker:function(self){let matrices=void 0;function sortSplats(matrices,view){var vertexCount=matrices.length/16;let maxDepth=-1/0,minDepth=1/0;var depthList=new Float32Array(vertexCount),sizeList=new Int32Array(depthList.buffer),validIndexList=new Int32Array(vertexCount);let validCount=0;for(let i=0;i<vertexCount;i++){var depth=view[0]*matrices[16*i+12]+view[1]*matrices[16*i+13]+view[2]*matrices[16*i+14]+view[3];depth<0&&matrices[16*i+15]>-.001*depth&&(depthList[validCount]=depth,validIndexList[validCount]=i,validCount++,depth>maxDepth&&(maxDepth=depth),depth<minDepth)&&(minDepth=depth)}var depthInv=65535/(maxDepth-minDepth),counts0=new Uint32Array(65536);for(let i=0;i<validCount;i++)sizeList[i]=(depthList[i]-minDepth)*depthInv|0,counts0[sizeList[i]]++;var starts0=new Uint32Array(65536);for(let i=1;i<65536;i++)starts0[i]=starts0[i-1]+counts0[i-1];var depthIndex=new Uint32Array(validCount);for(let i=0;i<validCount;i++)depthIndex[starts0[sizeList[i]]++]=validIndexList[i];return depthIndex}self.onmessage=e=>{if("clear"==e.data.method&&(matrices=void 0),"push"==e.data.method&&(new_matrices=new Float32Array(e.data.matrices),matrices=void 0===matrices?new_matrices:((resized=new Float32Array(matrices.length+new_matrices.length)).set(matrices),resized.set(new_matrices,matrices.length),resized)),"sort"==e.data.method)if(void 0===matrices){var sortedIndexes=new Uint32Array(1);self.postMessage({sortedIndexes:sortedIndexes},[sortedIndexes.buffer])}else{e=new Float32Array(e.data.view);const sortedIndexes=sortSplats(matrices,e);self.postMessage({sortedIndexes:sortedIndexes},[sortedIndexes.buffer])}}},processPlyBuffer:function(inputBuffer){var ubuf=new Uint8Array(inputBuffer),ubuf=(new TextDecoder).decode(ubuf.slice(0,10240)),header_end_index=ubuf.indexOf("end_header\n");if(header_end_index<0)throw new Error("Unable to read .ply file header");var vertexCount=parseInt(/element vertex (\d+)\n/.exec(ubuf)[1]);console.log("Vertex Count",vertexCount);let row_offset=0,offsets={},types={};var prop,TYPE_MAP={double:"getFloat64",int:"getInt32",uint:"getUint32",float:"getFloat32",short:"getInt16",ushort:"getUint16",uchar:"getUint8"};for(prop of ubuf.slice(0,header_end_index).split("\n").filter(k=>k.startsWith("property "))){var[,type,name]=prop.split(" "),type=TYPE_MAP[type]||"getInt8";types[name]=type,offsets[name]=row_offset,row_offset+=parseInt(type.replace(/[^\d]/g,""))/8}console.log("Bytes per row",row_offset,types,offsets);let dataView=new DataView(inputBuffer,header_end_index+"end_header\n".length),row=0;var attrs=new Proxy({},{get(target,prop){if(types[prop])return dataView[types[prop]](row*row_offset+offsets[prop],!0);throw new Error(prop+" not found")}});console.time("calculate importance");let sizeList=new Float32Array(vertexCount);var size,opacity,sizeIndex=new Uint32Array(vertexCount);for(row=0;row<vertexCount;row++)sizeIndex[row]=row,types.scale_0&&(size=Math.exp(attrs.scale_0)*Math.exp(attrs.scale_1)*Math.exp(attrs.scale_2),opacity=1/(1+Math.exp(-attrs.opacity)),sizeList[row]=size*opacity);console.timeEnd("calculate importance"),console.time("sort"),sizeIndex.sort((b,a)=>sizeList[a]-sizeList[b]),console.timeEnd("sort");var buffer=new ArrayBuffer(32*vertexCount);console.time("build buffer");for(let j=0;j<vertexCount;j++){row=sizeIndex[j];var qlen,position=new Float32Array(buffer,32*j,3),scales=new Float32Array(buffer,32*j+12,3),rgba=new Uint8ClampedArray(buffer,32*j+12+12,4),rot=new Uint8ClampedArray(buffer,32*j+12+12+4,4);types.scale_0?(qlen=Math.sqrt(attrs.rot_0**2+attrs.rot_1**2+attrs.rot_2**2+attrs.rot_3**2),rot[0]=attrs.rot_0/qlen*128+128,rot[1]=attrs.rot_1/qlen*128+128,rot[2]=attrs.rot_2/qlen*128+128,rot[3]=attrs.rot_3/qlen*128+128,scales[0]=Math.exp(attrs.scale_0),scales[1]=Math.exp(attrs.scale_1),scales[2]=Math.exp(attrs.scale_2)):(scales[0]=.01,scales[1]=.01,scales[2]=.01,rot[0]=255,rot[1]=0,rot[2]=0,rot[3]=0),position[0]=attrs.x,position[1]=attrs.y,position[2]=attrs.z,types.f_dc_0?(rgba[0]=255*(.5+(qlen=.28209479177387814)*attrs.f_dc_0),rgba[1]=255*(.5+qlen*attrs.f_dc_1),rgba[2]=255*(.5+qlen*attrs.f_dc_2)):(rgba[0]=attrs.red,rgba[1]=attrs.green,rgba[2]=attrs.blue),types.opacity?rgba[3]=1/(1+Math.exp(-attrs.opacity))*255:rgba[3]=255}return console.timeEnd("build buffer"),buffer}});