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

4 files changed, 213 insertions, 74 deletions

diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channel_target.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channel_target.py
index 2e4ac1d7..5b2f28e8 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channel_target.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channel_target.py
@@ -24,37 +24,47 @@ from io_scene_gltf2.blender.exp import gltf2_blender_gather_joints
 @cached
 def gather_animation_channel_target(channels: typing.Tuple[bpy.types.FCurve],
                                     blender_object: bpy.types.Object,
+                                    bake_bone: typing.Union[str, None],
+                                    bake_channel: typing.Union[str, None],
                                     export_settings
                                     ) -> gltf2_io.AnimationChannelTarget:
-    return gltf2_io.AnimationChannelTarget(
-        extensions=__gather_extensions(channels, blender_object, export_settings),
-        extras=__gather_extras(channels, blender_object, export_settings),
-        node=__gather_node(channels, blender_object, export_settings),
-        path=__gather_path(channels, blender_object, export_settings)
-    )
 
+        return gltf2_io.AnimationChannelTarget(
+            extensions=__gather_extensions(channels, blender_object, export_settings, bake_bone),
+            extras=__gather_extras(channels, blender_object, export_settings, bake_bone),
+            node=__gather_node(channels, blender_object, export_settings, bake_bone),
+            path=__gather_path(channels, blender_object, export_settings, bake_bone, bake_channel)
+        )
 
 def __gather_extensions(channels: typing.Tuple[bpy.types.FCurve],
                         blender_object: bpy.types.Object,
-                        export_settings
+                        export_settings,
+                        bake_bone: typing.Union[str, None]
                         ) -> typing.Any:
     return None
 
 
 def __gather_extras(channels: typing.Tuple[bpy.types.FCurve],
                     blender_object: bpy.types.Object,
-                    export_settings
+                    export_settings,
+                    bake_bone: typing.Union[str, None]
                     ) -> typing.Any:
     return None
 
 
 def __gather_node(channels: typing.Tuple[bpy.types.FCurve],
                   blender_object: bpy.types.Object,
-                  export_settings
+                  export_settings,
+                  bake_bone: typing.Union[str, None]
                   ) -> gltf2_io.Node:
     if blender_object.type == "ARMATURE":
         # TODO: get joint from fcurve data_path and gather_joint
-        blender_bone = blender_object.path_resolve(channels[0].data_path.rsplit('.', 1)[0])
+
+        if bake_bone is not None:
+            blender_bone = blender_object.pose.bones[bake_bone]
+        else:
+            blender_bone = blender_object.path_resolve(channels[0].data_path.rsplit('.', 1)[0])
+
         if isinstance(blender_bone, bpy.types.PoseBone):
             return gltf2_blender_gather_joints.gather_joint(blender_bone, export_settings)
 
