Adding Mixamo Animations to Studio Headless (#50)

This PR adds the ability to render out mixamo animations from the
headless mode. In addition to that, there were some bugs that have been
managed and some readability changes. Let me know if you have any
questions!

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<MixamoImportEvent>,
+                      mut load_scene_event: EventWriter<LoadSceneEvent>| {
+                    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<StudioMode>) -> 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<AppExit>,
-    main_camera: Query<(Entity, &ScreenCapture), With<MainCamera>>,
+    mut main_camera: Query<&mut ScreenCapture, With<MainCamera>>,
     bvh_event_reader: EventReader<BvhAnimCompleteEvent>,
 ) {
-    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::<ScreenCapture>();
+    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<AppExit>,
+    mut main_camera: Query<&mut ScreenCapture, With<MainCamera>>,
+    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);
+        }
+    }
+}

studio/src/main-headless.rs

@@ -43,6 +43,19 @@ enum Subcommands {
         #[arg(long, value_parser = clap::value_parser!(u32).range(1..))]
         frames: Option<u32>,
     },
+    Mixamo {
+        mixamo_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)]
@@ -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();

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};

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<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,
+    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::<ScreenCapture>();
-        }
+    for mut screen_capture in screen_captures.iter_mut() {
+        screen_capture.frames += 1;
+        // if screen_capture.number_of_frames == 0 {
+        //     commands.entity(entity).remove::<ScreenCapture>();
+        // }
     }
 }
 
@@ -128,7 +140,9 @@ fn extract_screen_capture_components(
     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());
+        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<RenderDevice>,
 ) {
     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::<Rgba<u8>>(&main_image_values, &buffers.main_buffer_extent, &path);