PostProcessingUtils: Add method for normal reconstruction. (#29703)

* PostProcessingUtils: Add method for normal reconstruction.

* DenoiseNode: Clean up.

* DenoiseNode: More clean up.

examples/jsm/tsl/display/DenoiseNode.js

@@ -1,5 +1,5 @@
 import { Vector2, Vector3 } from 'three';
-import { getViewPosition, convertToTexture, TempNode, nodeObject, Fn, float, NodeUpdateType, uv, uniform, Loop, luminance, vec2, vec3, vec4, uniformArray, int, dot, max, pow, abs, If, textureSize, sin, cos, mat2, PI } from 'three/tsl';
+import { getNormalFromDepth, getViewPosition, convertToTexture, TempNode, nodeObject, Fn, float, NodeUpdateType, uv, uniform, Loop, luminance, vec2, vec3, vec4, uniformArray, int, dot, max, pow, abs, If, textureSize, sin, cos, mat2, PI } from 'three/tsl';
 
 class DenoiseNode extends TempNode {
 
@@ -46,16 +46,16 @@ class DenoiseNode extends TempNode {
 
 		const sampleTexture = ( uv ) => this.textureNode.uv( uv );
 		const sampleDepth = ( uv ) => this.depthNode.uv( uv ).x;
-		const sampleNormal = ( uv ) => this.normalNode.uv( uv );
+		const sampleNormal = ( uv ) => ( this.normalNode !== null ) ? this.normalNode.uv( uv ).rgb.normalize() : getNormalFromDepth( uv, this.depthNode.value, this.cameraProjectionMatrixInverse );
 		const sampleNoise = ( uv ) => this.noiseNode.uv( uv );
 
 		const denoiseSample = Fn( ( [ center, viewNormal, viewPosition, sampleUv ] ) => {
 
-			const texel = sampleTexture( sampleUv );
-			const depth = sampleDepth( sampleUv );
-			const normal = sampleNormal( sampleUv ).rgb.normalize();
+			const texel = sampleTexture( sampleUv ).toVar();
+			const depth = sampleDepth( sampleUv ).toVar();
+			const normal = sampleNormal( sampleUv ).toVar();
 			const neighborColor = texel.rgb;
-			const viewPos = getViewPosition( sampleUv, depth, this.cameraProjectionMatrixInverse );
+			const viewPos = getViewPosition( sampleUv, depth, this.cameraProjectionMatrixInverse ).toVar();
 
 			const normalDiff = dot( viewNormal, normal ).toVar();
 			const normalSimilarity = pow( max( normalDiff, 0 ), this.normalPhi ).toVar();
@@ -71,10 +71,10 @@ class DenoiseNode extends TempNode {
 
 		const denoise = Fn( ( [ uvNode ] ) => {
 
-			const depth = sampleDepth( uvNode );
-			const viewNormal = sampleNormal( uvNode ).rgb.normalize();
+			const depth = sampleDepth( uvNode ).toVar();
+			const viewNormal = sampleNormal( uvNode ).toVar();
 
-			const texel = sampleTexture( uvNode );
+			const texel = sampleTexture( uvNode ).toVar();
 
 			If( depth.greaterThanEqual( 1.0 ).or( dot( viewNormal, viewNormal ).equal( 0.0 ) ), () => {
 
@@ -82,20 +82,20 @@ class DenoiseNode extends TempNode {
 
 			} );
 
-			const center = vec3( texel.rgb );
+			const center = vec3( texel.rgb ).toVar();
 
-			const viewPosition = getViewPosition( uvNode, depth, this.cameraProjectionMatrixInverse );
+			const viewPosition = getViewPosition( uvNode, depth, this.cameraProjectionMatrixInverse ).toVar();
 
 			const noiseResolution = textureSize( this.noiseNode, 0 );
 			let noiseUv = vec2( uvNode.x, uvNode.y.oneMinus() );
 			noiseUv = noiseUv.mul( this._resolution.div( noiseResolution ) );
-			const noiseTexel = sampleNoise( noiseUv );
+			const noiseTexel = sampleNoise( noiseUv ).toVar();
 
-			const x = sin( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) );
-			const y = cos( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) );
+			const x = sin( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) ).toVar();
+			const y = cos( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) ).toVar();
 
-			const noiseVec = vec2( x, y );
-			const rotationMatrix = mat2( noiseVec.x, noiseVec.y.negate(), noiseVec.x, noiseVec.y );
+			const noiseVec = vec2( x, y ).toVar();
+			const rotationMatrix = mat2( noiseVec.x, noiseVec.y.negate(), noiseVec.x, noiseVec.y ).toVar();
 
 			const totalWeight = float( 1.0 ).toVar();
 			const denoised = vec3( texel.rgb ).toVar();

examples/jsm/tsl/display/GTAONode.js

@@ -1,5 +1,5 @@
 import { DataTexture, RenderTarget, RepeatWrapping, Vector2, Vector3, PostProcessingUtils } from 'three';
-import { getScreenPosition, getViewPosition, QuadMesh, TempNode, nodeObject, Fn, float, NodeUpdateType, uv, uniform, Loop, vec2, vec3, vec4, int, dot, max, pow, abs, If, textureSize, sin, cos, PI, texture, passTexture, mat3, add, normalize, mul, cross, div, mix, sqrt, sub, acos, clamp, NodeMaterial } from 'three/tsl';
+import { getNormalFromDepth, getScreenPosition, getViewPosition, QuadMesh, TempNode, nodeObject, Fn, float, NodeUpdateType, uv, uniform, Loop, vec2, vec3, vec4, int, dot, max, pow, abs, If, textureSize, sin, cos, PI, texture, passTexture, mat3, add, normalize, mul, cross, div, mix, sqrt, sub, acos, clamp, NodeMaterial } from 'three/tsl';
 
 const _quadMesh = /*@__PURE__*/ new QuadMesh();
 const _size = /*@__PURE__*/ new Vector2();
@@ -91,6 +91,7 @@ class GTAONode extends TempNode {
 
 		const sampleDepth = ( uv ) => this.depthNode.uv( uv ).x;
 		const sampleNoise = ( uv ) => this.noiseNode.uv( uv );
+		const sampleNormal = ( uv ) => ( this.normalNode !== null ) ? this.normalNode.uv( uv ).rgb.normalize() : getNormalFromDepth( uv, this.depthNode.value, this.cameraProjectionMatrixInverse );
 
 		const ao = Fn( () => {
 
@@ -99,7 +100,7 @@ class GTAONode extends TempNode {
 			depth.greaterThanEqual( 1.0 ).discard();
 
 			const viewPosition = getViewPosition( uvNode, depth, this.cameraProjectionMatrixInverse ).toVar();
-			const viewNormal = this.normalNode.rgb.normalize().toVar();
+			const viewNormal = sampleNormal( uvNode ).toVar();
 
 			const radiusToUse = this.radius;
 

src/nodes/utils/PostProcessingUtils.js

@@ -1,4 +1,6 @@
-import { Fn, vec2, vec3, vec4 } from '../tsl/TSLBase.js';
+import { abs, cross, float, Fn, normalize, ivec2, sub, vec2, vec3, vec4 } from '../tsl/TSLBase.js';
+import { textureSize } from '../accessors/TextureSizeNode.js';
+import { textureLoad } from '../accessors/TextureNode.js';
 import { WebGPUCoordinateSystem } from '../../constants.js';
 
 /**
@@ -46,3 +48,42 @@ export const getScreenPosition = /*@__PURE__*/ Fn( ( [ viewPosition, projectionM
 	return vec2( sampleUv.x, sampleUv.y.oneMinus() );
 
 } );
+
+/**
+* Computes a normal vector based on depth data. Can be used as a fallback when no normal render
+* target is available or if flat surface normals are required.
+*
+* @param {vec2} uv - The texture coordinate.
+* @param {DepthTexture} depthTexture - The depth texture.
+* @param {mat4} projectionMatrixInverse - The camera's inverse projection matrix.
+* @return {vec3} The computed normal vector.
+*/
+export const getNormalFromDepth = /*@__PURE__*/ Fn( ( [ uv, depthTexture, projectionMatrixInverse ] ) => {
+
+	const size = textureSize( textureLoad( depthTexture ) );
+	const p = ivec2( uv.mul( size ) ).toVar();
+
+	const c0 = textureLoad( depthTexture, p ).toVar();
+
+	const l2 = textureLoad( depthTexture, p.sub( ivec2( 2, 0 ) ) ).toVar();
+	const l1 = textureLoad( depthTexture, p.sub( ivec2( 1, 0 ) ) ).toVar();
+	const r1 = textureLoad( depthTexture, p.add( ivec2( 1, 0 ) ) ).toVar();
+	const r2 = textureLoad( depthTexture, p.add( ivec2( 2, 0 ) ) ).toVar();
+	const b2 = textureLoad( depthTexture, p.add( ivec2( 0, 2 ) ) ).toVar();
+	const b1 = textureLoad( depthTexture, p.add( ivec2( 0, 1 ) ) ).toVar();
+	const t1 = textureLoad( depthTexture, p.sub( ivec2( 0, 1 ) ) ).toVar();
+	const t2 = textureLoad( depthTexture, p.sub( ivec2( 0, 2 ) ) ).toVar();
+
+	const dl = abs( sub( float( 2 ).mul( l1 ).sub( l2 ), c0 ) ).toVar();
+	const dr = abs( sub( float( 2 ).mul( r1 ).sub( r2 ), c0 ) ).toVar();
+	const db = abs( sub( float( 2 ).mul( b1 ).sub( b2 ), c0 ) ).toVar();
+	const dt = abs( sub( float( 2 ).mul( t1 ).sub( t2 ), c0 ) ).toVar();
+
+	const ce = getViewPosition( uv, c0, projectionMatrixInverse ).toVar();
+
+	const dpdx = dl.lessThan( dr ).select( ce.sub( getViewPosition( uv.sub( vec2( float( 1 ).div( size.x ), 0 ) ), l1, projectionMatrixInverse ) ), ce.negate().add( getViewPosition( uv.add( vec2( float( 1 ).div( size.x ), 0 ) ), r1, projectionMatrixInverse ) ) );
+	const dpdy = db.lessThan( dt ).select( ce.sub( getViewPosition( uv.add( vec2( 0, float( 1 ).div( size.y ) ) ), b1, projectionMatrixInverse ) ), ce.negate().add( getViewPosition( uv.sub( vec2( 0, float( 1 ).div( size.y ) ) ), t1, projectionMatrixInverse ) ) );
+
+	return normalize( cross( dpdx, dpdy ) );
+
+} );