@babylonjs/core
Version:
Getting started? Play directly with the Babylon.js API using our [playground](https://playground.babylonjs.com/). It also contains a lot of samples to learn how to use it.
232 lines • 8.17 kB
JavaScript
import { MaterialPluginBase } from "../materialPluginBase.js";
import { RegisterClass } from "../../Misc/typeStore.js";
import { GetGaussianSplattingMaxPartCount } from "./gaussianSplattingMaterial.js";
/**
* Plugin for GaussianSplattingMaterial that replaces per-splat color output with
* a pre-computed picking color for GPU-based hit testing.
*
* The picking color is computed on the CPU by encoding a 24-bit picking ID as RGB
* (matching the readback decoding in GPUPicker).
* @experimental
*/
export class GaussianSplattingGpuPickingMaterialPlugin extends MaterialPluginBase {
/**
* Creates a new GaussianSplattingGpuPickingMaterialPlugin.
* @param material The GaussianSplattingMaterial to attach the plugin to.
* @param maxPartCount The maximum number of parts supported for compound meshes.
*/
constructor(material, maxPartCount) {
super(material, "GaussianSplatGpuPicking", 200);
this._pickingColor = [0, 0, 0];
this._isCompound = false;
this._partPickingColors = [];
this._partVisibility = [];
this._maxPartCount = maxPartCount ?? GetGaussianSplattingMaxPartCount(material.getScene().getEngine());
this._enable(true);
}
/**
* Encodes a 24-bit picking ID into normalized RGB components.
* @param id The picking ID to encode
* @returns A tuple [r, g, b] with values in [0, 1]
*/
static EncodeIdToColor(id) {
return [((id >> 16) & 0xff) / 255, ((id >> 8) & 0xff) / 255, (id & 0xff) / 255];
}
/**
* Sets the picking color for a non-compound mesh from a picking ID.
* The ID is encoded into an RGB color on the CPU.
* @param id The 24-bit picking ID.
*/
set meshId(id) {
this._pickingColor = GaussianSplattingGpuPickingMaterialPlugin.EncodeIdToColor(id);
}
/**
* Sets whether this material is for a compound mesh with per-part picking.
*/
set isCompound(value) {
this._isCompound = value;
this.markAllDefinesAsDirty();
}
/**
* Gets whether this material is for a compound mesh with per-part picking.
*/
get isCompound() {
return this._isCompound;
}
/**
* Sets the per-part picking colors from an array of picking IDs.
* Each ID is encoded into an RGB color on the CPU.
* @param ids Array mapping part index to picking ID.
*/
set partMeshIds(ids) {
const colors = [];
for (let i = 0; i < this._maxPartCount; i++) {
const c = i < ids.length ? GaussianSplattingGpuPickingMaterialPlugin.EncodeIdToColor(ids[i]) : [0, 0, 0];
colors.push(c[0], c[1], c[2]);
}
this._partPickingColors = colors;
}
/**
* Sets which parts are active (pickable) for the compound picking pass.
* Parts not in the set are discarded in the shader by overriding partVisibility to 0.
* @param activeParts Array of part indices that should be pickable.
*/
setPartActive(activeParts) {
const visibility = new Array(this._maxPartCount).fill(0.0);
for (const index of activeParts) {
if (index >= 0 && index < this._maxPartCount) {
visibility[index] = 1.0;
}
}
this._partVisibility = visibility;
}
/**
* @returns the class name
*/
getClassName() {
return "GaussianSplattingGpuPickingMaterialPlugin";
}
/**
* Indicates this plugin supports both GLSL and WGSL.
* @param shaderLanguage the shader language to check
* @returns true for GLSL and WGSL
*/
isCompatible(shaderLanguage) {
switch (shaderLanguage) {
case 0 /* ShaderLanguage.GLSL */:
case 1 /* ShaderLanguage.WGSL */:
return true;
default:
return false;
}
}
/**
* Always ready — no textures or async resources to wait on.
* @param _defines the defines
* @param _scene the scene
* @param _engine the engine
* @param _subMesh the submesh
* @returns true
*/
isReadyForSubMesh(_defines, _scene, _engine, _subMesh) {
return true;
}
/**
* Returns custom shader code to inject GPU picking color output.
*
* @param shaderType "vertex" or "fragment"
* @param shaderLanguage the shader language to use (default: GLSL)
* @returns null or a map of injection point names to code strings
*/
getCustomCode(shaderType, shaderLanguage = 0 /* ShaderLanguage.GLSL */) {
if (shaderLanguage === 1 /* ShaderLanguage.WGSL */) {
return this._getCustomCodeWGSL(shaderType);
}
return this._getCustomCodeGLSL(shaderType);
}
_getCustomCodeGLSL(shaderType) {
if (shaderType === "vertex") {
return {
CUSTOM_VERTEX_DEFINITIONS: `varying float vPartIndex;`,
CUSTOM_VERTEX_UPDATE: `
#if IS_COMPOUND
vPartIndex = float(splat.partIndex);
#else
vPartIndex = 0.0;
#endif
`,
};
}
else if (shaderType === "fragment") {
return {
CUSTOM_FRAGMENT_DEFINITIONS: `
varying float vPartIndex;
#if IS_COMPOUND
uniform vec3 partPickingColors[${this._maxPartCount}];
#else
uniform vec3 pickingColor;
#endif
`,
CUSTOM_FRAGMENT_BEFORE_FRAGCOLOR: `
#if IS_COMPOUND
if (vColor.a < 0.001) discard;
finalColor = vec4(partPickingColors[int(vPartIndex + 0.5)], 1.0);
#else
if (vColor.a < 0.001) discard;
finalColor = vec4(pickingColor, 1.0);
#endif
`,
};
}
return null;
}
_getCustomCodeWGSL(shaderType) {
if (shaderType === "vertex") {
return {
CUSTOM_VERTEX_DEFINITIONS: `varying vPartIndex: f32;`,
CUSTOM_VERTEX_UPDATE: `
#if IS_COMPOUND
vertexOutputs.vPartIndex = f32(splat.partIndex);
#else
vertexOutputs.vPartIndex = 0.0;
#endif
`,
};
}
else if (shaderType === "fragment") {
return {
CUSTOM_FRAGMENT_DEFINITIONS: `
varying vPartIndex: f32;
#if IS_COMPOUND
uniform partPickingColors: array<vec3f, ${this._maxPartCount}>;
#else
uniform pickingColor: vec3f;
#endif
`,
CUSTOM_FRAGMENT_BEFORE_FRAGCOLOR: `
#if IS_COMPOUND
if (fragmentInputs.vColor.a < 0.001) { discard; }
finalColor = vec4f(uniforms.partPickingColors[i32(fragmentInputs.vPartIndex + 0.5)], 1.0);
#else
if (fragmentInputs.vColor.a < 0.001) { discard; }
finalColor = vec4f(uniforms.pickingColor, 1.0);
#endif
`,
};
}
return null;
}
/**
* Registers the picking uniforms with the engine.
* @returns uniform descriptions
*/
getUniforms() {
return {
externalUniforms: ["pickingColor", "partPickingColors"],
};
}
/**
* Binds the picking color uniform(s) each frame.
* @param _uniformBuffer the uniform buffer (unused — we bind directly on the effect)
* @param _scene the current scene
* @param _engine the current engine
* @param subMesh the submesh being rendered
*/
bindForSubMesh(_uniformBuffer, _scene, _engine, subMesh) {
const effect = subMesh.effect;
if (!effect) {
return;
}
if (this._isCompound) {
effect.setArray3("partPickingColors", this._partPickingColors);
// default all visible when setPartActive hasn't been called
const visibility = this._partVisibility.length > 0 ? this._partVisibility : new Array(this._maxPartCount).fill(1.0);
effect.setArray("partVisibility", visibility);
}
else {
effect.setFloat3("pickingColor", this._pickingColor[0], this._pickingColor[1], this._pickingColor[2]);
}
}
}
RegisterClass("BABYLON.GaussianSplattingGpuPickingMaterialPlugin", GaussianSplattingGpuPickingMaterialPlugin);
//# sourceMappingURL=gaussianSplattingGpuPickingMaterialPlugin.js.map