From 7311b147565c8974f61d5e3d15cf57c734cbf7c8 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Wed, 11 Jan 2023 22:21:24 +0300 Subject: [PATCH 01/17] Add screenshot api --- crates/bevy_render/Cargo.toml | 2 + crates/bevy_render/src/lib.rs | 1 + .../src/render_resource/texture.rs | 8 + .../bevy_render/src/renderer/graph_runner.rs | 2 + crates/bevy_render/src/renderer/mod.rs | 6 + .../src/texture/image_texture_conversion.rs | 13 ++ crates/bevy_render/src/view/window.rs | 40 +++- .../bevy_render/src/view/window/screenshot.rs | 175 ++++++++++++++++++ 8 files changed, 243 insertions(+), 4 deletions(-) create mode 100644 crates/bevy_render/src/view/window/screenshot.rs diff --git a/crates/bevy_render/Cargo.toml b/crates/bevy_render/Cargo.toml index eaa4d09af6bac..ceb49a04dff6f 100644 --- a/crates/bevy_render/Cargo.toml +++ b/crates/bevy_render/Cargo.toml @@ -44,6 +44,7 @@ bevy_time = { path = "../bevy_time", version = "0.9.0" } bevy_transform = { path = "../bevy_transform", version = "0.9.0" } bevy_window = { path = "../bevy_window", version = "0.9.0" } bevy_utils = { path = "../bevy_utils", version = "0.9.0" } +bevy_tasks = { path = "../bevy_tasks", version = "0.9.0" } # rendering image = { version = "0.24", default-features = false } @@ -76,3 +77,4 @@ basis-universal = { version = "0.2.0", optional = true } encase = { version = "0.4", features = ["glam"] } # For wgpu profiling using tracing. Use `RUST_LOG=info` to also capture the wgpu spans. profiling = { version = "1", features = ["profile-with-tracing"], optional = true } +async-channel = "1.4" diff --git a/crates/bevy_render/src/lib.rs b/crates/bevy_render/src/lib.rs index 0ba66415c35c8..06326d69ad0ad 100644 --- a/crates/bevy_render/src/lib.rs +++ b/crates/bevy_render/src/lib.rs @@ -137,6 +137,7 @@ impl Plugin for RenderPlugin { app.add_asset::() .add_debug_asset::() .init_asset_loader::() + .init_resource::() .init_debug_asset_loader::(); if let Some(backends) = self.wgpu_settings.backends { diff --git a/crates/bevy_render/src/render_resource/texture.rs b/crates/bevy_render/src/render_resource/texture.rs index 02cc818408073..108f5e0259393 100644 --- a/crates/bevy_render/src/render_resource/texture.rs +++ b/crates/bevy_render/src/render_resource/texture.rs @@ -93,6 +93,14 @@ impl TextureView { TextureViewValue::SurfaceTexture { texture, .. } => texture.try_unwrap(), } } + + #[inline] + pub(crate) fn get_surface_texture(&self) -> Option<&wgpu::SurfaceTexture> { + match &self.value { + TextureViewValue::TextureView(_) => None, + TextureViewValue::SurfaceTexture { texture, .. } => Some(&*texture), + } + } } impl From for TextureView { diff --git a/crates/bevy_render/src/renderer/graph_runner.rs b/crates/bevy_render/src/renderer/graph_runner.rs index 1513d1c4e5697..316fce06e3ac6 100644 --- a/crates/bevy_render/src/renderer/graph_runner.rs +++ b/crates/bevy_render/src/renderer/graph_runner.rs @@ -57,6 +57,7 @@ impl RenderGraphRunner { render_device: RenderDevice, queue: &wgpu::Queue, world: &World, + finalizer: impl FnOnce(&mut wgpu::CommandEncoder), ) -> Result<(), RenderGraphRunnerError> { let command_encoder = render_device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default()); @@ -66,6 +67,7 @@ impl RenderGraphRunner { }; Self::run_graph(graph, None, &mut render_context, world, &[])?; + finalizer(&mut 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 19cee9c9585a5..14c456c993c77 100644 --- a/crates/bevy_render/src/renderer/mod.rs +++ b/crates/bevy_render/src/renderer/mod.rs @@ -27,12 +27,16 @@ pub fn render_system(world: &mut World) { let graph = world.resource::(); let render_device = world.resource::(); let render_queue = world.resource::(); + let windows = world.resource::(); if let Err(e) = RenderGraphRunner::run( graph, render_device.clone(), // TODO: is this clone really necessary? &render_queue.0, world, + |encoder| { + crate::view::screenshot::submit_screenshot_commands(windows, encoder); + }, ) { error!("Error running render graph:"); { @@ -79,6 +83,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 de8ed285f724e..a8debeae2aab6 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,18 @@ impl Image { self.data, ) .map(DynamicImage::ImageRgba8), + 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 018f193094522..bfd146ecd7fff 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -1,6 +1,9 @@ +pub mod screenshot; + use crate::{ - render_resource::TextureView, + render_resource::{Buffer, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, + texture::TextureFormatPixelInfo, Extract, RenderApp, RenderStage, }; use bevy_app::{App, Plugin}; @@ -10,7 +13,9 @@ use bevy_window::{ CompositeAlphaMode, PresentMode, RawHandleWrapper, WindowClosed, WindowId, Windows, }; use std::ops::{Deref, DerefMut}; -use wgpu::TextureFormat; +use wgpu::{BufferUsages, TextureFormat}; + +use self::screenshot::ScreenshotManager; /// Token to ensure a system runs on the main thread. #[derive(Resource, Default)] @@ -50,6 +55,8 @@ pub struct ExtractedWindow { pub size_changed: bool, pub present_mode_changed: bool, pub alpha_mode: CompositeAlphaMode, + pub screenshot_func: Option, + pub screenshot_buffer: Option, } #[derive(Default, Resource)] @@ -73,6 +80,7 @@ impl DerefMut for ExtractedWindows { fn extract_windows( mut extracted_windows: ResMut, + screenshot_manager: Extract>, mut closed: Extract>, windows: Extract>, ) { @@ -97,6 +105,8 @@ fn extract_windows( size_changed: false, present_mode_changed: false, alpha_mode: window.alpha_mode(), + screenshot_func: None, + screenshot_buffer: None, }); // NOTE: Drop the swap chain frame here @@ -128,6 +138,11 @@ fn extract_windows( for closed_window in closed.iter() { extracted_windows.remove(&closed_window.id); } + 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 { @@ -198,7 +213,7 @@ pub fn prepare_windows( SurfaceData { surface, format } }); - let surface_configuration = wgpu::SurfaceConfiguration { + let mut surface_configuration = wgpu::SurfaceConfiguration { format: surface_data.format, width: window.physical_width, height: window.physical_height, @@ -218,6 +233,19 @@ pub fn prepare_windows( CompositeAlphaMode::Inherit => wgpu::CompositeAlphaMode::Inherit, }, }; + if window.screenshot_func.is_some() { + surface_configuration.usage |= wgpu::TextureUsages::COPY_SRC; + window.screenshot_buffer = Some(render_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + 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, + })); + } // A recurring issue is hitting `wgpu::SurfaceError::Timeout` on certain Linux // mesa driver implementations. This seems to be a quirk of some drivers. @@ -239,7 +267,11 @@ pub fn prepare_windows( let not_already_configured = window_surfaces.configured_windows.insert(window.id); let surface = &surface_data.surface; - if not_already_configured || window.size_changed || window.present_mode_changed { + if not_already_configured + || window.size_changed + || window.present_mode_changed + || window.screenshot_func.is_some() + { render_device.configure_surface(surface, &surface_configuration); let frame = surface .get_current_texture() 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 0000000000000..a3831c4fe5d27 --- /dev/null +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -0,0 +1,175 @@ +use std::{num::NonZeroU32, path::Path}; + +use bevy_ecs::prelude::*; +use bevy_log::info_span; +use bevy_tasks::AsyncComputeTaskPool; +use bevy_utils::HashMap; +use bevy_window::WindowId; +use parking_lot::Mutex; +use thiserror::Error; +use wgpu::{ + CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, +}; + +use crate::{prelude::Image, texture::TextureFormatPixelInfo}; + +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: WindowId, + 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: WindowId, + path: impl AsRef, + ) -> Result<(), ScreenshotAlreadyRequestedError> { + let path = path.as_ref().to_owned(); + self.take_screenshot(window, |image| { + image.try_into_dynamic().unwrap().save(path).unwrap(); + }) + } +} + +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() + } +} + +pub(crate) fn submit_screenshot_commands(windows: &ExtractedWindows, encoder: &mut CommandEncoder) { + for (window, texture) in windows + .values() + .filter_map(|w| w.swap_chain_texture.as_ref().map(|t| (w, t))) + { + if let Some(screenshot_buffer) = &window.screenshot_buffer { + let width = window.physical_width; + let height = window.physical_height; + let texture_format = window.swap_chain_texture_format.unwrap(); + let texture = &texture.get_surface_texture().unwrap().texture; + + encoder.copy_texture_to_buffer( + texture.as_image_copy(), + wgpu::ImageCopyBuffer { + buffer: &screenshot_buffer, + layout: crate::view::screenshot::layout_data(width, height, texture_format), + }, + Extent3d { + width, + height, + ..Default::default() + }, + ); + } + } +} + +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 buffer = window.screenshot_buffer.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_slice); + 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(); + } + } +} From 539370ce6a63723773369bc1e7bc0c69e27124d4 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Wed, 11 Jan 2023 23:12:39 +0300 Subject: [PATCH 02/17] fix clippy --- crates/bevy_render/src/render_resource/texture.rs | 2 +- crates/bevy_render/src/view/window/screenshot.rs | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/crates/bevy_render/src/render_resource/texture.rs b/crates/bevy_render/src/render_resource/texture.rs index 108f5e0259393..78d8214e112c7 100644 --- a/crates/bevy_render/src/render_resource/texture.rs +++ b/crates/bevy_render/src/render_resource/texture.rs @@ -98,7 +98,7 @@ impl TextureView { pub(crate) fn get_surface_texture(&self) -> Option<&wgpu::SurfaceTexture> { match &self.value { TextureViewValue::TextureView(_) => None, - TextureViewValue::SurfaceTexture { texture, .. } => Some(&*texture), + TextureViewValue::SurfaceTexture { texture, .. } => Some(texture), } } } diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index a3831c4fe5d27..ec6f83ae1433b 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -94,7 +94,7 @@ pub(crate) fn submit_screenshot_commands(windows: &ExtractedWindows, encoder: &m encoder.copy_texture_to_buffer( texture.as_image_copy(), wgpu::ImageCopyBuffer { - buffer: &screenshot_buffer, + buffer: screenshot_buffer, layout: crate::view::screenshot::layout_data(width, height, texture_format), }, Extent3d { @@ -135,7 +135,6 @@ pub(crate) fn collect_screenshots(world: &mut World) { // 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_slice); drop(buffer); if result.len() != ((width * height) as usize * pixel_size) { From ed06f126deb74a1b6ef827b1b96cd22a793e2bf2 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Thu, 12 Jan 2023 20:15:24 +0300 Subject: [PATCH 03/17] Comment updates --- crates/bevy_render/src/texture/image_texture_conversion.rs | 2 ++ crates/bevy_render/src/view/window.rs | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/crates/bevy_render/src/texture/image_texture_conversion.rs b/crates/bevy_render/src/texture/image_texture_conversion.rs index a8debeae2aab6..2d262a6e050ab 100644 --- a/crates/bevy_render/src/texture/image_texture_conversion.rs +++ b/crates/bevy_render/src/texture/image_texture_conversion.rs @@ -197,6 +197,8 @@ 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, diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index bfd146ecd7fff..3d54d088b4150 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -138,6 +138,10 @@ fn extract_windows( for closed_window in closed.iter() { extracted_windows.remove(&closed_window.id); } + // 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); From a3208de5d1804a659859b30b8d67c5b09762b266 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Thu, 12 Jan 2023 22:33:20 +0300 Subject: [PATCH 04/17] Switch to using custom texture for screenshots --- crates/bevy_render/src/camera/camera.rs | 2 +- .../src/camera/camera_driver_node.rs | 2 +- .../src/render_resource/texture.rs | 75 ++++------- crates/bevy_render/src/renderer/mod.rs | 7 +- crates/bevy_render/src/view/window.rs | 100 ++++++++++---- .../bevy_render/src/view/window/screenshot.rs | 127 ++++++++++++++++-- .../src/view/window/screenshot.wgsl | 14 ++ 7 files changed, 230 insertions(+), 97 deletions(-) create mode 100644 crates/bevy_render/src/view/window/screenshot.wgsl diff --git a/crates/bevy_render/src/camera/camera.rs b/crates/bevy_render/src/camera/camera.rs index 72bcfa3556df0..e5927d74f960c 100644 --- a/crates/bevy_render/src/camera/camera.rs +++ b/crates/bevy_render/src/camera/camera.rs @@ -329,7 +329,7 @@ impl RenderTarget { match self { RenderTarget::Window(window_id) => windows .get(window_id) - .and_then(|window| window.swap_chain_texture.as_ref()), + .and_then(|window| window.swap_chain_texture_view.as_ref()), RenderTarget::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 f57929f30caeb..024c02e0c0474 100644 --- a/crates/bevy_render/src/camera/camera_driver_node.rs +++ b/crates/bevy_render/src/camera/camera_driver_node.rs @@ -78,7 +78,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 78d8214e112c7..df0df616a6f6c 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,42 +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(), - } - } - - #[inline] - pub(crate) fn get_surface_texture(&self) -> Option<&wgpu::SurfaceTexture> { - match &self.value { - TextureViewValue::TextureView(_) => None, - TextureViewValue::SurfaceTexture { texture, .. } => Some(texture), - } - } } 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), } } } @@ -129,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/mod.rs b/crates/bevy_render/src/renderer/mod.rs index 14c456c993c77..c5e5f31216f5a 100644 --- a/crates/bevy_render/src/renderer/mod.rs +++ b/crates/bevy_render/src/renderer/mod.rs @@ -27,7 +27,6 @@ pub fn render_system(world: &mut World) { let graph = world.resource::(); let render_device = world.resource::(); let render_queue = world.resource::(); - let windows = world.resource::(); if let Err(e) = RenderGraphRunner::run( graph, @@ -35,7 +34,7 @@ pub fn render_system(world: &mut World) { &render_queue.0, world, |encoder| { - crate::view::screenshot::submit_screenshot_commands(windows, encoder); + crate::view::screenshot::submit_screenshot_commands(world, encoder); }, ) { error!("Error running render graph:"); @@ -68,8 +67,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(); } } diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 3d54d088b4150..96b924090b8f7 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -1,7 +1,7 @@ pub mod screenshot; use crate::{ - render_resource::{Buffer, TextureView}, + render_resource::{SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, texture::TextureFormatPixelInfo, Extract, RenderApp, RenderStage, @@ -13,9 +13,9 @@ use bevy_window::{ CompositeAlphaMode, PresentMode, RawHandleWrapper, WindowClosed, WindowId, Windows, }; use std::ops::{Deref, DerefMut}; -use wgpu::{BufferUsages, TextureFormat}; +use wgpu::{BufferUsages, TextureFormat, TextureUsages}; -use self::screenshot::ScreenshotManager; +use self::screenshot::{ScreenshotGpuMemory, ScreenshotManager, ScreenshotToScreenPipeline}; /// Token to ensure a system runs on the main thread. #[derive(Resource, Default)] @@ -34,6 +34,7 @@ impl Plugin for WindowRenderPlugin { render_app .init_resource::() .init_resource::() + .init_resource::() .init_non_send_resource::() .add_system_to_stage(RenderStage::Extract, extract_windows) .add_system_to_stage( @@ -50,13 +51,17 @@ 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, - pub screenshot_buffer: Option, } #[derive(Default, Resource)] @@ -100,17 +105,18 @@ fn extract_windows( physical_width: new_width, physical_height: new_height, present_mode: window.present_mode(), + swap_chain_texture_view: None, swap_chain_texture: None, + screenshot_memory: None, swap_chain_texture_format: None, size_changed: false, present_mode_changed: false, alpha_mode: window.alpha_mode(), screenshot_func: None, - screenshot_buffer: 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 = new_present_mode != extracted_window.present_mode; @@ -191,6 +197,7 @@ pub fn prepare_windows( render_device: Res, render_instance: Res, render_adapter: Res, + screenshot_pipeline: Res, ) { for window in windows .windows @@ -217,7 +224,7 @@ pub fn prepare_windows( SurfaceData { surface, format } }); - let mut surface_configuration = wgpu::SurfaceConfiguration { + let surface_configuration = wgpu::SurfaceConfiguration { format: surface_data.format, width: window.physical_width, height: window.physical_height, @@ -237,19 +244,6 @@ pub fn prepare_windows( CompositeAlphaMode::Inherit => wgpu::CompositeAlphaMode::Inherit, }, }; - if window.screenshot_func.is_some() { - surface_configuration.usage |= wgpu::TextureUsages::COPY_SRC; - window.screenshot_buffer = Some(render_device.create_buffer(&wgpu::BufferDescriptor { - label: None, - 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, - })); - } // A recurring issue is hitting `wgpu::SurfaceError::Timeout` on certain Linux // mesa driver implementations. This seems to be a quirk of some drivers. @@ -271,27 +265,32 @@ pub fn prepare_windows( let not_already_configured = window_surfaces.configured_windows.insert(window.id); let surface = &surface_data.surface; - if not_already_configured - || window.size_changed - || window.present_mode_changed - || window.screenshot_func.is_some() - { + if not_already_configured || window.size_changed || window.present_mode_changed { render_device.configure_surface(surface, &surface_configuration); let frame = surface .get_current_texture() .expect("Error configuring surface"); - window.swap_chain_texture = Some(TextureView::from(frame)); + window.swap_chain_texture_view = Some(TextureView::from( + frame.texture.create_view(&Default::default()), + )); + window.swap_chain_texture = Some(SurfaceTexture::from(frame)); } else { match surface.get_current_texture() { Ok(frame) => { - window.swap_chain_texture = Some(TextureView::from(frame)); + window.swap_chain_texture_view = Some(TextureView::from( + frame.texture.create_view(&Default::default()), + )); + window.swap_chain_texture = Some(SurfaceTexture::from(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.swap_chain_texture_view = Some(TextureView::from( + frame.texture.create_view(&Default::default()), + )); + window.swap_chain_texture = Some(SurfaceTexture::from(frame)); } #[cfg(target_os = "linux")] Err(wgpu::SurfaceError::Timeout) if may_erroneously_timeout() => { @@ -306,5 +305,48 @@ 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: None, + 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, + }); + let texture_view = texture.create_view(&Default::default()); + let buffer = render_device.create_buffer(&wgpu::BufferDescriptor { + label: None, + 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: None, + layout: &screenshot_pipeline.bind_group_layout, + entries: &[wgpu::BindGroupEntry { + binding: 0, + resource: wgpu::BindingResource::TextureView(&texture_view), + }], + }); + window.swap_chain_texture_view = Some(texture_view); + window.screenshot_memory = Some(ScreenshotGpuMemory { + texture, + buffer, + bind_group, + }) + } } } diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index ec6f83ae1433b..9384890bc35c6 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -1,4 +1,4 @@ -use std::{num::NonZeroU32, path::Path}; +use std::{borrow::Cow, num::NonZeroU32, path::Path}; use bevy_ecs::prelude::*; use bevy_log::info_span; @@ -11,7 +11,12 @@ use wgpu::{ CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, }; -use crate::{prelude::Image, texture::TextureFormatPixelInfo}; +use crate::{ + prelude::Image, + render_resource::{BindGroup, BindGroupLayout, Buffer, RenderPipeline, Texture}, + renderer::RenderDevice, + texture::TextureFormatPixelInfo, +}; use super::ExtractedWindows; @@ -80,21 +85,98 @@ pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> Ima } } -pub(crate) fn submit_screenshot_commands(windows: &ExtractedWindows, encoder: &mut CommandEncoder) { - for (window, texture) in windows - .values() - .filter_map(|w| w.swap_chain_texture.as_ref().map(|t| (w, t))) - { - if let Some(screenshot_buffer) = &window.screenshot_buffer { +#[derive(Resource)] +pub struct ScreenshotToScreenPipeline { + pub pipeline: RenderPipeline, + pub bind_group_layout: BindGroupLayout, +} + +impl FromWorld for ScreenshotToScreenPipeline { + fn from_world(render_world: &mut World) -> Self { + let device = render_world.resource::(); + + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: None, + source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("screenshot.wgsl"))), + }); + + 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, + }], + }); + + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("screenshot-to-screen-pgl"), + bind_group_layouts: &[&bind_group_layout], + push_constant_ranges: &[], + }); + + let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { + label: Some("screenshot-to-screen"), + layout: Some(&pipeline_layout), + vertex: wgpu::VertexState { + buffers: &[], + entry_point: "vs_main", + module: &module, + }, + 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(wgpu::FragmentState { + module: &module, + entry_point: "fs_main", + targets: &[Some(wgpu::ColorTargetState { + format: TextureFormat::Bgra8UnormSrgb, + blend: None, + write_mask: wgpu::ColorWrites::ALL, + })], + }), + multiview: None, + }); + + Self { + pipeline, + bind_group_layout, + } + } +} + +pub struct ScreenshotGpuMemory { + pub texture: Texture, + pub buffer: Buffer, + pub bind_group: BindGroup, +} + +pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) { + let windows = world.resource::(); + let pipeline = 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(); - let texture = &texture.get_surface_texture().unwrap().texture; encoder.copy_texture_to_buffer( - texture.as_image_copy(), + memory.texture.as_image_copy(), wgpu::ImageCopyBuffer { - buffer: screenshot_buffer, + buffer: &memory.buffer, layout: crate::view::screenshot::layout_data(width, height, texture_format), }, Extent3d { @@ -103,6 +185,27 @@ pub(crate) fn submit_screenshot_commands(windows: &ExtractedWindows, encoder: &m ..Default::default() }, ); + 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.pipeline); + pass.set_bind_group(0, &memory.bind_group, &[]); + pass.draw(0..3, 0..1); } } } @@ -117,7 +220,7 @@ pub(crate) fn collect_screenshots(world: &mut World) { let height = window.physical_height; let texture_format = window.swap_chain_texture_format.unwrap(); let pixel_size = texture_format.pixel_size(); - let buffer = window.screenshot_buffer.take().unwrap(); + let ScreenshotGpuMemory { buffer, .. } = window.screenshot_memory.take().unwrap(); let finish = async move { let (tx, rx) = async_channel::bounded(1); 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 0000000000000..67a6f3c58af22 --- /dev/null +++ b/crates/bevy_render/src/view/window/screenshot.wgsl @@ -0,0 +1,14 @@ +@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); +} From 6707d1967acdd1811b171297bea4bfd9fa0b8c4b Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Thu, 12 Jan 2023 22:35:59 +0300 Subject: [PATCH 05/17] appease clippy --- crates/bevy_render/src/view/window.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 96b924090b8f7..035a761ebc0d2 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -346,7 +346,7 @@ pub fn prepare_windows( texture, buffer, bind_group, - }) + }); } } } From 6487690a003d1222b9a21e2baa4aa3c770b738fb Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 13 Jan 2023 17:08:03 +0300 Subject: [PATCH 06/17] Adress more feedback --- crates/bevy_render/src/view/window.rs | 30 +++++++++---------- .../bevy_render/src/view/window/screenshot.rs | 14 +++++++-- .../src/view/window/screenshot.wgsl | 2 ++ 3 files changed, 28 insertions(+), 18 deletions(-) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 035a761ebc0d2..4f8c93e72b947 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -64,6 +64,15 @@ pub struct ExtractedWindow { 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)] pub struct ExtractedWindows { pub windows: HashMap, @@ -270,27 +279,18 @@ pub fn prepare_windows( let frame = surface .get_current_texture() .expect("Error configuring surface"); - window.swap_chain_texture_view = Some(TextureView::from( - frame.texture.create_view(&Default::default()), - )); - window.swap_chain_texture = Some(SurfaceTexture::from(frame)); + window.set_swapchain_texture(frame); } else { match surface.get_current_texture() { Ok(frame) => { - window.swap_chain_texture_view = Some(TextureView::from( - frame.texture.create_view(&Default::default()), - )); - window.swap_chain_texture = Some(SurfaceTexture::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_view = Some(TextureView::from( - frame.texture.create_view(&Default::default()), - )); - window.swap_chain_texture = Some(SurfaceTexture::from(frame)); + window.set_swapchain_texture(frame); } #[cfg(target_os = "linux")] Err(wgpu::SurfaceError::Timeout) if may_erroneously_timeout() => { @@ -308,7 +308,7 @@ pub fn prepare_windows( if window.screenshot_func.is_some() { let texture = render_device.create_texture(&wgpu::TextureDescriptor { - label: None, + label: Some("screenshot-capture-rendertarget"), size: wgpu::Extent3d { width: surface_configuration.width, height: surface_configuration.height, @@ -324,7 +324,7 @@ pub fn prepare_windows( }); let texture_view = texture.create_view(&Default::default()); let buffer = render_device.create_buffer(&wgpu::BufferDescriptor { - label: None, + label: Some("screenshot-transfer-buffer"), size: screenshot::get_aligned_size( window.physical_width, window.physical_height, @@ -334,7 +334,7 @@ pub fn prepare_windows( mapped_at_creation: false, }); let bind_group = render_device.create_bind_group(&wgpu::BindGroupDescriptor { - label: None, + label: Some("screenshot-to-screen-bind-group"), layout: &screenshot_pipeline.bind_group_layout, entries: &[wgpu::BindGroupEntry { binding: 0, diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index 9384890bc35c6..56556b2829830 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -1,7 +1,7 @@ use std::{borrow::Cow, num::NonZeroU32, path::Path}; use bevy_ecs::prelude::*; -use bevy_log::info_span; +use bevy_log::{error, info_span}; use bevy_tasks::AsyncComputeTaskPool; use bevy_utils::HashMap; use bevy_window::WindowId; @@ -58,8 +58,16 @@ impl ScreenshotManager { path: impl AsRef, ) -> Result<(), ScreenshotAlreadyRequestedError> { let path = path.as_ref().to_owned(); - self.take_screenshot(window, |image| { - image.try_into_dynamic().unwrap().save(path).unwrap(); + self.take_screenshot(window, move |img| match img.try_into_dynamic() { + Ok(dyn_img) => match image::ImageFormat::from_path(&path) { + Ok(format) => { + if let Err(e) = dyn_img.save_with_format(path, format) { + 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}"), }) } } diff --git a/crates/bevy_render/src/view/window/screenshot.wgsl b/crates/bevy_render/src/view/window/screenshot.wgsl index 67a6f3c58af22..2743fa1d950bd 100644 --- a/crates/bevy_render/src/view/window/screenshot.wgsl +++ b/crates/bevy_render/src/view/window/screenshot.wgsl @@ -1,3 +1,5 @@ +// 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); From bfcd4730b28e53407e812f76dc6a2b6dee360948 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 13 Jan 2023 17:43:00 +0300 Subject: [PATCH 07/17] Adress WASM panic --- crates/bevy_render/src/lib.rs | 1 - crates/bevy_render/src/view/window.rs | 24 +++- .../bevy_render/src/view/window/screenshot.rs | 136 +++++++++++------- 3 files changed, 99 insertions(+), 62 deletions(-) diff --git a/crates/bevy_render/src/lib.rs b/crates/bevy_render/src/lib.rs index 06326d69ad0ad..0ba66415c35c8 100644 --- a/crates/bevy_render/src/lib.rs +++ b/crates/bevy_render/src/lib.rs @@ -137,7 +137,6 @@ impl Plugin for RenderPlugin { app.add_asset::() .add_debug_asset::() .init_asset_loader::() - .init_resource::() .init_debug_asset_loader::(); if let Some(backends) = self.wgpu_settings.backends { diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 4f8c93e72b947..c47c6cc0a132d 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -1,7 +1,5 @@ -pub mod screenshot; - use crate::{ - render_resource::{SurfaceTexture, TextureView}, + render_resource::{PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, texture::TextureFormatPixelInfo, Extract, RenderApp, RenderStage, @@ -15,7 +13,11 @@ use bevy_window::{ use std::ops::{Deref, DerefMut}; use wgpu::{BufferUsages, TextureFormat, TextureUsages}; -use self::screenshot::{ScreenshotGpuMemory, ScreenshotManager, ScreenshotToScreenPipeline}; +pub mod screenshot; + +use screenshot::{ + ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline, +}; /// Token to ensure a system runs on the main thread. #[derive(Resource, Default)] @@ -30,6 +32,8 @@ 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::() @@ -57,7 +61,7 @@ pub struct ExtractedWindow { pub swap_chain_texture_view: Option, pub swap_chain_texture: Option, pub swap_chain_texture_format: Option, - pub screenshot_memory: Option, + pub screenshot_memory: Option, pub size_changed: bool, pub present_mode_changed: bool, pub alpha_mode: CompositeAlphaMode, @@ -207,6 +211,8 @@ pub fn prepare_windows( render_instance: Res, render_adapter: Res, screenshot_pipeline: Res, + mut pipeline_cache: ResMut, + mut pipelines: ResMut>, ) { for window in windows .windows @@ -341,11 +347,17 @@ pub fn prepare_windows( resource: wgpu::BindingResource::TextureView(&texture_view), }], }); + let pipeline_id = pipelines.specialize( + &mut pipeline_cache, + &screenshot_pipeline, + surface_configuration.format, + ); window.swap_chain_texture_view = Some(texture_view); - window.screenshot_memory = Some(ScreenshotGpuMemory { + 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 index 56556b2829830..b5a9834f08ab8 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -1,7 +1,10 @@ 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_span}; +use bevy_reflect::TypeUuid; use bevy_tasks::AsyncComputeTaskPool; use bevy_utils::HashMap; use bevy_window::WindowId; @@ -12,10 +15,15 @@ use wgpu::{ }; use crate::{ - prelude::Image, - render_resource::{BindGroup, BindGroupLayout, Buffer, RenderPipeline, Texture}, + prelude::{Image, Shader}, + render_resource::{ + BindGroup, BindGroupLayout, Buffer, CachedRenderPipelineId, FragmentState, PipelineCache, + RenderPipelineDescriptor, SpecializedRenderPipeline, SpecializedRenderPipelines, Texture, + VertexState, + }, renderer::RenderDevice, texture::TextureFormatPixelInfo, + RenderApp, }; use super::ExtractedWindows; @@ -72,6 +80,28 @@ impl ScreenshotManager { } } +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)) } @@ -95,7 +125,6 @@ pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> Ima #[derive(Resource)] pub struct ScreenshotToScreenPipeline { - pub pipeline: RenderPipeline, pub bind_group_layout: BindGroupLayout, } @@ -103,11 +132,6 @@ impl FromWorld for ScreenshotToScreenPipeline { fn from_world(render_world: &mut World) -> Self { let device = render_world.resource::(); - let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { - label: None, - source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("screenshot.wgsl"))), - }); - let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { label: Some("screenshot-to-screen-bgl"), entries: &[wgpu::BindGroupLayoutEntry { @@ -122,19 +146,22 @@ impl FromWorld for ScreenshotToScreenPipeline { }], }); - let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { - label: Some("screenshot-to-screen-pgl"), - bind_group_layouts: &[&bind_group_layout], - push_constant_ranges: &[], - }); + Self { bind_group_layout } + } +} - let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor { - label: Some("screenshot-to-screen"), - layout: Some(&pipeline_layout), - vertex: wgpu::VertexState { - buffers: &[], - entry_point: "vs_main", - module: &module, +impl SpecializedRenderPipeline for ScreenshotToScreenPipeline { + type Key = TextureFormat; + + fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor { + RenderPipelineDescriptor { + label: Some(Cow::Borrowed("screenshot-to-screen")), + layout: Some(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, @@ -147,34 +174,31 @@ impl FromWorld for ScreenshotToScreenPipeline { }, depth_stencil: None, multisample: wgpu::MultisampleState::default(), - fragment: Some(wgpu::FragmentState { - module: &module, - entry_point: "fs_main", - targets: &[Some(wgpu::ColorTargetState { - format: TextureFormat::Bgra8UnormSrgb, + 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, })], }), - multiview: None, - }); - - Self { - pipeline, - bind_group_layout, } } } -pub struct ScreenshotGpuMemory { +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 pipeline = world.resource::(); + let pipelines = world.resource::(); + for window in windows.values() { if let Some(memory) = &window.screenshot_memory { let width = window.physical_width; @@ -193,27 +217,29 @@ pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEnc ..Default::default() }, ); - 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.pipeline); - pass.set_bind_group(0, &memory.bind_group, &[]); - pass.draw(0..3, 0..1); + 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); + } } } } @@ -228,7 +254,7 @@ pub(crate) fn collect_screenshots(world: &mut World) { let height = window.physical_height; let texture_format = window.swap_chain_texture_format.unwrap(); let pixel_size = texture_format.pixel_size(); - let ScreenshotGpuMemory { buffer, .. } = window.screenshot_memory.take().unwrap(); + let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap(); let finish = async move { let (tx, rx) = async_channel::bounded(1); From 6ba02287e3aa228ff427cd46bde34895f3047a50 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 13 Jan 2023 17:44:12 +0300 Subject: [PATCH 08/17] clippy --- crates/bevy_render/src/view/window/screenshot.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index b5a9834f08ab8..5a4c075ace88a 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -70,7 +70,7 @@ impl ScreenshotManager { Ok(dyn_img) => match image::ImageFormat::from_path(&path) { Ok(format) => { if let Err(e) = dyn_img.save_with_format(path, format) { - error!("Cannot save screenshot, IO error: {e}") + error!("Cannot save screenshot, IO error: {e}"); } } Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"), From 90f5fed2ea133c10ae20423825668776fdcfa9a2 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 13 Jan 2023 18:01:21 +0300 Subject: [PATCH 09/17] fix hdr issue --- crates/bevy_render/src/view/window/screenshot.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index 5a4c075ace88a..a7a3329292e26 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -69,7 +69,10 @@ impl ScreenshotManager { self.take_screenshot(window, move |img| match img.try_into_dynamic() { Ok(dyn_img) => match image::ImageFormat::from_path(&path) { Ok(format) => { - if let Err(e) = dyn_img.save_with_format(path, 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(); + if let Err(e) = img.save_with_format(path, format) { error!("Cannot save screenshot, IO error: {e}"); } } From baabde3c946708a8c67d2546071a0972e611d1c0 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 13 Jan 2023 18:02:45 +0300 Subject: [PATCH 10/17] clippy, again --- crates/bevy_render/src/view/window.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index c47c6cc0a132d..e3dfe02861040 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -201,6 +201,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 From 1941274274325df23dac63cddec32e91e8386d42 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 13 Jan 2023 19:12:09 +0300 Subject: [PATCH 11/17] Add standalone example --- Cargo.toml | 10 ++++++ examples/README.md | 1 + examples/window/screenshot.rs | 62 +++++++++++++++++++++++++++++++++++ 3 files changed, 73 insertions(+) create mode 100644 examples/window/screenshot.rs diff --git a/Cargo.toml b/Cargo.toml index 48b7e55bb5d88..144b640e97213 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1556,6 +1556,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/examples/README.md b/examples/README.md index 7817465d9e8b8..2f6b4841d313a 100644 --- a/examples/README.md +++ b/examples/README.md @@ -328,6 +328,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 0000000000000..6ccf82317ec5f --- /dev/null +++ b/examples/window/screenshot.rs @@ -0,0 +1,62 @@ +//! An example showing how to save screenshots to disk + +use bevy::prelude::*; +use bevy::render::view::screenshot::ScreenshotManager; + +fn main() { + App::new() + .add_plugins(DefaultPlugins) + .add_startup_system(setup) + .add_system(screenshot_on_f12) + .run(); +} + +fn screenshot_on_f12( + input: Res>, + 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(bevy::window::WindowId::primary(), 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(Mesh::from(shape::Plane { size: 5.0 })), + 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() + }); +} From 86798ad7e01dff5f157c43dc068310e988286c72 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Mon, 23 Jan 2023 16:17:59 +0300 Subject: [PATCH 12/17] appease clippy --- crates/bevy_render/src/view/window.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 30fafff7aad2c..e087eec0343f1 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -216,7 +216,7 @@ pub fn prepare_windows( render_instance: Res, render_adapter: Res, screenshot_pipeline: Res, - mut pipeline_cache: ResMut, + pipeline_cache: Res, mut pipelines: ResMut>, ) { for window in windows.windows.values_mut() { @@ -347,7 +347,7 @@ pub fn prepare_windows( }], }); let pipeline_id = pipelines.specialize( - &mut pipeline_cache, + &pipeline_cache, &screenshot_pipeline, surface_configuration.format, ); From edb1e57a30fac736dcc3325689864ff8a4448e38 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Wed, 1 Feb 2023 02:29:48 +0300 Subject: [PATCH 13/17] fix bug --- crates/bevy_render/src/view/window.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 7e2ef3b889fcb..f09af443525f2 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -360,6 +360,7 @@ pub fn prepare_windows( 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 { From 6e0c39b2150fc10d0c050e6d4d138eb4712e1026 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Wed, 15 Feb 2023 20:32:35 +0300 Subject: [PATCH 14/17] fix stuff --- crates/bevy_render/src/view/window.rs | 2 +- examples/window/screenshot.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index 7c1f903de9dbd..b30be1a973c94 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -1,7 +1,7 @@ use crate::{ render_resource::{PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, - Extract, ExtractSchedule, RenderApp, RenderSet, + Extract, ExtractSchedule, RenderApp, RenderSet, texture::TextureFormatPixelInfo, }; use bevy_app::{App, Plugin}; use bevy_ecs::prelude::*; diff --git a/examples/window/screenshot.rs b/examples/window/screenshot.rs index 5d4e8428253d6..ee96b1b37991b 100644 --- a/examples/window/screenshot.rs +++ b/examples/window/screenshot.rs @@ -35,7 +35,7 @@ fn setup( ) { // plane commands.spawn(PbrBundle { - mesh: meshes.add(Mesh::from(shape::Plane { size: 5.0 })), + mesh: meshes.add(shape::Plane::from_size(5.0).into()), material: materials.add(Color::rgb(0.3, 0.5, 0.3).into()), ..default() }); From ef26a660fd4f88fc194bb53d490231c89dfd59bf Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Wed, 15 Feb 2023 20:34:50 +0300 Subject: [PATCH 15/17] formatting --- crates/bevy_render/src/view/window.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/bevy_render/src/view/window.rs b/crates/bevy_render/src/view/window.rs index b30be1a973c94..156395cc9c3d6 100644 --- a/crates/bevy_render/src/view/window.rs +++ b/crates/bevy_render/src/view/window.rs @@ -1,7 +1,8 @@ use crate::{ render_resource::{PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, - Extract, ExtractSchedule, RenderApp, RenderSet, texture::TextureFormatPixelInfo, + texture::TextureFormatPixelInfo, + Extract, ExtractSchedule, RenderApp, RenderSet, }; use bevy_app::{App, Plugin}; use bevy_ecs::prelude::*; From cf5fd040bd0bfa4928f9b191c58693e35760ad1f Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Fri, 24 Feb 2023 08:32:44 +0300 Subject: [PATCH 16/17] fix things up --- crates/bevy_render/src/view/window/screenshot.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index d253f3feb993f..214a95d8e9efa 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -158,7 +158,7 @@ impl SpecializedRenderPipeline for ScreenshotToScreenPipeline { fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor { RenderPipelineDescriptor { label: Some(Cow::Borrowed("screenshot-to-screen")), - layout: Some(vec![self.bind_group_layout.clone()]), + layout: vec![self.bind_group_layout.clone()], vertex: VertexState { buffers: vec![], shader_defs: vec![], @@ -186,6 +186,7 @@ impl SpecializedRenderPipeline for ScreenshotToScreenPipeline { write_mask: wgpu::ColorWrites::ALL, })], }), + push_constant_ranges: Vec::new(), } } } From 23bad7a91c56b2b550e6231cbf9da0d8b3e16d16 Mon Sep 17 00:00:00 2001 From: TheRawMeatball Date: Wed, 19 Apr 2023 07:54:56 +0300 Subject: [PATCH 17/17] print to log when a screenshot is taken --- crates/bevy_render/src/view/window/screenshot.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/bevy_render/src/view/window/screenshot.rs b/crates/bevy_render/src/view/window/screenshot.rs index 214a95d8e9efa..2a67f43005a7a 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -3,7 +3,7 @@ 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_span}; +use bevy_log::{error, info, info_span}; use bevy_reflect::TypeUuid; use bevy_tasks::AsyncComputeTaskPool; use bevy_utils::HashMap; @@ -71,8 +71,9 @@ impl ScreenshotManager { // 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(); - if let Err(e) = img.save_with_format(path, format) { - error!("Cannot save screenshot, IO error: {e}"); + 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}"),