Advanced Retargeting option: If the end user armature is complex, on the level of...
[blender.git] / release / scripts / startup / ui_mocap.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 import bpy
22
23 from bpy.props import *
24 from bpy import *
25 import mocap_constraints
26 import retarget
27 import mocap_tools
28 ### reloads modules (for testing purposes only)
29 from imp import reload
30 reload(mocap_constraints)
31 reload(retarget)
32 reload(mocap_tools)
33
34 from mocap_constraints import *
35
36 # MocapConstraint class
37 # Defines MocapConstraint datatype, used to add and configute mocap constraints
38 # Attached to Armature data
39
40
41 class MocapConstraint(bpy.types.PropertyGroup):
42     name = bpy.props.StringProperty(name="Name",
43         default="Mocap Fix",
44         description="Name of Mocap Fix",
45         update=setConstraint)
46     constrained_bone = bpy.props.StringProperty(name="Bone",
47         default="",
48         description="Constrained Bone",
49         update=updateConstraintBoneType)
50     constrained_boneB = bpy.props.StringProperty(name="Bone (2)",
51         default="",
52         description="Other Constrained Bone (optional, depends on type)",
53         update=setConstraint)
54     s_frame = bpy.props.IntProperty(name="S",
55         default=0,
56         description="Start frame of Fix",
57         update=setConstraint)
58     e_frame = bpy.props.IntProperty(name="E",
59         default=100,
60         description="End frame of Fix",
61         update=setConstraint)
62     smooth_in = bpy.props.IntProperty(name="In",
63         default=10,
64         description="Amount of frames to smooth in",
65         update=setConstraint,
66         min=0)
67     smooth_out = bpy.props.IntProperty(name="Out",
68         default=10,
69         description="Amount of frames to smooth out",
70         update=setConstraint,
71         min=0)
72     targetMesh = bpy.props.StringProperty(name="Mesh",
73         default="",
74         description="Target of Fix - Mesh (optional, depends on type)",
75         update=setConstraint)
76     active = bpy.props.BoolProperty(name="Active",
77         default=True,
78         description="Fix is active",
79         update=setConstraint)
80     show_expanded = bpy.props.BoolProperty(name="Show Expanded",
81         default=True,
82         description="Fix is fully shown")
83     targetPoint = bpy.props.FloatVectorProperty(name="Point", size=3,
84         subtype="XYZ", default=(0.0, 0.0, 0.0),
85         description="Target of Fix - Point",
86         update=setConstraint)
87     targetDist = bpy.props.FloatProperty(name="Offset",
88         default=0.0,
89         description="Distance and Floor Fixes - Desired offset",
90         update=setConstraint)
91     targetSpace = bpy.props.EnumProperty(
92         items=[("WORLD", "World Space", "Evaluate target in global space"),
93             ("LOCAL", "Object space", "Evaluate target in object space"),
94             ("constrained_boneB", "Other Bone Space", "Evaluate target in specified other bone space")],
95         name="Space",
96         description="In which space should Point type target be evaluated",
97         update=setConstraint)
98     type = bpy.props.EnumProperty(name="Type of constraint",
99         items=[("point", "Maintain Position", "Bone is at a specific point"),
100             ("freeze", "Maintain Position at frame", "Bone does not move from location specified in target frame"),
101             ("floor", "Stay above", "Bone does not cross specified mesh object eg floor"),
102             ("distance", "Maintain distance", "Target bones maintained specified distance")],
103         description="Type of Fix",
104         update=updateConstraintBoneType)
105     real_constraint = bpy.props.StringProperty()
106     real_constraint_bone = bpy.props.StringProperty()
107
108
109 bpy.utils.register_class(MocapConstraint)
110
111 bpy.types.Armature.mocap_constraints = bpy.props.CollectionProperty(type=MocapConstraint)
112
113
114 class AnimationStitchSettings(bpy.types.PropertyGroup):
115     first_action = bpy.props.StringProperty(name="Action 1",
116             description="First action in stitch")
117     second_action = bpy.props.StringProperty(name="Action 2",
118             description="Second action in stitch")
119     blend_frame = bpy.props.IntProperty(name="Stitch frame",
120             description="Frame to locate stitch on")
121     blend_amount = bpy.props.IntProperty(name="Blend amount",
122             description="Size of blending transitiion, on both sides of the stitch",
123             default=10)
124     second_offset = bpy.props.IntProperty(name="Second offset",
125             description="Frame offset for 2nd animation, where it should start",
126             default=10)
127     stick_bone = bpy.props.StringProperty(name="Stick Bone",
128             description="Bone to freeze during transition",
129             default="")
130
131 bpy.utils.register_class(AnimationStitchSettings)
132
133
134 class MocapNLATracks(bpy.types.PropertyGroup):
135     name = bpy.props.StringProperty()
136     base_track = bpy.props.StringProperty()
137     auto_fix_track = bpy.props.StringProperty()
138     manual_fix_track = bpy.props.StringProperty()
139     stride_action = bpy.props.StringProperty()
140
141 bpy.utils.register_class(MocapNLATracks)
142
143                     
144 def advancedRetargetToggle(self, context):
145     enduser_obj = context.active_object
146     performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
147     if enduser_obj is None or len(performer_obj) != 1:
148         print("Need active and selected armatures")
149         return
150     else:
151         performer_obj = performer_obj[0]
152     if self.advancedRetarget:
153         retarget.preAdvancedRetargeting(performer_obj, enduser_obj)
154     else:
155         retarget.cleanTempConstraints(enduser_obj)
156
157
158
159 bpy.types.Armature.stitch_settings = bpy.props.PointerProperty(type=AnimationStitchSettings)
160 bpy.types.Armature.active_mocap =  bpy.props.StringProperty(update=retarget.NLASystemInitialize)
161 bpy.types.Armature.mocapNLATracks = bpy.props.CollectionProperty(type=MocapNLATracks)
162 bpy.types.Armature.advancedRetarget = bpy.props.BoolProperty(default=False, update=advancedRetargetToggle)
163
164 #Update function for IK functionality. Is called when IK prop checkboxes are toggled.
165
166
167 def toggleIKBone(self, context):
168     if self.IKRetarget:
169         if not self.is_in_ik_chain:
170             print(self.name + " IK toggled ON!")
171             ik = self.constraints.new('IK')
172             #ik the whole chain up to the root, excluding
173             chainLen = 0
174             for parent_bone in self.parent_recursive:
175                 chainLen += 1
176                 if hasIKConstraint(parent_bone):
177                     break
178                 deformer_children = [child for child in parent_bone.children if child.bone.use_deform]
179                 #~ if len(deformer_children) > 1:
180                     #~ break
181             ik.chain_count = chainLen
182             for bone in self.parent_recursive:
183                 if bone.is_in_ik_chain:
184                     bone.IKRetarget = True
185     else:
186         print(self.name + " IK toggled OFF!")
187         cnstrn_bones = []
188         newChainLength = []
189         if hasIKConstraint(self):
190             cnstrn_bones = [self]
191         elif self.is_in_ik_chain:
192             cnstrn_bones = [child for child in self.children_recursive if hasIKConstraint(child)]
193             for cnstrn_bone in cnstrn_bones:
194                 newChainLength.append(cnstrn_bone.parent_recursive.index(self) + 1)
195         if cnstrn_bones:
196             # remove constraint, and update IK retarget for all parents of cnstrn_bone up to chain_len
197             for i, cnstrn_bone in enumerate(cnstrn_bones):
198                 print(cnstrn_bone.name)
199                 if newChainLength:
200                     ik = hasIKConstraint(cnstrn_bone)
201                     ik.chain_count = newChainLength[i]
202                 else:
203                     ik = hasIKConstraint(cnstrn_bone)
204                     cnstrn_bone.constraints.remove(ik)
205                     cnstrn_bone.IKRetarget = False
206             for bone in cnstrn_bone.parent_recursive:
207                 if not bone.is_in_ik_chain:
208                     bone.IKRetarget = False
209         
210
211 class MocapMapping(bpy.types.PropertyGroup):
212     name = bpy.props.StringProperty()
213
214 bpy.utils.register_class(MocapMapping)
215
216 bpy.types.Bone.map = bpy.props.StringProperty()
217 bpy.types.Bone.reverseMap = bpy.props.CollectionProperty(type=MocapMapping)
218 bpy.types.Bone.foot = bpy.props.BoolProperty(name="Foot",
219     description="Marks this bone as a 'foot', which determines retargeted animation's translation",
220     default=False)
221 bpy.types.PoseBone.IKRetarget = bpy.props.BoolProperty(name="IK",
222     description="Toggles IK Retargeting method for given bone",
223     update=toggleIKBone, default=False)
224
225
226 def updateIKRetarget():
227     # ensures that Blender constraints and IK properties are in sync
228     # currently runs when module is loaded, should run when scene is loaded
229     # or user adds a constraint to armature. Will be corrected in the future,
230     # once python callbacks are implemented
231     for obj in bpy.data.objects:
232         if obj.pose:
233             bones = obj.pose.bones
234             for pose_bone in bones:
235                 if pose_bone.is_in_ik_chain or hasIKConstraint(pose_bone):
236                     pose_bone.IKRetarget = True
237                 else:
238                     pose_bone.IKRetarget = False
239
240 updateIKRetarget()
241
242
243 class MocapPanel(bpy.types.Panel):
244     # Motion capture retargeting panel
245     bl_label = "Mocap tools"
246     bl_space_type = "PROPERTIES"
247     bl_region_type = "WINDOW"
248     bl_context = "object"
249
250     def draw(self, context):
251         self.layout.label("Preprocessing")
252         row = self.layout.row(align=True)
253         row.alignment = 'EXPAND'
254         row.operator("mocap.samples", text='Samples to Beziers')
255         row.operator("mocap.denoise", text='Clean noise')
256         row.operator("mocap.rotate_fix", text='Fix BVH Axis Orientation')
257         row.operator("mocap.scale_fix", text='Auto scale Performer')
258         row2 = self.layout.row(align=True)
259         row2.operator("mocap.looper", text='Loop animation')
260         row2.operator("mocap.limitdof", text='Constrain Rig')
261         row2.operator("mocap.removelimitdof", text='Unconstrain Rig')
262         self.layout.label("Retargeting")
263         enduser_obj = bpy.context.active_object
264         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj]
265         if enduser_obj is None or len(performer_obj) != 1:
266             self.layout.label("Select performer rig and target rig (as active)")
267         else:
268             self.layout.operator("mocap.guessmapping", text="Guess Hiearchy Mapping")
269             row3 = self.layout.row(align=True)
270             column1 = row3.column(align=True)
271             column1.label("Performer Rig")
272             column2 = row3.column(align=True)
273             column2.label("Enduser Rig")
274             performer_obj = performer_obj[0]
275             if performer_obj.data and enduser_obj.data:
276                 if performer_obj.data.name in bpy.data.armatures and enduser_obj.data.name in bpy.data.armatures:
277                     perf = performer_obj.data
278                     enduser_arm = enduser_obj.data
279                     perf_pose_bones = enduser_obj.pose.bones
280                     for bone in perf.bones:
281                         row = self.layout.row()
282                         row.prop(data=bone, property='foot', text='', icon='POSE_DATA')
283                         row.label(bone.name)
284                         row.prop_search(bone, "map", enduser_arm, "bones")
285                         label_mod = "FK"
286                         if bone.map:
287                             pose_bone = perf_pose_bones[bone.map]
288                             if pose_bone.is_in_ik_chain:
289                                 label_mod = "ik chain"
290                             if hasIKConstraint(pose_bone):
291                                 label_mod = "ik end"
292                             row.prop(pose_bone, 'IKRetarget')
293                             row.label(label_mod)
294                         else:
295                             row.label(" ")
296                             row.label(" ")
297                     mapRow = self.layout.row()
298                     mapRow.operator("mocap.savemapping", text='Save mapping')
299                     mapRow.operator("mocap.loadmapping", text='Load mapping')
300                     self.layout.prop(data=performer_obj.animation_data.action, property='name', text='Action Name')
301                     self.layout.prop(enduser_arm, "advancedRetarget", text='Advanced Retarget')
302                     self.layout.operator("mocap.retarget", text='RETARGET!')
303
304
305 class MocapConstraintsPanel(bpy.types.Panel):
306     #Motion capture constraints panel
307     bl_label = "Mocap Fixes"
308     bl_space_type = "PROPERTIES"
309     bl_region_type = "WINDOW"
310     bl_context = "object"
311
312     def draw(self, context):
313         layout = self.layout
314         if context.active_object:
315             if context.active_object.data:
316                 if context.active_object.data.name in bpy.data.armatures:
317                     enduser_obj = context.active_object
318                     enduser_arm = enduser_obj.data
319                     layout.operator_menu_enum("mocap.addmocapfix", "type")
320                     layout.operator("mocap.updateconstraints", text='Update Fixes')
321                     bakeRow = layout.row()
322                     bakeRow.operator("mocap.bakeconstraints", text='Bake Fixes')
323                     bakeRow.operator("mocap.unbakeconstraints", text='Unbake Fixes')
324                     layout.separator()
325                     for i, m_constraint in enumerate(enduser_arm.mocap_constraints):
326                         box = layout.box()
327                         headerRow = box.row()
328                         headerRow.prop(m_constraint, 'show_expanded', text='', icon='TRIA_DOWN' if m_constraint.show_expanded else 'TRIA_RIGHT', emboss=False)
329                         headerRow.prop(m_constraint, 'type', text='')
330                         headerRow.prop(m_constraint, 'name', text='')
331                         headerRow.prop(m_constraint, 'active', icon='MUTE_IPO_ON' if not m_constraint.active else'MUTE_IPO_OFF', text='', emboss=False)
332                         headerRow.operator("mocap.removeconstraint", text="", icon='X', emboss=False).constraint = i
333                         if m_constraint.show_expanded:
334                             box.separator()
335                             box.prop_search(m_constraint, 'constrained_bone', enduser_obj.pose, "bones", icon='BONE_DATA')
336                             if m_constraint.type == "distance" or m_constraint.type == "point":
337                                 box.prop_search(m_constraint, 'constrained_boneB', enduser_obj.pose, "bones", icon='CONSTRAINT_BONE')
338                             frameRow = box.row()
339                             frameRow.label("Frame Range:")
340                             frameRow.prop(m_constraint, 's_frame')
341                             frameRow.prop(m_constraint, 'e_frame')
342                             smoothRow = box.row()
343                             smoothRow.label("Smoothing:")
344                             smoothRow.prop(m_constraint, 'smooth_in')
345                             smoothRow.prop(m_constraint, 'smooth_out')
346                             targetRow = box.row()
347                             targetLabelCol = targetRow.column()
348                             targetLabelCol.label("Target settings:")
349                             targetPropCol = targetRow.column()
350                             if m_constraint.type == "floor":
351                                 targetPropCol.prop_search(m_constraint, 'targetMesh', bpy.data, "objects")
352                             if m_constraint.type == "point" or m_constraint.type == "freeze":
353                                 box.prop(m_constraint, 'targetSpace')
354                             if m_constraint.type == "point":
355                                 targetPropCol.prop(m_constraint, 'targetPoint')
356                             if m_constraint.type == "distance" or m_constraint.type == "floor":
357                                 targetPropCol.prop(m_constraint, 'targetDist')
358                             layout.separator()
359
360
361 class ExtraToolsPanel(bpy.types.Panel):
362     # Motion capture retargeting panel
363     bl_label = "Extra Mocap Tools"
364     bl_space_type = "PROPERTIES"
365     bl_region_type = "WINDOW"
366     bl_context = "object"
367
368     def draw(self, context):
369         layout = self.layout
370         layout.operator('mocap.pathediting', text="Follow Path")
371         layout.label("Animation Stitching")
372         activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
373         if activeIsArmature:
374             enduser_arm = context.active_object.data
375             layout.label("Retargeted Animations:")
376             layout.prop_search(enduser_arm, "active_mocap",enduser_arm, "mocapNLATracks")
377             settings = enduser_arm.stitch_settings
378             layout.prop_search(settings, "first_action", enduser_arm, "mocapNLATracks")
379             layout.prop_search(settings, "second_action", enduser_arm, "mocapNLATracks")
380             layout.prop(settings, "blend_frame")
381             layout.prop(settings, "blend_amount")
382             layout.prop(settings, "second_offset")
383             layout.prop_search(settings, "stick_bone", context.active_object.pose, "bones")
384             layout.operator('mocap.animstitch', text="Stitch Animations")
385
386
387 class OBJECT_OT_RetargetButton(bpy.types.Operator):
388     '''Retarget animation from selected armature to active armature '''
389     bl_idname = "mocap.retarget"
390     bl_label = "Retargets active action from Performer to Enduser"
391
392     def execute(self, context):
393         scene = context.scene
394         s_frame = scene.frame_start
395         e_frame = scene.frame_end
396         enduser_obj = context.active_object
397         performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
398         if enduser_obj is None or len(performer_obj) != 1:
399             print("Need active and selected armatures")
400         else:
401             performer_obj = performer_obj[0]
402             s_frame, e_frame = performer_obj.animation_data.action.frame_range
403             s_frame = int(s_frame)
404             e_frame = int(e_frame)
405         retarget.totalRetarget(performer_obj, enduser_obj, scene, s_frame, e_frame)
406         return {"FINISHED"}
407
408     @classmethod
409     def poll(cls, context):
410         if context.active_object:
411             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
412         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
413         if performer_obj:
414             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
415         else:
416             return False
417             
418     
419     #~ class OBJECT_OT_AdvancedRetargetButton(bpy.types.Operator):
420         #~ '''Prepare for advanced retargeting '''
421         #~ bl_idname = "mocap.preretarget"
422         #~ bl_label = "Prepares retarget of active action from Performer to Enduser"
423
424         #~ def execute(self, context):
425             #~ scene = context.scene
426             #~ s_frame = scene.frame_start
427             #~ e_frame = scene.frame_end
428             #~ enduser_obj = context.active_object
429             #~ performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
430             #~ if enduser_obj is None or len(performer_obj) != 1:
431                 #~ print("Need active and selected armatures")
432             #~ else:
433                 #~ performer_obj = performer_obj[0]
434             #~ retarget.preAdvancedRetargeting(performer_obj, enduser_obj)
435             #~ return {"FINISHED"}
436
437         #~ @classmethod
438         #~ def poll(cls, context):
439             #~ if context.active_object:
440                 #~ activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
441             #~ performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
442             #~ if performer_obj:
443                 #~ return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
444             #~ else:
445                 #~ return False
446
447
448 class OBJECT_OT_SaveMappingButton(bpy.types.Operator):
449     '''Save mapping to active armature (for future retargets) '''
450     bl_idname = "mocap.savemapping"
451     bl_label = "Saves user generated mapping from Performer to Enduser"
452
453     def execute(self, context):
454         enduser_obj = bpy.context.active_object
455         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
456         retarget.createDictionary(performer_obj.data, enduser_obj.data)
457         return {"FINISHED"}
458
459     @classmethod
460     def poll(cls, context):
461         if context.active_object:
462             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
463         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
464         if performer_obj:
465             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
466         else:
467             return False
468
469
470 class OBJECT_OT_LoadMappingButton(bpy.types.Operator):
471     '''Load saved mapping from active armature'''
472     bl_idname = "mocap.loadmapping"
473     bl_label = "Loads user generated mapping from Performer to Enduser"
474
475     def execute(self, context):
476         enduser_obj = bpy.context.active_object
477         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
478         retarget.loadMapping(performer_obj.data, enduser_obj.data)
479         return {"FINISHED"}
480
481     @classmethod
482     def poll(cls, context):
483         if context.active_object:
484             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
485         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
486         if performer_obj:
487             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
488         else:
489             return False
490
491
492 class OBJECT_OT_ConvertSamplesButton(bpy.types.Operator):
493     '''Convert active armature's sampled keyframed to beziers'''
494     bl_idname = "mocap.samples"
495     bl_label = "Converts samples / simplifies keyframes to beziers"
496
497     def execute(self, context):
498         mocap_tools.fcurves_simplify(context, context.active_object)
499         return {"FINISHED"}
500
501     @classmethod
502     def poll(cls, context):
503         return context.active_object.animation_data
504
505
506 class OBJECT_OT_LooperButton(bpy.types.Operator):
507     '''Trim active armature's animation to a single cycle, given a cyclic animation (such as a walk cycle)'''
508     bl_idname = "mocap.looper"
509     bl_label = "loops animation / sampled mocap data"
510
511     def execute(self, context):
512         mocap_tools.autoloop_anim()
513         return {"FINISHED"}
514
515     @classmethod
516     def poll(cls, context):
517         return context.active_object.animation_data
518
519
520 class OBJECT_OT_DenoiseButton(bpy.types.Operator):
521     '''Denoise active armature's animation. Good for dealing with 'bad' frames inherent in mocap animation'''
522     bl_idname = "mocap.denoise"
523     bl_label = "Denoises sampled mocap data "
524
525     def execute(self, context):
526         mocap_tools.denoise_median()
527         return {"FINISHED"}
528
529     @classmethod
530     def poll(cls, context):
531         return context.active_object
532
533     @classmethod
534     def poll(cls, context):
535         return context.active_object.animation_data
536
537
538 class OBJECT_OT_LimitDOFButton(bpy.types.Operator):
539     '''Create limit constraints on the active armature from the selected armature's animation's range of motion'''
540     bl_idname = "mocap.limitdof"
541     bl_label = "Analyzes animations Max/Min DOF and adds hard/soft constraints"
542
543     def execute(self, context):
544         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object][0]
545         mocap_tools.limit_dof(context, performer_obj, context.active_object)
546         return {"FINISHED"}
547
548     @classmethod
549     def poll(cls, context):
550         if context.active_object:
551             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
552         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
553         if performer_obj:
554             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
555         else:
556             return False
557
558
559 class OBJECT_OT_RemoveLimitDOFButton(bpy.types.Operator):
560     '''Removes previously created limit constraints on the active armature'''
561     bl_idname = "mocap.removelimitdof"
562     bl_label = "Removes previously created limit constraints on the active armature"
563
564     def execute(self, context):
565         mocap_tools.limit_dof_toggle_off(context, context.active_object)
566         return {"FINISHED"}
567
568     @classmethod
569     def poll(cls, context):
570         activeIsArmature = False
571         if context.active_object:
572             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
573         return activeIsArmature
574
575
576 class OBJECT_OT_RotateFixArmature(bpy.types.Operator):
577     '''Realign the active armature's axis system to match Blender (Commonly needed after bvh import)'''
578     bl_idname = "mocap.rotate_fix"
579     bl_label = "Rotates selected armature 90 degrees (fix for bvh import)"
580
581     def execute(self, context):
582         mocap_tools.rotate_fix_armature(context.active_object.data)
583         return {"FINISHED"}
584
585     @classmethod
586     def poll(cls, context):
587         if context.active_object:
588             return isinstance(context.active_object.data, bpy.types.Armature)
589
590
591 class OBJECT_OT_ScaleFixArmature(bpy.types.Operator):
592     '''Rescale selected armature to match the active animation, for convienence'''
593     bl_idname = "mocap.scale_fix"
594     bl_label = "Scales performer armature to match target armature"
595
596     def execute(self, context):
597         enduser_obj = bpy.context.active_object
598         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
599         mocap_tools.scale_fix_armature(performer_obj, enduser_obj)
600         return {"FINISHED"}
601
602     @classmethod
603     def poll(cls, context):
604         if context.active_object:
605             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
606         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
607         if performer_obj:
608             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
609         else:
610             return False
611
612
613 class MOCAP_OT_AddMocapFix(bpy.types.Operator):
614     '''Add a post-retarget fix - useful for fixing certain artifacts following the retarget'''
615     bl_idname = "mocap.addmocapfix"
616     bl_label = "Add Mocap Fix to target armature"
617     type = bpy.props.EnumProperty(name="Type of Fix",
618     items=[("point", "Maintain Position", "Bone is at a specific point"),
619         ("freeze", "Maintain Position at frame", "Bone does not move from location specified in target frame"),
620         ("floor", "Stay above", "Bone does not cross specified mesh object eg floor"),
621         ("distance", "Maintain distance", "Target bones maintained specified distance")],
622     description="Type of fix")
623
624     def execute(self, context):
625         enduser_obj = bpy.context.active_object
626         enduser_arm = enduser_obj.data
627         new_mcon = enduser_arm.mocap_constraints.add()
628         new_mcon.type = self.type
629         return {"FINISHED"}
630
631     @classmethod
632     def poll(cls, context):
633         if context.active_object:
634             return isinstance(context.active_object.data, bpy.types.Armature)
635
636
637 class OBJECT_OT_RemoveMocapConstraint(bpy.types.Operator):
638     '''Remove this post-retarget fix'''
639     bl_idname = "mocap.removeconstraint"
640     bl_label = "Removes fixes from target armature"
641     constraint = bpy.props.IntProperty()
642
643     def execute(self, context):
644         enduser_obj = bpy.context.active_object
645         enduser_arm = enduser_obj.data
646         m_constraints = enduser_arm.mocap_constraints
647         m_constraint = m_constraints[self.constraint]
648         if m_constraint.real_constraint:
649             bone = enduser_obj.pose.bones[m_constraint.real_constraint_bone]
650             cons_obj = getConsObj(bone)
651             removeConstraint(m_constraint, cons_obj)
652         m_constraints.remove(self.constraint)
653         return {"FINISHED"}
654
655     @classmethod
656     def poll(cls, context):
657         if context.active_object:
658             return isinstance(context.active_object.data, bpy.types.Armature)
659
660
661 class OBJECT_OT_BakeMocapConstraints(bpy.types.Operator):
662     '''Bake all post-retarget fixes to the Retarget Fixes NLA Track'''
663     bl_idname = "mocap.bakeconstraints"
664     bl_label = "Bake all fixes to target armature"
665
666     def execute(self, context):
667         bakeConstraints(context)
668         return {"FINISHED"}
669
670     @classmethod
671     def poll(cls, context):
672         if context.active_object:
673             return isinstance(context.active_object.data, bpy.types.Armature)
674
675
676 class OBJECT_OT_UnbakeMocapConstraints(bpy.types.Operator):
677     '''Unbake all post-retarget fixes - removes the baked data from the Retarget Fixes NLA Track'''
678     bl_idname = "mocap.unbakeconstraints"
679     bl_label = "Unbake all fixes to target armature"
680
681     def execute(self, context):
682         unbakeConstraints(context)
683         return {"FINISHED"}
684
685     @classmethod
686     def poll(cls, context):
687         if context.active_object:
688             return isinstance(context.active_object.data, bpy.types.Armature)
689
690
691 class OBJECT_OT_UpdateMocapConstraints(bpy.types.Operator):
692     '''Updates all post-retarget fixes - needed after changes to armature object or pose'''
693     bl_idname = "mocap.updateconstraints"
694     bl_label = "Updates all fixes to target armature - neccesary to take under consideration changes to armature object or pose"
695
696     def execute(self, context):
697         updateConstraints(context.active_object, context)
698         return {"FINISHED"}
699
700     @classmethod
701     def poll(cls, context):
702         if context.active_object:
703             return isinstance(context.active_object.data, bpy.types.Armature)
704
705
706 class OBJECT_OT_GuessHierachyMapping(bpy.types.Operator):
707     '''Attemps to auto figure out hierarchy mapping'''
708     bl_idname = "mocap.guessmapping"
709     bl_label = "Attemps to auto figure out hierarchy mapping"
710
711     def execute(self, context):
712         enduser_obj = bpy.context.active_object
713         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
714         mocap_tools.guessMapping(performer_obj, enduser_obj)
715         return {"FINISHED"}
716
717     @classmethod
718     def poll(cls, context):
719         if context.active_object:
720             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
721         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
722         if performer_obj:
723             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
724         else:
725             return False
726
727
728 class OBJECT_OT_PathEditing(bpy.types.Operator):
729     '''Sets active object (stride object) to follow the selected curve'''
730     bl_idname = "mocap.pathediting"
731     bl_label = "Sets active object (stride object) to follow the selected curve"
732
733     def execute(self, context):
734         path = [obj for obj in context.selected_objects if obj != context.active_object][0]
735         mocap_tools.path_editing(context, context.active_object, path)
736         return {"FINISHED"}
737
738     @classmethod
739     def poll(cls, context):
740         if context.active_object:
741             selected_objs = [obj for obj in context.selected_objects if obj != context.active_object and isinstance(obj.data, bpy.types.Curve)]
742             return selected_objs
743         else:
744             return False
745
746
747 class OBJECT_OT_AnimationStitchingButton(bpy.types.Operator):
748     '''Stitches two defined animations into a single one via alignment of NLA Tracks'''
749     bl_idname = "mocap.animstitch"
750     bl_label = "Stitches two defined animations into a single one via alignment of NLA Tracks"
751
752     def execute(self, context):
753         mocap_tools.anim_stitch(context, context.active_object)
754         return {"FINISHED"}
755
756     @classmethod
757     def poll(cls, context):
758         activeIsArmature = False
759         if context.active_object:
760             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
761             if activeIsArmature:
762                 stitch_settings = context.active_object.data.stitch_settings
763                 return (stitch_settings.first_action and stitch_settings.second_action)
764         return False
765     
766
767 def register():
768     bpy.utils.register_module(__name__)
769
770
771 def unregister():
772     bpy.utils.unregister_module(__name__)
773
774 if __name__ == "__main__":
775     register()