docs/bevy/examples/app/headless_renderer.rs 20.1 K · 528 lines · raw · history

1 //! This example illustrates how to make a headless renderer.
2 //! Derived from: <https://sotrh.github.io/learn-wgpu/showcase/windowless/#a-triangle-without-a-window>
3 //! It follows these steps:
4 //!
5 //! 1. Render from camera to gpu-image render target
6 //! 2. Copy from gpu image to buffer using `ImageCopyDriver` node in `RenderGraph`
7 //! 3. Copy from buffer to channel using `receive_image_from_buffer` after `RenderSystems::Render`
8 //! 4. Save from channel to random named file using `scene::update` at `PostUpdate` in `MainWorld`
9 //! 5. Exit if `single_image` setting is set
10 //!
11 //! If your goal is to capture a single “screenshot” as opposed to every single rendered frame
12 //! without gaps, it is simpler to use [`bevy::render::view::window::screenshot::Screenshot`]
13 //! than this approach.
14
15 use bevy::{
16 app::{AppExit, ScheduleRunnerPlugin},
17 camera::RenderTarget,
18 core_pipeline::tonemapping::Tonemapping,
19 image::TextureFormatPixelInfo,
20 prelude::*,
21 render::{
22 render_asset::RenderAssets,
23 render_resource::{
24 Buffer, BufferDescriptor, BufferUsages, CommandEncoderDescriptor, Extent3d, MapMode,
25 PollType, TexelCopyBufferInfo, TexelCopyBufferLayout, TextureFormat, TextureUsages,
26 },
27 renderer::{RenderContext, RenderDevice, RenderGraph, RenderGraphSystems, RenderQueue},
28 Extract, Render, RenderApp, RenderSystems,
29 },
30 window::ExitCondition,
31 winit::WinitPlugin,
32 };
33 use crossbeam_channel::{Receiver, Sender};
34 use std::{
35 ops::{Deref, DerefMut},
36 path::PathBuf,
37 sync::{
38 atomic::{AtomicBool, Ordering},
39 Arc,
40 },
41 time::Duration,
42 };
43 // To communicate between the main world and the render world we need a channel.
44 // Since the main world and render world run in parallel, there will always be a frame of latency
45 // between the data sent from the render world and the data received in the main world
46 //
47 // frame n => render world sends data through the channel at the end of the frame
48 // frame n + 1 => main world receives the data
49 //
50 // Receiver and Sender are kept in resources because there is single camera and single target
51 // That's why there is single images role, if you want to differentiate images
52 // from different cameras, you should keep Receiver in ImageCopier and Sender in ImageToSave
53 // or send some id with data
54
55 /// This will receive asynchronously any data sent from the render world
56 #[derive(Resource, Deref)]
57 struct MainWorldReceiver(Receiver<Vec<u8>>);
58
59 /// This will send asynchronously any data to the main world
60 #[derive(Resource, Deref)]
61 struct RenderWorldSender(Sender<Vec<u8>>);
62
63 // Parameters of resulting image
64 struct AppConfig {
65 width: u32,
66 height: u32,
67 single_image: bool,
68 }
69
70 fn main() {
71 let config = AppConfig {
72 width: 1920,
73 height: 1080,
74 single_image: true,
75 };
76
77 // setup frame capture
78 App::new()
79 .insert_resource(SceneController::new(
80 config.width,
81 config.height,
82 config.single_image,
83 ))
84 .insert_resource(ClearColor(Color::srgb_u8(0, 0, 0)))
85 .add_plugins(
86 DefaultPlugins
87 .set(ImagePlugin::default_nearest())
88 // Not strictly necessary, as the inclusion of ScheduleRunnerPlugin below
89 // replaces the bevy_winit app runner and so a window is never created.
90 .set(WindowPlugin {
91 primary_window: None,
92 // Don’t automatically exit due to having no windows.
93 // Instead, the code in `update()` will explicitly produce an `AppExit` event.
94 exit_condition: ExitCondition::DontExit,
95 ..default()
96 })
97 // WinitPlugin will panic in environments without a display server.
98 .disable::<WinitPlugin>(),
99 )
100 .add_plugins(ImageCopyPlugin)
101 // headless frame capture
102 .add_plugins(CaptureFramePlugin)
103 // ScheduleRunnerPlugin provides an alternative to the default bevy_winit app runner, which
104 // manages the loop without creating a window.
105 .add_plugins(ScheduleRunnerPlugin::run_loop(
106 // Run 60 times per second.
107 Duration::from_secs_f64(1.0 / 60.0),
108 ))
109 .init_resource::<SceneController>()
110 .add_systems(Startup, setup)
111 .run();
112 }
113
114 /// Capture image settings and state
115 #[derive(Debug, Default, Resource)]
116 struct SceneController {
117 state: SceneState,
118 name: String,
119 width: u32,
120 height: u32,
121 single_image: bool,
122 }
123
124 impl SceneController {
125 pub fn new(width: u32, height: u32, single_image: bool) -> SceneController {
126 SceneController {
127 state: SceneState::BuildScene,
128 name: String::from(""),
129 width,
130 height,
131 single_image,
132 }
133 }
134 }
135
136 /// Capture image state
137 #[derive(Debug, Default)]
138 enum SceneState {
139 #[default]
140 // State before any rendering
141 BuildScene,
142 // Rendering state, stores the number of frames remaining before saving the image
143 Render(u32),
144 }
145
146 fn setup(
147 mut commands: Commands,
148 mut meshes: ResMut<Assets<Mesh>>,
149 mut materials: ResMut<Assets<StandardMaterial>>,
150 mut images: ResMut<Assets<Image>>,
151 mut scene_controller: ResMut<SceneController>,
152 render_device: Res<RenderDevice>,
153 ) {
154 let render_target = setup_render_target(
155 &mut commands,
156 &mut images,
157 &render_device,
158 &mut scene_controller,
159 // pre_roll_frames should be big enough for full scene render,
160 // but the bigger it is, the longer example will run.
161 // To visualize stages of scene rendering change this param to 0
162 // and change AppConfig::single_image to false in main
163 // Stages are:
164 // 1. Transparent image
165 // 2. Few black box images
166 // 3. Fully rendered scene images
167 // Exact number depends on device speed, device load and scene size
168 40,
169 "main_scene".into(),
170 );
171
172 // Scene example for non black box picture
173 // circular base
174 commands.spawn((
175 Mesh3d(meshes.add(Circle::new(4.0))),
176 MeshMaterial3d(materials.add(Color::WHITE)),
177 Transform::from_rotation(Quat::from_rotation_x(-std::f32::consts::FRAC_PI_2)),
178 ));
179 // cube
180 commands.spawn((
181 Mesh3d(meshes.add(Cuboid::new(1.0, 1.0, 1.0))),
182 MeshMaterial3d(materials.add(Color::srgb_u8(124, 144, 255))),
183 Transform::from_xyz(0.0, 0.5, 0.0),
184 ));
185 // light
186 commands.spawn((
187 PointLight {
188 shadow_maps_enabled: true,
189 ..default()
190 },
191 Transform::from_xyz(4.0, 8.0, 4.0),
192 ));
193
194 commands.spawn((
195 Camera3d::default(),
196 render_target,
197 Tonemapping::Linear,
198 Transform::from_xyz(-2.5, 4.5, 9.0).looking_at(Vec3::ZERO, Vec3::Y),
199 ));
200 }
201
202 /// Plugin for Render world part of work
203 pub struct ImageCopyPlugin;
204 impl Plugin for ImageCopyPlugin {
205 fn build(&self, app: &mut App) {
206 let (s, r) = crossbeam_channel::unbounded();
207
208 let render_app = app
209 .insert_resource(MainWorldReceiver(r))
210 .sub_app_mut(RenderApp);
211
212 render_app
213 .insert_resource(RenderWorldSender(s))
214 // Make ImageCopiers accessible in RenderWorld system and plugin
215 .add_systems(ExtractSchedule, image_copy_extract)
216 // Receives image data from buffer to channel
217 // so we need to run it after the render graph is done
218 .add_systems(
219 Render,
220 receive_image_from_buffer.after(RenderSystems::Render),
221 )
222 .add_systems(
223 RenderGraph,
224 image_copy_driver.after(RenderGraphSystems::Submit),
225 );
226 }
227 }
228
229 /// Setups render target and cpu image for saving, changes scene state into render mode
230 fn setup_render_target(
231 commands: &mut Commands,
232 images: &mut ResMut<Assets<Image>>,
233 render_device: &Res<RenderDevice>,
234 scene_controller: &mut ResMut<SceneController>,
235 pre_roll_frames: u32,
236 scene_name: String,
237 ) -> RenderTarget {
238 let size = Extent3d {
239 width: scene_controller.width,
240 height: scene_controller.height,
241 ..Default::default()
242 };
243
244 // This is the texture that will be rendered to.
245 let mut render_target_image =
246 Image::new_target_texture(size.width, size.height, TextureFormat::Rgba8UnormSrgb, None);
247 render_target_image.texture_descriptor.usage |= TextureUsages::COPY_SRC;
248 let render_target_image_handle = images.add(render_target_image);
249
250 // This is the texture that will be copied to.
251 let cpu_image =
252 Image::new_target_texture(size.width, size.height, TextureFormat::Rgba8UnormSrgb, None);
253 let cpu_image_handle = images.add(cpu_image);
254
255 commands.spawn(ImageCopier::new(
256 render_target_image_handle.clone(),
257 size,
258 render_device,
259 ));
260
261 commands.spawn(ImageToSave(cpu_image_handle));
262
263 scene_controller.state = SceneState::Render(pre_roll_frames);
264 scene_controller.name = scene_name;
265 RenderTarget::Image(render_target_image_handle.into())
266 }
267
268 /// Setups image saver
269 pub struct CaptureFramePlugin;
270 impl Plugin for CaptureFramePlugin {
271 fn build(&self, app: &mut App) {
272 info!("Adding CaptureFramePlugin");
273 app.add_systems(PostUpdate, update);
274 }
275 }
276
277 /// `ImageCopier` aggregator in `RenderWorld`
278 #[derive(Clone, Default, Resource, Deref, DerefMut)]
279 struct ImageCopiers(pub Vec<ImageCopier>);
280
281 /// Used by `ImageCopyDriver` for copying from render target to buffer
282 #[derive(Clone, Component)]
283 struct ImageCopier {
284 buffer: Buffer,
285 enabled: Arc<AtomicBool>,
286 src_image: Handle<Image>,
287 }
288
289 impl ImageCopier {
290 pub fn new(
291 src_image: Handle<Image>,
292 size: Extent3d,
293 render_device: &RenderDevice,
294 ) -> ImageCopier {
295 let padded_bytes_per_row = RenderDevice::align_copy_bytes_per_row(size.width as usize * 4);
296 let cpu_buffer = render_device.create_buffer(&BufferDescriptor {
297 label: None,
298 size: padded_bytes_per_row as u64 * size.height as u64,
299 usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
300 mapped_at_creation: false,
301 });
302
303 ImageCopier {
304 buffer: cpu_buffer,
305 src_image,
306 enabled: Arc::new(AtomicBool::new(true)),
307 }
308 }
309
310 pub fn enabled(&self) -> bool {
311 self.enabled.load(Ordering::Relaxed)
312 }
313 }
314
315 /// Extracting `ImageCopier`s into render world, because `ImageCopyDriver` accesses them
316 fn image_copy_extract(mut commands: Commands, image_copiers: Extract<Query<&ImageCopier>>) {
317 commands.insert_resource(ImageCopiers(
318 image_copiers.iter().cloned().collect::<Vec<ImageCopier>>(),
319 ));
320 }
321
322 // Copies image content from render target to buffer
323 fn image_copy_driver(
324 render_context: RenderContext,
325 image_copiers: Res<ImageCopiers>,
326 render_queue: Res<RenderQueue>,
327 gpu_images: Res<RenderAssets<bevy::render::texture::GpuImage>>,
328 ) {
329 for image_copier in image_copiers.iter() {
330 if !image_copier.enabled() {
331 continue;
332 }
333
334 let src_image = gpu_images.get(&image_copier.src_image).unwrap();
335
336 let mut encoder = render_context
337 .render_device()
338 .create_command_encoder(&CommandEncoderDescriptor::default());
339
340 let block_dimensions = src_image.texture_descriptor.format.block_dimensions();
341 let block_size = src_image
342 .texture_descriptor
343 .format
344 .block_copy_size(None)
345 .unwrap();
346
347 // Calculating correct size of image row because
348 // copy_texture_to_buffer can copy image only by rows aligned wgpu::COPY_BYTES_PER_ROW_ALIGNMENT
349 // That's why image in buffer can be little bit wider
350 // This should be taken into account at copy from buffer stage
351 let padded_bytes_per_row = RenderDevice::align_copy_bytes_per_row(
352 (src_image.texture_descriptor.size.width as usize / block_dimensions.0 as usize)
353 * block_size as usize,
354 );
355
356 encoder.copy_texture_to_buffer(
357 src_image.texture.as_image_copy(),
358 TexelCopyBufferInfo {
359 buffer: &image_copier.buffer,
360 layout: TexelCopyBufferLayout {
361 offset: 0,
362 bytes_per_row: Some(
363 std::num::NonZero::<u32>::new(padded_bytes_per_row as u32)
364 .unwrap()
365 .into(),
366 ),
367 rows_per_image: None,
368 },
369 },
370 src_image.texture_descriptor.size,
371 );
372
373 render_queue.submit(std::iter::once(encoder.finish()));
374 }
375 }
376
377 /// runs in render world after Render stage to send image from buffer via channel (receiver is in main world)
378 fn receive_image_from_buffer(
379 image_copiers: Res<ImageCopiers>,
380 render_device: Res<RenderDevice>,
381 sender: Res<RenderWorldSender>,
382 ) {
383 for image_copier in image_copiers.0.iter() {
384 if !image_copier.enabled() {
385 continue;
386 }
387
388 // Finally time to get our data back from the gpu.
389 // First we get a buffer slice which represents a chunk of the buffer (which we
390 // can't access yet).
391 // We want the whole thing so use unbounded range.
392 let buffer_slice = image_copier.buffer.slice(..);
393
394 // Now things get complicated. WebGPU, for safety reasons, only allows either the GPU
395 // or CPU to access a buffer's contents at a time. We need to "map" the buffer which means
396 // flipping ownership of the buffer over to the CPU and making access legal. We do this
397 // with `BufferSlice::map_async`.
398 //
399 // The problem is that map_async is not an async function so we can't await it. What
400 // we need to do instead is pass in a closure that will be executed when the slice is
401 // either mapped or the mapping has failed.
402 //
403 // The problem with this is that we don't have a reliable way to wait in the main
404 // code for the buffer to be mapped and even worse, calling get_mapped_range or
405 // get_mapped_range_mut prematurely will cause a panic, not return an error.
406 //
407 // Using channels solves this as awaiting the receiving of a message from
408 // the passed closure will force the outside code to wait. It also doesn't hurt
409 // if the closure finishes before the outside code catches up as the message is
410 // buffered and receiving will just pick that up.
411 //
412 // It may also be worth noting that although on native, the usage of asynchronous
413 // channels is wholly unnecessary, for the sake of portability to Wasm
414 // we'll use async channels that work on both native and Wasm.
415
416 let (s, r) = crossbeam_channel::bounded(1);
417
418 // Maps the buffer so it can be read on the cpu
419 buffer_slice.map_async(MapMode::Read, move |r| match r {
420 // This will execute once the gpu is ready, so after the call to poll()
421 Ok(r) => s.send(r).expect("Failed to send map update"),
422 Err(err) => panic!("Failed to map buffer {err}"),
423 });
424
425 // In order for the mapping to be completed, one of three things must happen.
426 // One of those can be calling `Device::poll`. This isn't necessary on the web as devices
427 // are polled automatically but natively, we need to make sure this happens manually.
428 // `Maintain::Wait` will cause the thread to wait on native but not on WebGpu.
429
430 // This blocks until the gpu is done executing everything
431 render_device
432 .poll(PollType::wait_indefinitely())
433 .expect("Failed to poll device for map async");
434
435 // This blocks until the buffer is mapped
436 r.recv().expect("Failed to receive the map_async message");
437
438 // This could fail on app exit, if Main world clears resources (including receiver) while Render world still renders
439 let _ = sender.send(buffer_slice.get_mapped_range().unwrap().to_vec());
440
441 // We need to make sure all `BufferView`'s are dropped before we do what we're about
442 // to do.
443 // Unmap so that we can copy to the staging buffer in the next iteration.
444 image_copier.buffer.unmap();
445 }
446 }
447
448 /// CPU-side image for saving
449 #[derive(Component, Deref, DerefMut)]
450 struct ImageToSave(Handle<Image>);
451
452 // Takes from channel image content sent from render world and saves it to disk
453 fn update(
454 images_to_save: Query<&ImageToSave>,
455 receiver: Res<MainWorldReceiver>,
456 mut images: ResMut<Assets<Image>>,
457 mut scene_controller: ResMut<SceneController>,
458 mut app_exit_writer: MessageWriter<AppExit>,
459 mut file_number: Local<u32>,
460 ) {
461 if let SceneState::Render(n) = scene_controller.state {
462 if n < 1 {
463 // We don't want to block the main world on this,
464 // so we use try_recv which attempts to receive without blocking
465 let mut image_data = Vec::new();
466 while let Ok(data) = receiver.try_recv() {
467 // image generation could be faster than saving to fs,
468 // that's why use only last of them
469 image_data = data;
470 }
471 if !image_data.is_empty() {
472 for image in images_to_save.iter() {
473 // Fill correct data from channel to image
474 let mut img_bytes = images.get_mut(image.id()).unwrap();
475
476 // We need to ensure that this works regardless of the image dimensions
477 // If the image became wider when copying from the texture to the buffer,
478 // then the data is reduced to its original size when copying from the buffer to the image.
479 let row_bytes = img_bytes.width() as usize
480 * img_bytes.texture_descriptor.format.pixel_size().unwrap();
481 let aligned_row_bytes = RenderDevice::align_copy_bytes_per_row(row_bytes);
482 if row_bytes == aligned_row_bytes {
483 img_bytes.data.as_mut().unwrap().clone_from(&image_data);
484 } else {
485 // shrink data to original image size
486 img_bytes.data = Some(
487 image_data
488 .chunks(aligned_row_bytes)
489 .take(img_bytes.height() as usize)
490 .flat_map(|row| &row[..row_bytes.min(row.len())])
491 .cloned()
492 .collect(),
493 );
494 }
495
496 // Create RGBA Image Buffer
497 let img = match img_bytes.clone().try_into_dynamic() {
498 Ok(img) => img.to_rgba8(),
499 Err(e) => panic!("Failed to create image buffer {e:?}"),
500 };
501
502 // Prepare directory for images, test_images in bevy folder is used here for example
503 // You should choose the path depending on your needs
504 let images_dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("test_images");
505 info!("Saving image to: {images_dir:?}");
506 std::fs::create_dir_all(&images_dir).unwrap();
507
508 // Choose filename starting from 000.png
509 let image_path = images_dir.join(format!("{:03}.png", file_number.deref()));
510 *file_number.deref_mut() += 1;
511
512 // Finally saving image to file, this heavy blocking operation is kept here
513 // for example simplicity, but in real app you should move it to a separate task
514 if let Err(e) = img.save(image_path) {
515 panic!("Failed to save image: {e}");
516 };
517 }
518 if scene_controller.single_image {
519 app_exit_writer.write(AppExit::Success);
520 }
521 }
522 } else {
523 // clears channel for skipped frames
524 while receiver.try_recv().is_ok() {}
525 scene_controller.state = SceneState::Render(n - 1);
526 }
527 }
528 }