Bugfix: Baking mocap constraints now works for user created IK bones
[blender.git] / release / scripts / startup / ui_mocap.py
1 # ##### BEGIN GPL LICENSE BLOCK #####
2 #
3 #  This program is free software; you can redistribute it and/or
4 #  modify it under the terms of the GNU General Public License
5 #  as published by the Free Software Foundation; either version 2
6 #  of the License, or (at your option) any later version.
7 #
8 #  This program is distributed in the hope that it will be useful,
9 #  but WITHOUT ANY WARRANTY; without even the implied warranty of
10 #  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
11 #  GNU General Public License for more details.
12 #
13 #  You should have received a copy of the GNU General Public License
14 #  along with this program; if not, write to the Free Software Foundation,
15 #  Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 # <pep8 compliant>
20
21 import bpy
22
23 from bpy.props import *
24 from bpy import *
25 import mocap_constraints
26 import retarget
27 import mocap_tools
28 ### reloads modules (for testing purposes only)
29 from imp import reload
30 reload(mocap_constraints)
31 reload(retarget)
32 reload(mocap_tools)
33
34 from mocap_constraints import *
35
36 # MocapConstraint class
37 # Defines MocapConstraint datatype, used to add and configute mocap constraints
38 # Attached to Armature data
39
40
41 class MocapConstraint(bpy.types.PropertyGroup):
42     name = bpy.props.StringProperty(name="Name",
43         default="Mocap Constraint",
44         description="Name of Mocap Constraint",
45         update=setConstraint)
46     constrained_bone = bpy.props.StringProperty(name="Bone",
47         default="",
48         description="Constrained Bone",
49         update=updateConstraintBoneType)
50     constrained_boneB = bpy.props.StringProperty(name="Bone (2)",
51         default="",
52         description="Other Constrained Bone (optional, depends on type)",
53         update=setConstraint)
54     s_frame = bpy.props.IntProperty(name="S",
55         default=1,
56         description="Start frame of constraint",
57         update=setConstraint)
58     e_frame = bpy.props.IntProperty(name="E",
59         default=500,
60         description="End frame of constrain",
61         update=setConstraint)
62     smooth_in = bpy.props.IntProperty(name="In",
63         default=10,
64         description="Amount of frames to smooth in",
65         update=setConstraint,
66         min=0)
67     smooth_out = bpy.props.IntProperty(name="Out",
68         default=10,
69         description="Amount of frames to smooth out",
70         update=setConstraint,
71         min=0)
72     targetMesh = bpy.props.StringProperty(name="Mesh",
73         default="",
74         description="Target of Constraint - Mesh (optional, depends on type)",
75         update=setConstraint)
76     active = bpy.props.BoolProperty(name="Active",
77         default=True,
78         description="Constraint is active",
79         update=setConstraint)
80     targetPoint = bpy.props.FloatVectorProperty(name="Point", size=3,
81         subtype="XYZ", default=(0.0, 0.0, 0.0),
82         description="Target of Constraint - Point",
83         update=setConstraint)
84     targetDist = bpy.props.FloatProperty(name="Dist",
85         default=1,
86         description="Distance Constraint - Desired distance",
87         update=setConstraint)
88     targetSpace = bpy.props.EnumProperty(
89         items=[("WORLD", "World Space", "Evaluate target in global space"),
90             ("LOCAL", "Object space", "Evaluate target in object space"),
91             ("constrained_boneB", "Other Bone Space", "Evaluate target in specified other bone space")],
92         name="Space",
93         description="In which space should Point type target be evaluated",
94         update=setConstraint)
95     type = bpy.props.EnumProperty(name="Type of constraint",
96         items=[("point", "Maintain Position", "Bone is at a specific point"),
97             ("freeze", "Maintain Position at frame", "Bone does not move from location specified in target frame"),
98             ("floor", "Stay above", "Bone does not cross specified mesh object eg floor"),
99             ("distance", "Maintain distance", "Target bones maintained specified distance")],
100         description="Type of constraint",
101         update=updateConstraintBoneType)
102     real_constraint = bpy.props.StringProperty()
103     real_constraint_bone = bpy.props.StringProperty()
104
105
106 bpy.utils.register_class(MocapConstraint)
107
108 bpy.types.Armature.mocap_constraints = bpy.props.CollectionProperty(type=MocapConstraint)
109
110 #Update function for IK functionality. Is called when IK prop checkboxes are toggled.
111
112
113 def toggleIKBone(self, context):
114     if self.IKRetarget:
115         if not self.is_in_ik_chain:
116             print(self.name + " IK toggled ON!")
117             ik = self.constraints.new('IK')
118             #ik the whole chain up to the root, excluding
119             chainLen = 0
120             for parent_bone in self.parent_recursive:
121                 chainLen += 1
122                 if hasIKConstraint(parent_bone):
123                     break
124                 deformer_children = [child for child in parent_bone.children if child.bone.use_deform]
125                 if len(deformer_children) > 1:
126                     break
127             ik.chain_count = chainLen
128             for bone in self.parent_recursive:
129                 if bone.is_in_ik_chain:
130                     bone.IKRetarget = True
131     else:
132         print(self.name + " IK toggled OFF!")
133         cnstrn_bone = False
134         if hasIKConstraint(self):
135             cnstrn_bone = self
136         elif self.is_in_ik_chain:
137             cnstrn_bone = [child for child in self.children_recursive if hasIKConstraint(child)][0]
138         if cnstrn_bone:
139             # remove constraint, and update IK retarget for all parents of cnstrn_bone up to chain_len
140             ik = [constraint for constraint in cnstrn_bone.constraints if constraint.type == "IK"][0]
141             cnstrn_bone.constraints.remove(ik)
142             cnstrn_bone.IKRetarget = False
143             for bone in cnstrn_bone.parent_recursive:
144                 if not bone.is_in_ik_chain:
145                     bone.IKRetarget = False
146
147
148 class MocapMapping(bpy.types.PropertyGroup):
149     name = bpy.props.StringProperty()
150
151 bpy.utils.register_class(MocapMapping)
152
153 bpy.types.Bone.map = bpy.props.StringProperty()
154 bpy.types.Bone.reverseMap = bpy.props.CollectionProperty(type=MocapMapping)
155 bpy.types.Bone.foot = bpy.props.BoolProperty(name="Foot",
156     description="Marks this bone as a 'foot', which determines retargeted animation's translation",
157     default=False)
158 bpy.types.PoseBone.IKRetarget = bpy.props.BoolProperty(name="IK",
159     description="Toggles IK Retargeting method for given bone",
160     update=toggleIKBone, default=False)
161
162
163 def updateIKRetarget():
164     # ensures that Blender constraints and IK properties are in sync
165     # currently runs when module is loaded, should run when scene is loaded
166     # or user adds a constraint to armature. Will be corrected in the future,
167     # once python callbacks are implemented
168     for obj in bpy.data.objects:
169         if obj.pose:
170             bones = obj.pose.bones
171             for pose_bone in bones:
172                 if pose_bone.is_in_ik_chain or hasIKConstraint(pose_bone):
173                     pose_bone.IKRetarget = True
174                 else:
175                     pose_bone.IKRetarget = False
176
177 updateIKRetarget()
178
179
180 class MocapPanel(bpy.types.Panel):
181     # Motion capture retargeting panel
182     bl_label = "Mocap tools"
183     bl_space_type = "PROPERTIES"
184     bl_region_type = "WINDOW"
185     bl_context = "object"
186
187     def draw(self, context):
188         self.layout.label("Preprocessing")
189         row = self.layout.row(align=True)
190         row.alignment = 'EXPAND'
191         row.operator("mocap.samples", text='Samples to Beziers')
192         row.operator("mocap.denoise", text='Clean noise')
193         row.operator("mocap.rotate_fix", text='Fix BVH Axis Orientation')
194         row2 = self.layout.row(align=True)
195         row2.operator("mocap.looper", text='Loop animation')
196         row2.operator("mocap.limitdof", text='Constrain Rig')
197         self.layout.label("Retargeting")
198         row3 = self.layout.row(align=True)
199         column1 = row3.column(align=True)
200         column1.label("Performer Rig")
201         column2 = row3.column(align=True)
202         column2.label("Enduser Rig")
203         enduser_obj = bpy.context.active_object
204         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj]
205         if enduser_obj is None or len(performer_obj) != 1:
206             self.layout.label("Select performer rig and target rig (as active)")
207         else:
208             performer_obj = performer_obj[0]
209             if performer_obj.data and enduser_obj.data:
210                 if performer_obj.data.name in bpy.data.armatures and enduser_obj.data.name in bpy.data.armatures:
211                     perf = performer_obj.data
212                     enduser_arm = enduser_obj.data
213                     perf_pose_bones = enduser_obj.pose.bones
214                     for bone in perf.bones:
215                         row = self.layout.row()
216                         row.prop(data=bone, property='foot', text='', icon='POSE_DATA')
217                         row.label(bone.name)
218                         row.prop_search(bone, "map", enduser_arm, "bones")
219                         label_mod = "FK"
220                         if bone.map:
221                             pose_bone = perf_pose_bones[bone.map]
222                             if pose_bone.is_in_ik_chain:
223                                 label_mod = "ik chain"
224                             if hasIKConstraint(pose_bone):
225                                 label_mod = "ik end"
226                             row.prop(pose_bone, 'IKRetarget')
227                             row.label(label_mod)
228                         else:
229                             row.label(" ")
230                             row.label(" ")
231                     mapRow = self.layout.row()
232                     mapRow.operator("mocap.savemapping", text='Save mapping')
233                     mapRow.operator("mocap.loadmapping", text='Load mapping')
234                     self.layout.operator("mocap.retarget", text='RETARGET!')
235
236
237 class MocapConstraintsPanel(bpy.types.Panel):
238     #Motion capture constraints panel
239     bl_label = "Mocap constraints"
240     bl_space_type = "PROPERTIES"
241     bl_region_type = "WINDOW"
242     bl_context = "object"
243
244     def draw(self, context):
245         layout = self.layout
246         if context.active_object:
247             if context.active_object.data:
248                 if context.active_object.data.name in bpy.data.armatures:
249                     enduser_obj = context.active_object
250                     enduser_arm = enduser_obj.data
251                     layout.operator("mocap.addconstraint")
252                     layout.operator("mocap.bakeconstraints")
253                     layout.operator("mocap.unbakeconstraints")
254                     layout.separator()
255                     for i, m_constraint in enumerate(enduser_arm.mocap_constraints):
256                         box = layout.box()
257                         box.prop(m_constraint, 'name')
258                         box.prop(m_constraint, 'type')
259                         box.prop_search(m_constraint, 'constrained_bone', enduser_obj.pose, "bones")
260                         if m_constraint.type == "distance" or m_constraint.type == "point":
261                             box.prop_search(m_constraint, 'constrained_boneB', enduser_obj.pose, "bones")
262                         frameRow = box.row()
263                         frameRow.label("Frame Range:")
264                         frameRow.prop(m_constraint, 's_frame')
265                         frameRow.prop(m_constraint, 'e_frame')
266                         smoothRow = box.row()
267                         smoothRow.label("Smoothing:")
268                         smoothRow.prop(m_constraint, 'smooth_in')
269                         smoothRow.prop(m_constraint, 'smooth_out')
270                         targetRow = box.row()
271                         targetLabelCol = targetRow.column()
272                         targetLabelCol.label("Target settings:")
273                         targetPropCol = targetRow.column()
274                         if m_constraint.type == "floor":
275                             targetPropCol.prop_search(m_constraint, 'targetMesh', bpy.data, "objects")
276                         if m_constraint.type == "point" or m_constraint.type == "freeze":
277                             box.prop(m_constraint, 'targetSpace')
278                         if m_constraint.type == "point":
279                             targetPropCol.prop(m_constraint, 'targetPoint')
280                         if m_constraint.type == "distance":
281                             targetPropCol.prop(m_constraint, 'targetDist')
282                         checkRow = box.row()
283                         checkRow.prop(m_constraint, 'active')
284                         layout.operator("mocap.removeconstraint", text="Remove constraint").constraint = i
285                         layout.separator()
286
287
288 class OBJECT_OT_RetargetButton(bpy.types.Operator):
289     bl_idname = "mocap.retarget"
290     bl_label = "Retargets active action from Performer to Enduser"
291
292     def execute(self, context):
293         enduser_obj = context.active_object
294         performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
295         if enduser_obj is None or len(performer_obj) != 1:
296             print("Need active and selected armatures")
297         else:
298             performer_obj = performer_obj[0]
299         scene = context.scene
300         s_frame = scene.frame_start
301         e_frame = scene.frame_end
302         retarget.totalRetarget(performer_obj, enduser_obj, scene, s_frame, e_frame)
303         return {"FINISHED"}
304
305     @classmethod
306     def poll(cls, context):
307         if context.active_object:
308             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
309         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
310         if performer_obj:
311             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
312         else:
313             return False
314
315
316 class OBJECT_OT_SaveMappingButton(bpy.types.Operator):
317     bl_idname = "mocap.savemapping"
318     bl_label = "Saves user generated mapping from Performer to Enduser"
319
320     def execute(self, context):
321         enduser_obj = bpy.context.active_object
322         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
323         retarget.createDictionary(performer_obj.data, enduser_obj.data)
324         return {"FINISHED"}
325
326     @classmethod
327     def poll(cls, context):
328         if context.active_object:
329             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
330         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
331         if performer_obj:
332             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
333         else:
334             return False
335
336
337 class OBJECT_OT_LoadMappingButton(bpy.types.Operator):
338     bl_idname = "mocap.loadmapping"
339     bl_label = "Loads user generated mapping from Performer to Enduser"
340
341     def execute(self, context):
342         enduser_obj = bpy.context.active_object
343         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
344         retarget.loadMapping(performer_obj.data, enduser_obj.data)
345         return {"FINISHED"}
346
347     @classmethod
348     def poll(cls, context):
349         if context.active_object:
350             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
351         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
352         if performer_obj:
353             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
354         else:
355             return False
356
357
358 class OBJECT_OT_ConvertSamplesButton(bpy.types.Operator):
359     bl_idname = "mocap.samples"
360     bl_label = "Converts samples / simplifies keyframes to beziers"
361
362     def execute(self, context):
363         mocap_tools.fcurves_simplify()
364         return {"FINISHED"}
365
366     @classmethod
367     def poll(cls, context):
368         return context.active_object.animation_data
369
370
371 class OBJECT_OT_LooperButton(bpy.types.Operator):
372     bl_idname = "mocap.looper"
373     bl_label = "loops animation / sampled mocap data"
374
375     def execute(self, context):
376         mocap_tools.autoloop_anim()
377         return {"FINISHED"}
378
379     @classmethod
380     def poll(cls, context):
381         return context.active_object.animation_data
382
383
384 class OBJECT_OT_DenoiseButton(bpy.types.Operator):
385     bl_idname = "mocap.denoise"
386     bl_label = "Denoises sampled mocap data "
387
388     def execute(self, context):
389         mocap_tools.denoise_median()
390         return {"FINISHED"}
391
392     @classmethod
393     def poll(cls, context):
394         return context.active_object
395
396     @classmethod
397     def poll(cls, context):
398         return context.active_object.animation_data
399
400
401 class OBJECT_OT_LimitDOFButton(bpy.types.Operator):
402     bl_idname = "mocap.limitdof"
403     bl_label = "Analyzes animations Max/Min DOF and adds hard/soft constraints"
404
405     def execute(self, context):
406         return {"FINISHED"}
407
408     @classmethod
409     def poll(cls, context):
410         if context.active_object:
411             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
412         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
413         if performer_obj:
414             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
415         else:
416             return False
417
418
419 class OBJECT_OT_RotateFixArmature(bpy.types.Operator):
420     bl_idname = "mocap.rotate_fix"
421     bl_label = "Rotates selected armature 90 degrees (fix for bvh import)"
422
423     def execute(self, context):
424         mocap_tools.rotate_fix_armature(context.active_object.data)
425         return {"FINISHED"}
426
427     @classmethod
428     def poll(cls, context):
429         if context.active_object:
430             return isinstance(context.active_object.data, bpy.types.Armature)
431
432
433 class OBJECT_OT_AddMocapConstraint(bpy.types.Operator):
434     bl_idname = "mocap.addconstraint"
435     bl_label = "Add constraint to target armature"
436
437     def execute(self, context):
438         enduser_obj = bpy.context.active_object
439         enduser_arm = enduser_obj.data
440         new_mcon = enduser_arm.mocap_constraints.add()
441         return {"FINISHED"}
442
443     @classmethod
444     def poll(cls, context):
445         if context.active_object:
446             return isinstance(context.active_object.data, bpy.types.Armature)
447
448
449 class OBJECT_OT_RemoveMocapConstraint(bpy.types.Operator):
450     bl_idname = "mocap.removeconstraint"
451     bl_label = "Removes constraints from target armature"
452     constraint = bpy.props.IntProperty()
453
454     def execute(self, context):
455         enduser_obj = bpy.context.active_object
456         enduser_arm = enduser_obj.data
457         m_constraints = enduser_arm.mocap_constraints
458         m_constraint = m_constraints[self.constraint]
459         if m_constraint.real_constraint:
460             bone = enduser_obj.pose.bones[m_constraint.real_constraint_bone]
461             cons_obj = getConsObj(bone)
462             removeConstraint(m_constraint, cons_obj)
463         m_constraints.remove(self.constraint)
464         return {"FINISHED"}
465
466     @classmethod
467     def poll(cls, context):
468         if context.active_object:
469             return isinstance(context.active_object.data, bpy.types.Armature)
470
471
472 class OBJECT_OT_BakeMocapConstraints(bpy.types.Operator):
473     bl_idname = "mocap.bakeconstraints"
474     bl_label = "Bake all constraints to target armature"
475
476     def execute(self, context):
477         bakeConstraints(context)
478         return {"FINISHED"}
479
480     @classmethod
481     def poll(cls, context):
482         if context.active_object:
483             return isinstance(context.active_object.data, bpy.types.Armature)
484
485
486 class OBJECT_OT_UnbakeMocapConstraints(bpy.types.Operator):
487     bl_idname = "mocap.unbakeconstraints"
488     bl_label = "Unbake all constraints to target armature"
489
490     def execute(self, context):
491         unbakeConstraints(context)
492         return {"FINISHED"}
493
494     @classmethod
495     def poll(cls, context):
496         if context.active_object:
497             return isinstance(context.active_object.data, bpy.types.Armature)
498
499
500 def register():
501     bpy.utils.register_module(__name__)
502
503
504 def unregister():
505     bpy.utils.unregister_module(__name__)
506
507 if __name__ == "__main__":
508     register()