From d9527c101c2f49c6884763cce36ea1d27dd6a597 Mon Sep 17 00:00:00 2001 From: charlotte Date: Sun, 25 Aug 2024 07:14:32 -0700 Subject: [PATCH] Rewrite screenshots. (#14833) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # Objective Rewrite screenshotting to be able to accept any `RenderTarget`. Closes #12478 ## Solution Previously, screenshotting relied on setting a variety of state on the requested window. When extracted, the window's `swap_chain_texture_view` property would be swapped out with a texture_view created that frame for the screenshot pipeline to write back to the cpu. Besides being tightly coupled to window in a way that prevented screenshotting other render targets, this approach had the drawback of relying on the implicit state of `swap_chain_texture_view` being returned from a `NormalizedRenderTarget` when view targets were prepared. Because property is set every frame for windows, that wasn't a problem, but poses a problem for render target images. Namely, to do the equivalent trick, we'd have to replace the `GpuImage`'s texture view, and somehow restore it later. As such, this PR creates a new `prepare_view_textures` system which runs before `prepare_view_targets` that allows a new `prepare_screenshots` system to be sandwiched between and overwrite the render targets texture view if a screenshot has been requested that frame for the given target. Additionally, screenshotting itself has been changed to use a component + observer pattern. We now spawn a `Screenshot` component into the world, whose lifetime is tracked with a series of marker components. When the screenshot is read back to the CPU, we send the image over a channel back to the main world where an observer fires on the screenshot entity before being despawned the next frame. This allows the user to access resources in their save callback that might be useful (e.g. uploading the screenshot over the network, etc.). ## Testing ![image](https://github.com/user-attachments/assets/48f19aed-d9e1-4058-bb17-82b37f992b7b) TODO: - [x] Web - [ ] Manual texture view --- ## Showcase render to texture example: web saving still works: ## Migration Guide `ScreenshotManager` has been removed. To take a screenshot, spawn a `Screenshot` entity with the specified render target and provide an observer targeting the `ScreenshotCaptured` event. See the `window/screenshot` example to see an example. --------- Co-authored-by: Kristoffer Søholm --- crates/bevy_dev_tools/src/ci_testing/mod.rs | 5 +- .../bevy_dev_tools/src/ci_testing/systems.rs | 22 +- crates/bevy_render/src/view/mod.rs | 68 +- crates/bevy_render/src/view/window/mod.rs | 83 +-- .../bevy_render/src/view/window/screenshot.rs | 669 ++++++++++++++---- examples/window/screenshot.rs | 35 +- 6 files changed, 612 insertions(+), 270 deletions(-) diff --git a/crates/bevy_dev_tools/src/ci_testing/mod.rs b/crates/bevy_dev_tools/src/ci_testing/mod.rs index f6510f68a1..59949fa0bf 100644 --- a/crates/bevy_dev_tools/src/ci_testing/mod.rs +++ b/crates/bevy_dev_tools/src/ci_testing/mod.rs @@ -7,6 +7,7 @@ pub use self::config::*; use bevy_app::prelude::*; use bevy_ecs::schedule::IntoSystemConfigs; +use bevy_render::view::screenshot::trigger_screenshots; use bevy_time::TimeUpdateStrategy; use std::time::Duration; @@ -51,7 +52,9 @@ impl Plugin for CiTestingPlugin { .insert_resource(config) .add_systems( Update, - systems::send_events.before(bevy_window::close_when_requested), + systems::send_events + .before(trigger_screenshots) + .before(bevy_window::close_when_requested), ); } } diff --git a/crates/bevy_dev_tools/src/ci_testing/systems.rs b/crates/bevy_dev_tools/src/ci_testing/systems.rs index abb2b28bcc..2b652abf96 100644 --- a/crates/bevy_dev_tools/src/ci_testing/systems.rs +++ b/crates/bevy_dev_tools/src/ci_testing/systems.rs @@ -1,9 +1,8 @@ use super::config::*; use bevy_app::AppExit; use bevy_ecs::prelude::*; -use bevy_render::view::screenshot::ScreenshotManager; -use bevy_utils::tracing::{debug, info, warn}; -use bevy_window::PrimaryWindow; +use bevy_render::view::screenshot::{save_to_disk, Screenshot}; +use bevy_utils::tracing::{debug, info}; pub(crate) fn send_events(world: &mut World, mut current_frame: Local) { let mut config = world.resource_mut::(); @@ -23,21 +22,10 @@ pub(crate) fn send_events(world: &mut World, mut current_frame: Local) { info!("Exiting after {} frames. Test successful!", *current_frame); } CiTestingEvent::Screenshot => { - let mut primary_window_query = - world.query_filtered::>(); - let Ok(main_window) = primary_window_query.get_single(world) else { - warn!("Requesting screenshot, but PrimaryWindow is not available"); - continue; - }; - let Some(mut screenshot_manager) = world.get_resource_mut::() - else { - warn!("Requesting screenshot, but ScreenshotManager is not available"); - continue; - }; let path = format!("./screenshot-{}.png", *current_frame); - screenshot_manager - .save_screenshot_to_disk(main_window, path) - .unwrap(); + world + .spawn(Screenshot::primary_window()) + .observe(save_to_disk(path)); info!("Took a screenshot at frame {}.", *current_frame); } // Custom events are forwarded to the world. diff --git a/crates/bevy_render/src/view/mod.rs b/crates/bevy_render/src/view/mod.rs index f9c4a5c59b..3729f4048f 100644 --- a/crates/bevy_render/src/view/mod.rs +++ b/crates/bevy_render/src/view/mod.rs @@ -5,6 +5,7 @@ use bevy_asset::{load_internal_asset, Handle}; pub use visibility::*; pub use window::*; +use crate::camera::NormalizedRenderTarget; use crate::extract_component::ExtractComponentPlugin; use crate::{ camera::{ @@ -25,12 +26,13 @@ use crate::{ }; use bevy_app::{App, Plugin}; use bevy_color::LinearRgba; +use bevy_derive::{Deref, DerefMut}; use bevy_ecs::prelude::*; use bevy_math::{mat3, vec2, vec3, Mat3, Mat4, UVec4, Vec2, Vec3, Vec4, Vec4Swizzles}; use bevy_reflect::{std_traits::ReflectDefault, Reflect}; use bevy_render_macros::ExtractComponent; use bevy_transform::components::GlobalTransform; -use bevy_utils::HashMap; +use bevy_utils::{hashbrown::hash_map::Entry, HashMap}; use std::{ ops::Range, sync::{ @@ -119,6 +121,9 @@ impl Plugin for ViewPlugin { render_app.add_systems( Render, ( + prepare_view_attachments + .in_set(RenderSet::ManageViews) + .before(prepare_view_targets), prepare_view_targets .in_set(RenderSet::ManageViews) .after(prepare_windows) @@ -132,7 +137,9 @@ impl Plugin for ViewPlugin { fn finish(&self, app: &mut App) { if let Some(render_app) = app.get_sub_app_mut(RenderApp) { - render_app.init_resource::(); + render_app + .init_resource::() + .init_resource::(); } } } @@ -458,6 +465,13 @@ pub struct ViewTarget { out_texture: OutputColorAttachment, } +/// Contains [`OutputColorAttachment`] used for each target present on any view in the current +/// frame, after being prepared by [`prepare_view_attachments`]. Users that want to override +/// the default output color attachment for a specific target can do so by adding a +/// [`OutputColorAttachment`] to this resource before [`prepare_view_targets`] is called. +#[derive(Resource, Default, Deref, DerefMut)] +pub struct ViewTargetAttachments(HashMap); + pub struct PostProcessWrite<'a> { pub source: &'a TextureView, pub destination: &'a TextureView, @@ -794,11 +808,41 @@ struct MainTargetTextures { main_texture: Arc, } -#[allow(clippy::too_many_arguments)] -pub fn prepare_view_targets( - mut commands: Commands, +/// Prepares the view target [`OutputColorAttachment`] for each view in the current frame. +pub fn prepare_view_attachments( windows: Res, images: Res>, + manual_texture_views: Res, + cameras: Query<&ExtractedCamera>, + mut view_target_attachments: ResMut, +) { + view_target_attachments.clear(); + for camera in cameras.iter() { + let Some(target) = &camera.target else { + continue; + }; + + match view_target_attachments.entry(target.clone()) { + Entry::Occupied(_) => {} + Entry::Vacant(entry) => { + let Some(attachment) = target + .get_texture_view(&windows, &images, &manual_texture_views) + .cloned() + .zip(target.get_texture_format(&windows, &images, &manual_texture_views)) + .map(|(view, format)| { + OutputColorAttachment::new(view.clone(), format.add_srgb_suffix()) + }) + else { + continue; + }; + entry.insert(attachment); + } + }; + } +} + +pub fn prepare_view_targets( + mut commands: Commands, clear_color_global: Res, render_device: Res, mut texture_cache: ResMut, @@ -809,24 +853,16 @@ pub fn prepare_view_targets( &CameraMainTextureUsages, &Msaa, )>, - manual_texture_views: Res, + view_target_attachments: Res, ) { let mut textures = HashMap::default(); - let mut output_textures = HashMap::default(); for (entity, camera, view, texture_usage, msaa) in cameras.iter() { let (Some(target_size), Some(target)) = (camera.physical_target_size, &camera.target) else { continue; }; - let Some(out_texture) = output_textures.entry(target.clone()).or_insert_with(|| { - target - .get_texture_view(&windows, &images, &manual_texture_views) - .zip(target.get_texture_format(&windows, &images, &manual_texture_views)) - .map(|(view, format)| { - OutputColorAttachment::new(view.clone(), format.add_srgb_suffix()) - }) - }) else { + let Some(out_attachment) = view_target_attachments.get(target) else { continue; }; @@ -913,7 +949,7 @@ pub fn prepare_view_targets( main_texture: main_textures.main_texture.clone(), main_textures, main_texture_format, - out_texture: out_texture.clone(), + out_texture: out_attachment.clone(), }); } } diff --git a/crates/bevy_render/src/view/window/mod.rs b/crates/bevy_render/src/view/window/mod.rs index 027a15e467..816d8c4e8d 100644 --- a/crates/bevy_render/src/view/window/mod.rs +++ b/crates/bevy_render/src/view/window/mod.rs @@ -1,9 +1,6 @@ use crate::{ - render_resource::{ - BindGroupEntries, PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView, - }, + render_resource::{SurfaceTexture, TextureView}, renderer::{RenderAdapter, RenderDevice, RenderInstance}, - texture::TextureFormatPixelInfo, Extract, ExtractSchedule, Render, RenderApp, RenderSet, WgpuWrapper, }; use bevy_app::{App, Last, Plugin}; @@ -18,21 +15,16 @@ use bevy_winit::CustomCursorCache; use std::{ num::NonZeroU32, ops::{Deref, DerefMut}, - sync::PoisonError, }; use wgpu::{ - BufferUsages, SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages, - TextureViewDescriptor, + SurfaceConfiguration, SurfaceTargetUnsafe, TextureFormat, TextureUsages, TextureViewDescriptor, }; pub mod cursor; pub mod screenshot; -use screenshot::{ - ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline, -}; - use self::cursor::update_cursors; +use screenshot::{ScreenshotPlugin, ScreenshotToScreenPipeline}; pub struct WindowRenderPlugin; @@ -78,11 +70,9 @@ 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 size_changed: bool, pub present_mode_changed: bool, pub alpha_mode: CompositeAlphaMode, - pub screenshot_func: Option, } impl ExtractedWindow { @@ -120,7 +110,6 @@ impl DerefMut for ExtractedWindows { fn extract_windows( mut extracted_windows: ResMut, - screenshot_manager: Extract>, mut closing: Extract>, windows: Extract)>>, mut removed: Extract>, @@ -149,8 +138,6 @@ fn extract_windows( 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 @@ -189,20 +176,6 @@ fn extract_windows( extracted_windows.remove(&removed_window); window_surfaces.remove(&removed_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() - .unwrap_or_else(PoisonError::into_inner) - .drain() - { - if let Some(window) = extracted_windows.get_mut(&window) { - window.screenshot_func = Some(screenshot_func); - } - } } struct SurfaceData { @@ -254,9 +227,6 @@ pub fn prepare_windows( mut windows: ResMut, mut window_surfaces: ResMut, render_device: Res, - screenshot_pipeline: Res, - pipeline_cache: Res, - mut pipelines: ResMut>, #[cfg(target_os = "linux")] render_instance: Res, ) { for window in windows.windows.values_mut() { @@ -340,53 +310,6 @@ pub fn prepare_windows( } }; window.swap_chain_texture_format = Some(surface_data.configuration.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_data.configuration.width, - height: surface_data.configuration.height, - depth_or_array_layers: 1, - }, - mip_level_count: 1, - sample_count: 1, - dimension: wgpu::TextureDimension::D2, - format: surface_data.configuration.format.add_srgb_suffix(), - 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.configuration.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( - "screenshot-to-screen-bind-group", - &screenshot_pipeline.bind_group_layout, - &BindGroupEntries::single(&texture_view), - ); - let pipeline_id = pipelines.specialize( - &pipeline_cache, - &screenshot_pipeline, - surface_data.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 index 21f6f6a6bf..bc53fc9880 100644 --- a/crates/bevy_render/src/view/window/screenshot.rs +++ b/crates/bevy_render/src/view/window/screenshot.rs @@ -1,16 +1,13 @@ -use std::{borrow::Cow, path::Path, sync::PoisonError}; - -use bevy_app::Plugin; -use bevy_asset::{load_internal_asset, Handle}; -use bevy_ecs::{entity::EntityHashMap, prelude::*}; -use bevy_tasks::AsyncComputeTaskPool; -use bevy_utils::tracing::{error, info, info_span}; -use std::sync::Mutex; -use thiserror::Error; -use wgpu::{ - CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, +use super::ExtractedWindows; +use crate::camera::{ + ManualTextureViewHandle, ManualTextureViews, NormalizedRenderTarget, RenderTarget, +}; +use crate::render_asset::RenderAssets; +use crate::render_resource::{BindGroupEntries, BufferUsages, TextureUsages, TextureView}; +use crate::texture::{GpuImage, OutputColorAttachment}; +use crate::view::{ + prepare_view_attachments, prepare_view_targets, ViewTargetAttachments, WindowSurfaces, }; - use crate::{ prelude::{Image, Shader}, render_asset::RenderAssetUsages, @@ -21,51 +18,116 @@ use crate::{ }, renderer::RenderDevice, texture::TextureFormatPixelInfo, - RenderApp, + ExtractSchedule, MainWorld, Render, RenderApp, RenderSet, +}; +use bevy_app::{First, Plugin, Update}; +use bevy_asset::{load_internal_asset, Handle}; +use bevy_derive::{Deref, DerefMut}; +use bevy_ecs::event::event_update_system; +use bevy_ecs::system::SystemState; +use bevy_ecs::{entity::EntityHashMap, prelude::*}; +use bevy_hierarchy::DespawnRecursiveExt; +use bevy_reflect::Reflect; +use bevy_tasks::AsyncComputeTaskPool; +use bevy_utils::tracing::{error, info, warn}; +use bevy_utils::{default, HashSet}; +use bevy_window::{PrimaryWindow, WindowRef}; +use std::ops::Deref; +use std::sync::mpsc::{Receiver, Sender}; +use std::sync::{Arc, Mutex}; +use std::{borrow::Cow, path::Path}; +use wgpu::{ + CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT, }; -use super::ExtractedWindows; +#[derive(Event, Deref, DerefMut, Reflect, Debug)] +#[reflect(Debug)] +pub struct ScreenshotCaptured(pub Image); -pub type ScreenshotFn = Box; +/// A component that signals to the renderer to capture a screenshot this frame. +/// +/// This component should be spawned on a new entity with an observer that will trigger +/// with [`ScreenshotCaptured`] when the screenshot is ready. +/// +/// Screenshots are captured asynchronously and may not be available immediately after the frame +/// that the component is spawned on. The observer should be used to handle the screenshot when it +/// is ready. +/// +/// Note that the screenshot entity will be despawned after the screenshot is captured and the +/// observer is triggered. +/// +/// # Usage +/// +/// ``` +/// # use bevy_ecs::prelude::*; +/// # use bevy_render::view::screenshot::{save_to_disk, Screenshot}; +/// +/// fn take_screenshot(mut commands: Commands) { +/// commands.spawn(Screenshot::primary_window()) +/// .observe(save_to_disk("screenshot.png")); +/// } +/// ``` +#[derive(Component, Deref, DerefMut, Reflect, Debug)] +#[reflect(Component, Debug)] +pub struct Screenshot(pub RenderTarget); -/// 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>, -} +/// A marker component that indicates that a screenshot is currently being captured. +#[derive(Component)] +pub struct Capturing; -#[derive(Error, Debug)] -#[error("A screenshot for this window has already been requested.")] -pub struct ScreenshotAlreadyRequestedError; +/// A marker component that indicates that a screenshot has been captured, the image is ready, and +/// the screenshot entity can be despawned. +#[derive(Component)] +pub struct Captured; -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() - .unwrap_or_else(PoisonError::into_inner) - .try_insert(window, Box::new(callback)) - .map(|_| ()) - .map_err(|_| ScreenshotAlreadyRequestedError) +impl Screenshot { + /// Capture a screenshot of the provided window entity. + pub fn window(window: Entity) -> Self { + Self(RenderTarget::Window(WindowRef::Entity(window))) } - /// 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() { + /// Capture a screenshot of the primary window, if one exists. + pub fn primary_window() -> Self { + Self(RenderTarget::Window(WindowRef::Primary)) + } + + /// Capture a screenshot of the provided render target image. + pub fn image(image: Handle) -> Self { + Self(RenderTarget::Image(image)) + } + + /// Capture a screenshot of the provided manual texture view. + pub fn texture_view(texture_view: ManualTextureViewHandle) -> Self { + Self(RenderTarget::TextureView(texture_view)) + } +} + +struct ScreenshotPreparedState { + pub texture: Texture, + pub buffer: Buffer, + pub bind_group: BindGroup, + pub pipeline_id: CachedRenderPipelineId, + pub size: Extent3d, +} + +#[derive(Resource, Deref, DerefMut)] +pub struct CapturedScreenshots(pub Arc>>); + +#[derive(Resource, Deref, DerefMut, Default)] +struct RenderScreenshotTargets(EntityHashMap); + +#[derive(Resource, Deref, DerefMut, Default)] +struct RenderScreenshotsPrepared(EntityHashMap); + +#[derive(Resource, Deref, DerefMut)] +struct RenderScreenshotsSender(Sender<(Entity, Image)>); + +/// Saves the captured screenshot to disk at the provided path. +pub fn save_to_disk(path: impl AsRef) -> impl FnMut(Trigger) { + let path = path.as_ref().to_owned(); + move |trigger| { + let img = trigger.event().deref().clone(); + 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 @@ -118,17 +180,238 @@ impl ScreenshotManager { Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"), }, Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"), - }) + } } } +fn clear_screenshots(mut commands: Commands, screenshots: Query>) { + for entity in screenshots.iter() { + commands.entity(entity).despawn_recursive(); + } +} + +pub fn trigger_screenshots( + mut commands: Commands, + captured_screenshots: ResMut, +) { + let captured_screenshots = captured_screenshots.lock().unwrap(); + while let Ok((entity, image)) = captured_screenshots.try_recv() { + commands.entity(entity).insert(Captured); + commands.trigger_targets(ScreenshotCaptured(image), entity); + } +} + +fn extract_screenshots( + mut targets: ResMut, + mut main_world: ResMut, + mut system_state: Local< + Option< + SystemState<( + Commands, + Query>, + Query<(Entity, &Screenshot), Without>, + )>, + >, + >, + mut seen_targets: Local>, +) { + if system_state.is_none() { + *system_state = Some(SystemState::new(&mut main_world)); + } + let system_state = system_state.as_mut().unwrap(); + let (mut commands, primary_window, screenshots) = system_state.get_mut(&mut main_world); + + targets.clear(); + seen_targets.clear(); + + let primary_window = primary_window.iter().next(); + + for (entity, screenshot) in screenshots.iter() { + let render_target = screenshot.0.clone(); + let Some(render_target) = render_target.normalize(primary_window) else { + warn!( + "Unknown render target for screenshot, skipping: {:?}", + render_target + ); + continue; + }; + if seen_targets.contains(&render_target) { + warn!( + "Duplicate render target for screenshot, skipping entity {:?}: {:?}", + entity, render_target + ); + // If we don't despawn the entity here, it will be captured again in the next frame + commands.entity(entity).despawn_recursive(); + continue; + } + seen_targets.insert(render_target.clone()); + targets.insert(entity, render_target); + commands.entity(entity).insert(Capturing); + } + + system_state.apply(&mut main_world); +} + +#[allow(clippy::too_many_arguments)] +fn prepare_screenshots( + targets: Res, + mut prepared: ResMut, + window_surfaces: Res, + render_device: Res, + screenshot_pipeline: Res, + pipeline_cache: Res, + mut pipelines: ResMut>, + images: Res>, + manual_texture_views: Res, + mut view_target_attachments: ResMut, +) { + prepared.clear(); + for (entity, target) in targets.iter() { + match target { + NormalizedRenderTarget::Window(window) => { + let window = window.entity(); + let Some(surface_data) = window_surfaces.surfaces.get(&window) else { + warn!("Unknown window for screenshot, skipping: {:?}", window); + continue; + }; + let format = surface_data.configuration.format.add_srgb_suffix(); + let size = Extent3d { + width: surface_data.configuration.width, + height: surface_data.configuration.height, + ..default() + }; + let (texture_view, state) = prepare_screenshot_state( + size, + format, + &render_device, + &screenshot_pipeline, + &pipeline_cache, + &mut pipelines, + ); + prepared.insert(*entity, state); + view_target_attachments.insert( + target.clone(), + OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()), + ); + } + NormalizedRenderTarget::Image(image) => { + let Some(gpu_image) = images.get(image) else { + warn!("Unknown image for screenshot, skipping: {:?}", image); + continue; + }; + let format = gpu_image.texture_format; + let size = Extent3d { + width: gpu_image.size.x, + height: gpu_image.size.y, + ..default() + }; + let (texture_view, state) = prepare_screenshot_state( + size, + format, + &render_device, + &screenshot_pipeline, + &pipeline_cache, + &mut pipelines, + ); + prepared.insert(*entity, state); + view_target_attachments.insert( + target.clone(), + OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()), + ); + } + NormalizedRenderTarget::TextureView(texture_view) => { + let Some(manual_texture_view) = manual_texture_views.get(texture_view) else { + warn!( + "Unknown manual texture view for screenshot, skipping: {:?}", + texture_view + ); + continue; + }; + let format = manual_texture_view.format; + let size = Extent3d { + width: manual_texture_view.size.x, + height: manual_texture_view.size.y, + ..default() + }; + let (texture_view, state) = prepare_screenshot_state( + size, + format, + &render_device, + &screenshot_pipeline, + &pipeline_cache, + &mut pipelines, + ); + prepared.insert(*entity, state); + view_target_attachments.insert( + target.clone(), + OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()), + ); + } + } + } +} + +fn prepare_screenshot_state( + size: Extent3d, + format: TextureFormat, + render_device: &RenderDevice, + pipeline: &ScreenshotToScreenPipeline, + pipeline_cache: &PipelineCache, + pipelines: &mut SpecializedRenderPipelines, +) -> (TextureView, ScreenshotPreparedState) { + let texture = render_device.create_texture(&wgpu::TextureDescriptor { + label: Some("screenshot-capture-rendertarget"), + size, + mip_level_count: 1, + sample_count: 1, + dimension: wgpu::TextureDimension::D2, + 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: get_aligned_size(size.width, size.height, 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( + "screenshot-to-screen-bind-group", + &pipeline.bind_group_layout, + &BindGroupEntries::single(&texture_view), + ); + let pipeline_id = pipelines.specialize(pipeline_cache, pipeline, format); + + ( + texture_view, + ScreenshotPreparedState { + texture, + buffer, + bind_group, + pipeline_id, + size, + }, + ) +} + pub struct ScreenshotPlugin; const SCREENSHOT_SHADER_HANDLE: Handle = Handle::weak_from_u128(11918575842344596158); impl Plugin for ScreenshotPlugin { fn build(&self, app: &mut bevy_app::App) { - app.init_resource::(); + app.add_systems( + First, + clear_screenshots + .after(event_update_system) + .before(apply_deferred), + ) + .add_systems(Update, trigger_screenshots) + .register_type::() + .register_type::(); load_internal_asset!( app, @@ -139,8 +422,23 @@ impl Plugin for ScreenshotPlugin { } fn finish(&self, app: &mut bevy_app::App) { + let (tx, rx) = std::sync::mpsc::channel(); + app.insert_resource(CapturedScreenshots(Arc::new(Mutex::new(rx)))); + if let Some(render_app) = app.get_sub_app_mut(RenderApp) { - render_app.init_resource::>(); + render_app + .insert_resource(RenderScreenshotsSender(tx)) + .init_resource::() + .init_resource::() + .init_resource::>() + .add_systems(ExtractSchedule, extract_screenshots.ambiguous_with_all()) + .add_systems( + Render, + prepare_screenshots + .after(prepare_view_attachments) + .before(prepare_view_targets) + .in_set(RenderSet::ManageViews), + ); } } } @@ -221,114 +519,187 @@ impl SpecializedRenderPipeline for ScreenshotToScreenPipeline { } } -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 targets = world.resource::(); + let prepared = world.resource::(); let pipelines = world.resource::(); + let gpu_images = world.resource::>(); + let windows = world.resource::(); + let manual_texture_views = 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: layout_data(width, height, texture_format), - }, - Extent3d { + for (entity, render_target) in targets.iter() { + match render_target { + NormalizedRenderTarget::Window(window) => { + let window = window.entity(); + let Some(window) = windows.get(&window) else { + continue; + }; + let width = window.physical_width; + let height = window.physical_height; + let Some(texture_format) = window.swap_chain_texture_format else { + continue; + }; + let Some(swap_chain_texture) = window.swap_chain_texture.as_ref() else { + continue; + }; + let texture_view = swap_chain_texture.texture.create_view(&Default::default()); + render_screenshot( + encoder, + prepared, + pipelines, + entity, 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: wgpu::StoreOp::Store, - }, - })], - depth_stencil_attachment: None, - timestamp_writes: None, - occlusion_query_set: None, - }); - pass.set_pipeline(pipeline); - pass.set_bind_group(0, &memory.bind_group, &[]); - pass.draw(0..3, 0..1); + texture_format, + &texture_view, + ); } + NormalizedRenderTarget::Image(image) => { + let Some(gpu_image) = gpu_images.get(image) else { + warn!("Unknown image for screenshot, skipping: {:?}", image); + continue; + }; + let width = gpu_image.size.x; + let height = gpu_image.size.y; + let texture_format = gpu_image.texture_format; + let texture_view = gpu_image.texture_view.deref(); + render_screenshot( + encoder, + prepared, + pipelines, + entity, + width, + height, + texture_format, + texture_view, + ); + } + NormalizedRenderTarget::TextureView(texture_view) => { + let Some(texture_view) = manual_texture_views.get(texture_view) else { + warn!( + "Unknown manual texture view for screenshot, skipping: {:?}", + texture_view + ); + continue; + }; + let width = texture_view.size.x; + let height = texture_view.size.y; + let texture_format = texture_view.format; + let texture_view = texture_view.texture_view.deref(); + render_screenshot( + encoder, + prepared, + pipelines, + entity, + width, + height, + texture_format, + texture_view, + ); + } + }; + } +} + +#[allow(clippy::too_many_arguments)] +fn render_screenshot( + encoder: &mut CommandEncoder, + prepared: &RenderScreenshotsPrepared, + pipelines: &PipelineCache, + entity: &Entity, + width: u32, + height: u32, + texture_format: TextureFormat, + texture_view: &wgpu::TextureView, +) { + if let Some(prepared_state) = &prepared.get(entity) { + encoder.copy_texture_to_buffer( + prepared_state.texture.as_image_copy(), + wgpu::ImageCopyBuffer { + buffer: &prepared_state.buffer, + layout: layout_data(width, height, texture_format), + }, + Extent3d { + width, + height, + ..Default::default() + }, + ); + + if let Some(pipeline) = pipelines.get_render_pipeline(prepared_state.pipeline_id) { + let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor { + label: Some("screenshot_to_screen_pass"), + color_attachments: &[Some(wgpu::RenderPassColorAttachment { + view: texture_view, + resolve_target: None, + ops: wgpu::Operations { + load: wgpu::LoadOp::Load, + store: wgpu::StoreOp::Store, + }, + })], + depth_stencil_attachment: None, + timestamp_writes: None, + occlusion_query_set: None, + }); + pass.set_pipeline(pipeline); + pass.set_bind_group(0, &prepared_state.bind_group, &[]); + pass.draw(0..3, 0..1); } } } pub(crate) fn collect_screenshots(world: &mut World) { - let _span = info_span!("collect_screenshots"); + #[cfg(feature = "trace")] + let _span = bevy_utils::tracing::info_span!("collect_screenshots").entered(); - 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 sender = world.resource::().deref().clone(); + let prepared = world.resource::(); - 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); + for (entity, prepared) in prepared.iter() { + let entity = *entity; + let sender = sender.clone(); + let width = prepared.size.width; + let height = prepared.size.height; + let texture_format = prepared.texture.format(); + let pixel_size = texture_format.pixel_size(); + let buffer = prepared.buffer.clone(); - 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); + 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); - screenshot_func(Image::new( + 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); + } + + if let Err(e) = sender.send(( + entity, + Image::new( Extent3d { width, height, @@ -338,10 +709,12 @@ pub(crate) fn collect_screenshots(world: &mut World) { result, texture_format, RenderAssetUsages::RENDER_WORLD, - )); - }; + ), + )) { + error!("Failed to send screenshot: {:?}", e); + } + }; - AsyncComputeTaskPool::get().spawn(finish).detach(); - } + AsyncComputeTaskPool::get().spawn(finish).detach(); } } diff --git a/examples/window/screenshot.rs b/examples/window/screenshot.rs index 66591ac053..1bf800d18c 100644 --- a/examples/window/screenshot.rs +++ b/examples/window/screenshot.rs @@ -1,29 +1,48 @@ //! An example showing how to save screenshots to disk use bevy::prelude::*; -use bevy::render::view::screenshot::ScreenshotManager; -use bevy::window::PrimaryWindow; +use bevy::window::SystemCursorIcon; +use bevy_render::view::cursor::CursorIcon; +use bevy_render::view::screenshot::{save_to_disk, Capturing, Screenshot}; fn main() { App::new() .add_plugins(DefaultPlugins) .add_systems(Startup, setup) - .add_systems(Update, screenshot_on_spacebar) + .add_systems(Update, (screenshot_on_spacebar, screenshot_saving)) .run(); } fn screenshot_on_spacebar( + mut commands: Commands, input: Res>, - main_window: Query>, - mut screenshot_manager: ResMut, mut counter: Local, ) { if input.just_pressed(KeyCode::Space) { let path = format!("./screenshot-{}.png", *counter); *counter += 1; - screenshot_manager - .save_screenshot_to_disk(main_window.single(), path) - .unwrap(); + commands + .spawn(Screenshot::primary_window()) + .observe(save_to_disk(path)); + } +} + +fn screenshot_saving( + mut commands: Commands, + screenshot_saving: Query>, + windows: Query>, +) { + let window = windows.single(); + match screenshot_saving.iter().count() { + 0 => { + commands.entity(window).remove::(); + } + x if x > 0 => { + commands + .entity(window) + .insert(CursorIcon::from(SystemCursorIcon::Progress)); + } + _ => {} } }