dd5e6fa5d6d2f5ff2906f6f14d5dc6f01ce999de
[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.animstitch', text="Stitch Animations")
386
387
388 class OBJECT_OT_RetargetButton(bpy.types.Operator):
389     '''Retarget animation from selected armature to active armature '''
390     bl_idname = "mocap.retarget"
391     bl_label = "Retargets active action from Performer to Enduser"
392
393     def execute(self, context):
394         scene = context.scene
395         s_frame = scene.frame_start
396         e_frame = scene.frame_end
397         enduser_obj = context.active_object
398         performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
399         if enduser_obj is None or len(performer_obj) != 1:
400             print("Need active and selected armatures")
401         else:
402             performer_obj = performer_obj[0]
403             s_frame, e_frame = performer_obj.animation_data.action.frame_range
404             s_frame = int(s_frame)
405             e_frame = int(e_frame)
406         retarget.totalRetarget(performer_obj, enduser_obj, scene, s_frame, e_frame)
407         return {"FINISHED"}
408
409     @classmethod
410     def poll(cls, context):
411         if context.active_object:
412             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
413         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
414         if performer_obj:
415             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
416         else:
417             return False
418             
419     
420     #~ class OBJECT_OT_AdvancedRetargetButton(bpy.types.Operator):
421         #~ '''Prepare for advanced retargeting '''
422         #~ bl_idname = "mocap.preretarget"
423         #~ bl_label = "Prepares retarget of active action from Performer to Enduser"
424
425         #~ def execute(self, context):
426             #~ scene = context.scene
427             #~ s_frame = scene.frame_start
428             #~ e_frame = scene.frame_end
429             #~ enduser_obj = context.active_object
430             #~ performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
431             #~ if enduser_obj is None or len(performer_obj) != 1:
432                 #~ print("Need active and selected armatures")
433             #~ else:
434                 #~ performer_obj = performer_obj[0]
435             #~ retarget.preAdvancedRetargeting(performer_obj, enduser_obj)
436             #~ return {"FINISHED"}
437
438         #~ @classmethod
439         #~ def poll(cls, context):
440             #~ if context.active_object:
441                 #~ activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
442             #~ performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
443             #~ if performer_obj:
444                 #~ return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
445             #~ else:
446                 #~ return False
447
448
449 class OBJECT_OT_SaveMappingButton(bpy.types.Operator):
450     '''Save mapping to active armature (for future retargets) '''
451     bl_idname = "mocap.savemapping"
452     bl_label = "Saves user generated mapping from Performer to Enduser"
453
454     def execute(self, context):
455         enduser_obj = bpy.context.active_object
456         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
457         retarget.createDictionary(performer_obj.data, enduser_obj.data)
458         return {"FINISHED"}
459
460     @classmethod
461     def poll(cls, context):
462         if context.active_object:
463             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
464         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
465         if performer_obj:
466             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
467         else:
468             return False
469
470
471 class OBJECT_OT_LoadMappingButton(bpy.types.Operator):
472     '''Load saved mapping from active armature'''
473     bl_idname = "mocap.loadmapping"
474     bl_label = "Loads user generated mapping from Performer to Enduser"
475
476     def execute(self, context):
477         enduser_obj = bpy.context.active_object
478         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
479         retarget.loadMapping(performer_obj.data, enduser_obj.data)
480         return {"FINISHED"}
481
482     @classmethod
483     def poll(cls, context):
484         if context.active_object:
485             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
486         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
487         if performer_obj:
488             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
489         else:
490             return False
491
492
493 class OBJECT_OT_ConvertSamplesButton(bpy.types.Operator):
494     '''Convert active armature's sampled keyframed to beziers'''
495     bl_idname = "mocap.samples"
496     bl_label = "Converts samples / simplifies keyframes to beziers"
497
498     def execute(self, context):
499         mocap_tools.fcurves_simplify(context, context.active_object)
500         return {"FINISHED"}
501
502     @classmethod
503     def poll(cls, context):
504         return context.active_object.animation_data
505
506
507 class OBJECT_OT_LooperButton(bpy.types.Operator):
508     '''Trim active armature's animation to a single cycle, given a cyclic animation (such as a walk cycle)'''
509     bl_idname = "mocap.looper"
510     bl_label = "loops animation / sampled mocap data"
511
512     def execute(self, context):
513         mocap_tools.autoloop_anim()
514         return {"FINISHED"}
515
516     @classmethod
517     def poll(cls, context):
518         return context.active_object.animation_data
519
520
521 class OBJECT_OT_DenoiseButton(bpy.types.Operator):
522     '''Denoise active armature's animation. Good for dealing with 'bad' frames inherent in mocap animation'''
523     bl_idname = "mocap.denoise"
524     bl_label = "Denoises sampled mocap data "
525
526     def execute(self, context):
527         mocap_tools.denoise_median()
528         return {"FINISHED"}
529
530     @classmethod
531     def poll(cls, context):
532         return context.active_object
533
534     @classmethod
535     def poll(cls, context):
536         return context.active_object.animation_data
537
538
539 class OBJECT_OT_LimitDOFButton(bpy.types.Operator):
540     '''Create limit constraints on the active armature from the selected armature's animation's range of motion'''
541     bl_idname = "mocap.limitdof"
542     bl_label = "Analyzes animations Max/Min DOF and adds hard/soft constraints"
543
544     def execute(self, context):
545         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object][0]
546         mocap_tools.limit_dof(context, performer_obj, context.active_object)
547         return {"FINISHED"}
548
549     @classmethod
550     def poll(cls, context):
551         if context.active_object:
552             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
553         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
554         if performer_obj:
555             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
556         else:
557             return False
558
559
560 class OBJECT_OT_RemoveLimitDOFButton(bpy.types.Operator):
561     '''Removes previously created limit constraints on the active armature'''
562     bl_idname = "mocap.removelimitdof"
563     bl_label = "Removes previously created limit constraints on the active armature"
564
565     def execute(self, context):
566         mocap_tools.limit_dof_toggle_off(context, context.active_object)
567         return {"FINISHED"}
568
569     @classmethod
570     def poll(cls, context):
571         activeIsArmature = False
572         if context.active_object:
573             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
574         return activeIsArmature
575
576
577 class OBJECT_OT_RotateFixArmature(bpy.types.Operator):
578     '''Realign the active armature's axis system to match Blender (Commonly needed after bvh import)'''
579     bl_idname = "mocap.rotate_fix"
580     bl_label = "Rotates selected armature 90 degrees (fix for bvh import)"
581
582     def execute(self, context):
583         mocap_tools.rotate_fix_armature(context.active_object.data)
584         return {"FINISHED"}
585
586     @classmethod
587     def poll(cls, context):
588         if context.active_object:
589             return isinstance(context.active_object.data, bpy.types.Armature)
590
591
592 class OBJECT_OT_ScaleFixArmature(bpy.types.Operator):
593     '''Rescale selected armature to match the active animation, for convienence'''
594     bl_idname = "mocap.scale_fix"
595     bl_label = "Scales performer armature to match target armature"
596
597     def execute(self, context):
598         enduser_obj = bpy.context.active_object
599         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
600         mocap_tools.scale_fix_armature(performer_obj, enduser_obj)
601         return {"FINISHED"}
602
603     @classmethod
604     def poll(cls, context):
605         if context.active_object:
606             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
607         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
608         if performer_obj:
609             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
610         else:
611             return False
612
613
614 class MOCAP_OT_AddMocapFix(bpy.types.Operator):
615     '''Add a post-retarget fix - useful for fixing certain artifacts following the retarget'''
616     bl_idname = "mocap.addmocapfix"
617     bl_label = "Add Mocap Fix to target armature"
618     type = bpy.props.EnumProperty(name="Type of Fix",
619     items=[("point", "Maintain Position", "Bone is at a specific point"),
620         ("freeze", "Maintain Position at frame", "Bone does not move from location specified in target frame"),
621         ("floor", "Stay above", "Bone does not cross specified mesh object eg floor"),
622         ("distance", "Maintain distance", "Target bones maintained specified distance")],
623     description="Type of fix")
624
625     def execute(self, context):
626         enduser_obj = bpy.context.active_object
627         enduser_arm = enduser_obj.data
628         new_mcon = enduser_arm.mocap_constraints.add()
629         new_mcon.type = self.type
630         return {"FINISHED"}
631
632     @classmethod
633     def poll(cls, context):
634         if context.active_object:
635             return isinstance(context.active_object.data, bpy.types.Armature)
636
637
638 class OBJECT_OT_RemoveMocapConstraint(bpy.types.Operator):
639     '''Remove this post-retarget fix'''
640     bl_idname = "mocap.removeconstraint"
641     bl_label = "Removes fixes from target armature"
642     constraint = bpy.props.IntProperty()
643
644     def execute(self, context):
645         enduser_obj = bpy.context.active_object
646         enduser_arm = enduser_obj.data
647         m_constraints = enduser_arm.mocap_constraints
648         m_constraint = m_constraints[self.constraint]
649         if m_constraint.real_constraint:
650             bone = enduser_obj.pose.bones[m_constraint.real_constraint_bone]
651             cons_obj = getConsObj(bone)
652             removeConstraint(m_constraint, cons_obj)
653         m_constraints.remove(self.constraint)
654         return {"FINISHED"}
655
656     @classmethod
657     def poll(cls, context):
658         if context.active_object:
659             return isinstance(context.active_object.data, bpy.types.Armature)
660
661
662 class OBJECT_OT_BakeMocapConstraints(bpy.types.Operator):
663     '''Bake all post-retarget fixes to the Retarget Fixes NLA Track'''
664     bl_idname = "mocap.bakeconstraints"
665     bl_label = "Bake all fixes to target armature"
666
667     def execute(self, context):
668         bakeConstraints(context)
669         return {"FINISHED"}
670
671     @classmethod
672     def poll(cls, context):
673         if context.active_object:
674             return isinstance(context.active_object.data, bpy.types.Armature)
675
676
677 class OBJECT_OT_UnbakeMocapConstraints(bpy.types.Operator):
678     '''Unbake all post-retarget fixes - removes the baked data from the Retarget Fixes NLA Track'''
679     bl_idname = "mocap.unbakeconstraints"
680     bl_label = "Unbake all fixes to target armature"
681
682     def execute(self, context):
683         unbakeConstraints(context)
684         return {"FINISHED"}
685
686     @classmethod
687     def poll(cls, context):
688         if context.active_object:
689             return isinstance(context.active_object.data, bpy.types.Armature)
690
691
692 class OBJECT_OT_UpdateMocapConstraints(bpy.types.Operator):
693     '''Updates all post-retarget fixes - needed after changes to armature object or pose'''
694     bl_idname = "mocap.updateconstraints"
695     bl_label = "Updates all fixes to target armature - neccesary to take under consideration changes to armature object or pose"
696
697     def execute(self, context):
698         updateConstraints(context.active_object, context)
699         return {"FINISHED"}
700
701     @classmethod
702     def poll(cls, context):
703         if context.active_object:
704             return isinstance(context.active_object.data, bpy.types.Armature)
705
706
707 class OBJECT_OT_GuessHierachyMapping(bpy.types.Operator):
708     '''Attemps to auto figure out hierarchy mapping'''
709     bl_idname = "mocap.guessmapping"
710     bl_label = "Attemps to auto figure out hierarchy mapping"
711
712     def execute(self, context):
713         enduser_obj = bpy.context.active_object
714         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
715         mocap_tools.guessMapping(performer_obj, enduser_obj)
716         return {"FINISHED"}
717
718     @classmethod
719     def poll(cls, context):
720         if context.active_object:
721             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
722         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
723         if performer_obj:
724             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
725         else:
726             return False
727
728
729 class OBJECT_OT_PathEditing(bpy.types.Operator):
730     '''Sets active object (stride object) to follow the selected curve'''
731     bl_idname = "mocap.pathediting"
732     bl_label = "Sets active object (stride object) to follow the selected curve"
733
734     def execute(self, context):
735         path = [obj for obj in context.selected_objects if obj != context.active_object][0]
736         mocap_tools.path_editing(context, context.active_object, path)
737         return {"FINISHED"}
738
739     @classmethod
740     def poll(cls, context):
741         if context.active_object:
742             selected_objs = [obj for obj in context.selected_objects if obj != context.active_object and isinstance(obj.data, bpy.types.Curve)]
743             return selected_objs
744         else:
745             return False
746
747
748 class OBJECT_OT_AnimationStitchingButton(bpy.types.Operator):
749     '''Stitches two defined animations into a single one via alignment of NLA Tracks'''
750     bl_idname = "mocap.animstitch"
751     bl_label = "Stitches two defined animations into a single one via alignment of NLA Tracks"
752
753     def execute(self, context):
754         mocap_tools.anim_stitch(context, context.active_object)
755         return {"FINISHED"}
756
757     @classmethod
758     def poll(cls, context):
759         activeIsArmature = False
760         if context.active_object:
761             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
762             if activeIsArmature:
763                 stitch_settings = context.active_object.data.stitch_settings
764                 return (stitch_settings.first_action and stitch_settings.second_action)
765         return False
766     
767
768 def register():
769     bpy.utils.register_module(__name__)
770
771
772 def unregister():
773     bpy.utils.unregister_module(__name__)
774
775 if __name__ == "__main__":
776     register()