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