playcanvas/build/playcanvas/src/scene/gsplat/gsplat-instance.js

Open-source WebGL/WebGPU 3D engine for the web

import { Mat4 } from '../../core/math/mat4.js';
import { Vec3 } from '../../core/math/vec3.js';
import { CULLFACE_NONE, BUFFERUSAGE_COPY_DST, PIXELFORMAT_R32U, SEMANTIC_POSITION } from '../../platform/graphics/constants.js';
import { StorageBuffer } from '../../platform/graphics/storage-buffer.js';
import { MeshInstance } from '../mesh-instance.js';
import { GSplatResolveSH } from './gsplat-resolve-sh.js';
import { GSplatSorter } from './gsplat-sorter.js';
import { GSplatSogData } from './gsplat-sog-data.js';
import { GSplatResourceBase } from './gsplat-resource-base.js';
import { ShaderMaterial } from '../materials/shader-material.js';
import { BLEND_NONE, BLEND_PREMULTIPLIED } from '../constants.js';

const mat = new Mat4();
const cameraPosition = new Vec3();
const cameraDirection = new Vec3();
class GSplatInstance {
		destroy() {
				this.resource?.releaseMesh();
				this.orderTexture?.destroy();
				this.orderBuffer?.destroy();
				this.resolveSH?.destroy();
				this.material?.destroy();
				this.meshInstance?.destroy();
				this.sorter?.destroy();
		}
		setMaterialOrderData(material) {
				if (this.orderBuffer) {
						material.setParameter('splatOrder', this.orderBuffer);
				} else {
						material.setParameter('splatOrder', this.orderTexture);
						material.setParameter('splatTextureSize', this.orderTexture.width);
				}
		}
		set material(value) {
				if (this._material !== value) {
						this._material = value;
						this._material.setDefine('{GSPLAT_INSTANCE_SIZE}', String(GSplatResourceBase.instanceSize));
						this.setMaterialOrderData(this._material);
						if (this.meshInstance) {
								this.meshInstance.material = value;
						}
				}
		}
		get material() {
				return this._material;
		}
		configureMaterial(material, options = {}) {
				this.resource.configureMaterial(material, null, this.resource.format.getInputDeclarations());
				material.setDefine('{GSPLAT_INSTANCE_SIZE}', GSplatResourceBase.instanceSize);
				material.setParameter('numSplats', 0);
				this.setMaterialOrderData(material);
				material.setParameter('alphaClip', 0.3);
				material.setParameter('minPixelSize', 2.0);
				material.setDefine(`DITHER_${options.dither ? 'BLUENOISE' : 'NONE'}`, '');
				material.cull = CULLFACE_NONE;
				material.blendType = options.dither ? BLEND_NONE : BLEND_PREMULTIPLIED;
				material.depthWrite = !!options.dither;
		}
		sort(cameraNode) {
				if (this.sorter) {
						const cameraMat = cameraNode.getWorldTransform();
						cameraMat.getTranslation(cameraPosition);
						cameraMat.getZ(cameraDirection);
						const modelMat = this.meshInstance.node.getWorldTransform();
						const invModelMat = mat.invert(modelMat);
						invModelMat.transformPoint(cameraPosition, cameraPosition);
						invModelMat.transformVector(cameraDirection, cameraDirection);
						if (!cameraPosition.equalsApprox(this.lastCameraPosition) || !cameraDirection.equalsApprox(this.lastCameraDirection)) {
								this.lastCameraPosition.copy(cameraPosition);
								this.lastCameraDirection.copy(cameraDirection);
								this.sorter.setCamera(cameraPosition, cameraDirection);
						}
				}
		}
		update() {
				const count = this.sorter?.applyPendingSorted() ?? -1;
				if (count >= 0) {
						this.meshInstance.instancingCount = Math.ceil(count / GSplatResourceBase.instanceSize);
						this.material.setParameter('numSplats', count);
				}
				if (this.cameras.length > 0) {
						const camera = this.cameras[0];
						this.sort(camera._node);
						this.resolveSH?.render(camera._node, this.meshInstance.node.getWorldTransform());
						this.cameras.length = 0;
				}
		}
		setHighQualitySH(value) {
				const { resource } = this;
				const { gsplatData } = resource;
				if (gsplatData instanceof GSplatSogData && gsplatData.shBands > 0 && value === !!this.resolveSH) {
						if (this.resolveSH) {
								this.resolveSH.destroy();
								this.resolveSH = null;
						} else {
								this.resolveSH = new GSplatResolveSH(resource.device, this);
						}
				}
		}
		constructor(resource, options = {}){
				this.options = {};
				this.sorter = null;
				this.lastCameraPosition = new Vec3();
				this.lastCameraDirection = new Vec3();
				this.resolveSH = null;
				this.cameras = [];
				this.resource = resource;
				const device = resource.device;
				const dims = resource.streams.textureDimensions;
				const numSplats = dims.x * dims.y;
				if (device.isWebGPU) {
						this.orderBuffer = new StorageBuffer(device, numSplats * 4, BUFFERUSAGE_COPY_DST);
				} else {
						this.orderTexture = resource.streams.createTexture('splatOrder', PIXELFORMAT_R32U, dims);
				}
				if (options.material) {
						this._material = options.material;
						this._material.setDefine('{GSPLAT_INSTANCE_SIZE}', String(GSplatResourceBase.instanceSize));
						this.setMaterialOrderData(this._material);
				} else {
						this._material = new ShaderMaterial({
								uniqueName: 'SplatMaterial',
								vertexGLSL: '#include "gsplatVS"',
								fragmentGLSL: '#include "gsplatPS"',
								vertexWGSL: '#include "gsplatVS"',
								fragmentWGSL: '#include "gsplatPS"',
								attributes: {
										vertex_position: SEMANTIC_POSITION
								}
						});
						this.configureMaterial(this._material);
						this._material.update();
				}
				resource.ensureMesh();
				this.meshInstance = new MeshInstance(resource.mesh, this._material);
				this.meshInstance.setInstancing(true, true);
				this.meshInstance.gsplatInstance = this;
				this.meshInstance.instancingCount = 0;
				const centers = resource.centers.slice();
				const chunks = resource.chunks?.slice();
				const orderTarget = this.orderBuffer ?? this.orderTexture;
				this.sorter = new GSplatSorter(device, options.scene);
				this.sorter.init(orderTarget, numSplats, centers, chunks);
				this.setHighQualitySH(options.highQualitySH ?? false);
		}
}

export { GSplatInstance };