playcanvas
Version:
PlayCanvas WebGL game engine
223 lines (220 loc) • 8.08 kB
JavaScript
import { Vec3 } from '../../core/math/vec3.js';
import { Mat4 } from '../../core/math/mat4.js';
import { BlendState } from '../../platform/graphics/blend-state.js';
import { SEMANTIC_POSITION, PIXELFORMAT_RGBA8, CULLFACE_NONE } from '../../platform/graphics/constants.js';
import { DepthState } from '../../platform/graphics/depth-state.js';
import { RenderTarget } from '../../platform/graphics/render-target.js';
import { ShaderUtils } from '../shader-lib/shader-utils.js';
import { ShaderChunks } from '../shader-lib/shader-chunks.js';
import { Vec2 } from '../../core/math/vec2.js';
import { RenderPass } from '../../platform/graphics/render-pass.js';
import { QuadRender } from '../graphics/quad-render.js';
const vertexGLSL = `
attribute vec2 vertex_position;
void main(void) {
gl_Position = vec4(vertex_position, 0.0, 1.0);
}
`;
const fragmentGLSL = `
#include "gsplatEvalSHVS"
vec4 packRgb(vec3 v) {
uvec3 vb = uvec3(clamp(v, vec3(0.0), vec3(1.0)) * vec3(2047.0, 2047.0, 1023.0));
uint bits = (vb.x << 21) | (vb.y << 10) | vb.z;
return vec4((uvec4(bits) >> uvec4(24, 16, 8, 0)) & uvec4(0xff)) / vec4(255.0);
}
uniform mediump vec3 dir;
uniform mediump sampler2D centroids;
uniform mediump float shN_mins;
uniform mediump float shN_maxs;
void main(void) {
ivec2 uv = ivec2(gl_FragCoord.xy) * ivec2(SH_COEFFS, 1);
mediump vec3 coefficients[SH_COEFFS];
for (int i = 0; i < SH_COEFFS; i++) {
vec3 s = texelFetch(centroids, ivec2(uv.x + i, uv.y), 0).xyz;
coefficients[i] = mix(vec3(shN_mins), vec3(shN_maxs), s);
}
gl_FragColor = packRgb(evalSH(coefficients, dir) * 0.25 + 0.5);
}
`;
const vertexWGSL = `
attribute vertex_position: vec2f;
fn vertexMain(input: VertexInput) -> VertexOutput {
var output: VertexOutput;
output.position = vec4f(vertex_position, 0.0, 1.0);
return output;
}
`;
const fragmentWGSL = `
#include "gsplatEvalSHVS"
fn packRgb(v: vec3f) -> vec4f {
let vb = vec3u(clamp(v, vec3f(0.0), vec3f(1.0)) * vec3f(2047.0, 2047.0, 1023.0));
let bits = dot(vb, vec3u(1 << 21, 1 << 10, 1));
return vec4f((vec4u(bits) >> vec4u(24, 16, 8, 0)) & vec4u(0xff)) / vec4f(255.0);
}
uniform dir: vec3f;
uniform shN_mins: f32;
uniform shN_maxs: f32;
var centroids: texture_2d<f32>;
fn fragmentMain(input: FragmentInput) -> FragmentOutput {
var output: FragmentOutput;
var uv = vec2i(input.position.xy) * vec2i(SH_COEFFS, 1);
var coefficients: array<vec3f, SH_COEFFS>;
for (var i: i32 = 0; i < SH_COEFFS; i++) {
let s: vec3f = textureLoad(centroids, vec2i(uv.x + i, uv.y), 0).xyz;
coefficients[i] = mix(vec3f(uniform.shN_mins), vec3f(uniform.shN_maxs), s);
}
output.color = packRgb(evalSH(&coefficients, uniform.dir) * 0.25 + 0.5);
return output;
}
`;
const gsplatSogsColorGLSL = `
uniform mediump sampler2D sh0;
uniform highp sampler2D sh_labels;
uniform mediump sampler2D sh_result;
uniform vec4 sh0_mins;
uniform vec4 sh0_maxs;
float SH_C0 = 0.28209479177387814;
vec3 unpackRgb(vec4 v) {
uvec4 uv = uvec4(v * 255.0);
uint bits = (uv.x << 24) | (uv.y << 16) | (uv.z << 8) | uv.w;
uvec3 vb = (uvec3(bits) >> uvec3(21, 10, 0)) & uvec3(0x7ffu, 0x7ffu, 0x3ffu);
return vec3(vb) / vec3(2047.0, 2047.0, 1023.0);
}
vec4 readColor(in SplatSource source) {
vec4 baseSample = mix(sh0_mins, sh0_maxs, texelFetch(sh0, source.uv, 0));
vec4 base = vec4(vec3(0.5) + baseSample.xyz * SH_C0, 1.0 / (1.0 + exp(-baseSample.w)));
ivec2 labelSample = ivec2(texelFetch(sh_labels, source.uv, 0).xy * 255.0);
int n = labelSample.x + labelSample.y * 256;
vec4 shSample = texelFetch(sh_result, ivec2(n % 64, n / 64), 0);
vec3 sh = (unpackRgb(shSample) - vec3(0.5)) * 4.0;
return vec4(base.xyz + sh, base.w);
}
`;
const gsplatSogsColorWGSL = `
var sh0: texture_2d<f32>;
var sh_labels: texture_2d<f32>;
var sh_result: texture_2d<f32>;
uniform sh0_mins: vec4f;
uniform sh0_maxs: vec4f;
const SH_C0: f32 = 0.28209479177387814;
fn unpackRgb(v: vec4f) -> vec3f {
let bits = dot(vec4u(v * 255.0), vec4u(1u << 24, 1u << 16, 1u << 8, 1u));
let vb = (vec3u(bits) >> vec3u(21, 10, 0)) & vec3u(0x7ffu, 0x7ffu, 0x3ffu);
return vec3f(vb) / vec3f(2047.0, 2047.0, 1023.0);
}
fn readColor(source: ptr<function, SplatSource>) -> vec4f {
let baseSample: vec4f = mix(uniform.sh0_mins, uniform.sh0_maxs, textureLoad(sh0, source.uv, 0));
let base = vec4f(vec3f(0.5) + baseSample.xyz * SH_C0, 1.0 / (1.0 + exp(-baseSample.w)));
let labelSample: vec2i = vec2i(textureLoad(sh_labels, source.uv, 0).xy * 255.0);
let n = labelSample.x + labelSample.y * 256;
let shSample: vec4f = textureLoad(sh_result, vec2i(n % 64, n / 64), 0);
let sh: vec3f = (unpackRgb(shSample) - vec3f(0.5)) * 4.0;
return vec4f(base.xyz + sh, base.w);
}
`;
const resolve = (scope, values)=>{
for(const key in values){
scope.resolve(key).setValue(values[key]);
}
};
class CustomRenderPass extends RenderPass {
execute() {
this.executeCallback?.();
}
constructor(...args){
super(...args), this.executeCallback = null;
}
}
const invModelMat = new Mat4();
const dir = new Vec3();
class GSplatResolveSH {
constructor(device, gsplatInstance){
this.prevDir = new Vec3();
this.updateMode = 'enable';
this.device = device;
this.gsplatInstance = gsplatInstance;
const { resource } = gsplatInstance;
const includes = new Map(ShaderChunks.get(device, device.isWebGPU ? 'wgsl' : 'glsl'));
this.shader = ShaderUtils.createShader(device, {
uniqueName: 'gsplatResolveSH',
vertexGLSL,
fragmentGLSL,
vertexWGSL,
fragmentWGSL,
vertexIncludes: includes,
fragmentIncludes: includes,
fragmentDefines: new Map([
[
'SH_BANDS',
resource.gsplatData.shBands.toString()
]
]),
attributes: {
vertex_position: SEMANTIC_POSITION
}
});
this.texture = resource.createTexture('centroids', PIXELFORMAT_RGBA8, new Vec2(64, 1024));
this.renderTarget = new RenderTarget({
colorBuffer: this.texture,
depth: false
});
this.renderPass = new CustomRenderPass(device);
this.renderPass.init(this.renderTarget, {});
this.renderPass.colorOps.clear = true;
this.quadRender = new QuadRender(this.shader);
const { material } = gsplatInstance;
material.setDefine('SH_BANDS', '0');
const { shaderChunks } = material;
shaderChunks.glsl.set('gsplatSogsColorVS', gsplatSogsColorGLSL);
shaderChunks.wgsl.set('gsplatSogsColorVS', gsplatSogsColorWGSL);
material.update();
device.scope.resolve('sh_result').setValue(this.texture);
}
destroy() {
const { gsplatInstance } = this;
const { material } = gsplatInstance;
material.setDefine('SH_BANDS', gsplatInstance.resource.gsplatData.shBands.toString());
const { shaderChunks } = material;
shaderChunks.glsl.delete('gsplatSogsColorVS');
shaderChunks.wgsl.delete('gsplatSogsColorVS');
material.update();
this.quadRender.destroy();
this.renderPass.destroy();
this.renderTarget.destroy();
this.texture.destroy();
this.shader.destroy();
}
render(camera, modelMat) {
const { prevDir, updateMode } = this;
if (updateMode === 'disable') {
return;
}
invModelMat.invert(modelMat);
invModelMat.transformVector(camera.forward, dir);
dir.normalize();
if (updateMode === 'enable' && dir.equalsApprox(prevDir, 1e-3)) {
return;
}
prevDir.copy(dir);
const execute = ()=>{
const { device } = this;
const { sh_centroids, meta } = this.gsplatInstance.resource.gsplatData;
resolve(device.scope, {
dir: dir.toArray(),
centroids: sh_centroids,
shN_mins: meta.shN.mins,
shN_maxs: meta.shN.maxs
});
device.setCullMode(CULLFACE_NONE);
device.setDepthState(DepthState.NODEPTH);
device.setStencilState(null, null);
device.setBlendState(BlendState.NOBLEND);
this.quadRender.render();
};
this.renderPass.executeCallback = execute;
this.renderPass.render();
}
}
export { GSplatResolveSH };