mirror of
https://github.com/storytold/storyteller-bevy.git
synced 2026-10-18 07:35:19 +00:00
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>
This commit is contained in:
Brooks Palin
@@ -24,6 +24,7 @@ pub struct BvhPlugin;
|
||||
impl Plugin for BvhPlugin {
|
||||
fn build(&self, app: &mut App) {
|
||||
app.init_asset::<Bvh>()
|
||||
.add_event::<BvhAnimCompleteEvent>()
|
||||
.register_asset_loader(BvhLoader)
|
||||
// FIXME: Need to add a state for this module (similar to scene
|
||||
// module's `SceneState`) to manage the order/frequency at
|
||||
@@ -106,6 +107,7 @@ fn compute_inverse_bindposes(
|
||||
pub fn animate_bvh(
|
||||
mut q_bvh_anims: Query<(&mut BvhAnim, &Children)>,
|
||||
mut q_bvh_joints: Query<JointQuery, Without<BvhAnim>>,
|
||||
mut bvh_event_writer: EventWriter<BvhAnimCompleteEvent>,
|
||||
r_bvh_assets: Res<Assets<Bvh>>,
|
||||
r_time: Res<Time>,
|
||||
) {
|
||||
@@ -122,6 +124,9 @@ pub fn animate_bvh(
|
||||
|
||||
if anim.timer.just_finished() {
|
||||
anim.frame = (anim.frame + 1) % bvh_len;
|
||||
if anim.frame == 0 {
|
||||
bvh_event_writer.send(BvhAnimCompleteEvent);
|
||||
}
|
||||
}
|
||||
|
||||
anim.timer.percent()
|
||||
@@ -204,6 +209,16 @@ fn lerp_frames<T: ReadOnlyWorldQuery>(
|
||||
}
|
||||
}
|
||||
|
||||
//==============================================================================
|
||||
// Bvh Anim Events
|
||||
//==============================================================================
|
||||
|
||||
// This may seem weird, but this is for the Headless mode to
|
||||
// know when to stop recording.
|
||||
|
||||
#[derive(Event)]
|
||||
pub struct BvhAnimCompleteEvent;
|
||||
|
||||
//==============================================================================
|
||||
// Bvh Anim Component
|
||||
//==============================================================================
|
||||
|
||||
+119
-2
@@ -1,11 +1,20 @@
|
||||
#![allow(clippy::type_complexity, clippy::too_many_arguments)]
|
||||
|
||||
use bevy::{app::PluginGroupBuilder, prelude::*};
|
||||
use std::time::Duration;
|
||||
|
||||
use bevy::{
|
||||
app::{AppExit, PluginGroupBuilder, ScheduleRunnerPlugin},
|
||||
prelude::*,
|
||||
};
|
||||
use bevy_mod_picking::DefaultPickingPlugins;
|
||||
use bevy_web_asset::WebAssetPlugin;
|
||||
use bvh::{BvhAnimCompleteEvent, TestBvh};
|
||||
use screenshot::{ScreenCapture, ScreenShotPlugin};
|
||||
use space_editor::prelude::{EditorEvent, EditorPrefabPath};
|
||||
|
||||
#[cfg(feature = "wasm")]
|
||||
use wasm_bindgen::prelude::*;
|
||||
use wgpu::Extent3d;
|
||||
|
||||
mod anim;
|
||||
mod bvh;
|
||||
@@ -20,6 +29,7 @@ pub mod math;
|
||||
mod mixamo;
|
||||
mod persistance;
|
||||
mod scene;
|
||||
pub mod screenshot;
|
||||
mod skybox;
|
||||
mod ui;
|
||||
pub mod viewport;
|
||||
@@ -53,6 +63,7 @@ pub(crate) struct MainCamera;
|
||||
#[derive(Resource, Clone, Copy, PartialEq, Eq, Debug)]
|
||||
pub enum StudioMode {
|
||||
Editor,
|
||||
Headless,
|
||||
Viewer,
|
||||
}
|
||||
|
||||
@@ -62,7 +73,7 @@ pub(crate) fn is_editor(r_studio_mode: Res<StudioMode>) -> bool {
|
||||
|
||||
/// The collection of plugins that should be loaded for both viewer-mode and
|
||||
/// editor-mode instances of the app.
|
||||
struct CommonPlugins;
|
||||
pub struct CommonPlugins;
|
||||
impl PluginGroup for CommonPlugins {
|
||||
fn build(self) -> PluginGroupBuilder {
|
||||
PluginGroupBuilder::start::<Self>()
|
||||
@@ -272,3 +283,109 @@ pub fn start_internal(
|
||||
|
||||
app
|
||||
}
|
||||
|
||||
pub enum HeadlessMode {
|
||||
Bvh(String, String),
|
||||
Scene(String, String),
|
||||
}
|
||||
|
||||
pub fn start_headless(
|
||||
mode: HeadlessMode,
|
||||
fps: u32,
|
||||
outout_dir: &str,
|
||||
frames_to_capture: u32,
|
||||
image_res: Extent3d,
|
||||
) -> App {
|
||||
let mut app = App::new();
|
||||
|
||||
app.insert_resource(StudioMode::Headless);
|
||||
app.insert_resource(HeadlessSettings {
|
||||
frames_to_capture,
|
||||
output_dims: image_res,
|
||||
output_dir: outout_dir.to_string(),
|
||||
});
|
||||
|
||||
app.add_plugins((
|
||||
ScheduleRunnerPlugin::run_loop(Duration::from_secs_f64(1.0 / fps as f64)),
|
||||
DefaultPlugins.set(WindowPlugin {
|
||||
primary_window: None,
|
||||
exit_condition: bevy::window::ExitCondition::DontExit,
|
||||
close_when_requested: true,
|
||||
}),
|
||||
ScreenShotPlugin,
|
||||
CommonPlugins,
|
||||
));
|
||||
|
||||
match mode {
|
||||
HeadlessMode::Bvh(bvh_dir, skybox) => {
|
||||
app.add_plugins(BvhPlugin);
|
||||
app.add_systems(
|
||||
Startup,
|
||||
move |mut cmd: Commands,
|
||||
asset_server: Res<AssetServer>,
|
||||
mut ew: EventWriter<LoadSceneEvent>| {
|
||||
cmd.insert_resource(SkyboxId(skybox.clone()));
|
||||
cmd.insert_resource(TestBvh(asset_server.load(bvh_dir.clone())));
|
||||
|
||||
ew.send(LoadSceneEvent("Mannequin.gltf".to_string()));
|
||||
},
|
||||
);
|
||||
|
||||
app.add_systems(PostUpdate, check_for_capture_complete);
|
||||
}
|
||||
HeadlessMode::Scene(scene, skybox) => {
|
||||
app.add_systems(
|
||||
Startup,
|
||||
move |mut cmd: Commands, mut ew: EventWriter<EditorEvent>| {
|
||||
cmd.insert_resource(SkyboxId(skybox.clone()));
|
||||
ew.send(EditorEvent::Load(EditorPrefabPath::File(scene.clone())))
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
app
|
||||
}
|
||||
|
||||
pub fn is_headless(studio_state: Res<StudioMode>) -> bool {
|
||||
*studio_state == StudioMode::Headless
|
||||
}
|
||||
|
||||
pub fn is_not_headless(studio_state: Res<StudioMode>) -> bool {
|
||||
*studio_state != StudioMode::Headless
|
||||
}
|
||||
|
||||
#[derive(Resource)]
|
||||
pub struct HeadlessSettings {
|
||||
pub output_dims: Extent3d,
|
||||
pub output_dir: String,
|
||||
pub frames_to_capture: u32,
|
||||
}
|
||||
|
||||
impl Default for HeadlessSettings {
|
||||
fn default() -> Self {
|
||||
HeadlessSettings {
|
||||
output_dims: Extent3d {
|
||||
width: 1280,
|
||||
height: 720,
|
||||
depth_or_array_layers: 1,
|
||||
},
|
||||
output_dir: "./screen_shots/".to_string(),
|
||||
frames_to_capture: 1,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn check_for_capture_complete(
|
||||
main_camera: Query<&ScreenCapture, With<MainCamera>>,
|
||||
mut close_app: EventWriter<AppExit>,
|
||||
bvh_event_reader: EventReader<BvhAnimCompleteEvent>,
|
||||
) {
|
||||
let Ok(screen_capture) = main_camera.get_single() else {
|
||||
return;
|
||||
};
|
||||
|
||||
if screen_capture.is_complete() || !bvh_event_reader.is_empty() {
|
||||
close_app.send(AppExit);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,137 @@
|
||||
use bevy::render::render_resource::Extent3d;
|
||||
use clap::{Args, Parser, Subcommand};
|
||||
use studio::{start_headless, HeadlessMode};
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "StorytellerStudioHeadless")]
|
||||
#[command(author = "The Storyteller Team")]
|
||||
#[command(version = "1.0")]
|
||||
#[command(
|
||||
about = "Executable for Storyteller Studio Headless. Allows you to render frames to images."
|
||||
)]
|
||||
#[command(propagate_version = true)]
|
||||
struct StudioArgs {
|
||||
#[command(subcommand)]
|
||||
subcommand: Subcommands,
|
||||
}
|
||||
|
||||
#[derive(Subcommand)]
|
||||
enum Subcommands {
|
||||
Scene {
|
||||
scene_dir: String,
|
||||
#[arg(long, short)]
|
||||
output_dir: Option<String>,
|
||||
#[arg(long)]
|
||||
skybox: Option<String>,
|
||||
#[arg(long = "res", value_parser = extent3d_parser)]
|
||||
output_resolution: Option<Extent3d>,
|
||||
#[arg(long, value_parser = clap::value_parser!(u32).range(1..))]
|
||||
fps: Option<u32>,
|
||||
#[arg(long, value_parser = clap::value_parser!(u32).range(1..))]
|
||||
frames: Option<u32>,
|
||||
},
|
||||
Bvh {
|
||||
bvh_dir: String,
|
||||
#[arg(long, short = 'o', short)]
|
||||
out: Option<String>,
|
||||
#[arg(long)]
|
||||
skybox: Option<String>,
|
||||
#[arg(long = "res", value_parser = extent3d_parser)]
|
||||
output_resolution: Option<Extent3d>,
|
||||
#[arg(long, value_parser = clap::value_parser!(u32).range(1..))]
|
||||
fps: Option<u32>,
|
||||
#[arg(long, value_parser = clap::value_parser!(u32).range(1..))]
|
||||
frames: Option<u32>,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Args, Clone)]
|
||||
struct CommonArgs {}
|
||||
|
||||
fn main() {
|
||||
let args = StudioArgs::parse();
|
||||
|
||||
let mut app = match args.subcommand {
|
||||
Subcommands::Scene {
|
||||
scene_dir,
|
||||
skybox,
|
||||
output_dir,
|
||||
output_resolution,
|
||||
fps,
|
||||
frames,
|
||||
} => {
|
||||
let output_dir = output_dir.unwrap_or("./screen_shots/".to_string());
|
||||
let output_res = output_resolution.unwrap_or(Extent3d {
|
||||
width: 1280,
|
||||
height: 720,
|
||||
depth_or_array_layers: 1,
|
||||
});
|
||||
let fps = fps.unwrap_or(24);
|
||||
let frames = frames.unwrap_or(30);
|
||||
let skybox = skybox.unwrap_or("gum_trees_4k".to_string());
|
||||
start_headless(
|
||||
HeadlessMode::Scene(scene_dir, skybox),
|
||||
fps,
|
||||
&output_dir,
|
||||
frames,
|
||||
output_res,
|
||||
)
|
||||
}
|
||||
Subcommands::Bvh {
|
||||
bvh_dir,
|
||||
skybox,
|
||||
out: output_dir,
|
||||
output_resolution,
|
||||
fps,
|
||||
frames,
|
||||
} => {
|
||||
let output_dir = output_dir.unwrap_or("./screen_shots/".to_string());
|
||||
let output_res = output_resolution.unwrap_or(Extent3d {
|
||||
width: 1280,
|
||||
height: 720,
|
||||
depth_or_array_layers: 1,
|
||||
});
|
||||
let fps = fps.unwrap_or(24);
|
||||
let frames = frames.unwrap_or(u32::MAX);
|
||||
let skybox = skybox.unwrap_or("gum_trees_4k".to_string());
|
||||
start_headless(
|
||||
HeadlessMode::Bvh(bvh_dir, skybox),
|
||||
fps,
|
||||
&output_dir,
|
||||
frames,
|
||||
output_res,
|
||||
)
|
||||
}
|
||||
};
|
||||
|
||||
app.run();
|
||||
}
|
||||
|
||||
//==============================================================================
|
||||
// Parsers
|
||||
//==============================================================================
|
||||
|
||||
fn extent3d_parser(input: &str) -> Result<Extent3d, String> {
|
||||
let split = input.split('x').collect::<Vec<_>>();
|
||||
|
||||
if split.len() != 2 {
|
||||
return Err(
|
||||
"Invalid format for the Resolution argument, must follow the following pattern: <WIDTH>x<HEIGHT>"
|
||||
.to_string()
|
||||
);
|
||||
}
|
||||
|
||||
let Ok(width) = split[0].parse::<u32>() else {
|
||||
return Err("Width must be a number when defining the output resolution.".to_string());
|
||||
};
|
||||
|
||||
let Ok(height) = split[1].parse::<u32>() else {
|
||||
return Err("Width must be a number when defining the output resolution.".to_string());
|
||||
};
|
||||
|
||||
Ok(Extent3d {
|
||||
width,
|
||||
height,
|
||||
depth_or_array_layers: 1,
|
||||
})
|
||||
}
|
||||
+21
-3
@@ -1,11 +1,28 @@
|
||||
use bevy::{app::Update, window::close_on_esc};
|
||||
use studio::{StudioMode, ViewportSize};
|
||||
use clap::Parser;
|
||||
use studio::{screenshot::ScreenShotPlugin, StudioMode, ViewportSize};
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "StorytellerStudio")]
|
||||
#[command(author = "The Storyteller Team")]
|
||||
#[command(version = "1.0")]
|
||||
#[command(about = "Executable for Storyteller Studio.")]
|
||||
#[command(propagate_version = true)]
|
||||
struct StudioArgs {
|
||||
model: Option<String>,
|
||||
skybox: Option<String>,
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let args = StudioArgs::parse();
|
||||
|
||||
let model_path = args.model.unwrap_or("Mannequin.gltf".to_string());
|
||||
let skybox_path = args.skybox.unwrap_or("gum_trees_4k".to_string());
|
||||
|
||||
// Edit the arguments here to change the default startup scene or skybox
|
||||
let mut app = studio::start_internal(
|
||||
"bookcase.gltf",
|
||||
"gum_trees_4k",
|
||||
model_path.as_str(),
|
||||
skybox_path.as_str(),
|
||||
"", // ignored for non-WASM
|
||||
"Storyteller Studio",
|
||||
ViewportSize::default(),
|
||||
@@ -14,6 +31,7 @@ fn main() {
|
||||
|
||||
// Add any additional systems, resources, etc. for local debugging
|
||||
app.add_systems(Update, close_on_esc);
|
||||
app.add_plugins(ScreenShotPlugin);
|
||||
|
||||
app.run();
|
||||
}
|
||||
|
||||
+61
-3
@@ -9,19 +9,22 @@ use bevy::{
|
||||
math::vec2,
|
||||
pbr::ShadowFilteringMethod,
|
||||
prelude::*,
|
||||
render::primitives::Aabb,
|
||||
render::{camera::RenderTarget, primitives::Aabb},
|
||||
scene::InstanceId,
|
||||
};
|
||||
|
||||
use space_editor::{prelude::PrefabMarker, space_prefab::component::SceneAutoChild};
|
||||
#[cfg(feature = "wasm")]
|
||||
use wasm_bindgen::prelude::wasm_bindgen;
|
||||
use wgpu::{TextureDescriptor, TextureUsages};
|
||||
|
||||
use crate::{
|
||||
camera_controller::{CameraController, CameraControllerBundle},
|
||||
debug::DebugPoint,
|
||||
is_headless, is_not_headless,
|
||||
math::{Combinable, Extrema},
|
||||
MainCamera,
|
||||
screenshot::ScreenCapture,
|
||||
HeadlessSettings, MainCamera,
|
||||
};
|
||||
|
||||
pub struct ScenePlugin;
|
||||
@@ -35,7 +38,13 @@ impl Plugin for ScenePlugin {
|
||||
.insert_resource(Msaa::Off)
|
||||
.add_state::<SceneState>()
|
||||
.add_event::<LoadSceneEvent>()
|
||||
.add_systems(Startup, initialize)
|
||||
.add_systems(
|
||||
Startup,
|
||||
(
|
||||
initialize.run_if(is_not_headless),
|
||||
initialiaze_headless.run_if(is_headless),
|
||||
),
|
||||
)
|
||||
.add_systems(Update, init_scene_loading.in_set(SceneSpawnerSystem))
|
||||
.add_systems(
|
||||
OnEnter(SceneState::Active),
|
||||
@@ -210,6 +219,55 @@ fn initialize(mut cmd: Commands) {
|
||||
));
|
||||
}
|
||||
|
||||
fn initialiaze_headless(
|
||||
mut cmd: Commands,
|
||||
mut image_assets: ResMut<Assets<Image>>,
|
||||
headless_settings: Res<HeadlessSettings>,
|
||||
) {
|
||||
// Camera, graphics settings
|
||||
let focus = Vec3::ZERO;
|
||||
let camera_xform = Transform::from_xyz(-2.5, 4.5, 5.0).looking_at(focus, Vec3::Y);
|
||||
|
||||
let mut test_image = Image {
|
||||
texture_descriptor: TextureDescriptor {
|
||||
label: Some("render_texture"),
|
||||
size: headless_settings.output_dims,
|
||||
mip_level_count: 1,
|
||||
sample_count: 1,
|
||||
dimension: wgpu::TextureDimension::D2,
|
||||
format: wgpu::TextureFormat::Rgba8UnormSrgb,
|
||||
usage: TextureUsages::TEXTURE_BINDING
|
||||
| TextureUsages::COPY_SRC
|
||||
| TextureUsages::COPY_DST
|
||||
| TextureUsages::RENDER_ATTACHMENT,
|
||||
view_formats: &[],
|
||||
},
|
||||
..Default::default()
|
||||
};
|
||||
|
||||
test_image.resize(headless_settings.output_dims);
|
||||
|
||||
cmd.spawn((
|
||||
Name::new("Main Camera"),
|
||||
MainCamera,
|
||||
Camera3dBundle {
|
||||
transform: camera_xform,
|
||||
tonemapping: Tonemapping::BlenderFilmic,
|
||||
camera: Camera {
|
||||
hdr: false,
|
||||
target: RenderTarget::Image(image_assets.add(test_image)),
|
||||
..Default::default()
|
||||
},
|
||||
..default()
|
||||
},
|
||||
CameraControllerBundle::new(focus, camera_xform.translation),
|
||||
ScreenCapture::number_of_frames(
|
||||
headless_settings.frames_to_capture,
|
||||
&headless_settings.output_dir,
|
||||
),
|
||||
));
|
||||
}
|
||||
|
||||
fn init_scene_loading(mut cmd: Commands, mut er: EventReader<LoadSceneEvent>) {
|
||||
if let Some(id) = er.read().last() {
|
||||
info!("Loading scene: {}", id.0);
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user