Improve TAA

crates/bevy_core_pipeline/src/taa/taa.wgsl

@@ -4,14 +4,6 @@
 // http://leiy.cc/publications/TAA/TAA_EG2020_Talk.pdf
 // https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING
 
-// Controls how much to blend between the current and past samples
-// Lower numbers = less of the current sample and more of the past sample = more smoothing
-// Values chosen empirically
-const DEFAULT_HISTORY_BLEND_RATE: f32 = 0.1; // Default blend rate to use when no confidence in history
-const MIN_HISTORY_BLEND_RATE: f32 = 0.015; // Minimum blend rate allowed, to ensure at least some of the current sample is used
-
-#import bevy_core_pipeline::fullscreen_vertex_shader
-
 @group(0) @binding(0) var view_target: texture_2d<f32>;
 @group(0) @binding(1) var history: texture_2d<f32>;
 @group(0) @binding(2) var motion_vectors: texture_2d<f32>;
@@ -65,43 +57,31 @@ fn clip_towards_aabb_center(history_color: vec3<f32>, current_color: vec3<f32>,
 }
 
 fn sample_history(u: f32, v: f32) -> vec3<f32> {
-    return textureSample(history, linear_sampler, vec2(u, v)).rgb;
+    return textureSampleLevel(history, linear_sampler, vec2(u, v), 0.0).rgb;
 }
 
 fn sample_view_target(uv: vec2<f32>) -> vec3<f32> {
-    var sample = textureSample(view_target, nearest_sampler, uv).rgb;
+    var sample = textureSampleLevel(view_target, nearest_sampler, uv, 0.0).rgb;
 #ifdef TONEMAP
     sample = tonemap(sample);
 #endif
     return RGB_to_YCoCg(sample);
 }
 
