io_anim_bvh/export_bvh.py

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