glTF importer: refactoring animation import

1 files changed, 130 insertions, 77 deletions

diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_animation_node.py b/io_scene_gltf2/blender/imp/gltf2_blender_animation_node.py
index b6369b8b..d41b7892 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_animation_node.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_animation_node.py
@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import bpy
 from mathutils import Vector
 
 from ...io.imp.gltf2_io_binary import BinaryData
-from .gltf2_blender_animation_utils import simulate_stash, make_fcurve
+from .gltf2_blender_animation_utils import make_fcurve
 from .gltf2_blender_vnode import VNode
 
 
@@ -27,92 +28,144 @@ class BlenderNodeAnim():
 
     @staticmethod
     def anim(gltf, anim_idx, node_idx):
-        """Manage animation."""
-        node = gltf.data.nodes[node_idx]
-        vnode = gltf.vnodes[node_idx]
-        obj = vnode.blender_object
-        fps = bpy.context.scene.render.fps
-
+        """Manage animation targeting a node's TRS."""
         animation = gltf.data.animations[anim_idx]
-
+        node = gltf.data.nodes[node_idx]
         if anim_idx not in node.animations.keys():
             return
 
         for channel_idx in node.animations[anim_idx]:
             channel = animation.channels[channel_idx]
-            if channel.target.path in ['translation', 'rotation', 'scale']:
-                break
-        else:
-            return
+            if channel.target.path not in ['translation', 'rotation', 'scale']:
+                continue
+
+            BlenderNodeAnim.do_channel(gltf, anim_idx, node_idx, channel)
+
+    @staticmethod
+    def do_channel(gltf, anim_idx, node_idx, channel):
+        animation = gltf.data.animations[anim_idx]
+        vnode = gltf.vnodes[node_idx]
+        path = channel.target.path
+
+        action = BlenderNodeAnim.get_or_create_action(gltf, node_idx, animation.track_name)
+
+        keys = BinaryData.get_data_from_accessor(gltf, animation.samplers[channel.sampler].input)
+        values = BinaryData.get_data_from_accessor(gltf, animation.samplers[channel.sampler].output)
+
+        if animation.samplers[channel.sampler].interpolation == "CUBICSPLINE":
+            # TODO manage tangent?
+            values = values[1::3]
+
+        # Convert the curve from glTF to Blender.
+
+        if path == "translation":
+            blender_path = "location"
+            group_name = "Location"
+            num_components = 3
+            values = [gltf.loc_gltf_to_blender(vals) for vals in values]
+            values = vnode.base_locs_to_final_locs(values)
+
+        elif path == "rotation":
+            blender_path = "rotation_quaternion"
+            group_name = "Rotation"
+            num_components = 4
+            values = [gltf.quaternion_gltf_to_blender(vals) for vals in values]
+            values = vnode.base_rots_to_final_rots(values)
+
+        elif path == "scale":
+            blender_path = "scale"
+            group_name = "Scale"
+            num_components = 3
+            values = [gltf.scale_gltf_to_blender(vals) for vals in values]
+            values = vnode.base_scales_to_final_scales(values)
+
+        # Objects parented to a bone are translated to the bone tip by default.
+        # Correct for this by translating backwards from the tip to the root.
+        if vnode.type == VNode.Object and path == "translation":
+            if vnode.parent is not None and gltf.vnodes[vnode.parent].type == VNode.Bone:
+                bone_length = gltf.vnodes[vnode.parent].bone_length
+                off = Vector((0, -bone_length, 0))
+                values = [vals + off for vals in values]
+
+        if vnode.type == VNode.Bone:
+            # Need to animate the pose bone when the node is a bone.
+            group_name = vnode.blender_bone_name
+            blender_path = "pose.bones[%s].%s" % (json.dumps(vnode.blender_bone_name), blender_path)
+
+            # We have the final TRS of the bone in values. We need to give
+            # the TRS of the pose bone though, which is relative to the edit
+            # bone.
+            #
+            #     Final = EditBone * PoseBone
+            #   where
+            #     Final =    Trans[ft] Rot[fr] Scale[fs]
+            #     EditBone = Trans[et] Rot[er]
+            #     PoseBone = Trans[pt] Rot[pr] Scale[ps]
+            #
+            # Solving for PoseBone gives
+            #
+            #     pt = Rot[er^{-1}] (ft - et)
+            #     pr = er^{-1} fr
+            #     ps = fs
+
+            if path == 'translation':
+                edit_trans, edit_rot = vnode.editbone_trans, vnode.editbone_rot
+                edit_rot_inv = edit_rot.conjugated()
+                values = [
+                    edit_rot_inv @ (trans - edit_trans)
+                    for trans in values
+                ]
+
+            elif path == 'rotation':
+                edit_rot = vnode.editbone_rot
+                edit_rot_inv = edit_rot.conjugated()
+                values = [
+                    edit_rot_inv @ rot
+                    for rot in values
+                ]
+
+            elif path == 'scale':
+                pass  # no change needed
+
+        # To ensure rotations always take the shortest path, we flip
+        # adjacent antipodal quaternions.
+        if path == 'rotation':
+            for i in range(1, len(values)):
+                if values[i].dot(values[i-1]) < 0:
+                    values[i] = -values[i]
+
+        fps = bpy.context.scene.render.fps
+
+        coords = [0] * (2 * len(keys))
+        coords[::2] = (key[0] * fps for key in keys)
+
+        for i in range(0, num_components):
+            coords[1::2] = (vals[i] for vals in values)
+            make_fcurve(
+                action,
+                coords,
+                data_path=blender_path,
+                index=i,
+                group_name=group_name,
+                interpolation=animation.samplers[channel.sampler].interpolation,
+            )
+
+    @staticmethod
+    def get_or_create_action(gltf, node_idx, anim_name):
+        vnode = gltf.vnodes[node_idx]
+
+        if vnode.type == VNode.Bone:
+            # For bones, the action goes on the armature.
+            vnode = gltf.vnodes[vnode.bone_arma]
+
+        obj = vnode.blender_object
 
         action = gltf.action_cache.get(obj.name)
         if not action:
-            name = animation.track_name + "_" + obj.name
+            name = anim_name + "_" + obj.name
             action = bpy.data.actions.new(name)
             action.id_root = 'OBJECT'
-            gltf.needs_stash.append((obj, animation.track_name, action))
+            gltf.needs_stash.append((obj, action))
             gltf.action_cache[obj.name] = action
 
-        if not obj.animation_data:
-            obj.animation_data_create()
-        obj.animation_data.action = action
-
-        for channel_idx in node.animations[anim_idx]:
-            channel = animation.channels[channel_idx]
-
-            keys = BinaryData.get_data_from_accessor(gltf, animation.samplers[channel.sampler].input)
-            values = BinaryData.get_data_from_accessor(gltf, animation.samplers[channel.sampler].output)
-
-            if channel.target.path not in ['translation', 'rotation', 'scale']:
-                continue
-
-            if animation.samplers[channel.sampler].interpolation == "CUBICSPLINE":
-                # TODO manage tangent?
-                values = [values[idx * 3 + 1] for idx in range(0, len(keys))]
-
-            if channel.target.path == "translation":
-                blender_path = "location"
-                group_name = "Location"
-                num_components = 3
-                values = [gltf.loc_gltf_to_blender(vals) for vals in values]
-                values = vnode.base_locs_to_final_locs(values)
-
-                if vnode.parent is not None and gltf.vnodes[vnode.parent].type == VNode.Bone:
-                    # Nodes with a bone parent need to be translated
-                    # backwards from the tip to the root
-                    bone_length = gltf.vnodes[vnode.parent].bone_length
-                    off = Vector((0, -bone_length, 0))
-                    values = [vals + off for vals in values]
-
-            elif channel.target.path == "rotation":
-                blender_path = "rotation_quaternion"
-                group_name = "Rotation"
-                num_components = 4
-                values = [gltf.quaternion_gltf_to_blender(vals) for vals in values]
-                values = vnode.base_rots_to_final_rots(values)
-
-                # Manage antipodal quaternions
-                for i in range(1, len(values)):
-                    if values[i].dot(values[i-1]) < 0:
-                        values[i] = -values[i]
-
-            elif channel.target.path == "scale":
-                blender_path = "scale"
-                group_name = "Scale"
-                num_components = 3
-                values = [gltf.scale_gltf_to_blender(vals) for vals in values]
-                values = vnode.base_scales_to_final_scales(values)
-
-            coords = [0] * (2 * len(keys))
-            coords[::2] = (key[0] * fps for key in keys)
-
-            for i in range(0, num_components):
-                coords[1::2] = (vals[i] for vals in values)
-                make_fcurve(
-                    action,
-                    coords,
-                    data_path=blender_path,
-                    index=i,
-                    group_name=group_name,
-                    interpolation=animation.samplers[channel.sampler].interpolation,
-                )
+        return action