@@ -63,9 +73,14 @@ def __gather_node(channels: typing.Tuple[bpy.types.FCurve],
 
 def __gather_path(channels: typing.Tuple[bpy.types.FCurve],
                   blender_object: bpy.types.Object,
-                  export_settings
+                  export_settings,
+                  bake_bone: typing.Union[str, None],
+                  bake_channel: typing.Union[str, None]
                   ) -> str:
-    target = channels[0].data_path.split('.')[-1]
+    if bake_channel is None:
+        target = channels[0].data_path.split('.')[-1]
+    else:
+        target = bake_channel
     path = {
         "delta_location": "translation",
         "delta_rotation_euler": "rotation",
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channels.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channels.py
index 1b76142c..aa81b073 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channels.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_channels.py
@@ -31,26 +31,61 @@ def gather_animation_channels(blender_action: bpy.types.Action,
                               ) -> typing.List[gltf2_io.AnimationChannel]:
     channels = []
 
-    for channel_group in __get_channel_groups(blender_action, blender_object):
-        channel = __gather_animation_channel(channel_group, blender_object, export_settings)
-        if channel is not None:
-            channels.append(channel)
+    if blender_object.type == "ARMATURE" and export_settings['gltf_force_sampling'] is True:
+        # We have to store sampled animation data for every deformation bones
+
+        # First calculate range of animation for baking
+        bake_range_start = None
+        bake_range_end = None
+        groups = __get_channel_groups(blender_action, blender_object, export_settings)
+        for chans in groups:
+            ranges = [channel.range() for channel in chans]
+            if bake_range_start is None:
+                bake_range_start = min([channel.range()[0] for channel in chans])
+            else:
+                bake_range_start = min(bake_range_start, min([channel.range()[0] for channel in chans]))
+            if bake_range_end is None:
+                bake_range_end = max([channel.range()[1] for channel in chans])
+            else:
+                bake_range_end = max(bake_range_end, max([channel.range()[1] for channel in chans]))
+
+        # Then bake all bones
+        for bone in blender_object.data.bones:
+            for p in ["location", "rotation_quaternion", "scale"]:
+                channel = __gather_animation_channel(
+                    (),
+                    blender_object,
+                    export_settings,
+                    bone.name,
+                    p,
+                    bake_range_start,
+                    bake_range_end)
+                channels.append(channel)
+    else:
+        for channel_group in __get_channel_groups(blender_action, blender_object, export_settings):
+            channel = __gather_animation_channel(channel_group, blender_object, export_settings, None, None, None, None)
+            if channel is not None:
+                channels.append(channel)
 
     return channels
 
 
 def __gather_animation_channel(channels: typing.Tuple[bpy.types.FCurve],
                                blender_object: bpy.types.Object,
-                               export_settings
+                               export_settings,
+                               bake_bone: typing.Union[str, None],
+                               bake_channel: typing.Union[str, None],
+                               bake_range_start,
+                               bake_range_end
                                ) -> typing.Union[gltf2_io.AnimationChannel, None]:
     if not __filter_animation_channel(channels, blender_object, export_settings):
         return None
 
     return gltf2_io.AnimationChannel(
-        extensions=__gather_extensions(channels, blender_object, export_settings),
-        extras=__gather_extras(channels, blender_object, export_settings),
-        sampler=__gather_sampler(channels, blender_object, export_settings),
-        target=__gather_target(channels, blender_object, export_settings)
+        extensions=__gather_extensions(channels, blender_object, export_settings, bake_bone),
+        extras=__gather_extras(channels, blender_object, export_settings, bake_bone),
+        sampler=__gather_sampler(channels, blender_object, export_settings, bake_bone, bake_channel, bake_range_start, bake_range_end),
+        target=__gather_target(channels, blender_object, export_settings, bake_bone, bake_channel)
     )
 
 
@@ -63,38 +98,50 @@ def __filter_animation_channel(channels: typing.Tuple[bpy.types.FCurve],
 
 def __gather_extensions(channels: typing.Tuple[bpy.types.FCurve],
                         blender_object: bpy.types.Object,
-                        export_settings
+                        export_settings,
+                        bake_bone: typing.Union[str, None]
                         ) -> typing.Any:
     return None
 
 
 def __gather_extras(channels: typing.Tuple[bpy.types.FCurve],
                     blender_object: bpy.types.Object,
-                    export_settings
+                    export_settings,
+                    bake_bone: typing.Union[str, None]
                     ) -> typing.Any:
     return None
 
 
 def __gather_sampler(channels: typing.Tuple[bpy.types.FCurve],
                      blender_object: bpy.types.Object,
-                     export_settings
+                     export_settings,
+                     bake_bone: typing.Union[str, None],
+                     bake_channel: typing.Union[str, None],
+                     bake_range_start,
+                     bake_range_end
                      ) -> gltf2_io.AnimationSampler:
     return gltf2_blender_gather_animation_samplers.gather_animation_sampler(
         channels,
         blender_object,
+        bake_bone,
+        bake_channel,
+        bake_range_start,
+        bake_range_end,
         export_settings
     )
 
 
 def __gather_target(channels: typing.Tuple[bpy.types.FCurve],
                     blender_object: bpy.types.Object,
-                    export_settings
+                    export_settings,
+                    bake_bone: typing.Union[str, None],
+                    bake_channel: typing.Union[str, None]
                     ) -> gltf2_io.AnimationChannelTarget:
     return gltf2_blender_gather_animation_channel_target.gather_animation_channel_target(
-        channels, blender_object, export_settings)
+        channels, blender_object, bake_bone, bake_channel, export_settings)
 
 
-def __get_channel_groups(blender_action: bpy.types.Action, blender_object: bpy.types.Object):
+def __get_channel_groups(blender_action: bpy.types.Action, blender_object: bpy.types.Object, export_settings):
     targets = {}
     for fcurve in blender_action.fcurves:
         target_property = get_target_property_name(fcurve.data_path)
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():
diff --git a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_samplers.py b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_samplers.py
index cdd4be7b..412db171 100755
--- a/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_samplers.py
+++ b/io_scene_gltf2/blender/exp/gltf2_blender_gather_animation_samplers.py
@@ -32,28 +32,40 @@ from . import gltf2_blender_export_keys
 @cached
 def gather_animation_sampler(channels: typing.Tuple[bpy.types.FCurve],
                              blender_object: bpy.types.Object,
+                             bake_bone: typing.Union[str, None],
+                             bake_channel: typing.Union[str, None],
+                             bake_range_start,
+                             bake_range_end,
                              export_settings
                              ) -> gltf2_io.AnimationSampler:
 
     blender_object_if_armature = blender_object if blender_object.type == "ARMATURE" else None
     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)
+        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.pose.bones[bake_bone]
     else:
         pose_bone_if_armature = None
     non_keyed_values = __gather_non_keyed_values(channels, blender_object,
                                                  blender_object_if_armature, pose_bone_if_armature,
+                                                 bake_channel,
                                                  export_settings)
 
 
     return gltf2_io.AnimationSampler(
-        extensions=__gather_extensions(channels, blender_object_if_armature, export_settings),
-        extras=__gather_extras(channels, blender_object_if_armature, export_settings),
-        input=__gather_input(channels, blender_object_if_armature, non_keyed_values, export_settings),
-        interpolation=__gather_interpolation(channels, blender_object_if_armature, export_settings),
+        extensions=__gather_extensions(channels, blender_object_if_armature, export_settings, bake_bone, bake_channel),
+        extras=__gather_extras(channels, blender_object_if_armature, export_settings, bake_bone, bake_channel),
+        input=__gather_input(channels, blender_object_if_armature, non_keyed_values, bake_bone, bake_channel, bake_range_start, bake_range_end, export_settings),
+        interpolation=__gather_interpolation(channels, blender_object_if_armature, export_settings, bake_bone, bake_channel),
         output=__gather_output(channels, blender_object.matrix_parent_inverse.copy().freeze(),
                                blender_object_if_armature,
                                non_keyed_values,
+                               bake_bone,
+                               bake_channel,
+                               bake_range_start,
+                               bake_range_end,
                                export_settings)
     )
 
@@ -61,12 +73,16 @@ def __gather_non_keyed_values(channels: typing.Tuple[bpy.types.FCurve],
                               blender_object: bpy.types.Object,
                               blender_object_if_armature: typing.Optional[bpy.types.Object],
                               pose_bone_if_armature: typing.Optional[bpy.types.PoseBone],
+                              bake_channel: typing.Union[str, None],
                               export_settings
                               ) ->  typing.Tuple[typing.Optional[float]]:
 
     non_keyed_values = []
 
-    target = channels[0].data_path.split('.')[-1]
+    if bake_channel is None:
+        target = channels[0].data_path.split('.')[-1]
+    else:
+        target = bake_channel
     if target == "value":
         return ()
 
@@ -84,7 +100,17 @@ def __gather_non_keyed_values(channels: typing.Tuple[bpy.types.FCurve],
     }.get(target)
 
     for i in range(0, length):
-        if i in indices:
+        if bake_channel is not None:
+            non_keyed_values.append({
+                "delta_location" : blender_object.delta_location,
+                "delta_rotation_euler" : blender_object.delta_rotation_euler,
+                "location" : blender_object.location,
+                "rotation_axis_angle" : blender_object.rotation_axis_angle,
+                "rotation_euler" : blender_object.rotation_euler,
+                "rotation_quaternion" : blender_object.rotation_quaternion,
+                "scale" : blender_object.scale
+            }[target][i])
+        elif i in indices:
             non_keyed_values.append(None)
         else:
             if blender_object_if_armature is None:
@@ -112,14 +138,18 @@ def __gather_non_keyed_values(channels: typing.Tuple[bpy.types.FCurve],
 
 def __gather_extensions(channels: typing.Tuple[bpy.types.FCurve],
                         blender_object_if_armature: typing.Optional[bpy.types.Object],
-                        export_settings
+                        export_settings,
+                        bake_bone: typing.Union[str, None],
+                        bake_channel: typing.Union[str, None]
                         ) -> typing.Any:
     return None
 
 
 def __gather_extras(channels: typing.Tuple[bpy.types.FCurve],
                     blender_object_if_armature: typing.Optional[bpy.types.Object],
-                    export_settings
+                    export_settings,
+                    bake_bone: typing.Union[str, None],
+                    bake_channel: typing.Union[str, None]
                     ) -> typing.Any:
     return None
 
@@ -128,12 +158,20 @@ def __gather_extras(channels: typing.Tuple[bpy.types.FCurve],
 def __gather_input(channels: typing.Tuple[bpy.types.FCurve],
                    blender_object_if_armature: typing.Optional[bpy.types.Object],
                    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
                    ) -> gltf2_io.Accessor:
     """Gather the key time codes."""
     keyframes = gltf2_blender_gather_animation_sampler_keyframes.gather_keyframes(blender_object_if_armature,
                                                                                   channels,
                                                                                   non_keyed_values,
+                                                                                  bake_bone,
+                                                                                  bake_channel,
+                                                                                  bake_range_start,
+                                                                                  bake_range_end,
                                                                                   export_settings)
     times = [k.seconds for k in keyframes]
 
@@ -150,14 +188,19 @@ def __gather_input(channels: typing.Tuple[bpy.types.FCurve],
 
 def __gather_interpolation(channels: typing.Tuple[bpy.types.FCurve],
                            blender_object_if_armature: typing.Optional[bpy.types.Object],
-                           export_settings
+                           export_settings,
+                           bake_bone: typing.Union[str, None],
+                           bake_channel: typing.Union[str, None]
                            ) -> str:
     if gltf2_blender_gather_animation_sampler_keyframes.needs_baking(blender_object_if_armature,
                                                                      channels,
                                                                      export_settings):
-        max_keyframes = max([len(ch.keyframe_points) for ch in channels])
-        # If only single keyframe revert to STEP
-        return 'STEP' if max_keyframes < 2 else 'LINEAR'
+        if bake_bone is not None:
+            return 'LINEAR'
+        else:
+            max_keyframes = max([len(ch.keyframe_points) for ch in channels])
+            # If only single keyframe revert to STEP
+            return 'STEP' if max_keyframes < 2 else 'LINEAR'
 
     blender_keyframe = channels[0].keyframe_points[0]
 
@@ -174,24 +217,40 @@ def __gather_output(channels: typing.Tuple[bpy.types.FCurve],
                     parent_inverse,
                     blender_object_if_armature: typing.Optional[bpy.types.Object],
                     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
                     ) -> gltf2_io.Accessor:
     """Gather the data of the keyframes."""
     keyframes = gltf2_blender_gather_animation_sampler_keyframes.gather_keyframes(blender_object_if_armature,
                                                                                   channels,
                                                                                   non_keyed_values,
+                                                                                  bake_bone,
+                                                                                  bake_channel,
+                                                                                  bake_range_start,
+                                                                                  bake_range_end,
                                                                                   export_settings)
-
-    target_datapath = channels[0].data_path
+    if bake_bone is not None:
+        target_datapath = "pose.bones['" + bake_bone + "']." + bake_channel
+    else:
+        target_datapath = channels[0].data_path
 
     is_yup = export_settings[gltf2_blender_export_keys.YUP]
 
     # bone animations need to be handled differently as they are in a different coordinate system
-    object_path = get_target_object_path(target_datapath)
-    is_armature_animation = blender_object_if_armature is not None and object_path != ""
+    if bake_bone is None:
+        object_path = get_target_object_path(target_datapath)
+    else:
+        object_path = None
+    is_armature_animation = bake_bone is not None or (blender_object_if_armature is not None and object_path != "")
 
     if is_armature_animation:
-        bone = gltf2_blender_get.get_object_from_datapath(blender_object_if_armature, object_path)
+        if bake_bone is None:
+            bone = gltf2_blender_get.get_object_from_datapath(blender_object_if_armature, object_path)
+        else:
+            bone = blender_object_if_armature.pose.bones[bake_bone]
         if isinstance(bone, bpy.types.PoseBone):
             if bone.parent is None:
                 axis_basis_change = mathutils.Matrix.Identity(4)