Add lens dirt support to Bloom

crates/bevy_post_process/src/bloom/bloom.wgsl

@@ -6,18 +6,50 @@
 // * [COD] - Next Generation Post Processing in Call of Duty - http://www.iryoku.com/next-generation-post-processing-in-call-of-duty-advanced-warfare
 // * [PBB] - Physically Based Bloom - https://learnopengl.com/Guest-Articles/2022/Phys.-Based-Bloom
 
+#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
+
 struct BloomUniforms {
     threshold_precomputations: vec4<f32>,
     viewport: vec4<f32>,
     scale: vec2<f32>,
     aspect: f32,
+    intensity: f32,
+    low_frequency_boost: f32,
+    low_frequency_boost_curvature: f32,
+    high_pass_frequency: f32,
+    // BloomCompositeMode::EnergyConserving = 0, BloomCompositeMode::Additive = 1.
+    composite_mode: u32,
+    max_mip: f32,
+    lens_dirt_intensity: f32,
+    lens_dirt_tint: vec3<f32>,
+    padding: u32,
 };
 
 @group(0) @binding(0) var input_texture: texture_2d<f32>;
 @group(0) @binding(1) var s: sampler;
 
 @group(0) @binding(2) var<uniform> uniforms: BloomUniforms;
 
+#ifdef LENS_DIRT
+@group(0) @binding(3) var dirt_texture: texture_2d<f32>;
+@group(0) @binding(4) var dirt_sampler: sampler;
+
+// Shader version of `compute_blend_factor`.
+fn compute_blend_factor_gpu(mip: f32, max_mip: f32) -> f32 {
+    var lf_boost = (1.0 - pow(
+        1.0 - (mip / max_mip),
+        1.0 / (1.0 - uniforms.low_frequency_boost_curvature)
+        )) * uniforms.low_frequency_boost;
+    var high_pass_lq = 1.0 - clamp(
+        ((mip / max_mip) - uniforms.high_pass_frequency) / uniforms.high_pass_frequency,
+        0.0, 1.0);
+    if uniforms.composite_mode == 0u {
+        lf_boost *= 1.0 - uniforms.intensity;
+    }
+    return (uniforms.intensity + lf_boost) * high_pass_lq;
+}
+#endif
+
 #ifdef FIRST_DOWNSAMPLE
 // https://catlikecoding.com/unity/tutorials/advanced-rendering/bloom/#3.4
 fn soft_threshold(color: vec3<f32>) -> vec3<f32> {
@@ -155,9 +187,10 @@ fn sample_input_3x3_tent(uv: vec2<f32>) -> vec3<f32> {
 
 #ifdef FIRST_DOWNSAMPLE
 @fragment
-fn downsample_first(@location(0) output_uv: vec2<f32>) -> @location(0) vec4<f32> {
-    let sample_uv = uniforms.viewport.xy + output_uv * uniforms.viewport.zw;
+fn downsample_first(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    let sample_uv = uniforms.viewport.xy + in.uv * uniforms.viewport.zw;
     var sample = sample_input_13_tap(sample_uv);
+
     // Lower bound of 0.0001 is to avoid propagating multiplying by 0.0 through the
     // downscaling and upscaling which would result in black boxes.
     // The upper bound is to prevent NaNs.
@@ -172,12 +205,29 @@ fn downsample_first(@location(0) output_uv: vec2<f32>) -> @location(0) vec4<f32>
 }
 #endif
 
+#ifdef LENS_DIRT
+@fragment
+fn upsample_final(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    let bloom = sample_input_3x3_tent(in.uv);
+    let bloom_intensity = max(bloom.r, max(bloom.g, bloom.b));
+    let base_alpha = compute_blend_factor_gpu(0.0, uniforms.max_mip);
+
+    let dirt = textureSample(dirt_texture, dirt_sampler, in.uv).r;
+    let amount = clamp(dirt * uniforms.lens_dirt_intensity * bloom_intensity, 0.0, 1.0);
+
+    let result_bloom = mix(bloom.rgb, bloom.rgb * uniforms.lens_dirt_tint.rgb, amount);
+    let alpha = mix(base_alpha, 1.0, amount);
+
+    return vec4<f32>(result_bloom, alpha);
+}
+#endif
+
 @fragment
-fn downsample(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
-    return vec4<f32>(sample_input_13_tap(uv), 1.0);
+fn downsample(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    return vec4<f32>(sample_input_13_tap(in.uv), 1.0);
 }
 
 @fragment
-fn upsample(@location(0) uv: vec2<f32>) -> @location(0) vec4<f32> {
-    return vec4<f32>(sample_input_3x3_tent(uv), 1.0);
+fn upsample(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
+    return vec4<f32>(sample_input_3x3_tent(in.uv), 1.0);
 }

crates/bevy_post_process/src/bloom/downsampling_pipeline.rs

@@ -1,13 +1,13 @@
-use bevy_core_pipeline::FullscreenShader;
+use super::{settings::BloomUniforms, Bloom, BLOOM_TEXTURE_FORMAT};
 
-use super::{Bloom, BLOOM_TEXTURE_FORMAT};
 use bevy_asset::{load_embedded_asset, AssetServer, Handle};
+use bevy_core_pipeline::FullscreenShader;
 use bevy_ecs::{
     prelude::{Component, Entity},
     resource::Resource,
     system::{Commands, Query, Res, ResMut},
 };
-use bevy_math::{Vec2, Vec4};
+use bevy_math::Vec2;
 use bevy_render::{
     render_resource::{
         binding_types::{sampler, texture_2d, uniform_buffer},
@@ -42,17 +42,6 @@ pub struct BloomDownsamplingPipelineKeys {
     uniform_scale: bool,
 }
 
-/// The uniform struct extracted from [`Bloom`] attached to a Camera.
-/// Will be available for use in the Bloom shader.
-#[derive(Component, ShaderType, Clone)]
-pub struct BloomUniforms {
-    // Precomputed values used when thresholding, see https://catlikecoding.com/unity/tutorials/advanced-rendering/bloom/#3.4
-    pub threshold_precomputations: Vec4,
-    pub viewport: Vec4,
-    pub scale: Vec2,
-    pub aspect: f32,
-}
-
 pub fn init_bloom_downsampling_pipeline(
     mut commands: Commands,
     render_device: Res<RenderDevice>,

crates/bevy_post_process/src/bloom/mod.rs

@@ -2,13 +2,19 @@ mod downsampling_pipeline;
 mod settings;
 mod upsampling_pipeline;
 
-use bevy_image::ToExtents;
-pub use settings::{Bloom, BloomCompositeMode, BloomPrefilter};
-
 use crate::bloom::{
-    downsampling_pipeline::init_bloom_downsampling_pipeline,
-    upsampling_pipeline::init_bloom_upscaling_pipeline,
+    downsampling_pipeline::{
+        init_bloom_downsampling_pipeline, prepare_downsampling_pipeline, BloomDownsamplingPipeline,
+        BloomDownsamplingPipelineIds,
+    },
+    settings::BloomUniforms,
+    upsampling_pipeline::{
+        init_bloom_upscaling_pipeline, prepare_upsampling_pipeline, BloomUpsamplingPipeline,
+        UpsamplingPipelineIds,
+    },
 };
+pub use settings::{Bloom, BloomCompositeMode, BloomPrefilter, LensDirt};
+
 use bevy_app::{App, Plugin};
 use bevy_asset::embedded_asset;
 use bevy_color::{Gray, LinearRgba};
@@ -17,26 +23,21 @@ use bevy_core_pipeline::{
     tonemapping::tonemapping,
 };
 use bevy_ecs::prelude::*;
+use bevy_image::ToExtents;
 use bevy_math::{ops, UVec2};
 use bevy_render::{
     camera::ExtractedCamera,
     diagnostic::RecordDiagnostics,
     extract_component::{
         ComponentUniforms, DynamicUniformIndex, ExtractComponentPlugin, UniformComponentPlugin,
     },
+    render_asset::RenderAssets,
     render_resource::*,
     renderer::{RenderContext, RenderDevice, ViewQuery},
-    texture::{CachedTexture, TextureCache},
+    texture::{CachedTexture, FallbackImage, GpuImage, TextureCache},
     view::ViewTarget,
     GpuResourceAppExt, Render, RenderApp, RenderStartup, RenderSystems,
 };
-use downsampling_pipeline::{
-    prepare_downsampling_pipeline, BloomDownsamplingPipeline, BloomDownsamplingPipelineIds,
-    BloomUniforms,
-};
-use upsampling_pipeline::{
-    prepare_upsampling_pipeline, BloomUpsamplingPipeline, UpsamplingPipelineIds,
-};
 
 const BLOOM_TEXTURE_FORMAT: TextureFormat = TextureFormat::Rg11b10Ufloat;
 
@@ -97,8 +98,11 @@ pub fn bloom(
         &BloomDownsamplingPipelineIds,
     )>,
     downsampling_pipeline_res: Res<BloomDownsamplingPipeline>,
+    upsampling_pipeline_res: Res<BloomUpsamplingPipeline>,
     pipeline_cache: Res<PipelineCache>,
     uniforms: Res<ComponentUniforms<BloomUniforms>>,
+    gpu_images: Res<RenderAssets<GpuImage>>,
+    fallback: Res<FallbackImage>,
     mut ctx: RenderContext,
 ) {
     let (
@@ -147,6 +151,35 @@ pub fn bloom(
         )),
     );
 
+    // Use dirt pipeline if ready, otherwise standard final pipeline.
+    let (final_pipeline, use_dirt) = match upsampling_pipeline_ids.id_final_dirt {
+        Some(dirt_id) => match pipeline_cache.get_render_pipeline(dirt_id) {
+            Some(dirt_pipeline) => (dirt_pipeline, true),
+            None => (upsampling_final_pipeline, false),
+        },
+        None => (upsampling_final_pipeline, false),
+    };
+
+    // Choose the appropriate final bind group.
+    let final_bind_group = match (use_dirt, &bloom_settings.lens_dirt.texture) {
+        (true, Some(lens_dirt_texture)) => {
+            let dirt_image = gpu_images.get(lens_dirt_texture).unwrap_or(&fallback.d2);
+            ctx.render_device().create_bind_group(
+                "bloom_final_dirt_bind_group",
+                &pipeline_cache
+                    .get_bind_group_layout(&upsampling_pipeline_res.dirt_bind_group_layout),
+                &BindGroupEntries::sequential((
+                    &bloom_texture.view(0),
+                    &bind_groups.sampler,
+                    uniforms_binding.clone(),
+                    &dirt_image.texture_view,
+                    &bind_groups.sampler,
+                )),
+            )
+        }
+        _ => bind_groups.upsampling_bind_groups[(bloom_texture.mip_count - 1) as usize].clone(),
+    };
+
     let diagnostics = ctx.diagnostic_recorder();
     let diagnostics = diagnostics.as_deref();
     let time_span = diagnostics.time_span(ctx.command_encoder(), "bloom");
@@ -249,12 +282,8 @@ pub fn bloom(
             occlusion_query_set: None,
             multiview_mask: None,
         });
-        upsampling_final_pass.set_pipeline(upsampling_final_pipeline);
-        upsampling_final_pass.set_bind_group(
-            0,
-            &bind_groups.upsampling_bind_groups[(bloom_texture.mip_count - 1) as usize],
-            &[uniform_index.index()],
-        );
+        upsampling_final_pass.set_pipeline(final_pipeline);
+        upsampling_final_pass.set_bind_group(0, &final_bind_group, &[uniform_index.index()]);
         if let Some(viewport) = camera.viewport.as_ref() {
             upsampling_final_pass.set_viewport(
                 viewport.physical_position.x as f32,
@@ -265,8 +294,12 @@ pub fn bloom(
                 viewport.depth.end,
             );
         }
-        let blend = compute_blend_factor(bloom_settings, 0.0, (bloom_texture.mip_count - 1) as f32);
-        upsampling_final_pass.set_blend_constant(LinearRgba::gray(blend).into());
+
+        if !use_dirt {
+            let blend =
+                compute_blend_factor(bloom_settings, 0.0, (bloom_texture.mip_count - 1) as f32);
+            upsampling_final_pass.set_blend_constant(LinearRgba::gray(blend).into());
+        }
         upsampling_final_pass.draw(0..3, 0..1);
     }