913 lines
29 KiB
JavaScript
913 lines
29 KiB
JavaScript
import { abs, atan, bool, convertToTexture, cos, cross, Discard, dot, EPSILON, exp, float, Fn, getScreenPosition, getViewPosition, If, int, log, Loop, luminance, mat2, max, mix, nodeObject, NodeUpdateType, normalize, passTexture, PI, property, reflect, sin, smoothstep, sqrt, tan, texture, uniform, unpackRGBToNormal, uv, vec2, vec3, vec4 } from 'three/tsl';
|
||
import { HalfFloatType, MathUtils, Matrix4, NodeMaterial, QuadMesh, RendererUtils, RenderTarget, TempNode, Vector2 } from 'three/webgpu';
|
||
import { bindAnalyticNoise } from '../utils/RNoise.js';
|
||
import { ENV_RAY_LENGTH_THRESHOLD } from '../utils/SpecularHelpers.js';
|
||
|
||
const _quadMesh = /*@__PURE__*/ new QuadMesh();
|
||
const _size = /*@__PURE__*/ new Vector2();
|
||
|
||
let _rendererState;
|
||
|
||
const KERNEL_SAMPLES = 8;
|
||
const NOISE_ROTATION_SEED = 83;
|
||
const WORLD_RADIUS_SCALE = 0.1;
|
||
|
||
const AO_EDGE_STOPPING_BIAS = 0.05;
|
||
const AGGRESSIVITY_RADIUS_MIN = 0.001;
|
||
const DIFFUSE_CHROMA_WEIGHT = 2.0;
|
||
// Neighborhood luma coefficient-of-variation thresholds for gating the temporal inverse-luminance
|
||
// (firefly) suppression: below MIN the region is treated as flicker-free, above MAX as noisy.
|
||
const FLICKER_COV_GATE_MIN = 0.1;
|
||
const FLICKER_COV_GATE_MAX = 2;
|
||
|
||
/**
|
||
* Golden-angle Vogel disk offset.
|
||
*
|
||
* @tsl
|
||
*/
|
||
const vogelDisk = Fn( ( [ i, radius ] ) => {
|
||
|
||
const sampleCount = 8;
|
||
const theta = i.add( 0.5 ).mul( 2.399827721492203 );
|
||
const r = radius.mul( sqrt( i.add( 0.5 ).div( sampleCount ) ) );
|
||
return vec2( cos( theta ), sin( theta ) ).mul( r );
|
||
|
||
} ).setLayout( {
|
||
name: 'vogelDisk',
|
||
type: 'vec2',
|
||
inputs: [
|
||
{ name: 'i', type: 'float' },
|
||
{ name: 'radius', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* Chromatic color-similarity distance between two linear base colors (albedo).
|
||
*
|
||
* @tsl
|
||
*/
|
||
const diffuseColorDistance = Fn( ( [ a, b, compressLuma ] ) => {
|
||
|
||
const toYCoCg = ( c ) => vec3(
|
||
dot( c, vec3( 0.25, 0.5, 0.25 ) ),
|
||
c.r.sub( c.b ),
|
||
c.g.sub( c.r.add( c.b ).mul( 0.5 ) )
|
||
);
|
||
|
||
const ya = toYCoCg( a );
|
||
const yb = toYCoCg( b );
|
||
|
||
// `compressLuma` (0/1) range-compresses the luma term with log(1+L) so a fixed lumaPhi gives
|
||
// scale-invariant differences across the HDR range. 0 leaves luma linear (used for LDR albedo).
|
||
const compress = ( L ) => mix( L, log( L.add( 1 ) ), compressLuma );
|
||
|
||
const dLuma = abs( compress( ya.x ).sub( compress( yb.x ) ) );
|
||
const dChroma = vec2( ya.y.sub( yb.y ), ya.z.sub( yb.z ) ).length();
|
||
|
||
return dLuma.add( dChroma.mul( DIFFUSE_CHROMA_WEIGHT ) );
|
||
|
||
} ).setLayout( {
|
||
name: 'diffuseColorDistance',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'a', type: 'vec3' },
|
||
{ name: 'b', type: 'vec3' },
|
||
{ name: 'compressLuma', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
const _temporalWeight = Fn( ( [ x, strength ] ) => float( 1 ).div( x.pow( strength ) ) ).setLayout( {
|
||
name: 'temporalWeight',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'x', type: 'float' },
|
||
{ name: 'strength', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* Temporal accumulation variance factor in `[0, 1]`. Higher values mean more history confidence.
|
||
*
|
||
* @tsl
|
||
*/
|
||
const getTemporalVarianceFactor = Fn( ( [ frameNum, strength ] ) => {
|
||
|
||
return _temporalWeight( frameNum, strength ).max( 0.05 );
|
||
|
||
} ).setLayout( {
|
||
name: 'getTemporalVarianceFactor',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'frameNum', type: 'float' },
|
||
{ name: 'strength', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* World-space frustum height at `viewZ`. Algorithm originally from REBLUR (NRD).
|
||
* `tanHalfFovY` is `tan( verticalFov / 2 )`, hoisted by the caller since it is loop-invariant.
|
||
*
|
||
* @tsl
|
||
*/
|
||
const computeFrustumSize = Fn( ( [ viewZ, tanHalfFovY ] ) => {
|
||
|
||
return float( 2 ).mul( viewZ ).mul( tanHalfFovY );
|
||
|
||
} ).setLayout( {
|
||
name: 'computeFrustumSize',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'viewZ', type: 'float' },
|
||
{ name: 'tanHalfFovY', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* Maps world-space SSR ray length to `[0, 1]`. Environment rays (`worldRayLength == 0`) map to `1`.
|
||
* Algorithm originally from REBLUR (NRD).
|
||
*
|
||
* @tsl
|
||
*/
|
||
const computeHitDistFactor = Fn( ( [ worldRayLength, viewZ, tanHalfFovY ] ) => {
|
||
|
||
const frustumSize = computeFrustumSize( viewZ, tanHalfFovY );
|
||
const factor = worldRayLength.div( frustumSize.max( 1e-6 ) ).clamp( 0, 1 );
|
||
|
||
return factor;
|
||
|
||
} ).setLayout( {
|
||
name: 'computeHitDistFactor',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'worldRayLength', type: 'float' },
|
||
{ name: 'viewZ', type: 'float' },
|
||
{ name: 'tanHalfFovY', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* Maps an AO factor for edge-stopping comparisons.
|
||
*
|
||
* @tsl
|
||
*/
|
||
const mapAo = Fn( ( [ aoVal ] ) => aoVal.pow( 0.1 ) );
|
||
|
||
/**
|
||
* Specular dominant direction — smooth surfaces lean toward reflection, rough toward normal.
|
||
*
|
||
* @tsl
|
||
*/
|
||
const getSpecularDominantDirection = Fn( ( [ N, V, roughness ] ) => {
|
||
|
||
return normalize( mix( N, reflect( V.negate(), N ), roughness.oneMinus() ) );
|
||
|
||
} ).setLayout( {
|
||
name: 'getSpecularDominantDirection',
|
||
type: 'vec3',
|
||
inputs: [
|
||
{ name: 'N', type: 'vec3' },
|
||
{ name: 'V', type: 'vec3' },
|
||
{ name: 'roughness', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* GGX inverse-CDF: half-angle tangent enclosing `percent` of the specular lobe volume.
|
||
* `roughness` is perceptual (alpha = roughness²).
|
||
*
|
||
* @tsl
|
||
*/
|
||
const specularLobeTanHalfAngle = Fn( ( [ roughness, percent ] ) => {
|
||
|
||
const alpha = roughness.mul( roughness );
|
||
return alpha.mul( sqrt( percent.div( float( 1 ).sub( percent ).max( 1e-6 ) ) ) );
|
||
|
||
} ).setLayout( {
|
||
name: 'specularLobeTanHalfAngle',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'roughness', type: 'float' },
|
||
{ name: 'percent', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
const EXP_WEIGHT_SCALE = 4;
|
||
const NORMAL_ENCODING_ERROR = 1.5 / 255;
|
||
|
||
/**
|
||
* Loop-invariant part of the adaptive normal edge-stopping weight: the Gaussian falloff
|
||
* constant `2·EXP_WEIGHT_SCALE / lobeHalfAngle²`. `roughness`/`aggressivity`/`invNormalPhi`
|
||
* are constant across the kernel, so this is hoisted out of the tap loop and evaluated once
|
||
* per pixel. Lobe half-angle from REBLUR (NRD).
|
||
*
|
||
* @tsl
|
||
*/
|
||
const lobeNormalFalloff = Fn( ( [ roughness, aggressivity, invNormalPhi ] ) => {
|
||
|
||
const percent = mix( invNormalPhi.pow2(), float( 0 ), aggressivity.sqrt() ).clamp( 0.1, 0.99 );
|
||
const tanHalfAngle = specularLobeTanHalfAngle( roughness, percent );
|
||
const lobeHalfAngle = max( atan( tanHalfAngle ), float( NORMAL_ENCODING_ERROR ) );
|
||
|
||
const invHalfAngle = float( 1 ).div( lobeHalfAngle );
|
||
return invHalfAngle.mul( invHalfAngle ).mul( 2 * EXP_WEIGHT_SCALE );
|
||
|
||
} ).setLayout( {
|
||
name: 'lobeNormalFalloff',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'roughness', type: 'float' },
|
||
{ name: 'aggressivity', type: 'float' },
|
||
{ name: 'invNormalPhi', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* Adaptive lobe normal edge-stopping weight
|
||
*
|
||
* Evaluated entirely in cosine space: with `angle² ≈ 2(1 − cosθ)`, the original
|
||
* `exp( −SCALE·angle/halfAngle )` becomes a Gaussian `exp( falloff·(cosθ − 1) )`, so a
|
||
* single `exp` replaces the per-tap `acos`. Matches the original at the half-angle for
|
||
* narrow lobes and is slightly more permissive for wide (diffuse) ones.
|
||
*
|
||
* @tsl
|
||
*/
|
||
const lobeNormalWeight = Fn( ( [ viewNormal, nNormalV, lobeFalloff ] ) => {
|
||
|
||
const cosA = dot( viewNormal, nNormalV );
|
||
|
||
return exp( cosA.sub( 1 ).mul( lobeFalloff ) );
|
||
|
||
} ).setLayout( {
|
||
name: 'lobeNormalWeight',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'viewNormal', type: 'vec3' },
|
||
{ name: 'nNormalV', type: 'vec3' },
|
||
{ name: 'lobeFalloff', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* View-space plane distance between two surface points (edge-stopping geometry term).
|
||
*
|
||
* @tsl
|
||
*/
|
||
const planeDistance = Fn( ( [ position, nPosition, normal ] ) => {
|
||
|
||
return abs( dot( position.sub( nPosition ), normal ) );
|
||
|
||
} ).setLayout( {
|
||
name: 'planeDistance',
|
||
type: 'float',
|
||
inputs: [
|
||
{ name: 'position', type: 'vec3' },
|
||
{ name: 'nPosition', type: 'vec3' },
|
||
{ name: 'normal', type: 'vec3' },
|
||
]
|
||
} );
|
||
|
||
/**
|
||
* Inverse-luminance temporal blend with optional adaptive trust (Karis-style).
|
||
*
|
||
* @tsl
|
||
*/
|
||
const karisTemporalBlend = Fn( ( [ denoisedRgb, denoisedRaw, a, flickerSuppression, adaptiveTrust, nbhdMeanLuma, nbhdStddevLuma ] ) => {
|
||
|
||
const localCoV = nbhdStddevLuma.div( nbhdMeanLuma.max( 1e-4 ) );
|
||
const trustSuppress = localCoV.mul( adaptiveTrust ).mul( a.oneMinus() ).clamp( 0, 0.9 );
|
||
const aTrust = a.mul( trustSuppress.oneMinus() );
|
||
|
||
// In flicker-free neighborhoods, back off the inverse-luminance weighting so valid bright highlights
|
||
// keep their energy. Scaled by adaptiveTrust so the default (0) path is unchanged.
|
||
const noisy = smoothstep( FLICKER_COV_GATE_MIN, FLICKER_COV_GATE_MAX, localCoV );
|
||
const effFlicker = flickerSuppression.mul( mix( adaptiveTrust.oneMinus(), float( 1 ), noisy ) );
|
||
|
||
const wHist = float( 1 ).sub( aTrust ).div( luminance( denoisedRgb ).mul( effFlicker ).mul( 10 ).add( 1 ) );
|
||
const wRaw = aTrust.div( luminance( denoisedRaw ).mul( effFlicker ).mul( 10 ).add( 1 ) );
|
||
return denoisedRgb.mul( wHist ).add( denoisedRaw.mul( wRaw ) ).div( wHist.add( wRaw ).max( EPSILON ) );
|
||
|
||
} ).setLayout( {
|
||
name: 'karisTemporalBlend',
|
||
type: 'vec3',
|
||
inputs: [
|
||
{ name: 'denoisedRgb', type: 'vec3' },
|
||
{ name: 'denoisedRaw', type: 'vec3' },
|
||
{ name: 'a', type: 'float' },
|
||
{ name: 'flickerSuppression', type: 'float' },
|
||
{ name: 'adaptiveTrust', type: 'float' },
|
||
{ name: 'nbhdMeanLuma', type: 'float' },
|
||
{ name: 'nbhdStddevLuma', type: 'float' }
|
||
]
|
||
} );
|
||
|
||
const toTextureNode = ( value ) => {
|
||
|
||
if ( value === null ) return null;
|
||
|
||
if ( value.isTexture === true ) return texture( value );
|
||
|
||
return convertToTexture( value.getTextureNode?.() ?? value );
|
||
|
||
};
|
||
|
||
/**
|
||
* @typedef {'diffuse'|'specular'} DenoiseMode
|
||
*/
|
||
|
||
/**
|
||
* @typedef {'raylength'|'ao'|'none'} DenoiseAlphaSource
|
||
*/
|
||
|
||
/**
|
||
* @typedef {Object} RecurrentDenoiseNodeOptions
|
||
* @property {?Node<float>} [depth=null] - Scene depth buffer for view-space edge stopping.
|
||
* @property {?Node<vec3>} [normal=null] - View-space normals for geometric edge stopping.
|
||
* @property {?Node<vec4>} [metalRoughness=null] - Roughness/metalness G-buffer for specular edge stopping.
|
||
* @property {?Node<vec4>} [diffuse=null] - Scene base color (albedo) G-buffer for chromatic edge stopping.
|
||
* @property {?Node<vec4>} [raw=null] - Unfiltered input (e.g. raw SSR/SSGI) for secondary sampling and temporal blend.
|
||
* @property {DenoiseMode} [mode='diffuse'] - Denoising kernel type.
|
||
* @property {boolean} [accumulate=true] - When `true`, temporally blend the spatially-denoised result
|
||
* (Karis-style) and write frame weight to alpha for feedback loops. When `false`, only spatial filtering is applied.
|
||
*/
|
||
|
||
/**
|
||
* Post processing node for denoising temporally-accumulated screen-space effects
|
||
* such as SSGI (ambient occlusion / indirect diffuse) and SSR (specular reflections).
|
||
*
|
||
* The denoising kernel is selected at construction time via `mode`:
|
||
* `'diffuse'` (SSGI) or `'specular'` (SSR). The kernel uses a fixed 8-sample Vogel disk.
|
||
*
|
||
* @augments TempNode
|
||
* @three_import import { recurrentDenoise } from 'three/addons/tsl/display/RecurrentDenoiseNode.js';
|
||
*/
|
||
class RecurrentDenoiseNode extends TempNode {
|
||
|
||
static get type() {
|
||
|
||
return 'RecurrentDenoiseNode';
|
||
|
||
}
|
||
|
||
/**
|
||
* @param {TextureNode} inputTexture - Temporally filtered input to denoise (e.g. TRAA output).
|
||
* @param {Camera} camera
|
||
* @param {RecurrentDenoiseNodeOptions} [options={}]
|
||
*/
|
||
constructor( inputTexture, camera, options = {} ) {
|
||
|
||
super( 'vec4' );
|
||
|
||
const {
|
||
depth = null,
|
||
normal = null,
|
||
metalRoughness = null,
|
||
diffuse = null,
|
||
raw = null,
|
||
mode = 'diffuse',
|
||
accumulate = true,
|
||
} = options;
|
||
|
||
this.isRecurrentDenoiseNode = true;
|
||
this.camera = camera;
|
||
|
||
/**
|
||
* Denoising kernel type.
|
||
*
|
||
* @type {DenoiseMode}
|
||
*/
|
||
this.mode = mode;
|
||
|
||
/**
|
||
* When `true`, apply temporal blending after spatial denoising. When `false`, output spatially
|
||
* filtered colour only (alpha is passed through from the input temporal pass).
|
||
*
|
||
* @type {boolean}
|
||
*/
|
||
this.accumulate = accumulate;
|
||
|
||
this.textureNode = inputTexture;
|
||
this.depthNode = depth !== null ? nodeObject( depth ) : null;
|
||
this.normalNode = normal !== null ? nodeObject( normal ) : null;
|
||
this.rawNode = toTextureNode( raw );
|
||
this.roughnessMetalnessNode = metalRoughness !== null ? nodeObject( metalRoughness ) : null;
|
||
this.diffuseNode = diffuse !== null ? nodeObject( diffuse ) : null;
|
||
|
||
this._noiseIndex = uniform( 0 );
|
||
|
||
this.lumaPhi = uniform( 5 );
|
||
this.depthPhi = uniform( 5 );
|
||
this.normalPhi = uniform( 5 );
|
||
this.radius = uniform( 5 );
|
||
this.alphaPhi = uniform( 1 );
|
||
this.roughnessPhi = uniform( 100 );
|
||
this.diffusePhi = uniform( 100 );
|
||
this.adapt = uniform( 0.5 );
|
||
this.smoothDisocclusions = uniform( true, 'bool' );
|
||
this.strength = uniform( 0.25 );
|
||
this.maxFrames = uniform( 32 );
|
||
|
||
/**
|
||
* Which channel of the raw texture drives alpha-based edge stopping.
|
||
* `'raylength'` — alpha encodes SSR ray length; `'ao'` — alpha encodes AO factor;
|
||
* `'none'` — skip alpha-based edge stopping.
|
||
*
|
||
* @type {DenoiseAlphaSource}
|
||
* @default 'raylength'
|
||
*/
|
||
this.alphaSource = 'raylength';
|
||
|
||
this.flickerSuppression = uniform( 1 );
|
||
this.adaptiveTrust = uniform( 0 );
|
||
|
||
this.updateBeforeType = NodeUpdateType.FRAME;
|
||
|
||
this._resolution = uniform( new Vector2() );
|
||
this._fovY = uniform( MathUtils.degToRad( camera.fov ) );
|
||
this._cameraProjectionMatrixInverse = uniform( new Matrix4().copy( camera.projectionMatrixInverse ) );
|
||
this._cameraProjectionMatrix = uniform( new Matrix4().copy( camera.projectionMatrix ) );
|
||
this._viewMatrix = uniform( new Matrix4().copy( camera.matrixWorldInverse ) );
|
||
|
||
this._renderTarget = new RenderTarget( 1, 1, { depthBuffer: false, type: HalfFloatType } );
|
||
this._renderTarget.texture.name = 'RecurrentDenoiseNode.output';
|
||
|
||
this._material = new NodeMaterial();
|
||
this._material.name = 'RecurrentDenoise';
|
||
|
||
this._textureNode = passTexture( this, this._renderTarget.texture );
|
||
|
||
}
|
||
|
||
setSize( width, height ) {
|
||
|
||
if ( width === null || height === null ) return;
|
||
|
||
this._renderTarget.setSize( width, height );
|
||
this._resolution.value.set( width, height );
|
||
|
||
}
|
||
|
||
getTextureNode() {
|
||
|
||
return this._textureNode;
|
||
|
||
}
|
||
|
||
/**
|
||
* Returns the internal output render target (e.g. for temporal reprojection/SSGI temporal feedback loops).
|
||
*
|
||
* @returns {RenderTarget}
|
||
*/
|
||
getRenderTarget() {
|
||
|
||
return this._renderTarget;
|
||
|
||
}
|
||
|
||
updateBefore( frame ) {
|
||
|
||
const { renderer } = frame;
|
||
|
||
const drawingBufferSize = renderer.getDrawingBufferSize( _size );
|
||
const width = drawingBufferSize.width;
|
||
const height = drawingBufferSize.height;
|
||
|
||
const needsRestart = this._renderTarget.width !== width || this._renderTarget.height !== height;
|
||
this.setSize( width, height );
|
||
|
||
this._cameraProjectionMatrix.value.copy( this.camera.projectionMatrix );
|
||
this._cameraProjectionMatrixInverse.value.copy( this.camera.projectionMatrixInverse );
|
||
this._viewMatrix.value.copy( this.camera.matrixWorldInverse );
|
||
|
||
if ( this.camera.isPerspectiveCamera ) {
|
||
|
||
this._fovY.value = MathUtils.degToRad( this.camera.fov );
|
||
|
||
}
|
||
|
||
if ( frame.frameId !== undefined ) this._noiseIndex.value = frame.frameId;
|
||
|
||
// Denoise renders via an internal _quadMesh, not through the RenderPipeline output graph.
|
||
// Upstream passes (e.g. TemporalReprojectNode) referenced by a PassTextureNode input are
|
||
// otherwise never scheduled, their updateBefore() would not run and this pass would sample
|
||
// a stale/empty render target.
|
||
if ( this.textureNode.isPassTextureNode === true ) frame.updateBeforeNode( this.textureNode.passNode );
|
||
|
||
_rendererState = RendererUtils.resetRendererState( renderer, _rendererState );
|
||
|
||
if ( needsRestart === true ) {
|
||
|
||
renderer.initRenderTarget( this._renderTarget );
|
||
renderer.setRenderTarget( this._renderTarget );
|
||
renderer.clear();
|
||
renderer.setRenderTarget( null );
|
||
|
||
}
|
||
|
||
renderer.setRenderTarget( this._renderTarget );
|
||
_quadMesh.material = this._material;
|
||
_quadMesh.name = 'RecurrentDenoise';
|
||
_quadMesh.render( renderer );
|
||
renderer.setRenderTarget( null );
|
||
|
||
RendererUtils.restoreRendererState( renderer, _rendererState );
|
||
|
||
}
|
||
|
||
setup( builder ) {
|
||
|
||
const sampleAnalyticNoise = bindAnalyticNoise( this._resolution, NOISE_ROTATION_SEED );
|
||
|
||
const noiseRotationMatrix = Fn( ( [ r ] ) => {
|
||
|
||
const angle = r.mul( 2 ).mul( PI );
|
||
return mat2( cos( angle ), sin( angle ).negate(), sin( angle ), cos( angle ) );
|
||
|
||
} );
|
||
|
||
const sampleTexture = ( uvCoord ) => texture( this.textureNode, uvCoord ).max( 0 );
|
||
const sampleRaw = ( uvCoord ) => this.rawNode?.sample( uvCoord )?.max( 0 ) ?? vec3( 0 ).max( 0 );
|
||
const sampleDepth = ( uvCoord ) => this.depthNode?.sample( uvCoord )?.x ?? float( 0.5 );
|
||
const sampleNormal = ( uvCoord ) => unpackRGBToNormal( this.normalNode?.sample( uvCoord )?.rgb ?? vec3( 0, 0, 1 ) );
|
||
const sampleRoughnessMetalness = ( uvCoord ) => this.roughnessMetalnessNode?.sample( uvCoord )?.rg ?? vec2( 0, 1 );
|
||
const sampleDiffuse = ( uvCoord ) => this.diffuseNode?.sample( uvCoord )?.rgb ?? vec3( 0 );
|
||
|
||
// Neighborhood luma moments for the adaptive-trust (firefly) gating of the temporal blend.
|
||
const getNeighborhoodStats = Fn( ( [ uvCoord, centerSample ] ) => {
|
||
|
||
const rlSum = float( 0 ).toVar();
|
||
const rlSumW = float( 0 ).toVar();
|
||
const meanLuma = float( 0 ).toVar();
|
||
const m2Luma = float( 0 ).toVar();
|
||
const lumaCount = float( 0 ).toVar();
|
||
const hasEnvRay = bool( false ).toVar();
|
||
|
||
// 4-tap cross (pre-sampled center + 4 axis neighbors) instead of a full 3×3 — about half the fetches.
|
||
// The center tap reuses the caller's already-sampled raw texel.
|
||
const accumulate = ( dx, dy, sample ) => {
|
||
|
||
const neighbor = sample !== undefined
|
||
? sample
|
||
: texture( this.rawNode, uvCoord.add( vec2( dx, dy ).div( this._resolution ) ) ).max( 0 ).toConst();
|
||
|
||
if ( this.alphaSource === 'raylength' ) {
|
||
|
||
const sampleRl = neighbor.a.toVar();
|
||
If( sampleRl.greaterThan( ENV_RAY_LENGTH_THRESHOLD ), () => {
|
||
|
||
sampleRl.assign( 0.25 );
|
||
hasEnvRay.assign( true );
|
||
|
||
} );
|
||
const w = float( 1 ).div( sampleRl.add( 0.001 ) );
|
||
rlSum.addAssign( sampleRl.mul( w ) );
|
||
rlSumW.addAssign( w );
|
||
|
||
}
|
||
|
||
If( this.adaptiveTrust.greaterThan( 0 ), () => {
|
||
|
||
const nLuma = luminance( neighbor.rgb );
|
||
lumaCount.addAssign( 1 );
|
||
const delta = nLuma.sub( meanLuma ).toConst();
|
||
meanLuma.addAssign( delta.div( lumaCount ) );
|
||
m2Luma.addAssign( delta.mul( nLuma.sub( meanLuma ) ) );
|
||
|
||
} );
|
||
|
||
};
|
||
|
||
accumulate( 0, 0, centerSample );
|
||
accumulate( - 1, 0 );
|
||
accumulate( 1, 0 );
|
||
accumulate( 0, - 1 );
|
||
accumulate( 0, 1 );
|
||
|
||
const avgRayLength = this.alphaSource === 'raylength' ? rlSum.div( rlSumW ) : float( 1 );
|
||
const stddevLuma = sqrt( m2Luma.div( lumaCount.max( 1 ) ) );
|
||
|
||
// vec3( avgRayLength, meanLuma, stddevLuma )
|
||
return vec4( avgRayLength, meanLuma, stddevLuma, hasEnvRay.toFloat() );
|
||
|
||
} ).setLayout( {
|
||
name: 'getNeighborhoodStats',
|
||
type: 'vec4',
|
||
inputs: [
|
||
{ name: 'uvCoord', type: 'vec2' },
|
||
{ name: 'centerSample', type: 'vec4' }
|
||
]
|
||
} );
|
||
|
||
const denoiseFn = Fn( ( [ uvCoord ] ) => {
|
||
|
||
const result = property( 'vec4' );
|
||
|
||
const depth = sampleDepth( uvCoord ).toConst();
|
||
|
||
const runDenoise = () => {
|
||
|
||
const viewNormal = sampleNormal( uvCoord ).toConst();
|
||
const worldNormal = viewNormal.transformDirection( this._viewMatrix ).toConst();
|
||
const texel = sampleTexture( uvCoord ).max( 0 ).toConst();
|
||
|
||
const viewPosition = getViewPosition( uvCoord, depth, this._cameraProjectionMatrixInverse ).toConst();
|
||
const roughnessMetalness = sampleRoughnessMetalness( uvCoord ).toConst();
|
||
const roughness = roughnessMetalness.g;
|
||
const metalness = roughnessMetalness.r;
|
||
|
||
const noiseTexel = sampleAnalyticNoise( uvCoord, this._noiseIndex );
|
||
const rotationMatrix = noiseRotationMatrix( noiseTexel.r );
|
||
|
||
const frameNum = float( 1 ).div( texel.a );
|
||
const varianceFactor = getTemporalVarianceFactor( frameNum, this.strength.oneMinus() );
|
||
const aggressivity = varianceFactor.oneMinus();
|
||
|
||
const raw = sampleRaw( uvCoord ).toConst();
|
||
|
||
const viewZ = abs( viewPosition.z );
|
||
const rl = float( 1 ).toVar();
|
||
const nbhdMeanLuma = float( 0 ).toVar();
|
||
const nbhdStddevLuma = float( 0 ).toVar();
|
||
const hasEnvRay = bool( false ).toVar();
|
||
|
||
if ( this.alphaSource === 'raylength' ) {
|
||
|
||
const stats = getNeighborhoodStats( uvCoord, raw );
|
||
rl.assign( stats.x );
|
||
nbhdMeanLuma.assign( stats.y );
|
||
nbhdStddevLuma.assign( stats.z );
|
||
hasEnvRay.assign( stats.w.greaterThan( 0.5 ) );
|
||
|
||
} else {
|
||
|
||
If( this.adaptiveTrust.greaterThan( 0 ), () => {
|
||
|
||
const stats = getNeighborhoodStats( uvCoord, raw );
|
||
nbhdMeanLuma.assign( stats.y );
|
||
nbhdStddevLuma.assign( stats.z );
|
||
|
||
} );
|
||
|
||
}
|
||
|
||
const tanHalfFovY = this.alphaSource === 'raylength' ? tan( this._fovY.mul( 0.5 ) ).toConst() : null;
|
||
const hitDistFactor = this.alphaSource === 'raylength'
|
||
? computeHitDistFactor( rl, viewZ, tanHalfFovY ).toConst()
|
||
: float( 1 );
|
||
|
||
const denoised = texel.rgb.toVar();
|
||
const totalWeight = float( 1 ).toVar();
|
||
const denoisedFrame = frameNum.toVar();
|
||
const totalFrameWeight = float( 1 ).toVar();
|
||
|
||
const denoisedRaw = raw.rgb.toVar();
|
||
const totalWeightRaw = float( 1 ).toVar();
|
||
|
||
If( raw.rgb.length().lessThan( 0.0001 ), () => {
|
||
|
||
denoisedRaw.assign( vec3( 0 ) );
|
||
totalWeightRaw.assign( 0 );
|
||
|
||
} );
|
||
|
||
const avgAo = this.alphaSource === 'ao' ? raw.a.toConst() : float( 1 );
|
||
const mappedAvgAo = this.alphaSource === 'ao' ? mapAo( avgAo ) : float( 0 );
|
||
|
||
const worldRadius = this.radius.mul( WORLD_RADIUS_SCALE ).toVar();
|
||
|
||
if ( this.mode === 'specular' ) {
|
||
|
||
worldRadius.mulAssign( rl.mul( viewPosition.z.abs() ) );
|
||
worldRadius.mulAssign( roughness.sqrt().max( 0.01 ) );
|
||
|
||
} else {
|
||
|
||
worldRadius.mulAssign( avgAo.pow( 2 ).mul( viewPosition.z.abs() ) );
|
||
|
||
}
|
||
|
||
worldRadius.mulAssign( mix( 1, AGGRESSIVITY_RADIUS_MIN, aggressivity ) );
|
||
|
||
const T = vec3( 0 ).toVar();
|
||
const B = vec3( 0 ).toVar();
|
||
|
||
if ( this.mode === 'specular' ) {
|
||
|
||
const V = normalize( viewPosition ).negate();
|
||
const D = getSpecularDominantDirection( viewNormal, V, roughness );
|
||
const R = reflect( D.negate(), viewNormal );
|
||
const Tv = normalize( cross( viewNormal, R ) );
|
||
const Bv = cross( R, Tv );
|
||
const viewAngle = abs( viewNormal.z ).acos().div( float( Math.PI * 0.5 ) ).clamp( 0, 1 );
|
||
const skewFactor = mix( 1.0, roughness, viewAngle );
|
||
T.assign( Tv.mul( skewFactor ) );
|
||
B.assign( Bv );
|
||
|
||
} else {
|
||
|
||
const up = vec3( 0, 0, 1 );
|
||
const Tv = cross( up, viewNormal ).normalize().toVar();
|
||
If( Tv.length().lessThan( EPSILON ), () => {
|
||
|
||
Tv.assign( cross( vec3( 0, 1, 0 ), viewNormal ).normalize() );
|
||
|
||
} );
|
||
T.assign( Tv );
|
||
B.assign( cross( viewNormal, Tv ).normalize() );
|
||
|
||
}
|
||
|
||
T.mulAssign( worldRadius );
|
||
B.mulAssign( worldRadius );
|
||
|
||
const centerDiffuse = sampleDiffuse( uvCoord ).toConst();
|
||
const radiusShrink = float( 1 ).toVar();
|
||
|
||
// Directional analog of radiusShrink: an accumulated tangent-space shift that skews
|
||
// subsequent taps toward directions that yielded high weight (related geometry).
|
||
const polarBias = vec2( 0 ).toVar();
|
||
|
||
const depthWeightScale = this.depthPhi.mul( 500 ).mul( viewNormal.z.abs() ).div( viewPosition.z.abs() );
|
||
|
||
// Lobe geometry depends only on per-pixel terms, so compute its falloff constant once here.
|
||
const lobeFalloff = lobeNormalFalloff( roughness, aggressivity, this.normalPhi.oneMinus() ).toConst();
|
||
|
||
Loop( { start: int( 0 ), end: int( KERNEL_SAMPLES ), type: 'int', condition: '<', name: 'i' }, ( { i } ) => {
|
||
|
||
const baseOffset = vogelDisk( float( i ), 1 ).toVar();
|
||
const sampleDir = baseOffset.normalize().toConst();
|
||
|
||
// Blend the tap direction toward the polar bias, then restore the Vogel radius and shrink.
|
||
const skewedDir = mix( sampleDir, polarBias.max( EPSILON ).normalize(), this.adapt.mul( aggressivity )
|
||
.mul( polarBias.dot( polarBias ).greaterThan( 0.001 ).select( 1, 0 ) ) );
|
||
const offset = rotationMatrix.mul( skewedDir.mul( baseOffset.length().mul( radiusShrink ) ) ).toVar();
|
||
|
||
// Exact per-sample view-space projection (both paths)
|
||
const sampleViewPos = viewPosition.add( B.mul( offset.x ).add( T.mul( offset.y ) ) );
|
||
const sampleUv = getScreenPosition( sampleViewPos, this._cameraProjectionMatrix ).toVar();
|
||
sampleUv.assign( sampleUv.abs().oneMinus().abs().oneMinus().clamp() );
|
||
|
||
const neighborColor = sampleTexture( sampleUv ).max( 0 ).toConst();
|
||
|
||
// When no raw texture is bound, sampleRaw falls back to the filtered texture at the same UV.
|
||
const rawNeighborColor = sampleRaw( sampleUv ).max( 0 ).toVar();
|
||
// if ( this.mode === 'diffuse' ) rawNeighborColor.rgb.assign( mix( neighborColor.rgb, rawNeighborColor.rgb, neighborColor.a ) );
|
||
|
||
const nDepth = sampleDepth( sampleUv );
|
||
const nViewPosition = getViewPosition( sampleUv, nDepth, this._cameraProjectionMatrixInverse ).toConst();
|
||
const nViewZ = abs( nViewPosition.z ).toConst();
|
||
|
||
const kernelDiff = float( 0 ).toVar();
|
||
|
||
// Luma edge stopping
|
||
kernelDiff.addAssign( luminance( rawNeighborColor.rgb ).sub( luminance( raw.rgb ) ).abs().mul( this.lumaPhi ).mul( 10 ) );
|
||
|
||
// Diffuse edge stopping (only relevant for specular mode)
|
||
if ( this.diffuseNode !== null ) {
|
||
|
||
kernelDiff.addAssign( ( diffuseColorDistance( centerDiffuse, sampleDiffuse( sampleUv ), float( 0 ) ).mul( this.diffusePhi ).mul( metalness ) ) );
|
||
|
||
}
|
||
|
||
// AO edge stopping
|
||
if ( this.alphaSource === 'ao' ) {
|
||
|
||
const neighborMappedAo = mapAo( rawNeighborColor.a );
|
||
// We multiply here with aggressivity as well, since early application of aoW yields noise
|
||
const aoW = mappedAvgAo.div( mappedAvgAo.add( neighborMappedAo ).add( AO_EDGE_STOPPING_BIAS ) ).mul( this.alphaPhi ).mul( aggressivity );
|
||
|
||
kernelDiff.addAssign( ( aoW ) );
|
||
|
||
|
||
} else if ( this.alphaSource === 'raylength' ) {
|
||
|
||
// Ray length edge stopping
|
||
|
||
const neighborHitDistFactor = computeHitDistFactor( rawNeighborColor.a, nViewZ, tanHalfFovY );
|
||
const hdfDiff = hitDistFactor.sub( neighborHitDistFactor ).abs();
|
||
|
||
const rayLengthFactor = hdfDiff.mul( this.alphaPhi ).div( viewPosition.z.abs() );
|
||
|
||
// Env rays are harder to compare so we accept if this sample is an env ray and there is an env ray in the neighborhood
|
||
kernelDiff.addAssign( rawNeighborColor.a.greaterThan( ENV_RAY_LENGTH_THRESHOLD ).and( hasEnvRay ).select( 1, rayLengthFactor ) );
|
||
|
||
}
|
||
|
||
// Roughness edge stopping
|
||
if ( this.mode === 'specular' ) kernelDiff.addAssign( ( abs( roughness.sub( sampleRoughnessMetalness( sampleUv ).g ) ).mul( this.roughnessPhi ) ) );
|
||
|
||
const nViewNormal = sampleNormal( sampleUv );
|
||
const nWorldNormal = nViewNormal.transformDirection( this._viewMatrix );
|
||
const distToPlane = planeDistance( viewPosition, nViewPosition, viewNormal );
|
||
|
||
// Geometric edge stopping (depth and normal)
|
||
const depthDiff = distToPlane.mul( depthWeightScale );
|
||
const normalW = lobeNormalWeight( worldNormal, nWorldNormal, lobeFalloff );
|
||
|
||
// Sum every negative-exponent edge-stopping term (kernel + depth/plane, plus the SSR hit-distance term)
|
||
const w = exp( kernelDiff.mul( aggressivity ).add( depthDiff ).negate() ).mul( normalW ).toVar();
|
||
|
||
// Feedback to shrink radius based on the weight
|
||
radiusShrink.assign( mix( radiusShrink, w, this.adapt ) );
|
||
|
||
// Polar feedback: skew subsequent taps toward high-weight directions (related geometry)
|
||
polarBias.assign( mix( polarBias, sampleDir.mul( w.sub( 0.5 ) ), 0.5 ) );
|
||
|
||
// to mitigate the effect of fireflies and high variance in recently disoccluded regions, we weigh by the inverse luminance for the first 5 frames
|
||
w.mulAssign( mix( float( 1 ).div( luminance( rawNeighborColor.rgb ).pow( 2 ).add( 0.01 ) ), 1, frameNum.div( 5 ).min( 1 ) ) );
|
||
|
||
denoisedRaw.addAssign( rawNeighborColor.rgb.mul( w ) );
|
||
totalWeightRaw.addAssign( w );
|
||
|
||
denoised.addAssign( neighborColor.rgb.mul( w ) );
|
||
totalWeight.addAssign( w );
|
||
|
||
// Denoising the alpha (accumulation speed), to get smoother disocclusion transitions
|
||
If( this.smoothDisocclusions, () => {
|
||
|
||
const neighborAWeight = neighborColor.a.greaterThan( texel.a ).select( w.mul( 0.33 ), 0 );
|
||
denoisedFrame.addAssign( float( 1 ).div( neighborColor.a ).mul( neighborAWeight ) );
|
||
totalFrameWeight.addAssign( neighborAWeight );
|
||
|
||
} );
|
||
|
||
} );
|
||
|
||
denoised.divAssign( totalWeight.max( EPSILON ) );
|
||
denoised.assign( denoised.max( EPSILON ) );
|
||
denoisedRaw.divAssign( totalWeightRaw.max( EPSILON ) );
|
||
|
||
if ( this.accumulate ) {
|
||
|
||
const computedFrame = denoisedFrame.div( totalFrameWeight.max( EPSILON ) );
|
||
const a = float( 1 ).div( computedFrame.max( EPSILON ) ).toConst();
|
||
|
||
if ( this.rawNode !== null ) {
|
||
|
||
const blended = karisTemporalBlend(
|
||
denoised,
|
||
denoisedRaw,
|
||
a,
|
||
this.flickerSuppression,
|
||
this.adaptiveTrust,
|
||
nbhdMeanLuma,
|
||
nbhdStddevLuma
|
||
);
|
||
|
||
result.assign( vec4( blended, a ) );
|
||
|
||
} else {
|
||
|
||
const finalDenoised = mix( denoised, denoisedRaw, a );
|
||
result.assign( vec4( finalDenoised, a ) );
|
||
|
||
}
|
||
|
||
} else {
|
||
|
||
result.assign( vec4( denoised, texel.a ) );
|
||
|
||
}
|
||
|
||
};
|
||
|
||
If( depth.greaterThanEqual( 1.0 ), () => {
|
||
|
||
Discard();
|
||
|
||
} ).Else( runDenoise );
|
||
|
||
return result;
|
||
|
||
} );
|
||
|
||
this._material.fragmentNode = denoiseFn( uv() ).context( builder.getSharedContext() );
|
||
this._material.needsUpdate = true;
|
||
|
||
return this._textureNode;
|
||
|
||
}
|
||
|
||
dispose() {
|
||
|
||
this._renderTarget.dispose();
|
||
this._material.dispose();
|
||
|
||
}
|
||
|
||
}
|
||
|
||
export default RecurrentDenoiseNode;
|
||
|
||
/**
|
||
* @tsl
|
||
* @param {Node} inputTexture - Temporally filtered input to denoise (e.g. TRAA output).
|
||
* @param {Camera} camera
|
||
* @param {RecurrentDenoiseNodeOptions} [options={}]
|
||
* @returns {RecurrentDenoiseNode}
|
||
*/
|
||
export const recurrentDenoise = ( inputTexture, camera, options = {} ) => nodeObject( new RecurrentDenoiseNode(
|
||
toTextureNode( inputTexture ),
|
||
camera,
|
||
options
|
||
) );
|