Skip to content

Commit 39370e3

Browse files
Mugen87pull[bot]
authored andcommitted
Addons: Add TRAAPassNode. (#29636)
* TRAAPass: Initial setup. * TRAAPassNode: Implement clamping, fix velocity. * TRAAPassNode: Finalize initial code. * TRAAPassNode: Clean up. * Examples: Clean up. * TRAAPassNode: Clean up. * TRAAPassNode: Fix dispose(). * TRAANodePass: Refactor MRT setup. * TRAAPassNode: Clean up. * TRAAPassNode: Make MRT handling more robust. * TRAAPassNode: Implement prepass. * Revert "TRAAPassNode: Implement prepass." This reverts commit 172c2ee. * TRAAPassNode: Clean up. * WebGLBackend: Support render target textures in copyTextureToTexture(). * TRAAPassNode: Update comment. * TRAAPassNode: Remove redundant clear().
1 parent 8879dbb commit 39370e3

File tree

8 files changed

+525
-10
lines changed

8 files changed

+525
-10
lines changed

examples/files.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -400,6 +400,7 @@
400400
"webgpu_postprocessing_sobel",
401401
"webgpu_postprocessing_ssaa",
402402
"webgpu_postprocessing_ssr",
403+
"webgpu_postprocessing_traa",
403404
"webgpu_postprocessing_transition",
404405
"webgpu_postprocessing",
405406
"webgpu_procedural_texture",
Lines changed: 347 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,347 @@
1+
import { Color, Vector2, PostProcessingUtils, NearestFilter, Matrix4 } from 'three';
2+
import { add, float, If, Loop, int, Fn, min, max, clamp, nodeObject, PassNode, QuadMesh, texture, NodeMaterial, uniform, uv, vec2, vec4, luminance } from 'three/tsl';
3+
4+
const _quadMesh = /*@__PURE__*/ new QuadMesh();
5+
const _size = /*@__PURE__*/ new Vector2();
6+
7+
let _rendererState;
8+
9+
/**
10+
* Temporal Reprojection Anti-Aliasing (TRAA).
11+
*
12+
* References:
13+
* https://alextardif.com/TAA.html
14+
* https://www.elopezr.com/temporal-aa-and-the-quest-for-the-holy-trail/
15+
*
16+
*/
17+
class TRAAPassNode extends PassNode {
18+
19+
static get type() {
20+
21+
return 'TRAAPassNode';
22+
23+
}
24+
25+
constructor( scene, camera ) {
26+
27+
super( PassNode.COLOR, scene, camera );
28+
29+
this.isTRAAPassNode = true;
30+
31+
this.clearColor = new Color( 0x000000 );
32+
this.clearAlpha = 0;
33+
34+
this._jitterIndex = 0;
35+
this._originalProjectionMatrix = new Matrix4();
36+
37+
// uniforms
38+
39+
this._invSize = uniform( new Vector2() );
40+
41+
// render targets
42+
43+
this._sampleRenderTarget = null;
44+
this._historyRenderTarget = null;
45+
46+
// materials
47+
48+
this._resolveMaterial = new NodeMaterial();
49+
this._resolveMaterial.name = 'TRAA.Resolve';
50+
51+
}
52+
53+
setSize( width, height ) {
54+
55+
super.setSize( width, height );
56+
57+
let needsRestart = false;
58+
59+
if ( this.renderTarget.width !== this._sampleRenderTarget.width || this.renderTarget.height !== this._sampleRenderTarget.height ) {
60+
61+
this._sampleRenderTarget.setSize( this.renderTarget.width, this.renderTarget.height );
62+
this._historyRenderTarget.setSize( this.renderTarget.width, this.renderTarget.height );
63+
64+
this._invSize.value.set( 1 / this.renderTarget.width, 1 / this.renderTarget.height );
65+
66+
needsRestart = true;
67+
68+
}
69+
70+
return needsRestart;
71+
72+
}
73+
74+
updateBefore( frame ) {
75+
76+
const { renderer } = frame;
77+
const { scene, camera } = this;
78+
79+
_rendererState = PostProcessingUtils.resetRendererAndSceneState( renderer, scene, _rendererState );
80+
81+
//
82+
83+
this._pixelRatio = renderer.getPixelRatio();
84+
const size = renderer.getSize( _size );
85+
86+
const needsRestart = this.setSize( size.width, size.height, renderer );
87+
88+
//
89+
90+
this._cameraNear.value = camera.near;
91+
this._cameraFar.value = camera.far;
92+
93+
const viewOffset = {
94+
95+
fullWidth: this.renderTarget.width,
96+
fullHeight: this.renderTarget.height,
97+
offsetX: 0,
98+
offsetY: 0,
99+
width: this.renderTarget.width,
100+
height: this.renderTarget.height
101+
102+
};
103+
104+
const originalViewOffset = Object.assign( {}, camera.view );
105+
106+
if ( originalViewOffset.enabled ) Object.assign( viewOffset, originalViewOffset );
107+
108+
const jitterOffset = _JitterVectors[ this._jitterIndex ];
109+
110+
camera.updateProjectionMatrix();
111+
this._originalProjectionMatrix.copy( camera.projectionMatrix );
112+
113+
camera.setViewOffset(
114+
115+
viewOffset.fullWidth, viewOffset.fullHeight,
116+
117+
viewOffset.offsetX + jitterOffset[ 0 ] * 0.0625, viewOffset.offsetY + jitterOffset[ 1 ] * 0.0625, // 0.0625 = 1 / 16
118+
119+
viewOffset.width, viewOffset.height
120+
121+
);
122+
123+
const mrt = this.getMRT();
124+
const velocityOutput = mrt.get( 'velocity' );
125+
126+
if ( velocityOutput !== undefined ) {
127+
128+
velocityOutput.setProjectionMatrix( this._originalProjectionMatrix );
129+
130+
} else {
131+
132+
throw new Error( 'THREE:TRAAPassNode: Missing velocity output in MRT configuration.' );
133+
134+
}
135+
136+
renderer.setMRT( mrt );
137+
138+
renderer.setClearColor( this.clearColor, this.clearAlpha );
139+
renderer.setRenderTarget( this._sampleRenderTarget );
140+
renderer.render( scene, camera );
141+
142+
renderer.setRenderTarget( null );
143+
renderer.setMRT( null );
144+
145+
// every time when the dimensions change we need fresh history data. Copy the sample
146+
// into the history and final render target (no AA happens at that point).
147+
148+
if ( needsRestart === true ) {
149+
150+
// bind and clear render target to make sure they are initialized after the resize which triggers a dispose()
151+
152+
renderer.setRenderTarget( this._historyRenderTarget );
153+
renderer.clear();
154+
155+
renderer.setRenderTarget( this.renderTarget );
156+
renderer.clear();
157+
158+
renderer.setRenderTarget( null );
159+
160+
renderer.copyTextureToTexture( this._sampleRenderTarget.texture, this._historyRenderTarget.texture );
161+
renderer.copyTextureToTexture( this._sampleRenderTarget.texture, this.renderTarget.texture );
162+
163+
} else {
164+
165+
// resolve
166+
167+
renderer.setRenderTarget( this.renderTarget );
168+
_quadMesh.material = this._resolveMaterial;
169+
_quadMesh.render( renderer );
170+
renderer.setRenderTarget( null );
171+
172+
// update history
173+
174+
renderer.copyTextureToTexture( this.renderTarget.texture, this._historyRenderTarget.texture );
175+
176+
}
177+
178+
// copy depth
179+
180+
renderer.copyTextureToTexture( this._sampleRenderTarget.depthTexture, this.renderTarget.depthTexture );
181+
182+
// update jitter index
183+
184+
this._jitterIndex ++;
185+
this._jitterIndex = this._jitterIndex % ( _JitterVectors.length - 1 );
186+
187+
// restore
188+
189+
if ( originalViewOffset.enabled ) {
190+
191+
camera.setViewOffset(
192+
193+
originalViewOffset.fullWidth, originalViewOffset.fullHeight,
194+
195+
originalViewOffset.offsetX, originalViewOffset.offsetY,
196+
197+
originalViewOffset.width, originalViewOffset.height
198+
199+
);
200+
201+
} else {
202+
203+
camera.clearViewOffset();
204+
205+
}
206+
207+
velocityOutput.setProjectionMatrix( null );
208+
209+
PostProcessingUtils.restoreRendererAndSceneState( renderer, scene, _rendererState );
210+
211+
}
212+
213+
setup( builder ) {
214+
215+
if ( this._sampleRenderTarget === null ) {
216+
217+
this._sampleRenderTarget = this.renderTarget.clone();
218+
this._historyRenderTarget = this.renderTarget.clone();
219+
220+
this._sampleRenderTarget.texture.minFiler = NearestFilter;
221+
this._sampleRenderTarget.texture.magFilter = NearestFilter;
222+
223+
const velocityTarget = this._sampleRenderTarget.texture.clone();
224+
velocityTarget.isRenderTargetTexture = true;
225+
velocityTarget.name = 'velocity';
226+
227+
this._sampleRenderTarget.textures.push( velocityTarget );
228+
229+
}
230+
231+
// textures
232+
233+
const historyTexture = texture( this._historyRenderTarget.texture );
234+
const sampleTexture = texture( this._sampleRenderTarget.textures[ 0 ] );
235+
const velocityTexture = texture( this._sampleRenderTarget.textures[ 1 ] );
236+
const depthTexture = texture( this._sampleRenderTarget.depthTexture );
237+
238+
const resolve = Fn( () => {
239+
240+
const uvNode = uv();
241+
242+
const minColor = vec4( 10000 ).toVar();
243+
const maxColor = vec4( - 10000 ).toVar();
244+
const closestDepth = float( 1 ).toVar();
245+
const closestDepthPixelPosition = vec2( 0 ).toVar();
246+
247+
// sample a 3x3 neighborhood to create a box in color space
248+
// clamping the history color with the resulting min/max colors mitigates ghosting
249+
250+
Loop( { start: int( - 1 ), end: int( 1 ), type: 'int', condition: '<=', name: 'x' }, ( { x } ) => {
251+
252+
Loop( { start: int( - 1 ), end: int( 1 ), type: 'int', condition: '<=', name: 'y' }, ( { y } ) => {
253+
254+
const uvNeighbor = uvNode.add( vec2( float( x ), float( y ) ).mul( this._invSize ) ).toVar();
255+
const colorNeighbor = max( vec4( 0 ), sampleTexture.uv( uvNeighbor ) ).toVar(); // use max() to avoid propagate garbage values
256+
257+
minColor.assign( min( minColor, colorNeighbor ) );
258+
maxColor.assign( max( maxColor, colorNeighbor ) );
259+
260+
const currentDepth = depthTexture.uv( uvNeighbor ).r.toVar();
261+
262+
// find the sample position of the closest depth in the neighborhood (used for velocity)
263+
264+
If( currentDepth.lessThan( closestDepth ), () => {
265+
266+
closestDepth.assign( currentDepth );
267+
closestDepthPixelPosition.assign( uvNeighbor );
268+
269+
} );
270+
271+
} );
272+
273+
} );
274+
275+
// sampling/reprojection
276+
277+
const offset = velocityTexture.uv( closestDepthPixelPosition ).xy.mul( vec2( 0.5, - 0.5 ) ); // NDC to uv offset
278+
279+
const currentColor = sampleTexture.uv( uvNode );
280+
const historyColor = historyTexture.uv( uvNode.sub( offset ) );
281+
282+
// clamping
283+
284+
const clampedHistoryColor = clamp( historyColor, minColor, maxColor );
285+
286+
// flicker reduction based on luminance weighing
287+
288+
const currentWeight = float( 0.05 ).toVar();
289+
const historyWeight = currentWeight.oneMinus().toVar();
290+
291+
const compressedCurrent = currentColor.mul( float( 1 ).div( ( max( max( currentColor.r, currentColor.g ), currentColor.b ).add( 1.0 ) ) ) );
292+
const compressedHistory = clampedHistoryColor.mul( float( 1 ).div( ( max( max( clampedHistoryColor.r, clampedHistoryColor.g ), clampedHistoryColor.b ).add( 1.0 ) ) ) );
293+
294+
const luminanceCurrent = luminance( compressedCurrent.rgb );
295+
const luminanceHistory = luminance( compressedHistory.rgb );
296+
297+
currentWeight.mulAssign( float( 1.0 ).div( luminanceCurrent.add( 1 ) ) );
298+
historyWeight.mulAssign( float( 1.0 ).div( luminanceHistory.add( 1 ) ) );
299+
300+
return add( currentColor.mul( currentWeight ), clampedHistoryColor.mul( historyWeight ) ).div( max( currentWeight.add( historyWeight ), 0.00001 ) );
301+
302+
} );
303+
304+
// materials
305+
306+
this._resolveMaterial.fragmentNode = resolve();
307+
308+
return super.setup( builder );
309+
310+
}
311+
312+
dispose() {
313+
314+
super.dispose();
315+
316+
if ( this._sampleRenderTarget !== null ) {
317+
318+
this._sampleRenderTarget.dispose();
319+
this._historyRenderTarget.dispose();
320+
321+
}
322+
323+
this._resolveMaterial.dispose();
324+
325+
}
326+
327+
}
328+
329+
export default TRAAPassNode;
330+
331+
// These jitter vectors are specified in integers because it is easier.
332+
// I am assuming a [-8,8) integer grid, but it needs to be mapped onto [-0.5,0.5)
333+
// before being used, thus these integers need to be scaled by 1/16.
334+
//
335+
// Sample patterns reference: https://msdn.microsoft.com/en-us/library/windows/desktop/ff476218%28v=vs.85%29.aspx?f=255&MSPPError=-2147217396
336+
const _JitterVectors = [
337+
[ - 4, - 7 ], [ - 7, - 5 ], [ - 3, - 5 ], [ - 5, - 4 ],
338+
[ - 1, - 4 ], [ - 2, - 2 ], [ - 6, - 1 ], [ - 4, 0 ],
339+
[ - 7, 1 ], [ - 1, 2 ], [ - 6, 3 ], [ - 3, 3 ],
340+
[ - 7, 6 ], [ - 3, 6 ], [ - 5, 7 ], [ - 1, 7 ],
341+
[ 5, - 7 ], [ 1, - 6 ], [ 6, - 5 ], [ 4, - 4 ],
342+
[ 2, - 3 ], [ 7, - 2 ], [ 1, - 1 ], [ 4, - 1 ],
343+
[ 2, 1 ], [ 6, 2 ], [ 0, 4 ], [ 4, 4 ],
344+
[ 2, 5 ], [ 7, 5 ], [ 5, 6 ], [ 3, 7 ]
345+
];
346+
347+
export const traaPass = ( scene, camera ) => nodeObject( new TRAAPassNode( scene, camera ) );
18.9 KB
Loading

0 commit comments

Comments
 (0)