bevy_render/view/window/
screenshot.rs

1use std::{borrow::Cow, path::Path, sync::PoisonError};
2
3use bevy_app::Plugin;
4use bevy_asset::{load_internal_asset, Handle};
5use bevy_ecs::{entity::EntityHashMap, prelude::*};
6use bevy_tasks::AsyncComputeTaskPool;
7use bevy_utils::tracing::{error, info, info_span};
8use std::sync::Mutex;
9use thiserror::Error;
10use wgpu::{
11    CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT,
12};
13
14use crate::{
15    prelude::{Image, Shader},
16    render_asset::RenderAssetUsages,
17    render_resource::{
18        binding_types::texture_2d, BindGroup, BindGroupLayout, BindGroupLayoutEntries, Buffer,
19        CachedRenderPipelineId, FragmentState, PipelineCache, RenderPipelineDescriptor,
20        SpecializedRenderPipeline, SpecializedRenderPipelines, Texture, VertexState,
21    },
22    renderer::RenderDevice,
23    texture::TextureFormatPixelInfo,
24    RenderApp,
25};
26
27use super::ExtractedWindows;
28
29pub type ScreenshotFn = Box<dyn FnOnce(Image) + Send + Sync>;
30
31/// A resource which allows for taking screenshots of the window.
32#[derive(Resource, Default)]
33pub struct ScreenshotManager {
34    // this is in a mutex to enable extraction with only an immutable reference
35    pub(crate) callbacks: Mutex<EntityHashMap<ScreenshotFn>>,
36}
37
38#[derive(Error, Debug)]
39#[error("A screenshot for this window has already been requested.")]
40pub struct ScreenshotAlreadyRequestedError;
41
42impl ScreenshotManager {
43    /// Signals the renderer to take a screenshot of this frame.
44    ///
45    /// The given callback will eventually be called on one of the [`AsyncComputeTaskPool`]s threads.
46    pub fn take_screenshot(
47        &mut self,
48        window: Entity,
49        callback: impl FnOnce(Image) + Send + Sync + 'static,
50    ) -> Result<(), ScreenshotAlreadyRequestedError> {
51        self.callbacks
52            .get_mut()
53            .unwrap_or_else(PoisonError::into_inner)
54            .try_insert(window, Box::new(callback))
55            .map(|_| ())
56            .map_err(|_| ScreenshotAlreadyRequestedError)
57    }
58
59    /// Signals the renderer to take a screenshot of this frame.
60    ///
61    /// The screenshot will eventually be saved to the given path, and the format will be derived from the extension.
62    pub fn save_screenshot_to_disk(
63        &mut self,
64        window: Entity,
65        path: impl AsRef<Path>,
66    ) -> Result<(), ScreenshotAlreadyRequestedError> {
67        let path = path.as_ref().to_owned();
68        self.take_screenshot(window, move |img| match img.try_into_dynamic() {
69            Ok(dyn_img) => match image::ImageFormat::from_path(&path) {
70                Ok(format) => {
71                    // discard the alpha channel which stores brightness values when HDR is enabled to make sure
72                    // the screenshot looks right
73                    let img = dyn_img.to_rgb8();
74                    #[cfg(not(target_arch = "wasm32"))]
75                    match img.save_with_format(&path, format) {
76                        Ok(_) => info!("Screenshot saved to {}", path.display()),
77                        Err(e) => error!("Cannot save screenshot, IO error: {e}"),
78                    }
79
80                    #[cfg(target_arch = "wasm32")]
81                    {
82                        let save_screenshot = || {
83                            use image::EncodableLayout;
84                            use wasm_bindgen::{JsCast, JsValue};
85
86                            let mut image_buffer = std::io::Cursor::new(Vec::new());
87                            img.write_to(&mut image_buffer, format)
88                                .map_err(|e| JsValue::from_str(&format!("{e}")))?;
89                            // SAFETY: `image_buffer` only exist in this closure, and is not used after this line
90                            let parts = js_sys::Array::of1(&unsafe {
91                                js_sys::Uint8Array::view(image_buffer.into_inner().as_bytes())
92                                    .into()
93                            });
94                            let blob = web_sys::Blob::new_with_u8_array_sequence(&parts)?;
95                            let url = web_sys::Url::create_object_url_with_blob(&blob)?;
96                            let window = web_sys::window().unwrap();
97                            let document = window.document().unwrap();
98                            let link = document.create_element("a")?;
99                            link.set_attribute("href", &url)?;
100                            link.set_attribute(
101                                "download",
102                                path.file_name()
103                                    .and_then(|filename| filename.to_str())
104                                    .ok_or_else(|| JsValue::from_str("Invalid filename"))?,
105                            )?;
106                            let html_element = link.dyn_into::<web_sys::HtmlElement>()?;
107                            html_element.click();
108                            web_sys::Url::revoke_object_url(&url)?;
109                            Ok::<(), JsValue>(())
110                        };
111
112                        match (save_screenshot)() {
113                            Ok(_) => info!("Screenshot saved to {}", path.display()),
114                            Err(e) => error!("Cannot save screenshot, error: {e:?}"),
115                        };
116                    }
117                }
118                Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"),
119            },
120            Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"),
121        })
122    }
123}
124
125pub struct ScreenshotPlugin;
126
127const SCREENSHOT_SHADER_HANDLE: Handle<Shader> = Handle::weak_from_u128(11918575842344596158);
128
129impl Plugin for ScreenshotPlugin {
130    fn build(&self, app: &mut bevy_app::App) {
131        app.init_resource::<ScreenshotManager>();
132
133        load_internal_asset!(
134            app,
135            SCREENSHOT_SHADER_HANDLE,
136            "screenshot.wgsl",
137            Shader::from_wgsl
138        );
139    }
140
141    fn finish(&self, app: &mut bevy_app::App) {
142        if let Some(render_app) = app.get_sub_app_mut(RenderApp) {
143            render_app.init_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>();
144        }
145    }
146}
147
148pub(crate) fn align_byte_size(value: u32) -> u32 {
149    value + (COPY_BYTES_PER_ROW_ALIGNMENT - (value % COPY_BYTES_PER_ROW_ALIGNMENT))
150}
151
152pub(crate) fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 {
153    height * align_byte_size(width * pixel_size)
154}
155
156pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> ImageDataLayout {
157    ImageDataLayout {
158        bytes_per_row: if height > 1 {
159            // 1 = 1 row
160            Some(get_aligned_size(width, 1, format.pixel_size() as u32))
161        } else {
162            None
163        },
164        rows_per_image: None,
165        ..Default::default()
166    }
167}
168
169#[derive(Resource)]
170pub struct ScreenshotToScreenPipeline {
171    pub bind_group_layout: BindGroupLayout,
172}
173
174impl FromWorld for ScreenshotToScreenPipeline {
175    fn from_world(render_world: &mut World) -> Self {
176        let device = render_world.resource::<RenderDevice>();
177
178        let bind_group_layout = device.create_bind_group_layout(
179            "screenshot-to-screen-bgl",
180            &BindGroupLayoutEntries::single(
181                wgpu::ShaderStages::FRAGMENT,
182                texture_2d(wgpu::TextureSampleType::Float { filterable: false }),
183            ),
184        );
185
186        Self { bind_group_layout }
187    }
188}
189
190impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
191    type Key = TextureFormat;
192
193    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
194        RenderPipelineDescriptor {
195            label: Some(Cow::Borrowed("screenshot-to-screen")),
196            layout: vec![self.bind_group_layout.clone()],
197            vertex: VertexState {
198                buffers: vec![],
199                shader_defs: vec![],
200                entry_point: Cow::Borrowed("vs_main"),
201                shader: SCREENSHOT_SHADER_HANDLE,
202            },
203            primitive: wgpu::PrimitiveState {
204                cull_mode: Some(wgpu::Face::Back),
205                ..Default::default()
206            },
207            depth_stencil: None,
208            multisample: Default::default(),
209            fragment: Some(FragmentState {
210                shader: SCREENSHOT_SHADER_HANDLE,
211                entry_point: Cow::Borrowed("fs_main"),
212                shader_defs: vec![],
213                targets: vec![Some(wgpu::ColorTargetState {
214                    format: key,
215                    blend: None,
216                    write_mask: wgpu::ColorWrites::ALL,
217                })],
218            }),
219            push_constant_ranges: Vec::new(),
220        }
221    }
222}
223
224pub struct ScreenshotPreparedState {
225    pub texture: Texture,
226    pub buffer: Buffer,
227    pub bind_group: BindGroup,
228    pub pipeline_id: CachedRenderPipelineId,
229}
230
231pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) {
232    let windows = world.resource::<ExtractedWindows>();
233    let pipelines = world.resource::<PipelineCache>();
234
235    for window in windows.values() {
236        if let Some(memory) = &window.screenshot_memory {
237            let width = window.physical_width;
238            let height = window.physical_height;
239            let texture_format = window.swap_chain_texture_format.unwrap();
240
241            encoder.copy_texture_to_buffer(
242                memory.texture.as_image_copy(),
243                wgpu::ImageCopyBuffer {
244                    buffer: &memory.buffer,
245                    layout: layout_data(width, height, texture_format),
246                },
247                Extent3d {
248                    width,
249                    height,
250                    ..Default::default()
251                },
252            );
253            if let Some(pipeline) = pipelines.get_render_pipeline(memory.pipeline_id) {
254                let true_swapchain_texture_view = window
255                    .swap_chain_texture
256                    .as_ref()
257                    .unwrap()
258                    .texture
259                    .create_view(&Default::default());
260                let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
261                    label: Some("screenshot_to_screen_pass"),
262                    color_attachments: &[Some(wgpu::RenderPassColorAttachment {
263                        view: &true_swapchain_texture_view,
264                        resolve_target: None,
265                        ops: wgpu::Operations {
266                            load: wgpu::LoadOp::Load,
267                            store: wgpu::StoreOp::Store,
268                        },
269                    })],
270                    depth_stencil_attachment: None,
271                    timestamp_writes: None,
272                    occlusion_query_set: None,
273                });
274                pass.set_pipeline(pipeline);
275                pass.set_bind_group(0, &memory.bind_group, &[]);
276                pass.draw(0..3, 0..1);
277            }
278        }
279    }
280}
281
282pub(crate) fn collect_screenshots(world: &mut World) {
283    let _span = info_span!("collect_screenshots");
284
285    let mut windows = world.resource_mut::<ExtractedWindows>();
286    for window in windows.values_mut() {
287        if let Some(screenshot_func) = window.screenshot_func.take() {
288            let width = window.physical_width;
289            let height = window.physical_height;
290            let texture_format = window.swap_chain_texture_format.unwrap();
291            let pixel_size = texture_format.pixel_size();
292            let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap();
293
294            let finish = async move {
295                let (tx, rx) = async_channel::bounded(1);
296                let buffer_slice = buffer.slice(..);
297                // The polling for this map call is done every frame when the command queue is submitted.
298                buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
299                    let err = result.err();
300                    if err.is_some() {
301                        panic!("{}", err.unwrap().to_string());
302                    }
303                    tx.try_send(()).unwrap();
304                });
305                rx.recv().await.unwrap();
306                let data = buffer_slice.get_mapped_range();
307                // we immediately move the data to CPU memory to avoid holding the mapped view for long
308                let mut result = Vec::from(&*data);
309                drop(data);
310                drop(buffer);
311
312                if result.len() != ((width * height) as usize * pixel_size) {
313                    // Our buffer has been padded because we needed to align to a multiple of 256.
314                    // We remove this padding here
315                    let initial_row_bytes = width as usize * pixel_size;
316                    let buffered_row_bytes = align_byte_size(width * pixel_size as u32) as usize;
317
318                    let mut take_offset = buffered_row_bytes;
319                    let mut place_offset = initial_row_bytes;
320                    for _ in 1..height {
321                        result.copy_within(
322                            take_offset..take_offset + buffered_row_bytes,
323                            place_offset,
324                        );
325                        take_offset += buffered_row_bytes;
326                        place_offset += initial_row_bytes;
327                    }
328                    result.truncate(initial_row_bytes * height as usize);
329                }
330
331                screenshot_func(Image::new(
332                    Extent3d {
333                        width,
334                        height,
335                        depth_or_array_layers: 1,
336                    },
337                    wgpu::TextureDimension::D2,
338                    result,
339                    texture_format,
340                    RenderAssetUsages::RENDER_WORLD,
341                ));
342            };
343
344            AsyncComputeTaskPool::get().spawn(finish).detach();
345        }
346    }
347}
bevy_render/view/window/screenshot.rs

bevy_render/view/window/
screenshot.rs
