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