Skip to content

Commit 4d08201

Browse files
Mugen87Mugen87
authored
PostProcessingUtils: Add method for normal reconstruction. (#29703)
* PostProcessingUtils: Add method for normal reconstruction. * DenoiseNode: Clean up. * DenoiseNode: More clean up.
1 parent cde4019 commit 4d08201

File tree

3 files changed

+61
-19
lines changed

3 files changed

+61
-19
lines changed

‎examples/jsm/tsl/display/DenoiseNode.js‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { Vector2, Vector3 } from 'three';
2-
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';
2+
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';
33

44
class DenoiseNode extends TempNode {
55

@@ -46,16 +46,16 @@ class DenoiseNode extends TempNode {
4646

4747
const sampleTexture = ( uv ) => this.textureNode.uv( uv );
4848
const sampleDepth = ( uv ) => this.depthNode.uv( uv ).x;
49-
const sampleNormal = ( uv ) => this.normalNode.uv( uv );
49+
const sampleNormal = ( uv ) => ( this.normalNode !== null ) ? this.normalNode.uv( uv ).rgb.normalize() : getNormalFromDepth( uv, this.depthNode.value, this.cameraProjectionMatrixInverse );
5050
const sampleNoise = ( uv ) => this.noiseNode.uv( uv );
5151

5252
const denoiseSample = Fn( ( [ center, viewNormal, viewPosition, sampleUv ] ) => {
5353

54-
const texel = sampleTexture( sampleUv );
55-
const depth = sampleDepth( sampleUv );
56-
const normal = sampleNormal( sampleUv ).rgb.normalize();
54+
const texel = sampleTexture( sampleUv ).toVar();
55+
const depth = sampleDepth( sampleUv ).toVar();
56+
const normal = sampleNormal( sampleUv ).toVar();
5757
const neighborColor = texel.rgb;
58-
const viewPos = getViewPosition( sampleUv, depth, this.cameraProjectionMatrixInverse );
58+
const viewPos = getViewPosition( sampleUv, depth, this.cameraProjectionMatrixInverse ).toVar();
5959

6060
const normalDiff = dot( viewNormal, normal ).toVar();
6161
const normalSimilarity = pow( max( normalDiff, 0 ), this.normalPhi ).toVar();
@@ -71,31 +71,31 @@ class DenoiseNode extends TempNode {
7171

7272
const denoise = Fn( ( [ uvNode ] ) => {
7373

74-
const depth = sampleDepth( uvNode );
75-
const viewNormal = sampleNormal( uvNode ).rgb.normalize();
74+
const depth = sampleDepth( uvNode ).toVar();
75+
const viewNormal = sampleNormal( uvNode ).toVar();
7676

77-
const texel = sampleTexture( uvNode );
77+
const texel = sampleTexture( uvNode ).toVar();
7878

7979
If( depth.greaterThanEqual( 1.0 ).or( dot( viewNormal, viewNormal ).equal( 0.0 ) ), () => {
8080

8181
return texel;
8282

8383
} );
8484

85-
const center = vec3( texel.rgb );
85+
const center = vec3( texel.rgb ).toVar();
8686

87-
const viewPosition = getViewPosition( uvNode, depth, this.cameraProjectionMatrixInverse );
87+
const viewPosition = getViewPosition( uvNode, depth, this.cameraProjectionMatrixInverse ).toVar();
8888

8989
const noiseResolution = textureSize( this.noiseNode, 0 );
9090
let noiseUv = vec2( uvNode.x, uvNode.y.oneMinus() );
9191
noiseUv = noiseUv.mul( this._resolution.div( noiseResolution ) );
92-
const noiseTexel = sampleNoise( noiseUv );
92+
const noiseTexel = sampleNoise( noiseUv ).toVar();
9393

94-
const x = sin( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) );
95-
const y = cos( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) );
94+
const x = sin( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) ).toVar();
95+
const y = cos( noiseTexel.element( this.index.mod( 4 ).mul( 2 ).mul( PI ) ) ).toVar();
9696

97-
const noiseVec = vec2( x, y );
98-
const rotationMatrix = mat2( noiseVec.x, noiseVec.y.negate(), noiseVec.x, noiseVec.y );
97+
const noiseVec = vec2( x, y ).toVar();
98+
const rotationMatrix = mat2( noiseVec.x, noiseVec.y.negate(), noiseVec.x, noiseVec.y ).toVar();
9999

100100
const totalWeight = float( 1.0 ).toVar();
101101
const denoised = vec3( texel.rgb ).toVar();

‎examples/jsm/tsl/display/GTAONode.js‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { DataTexture, RenderTarget, RepeatWrapping, Vector2, Vector3, PostProcessingUtils } from 'three';
2-
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';
2+
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';
33

44
const _quadMesh = /*@__PURE__*/ new QuadMesh();
55
const _size = /*@__PURE__*/ new Vector2();
@@ -91,6 +91,7 @@ class GTAONode extends TempNode {
9191

9292
const sampleDepth = ( uv ) => this.depthNode.uv( uv ).x;
9393
const sampleNoise = ( uv ) => this.noiseNode.uv( uv );
94+
const sampleNormal = ( uv ) => ( this.normalNode !== null ) ? this.normalNode.uv( uv ).rgb.normalize() : getNormalFromDepth( uv, this.depthNode.value, this.cameraProjectionMatrixInverse );
9495

9596
const ao = Fn( () => {
9697

@@ -99,7 +100,7 @@ class GTAONode extends TempNode {
99100
depth.greaterThanEqual( 1.0 ).discard();
100101

101102
const viewPosition = getViewPosition( uvNode, depth, this.cameraProjectionMatrixInverse ).toVar();
102-
const viewNormal = this.normalNode.rgb.normalize().toVar();
103+
const viewNormal = sampleNormal( uvNode ).toVar();
103104

104105
const radiusToUse = this.radius;
105106

‎src/nodes/utils/PostProcessingUtils.js‎

Lines changed: 42 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
import { Fn, vec2, vec3, vec4 } from '../tsl/TSLBase.js';
1+
import { abs, cross, float, Fn, normalize, ivec2, sub, vec2, vec3, vec4 } from '../tsl/TSLBase.js';
2+
import { textureSize } from '../accessors/TextureSizeNode.js';
3+
import { textureLoad } from '../accessors/TextureNode.js';
24
import { WebGPUCoordinateSystem } from '../../constants.js';
35

46
/**
@@ -46,3 +48,42 @@ export const getScreenPosition = /*@__PURE__*/ Fn( ( [ viewPosition, projectionM
4648
return vec2( sampleUv.x, sampleUv.y.oneMinus() );
4749

4850
} );
51+
52+
/**
53+
* Computes a normal vector based on depth data. Can be used as a fallback when no normal render
54+
* target is available or if flat surface normals are required.
55+
*
56+
* @param {vec2} uv - The texture coordinate.
57+
* @param {DepthTexture} depthTexture - The depth texture.
58+
* @param {mat4} projectionMatrixInverse - The camera's inverse projection matrix.
59+
* @return {vec3} The computed normal vector.
60+
*/
61+
export const getNormalFromDepth = /*@__PURE__*/ Fn( ( [ uv, depthTexture, projectionMatrixInverse ] ) => {
62+
63+
const size = textureSize( textureLoad( depthTexture ) );
64+
const p = ivec2( uv.mul( size ) ).toVar();
65+
66+
const c0 = textureLoad( depthTexture, p ).toVar();
67+
68+
const l2 = textureLoad( depthTexture, p.sub( ivec2( 2, 0 ) ) ).toVar();
69+
const l1 = textureLoad( depthTexture, p.sub( ivec2( 1, 0 ) ) ).toVar();
70+
const r1 = textureLoad( depthTexture, p.add( ivec2( 1, 0 ) ) ).toVar();
71+
const r2 = textureLoad( depthTexture, p.add( ivec2( 2, 0 ) ) ).toVar();
72+
const b2 = textureLoad( depthTexture, p.add( ivec2( 0, 2 ) ) ).toVar();
73+
const b1 = textureLoad( depthTexture, p.add( ivec2( 0, 1 ) ) ).toVar();
74+
const t1 = textureLoad( depthTexture, p.sub( ivec2( 0, 1 ) ) ).toVar();
75+
const t2 = textureLoad( depthTexture, p.sub( ivec2( 0, 2 ) ) ).toVar();
76+
77+
const dl = abs( sub( float( 2 ).mul( l1 ).sub( l2 ), c0 ) ).toVar();
78+
const dr = abs( sub( float( 2 ).mul( r1 ).sub( r2 ), c0 ) ).toVar();
79+
const db = abs( sub( float( 2 ).mul( b1 ).sub( b2 ), c0 ) ).toVar();
80+
const dt = abs( sub( float( 2 ).mul( t1 ).sub( t2 ), c0 ) ).toVar();
81+
82+
const ce = getViewPosition( uv, c0, projectionMatrixInverse ).toVar();
83+
84+
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 ) ) );
85+
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 ) ) );
86+
87+
return normalize( cross( dpdx, dpdy ) );
88+
89+
} );

0 commit comments

Comments
 (0)