glTF importer: prettify bones: manage more readable bone rotation

Thanks to scurest!

6 files changed, 226 insertions, 49 deletions

diff --git a/io_scene_gltf2/__init__.py b/io_scene_gltf2/__init__.py
index 95e24933..f9ad9b8d 100755
--- a/io_scene_gltf2/__init__.py
+++ b/io_scene_gltf2/__init__.py
@@ -15,7 +15,7 @@
 bl_info = {
     'name': 'glTF 2.0 format',
     'author': 'Julien Duroure, Norbert Nopper, Urs Hanselmann, Moritz Becher, Benjamin Schmithüsen, Jim Eckerlein, and many external contributors',
-    "version": (1, 2, 36),
+    "version": (1, 2, 37),
     'blender': (2, 82, 7),
     'location': 'File > Import-Export',
     'description': 'Import-Export as glTF 2.0',
@@ -856,11 +856,26 @@ class ImportGLTF2(Operator, ImportHelper):
         description="How normals are computed during import",
         default="NORMALS")
 
+    bone_heuristic: EnumProperty(
+        name="Bone Dir",
+        items=(
+            ("BLENDER", "Blender (+Y)",
+                "Round-trips bone directions in glTFs exported from Blender.\n"
+                "Bone tips are placed on their local +Y axis (in glTF space)"),
+            ("TEMPERANCE", "Temperance",
+                "Okay for many different models.\n"
+                "Bone tips are placed at a child's root")
+        ),
+        description="Heuristic for placing bones. Tries to make bones pretty",
+        default="TEMPERANCE",
+    )
+
     def draw(self, context):
         layout = self.layout
 
         layout.prop(self, 'import_pack_images')
         layout.prop(self, 'import_shading')
+        layout.prop(self, 'bone_heuristic')
 
     def execute(self, context):
         return self.import_gltf2(context)
diff --git a/io_scene_gltf2/blender/com/gltf2_blender_math.py b/io_scene_gltf2/blender/com/gltf2_blender_math.py
index fb342bc4..b3f13f65 100755
--- a/io_scene_gltf2/blender/com/gltf2_blender_math.py
+++ b/io_scene_gltf2/blender/com/gltf2_blender_math.py
@@ -170,3 +170,45 @@ def transform_value(value: Vector, _: Matrix = Matrix.Identity(4)) -> Vector:
 def round_if_near(value: float, target: float) -> float:
     """If value is very close to target, round to target."""
     return value if abs(value - target) > 2.0e-6 else target
+
+def scale_rot_swap_matrix(rot):
+    """Returns a matrix m st. Scale[s] Rot[rot] = Rot[rot] Scale[m s].
+    If rot.to_matrix() is a signed permutation matrix, works for any s.
+    Otherwise works only if s is a uniform scaling.
+    """
+    m = nearby_signed_perm_matrix(rot)  # snap to signed perm matrix
+    m.transpose()  # invert permutation
+    for i in range(3):
+        for j in range(3):
+            m[i][j] = abs(m[i][j])  # discard sign
+    return m
+
+def nearby_signed_perm_matrix(rot):
+    """Returns a signed permutation matrix close to rot.to_matrix().
+    (A signed permutation matrix is like a permutation matrix, except
+    the non-zero entries can be ±1.)
+    """
+    m = rot.to_matrix()
+    x, y, z = m[0], m[1], m[2]
+
+    # Set the largest entry in the first row to ±1
+    a, b, c = abs(x[0]), abs(x[1]), abs(x[2])
+    i = 0 if a >= b and a >= c else 1 if b >= c else 2
+    x[i] = 1 if x[i] > 0 else -1
+    x[(i+1) % 3] = 0
+    x[(i+2) % 3] = 0
+
+    # Same for second row: only two columns to consider now.
+    a, b = abs(y[(i+1) % 3]), abs(y[(i+2) % 3])
+    j = (i+1) % 3 if a >= b else (i+2) % 3
+    y[j] = 1 if y[j] > 0 else -1
+    y[(j+1) % 3] = 0
+    y[(j+2) % 3] = 0
+
+    # Same for third row: only one column left
+    k = (0 + 1 + 2) - i - j
+    z[k] = 1 if z[k] > 0 else -1
+    z[(k+1) % 3] = 0
+    z[(k+2) % 3] = 0
+
+    return m
diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_animation_bone.py b/io_scene_gltf2/blender/imp/gltf2_blender_animation_bone.py
index 62c179c1..6f65de43 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_animation_bone.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_animation_bone.py
@@ -58,12 +58,14 @@ class BlenderBoneAnim():
         else:
             translation_keyframes = (gltf.loc_gltf_to_blender(vals) for vals in values)
 