-@fragment
-fn taa(@location(0) uv: vec2<f32>) -> Output {
-    let texture_size = vec2<f32>(textureDimensions(view_target));
-    let texel_size = 1.0 / texture_size;
-
-    // Fetch the current sample
-    let original_color = textureSample(view_target, nearest_sampler, uv);
-    var current_color = original_color.rgb;
-#ifdef TONEMAP
-    current_color = tonemap(current_color);
-#endif
-
-#ifndef RESET
-    // Pick the closest motion_vector from 5 samples (reduces aliasing on the edges of moving entities)
-    // https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 27
+// Pick the closest motion_vector from 5 samples (reduces aliasing on the edges of moving entities)
+// https://advances.realtimerendering.com/s2014/index.html#_HIGH-QUALITY_TEMPORAL_SUPERSAMPLING, slide 27
+fn closest_motion_vector(uv: vec2<f32>, texel_size: vec2<f32>) -> vec2<f32> {
     let offset = texel_size * 2.0;
     let d_uv_tl = uv + vec2(-offset.x, offset.y);
     let d_uv_tr = uv + vec2(offset.x, offset.y);
     let d_uv_bl = uv + vec2(-offset.x, -offset.y);
     let d_uv_br = uv + vec2(offset.x, -offset.y);
     var closest_uv = uv;
-    let d_tl = textureSample(depth, nearest_sampler, d_uv_tl);
-    let d_tr = textureSample(depth, nearest_sampler, d_uv_tr);
-    var closest_depth = textureSample(depth, nearest_sampler, uv);
-    let d_bl = textureSample(depth, nearest_sampler, d_uv_bl);
-    let d_br = textureSample(depth, nearest_sampler, d_uv_br);
+    let d_tl = textureSampleLevel(depth, nearest_sampler, d_uv_tl, 0.0);
+    let d_tr = textureSampleLevel(depth, nearest_sampler, d_uv_tr, 0.0);
+    var closest_depth = textureSampleLevel(depth, nearest_sampler, uv, 0.0);
+    let d_bl = textureSampleLevel(depth, nearest_sampler, d_uv_bl, 0.0);
+    let d_br = textureSampleLevel(depth, nearest_sampler, d_uv_br, 0.0);
     if d_tl > closest_depth {
         closest_uv = d_uv_tl;
         closest_depth = d_tl;
@@ -117,14 +97,28 @@ fn taa(@location(0) uv: vec2<f32>) -> Output {
     if d_br > closest_depth {
         closest_uv = d_uv_br;
     }
-    let closest_motion_vector = textureSample(motion_vectors, nearest_sampler, closest_uv).rg;
+    return textureSampleLevel(motion_vectors, nearest_sampler, closest_uv, 0.0).rg;
+}
 
+@fragment
+fn taa(@location(0) uv: vec2<f32>) -> Output {
+    let texture_size = vec2<f32>(textureDimensions(view_target));
+    let texel_size = 1.0 / texture_size;
+
+    // Fetch the current sample
+    let original_color = textureSampleLevel(view_target, nearest_sampler, uv, 0.0);
+    var current_color = original_color.rgb;
+#ifdef TONEMAP
+    current_color = tonemap(current_color);
+#endif
+
+#ifndef RESET
     // Reproject to find the equivalent sample from the past
     // Uses 5-sample Catmull-Rom filtering (reduces blurriness)
     // Catmull-Rom filtering: https://gist.github.com/TheRealMJP/c83b8c0f46b63f3a88a5986f4fa982b1
     // Ignoring corners: https://www.activision.com/cdn/research/Dynamic_Temporal_Antialiasing_and_Upsampling_in_Call_of_Duty_v4.pdf#page=68
     // Technically we should renormalize the weights since we're skipping the corners, but it's basically the same result
-    let history_uv = uv - closest_motion_vector;
+    let history_uv = uv - closest_motion_vector(uv, texel_size);
     let sample_position = history_uv * texture_size;
     let texel_center = floor(sample_position - 0.5) + 0.5;
     let f = sample_position - texel_center;
@@ -163,31 +157,24 @@ fn taa(@location(0) uv: vec2<f32>) -> Output {
     history_color = clip_towards_aabb_center(history_color, s_mm, mean - std_deviation, mean + std_deviation);
     history_color = YCoCg_to_RGB(history_color);
 
-    // How confident we are that the history is representative of the current frame
-    var history_confidence = textureSample(history, nearest_sampler, uv).a;
-    let pixel_motion_vector = abs(closest_motion_vector) * texture_size;
-    if pixel_motion_vector.x < 0.01 && pixel_motion_vector.y < 0.01 {
-        // Increment when pixels are not moving
-        history_confidence += 10.0;
-    } else {
-        // Else reset
-        history_confidence = 1.0;
-    }
+    // Use more of the history if we're confident in it (reduces noise)
+    // https://github.com/zombye/spectrum/blob/2f856aedc42703b157181ee5b8aa91c1b0acb621/shaders/program/temporal.glsl#L277-L279
+    var accumulated_samples = textureSampleLevel(history, nearest_sampler, history_uv, 0.0).a;
+    accumulated_samples *= f32(all(saturate(history_uv) == history_uv));
+    let pixel_center_distance = 1.0 - abs(2.0 * fract(sample_position) - 1.0);
+    accumulated_samples *= sqrt(pixel_center_distance.x * pixel_center_distance.y) * 0.25 + 0.75;
+    accumulated_samples = max(accumulated_samples + 1.0, 16.0);
 
     // Blend current and past sample
-    // Use more of the history if we're confident in it (reduces noise when there is no motion)
-    // https://hhoppe.com/supersample.pdf, section 4.1
-    let current_color_factor = clamp(1.0 / history_confidence, MIN_HISTORY_BLEND_RATE, DEFAULT_HISTORY_BLEND_RATE);
-    current_color = mix(history_color, current_color, current_color_factor);
+    current_color = mix(history_color, current_color, 1.0 / accumulated_samples);
 #endif // #ifndef RESET
 
-
     // Write output to history and view target
     var out: Output;
 #ifdef RESET
-    let history_confidence = 1.0 / MIN_HISTORY_BLEND_RATE;
+    let accumulated_samples = 0.0;
 #endif
-    out.history = vec4(current_color, history_confidence);
+    out.history = vec4(current_color, accumulated_samples);
 #ifdef TONEMAP
     current_color = reverse_tonemap(current_color);
 #endif