Add screenshot api (#7163)

Fixes https://github.com/bevyengine/bevy/issues/1207

# Objective

Right now, it's impossible to capture a screenshot of the entire window
without forking bevy. This is because
- The swapchain texture never has the COPY_SRC usage
- It can't be accessed without taking ownership of it
- Taking ownership of it breaks *a lot* of stuff

## Solution

- Introduce a dedicated api for taking a screenshot of a given bevy
window, and guarantee this screenshot will always match up with what
gets put on the screen.

---

## Changelog

- Added the `ScreenshotManager` resource with two functions,
`take_screenshot` and `save_screenshot_to_disk`
This commit is contained in:
Nile 2023-04-20 00:28:42 +03:00 committed by GitHub
parent 9fd867aeba
commit 9db70da96f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
12 changed files with 557 additions and 53 deletions

View file

@ -1909,6 +1909,16 @@ description = "Illustrates how to customize the default window settings"
category = "Window" category = "Window"
wasm = true 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]] [[example]]
name = "transparent_window" name = "transparent_window"
path = "examples/window/transparent_window.rs" path = "examples/window/transparent_window.rs"

View file

@ -423,7 +423,7 @@ impl NormalizedRenderTarget {
match self { match self {
NormalizedRenderTarget::Window(window_ref) => windows NormalizedRenderTarget::Window(window_ref) => windows
.get(&window_ref.entity()) .get(&window_ref.entity())
.and_then(|window| window.swap_chain_texture.as_ref()), .and_then(|window| window.swap_chain_texture_view.as_ref()),
NormalizedRenderTarget::Image(image_handle) => { NormalizedRenderTarget::Image(image_handle) => {
images.get(image_handle).map(|image| &image.texture_view) images.get(image_handle).map(|image| &image.texture_view)
} }

View file

@ -52,7 +52,7 @@ impl Node for CameraDriverNode {
continue; continue;
} }
let Some(swap_chain_texture) = &window.swap_chain_texture else { let Some(swap_chain_texture) = &window.swap_chain_texture_view else {
continue; continue;
}; };

View file

@ -51,31 +51,21 @@ define_atomic_id!(TextureViewId);
render_resource_wrapper!(ErasedTextureView, wgpu::TextureView); render_resource_wrapper!(ErasedTextureView, wgpu::TextureView);
render_resource_wrapper!(ErasedSurfaceTexture, wgpu::SurfaceTexture); 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). /// 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)] #[derive(Clone, Debug)]
pub struct TextureView { pub struct TextureView {
id: TextureViewId, id: TextureViewId,
value: TextureViewValue, value: ErasedTextureView,
}
pub struct SurfaceTexture {
value: ErasedSurfaceTexture,
}
impl SurfaceTexture {
pub fn try_unwrap(self) -> Option<wgpu::SurfaceTexture> {
self.value.try_unwrap()
}
} }
impl TextureView { impl TextureView {
@ -84,34 +74,21 @@ impl TextureView {
pub fn id(&self) -> TextureViewId { pub fn id(&self) -> TextureViewId {
self.id 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<wgpu::SurfaceTexture> {
match self.value {
TextureViewValue::TextureView(_) => None,
TextureViewValue::SurfaceTexture { texture, .. } => texture.try_unwrap(),
}
}
} }
impl From<wgpu::TextureView> for TextureView { impl From<wgpu::TextureView> for TextureView {
fn from(value: wgpu::TextureView) -> Self { fn from(value: wgpu::TextureView) -> Self {
TextureView { TextureView {
id: TextureViewId::new(), id: TextureViewId::new(),
value: TextureViewValue::TextureView(ErasedTextureView::new(value)), value: ErasedTextureView::new(value),
} }
} }
} }
impl From<wgpu::SurfaceTexture> for TextureView { impl From<wgpu::SurfaceTexture> for SurfaceTexture {
fn from(value: wgpu::SurfaceTexture) -> Self { fn from(value: wgpu::SurfaceTexture) -> Self {
let view = ErasedTextureView::new(value.texture.create_view(&Default::default())); SurfaceTexture {
let texture = ErasedSurfaceTexture::new(value); value: ErasedSurfaceTexture::new(value),
TextureView {
id: TextureViewId::new(),
value: TextureViewValue::SurfaceTexture { texture, view },
} }
} }
} }
@ -121,11 +98,17 @@ impl Deref for TextureView {
#[inline] #[inline]
fn deref(&self) -> &Self::Target { fn deref(&self) -> &Self::Target {
match &self.value { &self.value
TextureViewValue::TextureView(value) => value,
TextureViewValue::SurfaceTexture { view, .. } => view,
} }
} }
impl Deref for SurfaceTexture {
type Target = wgpu::SurfaceTexture;
#[inline]
fn deref(&self) -> &Self::Target {
&self.value
}
} }
define_atomic_id!(SamplerId); define_atomic_id!(SamplerId);

View file

@ -57,9 +57,12 @@ impl RenderGraphRunner {
render_device: RenderDevice, render_device: RenderDevice,
queue: &wgpu::Queue, queue: &wgpu::Queue,
world: &World, world: &World,
finalizer: impl FnOnce(&mut wgpu::CommandEncoder),
) -> Result<(), RenderGraphRunnerError> { ) -> Result<(), RenderGraphRunnerError> {
let mut render_context = RenderContext::new(render_device); let mut render_context = RenderContext::new(render_device);
Self::run_graph(graph, None, &mut render_context, world, &[], None)?; Self::run_graph(graph, None, &mut render_context, world, &[], None)?;
finalizer(render_context.command_encoder());
{ {
#[cfg(feature = "trace")] #[cfg(feature = "trace")]
let _span = info_span!("submit_graph_commands").entered(); let _span = info_span!("submit_graph_commands").entered();

View file

@ -35,6 +35,9 @@ pub fn render_system(world: &mut World) {
render_device.clone(), // TODO: is this clone really necessary? render_device.clone(), // TODO: is this clone really necessary?
&render_queue.0, &render_queue.0,
world, world,
|encoder| {
crate::view::screenshot::submit_screenshot_commands(world, encoder);
},
) { ) {
error!("Error running render graph:"); error!("Error running render graph:");
{ {
@ -66,8 +69,8 @@ pub fn render_system(world: &mut World) {
let mut windows = world.resource_mut::<ExtractedWindows>(); let mut windows = world.resource_mut::<ExtractedWindows>();
for window in windows.values_mut() { for window in windows.values_mut() {
if let Some(texture_view) = window.swap_chain_texture.take() { if let Some(wrapped_texture) = window.swap_chain_texture.take() {
if let Some(surface_texture) = texture_view.take_surface_texture() { if let Some(surface_texture) = wrapped_texture.try_unwrap() {
surface_texture.present(); surface_texture.present();
} }
} }
@ -81,6 +84,8 @@ pub fn render_system(world: &mut World) {
); );
} }
crate::view::screenshot::collect_screenshots(world);
// update the time and send it to the app world // update the time and send it to the app world
let time_sender = world.resource::<TimeSender>(); let time_sender = world.resource::<TimeSender>();
time_sender.0.try_send(Instant::now()).expect( time_sender.0.try_send(Instant::now()).expect(

View file

@ -174,6 +174,7 @@ impl Image {
/// - `TextureFormat::R8Unorm` /// - `TextureFormat::R8Unorm`
/// - `TextureFormat::Rg8Unorm` /// - `TextureFormat::Rg8Unorm`
/// - `TextureFormat::Rgba8UnormSrgb` /// - `TextureFormat::Rgba8UnormSrgb`
/// - `TextureFormat::Bgra8UnormSrgb`
/// ///
/// To convert [`Image`] to a different format see: [`Image::convert`]. /// To convert [`Image`] to a different format see: [`Image::convert`].
pub fn try_into_dynamic(self) -> anyhow::Result<DynamicImage> { pub fn try_into_dynamic(self) -> anyhow::Result<DynamicImage> {
@ -196,6 +197,20 @@ impl Image {
self.data, self.data,
) )
.map(DynamicImage::ImageRgba8), .map(DynamicImage::ImageRgba8),
// This format is commonly used as the format for the swapchain texture
// This conversion is added here to support screenshots
TextureFormat::Bgra8UnormSrgb => ImageBuffer::from_raw(
self.texture_descriptor.size.width,
self.texture_descriptor.size.height,
{
let mut data = self.data;
for bgra in data.chunks_exact_mut(4) {
bgra.swap(0, 2);
}
data
},
)
.map(DynamicImage::ImageRgba8),
// Throw and error if conversion isn't supported // Throw and error if conversion isn't supported
texture_format => { texture_format => {
return Err(anyhow!( return Err(anyhow!(

View file

@ -1,6 +1,7 @@
use crate::{ use crate::{
render_resource::TextureView, render_resource::{PipelineCache, SpecializedRenderPipelines, SurfaceTexture, TextureView},
renderer::{RenderAdapter, RenderDevice, RenderInstance}, renderer::{RenderAdapter, RenderDevice, RenderInstance},
texture::TextureFormatPixelInfo,
Extract, ExtractSchedule, Render, RenderApp, RenderSet, Extract, ExtractSchedule, Render, RenderApp, RenderSet,
}; };
use bevy_app::{App, Plugin}; use bevy_app::{App, Plugin};
@ -10,7 +11,13 @@ use bevy_window::{
CompositeAlphaMode, PresentMode, PrimaryWindow, RawHandleWrapper, Window, WindowClosed, CompositeAlphaMode, PresentMode, PrimaryWindow, RawHandleWrapper, Window, WindowClosed,
}; };
use std::ops::{Deref, DerefMut}; use std::ops::{Deref, DerefMut};
use wgpu::TextureFormat; use wgpu::{BufferUsages, TextureFormat, TextureUsages};
pub mod screenshot;
use screenshot::{
ScreenshotManager, ScreenshotPlugin, ScreenshotPreparedState, ScreenshotToScreenPipeline,
};
use super::Msaa; use super::Msaa;
@ -27,10 +34,13 @@ pub enum WindowSystem {
impl Plugin for WindowRenderPlugin { impl Plugin for WindowRenderPlugin {
fn build(&self, app: &mut App) { fn build(&self, app: &mut App) {
app.add_plugin(ScreenshotPlugin);
if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
render_app render_app
.init_resource::<ExtractedWindows>() .init_resource::<ExtractedWindows>()
.init_resource::<WindowSurfaces>() .init_resource::<WindowSurfaces>()
.init_resource::<ScreenshotToScreenPipeline>()
.init_non_send_resource::<NonSendMarker>() .init_non_send_resource::<NonSendMarker>()
.add_systems(ExtractSchedule, extract_windows) .add_systems(ExtractSchedule, extract_windows)
.configure_set(Render, WindowSystem::Prepare.in_set(RenderSet::Prepare)) .configure_set(Render, WindowSystem::Prepare.in_set(RenderSet::Prepare))
@ -46,11 +56,26 @@ pub struct ExtractedWindow {
pub physical_width: u32, pub physical_width: u32,
pub physical_height: u32, pub physical_height: u32,
pub present_mode: PresentMode, pub present_mode: PresentMode,
pub swap_chain_texture: Option<TextureView>, /// 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<TextureView>,
pub swap_chain_texture: Option<SurfaceTexture>,
pub swap_chain_texture_format: Option<TextureFormat>, pub swap_chain_texture_format: Option<TextureFormat>,
pub screenshot_memory: Option<ScreenshotPreparedState>,
pub size_changed: bool, pub size_changed: bool,
pub present_mode_changed: bool, pub present_mode_changed: bool,
pub alpha_mode: CompositeAlphaMode, pub alpha_mode: CompositeAlphaMode,
pub screenshot_func: Option<screenshot::ScreenshotFn>,
}
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)] #[derive(Default, Resource)]
@ -75,6 +100,7 @@ impl DerefMut for ExtractedWindows {
fn extract_windows( fn extract_windows(
mut extracted_windows: ResMut<ExtractedWindows>, mut extracted_windows: ResMut<ExtractedWindows>,
screenshot_manager: Extract<Res<ScreenshotManager>>,
mut closed: Extract<EventReader<WindowClosed>>, mut closed: Extract<EventReader<WindowClosed>>,
windows: Extract<Query<(Entity, &Window, &RawHandleWrapper, Option<&PrimaryWindow>)>>, windows: Extract<Query<(Entity, &Window, &RawHandleWrapper, Option<&PrimaryWindow>)>>,
) { ) {
@ -95,14 +121,17 @@ fn extract_windows(
physical_height: new_height, physical_height: new_height,
present_mode: window.present_mode, present_mode: window.present_mode,
swap_chain_texture: None, swap_chain_texture: None,
swap_chain_texture_view: None,
size_changed: false, size_changed: false,
swap_chain_texture_format: None, swap_chain_texture_format: None,
present_mode_changed: false, present_mode_changed: false,
alpha_mode: window.composite_alpha_mode, alpha_mode: window.composite_alpha_mode,
screenshot_func: None,
screenshot_memory: None,
}); });
// NOTE: Drop the swap chain frame here // 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 extracted_window.size_changed = new_width != extracted_window.physical_width
|| new_height != extracted_window.physical_height; || new_height != extracted_window.physical_height;
extracted_window.present_mode_changed = extracted_window.present_mode_changed =
@ -132,6 +161,15 @@ fn extract_windows(
for closed_window in closed.iter() { for closed_window in closed.iter() {
extracted_windows.remove(&closed_window.window); extracted_windows.remove(&closed_window.window);
} }
// This lock will never block because `callbacks` is `pub(crate)` and this is the singular callsite where it's locked.
// Even if a user had multiple copies of this system, since the system has a mutable resource access the two systems would never run
// at the same time
// TODO: since this is guaranteed, should the lock be replaced with an UnsafeCell to remove the overhead, or is it minor enough to be ignored?
for (window, screenshot_func) in screenshot_manager.callbacks.lock().drain() {
if let Some(window) = extracted_windows.get_mut(&window) {
window.screenshot_func = Some(screenshot_func);
}
}
} }
struct SurfaceData { struct SurfaceData {
@ -167,6 +205,7 @@ pub struct WindowSurfaces {
/// another alternative is to try to use [`ANGLE`](https://github.com/gfx-rs/wgpu#angle) and /// 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 /// [`Backends::GL`](crate::settings::Backends::GL) if your GPU/drivers support `OpenGL 4.3` / `OpenGL ES 3.0` or
/// later. /// later.
#[allow(clippy::too_many_arguments)]
pub fn prepare_windows( pub fn prepare_windows(
// By accessing a NonSend resource, we tell the scheduler to put this system on the main thread, // By accessing a NonSend resource, we tell the scheduler to put this system on the main thread,
// which is necessary for some OS s // which is necessary for some OS s
@ -176,6 +215,9 @@ pub fn prepare_windows(
render_device: Res<RenderDevice>, render_device: Res<RenderDevice>,
render_instance: Res<RenderInstance>, render_instance: Res<RenderInstance>,
render_adapter: Res<RenderAdapter>, render_adapter: Res<RenderAdapter>,
screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
pipeline_cache: Res<PipelineCache>,
mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
mut msaa: ResMut<Msaa>, mut msaa: ResMut<Msaa>,
) { ) {
for window in windows.windows.values_mut() { for window in windows.windows.values_mut() {
@ -285,18 +327,18 @@ pub fn prepare_windows(
let frame = surface let frame = surface
.get_current_texture() .get_current_texture()
.expect("Error configuring surface"); .expect("Error configuring surface");
window.swap_chain_texture = Some(TextureView::from(frame)); window.set_swapchain_texture(frame);
} else { } else {
match surface.get_current_texture() { match surface.get_current_texture() {
Ok(frame) => { Ok(frame) => {
window.swap_chain_texture = Some(TextureView::from(frame)); window.set_swapchain_texture(frame);
} }
Err(wgpu::SurfaceError::Outdated) => { Err(wgpu::SurfaceError::Outdated) => {
render_device.configure_surface(surface, &surface_configuration); render_device.configure_surface(surface, &surface_configuration);
let frame = surface let frame = surface
.get_current_texture() .get_current_texture()
.expect("Error reconfiguring surface"); .expect("Error reconfiguring surface");
window.swap_chain_texture = Some(TextureView::from(frame)); window.set_swapchain_texture(frame);
} }
#[cfg(target_os = "linux")] #[cfg(target_os = "linux")]
Err(wgpu::SurfaceError::Timeout) if may_erroneously_timeout() => { Err(wgpu::SurfaceError::Timeout) if may_erroneously_timeout() => {
@ -311,5 +353,55 @@ pub fn prepare_windows(
} }
}; };
window.swap_chain_texture_format = Some(surface_data.format); window.swap_chain_texture_format = Some(surface_data.format);
if window.screenshot_func.is_some() {
let texture = render_device.create_texture(&wgpu::TextureDescriptor {
label: Some("screenshot-capture-rendertarget"),
size: wgpu::Extent3d {
width: surface_configuration.width,
height: surface_configuration.height,
depth_or_array_layers: 1,
},
mip_level_count: 1,
sample_count: 1,
dimension: wgpu::TextureDimension::D2,
format: surface_configuration.format,
usage: TextureUsages::RENDER_ATTACHMENT
| TextureUsages::COPY_SRC
| TextureUsages::TEXTURE_BINDING,
view_formats: &[],
});
let texture_view = texture.create_view(&Default::default());
let buffer = render_device.create_buffer(&wgpu::BufferDescriptor {
label: Some("screenshot-transfer-buffer"),
size: screenshot::get_aligned_size(
window.physical_width,
window.physical_height,
surface_data.format.pixel_size() as u32,
) as u64,
usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let bind_group = render_device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("screenshot-to-screen-bind-group"),
layout: &screenshot_pipeline.bind_group_layout,
entries: &[wgpu::BindGroupEntry {
binding: 0,
resource: wgpu::BindingResource::TextureView(&texture_view),
}],
});
let pipeline_id = pipelines.specialize(
&pipeline_cache,
&screenshot_pipeline,
surface_configuration.format,
);
window.swap_chain_texture_view = Some(texture_view);
window.screenshot_memory = Some(ScreenshotPreparedState {
texture,
buffer,
bind_group,
pipeline_id,
});
}
} }
} }

View file

@ -0,0 +1,315 @@
use std::{borrow::Cow, num::NonZeroU32, path::Path};
use bevy_app::Plugin;
use bevy_asset::{load_internal_asset, HandleUntyped};
use bevy_ecs::prelude::*;
use bevy_log::{error, info, info_span};
use bevy_reflect::TypeUuid;
use bevy_tasks::AsyncComputeTaskPool;
use bevy_utils::HashMap;
use parking_lot::Mutex;
use thiserror::Error;
use wgpu::{
CommandEncoder, Extent3d, ImageDataLayout, TextureFormat, COPY_BYTES_PER_ROW_ALIGNMENT,
};
use crate::{
prelude::{Image, Shader},
render_resource::{
BindGroup, BindGroupLayout, Buffer, CachedRenderPipelineId, FragmentState, PipelineCache,
RenderPipelineDescriptor, SpecializedRenderPipeline, SpecializedRenderPipelines, Texture,
VertexState,
},
renderer::RenderDevice,
texture::TextureFormatPixelInfo,
RenderApp,
};
use super::ExtractedWindows;
pub type ScreenshotFn = Box<dyn FnOnce(Image) + Send + Sync>;
/// 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<HashMap<Entity, ScreenshotFn>>,
}
#[derive(Error, Debug)]
#[error("A screenshot for this window has already been requested.")]
pub struct ScreenshotAlreadyRequestedError;
impl ScreenshotManager {
/// Signals the renderer to take a screenshot of this frame.
///
/// The given callback will eventually be called on one of the [`AsyncComputeTaskPool`]s threads.
pub fn take_screenshot(
&mut self,
window: Entity,
callback: impl FnOnce(Image) + Send + Sync + 'static,
) -> Result<(), ScreenshotAlreadyRequestedError> {
self.callbacks
.get_mut()
.try_insert(window, Box::new(callback))
.map(|_| ())
.map_err(|_| ScreenshotAlreadyRequestedError)
}
/// Signals the renderer to take a screenshot of this frame.
///
/// The screenshot will eventually be saved to the given path, and the format will be derived from the extension.
pub fn save_screenshot_to_disk(
&mut self,
window: Entity,
path: impl AsRef<Path>,
) -> Result<(), ScreenshotAlreadyRequestedError> {
let path = path.as_ref().to_owned();
self.take_screenshot(window, move |img| match img.try_into_dynamic() {
Ok(dyn_img) => match image::ImageFormat::from_path(&path) {
Ok(format) => {
// discard the alpha channel which stores brightness values when HDR is enabled to make sure
// the screenshot looks right
let img = dyn_img.to_rgb8();
match img.save_with_format(&path, format) {
Ok(_) => info!("Screenshot saved to {}", path.display()),
Err(e) => error!("Cannot save screenshot, IO error: {e}"),
}
}
Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"),
},
Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"),
})
}
}
pub struct ScreenshotPlugin;
const SCREENSHOT_SHADER_HANDLE: HandleUntyped =
HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 11918575842344596158);
impl Plugin for ScreenshotPlugin {
fn build(&self, app: &mut bevy_app::App) {
app.init_resource::<ScreenshotManager>();
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::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>();
}
}
}
pub(crate) fn align_byte_size(value: u32) -> u32 {
value + (COPY_BYTES_PER_ROW_ALIGNMENT - (value % COPY_BYTES_PER_ROW_ALIGNMENT))
}
pub(crate) fn get_aligned_size(width: u32, height: u32, pixel_size: u32) -> u32 {
height * align_byte_size(width * pixel_size)
}
pub(crate) fn layout_data(width: u32, height: u32, format: TextureFormat) -> ImageDataLayout {
ImageDataLayout {
bytes_per_row: if height > 1 {
// 1 = 1 row
NonZeroU32::new(get_aligned_size(width, 1, format.pixel_size() as u32))
} else {
None
},
rows_per_image: None,
..Default::default()
}
}
#[derive(Resource)]
pub struct ScreenshotToScreenPipeline {
pub bind_group_layout: BindGroupLayout,
}
impl FromWorld for ScreenshotToScreenPipeline {
fn from_world(render_world: &mut World) -> Self {
let device = render_world.resource::<RenderDevice>();
let bind_group_layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("screenshot-to-screen-bgl"),
entries: &[wgpu::BindGroupLayoutEntry {
binding: 0,
visibility: wgpu::ShaderStages::FRAGMENT,
ty: wgpu::BindingType::Texture {
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
}],
});
Self { bind_group_layout }
}
}
impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
type Key = TextureFormat;
fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
RenderPipelineDescriptor {
label: Some(Cow::Borrowed("screenshot-to-screen")),
layout: vec![self.bind_group_layout.clone()],
vertex: VertexState {
buffers: vec![],
shader_defs: vec![],
entry_point: Cow::Borrowed("vs_main"),
shader: SCREENSHOT_SHADER_HANDLE.typed(),
},
primitive: wgpu::PrimitiveState {
topology: wgpu::PrimitiveTopology::TriangleList,
strip_index_format: None,
front_face: wgpu::FrontFace::Ccw,
cull_mode: Some(wgpu::Face::Back),
polygon_mode: wgpu::PolygonMode::Fill,
conservative: false,
unclipped_depth: false,
},
depth_stencil: None,
multisample: wgpu::MultisampleState::default(),
fragment: Some(FragmentState {
shader: SCREENSHOT_SHADER_HANDLE.typed(),
entry_point: Cow::Borrowed("fs_main"),
shader_defs: vec![],
targets: vec![Some(wgpu::ColorTargetState {
format: key,
blend: None,
write_mask: wgpu::ColorWrites::ALL,
})],
}),
push_constant_ranges: Vec::new(),
}
}
}
pub struct ScreenshotPreparedState {
pub texture: Texture,
pub buffer: Buffer,
pub bind_group: BindGroup,
pub pipeline_id: CachedRenderPipelineId,
}
pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) {
let windows = world.resource::<ExtractedWindows>();
let pipelines = world.resource::<PipelineCache>();
for window in windows.values() {
if let Some(memory) = &window.screenshot_memory {
let width = window.physical_width;
let height = window.physical_height;
let texture_format = window.swap_chain_texture_format.unwrap();
encoder.copy_texture_to_buffer(
memory.texture.as_image_copy(),
wgpu::ImageCopyBuffer {
buffer: &memory.buffer,
layout: crate::view::screenshot::layout_data(width, height, texture_format),
},
Extent3d {
width,
height,
..Default::default()
},
);
if let Some(pipeline) = pipelines.get_render_pipeline(memory.pipeline_id) {
let true_swapchain_texture_view = window
.swap_chain_texture
.as_ref()
.unwrap()
.texture
.create_view(&Default::default());
let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
label: Some("screenshot_to_screen_pass"),
color_attachments: &[Some(wgpu::RenderPassColorAttachment {
view: &true_swapchain_texture_view,
resolve_target: None,
ops: wgpu::Operations {
load: wgpu::LoadOp::Load,
store: true,
},
})],
depth_stencil_attachment: None,
});
pass.set_pipeline(pipeline);
pass.set_bind_group(0, &memory.bind_group, &[]);
pass.draw(0..3, 0..1);
}
}
}
}
pub(crate) fn collect_screenshots(world: &mut World) {
let _span = info_span!("collect_screenshots");
let mut windows = world.resource_mut::<ExtractedWindows>();
for window in windows.values_mut() {
if let Some(screenshot_func) = window.screenshot_func.take() {
let width = window.physical_width;
let height = window.physical_height;
let texture_format = window.swap_chain_texture_format.unwrap();
let pixel_size = texture_format.pixel_size();
let ScreenshotPreparedState { buffer, .. } = window.screenshot_memory.take().unwrap();
let finish = async move {
let (tx, rx) = async_channel::bounded(1);
let buffer_slice = buffer.slice(..);
// The polling for this map call is done every frame when the command queue is submitted.
buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
let err = result.err();
if err.is_some() {
panic!("{}", err.unwrap().to_string());
}
tx.try_send(()).unwrap();
});
rx.recv().await.unwrap();
let data = buffer_slice.get_mapped_range();
// we immediately move the data to CPU memory to avoid holding the mapped view for long
let mut result = Vec::from(&*data);
drop(data);
drop(buffer);
if result.len() != ((width * height) as usize * pixel_size) {
// Our buffer has been padded because we needed to align to a multiple of 256.
// We remove this padding here
let initial_row_bytes = width as usize * pixel_size;
let buffered_row_bytes = align_byte_size(width * pixel_size as u32) as usize;
let mut take_offset = buffered_row_bytes;
let mut place_offset = initial_row_bytes;
for _ in 1..height {
result.copy_within(
take_offset..take_offset + buffered_row_bytes,
place_offset,
);
take_offset += buffered_row_bytes;
place_offset += initial_row_bytes;
}
result.truncate(initial_row_bytes * height as usize);
}
screenshot_func(Image::new(
Extent3d {
width,
height,
depth_or_array_layers: 1,
},
wgpu::TextureDimension::D2,
result,
texture_format,
));
};
AsyncComputeTaskPool::get().spawn(finish).detach();
}
}
}

View file

@ -0,0 +1,16 @@
// This vertex shader will create a triangle that will cover the entire screen
// with minimal effort, avoiding the need for a vertex buffer etc.
@vertex
fn vs_main(@builtin(vertex_index) in_vertex_index: u32) -> @builtin(position) vec4<f32> {
let x = f32((in_vertex_index & 1u) << 2u);
let y = f32((in_vertex_index & 2u) << 1u);
return vec4<f32>(x - 1.0, y - 1.0, 0.0, 1.0);
}
@group(0) @binding(0) var t: texture_2d<f32>;
@fragment
fn fs_main(@builtin(position) pos: vec4<f32>) -> @location(0) vec4<f32> {
let coords = floor(pos.xy);
return textureLoad(t, vec2<i32>(coords), 0i);
}

View file

@ -353,6 +353,7 @@ Example | Description
[Low Power](../examples/window/low_power.rs) | Demonstrates settings to reduce power use for bevy applications [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 [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 [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 [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 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 [Window Settings](../examples/window/window_settings.rs) | Demonstrates customizing default window settings

View file

@ -0,0 +1,64 @@
//! An example showing how to save screenshots to disk
use bevy::prelude::*;
use bevy::render::view::screenshot::ScreenshotManager;
use bevy::window::PrimaryWindow;
fn main() {
App::new()
.add_plugins(DefaultPlugins)
.add_systems(Startup, setup)
.add_systems(Update, screenshot_on_f12)
.run();
}
fn screenshot_on_f12(
input: Res<Input<KeyCode>>,
main_window: Query<Entity, With<PrimaryWindow>>,
mut screenshot_manager: ResMut<ScreenshotManager>,
mut counter: Local<u32>,
) {
if input.just_pressed(KeyCode::F12) {
let path = format!("./screenshot-{}.png", *counter);
*counter += 1;
screenshot_manager
.save_screenshot_to_disk(main_window.single(), path)
.unwrap();
}
}
/// set up a simple 3D scene
fn setup(
mut commands: Commands,
mut meshes: ResMut<Assets<Mesh>>,
mut materials: ResMut<Assets<StandardMaterial>>,
) {
// plane
commands.spawn(PbrBundle {
mesh: meshes.add(shape::Plane::from_size(5.0).into()),
material: materials.add(Color::rgb(0.3, 0.5, 0.3).into()),
..default()
});
// cube
commands.spawn(PbrBundle {
mesh: meshes.add(Mesh::from(shape::Cube { size: 1.0 })),
material: materials.add(Color::rgb(0.8, 0.7, 0.6).into()),
transform: Transform::from_xyz(0.0, 0.5, 0.0),
..default()
});
// light
commands.spawn(PointLightBundle {
point_light: PointLight {
intensity: 1500.0,
shadows_enabled: true,
..default()
},
transform: Transform::from_xyz(4.0, 8.0, 4.0),
..default()
});
// camera
commands.spawn(Camera3dBundle {
transform: Transform::from_xyz(-2.0, 2.5, 5.0).looking_at(Vec3::ZERO, Vec3::Y),
..default()
});
}