594 lines
17 KiB
JavaScript
594 lines
17 KiB
JavaScript
import { DataTexture, RenderTarget, RepeatWrapping, Vector2, Vector3, TempNode, QuadMesh, NodeMaterial, RendererUtils, RedFormat } from 'three/webgpu';
|
||
import { reference, logarithmicDepthToViewZ, viewZToPerspectiveDepth, getNormalFromDepth, getScreenPosition, getViewPosition, 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, acos, clamp } from 'three/tsl';
|
||
|
||
const _quadMesh = /*@__PURE__*/ new QuadMesh();
|
||
const _size = /*@__PURE__*/ new Vector2();
|
||
|
||
// From Activision GTAO paper: https://www.activision.com/cdn/research/s2016_pbs_activision_occlusion.pptx
|
||
const _temporalRotations = [ 60, 300, 180, 240, 120, 0 ];
|
||
|
||
let _rendererState;
|
||
|
||
/**
|
||
* Post processing node for applying Ground Truth Ambient Occlusion (GTAO) to a scene.
|
||
* ```js
|
||
* const renderPipeline = new THREE.RenderPipeline( renderer );
|
||
*
|
||
* const scenePass = pass( scene, camera );
|
||
* scenePass.setMRT( mrt( {
|
||
* output: output,
|
||
* normal: normalView
|
||
* } ) );
|
||
*
|
||
* const scenePassColor = scenePass.getTextureNode( 'output' );
|
||
* const scenePassNormal = scenePass.getTextureNode( 'normal' );
|
||
* const scenePassDepth = scenePass.getTextureNode( 'depth' );
|
||
*
|
||
* const aoPass = ao( scenePassDepth, scenePassNormal, camera );
|
||
* const aoPassOutput = aoPass.getTextureNode();
|
||
*
|
||
* renderPipeline.outputNode = scenePassColor.mul( vec4( vec3( aoPassOutput.r ), 1 ) );
|
||
* ```
|
||
*
|
||
* Reference: [Practical Real-Time Strategies for Accurate Indirect Occlusion](https://www.activision.com/cdn/research/Practical_Real_Time_Strategies_for_Accurate_Indirect_Occlusion_NEW%20VERSION_COLOR.pdf).
|
||
*
|
||
* @augments TempNode
|
||
* @three_import import { ao } from 'three/addons/tsl/display/GTAONode.js';
|
||
*/
|
||
class GTAONode extends TempNode {
|
||
|
||
static get type() {
|
||
|
||
return 'GTAONode';
|
||
|
||
}
|
||
|
||
/**
|
||
* Constructs a new GTAO node.
|
||
*
|
||
* @param {Node<float>} depthNode - A node that represents the scene's depth.
|
||
* @param {?Node<vec3>} normalNode - A node that represents the scene's normals.
|
||
* @param {Camera} camera - The camera the scene is rendered with.
|
||
*/
|
||
constructor( depthNode, normalNode, camera ) {
|
||
|
||
super( 'float' );
|
||
|
||
/**
|
||
* A node that represents the scene's depth.
|
||
*
|
||
* @type {Node<float>}
|
||
*/
|
||
this.depthNode = depthNode;
|
||
|
||
/**
|
||
* A node that represents the scene's normals. If no normals are passed to the
|
||
* constructor (because MRT is not available), normals can be automatically
|
||
* reconstructed from depth values in the shader.
|
||
*
|
||
* @type {?Node<vec3>}
|
||
*/
|
||
this.normalNode = normalNode;
|
||
|
||
/**
|
||
* The resolution scale. By default the effect is rendered in full resolution
|
||
* for best quality but a value of `0.5` should be sufficient for most scenes.
|
||
*
|
||
* @type {number}
|
||
* @default 1
|
||
*/
|
||
this.resolutionScale = 1;
|
||
|
||
/**
|
||
* The `updateBeforeType` is set to `NodeUpdateType.FRAME` since the node renders
|
||
* its effect once per frame in `updateBefore()`.
|
||
*
|
||
* @type {string}
|
||
* @default 'frame'
|
||
*/
|
||
this.updateBeforeType = NodeUpdateType.FRAME;
|
||
|
||
/**
|
||
* The render target the ambient occlusion is rendered into.
|
||
*
|
||
* @private
|
||
* @type {RenderTarget}
|
||
*/
|
||
this._aoRenderTarget = new RenderTarget( 1, 1, { depthBuffer: false, format: RedFormat } );
|
||
this._aoRenderTarget.texture.name = 'GTAONode.AO';
|
||
|
||
// uniforms
|
||
|
||
/**
|
||
* The radius of the ambient occlusion.
|
||
*
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this.radius = uniform( 0.25 );
|
||
|
||
/**
|
||
* The resolution of the effect. Can be scaled via
|
||
* `resolutionScale`.
|
||
*
|
||
* @type {UniformNode<vec2>}
|
||
*/
|
||
this.resolution = uniform( new Vector2() );
|
||
|
||
/**
|
||
* The thickness of the ambient occlusion.
|
||
*
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this.thickness = uniform( 1 );
|
||
|
||
/**
|
||
* Another option to tweak the occlusion. The recommended range is
|
||
* `[1,2]` for attenuating the AO.
|
||
*
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this.distanceExponent = uniform( 1 );
|
||
|
||
/**
|
||
* The distance fall off value of the ambient occlusion.
|
||
* A lower value leads to a larger AO effect. The value
|
||
* should lie in the range `[0,1]`.
|
||
*
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this.distanceFallOff = uniform( 1 );
|
||
|
||
/**
|
||
* The scale of the ambient occlusion.
|
||
*
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this.scale = uniform( 1 );
|
||
|
||
/**
|
||
* How many samples are used to compute the AO.
|
||
* A higher value results in better quality but also
|
||
* in a more expensive runtime behavior.
|
||
*
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this.samples = uniform( 16 );
|
||
|
||
/**
|
||
* Whether to use temporal filtering or not. Setting this property to
|
||
* `true` requires the usage of `TRAANode`. This will help to reduce noise
|
||
* although it introduces typical TAA artifacts like ghosting and temporal
|
||
* instabilities.
|
||
*
|
||
* If setting this property to `false`, a manual denoise via `DenoiseNode`
|
||
* might be required.
|
||
*
|
||
* @type {boolean}
|
||
* @default false
|
||
*/
|
||
this.useTemporalFiltering = false;
|
||
|
||
/**
|
||
* The node represents the internal noise texture used by the AO.
|
||
*
|
||
* @private
|
||
* @type {TextureNode}
|
||
*/
|
||
this._noiseNode = texture( generateMagicSquareNoise() );
|
||
|
||
/**
|
||
* Represents the projection matrix of the scene's camera.
|
||
*
|
||
* @private
|
||
* @type {UniformNode<mat4>}
|
||
*/
|
||
this._cameraProjectionMatrix = uniform( camera.projectionMatrix );
|
||
|
||
/**
|
||
* Represents the inverse projection matrix of the scene's camera.
|
||
*
|
||
* @private
|
||
* @type {UniformNode<mat4>}
|
||
*/
|
||
this._cameraProjectionMatrixInverse = uniform( camera.projectionMatrixInverse );
|
||
|
||
/**
|
||
* Represents the near value of the scene's camera.
|
||
*
|
||
* @private
|
||
* @type {ReferenceNode<float>}
|
||
*/
|
||
this._cameraNear = reference( 'near', 'float', camera );
|
||
|
||
/**
|
||
* Represents the far value of the scene's camera.
|
||
*
|
||
* @private
|
||
* @type {ReferenceNode<float>}
|
||
*/
|
||
this._cameraFar = reference( 'far', 'float', camera );
|
||
|
||
/**
|
||
* Temporal direction that influences the rotation angle for each slice.
|
||
*
|
||
* @private
|
||
* @type {UniformNode<float>}
|
||
*/
|
||
this._temporalDirection = uniform( 0 );
|
||
|
||
/**
|
||
* The material that is used to render the effect.
|
||
*
|
||
* @private
|
||
* @type {NodeMaterial}
|
||
*/
|
||
this._material = new NodeMaterial();
|
||
this._material.name = 'GTAO';
|
||
|
||
/**
|
||
* The result of the effect is represented as a separate texture node.
|
||
*
|
||
* @private
|
||
* @type {PassTextureNode}
|
||
*/
|
||
this._textureNode = passTexture( this, this._aoRenderTarget.texture );
|
||
|
||
}
|
||
|
||
/**
|
||
* Returns the result of the effect as a texture node.
|
||
*
|
||
* @return {PassTextureNode} A texture node that represents the result of the effect.
|
||
*/
|
||
getTextureNode() {
|
||
|
||
return this._textureNode;
|
||
|
||
}
|
||
|
||
/**
|
||
* Sets the size of the effect.
|
||
*
|
||
* @param {number} width - The width of the effect.
|
||
* @param {number} height - The height of the effect.
|
||
*/
|
||
setSize( width, height ) {
|
||
|
||
width = Math.round( this.resolutionScale * width );
|
||
height = Math.round( this.resolutionScale * height );
|
||
|
||
this.resolution.value.set( width, height );
|
||
this._aoRenderTarget.setSize( width, height );
|
||
|
||
}
|
||
|
||
/**
|
||
* This method is used to render the effect once per frame.
|
||
*
|
||
* @param {NodeFrame} frame - The current node frame.
|
||
*/
|
||
updateBefore( frame ) {
|
||
|
||
const { renderer } = frame;
|
||
|
||
_rendererState = RendererUtils.resetRendererState( renderer, _rendererState );
|
||
|
||
// update temporal uniforms
|
||
|
||
if ( this.useTemporalFiltering === true ) {
|
||
|
||
const frameId = frame.frameId;
|
||
|
||
this._temporalDirection.value = _temporalRotations[ frameId % 6 ] / 360;
|
||
|
||
} else {
|
||
|
||
this._temporalDirection.value = 0;
|
||
|
||
}
|
||
|
||
//
|
||
|
||
const size = renderer.getDrawingBufferSize( _size );
|
||
this.setSize( size.width, size.height );
|
||
|
||
_quadMesh.material = this._material;
|
||
_quadMesh.name = 'AO';
|
||
|
||
// clear
|
||
|
||
renderer.setClearColor( 0xffffff, 1 );
|
||
|
||
// ao
|
||
|
||
renderer.setRenderTarget( this._aoRenderTarget );
|
||
_quadMesh.render( renderer );
|
||
|
||
// restore
|
||
|
||
RendererUtils.restoreRendererState( renderer, _rendererState );
|
||
|
||
}
|
||
|
||
/**
|
||
* This method is used to setup the effect's TSL code.
|
||
*
|
||
* @param {NodeBuilder} builder - The current node builder.
|
||
* @return {PassTextureNode}
|
||
*/
|
||
setup( builder ) {
|
||
|
||
const uvNode = uv();
|
||
|
||
const sampleDepth = ( uv ) => {
|
||
|
||
const depth = this.depthNode.sample( uv ).r;
|
||
|
||
if ( builder.renderer.logarithmicDepthBuffer === true ) {
|
||
|
||
const viewZ = logarithmicDepthToViewZ( depth, this._cameraNear, this._cameraFar );
|
||
|
||
return viewZToPerspectiveDepth( viewZ, this._cameraNear, this._cameraFar );
|
||
|
||
}
|
||
|
||
return depth;
|
||
|
||
};
|
||
|
||
const sampleNoise = ( uv ) => this._noiseNode.sample( uv );
|
||
const sampleNormal = ( uv ) => ( this.normalNode !== null ) ? this.normalNode.sample( uv ).rgb.normalize() : getNormalFromDepth( uv, this.depthNode.value, this._cameraProjectionMatrixInverse );
|
||
|
||
const ao = Fn( () => {
|
||
|
||
const depth = sampleDepth( uvNode ).toVar();
|
||
|
||
depth.greaterThanEqual( 1.0 ).discard();
|
||
|
||
const viewPosition = getViewPosition( uvNode, depth, this._cameraProjectionMatrixInverse ).toVar();
|
||
const viewNormal = sampleNormal( uvNode ).toVar();
|
||
|
||
const radiusToUse = this.radius;
|
||
|
||
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 randomVec = noiseTexel.xyz.mul( 2.0 ).sub( 1.0 );
|
||
const tangent = vec3( randomVec.xy, 0.0 ).normalize();
|
||
const bitangent = vec3( tangent.y.mul( - 1.0 ), tangent.x, 0.0 );
|
||
const kernelMatrix = mat3( tangent, bitangent, vec3( 0.0, 0.0, 1.0 ) );
|
||
|
||
const DIRECTIONS = this.samples.lessThan( 30 ).select( 3, 5 ).toVar();
|
||
const STEPS = add( this.samples, DIRECTIONS.sub( 1 ) ).div( DIRECTIONS ).toVar();
|
||
|
||
const ao = float( 0 ).toVar();
|
||
|
||
// Each iteration analyzes one vertical "slice" of the 3D space around the fragment.
|
||
|
||
Loop( { start: int( 0 ), end: DIRECTIONS, type: 'int', condition: '<' }, ( { i } ) => {
|
||
|
||
const angle = float( i ).div( float( DIRECTIONS ) ).mul( PI ).add( this._temporalDirection ).toVar();
|
||
const sampleDir = vec4( cos( angle ), sin( angle ), 0., add( 0.5, mul( 0.5, noiseTexel.w ) ) );
|
||
sampleDir.xyz = normalize( kernelMatrix.mul( sampleDir.xyz ) );
|
||
|
||
const viewDir = normalize( viewPosition.xyz.negate() ).toVar();
|
||
const sliceBitangent = normalize( cross( sampleDir.xyz, viewDir ) ).toVar();
|
||
const sliceTangent = cross( sliceBitangent, viewDir ).toVar();
|
||
|
||
// Project the view normal onto the slice plane (remove component along sliceBitangent).
|
||
// The unnormalized length is the foreshortening weight applied at slice integration.
|
||
// (Activision GTAO paper, Section 3.2 "Per-pixel sampling".)
|
||
const projNRaw = viewNormal.sub( sliceBitangent.mul( dot( viewNormal, sliceBitangent ) ) ).toVar();
|
||
const projNLen = projNRaw.length().toVar();
|
||
const projN = projNRaw.div( max( projNLen, float( 0.0001 ) ) ).toVar();
|
||
|
||
// γ — angle of projN within the slice plane, signed by the tangent direction.
|
||
const nSin = dot( projN, sliceTangent ).toVar();
|
||
const nCos = clamp( dot( projN, viewDir ), 0, 1 ).toVar();
|
||
const signNSin = nSin.greaterThanEqual( 0 ).select( float( 1 ), float( - 1 ) );
|
||
const angleN = signNSin.mul( acos( nCos ) ).toVar();
|
||
|
||
const tangentToNormalInSlice = cross( projN, sliceBitangent ).toVar();
|
||
const cosHorizons = vec2( dot( viewDir, tangentToNormalInSlice ), dot( viewDir, tangentToNormalInSlice.negate() ) ).toVar();
|
||
|
||
// For each slice, the inner loop performs ray marching to find the horizons.
|
||
|
||
Loop( { end: STEPS, type: 'int', name: 'j', condition: '<' }, ( { j } ) => {
|
||
|
||
const sampleViewOffset = sampleDir.xyz.mul( radiusToUse ).mul( sampleDir.w ).mul( pow( div( float( j ).add( 1.0 ), float( STEPS ) ), this.distanceExponent ) );
|
||
|
||
// The loop marches in two opposite directions (x and y) along the slice's line to find the horizon on both sides.
|
||
|
||
// x
|
||
|
||
const sampleScreenPositionX = getScreenPosition( viewPosition.add( sampleViewOffset ), this._cameraProjectionMatrix ).toVar();
|
||
const sampleDepthX = sampleDepth( sampleScreenPositionX ).toVar();
|
||
const sampleSceneViewPositionX = getViewPosition( sampleScreenPositionX, sampleDepthX, this._cameraProjectionMatrixInverse ).toVar();
|
||
const viewDeltaX = sampleSceneViewPositionX.sub( viewPosition ).toVar();
|
||
|
||
If( abs( viewDeltaX.z ).lessThan( this.thickness ), () => {
|
||
|
||
const sampleCosHorizon = dot( viewDir, normalize( viewDeltaX ) );
|
||
cosHorizons.x.addAssign( max( 0, mul( sampleCosHorizon.sub( cosHorizons.x ), mix( 1.0, float( 2.0 ).div( float( j ).add( 2 ) ), this.distanceFallOff ) ) ) );
|
||
|
||
} );
|
||
|
||
// y
|
||
|
||
const sampleScreenPositionY = getScreenPosition( viewPosition.sub( sampleViewOffset ), this._cameraProjectionMatrix ).toVar();
|
||
const sampleDepthY = sampleDepth( sampleScreenPositionY ).toVar();
|
||
const sampleSceneViewPositionY = getViewPosition( sampleScreenPositionY, sampleDepthY, this._cameraProjectionMatrixInverse ).toVar();
|
||
const viewDeltaY = sampleSceneViewPositionY.sub( viewPosition ).toVar();
|
||
|
||
If( abs( viewDeltaY.z ).lessThan( this.thickness ), () => {
|
||
|
||
const sampleCosHorizon = dot( viewDir, normalize( viewDeltaY ) );
|
||
cosHorizons.y.addAssign( max( 0, mul( sampleCosHorizon.sub( cosHorizons.y ), mix( 1.0, float( 2.0 ).div( float( j ).add( 2 ) ), this.distanceFallOff ) ) ) );
|
||
|
||
} );
|
||
|
||
} );
|
||
|
||
// Cosine-weighted inner integral, closed-form (Activision GTAO paper, Eq. 7).
|
||
// Per horizon h_i: term_i = −cos( 2 h_i − γ ) + cos( γ ) + 2 h_i sin( γ )
|
||
// The 0.25 factor is ½ (integral normalization) × ½ (averaging the two horizons).
|
||
//
|
||
// In this slice setup `sliceTangent = cross( sliceBitangent, viewDir )` works out
|
||
// opposite to `sampleDir`, so the +sampleDir samples (cosHorizons.x) live on the
|
||
// −T side of the slice and −sampleDir samples (cosHorizons.y) on the +T side.
|
||
// γ is signed by +T (sliceTangent), so hPos must read from cosHorizons.y.
|
||
|
||
const hPos = acos( cosHorizons.y ).toVar();
|
||
const hNeg = acos( cosHorizons.x ).negate().toVar();
|
||
|
||
const termPos = cos( hPos.mul( 2 ).sub( angleN ) ).negate().add( nCos ).add( hPos.mul( 2 ).mul( nSin ) );
|
||
const termNeg = cos( hNeg.mul( 2 ).sub( angleN ) ).negate().add( nCos ).add( hNeg.mul( 2 ).mul( nSin ) );
|
||
const a = termPos.add( termNeg ).mul( 0.25 );
|
||
|
||
// |projN| is the foreshortening weight from the per-slice normal projection.
|
||
ao.addAssign( projNLen.mul( a ) );
|
||
|
||
} );
|
||
|
||
ao.assign( clamp( ao.div( DIRECTIONS ), 0, 1 ) );
|
||
ao.assign( pow( ao, this.scale ) );
|
||
|
||
return ao;
|
||
|
||
} );
|
||
|
||
this._material.fragmentNode = ao().context( builder.getSharedContext() );
|
||
this._material.needsUpdate = true;
|
||
|
||
//
|
||
|
||
return this._textureNode;
|
||
|
||
}
|
||
|
||
/**
|
||
* Frees internal resources. This method should be called
|
||
* when the effect is no longer required.
|
||
*/
|
||
dispose() {
|
||
|
||
this._aoRenderTarget.dispose();
|
||
|
||
this._material.dispose();
|
||
|
||
}
|
||
|
||
}
|
||
|
||
export default GTAONode;
|
||
|
||
/**
|
||
* Generates the AO's noise texture for the given size.
|
||
*
|
||
* @param {number} [size=5] - The noise size.
|
||
* @return {DataTexture} The generated noise texture.
|
||
*/
|
||
function generateMagicSquareNoise( size = 5 ) {
|
||
|
||
const noiseSize = Math.floor( size ) % 2 === 0 ? Math.floor( size ) + 1 : Math.floor( size );
|
||
const magicSquare = generateMagicSquare( noiseSize );
|
||
const noiseSquareSize = magicSquare.length;
|
||
const data = new Uint8Array( noiseSquareSize * 4 );
|
||
|
||
for ( let inx = 0; inx < noiseSquareSize; ++ inx ) {
|
||
|
||
const iAng = magicSquare[ inx ];
|
||
const angle = ( 2 * Math.PI * iAng ) / noiseSquareSize;
|
||
const randomVec = new Vector3(
|
||
Math.cos( angle ),
|
||
Math.sin( angle ),
|
||
0
|
||
).normalize();
|
||
data[ inx * 4 ] = ( randomVec.x * 0.5 + 0.5 ) * 255;
|
||
data[ inx * 4 + 1 ] = ( randomVec.y * 0.5 + 0.5 ) * 255;
|
||
data[ inx * 4 + 2 ] = 127;
|
||
data[ inx * 4 + 3 ] = 255;
|
||
|
||
}
|
||
|
||
const noiseTexture = new DataTexture( data, noiseSize, noiseSize );
|
||
noiseTexture.wrapS = RepeatWrapping;
|
||
noiseTexture.wrapT = RepeatWrapping;
|
||
noiseTexture.needsUpdate = true;
|
||
|
||
return noiseTexture;
|
||
|
||
}
|
||
|
||
/**
|
||
* Computes an array of magic square values required to generate the noise texture.
|
||
*
|
||
* @param {number} size - The noise size.
|
||
* @return {Array<number>} The magic square values.
|
||
*/
|
||
function generateMagicSquare( size ) {
|
||
|
||
const noiseSize = Math.floor( size ) % 2 === 0 ? Math.floor( size ) + 1 : Math.floor( size );
|
||
const noiseSquareSize = noiseSize * noiseSize;
|
||
const magicSquare = Array( noiseSquareSize ).fill( 0 );
|
||
let i = Math.floor( noiseSize / 2 );
|
||
let j = noiseSize - 1;
|
||
|
||
for ( let num = 1; num <= noiseSquareSize; ) {
|
||
|
||
if ( i === - 1 && j === noiseSize ) {
|
||
|
||
j = noiseSize - 2;
|
||
i = 0;
|
||
|
||
} else {
|
||
|
||
if ( j === noiseSize ) {
|
||
|
||
j = 0;
|
||
|
||
}
|
||
|
||
if ( i < 0 ) {
|
||
|
||
i = noiseSize - 1;
|
||
|
||
}
|
||
|
||
}
|
||
|
||
if ( magicSquare[ i * noiseSize + j ] !== 0 ) {
|
||
|
||
j -= 2;
|
||
i ++;
|
||
continue;
|
||
|
||
} else {
|
||
|
||
magicSquare[ i * noiseSize + j ] = num ++;
|
||
|
||
}
|
||
|
||
j ++;
|
||
i --;
|
||
|
||
}
|
||
|
||
return magicSquare;
|
||
|
||
}
|
||
|
||
/**
|
||
* TSL function for creating a Ground Truth Ambient Occlusion (GTAO) effect.
|
||
*
|
||
* @tsl
|
||
* @function
|
||
* @param {Node<float>} depthNode - A node that represents the scene's depth.
|
||
* @param {?Node<vec3>} normalNode - A node that represents the scene's normals.
|
||
* @param {Camera} camera - The camera the scene is rendered with.
|
||
* @returns {GTAONode}
|
||
*/
|
||
export const ao = ( depthNode, normalNode, camera ) => new GTAONode( nodeObject( depthNode ), nodeObject( normalNode ), camera );
|