feat: added extra information to output extracted keyframes

assets/newton/animations/1242921_newton.keyframes.yaml → assets/newton/animations/idle.keyframes.yaml

@@ -1,3 +1,6 @@
+animation: idle
+beginning: 0
+duration: 45
 keyframes:
 - data:
   - bone: FR_HAA

core/animation/animation_engine.py

@@ -1,6 +1,3 @@
 class AnimationEngine:
     def __init__(self):
         pass
-
-    def __del__(self):
-        pass

core/animation/blender/__init__.py

@@ -0,0 +1 @@
+# This module is mainly to keep all blender-related files organized, but nothing should be imported from it.

core/animation/blender/keyframe_extractor.py

@@ -51,11 +51,9 @@ def extract_keyframes_from_scene(
                 bone = armature.pose.bones.get(bone_name)
                 if bone:
                     world_matrix = armature.matrix_world @ bone.matrix
-                    rotation_quaternion = world_matrix.to_quaternion()
 
                     rotation_keyframe = {
                         "bone": bone_name,
-                        "quaternion": list(rotation_quaternion[:]),
                     }
 
                     if euler:
@@ -65,6 +63,10 @@ def extract_keyframes_from_scene(
                         rotation_euler[:] = [e * 180 / 3.14159 for e in rotation_euler]
 
                         rotation_keyframe["euler"] = list(rotation_euler[:])
+                    else:
+                        rotation_quaternion = world_matrix.to_quaternion()
+
+                        rotation_keyframe["quaternion"] = list(rotation_quaternion[:])
 
                     frame_data.append(rotation_keyframe)
 
@@ -77,12 +79,20 @@ def extract_keyframes_from_scene(
     else:
         print(f"Armature {armature_name} not found, or not an armature.")
 
-    return keyframes_data, end_frame - start_frame, len(bone_names)
+    return keyframes_data, start_frame, end_frame - start_frame, len(bone_names)
 
 
-def list_to_yaml(data: list) -> str:
+def keyframes_to_yaml(
+    data: list,
+    beginning: int,
+    duration: int,
+    animation: str,
+) -> str:
     return yaml.dump(
         {
+            "animation": animation,
+            "beginning": beginning,
+            "duration": duration,
             "keyframes": data,
         }
     )
@@ -97,13 +107,27 @@ def setup_args() -> argparse.Namespace:
         "--armature",
         required=False,
         default="Armature",
+        type=str,
         help="Name of the Armature to extract keyframes from.",
     )
     parser.add_argument(
-        "--keyframes",
+        "--output",
         required=False,
+        type=str,
         help="Path to the file to save the keyframes to.",
     )
+    parser.add_argument(
+        "--animation",
+        required=True,
+        type=str,
+        help="Name of the animation (to be used in the output YAML).",
+    )
+    parser.add_argument(
+        "--euler",
+        action="store_true",
+        default=False,
+        help="Use Euler angles (degrees) instead of Quaternions (WXYZ).",
+    )
 
     args = parser.parse_args(sys.argv[sys.argv.index("--") + 1 :])
 
@@ -113,23 +137,26 @@ def setup_args() -> argparse.Namespace:
 def main():
     args = setup_args()
     armature_name = args.armature
-    keyframes_path = args.keyframes
+    output_path = args.output
+    animation_name = args.animation
 
     # Extract keyframes
-    rotations_dict, length, num_bones = extract_keyframes_from_scene(armature_name)
+    rotations_dict, start_frame, length, num_bones = extract_keyframes_from_scene(
+        armature_name
+    )
 
     print(
         f"Extracted rotation keyframes:\n  Bones: {num_bones}\n  Frames: {length}\n  Total Keyframes: {length * num_bones}"
     )
 
     # Convert to YAML
-    yaml_data = list_to_yaml(rotations_dict)
+    yaml_data = keyframes_to_yaml(rotations_dict, start_frame, length, animation_name)
 
     # Save to file
-    with open(keyframes_path, "w") as file:
+    with open(output_path, "w") as file:
         file.write(yaml_data)
 
-    print(f"Keyframes saved to: {keyframes_path}")
+    print(f"Keyframes saved to: {output_path}")
 
 
 # Ensure the script only runs when executed directly

newton.py

@@ -52,6 +52,12 @@ def setup_argparser() -> argparse.ArgumentParser:
         help="Enable domain randomization.",
         default="configs/randomization.yaml",
     )
+    parser.add_argument(
+        "--animations-dir",
+        type=str,
+        help="Path to the directory containing animations.",
+        default="assets/newton/animations",
+    )
     parser.add_argument(
         "--checkpoint",
         type=str,
@@ -88,6 +94,8 @@ def setup_argparser() -> argparse.ArgumentParser:
     world_config = load_config(cli_args.world_config)
     randomization_config = load_config(cli_args.randomization_config)
 
+    animation_dir = cli_args.animations_dir
+
     physics_only = cli_args.physics
     exporting = cli_args.export_onnx
     pidding = not exporting and cli_args.pid_control

scripts/keyframe_extractor.sh

@@ -60,4 +60,5 @@ echo "Extracting keyframes from $BLEND_FILE..."
 BLEND_FILE_NAME=$(basename "$BLEND_FILE" | cut -d "." -f 1)
 
 OUTPUT_FILE="$OUTPUT_DIR/$BLEND_FILE_NAME.keyframes.yaml"
-exec $BLENDER_EXECUTABLE --python-use-system-env --background --quiet "$BLEND_FILE" --python $KEYFRAME_EXTRACTOR_SCRIPT -- --armature "Armature" --keyframes "$OUTPUT_FILE"
+exec $BLENDER_EXECUTABLE --python-use-system-env --background --quiet "$BLEND_FILE" --python $KEYFRAME_EXTRACTOR_SCRIPT \
+  -- --armature "Armature" --output "$OUTPUT_FILE" --animation "$BLEND_FILE_NAME"