From 9db70da96f465d670f21e981fe1388655ffbee68 Mon Sep 17 00:00:00 2001 From: Nile Date: Thu, 20 Apr 2023 00:28:42 +0300 Subject: [PATCH] Add screenshot api (#7163) Fixes https://github.com/bevyengine/bevy/issues/1207 # Objective Right now, it's impossible to capture a screenshot of the entire window without forking bevy. This is because - The swapchain texture never has the COPY_SRC usage - It can't be accessed without taking ownership of it - Taking ownership of it breaks *a lot* of stuff ## Solution - Introduce a dedicated api for taking a screenshot of a given bevy window, and guarantee this screenshot will always match up with what gets put on the screen. --- ## Changelog - Added the `ScreenshotManager` resource with two functions, `take_screenshot` and `save_screenshot_to_disk` --- Cargo.toml | 10 + crates/bevy_render/src/camera/camera.rs | 2 +- .../src/camera/camera_driver_node.rs | 2 +- .../src/render_resource/texture.rs | 67 ++-- .../bevy_render/src/renderer/graph_runner.rs | 3 + crates/bevy_render/src/renderer/mod.rs | 9 +- .../src/texture/image_texture_conversion.rs | 15 + crates/bevy_render/src/view/window.rs | 106 +++++- .../bevy_render/src/view/window/screenshot.rs | 315 ++++++++++++++++++ .../src/view/window/screenshot.wgsl | 16 + examples/README.md | 1 + examples/window/screenshot.rs | 64 ++++ 12 files changed, 557 insertions(+), 53 deletions(-) create mode 100644 crates/bevy_render/src/view/window/screenshot.rs create mode 100644 crates/bevy_render/src/view/window/screenshot.wgsl create mode 100644 examples/window/screenshot.rs diff --git a/Cargo.toml b/Cargo.toml index 1110809061..ba49a55ecd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1909,6 +1909,16 @@ description = "Illustrates how to customize the default window settings" category = "Window" wasm = true +[[example]] +name = "screenshot" +path = "examples/window/screenshot.rs" + +[package.metadata.example.screenshot] +name = "Screenshot" +description = "Shows how to save screenshots to disk" +category = "Window" +wasm = false + [[example]] name = "transparent_window" path = "examples/window/transparent_window.rs" diff --git a/crates/bevy_render/src/camera/camera.rs b/crates/bevy_render/src/camera/camera.rs index 851783a3f1..ab3af2c242 100644 --- a/crates/bevy_render/src/camera/camera.rs +++ b/crates/bevy_render/src/camera/camera.rs @@ -423,7 +423,7 @@ impl NormalizedRenderTarget { match self { NormalizedRenderTarget::Window(window_ref) => windows .get(&window_ref.entity()) - .and_then(|window| window.swap_chain_texture.as_ref()), + .and_then(|window| window.swap_chain_texture_view.as_ref()), NormalizedRenderTarget::Image(image_handle) => { images.get(image_handle).map(|image| &image.texture_view) } diff --git a/crates/bevy_render/src/camera/camera_driver_node.rs b/crates/bevy_render/src/camera/camera_driver_node.rs index 84ff1896fc..b56bc52fc8 100644 --- a/crates/bevy_render/src/camera/camera_driver_node.rs +++ b/crates/bevy_render/src/camera/camera_driver_node.rs @@ -52,7 +52,7 @@ impl Node for CameraDriverNode { continue; } - let Some(swap_chain_texture) = &window.swap_chain_texture else { + let Some(swap_chain_texture) = &window.swap_chain_texture_view else { continue; }; diff --git a/crates/bevy_render/src/render_resource/texture.rs b/crates/bevy_render/src/render_resource/texture.rs index 02cc818408..df0df616a6 100644 --- a/crates/bevy_render/src/render_resource/texture.rs +++ b/crates/bevy_render/src/render_resource/texture.rs @@ -51,31 +51,21 @@ define_atomic_id!(TextureViewId); render_resource_wrapper!(ErasedTextureView, wgpu::TextureView); render_resource_wrapper!(ErasedSurfaceTexture, wgpu::SurfaceTexture); -/// This type combines wgpu's [`TextureView`](wgpu::TextureView) and -/// [`SurfaceTexture`](wgpu::SurfaceTexture) into the same interface. -#[derive(Clone, Debug)] -pub enum TextureViewValue { - /// The value is an actual wgpu [`TextureView`](wgpu::TextureView). - TextureView(ErasedTextureView), - - /// The value is a wgpu [`SurfaceTexture`](wgpu::SurfaceTexture), but dereferences to - /// a [`TextureView`](wgpu::TextureView). - SurfaceTexture { - // NOTE: The order of these fields is important because the view must be dropped before the - // frame is dropped - view: ErasedTextureView, - texture: ErasedSurfaceTexture, - }, -} - /// Describes a [`Texture`] with its associated metadata required by a pipeline or [`BindGroup`](super::BindGroup). -/// -/// May be converted from a [`TextureView`](wgpu::TextureView) or [`SurfaceTexture`](wgpu::SurfaceTexture) -/// or dereferences to a wgpu [`TextureView`](wgpu::TextureView). #[derive(Clone, Debug)] pub struct TextureView { id: TextureViewId, - value: TextureViewValue, + value: ErasedTextureView, +} + +pub struct SurfaceTexture { + value: ErasedSurfaceTexture, +} + +impl SurfaceTexture { + pub fn try_unwrap(self) -> Option { + self.value.try_unwrap() + } } impl TextureView { @@ -84,34 +74,21 @@ impl TextureView { pub fn id(&self) -> TextureViewId { self.id } - - /// Returns the [`SurfaceTexture`](wgpu::SurfaceTexture) of the texture view if it is of that type. - #[inline] - pub fn take_surface_texture(self) -> Option { - match self.value { - TextureViewValue::TextureView(_) => None, - TextureViewValue::SurfaceTexture { texture, .. } => texture.try_unwrap(), - } - } } impl From for TextureView { fn from(value: wgpu::TextureView) -> Self { TextureView { id: TextureViewId::new(), - value: TextureViewValue::TextureView(ErasedTextureView::new(value)), + value: ErasedTextureView::new(value), } } } -impl From for TextureView { +impl From for SurfaceTexture { fn from(value: wgpu::SurfaceTexture) -> Self { - let view = ErasedTextureView::new(value.texture.create_view(&Default::default())); - let texture = ErasedSurfaceTexture::new(value); - - TextureView { - id: TextureViewId::new(), - value: TextureViewValue::SurfaceTexture { texture, view }, + SurfaceTexture { + value: ErasedSurfaceTexture::new(value), } } } @@ -121,10 +98,16 @@ impl Deref for TextureView { #[inline] fn deref(&self) -> &Self::Target { - match &self.value { - TextureViewValue::TextureView(value) => value, - TextureViewValue::SurfaceTexture { view, .. } => view, - } + &self.value + } +} + +impl Deref for SurfaceTexture { + type Target = wgpu::SurfaceTexture; + + #[inline] + fn deref(&self) -> &Self::Target { + &self.value } } diff --git a/crates/bevy_render/src/renderer/graph_runner.rs b/crates/bevy_render/src/renderer/graph_runner.rs index 478298dd42..c046ef11a1 100644 --- a/crates/bevy_render/src/renderer/graph_runner.rs +++ b/crates/bevy_render/src/renderer/graph_runner.rs @@ -57,9 +57,12 @@ impl RenderGraphRunner { render_device: RenderDevice, queue: &wgpu::Queue, world: &World, + finalizer: impl FnOnce(&mut wgpu::CommandEncoder), ) -> Result<(), RenderGraphRunnerError> { let mut render_context = RenderContext::new(render_device); Self::run_graph(graph, None, &mut render_context, world, &[], None)?; + finalizer(render_context.command_encoder()); + { #[cfg(feature = "trace")] let _span = info_span!("submit_graph_commands").entered(); diff --git a/crates/bevy_render/src/renderer/mod.rs b/crates/bevy_render/src/renderer/mod.rs index b47f73b85e..a0fb731a6e 100644 --- a/crates/bevy_render/src/renderer/mod.rs +++ b/crates/bevy_render/src/renderer/mod.rs @@ -35,6 +35,9 @@ pub fn render_system(world: &mut World) { render_device.clone(), // TODO: is this clone really necessary? &render_queue.0, world, + |encoder| { + crate::view::screenshot::submit_screenshot_commands(world, encoder); + }, ) { error!("Error running render graph:"); { @@ -66,8 +69,8 @@ pub fn render_system(world: &mut World) { let mut windows = world.resource_mut::(); for window in windows.values_mut() { - if let Some(texture_view) = window.swap_chain_texture.take() { - if let Some(surface_texture) = texture_view.take_surface_texture() { + if let Some(wrapped_texture) = window.swap_chain_texture.take() { + if let Some(surface_texture) = wrapped_texture.try_unwrap() { surface_texture.present(); } } @@ -81,6 +84,8 @@ pub fn render_system(world: &mut World) { ); } + crate::view::screenshot::collect_screenshots(world); + // update the time and send it to the app world let time_sender = world.resource::(); time_sender.0.try_send(Instant::now()).expect( diff --git a/crates/bevy_render/src/texture/image_texture_conversion.rs b/crates/bevy_render/src/texture/image_texture_conversion.rs index 71eeff23e9..387d08b307 100644 --- a/crates/bevy_render/src/texture/image_texture_conversion.rs +++ b/crates/bevy_render/src/texture/image_texture_conversion.rs @@ -174,6 +174,7 @@ impl Image { /// - `TextureFormat::R8Unorm` /// - `TextureFormat::Rg8Unorm` /// - `TextureFormat::Rgba8UnormSrgb` + /// - `TextureFormat::Bgra8UnormSrgb` /// /// To convert [`Image`] to a different format see: [`Image::convert`]. pub fn try_into_dynamic(self) -> anyhow::Result { @@ -196,6 +197,20 @@ impl Image { self.data, ) .map(DynamicImage::ImageRgba8), + // This format is commonly used as the format for the swapchain texture + // This conversion is added here to support screenshots + TextureFormat::Bgra8UnormSrgb => ImageBuffer::from_raw( + self.texture_descriptor.size.width, + self.texture_descriptor.size.height, + { + let mut data = self.data; + for bgra in data.chunks_exact_mut(4) { + bgra.swap(0, 2); + } + data + }, + ) + .map(DynamicImage::ImageRgba8), // Throw and error if conversion isn't supported texture_format => { return Err(anyhow!( diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 8118e4fd06..0ca26adcf9 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -1,6 +1,7 @@ use crate::{ - render_resource::TextureView, + render_resource::{PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, + texture::TextureFormatPixelInfo, Extract, ExtractSchedule, Render, RenderApp, RenderSet, }; use bevy_app::{App, Plugin}; @@ -10,7 +11,13 @@ use bevy_window::{ CompositeAlphaMode, PresentMode, PrimaryWindow, RawHandleWrapper, Window, WindowClosed, }; use std::ops::{Deref, DerefMut}; -use wgpu::TextureFormat; +use wgpu::{BufferUsages, TextureFormat, TextureUsages}; + +pub mod screenshot; + +use screenshot::{ + ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline, +}; use super::Msaa; @@ -27,10 +34,13 @@ pub enum WindowSystem { impl Plugin for WindowRenderPlugin { fn build(&self, app: &mut App) { + app.add_plugin(ScreenshotPlugin); + if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { render_app .init_resource::() .init_resource::() + .init_resource::() .init_non_send_resource::() .add_systems(ExtractSchedule, extract_windows) .configure_set(Render, WindowSystem::Prepare.in_set(RenderSet::Prepare)) @@ -46,11 +56,26 @@ pub struct ExtractedWindow { pub physical_width: u32, pub physical_height: u32, pub present_mode: PresentMode, - pub swap_chain_texture: Option, + /// Note: this will not always be the swap chain texture view. When taking a screenshot, + /// this will point to an alternative texture instead to allow for copying the render result + /// to CPU memory. + pub swap_chain_texture_view: Option, + pub swap_chain_texture: Option, pub swap_chain_texture_format: Option, + pub screenshot_memory: Option, pub size_changed: bool, pub present_mode_changed: bool, pub alpha_mode: CompositeAlphaMode, + pub screenshot_func: Option, +} + +impl ExtractedWindow { + fn set_swapchain_texture(&mut self, frame: wgpu::SurfaceTexture) { + self.swap_chain_texture_view = Some(TextureView::from( + frame.texture.create_view(&Default::default()), + )); + self.swap_chain_texture = Some(SurfaceTexture::from(frame)); + } } #[derive(Default, Resource)] @@ -75,6 +100,7 @@ impl DerefMut for ExtractedWindows { fn extract_windows( mut extracted_windows: ResMut, + screenshot_manager: Extract>, mut closed: Extract>, windows: Extract)>>, ) { @@ -95,14 +121,17 @@ fn extract_windows( physical_height: new_height, present_mode: window.present_mode, swap_chain_texture: None, + swap_chain_texture_view: None, size_changed: false, swap_chain_texture_format: None, present_mode_changed: false, alpha_mode: window.composite_alpha_mode, + screenshot_func: None, + screenshot_memory: None, }); // NOTE: Drop the swap chain frame here - extracted_window.swap_chain_texture = None; + extracted_window.swap_chain_texture_view = None; extracted_window.size_changed = new_width != extracted_window.physical_width || new_height != extracted_window.physical_height; extracted_window.present_mode_changed = @@ -132,6 +161,15 @@ fn extract_windows( for closed_window in closed.iter() { extracted_windows.remove(&closed_window.window); } + // This lock will never block because `callbacks` is `pub(crate)` and this is the singular callsite where it's locked. + // Even if a user had multiple copies of this system, since the system has a mutable resource access the two systems would never run + // at the same time + // TODO: since this is guaranteed, should the lock be replaced with an UnsafeCell to remove the overhead, or is it minor enough to be ignored? + for (window, screenshot_func) in screenshot_manager.callbacks.lock().drain() { + if let Some(window) = extracted_windows.get_mut(&window) { + window.screenshot_func = Some(screenshot_func); + } + } } struct SurfaceData { @@ -167,6 +205,7 @@ pub struct WindowSurfaces { /// another alternative is to try to use [`ANGLE`](https://github.com/gfx-rs/wgpu#angle) and /// [`Backends::GL`](crate::settings::Backends::GL) if your GPU/drivers support `OpenGL 4.3` / `OpenGL ES 3.0` or /// later. +#[allow(clippy::too_many_arguments)] pub fn prepare_windows( // By accessing a NonSend resource, we tell the scheduler to put this system on the main thread, // which is necessary for some OS s @@ -176,6 +215,9 @@ pub fn prepare_windows( render_device: Res, render_instance: Res, render_adapter: Res, + screenshot_pipeline: Res, + pipeline_cache: Res, + mut pipelines: ResMut>, mut msaa: ResMut, ) { for window in windows.windows.values_mut() { @@ -285,18 +327,18 @@ pub fn prepare_windows( let frame = surface .get_current_texture() .expect("Error configuring surface"); - window.swap_chain_texture = Some(TextureView::from(frame)); + window.set_swapchain_texture(frame); } else { match surface.get_current_texture() { Ok(frame) => { - window.swap_chain_texture = Some(TextureView::from(frame)); + window.set_swapchain_texture(frame); } Err(wgpu::SurfaceError::Outdated) => { render_device.configure_surface(surface, &surface_configuration); let frame = surface .get_current_texture() .expect("Error reconfiguring surface"); - window.swap_chain_texture = Some(TextureView::from(frame)); + window.set_swapchain_texture(frame); } #[cfg(target_os = "linux")] Err(wgpu::SurfaceError::Timeout) if may_erroneously_timeout() => { @@ -311,5 +353,55 @@ pub fn prepare_windows( } }; window.swap_chain_texture_format = Some(surface_data.format); + + if window.screenshot_func.is_some() { + let texture = render_device.create_texture(&wgpu::TextureDescriptor { + label: Some("screenshot-capture-rendertarget"), + size: wgpu::Extent3d { + width: surface_configuration.width, + height: surface_configuration.height, + depth_or_array_layers: 1, + }, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + format: surface_configuration.format, + usage: TextureUsages::RENDER_ATTACHMENT + | TextureUsages::COPY_SRC + | TextureUsages::TEXTURE_BINDING, + view_formats: &[], + }); + let texture_view = texture.create_view(&Default::default()); + let buffer = render_device.create_buffer(&wgpu::BufferDescriptor { + label: Some("screenshot-transfer-buffer"), + size: screenshot::get_aligned_size( + window.physical_width, + window.physical_height, + surface_data.format.pixel_size() as u32, + ) as u64, + usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let bind_group = render_device.create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("screenshot-to-screen-bind-group"), + layout: &screenshot_pipeline.bind_group_layout, + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView(&texture_view), + }], + }); + let pipeline_id = pipelines.specialize( + &pipeline_cache, + &screenshot_pipeline, + surface_configuration.format, + ); + window.swap_chain_texture_view = Some(texture_view); + window.screenshot_memory = Some(ScreenshotPreparedState { + texture, + buffer, + bind_group, + pipeline_id, + }); + } } } diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs new file mode 100644 index 0000000000..2a67f43005 --- /dev/null +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -0,0 +1,315 @@ +use std::{borrow::Cow, num::NonZeroU32, path::Path}; + +use bevy_app::Plugin; +use bevy_asset::{load_internal_asset, HandleUntyped}; +use bevy_ecs::prelude::*; +use bevy_log::{error, info, info_span}; +use bevy_reflect::TypeUuid; +use bevy_tasks::AsyncComputeTaskPool; +use bevy_utils::HashMap; +use parking_lot::Mutex; +use thiserror::Error; +use wgpu::{ + CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, +}; + +use crate::{ + prelude::{Image, Shader}, + render_resource::{ + BindGroup, BindGroupLayout, Buffer, CachedRenderPipelineId, FragmentState, PipelineCache, + RenderPipelineDescriptor, SpecializedRenderPipeline, SpecializedRenderPipelines, Texture, + VertexState, + }, + renderer::RenderDevice, + texture::TextureFormatPixelInfo, + RenderApp, +}; + +use super::ExtractedWindows; + +pub type ScreenshotFn = Box; + +/// A resource which allows for taking screenshots of the window. +#[derive(Resource, Default)] +pub struct ScreenshotManager { + // this is in a mutex to enable extraction with only an immutable reference + pub(crate) callbacks: Mutex>, +} + +#[derive(Error, Debug)] +#[error("A screenshot for this window has already been requested.")] +pub struct ScreenshotAlreadyRequestedError; + +impl ScreenshotManager { + /// Signals the renderer to take a screenshot of this frame. + /// + /// The given callback will eventually be called on one of the [`AsyncComputeTaskPool`]s threads. + pub fn take_screenshot( + &mut self, + window: Entity, + callback: impl FnOnce(Image) + Send + Sync + 'static, + ) -> Result<(), ScreenshotAlreadyRequestedError> { + self.callbacks + .get_mut() + .try_insert(window, Box::new(callback)) + .map(|_| ()) + .map_err(|_| ScreenshotAlreadyRequestedError) + } + + /// Signals the renderer to take a screenshot of this frame. + /// + /// The screenshot will eventually be saved to the given path, and the format will be derived from the extension. + pub fn save_screenshot_to_disk( + &mut self, + window: Entity, + path: impl AsRef, + ) -> Result<(), ScreenshotAlreadyRequestedError> { + let path = path.as_ref().to_owned(); + self.take_screenshot(window, move |img| match img.try_into_dynamic() { + Ok(dyn_img) => match image::ImageFormat::from_path(&path) { + Ok(format) => { + // discard the alpha channel which stores brightness values when HDR is enabled to make sure + // the screenshot looks right + let img = dyn_img.to_rgb8(); + match img.save_with_format(&path, format) { + Ok(_) => info!("Screenshot saved to {}", path.display()), + Err(e) => error!("Cannot save screenshot, IO error: {e}"), + } + } + Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"), + }, + Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"), + }) + } +} + +pub struct ScreenshotPlugin; + +const SCREENSHOT_SHADER_HANDLE: HandleUntyped = + HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 11918575842344596158); + +impl Plugin for ScreenshotPlugin { + fn build(&self, app: &mut bevy_app::App) { + app.init_resource::(); + + load_internal_asset!( + app, + SCREENSHOT_SHADER_HANDLE, + "screenshot.wgsl", + Shader::from_wgsl + ); + + if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { + render_app.init_resource::>(); + } + } +} + +pub(crate) fn align_byte_size(value: u32) -> u32 { + value + (COPY_BYTES_PER_ROW_ALIGNMENT - (value % COPY_BYTES_PER_ROW_ALIGNMENT)) +} + +pub(crate) fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 { + height * align_byte_size(width * pixel_size) +} + +pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> ImageDataLayout { + ImageDataLayout { + bytes_per_row: if height > 1 { + // 1 = 1 row + NonZeroU32::new(get_aligned_size(width, 1, format.pixel_size() as u32)) + } else { + None + }, + rows_per_image: None, + ..Default::default() + } +} + +#[derive(Resource)] +pub struct ScreenshotToScreenPipeline { + pub bind_group_layout: BindGroupLayout, +} + +impl FromWorld for ScreenshotToScreenPipeline { + fn from_world(render_world: &mut World) -> Self { + let device = render_world.resource::(); + + let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("screenshot-to-screen-bgl"), + entries: &[wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::FRAGMENT, + ty: wgpu::BindingType::Texture { + sample_type: wgpu::TextureSampleType::Float { filterable: false }, + view_dimension: wgpu::TextureViewDimension::D2, + multisampled: false, + }, + count: None, + }], + }); + + Self { bind_group_layout } + } +} + +impl SpecializedRenderPipeline for ScreenshotToScreenPipeline { + type Key = TextureFormat; + + fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor { + RenderPipelineDescriptor { + label: Some(Cow::Borrowed("screenshot-to-screen")), + layout: vec![self.bind_group_layout.clone()], + vertex: VertexState { + buffers: vec![], + shader_defs: vec![], + entry_point: Cow::Borrowed("vs_main"), + shader: SCREENSHOT_SHADER_HANDLE.typed(), + }, + primitive: wgpu::PrimitiveState { + topology: wgpu::PrimitiveTopology::TriangleList, + strip_index_format: None, + front_face: wgpu::FrontFace::Ccw, + cull_mode: Some(wgpu::Face::Back), + polygon_mode: wgpu::PolygonMode::Fill, + conservative: false, + unclipped_depth: false, + }, + depth_stencil: None, + multisample: wgpu::MultisampleState::default(), + fragment: Some(FragmentState { + shader: SCREENSHOT_SHADER_HANDLE.typed(), + entry_point: Cow::Borrowed("fs_main"), + shader_defs: vec![], + targets: vec![Some(wgpu::ColorTargetState { + format: key, + blend: None, + write_mask: wgpu::ColorWrites::ALL, + })], + }), + push_constant_ranges: Vec::new(), + } + } +} + +pub struct ScreenshotPreparedState { + pub texture: Texture, + pub buffer: Buffer, + pub bind_group: BindGroup, + pub pipeline_id: CachedRenderPipelineId, +} + +pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) { + let windows = world.resource::(); + let pipelines = world.resource::(); + + for window in windows.values() { + if let Some(memory) = &window.screenshot_memory { + let width = window.physical_width; + let height = window.physical_height; + let texture_format = window.swap_chain_texture_format.unwrap(); + + encoder.copy_texture_to_buffer( + memory.texture.as_image_copy(), + wgpu::ImageCopyBuffer { + buffer: &memory.buffer, + layout: crate::view::screenshot::layout_data(width, height, texture_format), + }, + Extent3d { + width, + height, + ..Default::default() + }, + ); + if let Some(pipeline) = pipelines.get_render_pipeline(memory.pipeline_id) { + let true_swapchain_texture_view = window + .swap_chain_texture + .as_ref() + .unwrap() + .texture + .create_view(&Default::default()); + let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("screenshot_to_screen_pass"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: &true_swapchain_texture_view, + resolve_target: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Load, + store: true, + }, + })], + depth_stencil_attachment: None, + }); + pass.set_pipeline(pipeline); + pass.set_bind_group(0, &memory.bind_group, &[]); + pass.draw(0..3, 0..1); + } + } + } +} + +pub(crate) fn collect_screenshots(world: &mut World) { + let _span = info_span!("collect_screenshots"); + + let mut windows = world.resource_mut::(); + for window in windows.values_mut() { + if let Some(screenshot_func) = window.screenshot_func.take() { + let width = window.physical_width; + let height = window.physical_height; + let texture_format = window.swap_chain_texture_format.unwrap(); + let pixel_size = texture_format.pixel_size(); + let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap(); + + let finish = async move { + let (tx, rx) = async_channel::bounded(1); + let buffer_slice = buffer.slice(..); + // The polling for this map call is done every frame when the command queue is submitted. + buffer_slice.map_async(wgpu::MapMode::Read, move |result| { + let err = result.err(); + if err.is_some() { + panic!("{}", err.unwrap().to_string()); + } + tx.try_send(()).unwrap(); + }); + rx.recv().await.unwrap(); + let data = buffer_slice.get_mapped_range(); + // we immediately move the data to CPU memory to avoid holding the mapped view for long + let mut result = Vec::from(&*data); + drop(data); + drop(buffer); + + if result.len() != ((width * height) as usize * pixel_size) { + // Our buffer has been padded because we needed to align to a multiple of 256. + // We remove this padding here + let initial_row_bytes = width as usize * pixel_size; + let buffered_row_bytes = align_byte_size(width * pixel_size as u32) as usize; + + let mut take_offset = buffered_row_bytes; + let mut place_offset = initial_row_bytes; + for _ in 1..height { + result.copy_within( + take_offset..take_offset + buffered_row_bytes, + place_offset, + ); + take_offset += buffered_row_bytes; + place_offset += initial_row_bytes; + } + result.truncate(initial_row_bytes * height as usize); + } + + screenshot_func(Image::new( + Extent3d { + width, + height, + depth_or_array_layers: 1, + }, + wgpu::TextureDimension::D2, + result, + texture_format, + )); + }; + + AsyncComputeTaskPool::get().spawn(finish).detach(); + } + } +} diff --git a/crates/bevy_render/src/view/window/screenshot.wgsl b/crates/bevy_render/src/view/window/screenshot.wgsl new file mode 100644 index 0000000000..2743fa1d95 --- /dev/null +++ b/crates/bevy_render/src/view/window/screenshot.wgsl @@ -0,0 +1,16 @@ +// This vertex shader will create a triangle that will cover the entire screen +// with minimal effort, avoiding the need for a vertex buffer etc. +@vertex +fn vs_main(@builtin(vertex_index) in_vertex_index: u32) -> @builtin(position) vec4 { + let x = f32((in_vertex_index & 1u) << 2u); + let y = f32((in_vertex_index & 2u) << 1u); + return vec4(x - 1.0, y - 1.0, 0.0, 1.0); +} + +@group(0) @binding(0) var t: texture_2d; + +@fragment +fn fs_main(@builtin(position) pos: vec4) -> @location(0) vec4 { + let coords = floor(pos.xy); + return textureLoad(t, vec2(coords), 0i); +} diff --git a/examples/README.md b/examples/README.md index c04d30c330..380ed3974f 100644 --- a/examples/README.md +++ b/examples/README.md @@ -353,6 +353,7 @@ Example | Description [Low Power](../examples/window/low_power.rs) | Demonstrates settings to reduce power use for bevy applications [Multiple Windows](../examples/window/multiple_windows.rs) | Demonstrates creating multiple windows, and rendering to them [Scale Factor Override](../examples/window/scale_factor_override.rs) | Illustrates how to customize the default window settings +[Screenshot](../examples/window/screenshot.rs) | Shows how to save screenshots to disk [Transparent Window](../examples/window/transparent_window.rs) | Illustrates making the window transparent and hiding the window decoration [Window Resizing](../examples/window/window_resizing.rs) | Demonstrates resizing and responding to resizing a window [Window Settings](../examples/window/window_settings.rs) | Demonstrates customizing default window settings diff --git a/examples/window/screenshot.rs b/examples/window/screenshot.rs new file mode 100644 index 0000000000..951a46fd79 --- /dev/null +++ b/examples/window/screenshot.rs @@ -0,0 +1,64 @@ +//! An example showing how to save screenshots to disk + +use bevy::prelude::*; +use bevy::render::view::screenshot::ScreenshotManager; +use bevy::window::PrimaryWindow; + +fn main() { + App::new() + .add_plugins(DefaultPlugins) + .add_systems(Startup, setup) + .add_systems(Update, screenshot_on_f12) + .run(); +} + +fn screenshot_on_f12( + input: Res>, + main_window: Query>, + mut screenshot_manager: ResMut, + mut counter: Local, +) { + if input.just_pressed(KeyCode::F12) { + let path = format!("./screenshot-{}.png", *counter); + *counter += 1; + screenshot_manager + .save_screenshot_to_disk(main_window.single(), path) + .unwrap(); + } +} + +/// set up a simple 3D scene +fn setup( + mut commands: Commands, + mut meshes: ResMut>, + mut materials: ResMut>, +) { + // plane + commands.spawn(PbrBundle { + mesh: meshes.add(shape::Plane::from_size(5.0).into()), + material: materials.add(Color::rgb(0.3, 0.5, 0.3).into()), + ..default() + }); + // cube + commands.spawn(PbrBundle { + mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })), + material: materials.add(Color::rgb(0.8, 0.7, 0.6).into()), + transform: Transform::from_xyz(0.0, 0.5, 0.0), + ..default() + }); + // light + commands.spawn(PointLightBundle { + point_light: PointLight { + intensity: 1500.0, + shadows_enabled: true, + ..default() + }, + transform: Transform::from_xyz(4.0, 8.0, 4.0), + ..default() + }); + // camera + commands.spawn(Camera3dBundle { + transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y), + ..default() + }); +}