io_anim_bvh/export_bvh.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 # Script copyright (C) Campbell Barton
   4 # fixes from Andrea Rugliancich
   5
   6 import bpy
   7
   8
   9 def write_armature(
  10         context,
  11         filepath,
  12         frame_start,
  13         frame_end,
  14         global_scale=1.0,
  15         rotate_mode='NATIVE',
  16         root_transform_only=False,
  17 ):
  18
  19     def ensure_rot_order(rot_order_str):
  20         if set(rot_order_str) != {'X', 'Y', 'Z'}:
  21             rot_order_str = "XYZ"
  22         return rot_order_str
  23
  24     from mathutils import Matrix, Euler
  25     from math import degrees
  26
  27     file = open(filepath, "w", encoding="utf8", newline="\n")
  28
  29     obj = context.object
  30     arm = obj.data
  31
  32     # Build a dictionary of children.
  33     # None for parentless
  34     children = {None: []}
  35
  36     # initialize with blank lists
  37     for bone in arm.bones:
  38         children[bone.name] = []
  39
  40     # keep bone order from armature, no sorting, not esspential but means
  41     # we can maintain order from import -> export which secondlife incorrectly expects.
  42     for bone in arm.bones:
  43         children[getattr(bone.parent, "name", None)].append(bone.name)
  44
  45     # bone name list in the order that the bones are written
  46     serialized_names = []
  47
  48     node_locations = {}
  49
  50     file.write("HIERARCHY\n")
  51
  52     def write_recursive_nodes(bone_name, indent):
  53         my_children = children[bone_name]
  54
  55         indent_str = "\t" * indent
  56
  57         bone = arm.bones[bone_name]
  58         pose_bone = obj.pose.bones[bone_name]
  59         loc = bone.head_local
  60         node_locations[bone_name] = loc
  61
  62         if rotate_mode == "NATIVE":
  63             rot_order_str = ensure_rot_order(pose_bone.rotation_mode)
  64         else:
  65             rot_order_str = rotate_mode
  66
  67         # make relative if we can
  68         if bone.parent:
  69             loc = loc - node_locations[bone.parent.name]
  70
  71         if indent:
  72             file.write("%sJOINT %s\n" % (indent_str, bone_name))
  73         else:
  74             file.write("%sROOT %s\n" % (indent_str, bone_name))
  75
  76         file.write("%s{\n" % indent_str)
  77         file.write("%s\tOFFSET %.6f %.6f %.6f\n" % (indent_str, *(loc * global_scale)))
  78         if (bone.use_connect or root_transform_only) and bone.parent:
  79             file.write("%s\tCHANNELS 3 %srotation %srotation %srotation\n" % (indent_str, *rot_order_str))
  80         else:
  81             file.write("%s\tCHANNELS 6 Xposition Yposition Zposition %srotation %srotation %srotation\n" % (indent_str, *rot_order_str))
  82
  83         if my_children:
  84             # store the location for the children
  85             # to get their relative offset
  86
  87             # Write children
  88             for child_bone in my_children:
  89                 serialized_names.append(child_bone)
  90                 write_recursive_nodes(child_bone, indent + 1)
  91
  92         else:
  93             # Write the bone end.
  94             file.write("%s\tEnd Site\n" % indent_str)
  95             file.write("%s\t{\n" % indent_str)
  96             loc = bone.tail_local - node_locations[bone_name]
  97             file.write("%s\t\tOFFSET %.6f %.6f %.6f\n" % (indent_str, *(loc * global_scale)))
  98             file.write("%s\t}\n" % indent_str)
  99
 100         file.write("%s}\n" % indent_str)
 101
 102     if len(children[None]) == 1:
 103         key = children[None][0]
 104         serialized_names.append(key)
 105         indent = 0
 106
 107         write_recursive_nodes(key, indent)
 108
 109     else:
 110         # Write a dummy parent node, with a dummy key name
 111         # Just be sure it's not used by another bone!
 112         i = 0
 113         key = "__%d" % i
 114         while key in children:
 115             i += 1
 116             key = "__%d" % i
 117         file.write("ROOT %s\n" % key)
 118         file.write("{\n")
 119         file.write("\tOFFSET 0.0 0.0 0.0\n")
 120         file.write("\tCHANNELS 0\n")  # Xposition Yposition Zposition Xrotation Yrotation Zrotation
 121         indent = 1
 122
 123         # Write children
 124         for child_bone in children[None]:
 125             serialized_names.append(child_bone)
 126             write_recursive_nodes(child_bone, indent)
 127
 128         file.write("}\n")
 129
 130     # redefine bones as sorted by serialized_names
 131     # so we can write motion
 132
 133     class DecoratedBone:
 134         __slots__ = (
 135             # Bone name, used as key in many places.
 136             "name",
 137             "parent",  # decorated bone parent, set in a later loop
 138             # Blender armature bone.
 139             "rest_bone",
 140             # Blender pose bone.
 141             "pose_bone",
 142             # Blender pose matrix.
 143             "pose_mat",
 144             # Blender rest matrix (armature space).
 145             "rest_arm_mat",
 146             # Blender rest matrix (local space).
 147             "rest_local_mat",
 148             # Pose_mat inverted.
 149             "pose_imat",
 150             # Rest_arm_mat inverted.
 151             "rest_arm_imat",
 152             # Rest_local_mat inverted.
 153             "rest_local_imat",
 154             # Last used euler to preserve euler compatibility in between keyframes.
 155             "prev_euler",
 156             # Is the bone disconnected to the parent bone?
 157             "skip_position",
 158             "rot_order",
 159             "rot_order_str",
 160             # Needed for the euler order when converting from a matrix.
 161             "rot_order_str_reverse",
 162         )
 163
 164         _eul_order_lookup = {
 165             'XYZ': (0, 1, 2),
 166             'XZY': (0, 2, 1),
 167             'YXZ': (1, 0, 2),
 168             'YZX': (1, 2, 0),
 169             'ZXY': (2, 0, 1),
 170             'ZYX': (2, 1, 0),
 171         }
 172
 173         def __init__(self, bone_name):
 174             self.name = bone_name
 175             self.rest_bone = arm.bones[bone_name]
 176             self.pose_bone = obj.pose.bones[bone_name]
 177
 178             if rotate_mode == "NATIVE":
 179                 self.rot_order_str = ensure_rot_order(self.pose_bone.rotation_mode)
 180             else:
 181                 self.rot_order_str = rotate_mode
 182             self.rot_order_str_reverse = self.rot_order_str[::-1]
 183
 184             self.rot_order = DecoratedBone._eul_order_lookup[self.rot_order_str]
 185
 186             self.pose_mat = self.pose_bone.matrix
 187
 188             # mat = self.rest_bone.matrix  # UNUSED
 189             self.rest_arm_mat = self.rest_bone.matrix_local
 190             self.rest_local_mat = self.rest_bone.matrix
 191
 192             # inverted mats
 193             self.pose_imat = self.pose_mat.inverted()
 194             self.rest_arm_imat = self.rest_arm_mat.inverted()
 195             self.rest_local_imat = self.rest_local_mat.inverted()
 196
 197             self.parent = None
 198             self.prev_euler = Euler((0.0, 0.0, 0.0), self.rot_order_str_reverse)
 199             self.skip_position = ((self.rest_bone.use_connect or root_transform_only) and self.rest_bone.parent)
 200
 201         def update_posedata(self):
 202             self.pose_mat = self.pose_bone.matrix
 203             self.pose_imat = self.pose_mat.inverted()
 204
 205         def __repr__(self):
 206             if self.parent:
 207                 return "[\"%s\" child on \"%s\"]\n" % (self.name, self.parent.name)
 208             else:
 209                 return "[\"%s\" root bone]\n" % (self.name)
 210
 211     bones_decorated = [DecoratedBone(bone_name) for bone_name in serialized_names]
 212
 213     # Assign parents
 214     bones_decorated_dict = {dbone.name: dbone for dbone in bones_decorated}
 215     for dbone in bones_decorated:
 216         parent = dbone.rest_bone.parent
 217         if parent:
 218             dbone.parent = bones_decorated_dict[parent.name]
 219     del bones_decorated_dict
 220     # finish assigning parents
 221
 222     scene = context.scene
 223     frame_current = scene.frame_current
 224
 225     file.write("MOTION\n")
 226     file.write("Frames: %d\n" % (frame_end - frame_start + 1))
 227     file.write("Frame Time: %.6f\n" % (1.0 / (scene.render.fps / scene.render.fps_base)))
 228
 229     for frame in range(frame_start, frame_end + 1):
 230         scene.frame_set(frame)
 231
 232         for dbone in bones_decorated:
 233             dbone.update_posedata()
 234
 235         for dbone in bones_decorated:
 236             trans = Matrix.Translation(dbone.rest_bone.head_local)
 237             itrans = Matrix.Translation(-dbone.rest_bone.head_local)
 238
 239             if dbone.parent:
 240                 mat_final = dbone.parent.rest_arm_mat @ dbone.parent.pose_imat @ dbone.pose_mat @ dbone.rest_arm_imat
 241                 mat_final = itrans @ mat_final @ trans
 242                 loc = mat_final.to_translation() + (dbone.rest_bone.head_local - dbone.parent.rest_bone.head_local)
 243             else:
 244                 mat_final = dbone.pose_mat @ dbone.rest_arm_imat
 245                 mat_final = itrans @ mat_final @ trans
 246                 loc = mat_final.to_translation() + dbone.rest_bone.head
 247
 248             # keep eulers compatible, no jumping on interpolation.
 249             rot = mat_final.to_euler(dbone.rot_order_str_reverse, dbone.prev_euler)
 250
 251             if not dbone.skip_position:
 252                 file.write("%.6f %.6f %.6f " % (loc * global_scale)[:])
 253
 254             file.write("%.6f %.6f %.6f " % (degrees(rot[dbone.rot_order[0]]), degrees(rot[dbone.rot_order[1]]), degrees(rot[dbone.rot_order[2]])))
 255
 256             dbone.prev_euler = rot
 257
 258         file.write("\n")
 259
 260     file.close()
 261
 262     scene.frame_set(frame_current)
 263
 264     print("BVH Exported: %s frames:%d\n" % (filepath, frame_end - frame_start + 1))
 265
 266
 267 def save(
 268         context, filepath="",
 269         frame_start=-1,
 270         frame_end=-1,
 271         global_scale=1.0,
 272         rotate_mode="NATIVE",
 273         root_transform_only=False,
 274 ):
 275     write_armature(
 276         context, filepath,
 277         frame_start=frame_start,
 278         frame_end=frame_end,
 279         global_scale=global_scale,
 280         rotate_mode=rotate_mode,
 281         root_transform_only=root_transform_only,
 282     )
 283
 284     return {'FINISHED'}