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