diff --git a/studio/src/lib.rs b/studio/src/lib.rs index 62ee32d..a449892 100644 --- a/studio/src/lib.rs +++ b/studio/src/lib.rs @@ -244,6 +244,7 @@ pub fn start_internal( let mut app = App::new(); + app.insert_resource(HeadlessMode::None); app.insert_resource(studio_mode); app.add_plugins(( @@ -284,9 +285,12 @@ pub fn start_internal( app } +#[derive(Debug, Resource, Clone)] pub enum HeadlessMode { Bvh(String, String), Scene(String, String), + Mixamo(String, String), + None, } pub fn start_headless( @@ -298,9 +302,10 @@ pub fn start_headless( ) -> App { let mut app = App::new(); + app.insert_resource(mode.clone()); app.insert_resource(StudioMode::Headless); app.insert_resource(HeadlessSettings { - frames_to_capture, + max_frames: frames_to_capture, output_dims: image_res, output_dir: outout_dir.to_string(), }); @@ -330,7 +335,7 @@ pub fn start_headless( }, ); - app.add_systems(PostUpdate, check_for_capture_complete); + app.add_systems(PostUpdate, check_for_bvh_capture_complete); } HeadlessMode::Scene(scene, skybox) => { app.add_systems( @@ -341,6 +346,23 @@ pub fn start_headless( }, ); } + HeadlessMode::Mixamo(mixamo_path, skybox) => { + app.add_plugins(MixamoPlugin); + app.add_systems( + Startup, + move |mut commands: Commands, + mut ew: EventWriter, + mut load_scene_event: EventWriter| { + ew.send(MixamoImportEvent(mixamo_path.clone())); + load_scene_event.send(LoadSceneEvent("Mannequin.gltf".to_string())); + commands.insert_resource(SkyboxId(skybox.clone())); + }, + ); + app.add_systems(PostUpdate, check_for_mixamo_capture_complete); + + app.run(); + } + HeadlessMode::None => {} } app @@ -358,7 +380,7 @@ pub fn is_not_headless(studio_state: Res) -> bool { pub struct HeadlessSettings { pub output_dims: Extent3d, pub output_dir: String, - pub frames_to_capture: u32, + pub max_frames: u32, } impl Default for HeadlessSettings { @@ -370,23 +392,45 @@ impl Default for HeadlessSettings { depth_or_array_layers: 1, }, output_dir: "./screen_shots/".to_string(), - frames_to_capture: 1, + max_frames: u32::MAX, } } } -fn check_for_capture_complete( - mut commands: Commands, +fn check_for_bvh_capture_complete( mut close_app: EventWriter, - main_camera: Query<(Entity, &ScreenCapture), With>, + mut main_camera: Query<&mut ScreenCapture, With>, bvh_event_reader: EventReader, ) { - let Ok((entity, screen_capture)) = main_camera.get_single() else { + let Ok(mut screen_capture) = main_camera.get_single_mut() else { return; }; - if screen_capture.is_complete() || !bvh_event_reader.is_empty() { - commands.entity(entity).remove::(); + println!("Frames captured {}", screen_capture.frames_captured()); + if !screen_capture.should_capture() || !bvh_event_reader.is_empty() { + screen_capture.stop_capturing(); close_app.send(AppExit); } } + +fn check_for_mixamo_capture_complete( + mut close_app: EventWriter, + mut main_camera: Query<&mut ScreenCapture, With>, + mixamo_anims: Query<&AnimationPlayer>, +) { + let Ok(mut capture) = main_camera.get_single_mut() else { + return; + }; + + for anim in mixamo_anims.iter() { + println!( + "Time elapsed {} | Frames Captured : {}", + anim.elapsed(), + capture.frames_captured() + ); + if anim.completions() != 0 { + capture.stop_capturing(); + close_app.send(AppExit); + } + } +} diff --git a/studio/src/main-headless.rs b/studio/src/main-headless.rs index e8edcc0..da01023 100644 --- a/studio/src/main-headless.rs +++ b/studio/src/main-headless.rs @@ -43,6 +43,19 @@ enum Subcommands { #[arg(long, value_parser = clap::value_parser!(u32).range(1..))] frames: Option, }, + Mixamo { + mixamo_dir: String, + #[arg(long, short = 'o', short)] + out: Option, + #[arg(long)] + skybox: Option, + #[arg(long = "res", value_parser = extent3d_parser)] + output_resolution: Option, + #[arg(long, value_parser = clap::value_parser!(u32).range(1..))] + fps: Option, + #[arg(long, value_parser = clap::value_parser!(u32).range(1..))] + frames: Option, + }, } #[derive(Args, Clone)] @@ -102,6 +115,31 @@ fn main() { output_res, ) } + Subcommands::Mixamo { + mixamo_dir, + out, + skybox, + output_resolution, + fps, + frames, + } => { + let output_dir = out.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::Mixamo(mixamo_dir, skybox), + fps, + &output_dir, + frames, + output_res, + ) + } }; app.run(); diff --git a/studio/src/scene.rs b/studio/src/scene.rs index 2f04039..ed32e68 100644 --- a/studio/src/scene.rs +++ b/studio/src/scene.rs @@ -48,7 +48,12 @@ impl Plugin for ScenePlugin { .add_systems(Update, init_scene_loading.in_set(SceneSpawnerSystem)) .add_systems( OnEnter(SceneState::Active), - (tag_scene_elements, frame_scene_in_view).in_set(SceneSpawnerSystem), + ( + tag_scene_elements, + frame_scene_in_view, + start_capture.run_if(is_headless), + ) + .in_set(SceneSpawnerSystem), ); app.add_systems(Update, tag_scene_elements); @@ -261,10 +266,7 @@ fn initialiaze_headless( ..default() }, CameraControllerBundle::new(focus, camera_xform.translation), - ScreenCapture::number_of_frames( - headless_settings.frames_to_capture, - &headless_settings.output_dir, - ), + ScreenCapture::from_out(&headless_settings.output_dir), )); } @@ -429,6 +431,12 @@ fn frame_scene_in_view( cam_xform.translation = cam.focus + cam_xform.back() * radius; } +fn start_capture(mut captures: Query<&mut ScreenCapture>) { + for mut capture in captures.iter_mut() { + capture.start_capturing(); + } +} + #[cfg(feature = "wasm")] mod wasm { use bevy::{prelude::*, window::PrimaryWindow}; diff --git a/studio/src/screenshot.rs b/studio/src/screenshot.rs index 001691c..f773d13 100644 --- a/studio/src/screenshot.rs +++ b/studio/src/screenshot.rs @@ -30,8 +30,7 @@ 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)), + (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 { @@ -65,60 +64,73 @@ impl Plugin for ScreenShotPlugin { } } -fn test_system( - mut commands: Commands, - camera: Query>, - input: Res>, -) { - 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, + frames: u32, + max_frames: u32, to_directory: String, capture_depth: bool, + currently_captureing: bool, } impl ScreenCapture { - pub fn number_of_frames(frames: u32, output_dir: &str) -> Self { + pub fn from_out(output_dir: &str) -> Self { ScreenCapture { - number_of_frames: frames, + frames: 0, to_directory: output_dir.to_string(), capture_depth: false, + currently_captureing: false, + max_frames: u32::MAX, } } + pub fn with_max_frames(mut self, max_frames: u32) -> Self { + self.max_frames = max_frames; + self + } + 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 + pub fn is_capturing(&self) -> bool { + self.currently_captureing + } + + pub fn stop_capturing(&mut self) { + self.currently_captureing = false; + } + + pub fn start_capturing(&mut self) { + self.currently_captureing = true; + } + + pub fn should_capture(&self) -> bool { + self.frames < self.max_frames && self.is_capturing() + } + + pub fn frames_captured(&self) -> u32 { + self.frames + } + + pub fn max_frames(&self) -> u32 { + self.max_frames } } fn update_screen_capture_on_extract_step( - mut commands: Commands, - mut screen_captures: Query<(Entity, &mut ScreenCapture)>, + mut screen_captures: Query<&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::(); - } + for mut screen_capture in screen_captures.iter_mut() { + screen_capture.frames += 1; + // if screen_capture.number_of_frames == 0 { + // commands.entity(entity).remove::(); + // } } } @@ -128,7 +140,9 @@ fn extract_screen_capture_components( screen_capure_comps: Extract>>, ) { for (entity, screen_capture) in screen_capure_comps.iter() { - commands.get_or_spawn(entity).insert(screen_capture.clone()); + if screen_capture.should_capture() { + commands.get_or_spawn(entity).insert(screen_capture.clone()); + } } } @@ -152,7 +166,7 @@ fn prepare_screen_capture_buffers( device: Res, ) { for (entity, view_target, screen_capture) in views.iter() { - if screen_capture.number_of_frames > 0 { + if screen_capture.frames > 0 { let texture = view_target.main_texture(); let size = texture.size(); let format = texture.format(); @@ -174,8 +188,6 @@ fn prepare_screen_capture_buffers( main_bytes_per_row: bytes_per_row, main_padded_bytes_per_row: padded_bytes_per_row, }); - - println!("Buffers Created!"); } } } @@ -219,7 +231,6 @@ impl ViewNode for ScreenCaptureNode { }, view_query.1.main_texture().size(), ); - println!("Saved to buffer!"); Ok(()) } } @@ -255,7 +266,7 @@ fn save_buffer_to_file( let path = format!( "{}img_{:0.5}.png", - screen_capture.to_directory, screen_capture.number_of_frames + screen_capture.to_directory, screen_capture.frames ); save_buffer::>(&main_image_values, &buffers.main_buffer_extent, &path);