1use super::ExtractedWindows;
2use crate::{
3 gpu_readback,
4 render_asset::RenderAssets,
5 render_resource::{
6 binding_types::texture_2d, BindGroup, BindGroupEntries, BindGroupLayout,
7 BindGroupLayoutEntries, Buffer, BufferUsages, CachedRenderPipelineId, FragmentState,
8 PipelineCache, RenderPipelineDescriptor, SpecializedRenderPipeline,
9 SpecializedRenderPipelines, Texture, TextureUsages, TextureView, VertexState,
10 },
11 renderer::RenderDevice,
12 texture::{GpuImage, ManualTextureViews, OutputColorAttachment},
13 view::{prepare_view_attachments, prepare_view_targets, ViewTargetAttachments, WindowSurfaces},
14 ExtractSchedule, MainWorld, Render, RenderApp, RenderStartup, RenderSystems,
15};
16use alloc::{borrow::Cow, sync::Arc};
17use bevy_app::{First, Plugin, Update};
18use bevy_asset::{embedded_asset, load_embedded_asset, AssetServer, Handle, RenderAssetUsages};
19use bevy_camera::{ManualTextureViewHandle, NormalizedRenderTarget, RenderTarget};
20use bevy_derive::{Deref, DerefMut};
21use bevy_ecs::{
22 entity::EntityHashMap, message::message_update_system, prelude::*, system::SystemState,
23};
24use bevy_image::{Image, TextureFormatPixelInfo, ToExtents};
25use bevy_platform::collections::HashSet;
26use bevy_reflect::Reflect;
27use bevy_shader::Shader;
28use bevy_tasks::AsyncComputeTaskPool;
29use bevy_utils::default;
30use bevy_window::{PrimaryWindow, WindowRef};
31use core::ops::Deref;
32use std::{
33 path::Path,
34 sync::{
35 mpsc::{Receiver, Sender},
36 Mutex,
37 },
38};
39use tracing::{error, info, warn};
40use wgpu::{CommandEncoder, Extent3d, TextureFormat};
41
42#[derive(EntityEvent, Reflect, Deref, DerefMut, Debug)]
43#[reflect(Debug)]
44pub struct ScreenshotCaptured {
45 pub entity: Entity,
46 #[deref]
47 pub image: Image,
48}
49
50#[derive(Component, Deref, DerefMut, Reflect, Debug)]
74#[reflect(Component, Debug)]
75pub struct Screenshot(pub RenderTarget);
76
77#[derive(Component, Default)]
79pub struct Capturing;
80
81#[derive(Component, Default)]
84pub struct Captured;
85
86impl Screenshot {
87 pub fn window(window: Entity) -> Self {
89 Self(RenderTarget::Window(WindowRef::Entity(window)))
90 }
91
92 pub fn primary_window() -> Self {
94 Self(RenderTarget::Window(WindowRef::Primary))
95 }
96
97 pub fn image(image: Handle<Image>) -> Self {
99 Self(RenderTarget::Image(image.into()))
100 }
101
102 pub fn texture_view(texture_view: ManualTextureViewHandle) -> Self {
104 Self(RenderTarget::TextureView(texture_view))
105 }
106}
107
108struct ScreenshotPreparedState {
109 pub texture: Texture,
110 pub buffer: Buffer,
111 pub bind_group: BindGroup,
112 pub pipeline_id: CachedRenderPipelineId,
113 pub size: Extent3d,
114}
115
116#[derive(Resource, Deref, DerefMut)]
117pub struct CapturedScreenshots(pub Arc<Mutex<Receiver<(Entity, Image)>>>);
118
119#[derive(Resource, Deref, DerefMut, Default)]
120struct RenderScreenshotTargets(EntityHashMap<NormalizedRenderTarget>);
121
122#[derive(Resource, Deref, DerefMut, Default)]
123struct RenderScreenshotsPrepared(EntityHashMap<ScreenshotPreparedState>);
124
125#[derive(Resource, Deref, DerefMut)]
126struct RenderScreenshotsSender(Sender<(Entity, Image)>);
127
128pub fn save_to_disk(path: impl AsRef<Path>) -> impl FnMut(On<ScreenshotCaptured>) {
130 let path = path.as_ref().to_owned();
131 move |screenshot_captured| {
132 let img = screenshot_captured.image.clone();
133 match img.try_into_dynamic() {
134 Ok(dyn_img) => match image::ImageFormat::from_path(&path) {
135 Ok(format) => {
136 let img = dyn_img.to_rgb8();
139 #[cfg(not(target_arch = "wasm32"))]
140 match img.save_with_format(&path, format) {
141 Ok(_) => info!("Screenshot saved to {}", path.display()),
142 Err(e) => error!("Cannot save screenshot, IO error: {e}"),
143 }
144
145 #[cfg(target_arch = "wasm32")]
146 {
147 let save_screenshot = || {
148 use image::EncodableLayout;
149 use wasm_bindgen::{JsCast, JsValue};
150
151 let mut image_buffer = std::io::Cursor::new(Vec::new());
152 img.write_to(&mut image_buffer, format)
153 .map_err(|e| JsValue::from_str(&format!("{e}")))?;
154 let parts = js_sys::Array::of1(&unsafe {
156 js_sys::Uint8Array::view(image_buffer.into_inner().as_bytes())
157 .into()
158 });
159 let blob = web_sys::Blob::new_with_u8_array_sequence(&parts)?;
160 let url = web_sys::Url::create_object_url_with_blob(&blob)?;
161 let window = web_sys::window().unwrap();
162 let document = window.document().unwrap();
163 let link = document.create_element("a")?;
164 link.set_attribute("href", &url)?;
165 link.set_attribute(
166 "download",
167 path.file_name()
168 .and_then(|filename| filename.to_str())
169 .ok_or_else(|| JsValue::from_str("Invalid filename"))?,
170 )?;
171 let html_element = link.dyn_into::<web_sys::HtmlElement>()?;
172 html_element.click();
173 web_sys::Url::revoke_object_url(&url)?;
174 Ok::<(), JsValue>(())
175 };
176
177 match (save_screenshot)() {
178 Ok(_) => info!("Screenshot saved to {}", path.display()),
179 Err(e) => error!("Cannot save screenshot, error: {e:?}"),
180 };
181 }
182 }
183 Err(e) => error!("Cannot save screenshot, requested format not recognized: {e}"),
184 },
185 Err(e) => error!("Cannot save screenshot, screen format cannot be understood: {e}"),
186 }
187 }
188}
189
190fn clear_screenshots(mut commands: Commands, screenshots: Query<Entity, With<Captured>>) {
191 for entity in screenshots.iter() {
192 commands.entity(entity).despawn();
193 }
194}
195
196pub fn trigger_screenshots(
197 mut commands: Commands,
198 captured_screenshots: ResMut<CapturedScreenshots>,
199) {
200 let captured_screenshots = captured_screenshots.lock().unwrap();
201 while let Ok((entity, image)) = captured_screenshots.try_recv() {
202 commands.entity(entity).insert(Captured);
203 commands.trigger(ScreenshotCaptured { image, entity });
204 }
205}
206
207fn extract_screenshots(
208 mut targets: ResMut<RenderScreenshotTargets>,
209 mut main_world: ResMut<MainWorld>,
210 mut system_state: Local<
211 Option<
212 SystemState<(
213 Commands,
214 Query<Entity, With<PrimaryWindow>>,
215 Query<(Entity, &Screenshot), Without<Capturing>>,
216 )>,
217 >,
218 >,
219 mut seen_targets: Local<HashSet<NormalizedRenderTarget>>,
220) {
221 if system_state.is_none() {
222 *system_state = Some(SystemState::new(&mut main_world));
223 }
224 let system_state = system_state.as_mut().unwrap();
225 let (mut commands, primary_window, screenshots) = system_state.get_mut(&mut main_world);
226
227 targets.clear();
228 seen_targets.clear();
229
230 let primary_window = primary_window.iter().next();
231
232 for (entity, screenshot) in screenshots.iter() {
233 let render_target = screenshot.0.clone();
234 let Some(render_target) = render_target.normalize(primary_window) else {
235 warn!(
236 "Unknown render target for screenshot, skipping: {:?}",
237 render_target
238 );
239 continue;
240 };
241 if seen_targets.contains(&render_target) {
242 warn!(
243 "Duplicate render target for screenshot, skipping entity {}: {:?}",
244 entity, render_target
245 );
246 commands.entity(entity).despawn();
248 continue;
249 }
250 seen_targets.insert(render_target.clone());
251 targets.insert(entity, render_target);
252 commands.entity(entity).insert(Capturing);
253 }
254
255 system_state.apply(&mut main_world);
256}
257
258fn prepare_screenshots(
259 targets: Res<RenderScreenshotTargets>,
260 mut prepared: ResMut<RenderScreenshotsPrepared>,
261 window_surfaces: Res<WindowSurfaces>,
262 render_device: Res<RenderDevice>,
263 screenshot_pipeline: Res<ScreenshotToScreenPipeline>,
264 pipeline_cache: Res<PipelineCache>,
265 mut pipelines: ResMut<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>,
266 images: Res<RenderAssets<GpuImage>>,
267 manual_texture_views: Res<ManualTextureViews>,
268 mut view_target_attachments: ResMut<ViewTargetAttachments>,
269) {
270 prepared.clear();
271 for (entity, target) in targets.iter() {
272 match target {
273 NormalizedRenderTarget::Window(window) => {
274 let window = window.entity();
275 let Some(surface_data) = window_surfaces.surfaces.get(&window) else {
276 warn!("Unknown window for screenshot, skipping: {}", window);
277 continue;
278 };
279 let format = surface_data.configuration.format.add_srgb_suffix();
280 let size = Extent3d {
281 width: surface_data.configuration.width,
282 height: surface_data.configuration.height,
283 ..default()
284 };
285 let (texture_view, state) = prepare_screenshot_state(
286 size,
287 format,
288 &render_device,
289 &screenshot_pipeline,
290 &pipeline_cache,
291 &mut pipelines,
292 );
293 prepared.insert(*entity, state);
294 view_target_attachments.insert(
295 target.clone(),
296 OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()),
297 );
298 }
299 NormalizedRenderTarget::Image(image) => {
300 let Some(gpu_image) = images.get(&image.handle) else {
301 warn!("Unknown image for screenshot, skipping: {:?}", image);
302 continue;
303 };
304 let format = gpu_image.texture_format;
305 let (texture_view, state) = prepare_screenshot_state(
306 gpu_image.size,
307 format,
308 &render_device,
309 &screenshot_pipeline,
310 &pipeline_cache,
311 &mut pipelines,
312 );
313 prepared.insert(*entity, state);
314 view_target_attachments.insert(
315 target.clone(),
316 OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()),
317 );
318 }
319 NormalizedRenderTarget::TextureView(texture_view) => {
320 let Some(manual_texture_view) = manual_texture_views.get(texture_view) else {
321 warn!(
322 "Unknown manual texture view for screenshot, skipping: {:?}",
323 texture_view
324 );
325 continue;
326 };
327 let format = manual_texture_view.format;
328 let size = manual_texture_view.size.to_extents();
329 let (texture_view, state) = prepare_screenshot_state(
330 size,
331 format,
332 &render_device,
333 &screenshot_pipeline,
334 &pipeline_cache,
335 &mut pipelines,
336 );
337 prepared.insert(*entity, state);
338 view_target_attachments.insert(
339 target.clone(),
340 OutputColorAttachment::new(texture_view.clone(), format.add_srgb_suffix()),
341 );
342 }
343 NormalizedRenderTarget::None { .. } => {
344 }
346 }
347 }
348}
349
350fn prepare_screenshot_state(
351 size: Extent3d,
352 format: TextureFormat,
353 render_device: &RenderDevice,
354 pipeline: &ScreenshotToScreenPipeline,
355 pipeline_cache: &PipelineCache,
356 pipelines: &mut SpecializedRenderPipelines<ScreenshotToScreenPipeline>,
357) -> (TextureView, ScreenshotPreparedState) {
358 let texture = render_device.create_texture(&wgpu::TextureDescriptor {
359 label: Some("screenshot-capture-rendertarget"),
360 size,
361 mip_level_count: 1,
362 sample_count: 1,
363 dimension: wgpu::TextureDimension::D2,
364 format,
365 usage: TextureUsages::RENDER_ATTACHMENT
366 | TextureUsages::COPY_SRC
367 | TextureUsages::TEXTURE_BINDING,
368 view_formats: &[],
369 });
370 let texture_view = texture.create_view(&Default::default());
371 let buffer = render_device.create_buffer(&wgpu::BufferDescriptor {
372 label: Some("screenshot-transfer-buffer"),
373 size: gpu_readback::get_aligned_size(size, format.pixel_size().unwrap_or(0) as u32) as u64,
374 usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
375 mapped_at_creation: false,
376 });
377 let bind_group = render_device.create_bind_group(
378 "screenshot-to-screen-bind-group",
379 &pipeline.bind_group_layout,
380 &BindGroupEntries::single(&texture_view),
381 );
382 let pipeline_id = pipelines.specialize(pipeline_cache, pipeline, format);
383
384 (
385 texture_view,
386 ScreenshotPreparedState {
387 texture,
388 buffer,
389 bind_group,
390 pipeline_id,
391 size,
392 },
393 )
394}
395
396pub struct ScreenshotPlugin;
397
398impl Plugin for ScreenshotPlugin {
399 fn build(&self, app: &mut bevy_app::App) {
400 embedded_asset!(app, "screenshot.wgsl");
401
402 let (tx, rx) = std::sync::mpsc::channel();
403 app.insert_resource(CapturedScreenshots(Arc::new(Mutex::new(rx))))
404 .add_systems(
405 First,
406 clear_screenshots
407 .after(message_update_system)
408 .before(ApplyDeferred),
409 )
410 .add_systems(Update, trigger_screenshots);
411
412 let Some(render_app) = app.get_sub_app_mut(RenderApp) else {
413 return;
414 };
415
416 render_app
417 .insert_resource(RenderScreenshotsSender(tx))
418 .init_resource::<RenderScreenshotTargets>()
419 .init_resource::<RenderScreenshotsPrepared>()
420 .init_resource::<SpecializedRenderPipelines<ScreenshotToScreenPipeline>>()
421 .add_systems(RenderStartup, init_screenshot_to_screen_pipeline)
422 .add_systems(ExtractSchedule, extract_screenshots.ambiguous_with_all())
423 .add_systems(
424 Render,
425 prepare_screenshots
426 .after(prepare_view_attachments)
427 .before(prepare_view_targets)
428 .in_set(RenderSystems::ManageViews),
429 );
430 }
431}
432
433#[derive(Resource)]
434pub struct ScreenshotToScreenPipeline {
435 pub bind_group_layout: BindGroupLayout,
436 pub shader: Handle<Shader>,
437}
438
439pub fn init_screenshot_to_screen_pipeline(
440 mut commands: Commands,
441 render_device: Res<RenderDevice>,
442 asset_server: Res<AssetServer>,
443) {
444 let bind_group_layout = render_device.create_bind_group_layout(
445 "screenshot-to-screen-bgl",
446 &BindGroupLayoutEntries::single(
447 wgpu::ShaderStages::FRAGMENT,
448 texture_2d(wgpu::TextureSampleType::Float { filterable: false }),
449 ),
450 );
451
452 let shader = load_embedded_asset!(asset_server.as_ref(), "screenshot.wgsl");
453
454 commands.insert_resource(ScreenshotToScreenPipeline {
455 bind_group_layout,
456 shader,
457 });
458}
459
460impl SpecializedRenderPipeline for ScreenshotToScreenPipeline {
461 type Key = TextureFormat;
462
463 fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
464 RenderPipelineDescriptor {
465 label: Some(Cow::Borrowed("screenshot-to-screen")),
466 layout: vec![self.bind_group_layout.clone()],
467 vertex: VertexState {
468 shader: self.shader.clone(),
469 ..default()
470 },
471 primitive: wgpu::PrimitiveState {
472 cull_mode: Some(wgpu::Face::Back),
473 ..Default::default()
474 },
475 multisample: Default::default(),
476 fragment: Some(FragmentState {
477 shader: self.shader.clone(),
478 targets: vec![Some(wgpu::ColorTargetState {
479 format: key,
480 blend: None,
481 write_mask: wgpu::ColorWrites::ALL,
482 })],
483 ..default()
484 }),
485 ..default()
486 }
487 }
488}
489
490pub(crate) fn submit_screenshot_commands(world: &World, encoder: &mut CommandEncoder) {
491 let targets = world.resource::<RenderScreenshotTargets>();
492 let prepared = world.resource::<RenderScreenshotsPrepared>();
493 let pipelines = world.resource::<PipelineCache>();
494 let gpu_images = world.resource::<RenderAssets<GpuImage>>();
495 let windows = world.resource::<ExtractedWindows>();
496 let manual_texture_views = world.resource::<ManualTextureViews>();
497
498 for (entity, render_target) in targets.iter() {
499 match render_target {
500 NormalizedRenderTarget::Window(window) => {
501 let window = window.entity();
502 let Some(window) = windows.get(&window) else {
503 continue;
504 };
505 let width = window.physical_width;
506 let height = window.physical_height;
507 let Some(texture_format) = window.swap_chain_texture_format else {
508 continue;
509 };
510 let Some(swap_chain_texture) = window.swap_chain_texture.as_ref() else {
511 continue;
512 };
513 let texture_view = swap_chain_texture.texture.create_view(&Default::default());
514 render_screenshot(
515 encoder,
516 prepared,
517 pipelines,
518 entity,
519 width,
520 height,
521 texture_format,
522 &texture_view,
523 );
524 }
525 NormalizedRenderTarget::Image(image) => {
526 let Some(gpu_image) = gpu_images.get(&image.handle) else {
527 warn!("Unknown image for screenshot, skipping: {:?}", image);
528 continue;
529 };
530 let width = gpu_image.size.width;
531 let height = gpu_image.size.height;
532 let texture_format = gpu_image.texture_format;
533 let texture_view = gpu_image.texture_view.deref();
534 render_screenshot(
535 encoder,
536 prepared,
537 pipelines,
538 entity,
539 width,
540 height,
541 texture_format,
542 texture_view,
543 );
544 }
545 NormalizedRenderTarget::TextureView(texture_view) => {
546 let Some(texture_view) = manual_texture_views.get(texture_view) else {
547 warn!(
548 "Unknown manual texture view for screenshot, skipping: {:?}",
549 texture_view
550 );
551 continue;
552 };
553 let width = texture_view.size.x;
554 let height = texture_view.size.y;
555 let texture_format = texture_view.format;
556 let texture_view = texture_view.texture_view.deref();
557 render_screenshot(
558 encoder,
559 prepared,
560 pipelines,
561 entity,
562 width,
563 height,
564 texture_format,
565 texture_view,
566 );
567 }
568 NormalizedRenderTarget::None { .. } => {
569 }
571 };
572 }
573}
574
575fn render_screenshot(
576 encoder: &mut CommandEncoder,
577 prepared: &RenderScreenshotsPrepared,
578 pipelines: &PipelineCache,
579 entity: &Entity,
580 width: u32,
581 height: u32,
582 texture_format: TextureFormat,
583 texture_view: &wgpu::TextureView,
584) {
585 if let Some(prepared_state) = &prepared.get(entity) {
586 let extent = Extent3d {
587 width,
588 height,
589 depth_or_array_layers: 1,
590 };
591 encoder.copy_texture_to_buffer(
592 prepared_state.texture.as_image_copy(),
593 wgpu::TexelCopyBufferInfo {
594 buffer: &prepared_state.buffer,
595 layout: gpu_readback::layout_data(extent, texture_format),
596 },
597 extent,
598 );
599
600 if let Some(pipeline) = pipelines.get_render_pipeline(prepared_state.pipeline_id) {
601 let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
602 label: Some("screenshot_to_screen_pass"),
603 color_attachments: &[Some(wgpu::RenderPassColorAttachment {
604 view: texture_view,
605 depth_slice: None,
606 resolve_target: None,
607 ops: wgpu::Operations {
608 load: wgpu::LoadOp::Load,
609 store: wgpu::StoreOp::Store,
610 },
611 })],
612 depth_stencil_attachment: None,
613 timestamp_writes: None,
614 occlusion_query_set: None,
615 });
616 pass.set_pipeline(pipeline);
617 pass.set_bind_group(0, &prepared_state.bind_group, &[]);
618 pass.draw(0..3, 0..1);
619 }
620 }
621}
622
623pub(crate) fn collect_screenshots(world: &mut World) {
624 #[cfg(feature = "trace")]
625 let _span = tracing::info_span!("collect_screenshots").entered();
626
627 let sender = world.resource::<RenderScreenshotsSender>().deref().clone();
628 let prepared = world.resource::<RenderScreenshotsPrepared>();
629
630 for (entity, prepared) in prepared.iter() {
631 let entity = *entity;
632 let sender = sender.clone();
633 let width = prepared.size.width;
634 let height = prepared.size.height;
635 let texture_format = prepared.texture.format();
636 let Ok(pixel_size) = texture_format.pixel_size() else {
637 continue;
638 };
639 let buffer = prepared.buffer.clone();
640
641 let finish = async move {
642 let (tx, rx) = async_channel::bounded(1);
643 let buffer_slice = buffer.slice(..);
644 buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
646 let err = result.err();
647 if err.is_some() {
648 panic!("{}", err.unwrap().to_string());
649 }
650 tx.try_send(()).unwrap();
651 });
652 rx.recv().await.unwrap();
653 let data = buffer_slice.get_mapped_range();
654 let mut result = Vec::from(&*data);
656 drop(data);
657
658 if result.len() != ((width * height) as usize * pixel_size) {
659 let initial_row_bytes = width as usize * pixel_size;
662 let buffered_row_bytes =
663 gpu_readback::align_byte_size(width * pixel_size as u32) as usize;
664
665 let mut take_offset = buffered_row_bytes;
666 let mut place_offset = initial_row_bytes;
667 for _ in 1..height {
668 result.copy_within(take_offset..take_offset + buffered_row_bytes, place_offset);
669 take_offset += buffered_row_bytes;
670 place_offset += initial_row_bytes;
671 }
672 result.truncate(initial_row_bytes * height as usize);
673 }
674
675 if let Err(e) = sender.send((
676 entity,
677 Image::new(
678 Extent3d {
679 width,
680 height,
681 depth_or_array_layers: 1,
682 },
683 wgpu::TextureDimension::D2,
684 result,
685 texture_format,
686 RenderAssetUsages::RENDER_WORLD,
687 ),
688 )) {
689 error!("Failed to send screenshot: {}", e);
690 }
691 };
692
693 AsyncComputeTaskPool::get().spawn(finish).detach();
694 }
695}