glTF exporter: Sampled animation now bake all bones, including constraints

Using "Always Sample Animation" can now export complex rigs, including constraints, like rigify for example

1 files changed, 46 insertions, 28 deletions

diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_sampler_keyframes.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_sampler_keyframes.py
index a6cd0f9e..fd0bb9c5 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_sampler_keyframes.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_sampler_keyframes.py
@@ -25,12 +25,19 @@ from io_scene_gltf2.io.com import gltf2_io_debug
 
 
 class Keyframe:
-    def __init__(self, channels: typing.Tuple[bpy.types.FCurve], frame: float):
+    def __init__(self, channels: typing.Tuple[bpy.types.FCurve], frame: float, bake_channel: typing.Union[str, None]):
         self.seconds = frame / bpy.context.scene.render.fps
         self.frame = frame
         self.fps = bpy.context.scene.render.fps
-        self.target = channels[0].data_path.split('.')[-1]
-        self.__indices = [c.array_index for c in channels]
+        if bake_channel is None:
+            self.target = channels[0].data_path.split('.')[-1]
+            self.__indices = [c.array_index for c in channels]
+        else:
+            self.target = bake_channel
+            self.__indices = []
+            for i in range(self.get_target_len()):
+                self.__indices.append(i)
+
 
         # Data holders for virtual properties
         self.__value = None
@@ -73,11 +80,6 @@ class Keyframe:
     def set_value_index(self, idx, val):
         self.__value[idx] = val
 
-    def set_full_value(self, val):
-        self.__value = [0.0] * self.get_target_len()
-        for i in range(0, self.get_target_len()):
-            self.set_value_index(i, val[i])
-
     @property
     def value(self) -> typing.Union[mathutils.Vector, mathutils.Euler, mathutils.Quaternion, typing.List[float]]:
         return self.__value
@@ -108,47 +110,62 @@ class Keyframe:
 def gather_keyframes(blender_object_if_armature: typing.Optional[bpy.types.Object],
                      channels: typing.Tuple[bpy.types.FCurve],
                      non_keyed_values: typing.Tuple[typing.Optional[float]],
+                     bake_bone: typing.Union[str, None],
+                     bake_channel: typing.Union[str, None],
+                     bake_range_start,
+                     bake_range_end,
                      export_settings
                      ) -> typing.List[Keyframe]:
     """Convert the blender action groups' fcurves to keyframes for use in glTF."""
-    # Find the start and end of the whole action group
-    ranges = [channel.range() for channel in channels]
-
-    start_frame = min([channel.range()[0] for channel in channels])
-    end_frame = max([channel.range()[1] for channel in channels])
+    if bake_bone is None:
+        # Find the start and end of the whole action group
+        ranges = [channel.range() for channel in channels]
 
-    keyframes = []
-
-    if blender_object_if_armature is not None:
-        pose_bone_if_armature = gltf2_blender_get.get_object_from_datapath(blender_object_if_armature,
-                                                                           channels[0].data_path)
+        start_frame = min([channel.range()[0] for channel in channels])
+        end_frame = max([channel.range()[1] for channel in channels])
     else:
-        pose_bone_if_armature = None
+        start_frame = bake_range_start
+        end_frame = bake_range_end
 
+    keyframes = []
     if needs_baking(blender_object_if_armature, channels, export_settings):
         # Bake the animation, by evaluating the animation for all frames
         # TODO: maybe baking can also be done with FCurve.convert_to_samples
 
+        if blender_object_if_armature is not None:
+            if bake_bone is None:
+                pose_bone_if_armature = gltf2_blender_get.get_object_from_datapath(blender_object_if_armature,
+                                                                               channels[0].data_path)
+            else:
+                pose_bone_if_armature = blender_object_if_armature.pose.bones[bake_bone]
+        else:
+            pose_bone_if_armature = None
+
         # sample all frames
         frame = start_frame
         step = export_settings['gltf_frame_step']
         while frame <= end_frame:
-            key = Keyframe(channels, frame)
+            key = Keyframe(channels, frame, bake_channel)
             if isinstance(pose_bone_if_armature, bpy.types.PoseBone):
                 # we need to bake in the constraints
                 bpy.context.scene.frame_set(frame)
-                # TODO, this is not working if the action is not active (NLA case for example)
-                trans, rot, scale = pose_bone_if_armature.matrix_basis.decompose()
-                target_property = channels[0].data_path.split('.')[-1]
-                # Store all values, not only the data from fcurve:
-                # All indices must be stored
-                key.set_full_value({
+                if bake_bone is None:
+                    trans, rot, scale = pose_bone_if_armature.matrix_basis.decompose()
+                else:
+                    matrix = pose_bone_if_armature.matrix
+                    new_matrix = blender_object_if_armature.convert_space(pose_bone=pose_bone_if_armature, matrix=matrix, from_space='POSE', to_space='LOCAL')
+                    trans, rot, scale = new_matrix.decompose()
+                if bake_channel is None:
+                    target_property = channels[0].data_path.split('.')[-1]
+                else:
+                    target_property = bake_channel
+                key.value = {
                     "location": trans,
                     "rotation_axis_angle": rot,
                     "rotation_euler": rot,
                     "rotation_quaternion": rot,
                     "scale": scale
-                }[target_property])
+                }[target_property]
             else:
                 key.value = [c.evaluate(frame) for c in channels]
                 complete_key(key, non_keyed_values)
@@ -158,7 +175,8 @@ def gather_keyframes(blender_object_if_armature: typing.Optional[bpy.types.Objec
         # Just use the keyframes as they are specified in blender
         frames = [keyframe.co[0] for keyframe in channels[0].keyframe_points]
         for i, frame in enumerate(frames):
-            key = Keyframe(channels, frame)
+            key = Keyframe(channels, frame, bake_channel)
+            # key.value = [c.keyframe_points[i].co[0] for c in action_group.channels]
             key.value = [c.evaluate(frame) for c in channels]
             # Complete key with non keyed values, if needed
             if len(channels) != key.get_target_len():