Studio Headless Mode (#36)

This PR adds the capaibility to run Stroyteller studio from the command line, rendering the result to images. This is the basic functionality, and there are somethings that I am saving for later when we make this a more fully fledged feature. Limitations: - Right now, you cannot do image output resolutions that aren't 2^n. This is a bug I plan on tackling later. If you have questions on how the command line works, use the --help argument. Let me know what you think! --------- Co-authored-by: Danny McGee <dannymcgee@gmail.com>
2026-10-28 08:58:07 +00:00 · 2024-02-19 16:00:30 -07:00
parent a1779fb1ec
commit f1c1b85a98
11 changed files with 1016 additions and 28 deletions
@@ -0,0 +1,285 @@
+use bevy::{
+    core_pipeline::{core_3d, prepass::ViewPrepassTextures},
+    prelude::*,
+    render::{
+        render_graph::{RenderGraphApp, ViewNode, ViewNodeRunner},
+        render_resource::{
+            Buffer, BufferDescriptor, BufferUsages, ImageCopyBuffer, ImageDataLayout, MapMode,
+        },
+        renderer::RenderDevice,
+        view::ViewTarget,
+        Extract, Render, RenderSet,
+    },
+};
+use bytemuck::AnyBitPattern;
+use image::{
+    error::UnsupportedErrorKind, EncodableLayout, ImageBuffer, ImageError, Pixel,
+    PixelWithColorType, Rgba,
+};
+use wgpu::{Extent3d, MaintainBase};
+
+use crate::scene::SceneState;
+
+//==============================================================================
+//                        Screen Shot Plugin
+//==============================================================================
+
+pub struct ScreenShotPlugin;
+
+impl Plugin for ScreenShotPlugin {
+    fn build(&self, app: &mut App) {
+        app.add_systems(
+            Update,
+            (test_system, update_screen_capture_on_extract_step)
+                .run_if(in_state(SceneState::Active)),
+        );
+
+        let Ok(render_world) = app.get_sub_app_mut(bevy::render::RenderApp) else {
+            return;
+        };
+
+        render_world
+            .add_systems(ExtractSchedule, extract_screen_capture_components)
+            .add_systems(
+                Render,
+                (
+                    prepare_screen_capture_buffers.in_set(RenderSet::PrepareResources),
+                    save_buffer_to_file
+                        .after(RenderSet::Render)
+                        .before(RenderSet::Cleanup),
+                ),
+            )
+            //Render Graph Additions
+            .add_render_graph_node::<ViewNodeRunner<ScreenCaptureNode>>(
+                core_3d::graph::NAME,
+                ScreenCaptureNode::NAME,
+            )
+            .add_render_graph_edges(
+                core_3d::graph::NAME,
+                &[
+                    core_3d::graph::node::TONEMAPPING,
+                    ScreenCaptureNode::NAME,
+                    core_3d::graph::node::END_MAIN_PASS_POST_PROCESSING,
+                ],
+            );
+    }
+}
+
+fn test_system(
+    mut commands: Commands,
+    camera: Query<Entity, With<Camera>>,
+    input: Res<Input<KeyCode>>,
+) {
+    if input.just_pressed(KeyCode::F10) {
+        for entity in camera.iter() {
+            commands
+                .entity(entity)
+                .insert(ScreenCapture::number_of_frames(16, "./scrnshot/"));
+        }
+    }
+}
+
+//==============================================================================
+//                        Sceen Shot Component
+//==============================================================================
+
+#[derive(Component, Clone)]
+pub struct ScreenCapture {
+    number_of_frames: u32,
+    to_directory: String,
+    capture_depth: bool,
+}
+
+impl ScreenCapture {
+    pub fn number_of_frames(frames: u32, output_dir: &str) -> Self {
+        ScreenCapture {
+            number_of_frames: frames,
+            to_directory: output_dir.to_string(),
+            capture_depth: false,
+        }
+    }
+
+    pub fn capture_depth(mut self) -> Self {
+        self.capture_depth = true;
+        self
+    }
+
+    pub fn is_complete(&self) -> bool {
+        println!("Number of Frames Left: {}", self.number_of_frames);
+        self.number_of_frames == 1
+    }
+}
+
+fn update_screen_capture_on_extract_step(
+    mut commands: Commands,
+    mut screen_captures: Query<(Entity, &mut ScreenCapture)>,
+) {
+    for (entity, mut screen_capture) in screen_captures.iter_mut() {
+        screen_capture.number_of_frames -= 1;
+        if screen_capture.number_of_frames == 0 {
+            commands.entity(entity).remove::<ScreenCapture>();
+        }
+    }
+}
+
+//This system extracts the ScreenCaptureComponent to the render world.
+fn extract_screen_capture_components(
+    mut commands: Commands,
+    screen_capure_comps: Extract<Query<(Entity, &ScreenCapture), With<Camera>>>,
+) {
+    for (entity, screen_capture) in screen_capure_comps.iter() {
+        commands.get_or_spawn(entity).insert(screen_capture.clone());
+    }
+}
+
+//==============================================================================
+//                        ScreenCaptureBuffers
+//==============================================================================
+
+#[derive(Component)]
+struct ScreenCaptureBuffers {
+    main_buffer: Buffer,
+    main_buffer_extent: Extent3d,
+    main_bytes_per_row: u32,
+    main_padded_bytes_per_row: u32,
+}
+
+//During this function, we are making the buffers for the screenshot to be saved in
+//This is where we are instatiating the ScreenCaptureBuffers component.
+fn prepare_screen_capture_buffers(
+    mut commands: Commands,
+    views: Query<(Entity, &ViewTarget, &ScreenCapture)>,
+    device: Res<RenderDevice>,
+) {
+    for (entity, view_target, screen_capture) in views.iter() {
+        if screen_capture.number_of_frames > 0 {
+            let texture = view_target.main_texture();
+            let size = texture.size();
+            let format = texture.format();
+            let bytes_per_row =
+                (size.width / format.block_dimensions().0) * format.block_size(None).unwrap();
+            let padded_bytes_per_row =
+                RenderDevice::align_copy_bytes_per_row(bytes_per_row as usize) as u32;
+
+            let main_buffer = device.create_buffer(&BufferDescriptor {
+                label: None,
+                size: (size.height * bytes_per_row) as u64,
+                usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
+                mapped_at_creation: false,
+            });
+
+            commands.entity(entity).insert(ScreenCaptureBuffers {
+                main_buffer,
+                main_buffer_extent: size,
+                main_bytes_per_row: bytes_per_row,
+                main_padded_bytes_per_row: padded_bytes_per_row,
+            });
+
+            println!("Buffers Created!");
+        }
+    }
+}
+
+//==============================================================================
+//                        Screen Shot View Node
+//==============================================================================
+
+#[derive(Default)]
+struct ScreenCaptureNode;
+
+impl ScreenCaptureNode {
+    pub const NAME: &'static str = "ScreenCaptureNode";
+}
+
+impl ViewNode for ScreenCaptureNode {
+    type ViewQuery = (
+        Entity,
+        &'static ViewTarget,
+        &'static ScreenCapture,
+        &'static ScreenCaptureBuffers,
+        Option<&'static ViewPrepassTextures>,
+    );
+
+    fn run(
+        &self,
+        _: &mut bevy::render::render_graph::RenderGraphContext,
+        render_context: &mut bevy::render::renderer::RenderContext,
+        view_query: bevy::ecs::query::QueryItem<Self::ViewQuery>,
+        _: &World,
+    ) -> Result<(), bevy::render::render_graph::NodeRunError> {
+        render_context.command_encoder().copy_texture_to_buffer(
+            view_query.1.main_texture().as_image_copy(),
+            ImageCopyBuffer {
+                buffer: &view_query.3.main_buffer,
+                layout: ImageDataLayout {
+                    offset: 0,
+                    bytes_per_row: Some(view_query.3.main_padded_bytes_per_row),
+                    rows_per_image: None,
+                },
+            },
+            view_query.1.main_texture().size(),
+        );
+        println!("Saved to buffer!");
+        Ok(())
+    }
+}
+
+//==============================================================================
+//                        Saving the buffer to file
+//==============================================================================
+
+fn save_buffer_to_file(
+    buffers: Query<(&ScreenCapture, &ScreenCaptureBuffers)>,
+    device: Res<RenderDevice>,
+) {
+    for (screen_capture, buffers) in buffers.iter() {
+        if buffers.main_bytes_per_row != buffers.main_padded_bytes_per_row {
+            println!("These don't match!");
+        }
+
+        let main_image_values = {
+            let main_buffer_slice = buffers.main_buffer.slice(..);
+
+            let (tx, rx) = futures::channel::oneshot::channel();
+            main_buffer_slice.map_async(MapMode::Read, move |result| {
+                tx.send(result).unwrap();
+            });
+            device.poll(MaintainBase::Wait);
+            futures_lite::future::block_on(rx).unwrap().unwrap();
+
+            main_buffer_slice.get_mapped_range().to_vec()
+        };
+
+        std::fs::create_dir_all(&screen_capture.to_directory)
+            .expect("Output path could not be created");
+
+        let path = format!(
+            "{}img_{:0.5}.png",
+            screen_capture.to_directory, screen_capture.number_of_frames
+        );
+
+        save_buffer::<Rgba<u8>>(&main_image_values, &buffers.main_buffer_extent, &path);
+    }
+}
+
+fn save_buffer<P: Pixel + PixelWithColorType>(
+    image_bytes: &[P::Subpixel],
+    source_size: &Extent3d,
+    path: &str,
+) where
+    P::Subpixel: AnyBitPattern,
+    [P::Subpixel]: EncodableLayout,
+{
+    match ImageBuffer::<P, _>::from_raw(source_size.width, source_size.height, image_bytes) {
+        Some(buffer) => {
+            if let Err(ImageError::Unsupported(err)) = buffer.save(path) {
+                if let UnsupportedErrorKind::Format(hint) = err.kind() {
+                    println!("Image format {} is not supported", hint);
+                }
+            }
+        }
+        None => {
+            println!("Failed creating image buffer for '{}'", path);
+        }
+    }
+}