+        final_translations = vnode.base_locs_to_final_locs(translation_keyframes)
+
         # Calculate pose bone trans from final bone trans
         edit_trans, edit_rot = vnode.editbone_trans, vnode.editbone_rot
         edit_rot_inv = edit_rot.conjugated()
         pose_translations = [
             edit_rot_inv @ (trans - edit_trans)
-            for trans in translation_keyframes
+            for trans in final_translations
         ]
 
         BlenderBoneAnim.fill_fcurves(
@@ -93,12 +95,14 @@ class BlenderBoneAnim():
         else:
             quat_keyframes = [gltf.quaternion_gltf_to_blender(vals) for vals in values]
 
+        final_rots = vnode.base_rots_to_final_rots(quat_keyframes)
+
         # Calculate pose bone rotation from final bone rotation
         edit_rot = vnode.editbone_rot
         edit_rot_inv = edit_rot.conjugated()
         pose_rots = [
             edit_rot_inv @ rot
-            for rot in quat_keyframes
+            for rot in final_rots
         ]
 
         # Manage antipodal quaternions
@@ -133,10 +137,13 @@ class BlenderBoneAnim():
         else:
             scale_keyframes = [gltf.scale_gltf_to_blender(vals) for vals in values]
 
+        final_scales = vnode.base_scales_to_final_scales(scale_keyframes)
+        pose_scales = final_scales  # no change needed
+
         BlenderBoneAnim.fill_fcurves(
             obj.animation_data.action,
             keys,
-            scale_keyframes,
+            pose_scales,
             group_name,
             blender_path,
             animation.samplers[channel.sampler].interpolation
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 a0205483..b6369b8b 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_animation_node.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_animation_node.py
@@ -74,20 +74,22 @@ class BlenderNodeAnim():
                 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 by their bone length (always 1 currently)
-                    off = Vector((0, -1, 0))
-                    values = [gltf.loc_gltf_to_blender(vals) + off for vals in values]
-                else:
-                    values = [gltf.loc_gltf_to_blender(vals) for vals in values]
+                    # 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)):
@@ -99,6 +101,7 @@ class BlenderNodeAnim():
                 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)
diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_node.py b/io_scene_gltf2/blender/imp/gltf2_blender_node.py
index d1ffdbe9..6ce91ea3 100755
--- a/io_scene_gltf2/blender/imp/gltf2_blender_node.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_node.py
@@ -78,7 +78,7 @@ class BlenderNode():
             set_extras(obj, pynode.extras)
 
         # Set transform
-        trans, rot, scale = vnode.trs
+        trans, rot, scale = vnode.trs()
         obj.location = trans
         obj.rotation_mode = 'QUATERNION'
         obj.rotation_quaternion = rot
@@ -96,8 +96,8 @@ class BlenderNode():
                 obj.parent_bone = parent_vnode.blender_bone_name
 
                 # Nodes with a bone parent need to be translated
-                # backwards by their bone length (always 1 currently)
-                obj.location += Vector((0, -1, 0))
+                # backwards from the tip to the root
+                obj.location += Vector((0, -parent_vnode.bone_length, 0))
 
         bpy.data.scenes[gltf.blender_scene].collection.objects.link(obj)
 
@@ -138,6 +138,7 @@ class BlenderNode():
             editbone.head = arma_mat @ Vector((0, 0, 0))
             editbone.tail = arma_mat @ Vector((0, 1, 0))
             editbone.align_roll(arma_mat @ Vector((0, 0, 1)) - editbone.head)
+            editbone.length = vnode.bone_length
 
             if isinstance(id, int):
                 pynode = gltf.data.nodes[id]
@@ -161,7 +162,7 @@ class BlenderNode():
 
             # BoneTRS = EditBone * PoseBone
             # Set PoseBone to make BoneTRS = vnode.trs.
-            t, r, s = vnode.trs
+            t, r, s = vnode.trs()
             et, er = vnode.editbone_trans, vnode.editbone_rot
             pose_bone.location = er.conjugated() @ (t - et)
             pose_bone.rotation_mode = 'QUATERNION'
diff --git a/io_scene_gltf2/blender/imp/gltf2_blender_vnode.py b/io_scene_gltf2/blender/imp/gltf2_blender_vnode.py
index bd5edcd1..554c09e9 100644
--- a/io_scene_gltf2/blender/imp/gltf2_blender_vnode.py
+++ b/io_scene_gltf2/blender/imp/gltf2_blender_vnode.py
@@ -15,6 +15,8 @@
 import bpy
 from mathutils import Vector, Quaternion, Matrix
 
+from ..com.gltf2_blender_math import scale_rot_swap_matrix, nearby_signed_perm_matrix
+
 def compute_vnodes(gltf):
     """Computes the tree of virtual nodes.
     Copies the glTF nodes into a tree of VNodes, then performs a series of
@@ -26,6 +28,7 @@ def compute_vnodes(gltf):
     fixup_multitype_nodes(gltf)
     correct_cameras_and_lights(gltf)
     pick_bind_pose(gltf)
+    prettify_bones(gltf)
     calc_bone_matrices(gltf)
 
 
@@ -45,17 +48,60 @@ class VNode:
         self.parent = None
         self.type = VNode.Object
         self.is_arma = False
-        self.trs = (
+        self.base_trs = (
             Vector((0, 0, 0)),
             Quaternion((1, 0, 0, 0)),
             Vector((1, 1, 1)),
         )
+        # Additional rotations before/after the base TRS.
+        # Allows per-vnode axis adjustment. See local_rotation.
+        self.rotation_after = Quaternion((1, 0, 0, 0))
+        self.rotation_before = Quaternion((1, 0, 0, 0))
+
         # Indices of the glTF node where the mesh, etc. came from.
         # (They can get moved around.)
         self.mesh_node_idx = None
         self.camera_node_idx = None
         self.light_node_idx = None
 
+    def trs(self):
+        # (final TRS) = (rotation after) (base TRS) (rotation before)
+        t, r, s = self.base_trs
+        m = scale_rot_swap_matrix(self.rotation_before)
+        return (
+            self.rotation_after @ t,
+            self.rotation_after @ r @ self.rotation_before,
+            m @ s,
+        )
+
+    def base_locs_to_final_locs(self, base_locs):
+        ra = self.rotation_after
+        return [ra @ loc for loc in base_locs]
+
+    def base_rots_to_final_rots(self, base_rots):
+        ra, rb = self.rotation_after, self.rotation_before
+        return [ra @ rot @ rb for rot in base_rots]
+
+    def base_scales_to_final_scales(self, base_scales):
+        m = scale_rot_swap_matrix(self.rotation_before)
+        return [m @ scale for scale in base_scales]
+
+def local_rotation(gltf, vnode_id, rot):
+    """Appends a local rotation to vnode's world transform:
+    (new world transform) = (old world transform) @ (rot)
+    without changing the world transform of vnode's children.
+
+    For correctness, rot must be a signed permutation of the axes
+    (eg. (X Y Z)->(X -Z Y)) OR vnode's scale must always be uniform.
+    """
+    gltf.vnodes[vnode_id].rotation_before @= rot
+
+    # Append the inverse rotation after children's TRS to cancel it out.
+    rot_inv = rot.conjugated()
+    for child in gltf.vnodes[vnode_id].children:
+        gltf.vnodes[child].rotation_after = \
+            rot_inv @ gltf.vnodes[child].rotation_after
+
 
 def init_vnodes(gltf):
     # Map of all VNodes. The keys are arbitrary IDs.
@@ -67,7 +113,7 @@ def init_vnodes(gltf):
         gltf.vnodes[i] = vnode
         vnode.name = pynode.name or 'Node_%d' % i
         vnode.children = list(pynode.children or [])
-        vnode.trs = get_node_trs(gltf, pynode)
+        vnode.base_trs = get_node_trs(gltf, pynode)
         if pynode.mesh is not None:
             vnode.mesh_node_idx = i
         if pynode.camera is not None:
@@ -216,7 +262,7 @@ def move_skinned_meshes(gltf):
         )
         if ok_to_move:
             reparent(gltf, id, new_parent=arma)
-            vnode.trs = (
+            vnode.base_trs = (
                 Vector((0, 0, 0)),
                 Quaternion((1, 0, 0, 0)),
                 Vector((1, 1, 1)),
@@ -304,35 +350,13 @@ def correct_cameras_and_lights(gltf):
     if gltf.camera_correction is None:
         return
 
-    trs = (Vector((0, 0, 0)), gltf.camera_correction, Vector((1, 1, 1)))
-
-    ids = list(gltf.vnodes.keys())
-    for id in ids:
-        vnode = gltf.vnodes[id]
-
-        # Move the camera/light onto a new child and set its rotation
-        # TODO: "hard apply" the rotation without creating a new node
-        #       (like we'll need to do for bones)
+    for id, vnode in gltf.vnodes.items():
+        needs_correction = \
+           vnode.camera_node_idx is not None or \
+           vnode.light_node_idx is not None
 
-        if vnode.camera_node_idx is not None:
-            new_id = str(id) + '.camera-correction'
-            gltf.vnodes[new_id] = VNode()
-            gltf.vnodes[new_id].name = vnode.name + ' Correction'
-            gltf.vnodes[new_id].trs = trs
-            gltf.vnodes[new_id].camera_node_idx = vnode.camera_node_idx
-            gltf.vnodes[new_id].parent = id
-            vnode.children.append(new_id)
-            vnode.camera_node_idx = None
-
-        if vnode.light_node_idx is not None:
-            new_id = str(id) + '.light-correction'
-            gltf.vnodes[new_id] = VNode()
-            gltf.vnodes[new_id].name = vnode.name + ' Correction'
-            gltf.vnodes[new_id].trs = trs
-            gltf.vnodes[new_id].light_node_idx = vnode.light_node_idx
-            gltf.vnodes[new_id].parent = id
-            vnode.children.append(new_id)
-            vnode.light_node_idx = None
+        if needs_correction:
+            local_rotation(gltf, id, gltf.camera_correction)
 
 
 def pick_bind_pose(gltf):
@@ -345,14 +369,102 @@ def pick_bind_pose(gltf):
         if vnode.type == VNode.Bone:
             # For now, use the node TR for bind pose.
             # TODO: try calculating from inverseBindMatices?
-            vnode.bind_trans = Vector(vnode.trs[0])
-            vnode.bind_rot = Quaternion(vnode.trs[1])
+            vnode.bind_trans = Vector(vnode.base_trs[0])
+            vnode.bind_rot = Quaternion(vnode.base_trs[1])
 
             # Initialize editbones to match bind pose
             vnode.editbone_trans = Vector(vnode.bind_trans)
             vnode.editbone_rot = Quaternion(vnode.bind_rot)
 
 
+def prettify_bones(gltf):
+    """
+    Prettify bone lengths/directions.
+    """
+    def visit(vnode_id, parent_rot=None):  # Depth-first walk
+        vnode = gltf.vnodes[vnode_id]
+        rot = None
+
+        if vnode.type == VNode.Bone:
+            vnode.bone_length = pick_bone_length(gltf, vnode_id)
+            rot = pick_bone_rotation(gltf, vnode_id, parent_rot)
+            if rot is not None:
+                rotate_edit_bone(gltf, vnode_id, rot)
+
+        for child in vnode.children:
+            visit(child, parent_rot=rot)
+
+    visit('root')
+
+MIN_BONE_LENGTH = 0.004  # too small and bones get deleted
+
+def pick_bone_length(gltf, bone_id):
+    """Heuristic for bone length."""
+    vnode = gltf.vnodes[bone_id]
+
+    child_locs = [
+        gltf.vnodes[child].editbone_trans
+        for child in vnode.children
+        if gltf.vnodes[child].type == VNode.Bone
+    ]
+    child_locs = [loc for loc in child_locs if loc.length > MIN_BONE_LENGTH]
+    if child_locs:
+        return min(loc.length for loc in child_locs)
+
+    if gltf.vnodes[vnode.parent].type == VNode.Bone:
+        return gltf.vnodes[vnode.parent].bone_length
+
+    if vnode.editbone_trans.length > MIN_BONE_LENGTH:
+        return vnode.editbone_trans.length
+
+    return 1
+
+def pick_bone_rotation(gltf, bone_id, parent_rot):
+    """Heuristic for bone rotation.
+    A bone's tip lies on its local +Y axis so rotating a bone let's us
+    adjust the bone direction.
+    """
+    if bpy.app.debug_value == 100:
+        return None
+
+    if gltf.import_settings['bone_heuristic'] == 'BLENDER':
+        return Quaternion((2**0.5/2, 2**0.5/2, 0, 0))
+    elif gltf.import_settings['bone_heuristic'] == 'TEMPERANCE':
+        return temperance(gltf, bone_id, parent_rot)
+
+def temperance(gltf, bone_id, parent_rot):
+    vnode = gltf.vnodes[bone_id]
+
+    # Try to put our tip at the centroid of our children
+    child_locs = [
+        gltf.vnodes[child].editbone_trans
+        for child in vnode.children
+        if gltf.vnodes[child].type == VNode.Bone
+    ]
+    child_locs = [loc for loc in child_locs if loc.length > MIN_BONE_LENGTH]
+    if child_locs:
+        centroid = sum(child_locs, Vector((0, 0, 0)))
+        rot = Vector((0, 1, 0)).rotation_difference(centroid)
+        rot = nearby_signed_perm_matrix(rot).to_quaternion()
+        return rot
+
+    return parent_rot
+
+def rotate_edit_bone(gltf, bone_id, rot):
+    """Rotate one edit bone without affecting anything else."""
+    gltf.vnodes[bone_id].editbone_rot @= rot
+    # Cancel out the rotation so children aren't affected.
+    rot_inv = rot.conjugated()
+    for child_id in gltf.vnodes[bone_id].children:
+        child = gltf.vnodes[child_id]
+        if child.type == VNode.Bone:
+            child.editbone_trans = rot_inv @ child.editbone_trans
+            child.editbone_rot = rot_inv @ child.editbone_rot
+    # Need to rotate the bone's final TRS by the same amount so skinning
+    # isn't affected.
+    local_rotation(gltf, bone_id, rot)
+
+
 def calc_bone_matrices(gltf):
     """
     Calculate the transformations from bone space to arma space in the bind
@@ -380,6 +492,3 @@ def calc_bone_matrices(gltf):
             visit(child)
 
     visit('root')
-
-
-# TODO: add pass to rotate/resize bones so they look pretty