@animech-public/playcanvas/build/playcanvas.prf/src/extras/render-passes/render-pass-taa.js

PlayCanvas WebGL game engine

import { FILTER_LINEAR, ADDRESS_CLAMP_TO_EDGE } from '../../platform/graphics/constants.js';
import { Texture } from '../../platform/graphics/texture.js';
import { shaderChunks } from '../../scene/shader-lib/chunks/chunks.js';
import { RenderPassShaderQuad } from '../../scene/graphics/render-pass-shader-quad.js';
import { RenderTarget } from '../../platform/graphics/render-target.js';

const fs = `
	uniform highp sampler2D uSceneDepthMap;
	uniform sampler2D sourceTexture;
	uniform sampler2D historyTexture;
	uniform mat4 matrix_viewProjectionPrevious;
	uniform mat4 matrix_viewProjectionInverse;
	uniform vec4 jitters;
	uniform vec2 textureSize;
	varying vec2 uv0;
	vec2 reproject(vec2 uv, float depth) {
		#ifndef WEBGPU
			depth = depth * 2.0 - 1.0;
		#endif
		vec4 ndc = vec4(uv * 2.0 - 1.0, depth, 1.0);
		ndc.xy -= jitters.xy;
		vec4 worldPosition = matrix_viewProjectionInverse * ndc;
		worldPosition /= worldPosition.w;
	
		vec4 screenPrevious = matrix_viewProjectionPrevious * worldPosition;
		return (screenPrevious.xy / screenPrevious.w) * 0.5 + 0.5;
	}
	vec4 colorClamp(vec2 uv, vec4 historyColor) {
		vec3 minColor = vec3(9999.0);
		vec3 maxColor = vec3(-9999.0);
 
		for(float x = -1.0; x <= 1.0; ++x)
		{
			for(float y = -1.0; y <= 1.0; ++y)
			{
				vec3 color = texture2D(sourceTexture, uv + vec2(x, y) / textureSize).rgb;
				minColor = min(minColor, color);
				maxColor = max(maxColor, color);
			}
		}
 
		vec3 clamped = clamp(historyColor.rgb, minColor, maxColor);
		return vec4(clamped, historyColor.a);
	}
	void main()
	{
		vec2 uv = uv0;
		#ifdef WEBGPU
			uv.y = 1.0 - uv.y;
		#endif
		vec4 srcColor = texture2D(sourceTexture, uv);
		float depth = texture2DLodEXT(uSceneDepthMap, uv, 0.0).r;
		vec2 historyUv = reproject(uv0, depth);
		#ifdef QUALITY_HIGH
			vec4 historyColor = SampleTextureCatmullRom(TEXTURE_PASS(historyTexture), historyUv, textureSize);
		#else
			vec4 historyColor = texture2D(historyTexture, historyUv);
		#endif
		vec4 historyColorClamped = colorClamp(uv, historyColor);
		float mixFactor = (historyUv.x < 0.0 || historyUv.x > 1.0 || historyUv.y < 0.0 || historyUv.y > 1.0) ?
			1.0 : 0.05;
		gl_FragColor = mix(historyColorClamped, srcColor, mixFactor);
	}
`;
class RenderPassTAA extends RenderPassShaderQuad {
	constructor(device, sourceTexture, cameraComponent) {
		super(device);
		this.historyIndex = 0;
		this.historyTexture = null;
		this.historyTextures = [];
		this.historyRenderTargets = [];
		this.sourceTexture = sourceTexture;
		this.cameraComponent = cameraComponent;
		const defines = `
			#define QUALITY_HIGH
		`;
		const fsChunks = shaderChunks.sampleCatmullRomPS;
		this.shader = this.createQuadShader('TaaResolveShader', defines + fsChunks + fs);
		const {
			scope
		} = device;
		this.sourceTextureId = scope.resolve('sourceTexture');
		this.textureSizeId = scope.resolve('textureSize');
		this.textureSize = new Float32Array(2);
		this.historyTextureId = scope.resolve('historyTexture');
		this.viewProjPrevId = scope.resolve('matrix_viewProjectionPrevious');
		this.viewProjInvId = scope.resolve('matrix_viewProjectionInverse');
		this.jittersId = scope.resolve('jitters');
		this.setup();
	}
	destroy() {
		if (this.renderTarget) {
			this.renderTarget.destroyTextureBuffers();
			this.renderTarget.destroy();
			this.renderTarget = null;
		}
	}
	setup() {
		for (let i = 0; i < 2; ++i) {
			this.historyTextures[i] = new Texture(this.device, {
				name: `TAA-History-${i}`,
				width: 4,
				height: 4,
				format: this.sourceTexture.format,
				mipmaps: false,
				minFilter: FILTER_LINEAR,
				magFilter: FILTER_LINEAR,
				addressU: ADDRESS_CLAMP_TO_EDGE,
				addressV: ADDRESS_CLAMP_TO_EDGE
			});
			this.historyRenderTargets[i] = new RenderTarget({
				colorBuffer: this.historyTextures[i],
				depth: false
			});
		}
		this.historyTexture = this.historyTextures[0];
		this.init(this.historyRenderTargets[0], {
			resizeSource: this.sourceTexture
		});
	}
	before() {
		this.sourceTextureId.setValue(this.sourceTexture);
		this.historyTextureId.setValue(this.historyTextures[1 - this.historyIndex]);
		this.textureSize[0] = this.sourceTexture.width;
		this.textureSize[1] = this.sourceTexture.height;
		this.textureSizeId.setValue(this.textureSize);
		const camera = this.cameraComponent.camera;
		this.viewProjPrevId.setValue(camera._viewProjPrevious.data);
		this.viewProjInvId.setValue(camera._viewProjInverse.data);
		this.jittersId.setValue(camera._jitters);
	}
	update() {
		this.historyIndex = 1 - this.historyIndex;
		this.historyTexture = this.historyTextures[this.historyIndex];
		this.renderTarget = this.historyRenderTargets[this.historyIndex];
		return this.historyTexture;
	}
}

export { RenderPassTAA };