TRAANode: Reduce Ghosting/Smearing (#31895)

* Solve for world position

* Make basic movement rejection a thing

* Fix the issue with edges

* Update TRAANode.js

* Adjust comment formatting

* Resolve unused vars, formatting, and debug nitpicks

* Update TRAANode.js

Clean up.

* Update TRAANode.js

- Store previous depth in history render target.
- Don't copy depth in the first frame.

---------

Co-authored-by: Michael Herzog <michael.herzog@human-interactive.org>

examples/jsm/tsl/display/TRAANode.js

@@ -1,5 +1,5 @@
-import { HalfFloatType, Vector2, RenderTarget, RendererUtils, QuadMesh, NodeMaterial, TempNode, NodeUpdateType, Matrix4 } from 'three/webgpu';
-import { add, float, If, Loop, int, Fn, min, max, clamp, nodeObject, texture, uniform, uv, vec2, vec4, luminance, convertToTexture, passTexture, velocity } from 'three/tsl';
+import { HalfFloatType, Vector2, RenderTarget, RendererUtils, QuadMesh, NodeMaterial, TempNode, NodeUpdateType, Matrix4, DepthTexture } from 'three/webgpu';
+import { add, float, If, Loop, int, Fn, min, max, clamp, nodeObject, texture, uniform, uv, vec2, vec4, luminance, convertToTexture, passTexture, velocity, getViewPosition, length, mat4 } from 'three/tsl';
 
 const _quadMesh = /*@__PURE__*/ new QuadMesh();
 const _size = /*@__PURE__*/ new Vector2();
@@ -100,13 +100,45 @@ class TRAANode extends TempNode {
 		 */
 		this._invSize = uniform( new Vector2() );
 
+		/**
+		 * A uniform node holding the camera world matrix.
+		 *
+		 * @private
+		 * @type {UniformNode<mat4>}
+		 */
+		this._cameraWorldMatrix = uniform( new Matrix4() );
+
+		/**
+		 * A uniform node holding the camera projection matrix inverse.
+		 *
+		 * @private
+		 * @type {UniformNode<mat4>}
+		 */
+		this._cameraProjectionMatrixInverse = uniform( new Matrix4() );
+
+		/**
+		 * A uniform node holding the previous frame's view matrix.
+		 *
+		 * @private
+		 * @type {UniformNode<mat4>}
+		 */
+		this._previousCameraWorldMatrix = uniform( new Matrix4() );
+
+		/**
+		 * A uniform node holding the previous frame's projection matrix inverse.
+		 *
+		 * @private
+		 * @type {UniformNode<mat4>}
+		 */
+		this._previousCameraProjectionMatrixInverse = uniform( new Matrix4() );
+
 		/**
 		 * The render target that represents the history of frame data.
 		 *
 		 * @private
 		 * @type {?RenderTarget}
 		 */
-		this._historyRenderTarget = new RenderTarget( 1, 1, { depthBuffer: false, type: HalfFloatType } );
+		this._historyRenderTarget = new RenderTarget( 1, 1, { depthBuffer: false, type: HalfFloatType, depthTexture: new DepthTexture() } );
 		this._historyRenderTarget.texture.name = 'TRAANode.history';
 
 		/**
@@ -143,6 +175,14 @@ class TRAANode extends TempNode {
 		 */
 		this._originalProjectionMatrix = new Matrix4();
 
+		/**
+		 * A texture node for the previous depth buffer.
+		 *
+		 * @private
+		 * @type {TextureNode}
+		 */
+		this._previousDepthNode = texture( new DepthTexture( 1, 1 ) );
+
 		/**
 		 * Sync the post processing stack with the TRAA node.
 		 * @private
@@ -150,6 +190,14 @@ class TRAANode extends TempNode {
 		 */
 		this._needsPostProcessingSync = false;
 
+		/**
+		 * Whether the current frame is the first one or not.
+		 *
+		 * @private
+		 * @type {boolean}
+		 */
+		this._firstFrame = true;
+
 	}
 
 	/**
@@ -246,6 +294,14 @@ class TRAANode extends TempNode {
 
 		const { renderer } = frame;
 
+		// Store previous frame matrices before updating current ones
+		this._previousCameraWorldMatrix.value.copy( this._cameraWorldMatrix.value );
+		this._previousCameraProjectionMatrixInverse.value.copy( this._cameraProjectionMatrixInverse.value );
+
+		// Update camera matrices uniforms
+		this._cameraWorldMatrix.value.copy( this.camera.matrixWorld );
+		this._cameraProjectionMatrixInverse.value.copy( this.camera.projectionMatrixInverse );
+
 		// keep the TRAA in sync with the dimensions of the beauty node
 
 		const beautyRenderTarget = ( this.beautyNode.isRTTNode ) ? this.beautyNode.renderTarget : this.beautyNode.passNode.renderTarget;
@@ -300,6 +356,21 @@ class TRAANode extends TempNode {
 
 		renderer.copyTextureToTexture( this._resolveRenderTarget.texture, this._historyRenderTarget.texture );
 
+		// Copy current depth to previous depth buffer
+
+		if ( this._firstFrame === true ) {
+
+			this._firstFrame = false;
+
+		} else {
+
+			const currentDepth = this.depthNode.value;
+			renderer.copyTextureToTexture( currentDepth, this._historyRenderTarget.depthTexture );
+			this._previousDepthNode.value = this._historyRenderTarget.depthTexture;
+
+		}
+
+
 		// restore
 
 		RendererUtils.restoreRendererState( renderer, _rendererState );
@@ -347,6 +418,7 @@ class TRAANode extends TempNode {
 			const minColor = vec4( 10000 ).toVar();
 			const maxColor = vec4( - 10000 ).toVar();
 			const closestDepth = float( 1 ).toVar();
+			const farthestDepth = float( 0 ).toVar();
 			const closestDepthPixelPosition = vec2( 0 ).toVar();
 
 			// sample a 3x3 neighborhood to create a box in color space
@@ -373,6 +445,14 @@ class TRAANode extends TempNode {
 
 					} );
 
+					// find the farthest depth in the neighborhood (used to preserve edge anti-aliasing)
+
+					If( currentDepth.greaterThan( farthestDepth ), () => {
+
+						farthestDepth.assign( currentDepth );
+
+					} );
+
 				} );
 
 			} );
@@ -388,11 +468,39 @@ class TRAANode extends TempNode {
 
 			const clampedHistoryColor = clamp( historyColor, minColor, maxColor );
 
-			// flicker reduction based on luminance weighing
+			// calculate current frame world position
+
+			const currentDepth = depthTexture.sample( uvNode ).r;
+			const currentViewPosition = getViewPosition( uvNode, currentDepth, this._cameraProjectionMatrixInverse );
+			const currentWorldPosition = this._cameraWorldMatrix.mul( vec4( currentViewPosition, 1.0 ) ).xyz;
+
+			// calculate previous frame world position from history UV and previous depth
+
+			const historyUV = uvNode.sub( offset );
+			const previousDepth = this._previousDepthNode.sample( historyUV ).r;
+			const previousViewPosition = getViewPosition( historyUV, previousDepth, this._previousCameraProjectionMatrixInverse );
+			const previousWorldPosition = this._previousCameraWorldMatrix.mul( vec4( previousViewPosition, 1.0 ) ).xyz;
+
+			// calculate difference in world positions
+
+			const worldPositionDifference = length( currentWorldPosition.sub( previousWorldPosition ) ).toVar();
+			worldPositionDifference.assign( min( max( worldPositionDifference.sub( 1.0 ), 0.0 ), 1.0 ) );
 
 			const currentWeight = float( 0.05 ).toVar();
 			const historyWeight = currentWeight.oneMinus().toVar();
 
+			// zero out history weight if world positions are different (indicating motion) except on edges
+
+			const rejectPixel = worldPositionDifference.greaterThan( 0.01 ).and( farthestDepth.sub( closestDepth ).lessThan( 0.0001 ) );
+			If( rejectPixel, () => {
+
+				currentWeight.assign( 1.0 );
+				historyWeight.assign( 0.0 );
+
+			} );
+
+			// flicker reduction based on luminance weighing
+
 			const compressedCurrent = currentColor.mul( float( 1 ).div( ( max( currentColor.r, currentColor.g, currentColor.b ).add( 1.0 ) ) ) );
 			const compressedHistory = clampedHistoryColor.mul( float( 1 ).div( ( max( clampedHistoryColor.r, clampedHistoryColor.g, clampedHistoryColor.b ).add( 1.0 ) ) ) );
 
@@ -402,7 +510,9 @@ class TRAANode extends TempNode {
 			currentWeight.mulAssign( float( 1.0 ).div( luminanceCurrent.add( 1 ) ) );
 			historyWeight.mulAssign( float( 1.0 ).div( luminanceHistory.add( 1 ) ) );
 
-			return add( currentColor.mul( currentWeight ), clampedHistoryColor.mul( historyWeight ) ).div( max( currentWeight.add( historyWeight ), 0.00001 ) );
+			const smoothedOutput = add( currentColor.mul( currentWeight ), clampedHistoryColor.mul( historyWeight ) ).div( max( currentWeight.add( historyWeight ), 0.00001 ) ).toVar();
+
+			return smoothedOutput;
 
 		} );