From 0e9f94d6fc63740d2a05dcff209a987ae5c763f8 Mon Sep 17 00:00:00 2001 From: David Catuhe Date: Wed, 5 Jun 2024 17:01:53 -0700 Subject: [PATCH] Create a new picking mechanism using GPU (#15166) * step 1 * step 2 * Final step for v1 * Adding support for Instances * Force removal of previous one * clean up the added vertex buffer --- packages/dev/core/src/Collisions/gpuPicker.ts | 314 ++++++++++++++++++ packages/dev/core/src/Collisions/index.ts | 1 + .../dev/core/src/Helpers/environmentHelper.ts | 2 + .../dev/core/src/Materials/shaderMaterial.ts | 11 +- .../ReferenceImages/gpuPicker.png | Bin 0 -> 2560 bytes .../tests/test/visualization/config.json | 6 + 6 files changed, 333 insertions(+), 1 deletion(-) create mode 100644 packages/dev/core/src/Collisions/gpuPicker.ts create mode 100644 packages/tools/tests/test/visualization/ReferenceImages/gpuPicker.png diff --git a/packages/dev/core/src/Collisions/gpuPicker.ts b/packages/dev/core/src/Collisions/gpuPicker.ts new file mode 100644 index 00000000000..c7188d27cc8 --- /dev/null +++ b/packages/dev/core/src/Collisions/gpuPicker.ts @@ -0,0 +1,314 @@ +import { Constants } from "core/Engines/constants"; +import type { Engine } from "core/Engines/engine"; +import type { WebGPUEngine } from "core/Engines/webgpuEngine"; +import { RenderTargetTexture } from "core/Materials/Textures/renderTargetTexture"; +import type { Material } from "core/Materials/material"; +import { ShaderMaterial } from "core/Materials/shaderMaterial"; +import { Color3, Color4 } from "core/Maths/math.color"; +import type { AbstractMesh } from "core/Meshes/abstractMesh"; +import { VertexBuffer } from "core/Meshes/buffer"; +import type { Mesh } from "core/Meshes/mesh"; +import type { Scene } from "core/scene"; +import type { Nullable } from "core/types"; + +/** + * Class used to perform a picking operation using GPU + * Please note that GPUPIcker cannot pick instances, only meshes + */ +export class GPUPicker { + private _pickingTexure: Nullable = null; + private _idMap: { [key: string]: number } = {}; + private _idColors: Array = []; + private _cachedMaterials: Array> = []; + private _cachedScene: Nullable; + private _renderMaterial: Nullable; + private _pickableMeshes: Array; + private _readbuffer: Uint8Array; + private _meshRenderingCount: number = 0; + private _userDefinedList = false; + + private _createRenderTarget(scene: Scene, width: number, height: number) { + this._pickingTexure = new RenderTargetTexture( + "pickingTexure", + { width: width, height: height }, + scene, + false, + undefined, + Constants.TEXTURETYPE_UNSIGNED_INT, + false, + Constants.TEXTURE_NEAREST_NEAREST + ); + } + + private _createColorMaterial(scene: Scene) { + if (this._renderMaterial) { + this._renderMaterial.dispose(); + } + + const defines: string[] = []; + const options = { + attributes: [VertexBuffer.PositionKind], + uniforms: ["world", "viewProjection", "color"], + needAlphaBlending: false, + defines: defines, + useClipPlane: null, + }; + + this._renderMaterial = new ShaderMaterial("colorShader", scene, "color", options, false); + + const callback = (mesh: AbstractMesh | undefined) => { + if (!mesh) { + return; + } + + const effect = this._renderMaterial!.getEffect(); + + if (!mesh.hasInstances && !mesh.isAnInstance) { + effect.setColor4("color", this._idColors[mesh.uniqueId], 1); + } + + this._meshRenderingCount++; + }; + + this._renderMaterial.customBindingObservable.add(callback); + } + + /** + * Set the list of meshes to pick from + * @param list defines the list of meshes to pick from + */ + public setPickingList(list: Nullable>) { + if (!list) { + if (this._userDefinedList) { + for (let index = 0; index < this._pickableMeshes.length; index++) { + const mesh = this._pickableMeshes[index]; + if (mesh.hasInstances) { + (mesh as Mesh).removeVerticesData(VertexBuffer.ColorKind); + } + } + } + this._userDefinedList = false; + this._pickableMeshes = []; + this._idMap = {}; + this._idColors = []; + return; + } + this._userDefinedList = true; + this._pickableMeshes = list; + + // We will affect colors and create vertex color buffers + let id = 1; + for (let index = 0; index < this._pickableMeshes.length; index++) { + const mesh = this._pickableMeshes[index]; + mesh.useVertexColors = false; + + if (mesh.isAnInstance) { + continue; // This will be handled by the source mesh + } + + const r = (id & 0xff0000) >> 16; + const g = (id & 0x00ff00) >> 8; + const b = (id & 0x0000ff) >> 0; + this._idMap[`${r}_${g}_${b}`] = index; + id++; + + if (mesh.hasInstances) { + const instances = (mesh as Mesh).instances; + const colorData = new Float32Array(4 * (instances.length + 1)); + const engine = mesh.getEngine(); + + colorData[0] = r / 255.0; + colorData[1] = g / 255.0; + colorData[2] = b / 255.0; + colorData[3] = 1.0; + for (let i = 0; i < instances.length; i++) { + const instance = instances[i]; + const r = (id & 0xff0000) >> 16; + const g = (id & 0x00ff00) >> 8; + const b = (id & 0x0000ff) >> 0; + this._idMap[`${r}_${g}_${b}`] = this._pickableMeshes.indexOf(instance); + + colorData[(i + 1) * 4] = r / 255.0; + colorData[(i + 1) * 4 + 1] = g / 255.0; + colorData[(i + 1) * 4 + 2] = b / 255.0; + colorData[(i + 1) * 4 + 3] = 1.0; + id++; + } + + const buffer = new VertexBuffer(engine, colorData, VertexBuffer.ColorKind, false, false, 4, true); + (mesh as Mesh).setVerticesBuffer(buffer, true); + } else { + this._idColors[mesh.uniqueId] = Color3.FromInts(r, g, b); + } + } + } + + /** + * Execute a picking operation + * @param x defines the X coordinates where to run the pick + * @param y defines the Y coordinates where to run the pick + * @param scene defines the scene to pick from + * @param disposeWhenDone defines a boolean indicating we do not want to keep resources alive + * @returns A promise with the picking results + */ + public pickAsync(x: number, y: number, scene: Scene, disposeWhenDone = true): Promise> { + const engine = scene.getEngine(); + const rttSizeW = engine.getRenderWidth(); + const rttSizeH = engine.getRenderHeight(); + + if (!this._readbuffer) { + this._readbuffer = new Uint8Array(engine.isWebGPU ? 256 : 4); // Because of block alignment in WebGPU + } + + this._meshRenderingCount = 0; + // Ensure ints + x = x >> 0; + y = y >> 0; + + if (x < 0 || y < 0 || x >= rttSizeW || y >= rttSizeH) { + return Promise.resolve(null); + } + + // Invert Y + y = rttSizeH - y; + + if (!this._pickingTexure) { + this._createRenderTarget(scene, rttSizeW, rttSizeH); + } else { + const size = this._pickingTexure.getSize(); + + if (size.width !== rttSizeW || size.height !== rttSizeH || this._cachedScene !== scene) { + this._pickingTexure.dispose(); + this._createRenderTarget(scene, rttSizeW, rttSizeH); + } + } + + if (!this._cachedScene || this._cachedScene !== scene) { + this._createColorMaterial(scene); + } + + this._cachedScene = scene; + scene.customRenderTargets.push(this._pickingTexure!); + this._pickingTexure!.renderList = []; + this._pickingTexure!.clearColor = new Color4(0, 0, 0, 0); + + // We need to give every mesh an unique color (when there is no picking list) + this._pickingTexure!.onBeforeRender = () => { + if (!this._userDefinedList) { + this._pickableMeshes = scene.meshes.filter((m) => (m as Mesh).geometry && m.isPickable && (m as Mesh).onBeforeRenderObservable) as Mesh[]; + } + + for (let index = 0; index < this._pickableMeshes.length; index++) { + const mesh = this._pickableMeshes[index]; + if (!mesh.isAnInstance) { + this._cachedMaterials[index] = mesh.material; + if (!this._userDefinedList) { + // We need to define the color for each mesh as it was not previously defined + const id = index + 1; + const r = (id & 0xff0000) >> 16; + const g = (id & 0x00ff00) >> 8; + const b = (id & 0x0000ff) >> 0; + this._idMap[`${r}_${g}_${b}`] = index; + + this._idColors[mesh.uniqueId] = Color3.FromInts(r, g, b); + } else { + mesh.useVertexColors = true; // In that case we will be using vertex colors and not an uniform to support instances + } + mesh.material = this._renderMaterial; + } + this._pickingTexure?.renderList?.push(mesh); + } + + // Enable scissor + if ((engine as WebGPUEngine | Engine).enableScissor) { + (engine as WebGPUEngine | Engine).enableScissor(x, y, 1, 1); + } + }; + + return new Promise((resolve, reject) => { + this._pickingTexure!.onAfterRender = async () => { + // Disable scissor + if ((engine as WebGPUEngine | Engine).disableScissor) { + (engine as WebGPUEngine | Engine).disableScissor(); + } + + if (!this._pickingTexure) { + reject(); + } + + let pickedMesh: Nullable = null; + const wasSuccessfull = this._meshRenderingCount > 0; + + // Restore materials + for (let index = 0; index < this._pickableMeshes.length; index++) { + const mesh = this._pickableMeshes[index]; + if (!mesh.isAnInstance) { + mesh.material = this._cachedMaterials[index]; + + if (this._userDefinedList) { + mesh.useVertexColors = false; + } + } + } + + if (wasSuccessfull) { + // Remove from the active RTTs + const index = scene.customRenderTargets.indexOf(this._pickingTexure!); + if (index > -1) { + scene.customRenderTargets.splice(index, 1); + } + + // Do the actual picking + if (await this._readTexturePixelsAsync(x, y)) { + const colorId = `${this._readbuffer[0]}_${this._readbuffer[1]}_${this._readbuffer[2]}`; + pickedMesh = this._pickableMeshes[this._idMap[colorId]]; + } + } + + // Clean-up + if (!this._userDefinedList) { + this._idMap = {}; + this._idColors = []; + + this._pickableMeshes = []; + } + this._pickingTexure!.renderList = []; + + if (!wasSuccessfull) { + this._meshRenderingCount = 0; + return; // We need to wait for the shaders to be ready + } else { + if (disposeWhenDone) { + this.dispose(); + } + resolve(pickedMesh); + } + }; + }); + } + + private async _readTexturePixelsAsync(x: number, y: number) { + if (!this._cachedScene || !this._pickingTexure?._texture) { + return false; + } + const engine = this._cachedScene.getEngine(); + await engine._readTexturePixels(this._pickingTexure._texture, 1, 1, -1, 0, this._readbuffer, true, true, x, y); + + return true; + } + + /** Release the resources */ + public dispose() { + if (this._userDefinedList) { + this.setPickingList(null); + } + this._pickableMeshes = []; + this._cachedScene = null; + + // Cleaning up + this._pickingTexure?.dispose(); + this._pickingTexure = null; + this._renderMaterial?.dispose(); + this._renderMaterial = null; + } +} diff --git a/packages/dev/core/src/Collisions/index.ts b/packages/dev/core/src/Collisions/index.ts index 1be2f23e24a..0c9b7e48444 100644 --- a/packages/dev/core/src/Collisions/index.ts +++ b/packages/dev/core/src/Collisions/index.ts @@ -3,3 +3,4 @@ export * from "./collisionCoordinator"; export * from "./pickingInfo"; export * from "./intersectionInfo"; export * from "./meshCollisionData"; +export * from "./gpuPicker"; diff --git a/packages/dev/core/src/Helpers/environmentHelper.ts b/packages/dev/core/src/Helpers/environmentHelper.ts index 0754ffc1aba..a2f746dfb67 100644 --- a/packages/dev/core/src/Helpers/environmentHelper.ts +++ b/packages/dev/core/src/Helpers/environmentHelper.ts @@ -539,6 +539,7 @@ export class EnvironmentHelper { if (!this._ground || this._ground.isDisposed()) { this._ground = CreatePlane("BackgroundPlane", { size: sceneSize.groundSize }, this._scene); this._ground.rotation.x = Math.PI / 2; // Face up by default. + this._ground.isPickable = false; this._ground.parent = this._rootMesh; this._ground.onDisposeObservable.add(() => { this._ground = null; @@ -646,6 +647,7 @@ export class EnvironmentHelper { private _setupSkybox(sceneSize: ISceneSize): void { if (!this._skybox || this._skybox.isDisposed()) { this._skybox = CreateBox("BackgroundSkybox", { size: sceneSize.skyboxSize, sideOrientation: Mesh.BACKSIDE }, this._scene); + this._skybox.isPickable = false; this._skybox.onDisposeObservable.add(() => { this._skybox = null; }); diff --git a/packages/dev/core/src/Materials/shaderMaterial.ts b/packages/dev/core/src/Materials/shaderMaterial.ts index 09f5c6b3e72..7ccfc0bc310 100644 --- a/packages/dev/core/src/Materials/shaderMaterial.ts +++ b/packages/dev/core/src/Materials/shaderMaterial.ts @@ -33,6 +33,7 @@ import { PrepareAttributesForBakedVertexAnimation, PushAttributesForInstances, } from "./materialHelper.functions"; +import { Observable } from "core/Misc/observable"; const onCreatedEffectParameters = { effect: null as unknown as Effect, subMesh: null as unknown as Nullable }; @@ -141,6 +142,11 @@ export class ShaderMaterial extends PushMaterial { private _cachedWorldViewProjectionMatrix = new Matrix(); private _multiview = false; + /** + * Gets an observable raised before binding of all entities to allow custom processing. + */ + public customBindingObservable = new Observable(); + /** * @internal */ @@ -1022,7 +1028,9 @@ export class ShaderMaterial extends PushMaterial { } } - const mustRebind = mesh && storeEffectOnSubMeshes ? this._mustRebind(scene, effect, subMesh, mesh.visibility) : scene.getCachedMaterial() !== this; + const mustRebind = true; //mesh && storeEffectOnSubMeshes ? this._mustRebind(scene, effect, subMesh, mesh.visibility) : scene.getCachedMaterial() !== this; + + this.customBindingObservable.notifyObservers(mesh); if (effect && mustRebind) { if (!useSceneUBO && this._options.uniforms.indexOf("view") !== -1) { @@ -1458,6 +1466,7 @@ export class ShaderMaterial extends PushMaterial { } this._textures = {}; + this.customBindingObservable.clear(); super.dispose(forceDisposeEffect, forceDisposeTextures, notBoundToMesh); } diff --git a/packages/tools/tests/test/visualization/ReferenceImages/gpuPicker.png b/packages/tools/tests/test/visualization/ReferenceImages/gpuPicker.png new file mode 100644 index 0000000000000000000000000000000000000000..2e1a5294c4648727814d00b08576a35e55f8a878 GIT binary patch literal 2560 zcmeH}T~HHO6oBt8StQwj5O5nQHbO#LOQI39;2;e`G(nZRT002C7ahi1ipdrNpqk9rJX$^&gEuJ?uGd}t?&T}mA0;M+ZC|(Q9{Wr2c@#+||kGq5D`^`N# z0OTPK1qb*FK%1t6X$b$BaIriMu&LFT$2!fD=7#EKjFO|c(bDv7n^T>R_eu5px!WTp zXYMWr`m6%%RnMnV`#tCB6fVs}F`yshJO`MVD#;FfvJK|Y_knN3H@`Nel2n(*uCZcyq>{rdw&S9)>k8=Y7~E2 zUn3n#$VPlcB$s-)*Cvm_C$btX!zY4U=-H5gHuJE8?tu#itXEKr80?j49>MkU2oN1l zA?DB_zs?T=Cra46^`-?CFQ*WG1vV;SunY~ssd>ba+Sftr=wzehA{cOvECGL|%5LY= zuZ&Tik%swC=?We(Y>}RO_Qd>E>|hW1G~q+4+2GUe^|99jNr$y#=!%GY%_9_NJ?fP| z-5`2}jj^us6Ot$$tK;z*NSF^vTTrMjH7^J^L9sM^2>Ft= zNt_(=p)aeP+I?_MIwF1JNmq6LTKC`JmdqvCXSr8(-KE_;Lm7wiCbA)`e9AZ>H1&CF zoT+hC$oa##xq`|2v2L$(?5@D{^BrSUs&)Zt>JHhb7&{Zu_TpL5ER~eOesgqL+QJX! zHZ3EOuzc80YIXB*YL7M!(f=VNR96PK)yeOI%_uI01&2VN{4@kBVzBtw4CIVA$!&nQ zD1!lSUy|Dye~xaNv^xBPTlg=by1f4P*u`>^b7b-R&i;_t@u4 zju#QCS