Bugfixing for retargeting - unconnected bones now retarget alot better. Also some...
[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="Offset",
85         default=0.0,
86         description="Distance and Floor Constraints - Desired offset",
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         row.operator("mocap.scale_fix", text='Auto scale Performer')
195         row2 = self.layout.row(align=True)
196         row2.operator("mocap.looper", text='Loop animation')
197         row2.operator("mocap.limitdof", text='Constrain Rig')
198         self.layout.label("Retargeting")
199         row3 = self.layout.row(align=True)
200         column1 = row3.column(align=True)
201         column1.label("Performer Rig")
202         column2 = row3.column(align=True)
203         column2.label("Enduser Rig")
204         enduser_obj = bpy.context.active_object
205         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj]
206         if enduser_obj is None or len(performer_obj) != 1:
207             self.layout.label("Select performer rig and target rig (as active)")
208         else:
209             performer_obj = performer_obj[0]
210             if performer_obj.data and enduser_obj.data:
211                 if performer_obj.data.name in bpy.data.armatures and enduser_obj.data.name in bpy.data.armatures:
212                     perf = performer_obj.data
213                     enduser_arm = enduser_obj.data
214                     perf_pose_bones = enduser_obj.pose.bones
215                     for bone in perf.bones:
216                         row = self.layout.row()
217                         row.prop(data=bone, property='foot', text='', icon='POSE_DATA')
218                         row.label(bone.name)
219                         row.prop_search(bone, "map", enduser_arm, "bones")
220                         label_mod = "FK"
221                         if bone.map:
222                             pose_bone = perf_pose_bones[bone.map]
223                             if pose_bone.is_in_ik_chain:
224                                 label_mod = "ik chain"
225                             if hasIKConstraint(pose_bone):
226                                 label_mod = "ik end"
227                             row.prop(pose_bone, 'IKRetarget')
228                             row.label(label_mod)
229                         else:
230                             row.label(" ")
231                             row.label(" ")
232                     mapRow = self.layout.row()
233                     mapRow.operator("mocap.savemapping", text='Save mapping')
234                     mapRow.operator("mocap.loadmapping", text='Load mapping')
235                     self.layout.operator("mocap.retarget", text='RETARGET!')
236
237
238 class MocapConstraintsPanel(bpy.types.Panel):
239     #Motion capture constraints panel
240     bl_label = "Mocap constraints"
241     bl_space_type = "PROPERTIES"
242     bl_region_type = "WINDOW"
243     bl_context = "object"
244
245     def draw(self, context):
246         layout = self.layout
247         if context.active_object:
248             if context.active_object.data:
249                 if context.active_object.data.name in bpy.data.armatures:
250                     enduser_obj = context.active_object
251                     enduser_arm = enduser_obj.data
252                     layout.operator("mocap.addconstraint")
253                     layout.operator("mocap.bakeconstraints")
254                     layout.operator("mocap.unbakeconstraints")
255                     layout.separator()
256                     for i, m_constraint in enumerate(enduser_arm.mocap_constraints):
257                         box = layout.box()
258                         box.prop(m_constraint, 'name')
259                         box.prop(m_constraint, 'type')
260                         box.prop_search(m_constraint, 'constrained_bone', enduser_obj.pose, "bones")
261                         if m_constraint.type == "distance" or m_constraint.type == "point":
262                             box.prop_search(m_constraint, 'constrained_boneB', enduser_obj.pose, "bones")
263                         frameRow = box.row()
264                         frameRow.label("Frame Range:")
265                         frameRow.prop(m_constraint, 's_frame')
266                         frameRow.prop(m_constraint, 'e_frame')
267                         smoothRow = box.row()
268                         smoothRow.label("Smoothing:")
269                         smoothRow.prop(m_constraint, 'smooth_in')
270                         smoothRow.prop(m_constraint, 'smooth_out')
271                         targetRow = box.row()
272                         targetLabelCol = targetRow.column()
273                         targetLabelCol.label("Target settings:")
274                         targetPropCol = targetRow.column()
275                         if m_constraint.type == "floor":
276                             targetPropCol.prop_search(m_constraint, 'targetMesh', bpy.data, "objects")
277                         if m_constraint.type == "point" or m_constraint.type == "freeze":
278                             box.prop(m_constraint, 'targetSpace')
279                         if m_constraint.type == "point":
280                             targetPropCol.prop(m_constraint, 'targetPoint')
281                         if m_constraint.type == "distance" or m_constraint.type == "floor":
282                             targetPropCol.prop(m_constraint, 'targetDist')
283                         checkRow = box.row()
284                         checkRow.prop(m_constraint, 'active')
285                         layout.operator("mocap.removeconstraint", text="Remove constraint").constraint = i
286                         layout.separator()
287
288
289 class OBJECT_OT_RetargetButton(bpy.types.Operator):
290     bl_idname = "mocap.retarget"
291     bl_label = "Retargets active action from Performer to Enduser"
292
293     def execute(self, context):
294         enduser_obj = context.active_object
295         performer_obj = [obj for obj in context.selected_objects if obj != enduser_obj]
296         if enduser_obj is None or len(performer_obj) != 1:
297             print("Need active and selected armatures")
298         else:
299             performer_obj = performer_obj[0]
300         scene = context.scene
301         s_frame = scene.frame_start
302         e_frame = scene.frame_end
303         retarget.totalRetarget(performer_obj, enduser_obj, scene, s_frame, e_frame)
304         return {"FINISHED"}
305
306     @classmethod
307     def poll(cls, context):
308         if context.active_object:
309             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
310         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
311         if performer_obj:
312             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
313         else:
314             return False
315
316
317 class OBJECT_OT_SaveMappingButton(bpy.types.Operator):
318     bl_idname = "mocap.savemapping"
319     bl_label = "Saves user generated mapping from Performer to Enduser"
320
321     def execute(self, context):
322         enduser_obj = bpy.context.active_object
323         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
324         retarget.createDictionary(performer_obj.data, enduser_obj.data)
325         return {"FINISHED"}
326
327     @classmethod
328     def poll(cls, context):
329         if context.active_object:
330             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
331         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
332         if performer_obj:
333             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
334         else:
335             return False
336
337
338 class OBJECT_OT_LoadMappingButton(bpy.types.Operator):
339     bl_idname = "mocap.loadmapping"
340     bl_label = "Loads user generated mapping from Performer to Enduser"
341
342     def execute(self, context):
343         enduser_obj = bpy.context.active_object
344         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
345         retarget.loadMapping(performer_obj.data, enduser_obj.data)
346         return {"FINISHED"}
347
348     @classmethod
349     def poll(cls, context):
350         if context.active_object:
351             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
352         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
353         if performer_obj:
354             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
355         else:
356             return False
357
358
359 class OBJECT_OT_ConvertSamplesButton(bpy.types.Operator):
360     bl_idname = "mocap.samples"
361     bl_label = "Converts samples / simplifies keyframes to beziers"
362
363     def execute(self, context):
364         mocap_tools.fcurves_simplify()
365         return {"FINISHED"}
366
367     @classmethod
368     def poll(cls, context):
369         return context.active_object.animation_data
370
371
372 class OBJECT_OT_LooperButton(bpy.types.Operator):
373     bl_idname = "mocap.looper"
374     bl_label = "loops animation / sampled mocap data"
375
376     def execute(self, context):
377         mocap_tools.autoloop_anim()
378         return {"FINISHED"}
379
380     @classmethod
381     def poll(cls, context):
382         return context.active_object.animation_data
383
384
385 class OBJECT_OT_DenoiseButton(bpy.types.Operator):
386     bl_idname = "mocap.denoise"
387     bl_label = "Denoises sampled mocap data "
388
389     def execute(self, context):
390         mocap_tools.denoise_median()
391         return {"FINISHED"}
392
393     @classmethod
394     def poll(cls, context):
395         return context.active_object
396
397     @classmethod
398     def poll(cls, context):
399         return context.active_object.animation_data
400
401
402 class OBJECT_OT_LimitDOFButton(bpy.types.Operator):
403     bl_idname = "mocap.limitdof"
404     bl_label = "Analyzes animations Max/Min DOF and adds hard/soft constraints"
405
406     def execute(self, context):
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_RotateFixArmature(bpy.types.Operator):
421     bl_idname = "mocap.rotate_fix"
422     bl_label = "Rotates selected armature 90 degrees (fix for bvh import)"
423
424     def execute(self, context):
425         mocap_tools.rotate_fix_armature(context.active_object.data)
426         return {"FINISHED"}
427
428     @classmethod
429     def poll(cls, context):
430         if context.active_object:
431             return isinstance(context.active_object.data, bpy.types.Armature)
432
433
434 class OBJECT_OT_ScaleFixArmature(bpy.types.Operator):
435     bl_idname = "mocap.scale_fix"
436     bl_label = "Scales performer armature to match target armature"
437
438     def execute(self, context):
439         enduser_obj = bpy.context.active_object
440         performer_obj = [obj for obj in bpy.context.selected_objects if obj != enduser_obj][0]
441         mocap_tools.scale_fix_armature(performer_obj, enduser_obj)
442         return {"FINISHED"}
443
444     @classmethod
445     def poll(cls, context):
446         if context.active_object:
447             activeIsArmature = isinstance(context.active_object.data, bpy.types.Armature)
448         performer_obj = [obj for obj in context.selected_objects if obj != context.active_object]
449         if performer_obj:
450             return activeIsArmature and isinstance(performer_obj[0].data, bpy.types.Armature)
451         else:
452             return False
453
454
455 class OBJECT_OT_AddMocapConstraint(bpy.types.Operator):
456     bl_idname = "mocap.addconstraint"
457     bl_label = "Add constraint to target armature"
458
459     def execute(self, context):
460         enduser_obj = bpy.context.active_object
461         enduser_arm = enduser_obj.data
462         new_mcon = enduser_arm.mocap_constraints.add()
463         return {"FINISHED"}
464
465     @classmethod
466     def poll(cls, context):
467         if context.active_object:
468             return isinstance(context.active_object.data, bpy.types.Armature)
469
470
471 class OBJECT_OT_RemoveMocapConstraint(bpy.types.Operator):
472     bl_idname = "mocap.removeconstraint"
473     bl_label = "Removes constraints from target armature"
474     constraint = bpy.props.IntProperty()
475
476     def execute(self, context):
477         enduser_obj = bpy.context.active_object
478         enduser_arm = enduser_obj.data
479         m_constraints = enduser_arm.mocap_constraints
480         m_constraint = m_constraints[self.constraint]
481         if m_constraint.real_constraint:
482             bone = enduser_obj.pose.bones[m_constraint.real_constraint_bone]
483             cons_obj = getConsObj(bone)
484             removeConstraint(m_constraint, cons_obj)
485         m_constraints.remove(self.constraint)
486         return {"FINISHED"}
487
488     @classmethod
489     def poll(cls, context):
490         if context.active_object:
491             return isinstance(context.active_object.data, bpy.types.Armature)
492
493
494 class OBJECT_OT_BakeMocapConstraints(bpy.types.Operator):
495     bl_idname = "mocap.bakeconstraints"
496     bl_label = "Bake all constraints to target armature"
497
498     def execute(self, context):
499         bakeConstraints(context)
500         return {"FINISHED"}
501
502     @classmethod
503     def poll(cls, context):
504         if context.active_object:
505             return isinstance(context.active_object.data, bpy.types.Armature)
506
507
508 class OBJECT_OT_UnbakeMocapConstraints(bpy.types.Operator):
509     bl_idname = "mocap.unbakeconstraints"
510     bl_label = "Unbake all constraints to target armature"
511
512     def execute(self, context):
513         unbakeConstraints(context)
514         return {"FINISHED"}
515
516     @classmethod
517     def poll(cls, context):
518         if context.active_object:
519             return isinstance(context.active_object.data, bpy.types.Armature)
520
521
522 def register():
523     bpy.utils.register_module(__name__)
524
525
526 def unregister():
527     bpy.utils.unregister_module(__name__)
528
529 if __name__ == "__main__":
530     register()