befd6cdbe89de318c8e858e87475be61db6878b6
[blender.git] / release / scripts / op / io_anim_bvh / import_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
23 import math
24 from math import radians
25
26 import bpy
27 import mathutils
28 from mathutils import Vector, Euler, Matrix
29
30
31 class bvh_node_class(object):
32     __slots__ = (
33     'name',# bvh joint name
34     'parent',# bvh_node_class type or None for no parent
35     'children',# a list of children of this type.
36     'rest_head_world',# worldspace rest location for the head of this node
37     'rest_head_local',# localspace rest location for the head of this node
38     'rest_tail_world',# # worldspace rest location for the tail of this node
39     'rest_tail_local',# # worldspace rest location for the tail of this node
40     'channels',# list of 6 ints, -1 for an unused channel, otherwise an index for the BVH motion data lines, lock triple then rot triple
41     'rot_order',# a triple of indicies as to the order rotation is applied. [0,1,2] is x/y/z - [None, None, None] if no rotation.
42     'anim_data',# a list one tuple's one for each frame. (locx, locy, locz, rotx, roty, rotz)
43     'has_loc',# Conveinience function, bool, same as (channels[0]!=-1 or channels[1]!=-1 channels[2]!=-1)
44     'has_rot',# Conveinience function, bool, same as (channels[3]!=-1 or channels[4]!=-1 channels[5]!=-1)
45     'temp')# use this for whatever you want
46
47     def __init__(self, name, rest_head_world, rest_head_local, parent, channels, rot_order):
48         self.name = name
49         self.rest_head_world = rest_head_world
50         self.rest_head_local = rest_head_local
51         self.rest_tail_world = None
52         self.rest_tail_local = None
53         self.parent = parent
54         self.channels = channels
55         self.rot_order = rot_order
56
57         # convenience functions
58         self.has_loc = channels[0] != -1 or channels[1] != -1 or channels[2] != -1
59         self.has_rot = channels[3] != -1 or channels[4] != -1 or channels[5] != -1
60
61
62         self.children = []
63
64         # list of 6 length tuples: (lx,ly,lz, rx,ry,rz)
65         # even if the channels arnt used they will just be zero
66         #
67         self.anim_data = [(0, 0, 0, 0, 0, 0)]
68
69     def __repr__(self):
70         return 'BVH name:"%s", rest_loc:(%.3f,%.3f,%.3f), rest_tail:(%.3f,%.3f,%.3f)' %\
71         (self.name,\
72         self.rest_head_world.x, self.rest_head_world.y, self.rest_head_world.z,\
73         self.rest_head_world.x, self.rest_head_world.y, self.rest_head_world.z)
74
75
76 # Change the order rotation is applied.
77 MATRIX_IDENTITY_3x3 = Matrix([1, 0, 0], [0, 1, 0], [0, 0, 1])
78 MATRIX_IDENTITY_4x4 = Matrix([1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1])
79
80
81 def eulerRotate(x, y, z, rot_order):
82     # Clamp all values between 0 and 360, values outside this raise an error.
83     mats = [Matrix.Rotation(x, 3, 'X'), Matrix.Rotation(y, 3, 'Y'), Matrix.Rotation(z, 3, 'Z')]
84     return (MATRIX_IDENTITY_3x3 * mats[rot_order[0]] * (mats[rot_order[1]] * (mats[rot_order[2]]))).to_euler()
85
86     # Should work but doesnt!
87     '''
88     eul = Euler((x, y, z))
89     eul.order = "XYZ"[rot_order[0]] + "XYZ"[rot_order[1]] + "XYZ"[rot_order[2]]
90     return tuple(eul.to_matrix().to_euler())
91     '''
92
93
94 def read_bvh(context, file_path, ROT_MODE='XYZ', GLOBAL_SCALE=1.0):
95     # File loading stuff
96     # Open the file for importing
97     file = open(file_path, 'rU')
98
99     # Seperate into a list of lists, each line a list of words.
100     file_lines = file.readlines()
101     # Non standard carrage returns?
102     if len(file_lines) == 1:
103         file_lines = file_lines[0].split('\r')
104
105     # Split by whitespace.
106     file_lines = [ll for ll in [l.split() for l in file_lines] if ll]
107
108
109     # Create Hirachy as empties
110
111     if file_lines[0][0].lower() == 'hierarchy':
112         #print 'Importing the BVH Hierarchy for:', file_path
113         pass
114     else:
115         raise 'ERROR: This is not a BVH file'
116
117     bvh_nodes = {None: None}
118     bvh_nodes_serial = [None]
119
120     channelIndex = -1
121
122
123     lineIdx = 0 # An index for the file.
124     while lineIdx < len(file_lines) -1:
125         #...
126         if file_lines[lineIdx][0].lower() == 'root' or file_lines[lineIdx][0].lower() == 'joint':
127
128             # Join spaces into 1 word with underscores joining it.
129             if len(file_lines[lineIdx]) > 2:
130                 file_lines[lineIdx][1] = '_'.join(file_lines[lineIdx][1:])
131                 file_lines[lineIdx] = file_lines[lineIdx][:2]
132
133             # MAY NEED TO SUPPORT MULTIPLE ROOT's HERE!!!, Still unsure weather multiple roots are possible.??
134
135             # Make sure the names are unique- Object names will match joint names exactly and both will be unique.
136             name = file_lines[lineIdx][1]
137
138             #print '%snode: %s, parent: %s' % (len(bvh_nodes_serial) * '  ', name,  bvh_nodes_serial[-1])
139
140             lineIdx += 2 # Incriment to the next line (Offset)
141             rest_head_local = Vector((float(file_lines[lineIdx][1]), float(file_lines[lineIdx][2]), float(file_lines[lineIdx][3]))) * GLOBAL_SCALE
142             lineIdx += 1 # Incriment to the next line (Channels)
143
144             # newChannel[Xposition, Yposition, Zposition, Xrotation, Yrotation, Zrotation]
145             # newChannel references indecies to the motiondata,
146             # if not assigned then -1 refers to the last value that will be added on loading at a value of zero, this is appended
147             # We'll add a zero value onto the end of the MotionDATA so this is always refers to a value.
148             my_channel = [-1, -1, -1, -1, -1, -1]
149             my_rot_order = [None, None, None]
150             rot_count = 0
151             for channel in file_lines[lineIdx][2:]:
152                 channel = channel.lower()
153                 channelIndex += 1 # So the index points to the right channel
154                 if channel == 'xposition':
155                     my_channel[0] = channelIndex
156                 elif channel == 'yposition':
157                     my_channel[1] = channelIndex
158                 elif channel == 'zposition':
159                     my_channel[2] = channelIndex
160
161                 elif channel == 'xrotation':
162                     my_channel[3] = channelIndex
163                     my_rot_order[rot_count] = 0
164                     rot_count += 1
165                 elif channel == 'yrotation':
166                     my_channel[4] = channelIndex
167                     my_rot_order[rot_count] = 1
168                     rot_count += 1
169                 elif channel == 'zrotation':
170                     my_channel[5] = channelIndex
171                     my_rot_order[rot_count] = 2
172                     rot_count += 1
173
174             channels = file_lines[lineIdx][2:]
175
176             my_parent = bvh_nodes_serial[-1] # account for none
177
178
179             # Apply the parents offset accumletivly
180             if my_parent == None:
181                 rest_head_world = Vector(rest_head_local)
182             else:
183                 rest_head_world = my_parent.rest_head_world + rest_head_local
184
185             bvh_node = bvh_nodes[name] = bvh_node_class(name, rest_head_world, rest_head_local, my_parent, my_channel, my_rot_order)
186
187             # If we have another child then we can call ourselves a parent, else
188             bvh_nodes_serial.append(bvh_node)
189
190         # Account for an end node
191         if file_lines[lineIdx][0].lower() == 'end' and file_lines[lineIdx][1].lower() == 'site': # There is somtimes a name after 'End Site' but we will ignore it.
192             lineIdx += 2 # Incriment to the next line (Offset)
193             rest_tail = Vector((float(file_lines[lineIdx][1]), float(file_lines[lineIdx][2]), float(file_lines[lineIdx][3]))) * GLOBAL_SCALE
194
195             bvh_nodes_serial[-1].rest_tail_world = bvh_nodes_serial[-1].rest_head_world + rest_tail
196             bvh_nodes_serial[-1].rest_tail_local = bvh_nodes_serial[-1].rest_head_local + rest_tail
197
198
199             # Just so we can remove the Parents in a uniform way- End end never has kids
200             # so this is a placeholder
201             bvh_nodes_serial.append(None)
202
203         if len(file_lines[lineIdx]) == 1 and file_lines[lineIdx][0] == '}': # == ['}']
204             bvh_nodes_serial.pop() # Remove the last item
205
206         if len(file_lines[lineIdx]) == 1 and file_lines[lineIdx][0].lower() == 'motion':
207             #print '\nImporting motion data'
208             lineIdx += 3 # Set the cursor to the first frame
209             break
210
211         lineIdx += 1
212
213
214     # Remove the None value used for easy parent reference
215     del bvh_nodes[None]
216     # Dont use anymore
217     del bvh_nodes_serial
218
219     bvh_nodes_list = bvh_nodes.values()
220
221     while lineIdx < len(file_lines):
222         line = file_lines[lineIdx]
223         for bvh_node in bvh_nodes_list:
224             #for bvh_node in bvh_nodes_serial:
225             lx = ly = lz = rx = ry = rz = 0.0
226             channels = bvh_node.channels
227             anim_data = bvh_node.anim_data
228             if channels[0] != -1:
229                 lx = GLOBAL_SCALE * float(line[channels[0]])
230
231             if channels[1] != -1:
232                 ly = GLOBAL_SCALE * float(line[channels[1]])
233
234             if channels[2] != -1:
235                 lz = GLOBAL_SCALE * float(line[channels[2]])
236
237             if channels[3] != -1 or channels[4] != -1 or channels[5] != -1:
238                 rx, ry, rz = float(line[channels[3]]), float(line[channels[4]]), float(line[channels[5]])
239
240                 if ROT_MODE != 'NATIVE':
241                     rx, ry, rz = eulerRotate(radians(rx), radians(ry), radians(rz), bvh_node.rot_order)
242                 else:
243                     rx, ry, rz = radians(rx), radians(ry), radians(rz)
244
245             # Done importing motion data #
246             anim_data.append((lx, ly, lz, rx, ry, rz))
247         lineIdx += 1
248
249     # Assign children
250     for bvh_node in bvh_nodes.values():
251         bvh_node_parent = bvh_node.parent
252         if bvh_node_parent:
253             bvh_node_parent.children.append(bvh_node)
254
255     # Now set the tip of each bvh_node
256     for bvh_node in bvh_nodes.values():
257
258         if not bvh_node.rest_tail_world:
259             if len(bvh_node.children) == 0:
260                 # could just fail here, but rare BVH files have childless nodes
261                 bvh_node.rest_tail_world = Vector(bvh_node.rest_head_world)
262                 bvh_node.rest_tail_local = Vector(bvh_node.rest_head_local)
263             elif len(bvh_node.children) == 1:
264                 bvh_node.rest_tail_world = Vector(bvh_node.children[0].rest_head_world)
265                 bvh_node.rest_tail_local = bvh_node.rest_head_local + bvh_node.children[0].rest_head_local
266             else:
267                 # allow this, see above
268                 #if not bvh_node.children:
269                 #       raise 'error, bvh node has no end and no children. bad file'
270
271                 # Removed temp for now
272                 rest_tail_world = Vector((0.0, 0.0, 0.0))
273                 rest_tail_local = Vector((0.0, 0.0, 0.0))
274                 for bvh_node_child in bvh_node.children:
275                     rest_tail_world += bvh_node_child.rest_head_world
276                     rest_tail_local += bvh_node_child.rest_head_local
277
278                 bvh_node.rest_tail_world = rest_tail_world * (1.0 / len(bvh_node.children))
279                 bvh_node.rest_tail_local = rest_tail_local * (1.0 / len(bvh_node.children))
280
281         # Make sure tail isnt the same location as the head.
282         if (bvh_node.rest_tail_local - bvh_node.rest_head_local).length <= 0.001 * GLOBAL_SCALE:
283             bvh_node.rest_tail_local.y = bvh_node.rest_tail_local.y + GLOBAL_SCALE / 10
284             bvh_node.rest_tail_world.y = bvh_node.rest_tail_world.y + GLOBAL_SCALE / 10
285
286     return bvh_nodes
287
288
289 def bvh_node_dict2objects(context, bvh_nodes, IMPORT_START_FRAME=1, IMPORT_LOOP=False):
290
291     if IMPORT_START_FRAME < 1:
292         IMPORT_START_FRAME = 1
293
294     scn = context.scene
295     scn.objects.selected = []
296
297     objects = []
298
299     def add_ob(name):
300         ob = scn.objects.new('Empty', None)
301         objects.append(ob)
302         return ob
303
304     # Add objects
305     for name, bvh_node in bvh_nodes.items():
306         bvh_node.temp = add_ob(name)
307
308     # Parent the objects
309     for bvh_node in bvh_nodes.values():
310         bvh_node.temp.makeParent([bvh_node_child.temp for bvh_node_child in bvh_node.children], 1, 0) # ojbs, noninverse, 1 = not fast.
311
312     # Offset
313     for bvh_node in bvh_nodes.values():
314         # Make relative to parents offset
315         bvh_node.temp.loc = bvh_node.rest_head_local
316
317     # Add tail objects
318     for name, bvh_node in bvh_nodes.items():
319         if not bvh_node.children:
320             ob_end = add_ob(name + '_end')
321             bvh_node.temp.makeParent([ob_end], 1, 0) # ojbs, noninverse, 1 = not fast.
322             ob_end.loc = bvh_node.rest_tail_local
323
324
325     # Animate the data, the last used bvh_node will do since they all have the same number of frames
326     for frame_current in range(len(bvh_node.anim_data)):
327         Blender.Set('curframe', frame_current + IMPORT_START_FRAME)
328
329         for bvh_node in bvh_nodes.values():
330             lx, ly, lz, rx, ry, rz = bvh_node.anim_data[frame_current]
331
332             rest_head_local = bvh_node.rest_head_local
333             bvh_node.temp.loc = rest_head_local + Vector((lx, ly, lz))
334
335             bvh_node.temp.rot = rx, ry, rz
336
337             bvh_node.temp.insertIpoKey(Blender.Object.IpoKeyTypes.LOCROT) # XXX invalid
338
339     scn.update(1)
340     return objects
341
342
343 def bvh_node_dict2armature(context, bvh_nodes, ROT_MODE='XYZ', IMPORT_START_FRAME=1, IMPORT_LOOP=False):
344
345     if IMPORT_START_FRAME < 1:
346         IMPORT_START_FRAME = 1
347
348     # Add the new armature,
349     scn = context.scene
350 #XXX    scn.objects.selected = []
351     for ob in scn.objects:
352         ob.select = False
353
354     scn.set_frame(IMPORT_START_FRAME)
355
356     arm_data = bpy.data.armatures.new("MyBVH")
357     arm_ob = bpy.data.objects.new("MyBVH", arm_data)
358
359     scn.objects.link(arm_ob)
360
361     arm_ob.select = True
362     scn.objects.active = arm_ob
363     print(scn.objects.active)
364
365     bpy.ops.object.mode_set(mode='OBJECT', toggle=False)
366     bpy.ops.object.mode_set(mode='EDIT', toggle=False)
367
368
369     # Get the average bone length for zero length bones, we may not use this.
370     average_bone_length = 0.0
371     nonzero_count = 0
372     for bvh_node in bvh_nodes.values():
373         l = (bvh_node.rest_head_local - bvh_node.rest_tail_local).length
374         if l:
375             average_bone_length += l
376             nonzero_count += 1
377
378     # Very rare cases all bones couldbe zero length???
379     if not average_bone_length:
380         average_bone_length = 0.1
381     else:
382         # Normal operation
383         average_bone_length = average_bone_length / nonzero_count
384
385
386 #XXX - sloppy operator code
387
388     bpy.ops.armature.delete()
389     bpy.ops.armature.select_all()
390     bpy.ops.armature.delete()
391
392     ZERO_AREA_BONES = []
393     for name, bvh_node in bvh_nodes.items():
394         # New editbone
395         bpy.ops.armature.bone_primitive_add(name="Bone")
396
397         bone = bvh_node.temp = arm_data.edit_bones[-1]
398
399         bone.name = name
400 #               arm_data.bones[name]= bone
401
402         bone.head = bvh_node.rest_head_world
403         bone.tail = bvh_node.rest_tail_world
404
405         # ZERO AREA BONES.
406         if (bone.head - bone.tail).length < 0.001:
407             if bvh_node.parent:
408                 ofs = bvh_node.parent.rest_head_local - bvh_node.parent.rest_tail_local
409                 if ofs.length: # is our parent zero length also?? unlikely
410                     bone.tail = bone.tail + ofs
411                 else:
412                     bone.tail.y = bone.tail.y + average_bone_length
413             else:
414                 bone.tail.y = bone.tail.y + average_bone_length
415
416             ZERO_AREA_BONES.append(bone.name)
417
418
419     for bvh_node in bvh_nodes.values():
420         if bvh_node.parent:
421             # bvh_node.temp is the Editbone
422
423             # Set the bone parent
424             bvh_node.temp.parent = bvh_node.parent.temp
425
426             # Set the connection state
427             if not bvh_node.has_loc and\
428             bvh_node.parent and\
429             bvh_node.parent.temp.name not in ZERO_AREA_BONES and\
430             bvh_node.parent.rest_tail_local == bvh_node.rest_head_local:
431                 bvh_node.temp.use_connect = True
432
433     # Replace the editbone with the editbone name,
434     # to avoid memory errors accessing the editbone outside editmode
435     for bvh_node in bvh_nodes.values():
436         bvh_node.temp = bvh_node.temp.name
437
438 #XXX    arm_data.update()
439
440     # Now Apply the animation to the armature
441
442     # Get armature animation data
443     bpy.ops.object.mode_set(mode='OBJECT', toggle=False)
444     bpy.ops.object.mode_set(mode='POSE', toggle=False)
445
446     pose = arm_ob.pose
447     pose_bones = pose.bones
448
449     if ROT_MODE == 'NATIVE':
450         eul_order_lookup = {\
451             (0, 1, 2): 'XYZ',
452             (0, 2, 1): 'XZY',
453             (1, 0, 2): 'YXZ',
454             (1, 2, 0): 'YZX',
455             (2, 0, 1): 'ZXY',
456             (2, 1, 0): 'ZYX'}
457
458         for bvh_node in bvh_nodes.values():
459             bone_name = bvh_node.temp # may not be the same name as the bvh_node, could have been shortened.
460             pose_bone = pose_bones[bone_name]
461             pose_bone.rotation_mode = eul_order_lookup[tuple(bvh_node.rot_order)]
462
463     elif ROT_MODE != 'QUATERNION':
464         for pose_bone in pose_bones:
465             pose_bone.rotation_mode = ROT_MODE
466     else:
467         # Quats default
468         pass
469
470     context.scene.update()
471
472     bpy.ops.pose.select_all() # set
473     bpy.ops.anim.keyframe_insert_menu(type=-4) # XXX -     -4 ???
474
475
476 #XXX    action = Blender.Armature.NLA.NewAction("Action")
477 #XXX    action.setActive(arm_ob)
478
479     #bpy.ops.action.new()
480     #action = bpy.data.actions[-1]
481
482     # arm_ob.animation_data.action = action
483     action = arm_ob.animation_data.action
484
485     # Replace the bvh_node.temp (currently an editbone)
486     # With a tuple  (pose_bone, armature_bone, bone_rest_matrix, bone_rest_matrix_inv)
487     for bvh_node in bvh_nodes.values():
488         bone_name = bvh_node.temp # may not be the same name as the bvh_node, could have been shortened.
489         pose_bone = pose_bones[bone_name]
490         rest_bone = arm_data.bones[bone_name]
491         bone_rest_matrix = rest_bone.matrix_local.rotation_part()
492
493
494         bone_rest_matrix_inv = Matrix(bone_rest_matrix)
495         bone_rest_matrix_inv.invert()
496
497         bone_rest_matrix_inv.resize4x4()
498         bone_rest_matrix.resize4x4()
499         bvh_node.temp = (pose_bone, bone, bone_rest_matrix, bone_rest_matrix_inv)
500
501
502     # Make a dict for fast access without rebuilding a list all the time.
503
504     # KEYFRAME METHOD, SLOW, USE IPOS DIRECT
505     # TODO: use f-point samples instead (Aligorith)
506
507     if ROT_MODE != 'QUATERNION':
508         prev_euler = [Euler() for i in range(len(bvh_nodes))]
509
510     # Animate the data, the last used bvh_node will do since they all have the same number of frames
511     for frame_current in range(len(bvh_node.anim_data)-1): # skip the first frame (rest frame)
512         # print frame_current
513
514         # if frame_current==40: # debugging
515         #       break
516
517         # Dont neet to set the current frame
518         for i, bvh_node in enumerate(bvh_nodes.values()):
519             pose_bone, bone, bone_rest_matrix, bone_rest_matrix_inv = bvh_node.temp
520             lx, ly, lz, rx, ry, rz = bvh_node.anim_data[frame_current + 1]
521
522             if bvh_node.has_rot:
523                 bone_rotation_matrix = Euler((rx, ry, rz)).to_matrix().resize4x4()
524                 bone_rotation_matrix = bone_rest_matrix_inv * bone_rotation_matrix * bone_rest_matrix
525
526                 if ROT_MODE == 'QUATERNION':
527                     pose_bone.rotation_quaternion = bone_rotation_matrix.to_quat()
528                 else:
529                     euler = bone_rotation_matrix.to_euler(pose_bone.rotation_mode, prev_euler[i])
530                     pose_bone.rotation_euler = euler
531                     prev_euler[i] = euler
532
533             if bvh_node.has_loc:
534                 pose_bone.location = (bone_rest_matrix_inv * Matrix.Translation(Vector((lx, ly, lz)) - bvh_node.rest_head_local)).translation_part()
535
536             if bvh_node.has_loc:
537                 pose_bone.keyframe_insert("location")
538             if bvh_node.has_rot:
539                 if ROT_MODE == 'QUATERNION':
540                     pose_bone.keyframe_insert("rotation_quaternion")
541                 else:
542                     pose_bone.keyframe_insert("rotation_euler")
543
544
545         # bpy.ops.anim.keyframe_insert_menu(type=-4) # XXX -     -4 ???
546         bpy.ops.screen.frame_offset(delta=1)
547
548     for cu in action.fcurves:
549         if IMPORT_LOOP:
550             pass # 2.5 doenst have cyclic now?
551
552         for bez in cu.keyframe_points:
553             bez.interpolation = 'LINEAR'
554
555     return arm_ob
556
557
558 def load(operator, context, filepath="", rotate_mode='NATIVE', scale=1.0, use_cyclic=False, frame_start=1):
559     import time
560     t1 = time.time()
561     print('\tparsing bvh %r...' % filepath, end="")
562
563     bvh_nodes = read_bvh(context, filepath,
564             ROT_MODE=rotate_mode,
565             GLOBAL_SCALE=scale)
566
567     print('%.4f' % (time.time() - t1))
568     t1 = time.time()
569     print('\timporting to blender...', end="")
570
571     bvh_node_dict2armature(context, bvh_nodes,
572             ROT_MODE=rotate_mode,
573             IMPORT_START_FRAME=frame_start,
574             IMPORT_LOOP=use_cyclic)
575
576     print('Done in %.4f\n' % (time.time() - t1))
577     
578     return {'FINISHED'}