Skip to content

Commit

Permalink
PostProcessingUtils: Add method for normal reconstruction. (#29703)
Browse files Browse the repository at this point in the history 
* PostProcessingUtils: Add method for normal reconstruction.

* DenoiseNode: Clean up.

* DenoiseNode: More clean up.
  • Loading branch information
@Mugen87
Mugen87 authored Oct 20, 2024
1 parent cde4019 commit 4d08201
Show file tree
Hide file tree
Showing 3 changed files with 61 additions and 19 deletions.
32 changes: 16 additions & 16 deletions examples/jsm/tsl/display/DenoiseNode.js
View file
Original file line number Diff line number Diff line change
@@ -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 {

Expand Down Expand Up @@ -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();
Expand All @@ -71,31 +71,31 @@ 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 ) ), () => {

return texel;

} );

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();
Expand Down
5 changes: 3 additions & 2 deletions examples/jsm/tsl/display/GTAONode.js
View file
Original file line number Diff line number Diff line change
@@ -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();
Expand Down Expand Up @@ -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( () => {

Expand All @@ -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;

Expand Down
43 changes: 42 additions & 1 deletion src/nodes/utils/PostProcessingUtils.js
View file
Original file line number Diff line number Diff line change
@@ -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';

/**
Expand Down Expand Up @@ -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 ) );

} );

0 comments on commit 4d08201

Please sign in to comment.