Skip to main content

bevy_post_process/auto_exposure/
mod.rs

1use bevy_app::prelude::*;
2use bevy_asset::{embedded_asset, AssetApp, Assets, Handle};
3use bevy_ecs::prelude::*;
4use bevy_render::{
5    diagnostic::RecordDiagnostics,
6    extract_component::ExtractComponentPlugin,
7    globals::GlobalsBuffer,
8    render_asset::{RenderAssetPlugin, RenderAssets},
9    render_resource::{
10        BindGroupEntries, Buffer, BufferBinding, BufferDescriptor, BufferUsages,
11        ComputePassDescriptor, PipelineCache, ShaderType, SpecializedComputePipelines,
12    },
13    renderer::{RenderContext, RenderDevice, ViewQuery},
14    texture::{FallbackImage, GpuImage},
15    view::{ExtractedView, ViewTarget, ViewUniform, ViewUniformOffset, ViewUniforms},
16    ExtractSchedule, GpuResourceAppExt, Render, RenderApp, RenderStartup, RenderSystems,
17};
18
19mod buffers;
20mod compensation_curve;
21mod pipeline;
22mod settings;
23
24use buffers::{extract_buffers, prepare_buffers, AutoExposureBuffers};
25pub use compensation_curve::{AutoExposureCompensationCurve, AutoExposureCompensationCurveError};
26use pipeline::{AutoExposurePass, AutoExposurePipeline, ViewAutoExposurePipeline};
27pub use settings::AutoExposure;
28
29use crate::auto_exposure::{
30    compensation_curve::GpuAutoExposureCompensationCurve, pipeline::init_auto_exposure_pipeline,
31};
32use bevy_core_pipeline::{
33    schedule::{Core3d, Core3dSystems},
34    tonemapping::tonemapping,
35};
36
37/// Plugin for the auto exposure feature.
38///
39/// See [`AutoExposure`] for more details.
40pub struct AutoExposurePlugin;
41
42#[derive(Resource)]
43struct AutoExposureResources {
44    histogram: Buffer,
45}
46
47impl Plugin for AutoExposurePlugin {
48    fn build(&self, app: &mut App) {
49        embedded_asset!(app, "auto_exposure.wgsl");
50
51        app.add_plugins(RenderAssetPlugin::<GpuAutoExposureCompensationCurve>::default())
52            .init_asset::<AutoExposureCompensationCurve>()
53            .register_asset_reflect::<AutoExposureCompensationCurve>();
54        app.world_mut()
55            .resource_mut::<Assets<AutoExposureCompensationCurve>>()
56            .insert(&Handle::default(), AutoExposureCompensationCurve::default())
57            .unwrap();
58
59        app.add_plugins(ExtractComponentPlugin::<AutoExposure>::default());
60
61        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
62            return;
63        };
64
65        render_app
66            .init_gpu_resource::<SpecializedComputePipelines<AutoExposurePipeline>>()
67            .init_resource::<AutoExposureBuffers>()
68            .add_systems(
69                RenderStartup,
70                (init_auto_exposure_pipeline, init_auto_exposure_resources),
71            )
72            .add_systems(ExtractSchedule, extract_buffers)
73            .add_systems(
74                Render,
75                (
76                    prepare_buffers.in_set(RenderSystems::Prepare),
77                    queue_view_auto_exposure_pipelines.in_set(RenderSystems::Queue),
78                ),
79            )
80            .add_systems(
81                Core3d,
82                auto_exposure
83                    .before(tonemapping)
84                    .in_set(Core3dSystems::PostProcess),
85            );
86    }
87}
88
89pub fn init_auto_exposure_resources(mut commands: Commands, render_device: Res<RenderDevice>) {
90    commands.insert_resource(AutoExposureResources {
91        histogram: render_device.create_buffer(&BufferDescriptor {
92            label: Some("histogram buffer"),
93            size: pipeline::HISTOGRAM_BIN_COUNT * 4,
94            usage: BufferUsages::STORAGE,
95            mapped_at_creation: false,
96        }),
97    });
98}
99
100fn queue_view_auto_exposure_pipelines(
101    mut commands: Commands,
102    pipeline_cache: Res<PipelineCache>,
103    mut compute_pipelines: ResMut<SpecializedComputePipelines<AutoExposurePipeline>>,
104    pipeline: Res<AutoExposurePipeline>,
105    view_targets: Query<(Entity, &AutoExposure)>,
106) {
107    for (entity, auto_exposure) in view_targets.iter() {
108        let histogram_pipeline =
109            compute_pipelines.specialize(&pipeline_cache, &pipeline, AutoExposurePass::Histogram);
110        let average_pipeline =
111            compute_pipelines.specialize(&pipeline_cache, &pipeline, AutoExposurePass::Average);
112
113        commands.entity(entity).insert(ViewAutoExposurePipeline {
114            histogram_pipeline,
115            mean_luminance_pipeline: average_pipeline,
116            compensation_curve: auto_exposure.compensation_curve.clone(),
117            metering_mask: auto_exposure.metering_mask.clone(),
118        });
119    }
120}
121
122fn auto_exposure(
123    view: ViewQuery<(
124        &ViewUniformOffset,
125        &ViewTarget,
126        &ViewAutoExposurePipeline,
127        &ExtractedView,
128    )>,
129    pipeline_cache: Res<PipelineCache>,
130    pipeline: Res<AutoExposurePipeline>,
131    resources: Res<AutoExposureResources>,
132    view_uniforms: Res<ViewUniforms>,
133    globals_buffer: Res<GlobalsBuffer>,
134    auto_exposure_buffers: Res<AutoExposureBuffers>,
135    fallback: Res<FallbackImage>,
136    gpu_images: Res<RenderAssets<GpuImage>>,
137    compensation_curves: Res<RenderAssets<GpuAutoExposureCompensationCurve>>,
138    mut ctx: RenderContext,
139) {
140    let view_entity = view.entity();
141    let (view_uniform_offset, view_target, auto_exposure_pipeline, extracted_view) =
142        view.into_inner();
143
144    let Some(auto_exposure_buffer) = auto_exposure_buffers.buffers.get(&view_entity) else {
145        return;
146    };
147
148    let (Some(histogram_pipeline), Some(average_pipeline)) = (
149        pipeline_cache.get_compute_pipeline(auto_exposure_pipeline.histogram_pipeline),
150        pipeline_cache.get_compute_pipeline(auto_exposure_pipeline.mean_luminance_pipeline),
151    ) else {
152        return;
153    };
154
155    let view_uniforms_buffer = view_uniforms.uniforms.buffer().unwrap();
156    let source = view_target.main_texture_view();
157
158    let mask = gpu_images
159        .get(&auto_exposure_pipeline.metering_mask)
160        .map(|i| &i.texture_view)
161        .unwrap_or(&fallback.d2.texture_view);
162
163    let Some(compensation_curve) =
164        compensation_curves.get(&auto_exposure_pipeline.compensation_curve)
165    else {
166        return;
167    };
168
169    let compute_bind_group = ctx.render_device().create_bind_group(
170        None,
171        &pipeline_cache.get_bind_group_layout(&pipeline.histogram_layout),
172        &BindGroupEntries::sequential((
173            &globals_buffer.buffer,
174            &auto_exposure_buffer.settings,
175            source,
176            mask,
177            &compensation_curve.texture_view,
178            &compensation_curve.extents,
179            resources.histogram.as_entire_buffer_binding(),
180            &auto_exposure_buffer.state,
181            BufferBinding {
182                buffer: view_uniforms_buffer,
183                size: Some(ViewUniform::min_size()),
184                offset: 0,
185            },
186        )),
187    );
188
189    let diagnostics = ctx.diagnostic_recorder();
190    let diagnostics = diagnostics.as_deref();
191    let time_span = diagnostics.time_span(ctx.command_encoder(), "auto_exposure");
192
193    {
194        let mut compute_pass = ctx
195            .command_encoder()
196            .begin_compute_pass(&ComputePassDescriptor {
197                label: Some("auto_exposure"),
198                timestamp_writes: None,
199            });
200
201        compute_pass.set_bind_group(0, &compute_bind_group, &[view_uniform_offset.offset]);
202        compute_pass.set_pipeline(histogram_pipeline);
203        compute_pass.dispatch_workgroups(
204            extracted_view.viewport.z.div_ceil(16),
205            extracted_view.viewport.w.div_ceil(16),
206            1,
207        );
208        compute_pass.set_pipeline(average_pipeline);
209        compute_pass.dispatch_workgroups(1, 1, 1);
210    }
211
212    time_span.end(ctx.command_encoder());
213}