io_anim_bvh/export_bvh.py

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