bevy_render/view/window/
screenshot.rs

1use crate::{
2    gpu_readback,
3    render_asset::RenderAssets,
4    render_resource::{
5        BindGroup, BindGroupEntries, Buffer, BufferUsages, PipelineCache,
6        SpecializedRenderPipeline, SpecializedRenderPipelines, Texture, TextureUsages, TextureView,
7    },
8    renderer::RenderDevice,
9    sync_world::MainEntity,
10    texture::{GpuImage, ManualTextureViews, OutputColorAttachment},
11    view::{
12        prepare_view_attachments, prepare_view_targets, ExtractedWindow, SurfaceData,
13        ViewTargetAttachments,
14    },
15    ExtractSchedule, GpuResourceAppExt, MainWorld, Render, RenderApp, RenderStartup, RenderSystems,
16};
17use alloc::{borrow::Cow, sync::Arc};
18use bevy_app::{First, Plugin, Update};
19use bevy_asset::{embedded_asset, load_embedded_asset, AssetServer, Handle, RenderAssetUsages};
20use bevy_camera::{ManualTextureViewHandle, NormalizedRenderTarget, RenderTarget};
21use bevy_color::{ColorToPacked, LinearRgba, Srgba};
22use bevy_derive::{Deref, DerefMut};
23use bevy_ecs::{
24    entity::EntityHashMap, message::message_update_system, prelude::*, system::SystemState,
25};
26use bevy_image::{Image, TextureFormatPixelInfo, ToExtents};
27use bevy_log::{error, info, warn};
28use bevy_material::{
29    bind_group_layout_entries::{binding_types::texture_2d, BindGroupLayoutEntries},
30    descriptor::{
31        BindGroupLayoutDescriptor, CachedRenderPipelineId, FragmentState, RenderPipelineDescriptor,
32        VertexState,
33    },
34};
35use bevy_platform::collections::HashSet;
36use bevy_reflect::Reflect;
37use bevy_shader::Shader;
38use bevy_tasks::AsyncComputeTaskPool;
39use bevy_utils::default;
40use bevy_window::{PrimaryWindow, WindowRef};
41use core::ops::Deref;
42use std::{
43    path::Path,
44    sync::{
45        mpsc::{Receiver, Sender},
46        Mutex,
47    },
48};
49use wgpu::{CommandEncoder, Extent3d, TextureFormat};
50
51#[derive(EntityEvent, Reflect, Deref, DerefMut, Debug)]
52#[reflect(Debug, Event)]
53pub struct ScreenshotCaptured {
54    pub entity: Entity,
55    #[deref]
56    pub image: Image,
57}
58
59/// A component that signals to the renderer to capture a screenshot this frame.
60///
61/// This component should be spawned on a new entity with an observer that will trigger
62/// with [`ScreenshotCaptured`] when the screenshot is ready.
63///
64/// Screenshots are captured asynchronously and may not be available immediately after the frame
65/// that the component is spawned on. The observer should be used to handle the screenshot when it
66/// is ready.
67///
68/// Note that the screenshot entity will be despawned after the screenshot is captured and the
69/// observer is triggered.
70///
71/// # Usage
72///
73/// ```
74/// # use bevy_ecs::prelude::*;
75/// # use bevy_render::view::screenshot::{save_to_disk, Screenshot};
76///
77/// fn take_screenshot(mut commands: Commands) {
78///    commands.spawn(Screenshot::primary_window())
79///       .observe(save_to_disk("screenshot.png"));
80/// }
81/// ```
82#[derive(Component, Deref, DerefMut, Reflect, Debug)]
83#[reflect(Component, Debug)]
84pub struct Screenshot(pub RenderTarget);
85
86/// A marker component that indicates that a screenshot is currently being captured.
87#[derive(Component, Default)]
88pub struct Capturing;
89
90/// A marker component that indicates that a screenshot has been captured, the image is ready, and
91/// the screenshot entity can be despawned.
92#[derive(Component, Default)]
93pub struct Captured;
94
95impl Screenshot {
96    /// Capture a screenshot of the provided window entity.
97    pub fn window(window: Entity) -> Self {
98        Self(RenderTarget::Window(WindowRef::Entity(window)))
99    }
100
101    /// Capture a screenshot of the primary window, if one exists.
102    pub fn primary_window() -> Self {
103        Self(RenderTarget::Window(WindowRef::Primary))
104    }
105
106    /// Capture a screenshot of the provided render target image.
107    pub fn image(image: Handle<Image>) -> Self {
108        Self(RenderTarget::Image(image.into()))
109    }
110
111    /// Capture a screenshot of the provided manual texture view.
112    pub fn texture_view(texture_view: ManualTextureViewHandle) -> Self {
113        Self(RenderTarget::TextureView(texture_view))
114    }
115}
116
117struct ScreenshotPreparedState {
118    pub texture: Texture,
119    pub buffer: Buffer,
120    pub bind_group: BindGroup,
121    pub pipeline_id: CachedRenderPipelineId,
122    pub size: Extent3d,
123}
124
125#[derive(Resource, Deref, DerefMut)]
126pub struct CapturedScreenshots(pub Arc<Mutex<Receiver<(Entity, Image)>>>);
127
128#[derive(Resource, Deref, DerefMut, Default)]
129struct RenderScreenshotTargets(EntityHashMap<NormalizedRenderTarget>);
130
131#[derive(Resource, Deref, DerefMut, Default)]
132struct RenderScreenshotsPrepared(EntityHashMap<ScreenshotPreparedState>);
133
134#[derive(Resource, Deref, DerefMut)]
135struct RenderScreenshotsSender(Sender<(Entity, Image)>);
136
137/// Saves the captured screenshot to disk at the provided path.
138///
139/// A screenshot of an 8-bit render target saves to any format `image`
140/// supports. A screenshot of an `Rgba16Float` or `Rgba32Float` render target,
141/// for example an `Hdr` camera rendering to an `Image`, holds values outside
142/// `0..=1`. Those are kept in `.exr` and `.hdr`. Any other format clips to
143/// `0..=1` and converts to 8-bit sRGB, with a warning. `.exr` needs the `exr`
144/// cargo feature.
145pub fn save_to_disk(path: impl AsRef<Path>) -> impl FnMut(On<ScreenshotCaptured>) {
146    let path = path.as_ref().to_owned();
147    move |screenshot_captured| {
148        use image::{DynamicImage, ImageFormat};
149
150        let img = screenshot_captured.image.clone();
151        match img.try_into_dynamic() {
152            Ok(dyn_img) => match ImageFormat::from_path(&path) {
153                Ok(format) => {
154                    let is_float = matches!(
155                        dyn_img,
156                        DynamicImage::ImageRgb32F(_) | DynamicImage::ImageRgba32F(_)
157                    );
158                    let img = if is_float {
159                        match format {
160                            ImageFormat::OpenExr => dyn_img,
161                            ImageFormat::Hdr => {
162                                let mut rgb = dyn_img.into_rgb32f();
163                                // `image`'s Radiance encoder stores one shared exponent byte per
164                                // pixel, computed as `floor(log2(max)) + 1` in `i32` and stored as
165                                // `(exp + 128) as u8`, with no range checks. An infinite channel
166                                // overflows the `i32` add, and a channel at or above 2^127 or below
167                                // 2^-128 wraps the exponent byte. Clamp to `0.0..=2^126`, and
168                                // replace NaN and values below `f32::MIN_POSITIVE` with zero, so
169                                // the encoder accepts every value.
170                                let max = bevy_math::ops::exp2(126.0);
171                                for value in rgb.iter_mut() {
172                                    *value = value.clamp(0.0, max);
173                                    if !value.is_normal() {
174                                        *value = 0.0;
175                                    }
176                                }
177                                DynamicImage::ImageRgb32F(rgb)
178                            }
179                            _ => {
180                                warn!(
181                                    "Saving a floating point screenshot as 8-bit RGB. \
182                                    Values above 1.0 are clipped. \
183                                    Save to an .exr or .hdr path to keep them."
184                                );
185                                let rgb = dyn_img.into_rgb32f();
186                                let encoded =
187                                    image::RgbImage::from_fn(rgb.width(), rgb.height(), |x, y| {
188                                        let [r, g, b] = rgb.get_pixel(x, y).0;
189                                        image::Rgb(
190                                            Srgba::from(LinearRgba::rgb(r, g, b))
191                                                .to_u8_array_no_alpha(),
192                                        )
193                                    });
194                                DynamicImage::ImageRgb8(encoded)
195                            }
196                        }
197                    } else {
198                        // discard the alpha channel which stores brightness values when HDR is enabled to make sure
199                        // the screenshot looks right
200                        DynamicImage::ImageRgb8(dyn_img.to_rgb8())
201                    };
202                    #[cfg(not(target_arch = "wasm32"))]
203                    match img.save_with_format(&path, format) {
204                        Ok(_) => info!("Screenshot saved to {}", path.display()),
205                        Err(e) => error!("Cannot save screenshot, IO error: {e}"),
206                    }
207
208                    #[cfg(target_arch = "wasm32")]
209                    {
210                        let save_screenshot = || {
211                            use image::EncodableLayout;
212                            use wasm_bindgen::{JsCast, JsValue};
213
214                            let mut image_buffer = std::io::Cursor::new(Vec::new());
215                            img.write_to(&mut image_buffer, format)
216                                .map_err(|e| JsValue::from_str(&format!("{e}")))?;
217
218                            let parts = js_sys::Array::of1(
219                                &js_sys::Uint8Array::new_from_slice(
220                                    image_buffer.into_inner().as_bytes(),
221                                )
222                                .into(),
223                            );
224                            let blob = web_sys::Blob::new_with_u8_array_sequence(&parts)?;
225                            let url = web_sys::Url::create_object_url_with_blob(&blob)?;
226                            let window = web_sys::window().unwrap();
227                            let document = window.document().unwrap();
228                            let link = document.create_element("a")?;
229                            link.set_attribute("href", &url)?;
230                            link.set_attribute(
231                                "download",
232                                path.file_name()
233                                    .and_then(|filename| filename.to_str())
234                                    .ok_or_else(|| JsValue::from_str("Invalid filename"))?,
235                            )?;
236                            let html_element = link.dyn_into::<web_sys::HtmlElement>()?;
237                            html_element.click();
238                            web_sys::Url::revoke_object_url(&url)?;
239                            Ok::<(), JsValue>(())
240                        };
241
242                        match (save_screenshot)() {
243                            Ok(_) => info!("Screenshot saved to {}", path.display()),
244                            Err(e) => error!("Cannot save screenshot, error: {e:?}"),
245                        };
246                    }
247                }
248                Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"),
249            },
250            Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"),
251        }
252    }
253}
254
255fn clear_screenshots(mut commands: Commands) {
256    commands.despawn_all::<With<Captured>>();
257}
258
259pub fn trigger_screenshots(
260    mut commands: Commands,
261    captured_screenshots: ResMut<CapturedScreenshots>,
262) {
263    let captured_screenshots = captured_screenshots.lock().unwrap();
264    while let Ok((entity, image)) = captured_screenshots.try_recv() {
265        commands.entity(entity).insert(Captured);
266        commands.trigger(ScreenshotCaptured { image, entity });
267    }
268}
269
270fn extract_screenshots(
271    mut targets: ResMut<RenderScreenshotTargets>,
272    mut main_world: ResMut<MainWorld>,
273    mut system_state: Local<
274        Option<
275            SystemState<(
276                Commands,
277                Query<Entity, With<PrimaryWindow>>,
278                Query<(Entity, &Screenshot), Without<Capturing>>,
279            )>,
280        >,
281    >,
282    mut seen_targets: Local<HashSet<NormalizedRenderTarget>>,
283) {
284    if system_state.is_none() {
285        *system_state = Some(SystemState::new(&mut main_world));
286    }
287    let system_state = system_state.as_mut().unwrap();
288    let (mut commands, primary_window, screenshots) =
289        system_state.get_mut(&mut main_world).unwrap();
290
291    targets.clear();
292    seen_targets.clear();
293
294    let primary_window = primary_window.iter().next();
295
296    for (entity, screenshot) in screenshots.iter() {
297        let render_target = screenshot.0.clone();
298        let Some(render_target) = render_target.normalize(primary_window) else {
299            warn!(
300                "Unknown render target for screenshot, skipping: {:?}",
301                render_target
302            );
303            continue;
304        };
305        if seen_targets.contains(&render_target) {
306            warn!(
307                "Duplicate render target for screenshot, skipping entity {}: {:?}",
308                entity, render_target
309            );
310            // If we don't despawn the entity here, it will be captured again in the next frame
311            commands.entity(entity).despawn();
312            continue;
313        }
314        seen_targets.insert(render_target.clone());
315        targets.insert(entity, render_target);
316        commands.entity(entity).insert(Capturing);
317    }
318
319    system_state.apply(&mut main_world);
320}
321
322fn prepare_screenshots(
323    targets: Res<RenderScreenshotTargets>,
324    mut prepared: ResMut<RenderScreenshotsPrepared>,
325    window_surfaces: Query<(MainEntity, &SurfaceData)>,
326    render_device: Res<RenderDevice>,
327    screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
328    pipeline_cache: Res<PipelineCache>,
329    mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
330    images: Res<RenderAssets<GpuImage>>,
331    manual_texture_views: Res<ManualTextureViews>,
332    mut view_target_attachments: ResMut<ViewTargetAttachments>,
333) {
334    prepared.clear();
335    for (entity, target) in targets.iter() {
336        match target {
337            NormalizedRenderTarget::Window(window) => {
338                let window = window.entity();
339                let Some((_, surface_data)) = window_surfaces.iter().find(|(e, _)| *e == window)
340                else {
341                    warn!("Unknown window for screenshot, skipping: {}", window);
342                    continue;
343                };
344                let view_format = surface_data
345                    .texture_view_format
346                    .unwrap_or(surface_data.configuration.format);
347                let size = Extent3d {
348                    width: surface_data.configuration.width,
349                    height: surface_data.configuration.height,
350                    ..default()
351                };
352                let (texture_view, state) = prepare_screenshot_state(
353                    size,
354                    view_format,
355                    &render_device,
356                    &screenshot_pipeline,
357                    &pipeline_cache,
358                    &mut pipelines,
359                );
360                prepared.insert(*entity, state);
361                view_target_attachments.insert(
362                    target.clone(),
363                    OutputColorAttachment::new(texture_view.clone(), view_format),
364                );
365            }
366            NormalizedRenderTarget::Image(image) => {
367                let Some(gpu_image) = images.get(&image.handle) else {
368                    warn!("Unknown image for screenshot, skipping: {:?}", image);
369                    continue;
370                };
371                let view_format = gpu_image.view_format();
372                let (texture_view, state) = prepare_screenshot_state(
373                    gpu_image.texture_descriptor.size,
374                    view_format,
375                    &render_device,
376                    &screenshot_pipeline,
377                    &pipeline_cache,
378                    &mut pipelines,
379                );
380                prepared.insert(*entity, state);
381                view_target_attachments.insert(
382                    target.clone(),
383                    OutputColorAttachment::new(texture_view.clone(), view_format),
384                );
385            }
386            NormalizedRenderTarget::TextureView(texture_view) => {
387                let Some(manual_texture_view) = manual_texture_views.get(texture_view) else {
388                    warn!(
389                        "Unknown manual texture view for screenshot, skipping: {:?}",
390                        texture_view
391                    );
392                    continue;
393                };
394                let view_format = manual_texture_view.view_format;
395                let size = manual_texture_view.size.to_extents();
396                let (texture_view, state) = prepare_screenshot_state(
397                    size,
398                    view_format,
399                    &render_device,
400                    &screenshot_pipeline,
401                    &pipeline_cache,
402                    &mut pipelines,
403                );
404                prepared.insert(*entity, state);
405                view_target_attachments.insert(
406                    target.clone(),
407                    OutputColorAttachment::new(texture_view.clone(), view_format),
408                );
409            }
410            NormalizedRenderTarget::None { .. } => {
411                // Nothing to screenshot!
412            }
413        }
414    }
415}
416
417fn prepare_screenshot_state(
418    size: Extent3d,
419    format: TextureFormat,
420    render_device: &RenderDevice,
421    pipeline: &ScreenshotToScreenPipeline,
422    pipeline_cache: &PipelineCache,
423    pipelines: &mut SpecializedRenderPipelines<ScreenshotToScreenPipeline>,
424) -> (TextureView, ScreenshotPreparedState) {
425    let texture = render_device.create_texture(&wgpu::TextureDescriptor {
426        label: Some("screenshot-capture-rendertarget"),
427        size,
428        mip_level_count: 1,
429        sample_count: 1,
430        dimension: wgpu::TextureDimension::D2,
431        format,
432        usage: TextureUsages::RENDER_ATTACHMENT
433            | TextureUsages::COPY_SRC
434            | TextureUsages::TEXTURE_BINDING,
435        view_formats: &[],
436    });
437    let texture_view = texture.create_view(&Default::default());
438    let buffer = render_device.create_buffer(&wgpu::BufferDescriptor {
439        label: Some("screenshot-transfer-buffer"),
440        size: gpu_readback::get_aligned_size(size, format.pixel_size().unwrap_or(0) as u32) as u64,
441        usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
442        mapped_at_creation: false,
443    });
444    let bind_group = render_device.create_bind_group(
445        "screenshot-to-screen-bind-group",
446        &pipeline_cache.get_bind_group_layout(&pipeline.bind_group_layout),
447        &BindGroupEntries::single(&texture_view),
448    );
449    let pipeline_id = pipelines.specialize(pipeline_cache, pipeline, format);
450
451    (
452        texture_view,
453        ScreenshotPreparedState {
454            texture,
455            buffer,
456            bind_group,
457            pipeline_id,
458            size,
459        },
460    )
461}
462
463pub struct ScreenshotPlugin;
464
465impl Plugin for ScreenshotPlugin {
466    fn build(&self, app: &mut bevy_app::App) {
467        embedded_asset!(app, "screenshot.wesl");
468
469        let (tx, rx) = std::sync::mpsc::channel();
470        app.register_type::<Screenshot>()
471            .register_type::<ScreenshotCaptured>()
472            .insert_resource(CapturedScreenshots(Arc::new(Mutex::new(rx))))
473            .add_systems(
474                First,
475                clear_screenshots
476                    .after(message_update_system)
477                    .before(ApplyDeferred),
478            )
479            .add_systems(Update, trigger_screenshots);
480
481        let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
482            return;
483        };
484
485        render_app
486            .insert_resource(RenderScreenshotsSender(tx))
487            .init_resource::<RenderScreenshotTargets>()
488            .init_resource::<RenderScreenshotsPrepared>()
489            .init_gpu_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>()
490            .add_systems(RenderStartup, init_screenshot_to_screen_pipeline)
491            .add_systems(ExtractSchedule, extract_screenshots.ambiguous_with_all())
492            .add_systems(
493                Render,
494                prepare_screenshots
495                    .after(prepare_view_attachments)
496                    .before(prepare_view_targets)
497                    .in_set(RenderSystems::PrepareViews),
498            );
499    }
500}
501
502#[derive(Resource)]
503pub struct ScreenshotToScreenPipeline {
504    pub bind_group_layout: BindGroupLayoutDescriptor,
505    pub shader: Handle<Shader>,
506}
507
508pub fn init_screenshot_to_screen_pipeline(mut commands: Commands, asset_server: Res<AssetServer>) {
509    let bind_group_layout = BindGroupLayoutDescriptor::new(
510        "screenshot-to-screen-bgl",
511        &BindGroupLayoutEntries::single(
512            wgpu::ShaderStages::FRAGMENT,
513            texture_2d(wgpu::TextureSampleType::Float { filterable: false }),
514        ),
515    );
516
517    let shader = load_embedded_asset!(asset_server.as_ref(), "screenshot.wesl");
518
519    commands.insert_resource(ScreenshotToScreenPipeline {
520        bind_group_layout,
521        shader,
522    });
523}
524
525impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
526    type Key = TextureFormat;
527
528    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
529        RenderPipelineDescriptor {
530            label: Some(Cow::Borrowed("screenshot-to-screen")),
531            layout: vec![self.bind_group_layout.clone()],
532            vertex: VertexState {
533                shader: self.shader.clone(),
534                ..default()
535            },
536            primitive: wgpu::PrimitiveState {
537                cull_mode: Some(wgpu::Face::Back),
538                ..Default::default()
539            },
540            multisample: Default::default(),
541            fragment: Some(FragmentState {
542                shader: self.shader.clone(),
543                targets: vec![Some(wgpu::ColorTargetState {
544                    format: key,
545                    blend: None,
546                    write_mask: wgpu::ColorWrites::ALL,
547                })],
548                ..default()
549            }),
550            ..default()
551        }
552    }
553}
554
555pub(crate) type SubmitScreenshotCommandsState<'w, 's> =
556    Query<'w, 's, (MainEntity, &'s ExtractedWindow)>;
557
558pub(crate) fn submit_screenshot_commands(
559    world: &World,
560    state: &mut SystemState<SubmitScreenshotCommandsState>,
561    encoder: &mut CommandEncoder,
562) {
563    let targets = world.resource::<RenderScreenshotTargets>();
564    let prepared = world.resource::<RenderScreenshotsPrepared>();
565    let pipelines = world.resource::<PipelineCache>();
566    let gpu_images = world.resource::<RenderAssets<GpuImage>>();
567    let manual_texture_views = world.resource::<ManualTextureViews>();
568
569    let windows = state.get(world).unwrap();
570
571    for (entity, render_target) in targets.iter() {
572        match render_target {
573            NormalizedRenderTarget::Window(window) => {
574                let window = window.entity();
575                let Some((_, window)) = windows.iter().find(|(e, _)| *e == window) else {
576                    continue;
577                };
578                let width = window.physical_width;
579                let height = window.physical_height;
580                let Some(texture_format) = window.swap_chain_texture_view_format else {
581                    continue;
582                };
583                let Some(swap_chain_texture_view) = window.swap_chain_texture_view.as_ref() else {
584                    continue;
585                };
586                render_screenshot(
587                    encoder,
588                    prepared,
589                    pipelines,
590                    entity,
591                    width,
592                    height,
593                    texture_format,
594                    swap_chain_texture_view,
595                );
596            }
597            NormalizedRenderTarget::Image(image) => {
598                let Some(gpu_image) = gpu_images.get(&image.handle) else {
599                    warn!("Unknown image for screenshot, skipping: {:?}", image);
600                    continue;
601                };
602                let width = gpu_image.texture_descriptor.size.width;
603                let height = gpu_image.texture_descriptor.size.height;
604                let texture_format = gpu_image.texture_descriptor.format;
605                let texture_view = gpu_image.texture_view.deref();
606                render_screenshot(
607                    encoder,
608                    prepared,
609                    pipelines,
610                    entity,
611                    width,
612                    height,
613                    texture_format,
614                    texture_view,
615                );
616            }
617            NormalizedRenderTarget::TextureView(texture_view) => {
618                let Some(texture_view) = manual_texture_views.get(texture_view) else {
619                    warn!(
620                        "Unknown manual texture view for screenshot, skipping: {:?}",
621                        texture_view
622                    );
623                    continue;
624                };
625                let width = texture_view.size.x;
626                let height = texture_view.size.y;
627                let texture_format = texture_view.view_format;
628                let texture_view = texture_view.texture_view.deref();
629                render_screenshot(
630                    encoder,
631                    prepared,
632                    pipelines,
633                    entity,
634                    width,
635                    height,
636                    texture_format,
637                    texture_view,
638                );
639            }
640            NormalizedRenderTarget::None { .. } => {
641                // Nothing to screenshot!
642            }
643        };
644    }
645}
646
647fn render_screenshot(
648    encoder: &mut CommandEncoder,
649    prepared: &RenderScreenshotsPrepared,
650    pipelines: &PipelineCache,
651    entity: &Entity,
652    width: u32,
653    height: u32,
654    texture_format: TextureFormat,
655    texture_view: &wgpu::TextureView,
656) {
657    if let Some(prepared_state) = &prepared.get(entity) {
658        let extent = Extent3d {
659            width,
660            height,
661            depth_or_array_layers: 1,
662        };
663        encoder.copy_texture_to_buffer(
664            prepared_state.texture.as_image_copy(),
665            wgpu::TexelCopyBufferInfo {
666                buffer: &prepared_state.buffer,
667                layout: gpu_readback::layout_data(extent, texture_format),
668            },
669            extent,
670        );
671
672        if let Some(pipeline) = pipelines.get_render_pipeline(prepared_state.pipeline_id) {
673            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
674                label: Some("screenshot_to_screen_pass"),
675                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
676                    view: texture_view,
677                    depth_slice: None,
678                    resolve_target: None,
679                    ops: wgpu::Operations {
680                        load: wgpu::LoadOp::Load,
681                        store: wgpu::StoreOp::Store,
682                    },
683                })],
684                depth_stencil_attachment: None,
685                timestamp_writes: None,
686                occlusion_query_set: None,
687                multiview_mask: None,
688            });
689            pass.set_pipeline(pipeline);
690            pass.set_bind_group(0, &prepared_state.bind_group, &[]);
691            pass.draw(0..3, 0..1);
692        }
693    }
694}
695
696pub(crate) fn collect_screenshots(world: &mut World) {
697    #[cfg(feature = "trace")]
698    let _span = bevy_log::info_span!("collect_screenshots").entered();
699
700    let sender = world.resource::<RenderScreenshotsSender>().deref().clone();
701    let prepared = world.resource::<RenderScreenshotsPrepared>();
702
703    for (entity, prepared) in prepared.iter() {
704        let entity = *entity;
705        let sender = sender.clone();
706        let width = prepared.size.width;
707        let height = prepared.size.height;
708        let texture_format = prepared.texture.format();
709        let Ok(pixel_size) = texture_format.pixel_size() else {
710            continue;
711        };
712        let buffer = prepared.buffer.clone();
713
714        let finish = async move {
715            let (tx, rx) = async_channel::bounded(1);
716            let buffer_slice = buffer.slice(..);
717            // The polling for this map call is done every frame when the command queue is submitted.
718            buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
719                if let Err(err) = result {
720                    panic!("{}", err.to_string());
721                }
722                tx.try_send(()).unwrap();
723            });
724            rx.recv().await.unwrap();
725            let data = buffer_slice.get_mapped_range().unwrap();
726            // we immediately move the data to CPU memory to avoid holding the mapped view for long
727            let mut result = Vec::from(&*data);
728            drop(data);
729
730            if result.len() != ((width * height) as usize * pixel_size) {
731                // Our buffer has been padded because we needed to align to a multiple of 256.
732                // We remove this padding here
733                let initial_row_bytes = width as usize * pixel_size;
734                let buffered_row_bytes =
735                    gpu_readback::align_byte_size(width * pixel_size as u32) as usize;
736
737                let mut take_offset = buffered_row_bytes;
738                let mut place_offset = initial_row_bytes;
739                for _ in 1..height {
740                    result.copy_within(take_offset..take_offset + buffered_row_bytes, place_offset);
741                    take_offset += buffered_row_bytes;
742                    place_offset += initial_row_bytes;
743                }
744                result.truncate(initial_row_bytes * height as usize);
745            }
746
747            if let Err(e) = sender.send((
748                entity,
749                Image::new(
750                    Extent3d {
751                        width,
752                        height,
753                        depth_or_array_layers: 1,
754                    },
755                    wgpu::TextureDimension::D2,
756                    result,
757                    texture_format,
758                    RenderAssetUsages::MAIN_WORLD,
759                ),
760            )) {
761                error!("Failed to send screenshot: {}", e);
762            }
763        };
764
765        AsyncComputeTaskPool::get().spawn(finish).detach();
766    }
767}
bevy_render/view/window/screenshot.rs

bevy_render/view/window/
screenshot.rs
