rigify/rig_ui_template.py

   1 # SPDX-License-Identifier: GPL-2.0-or-later
   2
   3 import bpy
   4
   5 from collections import OrderedDict
   6
   7 from .utils.animation import SCRIPT_REGISTER_BAKE, SCRIPT_UTILITIES_BAKE
   8
   9 from . import base_generate
  10
  11 from rna_prop_ui import rna_idprop_quote_path
  12
  13
  14 UI_IMPORTS = [
  15     'import bpy',
  16     'import math',
  17     'import json',
  18     'import collections',
  19     'import traceback',
  20     'from math import pi',
  21     'from bpy.props import StringProperty',
  22     'from mathutils import Euler, Matrix, Quaternion, Vector',
  23     'from rna_prop_ui import rna_idprop_quote_path',
  24 ]
  25
  26 UI_BASE_UTILITIES = '''
  27 rig_id = "%s"
  28
  29
  30 ############################
  31 ## Math utility functions ##
  32 ############################
  33
  34 def perpendicular_vector(v):
  35     """ Returns a vector that is perpendicular to the one given.
  36         The returned vector is _not_ guaranteed to be normalized.
  37     """
  38     # Create a vector that is not aligned with v.
  39     # It doesn't matter what vector.  Just any vector
  40     # that's guaranteed to not be pointing in the same
  41     # direction.
  42     if abs(v[0]) < abs(v[1]):
  43         tv = Vector((1,0,0))
  44     else:
  45         tv = Vector((0,1,0))
  46
  47     # Use cross prouct to generate a vector perpendicular to
  48     # both tv and (more importantly) v.
  49     return v.cross(tv)
  50
  51
  52 def rotation_difference(mat1, mat2):
  53     """ Returns the shortest-path rotational difference between two
  54         matrices.
  55     """
  56     q1 = mat1.to_quaternion()
  57     q2 = mat2.to_quaternion()
  58     angle = math.acos(min(1,max(-1,q1.dot(q2)))) * 2
  59     if angle > pi:
  60         angle = -angle + (2*pi)
  61     return angle
  62
  63 def find_min_range(f,start_angle,delta=pi/8):
  64     """ finds the range where lies the minimum of function f applied on bone_ik and bone_fk
  65         at a certain angle.
  66     """
  67     angle = start_angle
  68     while (angle > (start_angle - 2*pi)) and (angle < (start_angle + 2*pi)):
  69         l_dist = f(angle-delta)
  70         c_dist = f(angle)
  71         r_dist = f(angle+delta)
  72         if min((l_dist,c_dist,r_dist)) == c_dist:
  73             return (angle-delta,angle+delta)
  74         else:
  75             angle=angle+delta
  76
  77 def ternarySearch(f, left, right, absolutePrecision):
  78     """
  79     Find minimum of unimodal function f() within [left, right]
  80     To find the maximum, revert the if/else statement or revert the comparison.
  81     """
  82     while True:
  83         #left and right are the current bounds; the maximum is between them
  84         if abs(right - left) < absolutePrecision:
  85             return (left + right)/2
  86
  87         leftThird = left + (right - left)/3
  88         rightThird = right - (right - left)/3
  89
  90         if f(leftThird) > f(rightThird):
  91             left = leftThird
  92         else:
  93             right = rightThird
  94 '''
  95
  96 UTILITIES_FUNC_COMMON_IKFK = ['''
  97 #########################################
  98 ## "Visual Transform" helper functions ##
  99 #########################################
 100
 101 def get_pose_matrix_in_other_space(mat, pose_bone):
 102     """ Returns the transform matrix relative to pose_bone's current
 103         transform space.  In other words, presuming that mat is in
 104         armature space, slapping the returned matrix onto pose_bone
 105         should give it the armature-space transforms of mat.
 106     """
 107     return pose_bone.id_data.convert_space(matrix=mat, pose_bone=pose_bone, from_space='POSE', to_space='LOCAL')
 108
 109
 110 def convert_pose_matrix_via_rest_delta(mat, from_bone, to_bone):
 111     """Convert pose of one bone to another bone, preserving the rest pose difference between them."""
 112     return mat @ from_bone.bone.matrix_local.inverted() @ to_bone.bone.matrix_local
 113
 114
 115 def convert_pose_matrix_via_pose_delta(mat, from_bone, to_bone):
 116     """Convert pose of one bone to another bone, preserving the current pose difference between them."""
 117     return mat @ from_bone.matrix.inverted() @ to_bone.matrix
 118
 119
 120 def get_local_pose_matrix(pose_bone):
 121     """ Returns the local transform matrix of the given pose bone.
 122     """
 123     return get_pose_matrix_in_other_space(pose_bone.matrix, pose_bone)
 124
 125
 126 def set_pose_translation(pose_bone, mat):
 127     """ Sets the pose bone's translation to the same translation as the given matrix.
 128         Matrix should be given in bone's local space.
 129     """
 130     pose_bone.location = mat.to_translation()
 131
 132
 133 def set_pose_rotation(pose_bone, mat):
 134     """ Sets the pose bone's rotation to the same rotation as the given matrix.
 135         Matrix should be given in bone's local space.
 136     """
 137     q = mat.to_quaternion()
 138
 139     if pose_bone.rotation_mode == 'QUATERNION':
 140         pose_bone.rotation_quaternion = q
 141     elif pose_bone.rotation_mode == 'AXIS_ANGLE':
 142         pose_bone.rotation_axis_angle[0] = q.angle
 143         pose_bone.rotation_axis_angle[1] = q.axis[0]
 144         pose_bone.rotation_axis_angle[2] = q.axis[1]
 145         pose_bone.rotation_axis_angle[3] = q.axis[2]
 146     else:
 147         pose_bone.rotation_euler = q.to_euler(pose_bone.rotation_mode)
 148
 149
 150 def set_pose_scale(pose_bone, mat):
 151     """ Sets the pose bone's scale to the same scale as the given matrix.
 152         Matrix should be given in bone's local space.
 153     """
 154     pose_bone.scale = mat.to_scale()
 155
 156
 157 def match_pose_translation(pose_bone, target_bone):
 158     """ Matches pose_bone's visual translation to target_bone's visual
 159         translation.
 160         This function assumes you are in pose mode on the relevant armature.
 161     """
 162     mat = get_pose_matrix_in_other_space(target_bone.matrix, pose_bone)
 163     set_pose_translation(pose_bone, mat)
 164
 165
 166 def match_pose_rotation(pose_bone, target_bone):
 167     """ Matches pose_bone's visual rotation to target_bone's visual
 168         rotation.
 169         This function assumes you are in pose mode on the relevant armature.
 170     """
 171     mat = get_pose_matrix_in_other_space(target_bone.matrix, pose_bone)
 172     set_pose_rotation(pose_bone, mat)
 173
 174
 175 def match_pose_scale(pose_bone, target_bone):
 176     """ Matches pose_bone's visual scale to target_bone's visual
 177         scale.
 178         This function assumes you are in pose mode on the relevant armature.
 179     """
 180     mat = get_pose_matrix_in_other_space(target_bone.matrix, pose_bone)
 181     set_pose_scale(pose_bone, mat)
 182
 183
 184 ##############################
 185 ## IK/FK snapping functions ##
 186 ##############################
 187
 188 def correct_rotation(view_layer, bone_ik, target_matrix, *, ctrl_ik=None):
 189     """ Corrects the ik rotation in ik2fk snapping functions
 190     """
 191
 192     axis = target_matrix.to_3x3().col[1].normalized()
 193     ctrl_ik = ctrl_ik or bone_ik
 194
 195     def distance(angle):
 196         # Rotate the bone and return the actual angle between bones
 197         ctrl_ik.rotation_euler[1] = angle
 198         view_layer.update()
 199
 200         return -(bone_ik.vector.normalized().dot(axis))
 201
 202     if ctrl_ik.rotation_mode in {'QUATERNION', 'AXIS_ANGLE'}:
 203         ctrl_ik.rotation_mode = 'ZXY'
 204
 205     start_angle = ctrl_ik.rotation_euler[1]
 206
 207     alfarange = find_min_range(distance, start_angle)
 208     alfamin = ternarySearch(distance, alfarange[0], alfarange[1], pi / 180)
 209
 210     ctrl_ik.rotation_euler[1] = alfamin
 211     view_layer.update()
 212
 213
 214 def correct_scale(view_layer, bone_ik, target_matrix, *, ctrl_ik=None):
 215     """ Correct the scale of the base IK bone. """
 216     input_scale = target_matrix.to_scale()
 217     ctrl_ik = ctrl_ik or bone_ik
 218
 219     for i in range(3):
 220         cur_scale = bone_ik.matrix.to_scale()
 221
 222         ctrl_ik.scale = [
 223             v * i / c for v, i, c in zip(bone_ik.scale, input_scale, cur_scale)
 224         ]
 225
 226         view_layer.update()
 227
 228         if all(abs((c - i)/i) < 0.01 for i, c in zip(input_scale, cur_scale)):
 229             break
 230
 231
 232 def match_pole_target(view_layer, ik_first, ik_last, pole, match_bone_matrix, length):
 233     """ Places an IK chain's pole target to match ik_first's
 234         transforms to match_bone.  All bones should be given as pose bones.
 235         You need to be in pose mode on the relevant armature object.
 236         ik_first: first bone in the IK chain
 237         ik_last:  last bone in the IK chain
 238         pole:  pole target bone for the IK chain
 239         match_bone:  bone to match ik_first to (probably first bone in a matching FK chain)
 240         length:  distance pole target should be placed from the chain center
 241     """
 242     a = ik_first.matrix.to_translation()
 243     b = ik_last.matrix.to_translation() + ik_last.vector
 244
 245     # Vector from the head of ik_first to the
 246     # tip of ik_last
 247     ikv = b - a
 248
 249     # Get a vector perpendicular to ikv
 250     pv = perpendicular_vector(ikv).normalized() * length
 251
 252     def set_pole(pvi):
 253         """ Set pole target's position based on a vector
 254             from the arm center line.
 255         """
 256         # Translate pvi into armature space
 257         ploc = a + (ikv/2) + pvi
 258
 259         # Set pole target to location
 260         mat = get_pose_matrix_in_other_space(Matrix.Translation(ploc), pole)
 261         set_pose_translation(pole, mat)
 262
 263         view_layer.update()
 264
 265     set_pole(pv)
 266
 267     # Get the rotation difference between ik_first and match_bone
 268     angle = rotation_difference(ik_first.matrix, match_bone_matrix)
 269
 270     # Try compensating for the rotation difference in both directions
 271     pv1 = Matrix.Rotation(angle, 4, ikv) @ pv
 272     set_pole(pv1)
 273     ang1 = rotation_difference(ik_first.matrix, match_bone_matrix)
 274
 275     pv2 = Matrix.Rotation(-angle, 4, ikv) @ pv
 276     set_pole(pv2)
 277     ang2 = rotation_difference(ik_first.matrix, match_bone_matrix)
 278
 279     # Do the one with the smaller angle
 280     if ang1 < ang2:
 281         set_pole(pv1)
 282
 283 ##########
 284 ## Misc ##
 285 ##########
 286
 287 def parse_bone_names(names_string):
 288     if names_string[0] == '[' and names_string[-1] == ']':
 289         return eval(names_string)
 290     else:
 291         return names_string
 292
 293 ''']
 294
 295 UTILITIES_FUNC_OLD_ARM_FKIK = ['''
 296 ######################
 297 ## IK Arm functions ##
 298 ######################
 299
 300 def fk2ik_arm(obj, fk, ik):
 301     """ Matches the fk bones in an arm rig to the ik bones.
 302         obj: armature object
 303         fk:  list of fk bone names
 304         ik:  list of ik bone names
 305     """
 306     view_layer = bpy.context.view_layer
 307     uarm  = obj.pose.bones[fk[0]]
 308     farm  = obj.pose.bones[fk[1]]
 309     hand  = obj.pose.bones[fk[2]]
 310     uarmi = obj.pose.bones[ik[0]]
 311     farmi = obj.pose.bones[ik[1]]
 312     handi = obj.pose.bones[ik[2]]
 313
 314     if 'auto_stretch' in handi.keys():
 315         # This is kept for compatibility with legacy rigify Human
 316         # Stretch
 317         if handi['auto_stretch'] == 0.0:
 318             uarm['stretch_length'] = handi['stretch_length']
 319         else:
 320             diff = (uarmi.vector.length + farmi.vector.length) / (uarm.vector.length + farm.vector.length)
 321             uarm['stretch_length'] *= diff
 322
 323         # Upper arm position
 324         match_pose_rotation(uarm, uarmi)
 325         match_pose_scale(uarm, uarmi)
 326         view_layer.update()
 327
 328         # Forearm position
 329         match_pose_rotation(farm, farmi)
 330         match_pose_scale(farm, farmi)
 331         view_layer.update()
 332
 333         # Hand position
 334         match_pose_rotation(hand, handi)
 335         match_pose_scale(hand, handi)
 336         view_layer.update()
 337     else:
 338         # Upper arm position
 339         match_pose_translation(uarm, uarmi)
 340         match_pose_rotation(uarm, uarmi)
 341         match_pose_scale(uarm, uarmi)
 342         view_layer.update()
 343
 344         # Forearm position
 345         #match_pose_translation(hand, handi)
 346         match_pose_rotation(farm, farmi)
 347         match_pose_scale(farm, farmi)
 348         view_layer.update()
 349
 350         # Hand position
 351         match_pose_translation(hand, handi)
 352         match_pose_rotation(hand, handi)
 353         match_pose_scale(hand, handi)
 354         view_layer.update()
 355
 356
 357 def ik2fk_arm(obj, fk, ik):
 358     """ Matches the ik bones in an arm rig to the fk bones.
 359         obj: armature object
 360         fk:  list of fk bone names
 361         ik:  list of ik bone names
 362     """
 363     view_layer = bpy.context.view_layer
 364     uarm  = obj.pose.bones[fk[0]]
 365     farm  = obj.pose.bones[fk[1]]
 366     hand  = obj.pose.bones[fk[2]]
 367     uarmi = obj.pose.bones[ik[0]]
 368     farmi = obj.pose.bones[ik[1]]
 369     handi = obj.pose.bones[ik[2]]
 370
 371     main_parent = obj.pose.bones[ik[4]]
 372
 373     if ik[3] != "" and main_parent['pole_vector']:
 374         pole  = obj.pose.bones[ik[3]]
 375     else:
 376         pole = None
 377
 378
 379     if pole:
 380         # Stretch
 381         # handi['stretch_length'] = uarm['stretch_length']
 382
 383         # Hand position
 384         match_pose_translation(handi, hand)
 385         match_pose_rotation(handi, hand)
 386         match_pose_scale(handi, hand)
 387         view_layer.update()
 388
 389         # Pole target position
 390         match_pole_target(view_layer, uarmi, farmi, pole, uarm.matrix, (uarmi.length + farmi.length))
 391
 392     else:
 393         # Hand position
 394         match_pose_translation(handi, hand)
 395         match_pose_rotation(handi, hand)
 396         match_pose_scale(handi, hand)
 397         view_layer.update()
 398
 399         # Upper Arm position
 400         match_pose_translation(uarmi, uarm)
 401         #match_pose_rotation(uarmi, uarm)
 402         set_pose_rotation(uarmi, Matrix())
 403         match_pose_scale(uarmi, uarm)
 404         view_layer.update()
 405
 406         # Rotation Correction
 407         correct_rotation(view_layer, uarmi, uarm.matrix)
 408
 409     correct_scale(view_layer, uarmi, uarm.matrix)
 410 ''']
 411
 412 UTILITIES_FUNC_OLD_LEG_FKIK = ['''
 413 ######################
 414 ## IK Leg functions ##
 415 ######################
 416
 417 def fk2ik_leg(obj, fk, ik):
 418     """ Matches the fk bones in a leg rig to the ik bones.
 419         obj: armature object
 420         fk:  list of fk bone names
 421         ik:  list of ik bone names
 422     """
 423     view_layer = bpy.context.view_layer
 424     thigh  = obj.pose.bones[fk[0]]
 425     shin   = obj.pose.bones[fk[1]]
 426     foot   = obj.pose.bones[fk[2]]
 427     mfoot  = obj.pose.bones[fk[3]]
 428     thighi = obj.pose.bones[ik[0]]
 429     shini  = obj.pose.bones[ik[1]]
 430     footi  = obj.pose.bones[ik[2]]
 431     mfooti = obj.pose.bones[ik[3]]
 432
 433     if 'auto_stretch' in footi.keys():
 434         # This is kept for compatibility with legacy rigify Human
 435         # Stretch
 436         if footi['auto_stretch'] == 0.0:
 437             thigh['stretch_length'] = footi['stretch_length']
 438         else:
 439             diff = (thighi.vector.length + shini.vector.length) / (thigh.vector.length + shin.vector.length)
 440             thigh['stretch_length'] *= diff
 441
 442         # Thigh position
 443         match_pose_rotation(thigh, thighi)
 444         match_pose_scale(thigh, thighi)
 445         view_layer.update()
 446
 447         # Shin position
 448         match_pose_rotation(shin, shini)
 449         match_pose_scale(shin, shini)
 450         view_layer.update()
 451
 452         # Foot position
 453         footmat = get_pose_matrix_in_other_space(mfooti.matrix, foot)
 454         footmat = convert_pose_matrix_via_rest_delta(footmat, mfoot, foot)
 455         set_pose_rotation(foot, footmat)
 456         set_pose_scale(foot, footmat)
 457         view_layer.update()
 458
 459     else:
 460         # Thigh position
 461         match_pose_translation(thigh, thighi)
 462         match_pose_rotation(thigh, thighi)
 463         match_pose_scale(thigh, thighi)
 464         view_layer.update()
 465
 466         # Shin position
 467         match_pose_rotation(shin, shini)
 468         match_pose_scale(shin, shini)
 469         view_layer.update()
 470
 471         # Foot position
 472         footmat = get_pose_matrix_in_other_space(mfooti.matrix, foot)
 473         footmat = convert_pose_matrix_via_rest_delta(footmat, mfoot, foot)
 474         set_pose_rotation(foot, footmat)
 475         set_pose_scale(foot, footmat)
 476         view_layer.update()
 477
 478
 479 def ik2fk_leg(obj, fk, ik):
 480     """ Matches the ik bones in a leg rig to the fk bones.
 481         obj: armature object
 482         fk:  list of fk bone names
 483         ik:  list of ik bone names
 484     """
 485     view_layer = bpy.context.view_layer
 486     thigh    = obj.pose.bones[fk[0]]
 487     shin     = obj.pose.bones[fk[1]]
 488     mfoot    = obj.pose.bones[fk[2]]
 489     if fk[3] != "":
 490         foot      = obj.pose.bones[fk[3]]
 491     else:
 492         foot = None
 493     thighi   = obj.pose.bones[ik[0]]
 494     shini    = obj.pose.bones[ik[1]]
 495     footi    = obj.pose.bones[ik[2]]
 496     footroll = obj.pose.bones[ik[3]]
 497
 498     main_parent = obj.pose.bones[ik[6]]
 499
 500     if ik[4] != "" and main_parent['pole_vector']:
 501         pole     = obj.pose.bones[ik[4]]
 502     else:
 503         pole = None
 504     mfooti   = obj.pose.bones[ik[5]]
 505
 506     if (not pole) and (foot):
 507
 508         # Clear footroll
 509         set_pose_rotation(footroll, Matrix())
 510         view_layer.update()
 511
 512         # Foot position
 513         footmat = get_pose_matrix_in_other_space(foot.matrix, footi)
 514         footmat = convert_pose_matrix_via_rest_delta(footmat, mfooti, footi)
 515         set_pose_translation(footi, footmat)
 516         set_pose_rotation(footi, footmat)
 517         set_pose_scale(footi, footmat)
 518         view_layer.update()
 519
 520         # Thigh position
 521         match_pose_translation(thighi, thigh)
 522         #match_pose_rotation(thighi, thigh)
 523         set_pose_rotation(thighi, Matrix())
 524         match_pose_scale(thighi, thigh)
 525         view_layer.update()
 526
 527         # Rotation Correction
 528         correct_rotation(view_layer, thighi, thigh.matrix)
 529
 530     else:
 531         # Stretch
 532         if 'stretch_length' in footi.keys() and 'stretch_length' in thigh.keys():
 533             # Kept for compat with legacy rigify Human
 534             footi['stretch_length'] = thigh['stretch_length']
 535
 536         # Clear footroll
 537         set_pose_rotation(footroll, Matrix())
 538         view_layer.update()
 539
 540         # Foot position
 541         footmat = get_pose_matrix_in_other_space(mfoot.matrix, footi)
 542         footmat = convert_pose_matrix_via_rest_delta(footmat, mfooti, footi)
 543         set_pose_translation(footi, footmat)
 544         set_pose_rotation(footi, footmat)
 545         set_pose_scale(footi, footmat)
 546         view_layer.update()
 547
 548         # Pole target position
 549         match_pole_target(view_layer, thighi, shini, pole, thigh.matrix, (thighi.length + shini.length))
 550
 551     correct_scale(view_layer, thighi, thigh.matrix)
 552 ''']
 553
 554 UTILITIES_FUNC_OLD_POLE = ['''
 555 ################################
 556 ## IK Rotation-Pole functions ##
 557 ################################
 558
 559 def rotPoleToggle(rig, limb_type, controls, ik_ctrl, fk_ctrl, parent, pole):
 560
 561     rig_id = rig.data['rig_id']
 562     leg_fk2ik = eval('bpy.ops.pose.rigify_leg_fk2ik_' + rig_id)
 563     arm_fk2ik = eval('bpy.ops.pose.rigify_arm_fk2ik_' + rig_id)
 564     leg_ik2fk = eval('bpy.ops.pose.rigify_leg_ik2fk_' + rig_id)
 565     arm_ik2fk = eval('bpy.ops.pose.rigify_arm_ik2fk_' + rig_id)
 566
 567     controls = parse_bone_names(controls)
 568     ik_ctrl = parse_bone_names(ik_ctrl)
 569     fk_ctrl = parse_bone_names(fk_ctrl)
 570     parent = parse_bone_names(parent)
 571     pole = parse_bone_names(pole)
 572
 573     pbones = bpy.context.selected_pose_bones
 574     bpy.ops.pose.select_all(action='DESELECT')
 575
 576     for b in pbones:
 577
 578         new_pole_vector_value = not rig.pose.bones[parent]['pole_vector']
 579
 580         if b.name in controls or b.name in ik_ctrl:
 581             if limb_type == 'arm':
 582                 func1 = arm_fk2ik
 583                 func2 = arm_ik2fk
 584                 rig.pose.bones[controls[0]].bone.select = not new_pole_vector_value
 585                 rig.pose.bones[controls[4]].bone.select = not new_pole_vector_value
 586                 rig.pose.bones[parent].bone.select = not new_pole_vector_value
 587                 rig.pose.bones[pole].bone.select = new_pole_vector_value
 588
 589                 kwargs1 = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
 590                           'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1],
 591                           'hand_ik': controls[4]}
 592                 kwargs2 = {'uarm_fk': controls[1], 'farm_fk': controls[2], 'hand_fk': controls[3],
 593                           'uarm_ik': controls[0], 'farm_ik': ik_ctrl[1], 'hand_ik': controls[4],
 594                           'pole': pole, 'main_parent': parent}
 595             else:
 596                 func1 = leg_fk2ik
 597                 func2 = leg_ik2fk
 598                 rig.pose.bones[controls[0]].bone.select = not new_pole_vector_value
 599                 rig.pose.bones[controls[6]].bone.select = not new_pole_vector_value
 600                 rig.pose.bones[controls[5]].bone.select = not new_pole_vector_value
 601                 rig.pose.bones[parent].bone.select = not new_pole_vector_value
 602                 rig.pose.bones[pole].bone.select = new_pole_vector_value
 603
 604                 kwargs1 = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
 605                           'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
 606                           'foot_ik': ik_ctrl[2], 'mfoot_ik': ik_ctrl[2]}
 607                 kwargs2 = {'thigh_fk': controls[1], 'shin_fk': controls[2], 'foot_fk': controls[3],
 608                           'mfoot_fk': controls[7], 'thigh_ik': controls[0], 'shin_ik': ik_ctrl[1],
 609                           'foot_ik': controls[6], 'pole': pole, 'footroll': controls[5], 'mfoot_ik': ik_ctrl[2],
 610                           'main_parent': parent}
 611
 612             func1(**kwargs1)
 613             rig.pose.bones[parent]['pole_vector'] = new_pole_vector_value
 614             func2(**kwargs2)
 615
 616             bpy.ops.pose.select_all(action='DESELECT')
 617 ''']
 618
 619 REGISTER_OP_OLD_ARM_FKIK = ['Rigify_Arm_FK2IK', 'Rigify_Arm_IK2FK']
 620
 621 UTILITIES_OP_OLD_ARM_FKIK = ['''
 622 ##################################
 623 ## IK/FK Arm snapping operators ##
 624 ##################################
 625
 626 class Rigify_Arm_FK2IK(bpy.types.Operator):
 627     """ Snaps an FK arm to an IK arm.
 628     """
 629     bl_idname = "pose.rigify_arm_fk2ik_" + rig_id
 630     bl_label = "Rigify Snap FK arm to IK"
 631     bl_options = {'UNDO', 'INTERNAL'}
 632
 633     uarm_fk: StringProperty(name="Upper Arm FK Name")
 634     farm_fk: StringProperty(name="Forerm FK Name")
 635     hand_fk: StringProperty(name="Hand FK Name")
 636
 637     uarm_ik: StringProperty(name="Upper Arm IK Name")
 638     farm_ik: StringProperty(name="Forearm IK Name")
 639     hand_ik: StringProperty(name="Hand IK Name")
 640
 641     @classmethod
 642     def poll(cls, context):
 643         return (context.active_object != None and context.mode == 'POSE')
 644
 645     def execute(self, context):
 646         fk2ik_arm(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik])
 647         return {'FINISHED'}
 648
 649
 650 class Rigify_Arm_IK2FK(bpy.types.Operator):
 651     """ Snaps an IK arm to an FK arm.
 652     """
 653     bl_idname = "pose.rigify_arm_ik2fk_" + rig_id
 654     bl_label = "Rigify Snap IK arm to FK"
 655     bl_options = {'UNDO', 'INTERNAL'}
 656
 657     uarm_fk: StringProperty(name="Upper Arm FK Name")
 658     farm_fk: StringProperty(name="Forerm FK Name")
 659     hand_fk: StringProperty(name="Hand FK Name")
 660
 661     uarm_ik: StringProperty(name="Upper Arm IK Name")
 662     farm_ik: StringProperty(name="Forearm IK Name")
 663     hand_ik: StringProperty(name="Hand IK Name")
 664     pole   : StringProperty(name="Pole IK Name")
 665
 666     main_parent: StringProperty(name="Main Parent", default="")
 667
 668     @classmethod
 669     def poll(cls, context):
 670         return (context.active_object != None and context.mode == 'POSE')
 671
 672     def execute(self, context):
 673         ik2fk_arm(context.active_object, fk=[self.uarm_fk, self.farm_fk, self.hand_fk], ik=[self.uarm_ik, self.farm_ik, self.hand_ik, self.pole, self.main_parent])
 674         return {'FINISHED'}
 675 ''']
 676
 677 REGISTER_OP_OLD_LEG_FKIK = ['Rigify_Leg_FK2IK', 'Rigify_Leg_IK2FK']
 678
 679 UTILITIES_OP_OLD_LEG_FKIK = ['''
 680 ##################################
 681 ## IK/FK Leg snapping operators ##
 682 ##################################
 683
 684 class Rigify_Leg_FK2IK(bpy.types.Operator):
 685     """ Snaps an FK leg to an IK leg.
 686     """
 687     bl_idname = "pose.rigify_leg_fk2ik_" + rig_id
 688     bl_label = "Rigify Snap FK leg to IK"
 689     bl_options = {'UNDO', 'INTERNAL'}
 690
 691     thigh_fk: StringProperty(name="Thigh FK Name")
 692     shin_fk:  StringProperty(name="Shin FK Name")
 693     foot_fk:  StringProperty(name="Foot FK Name")
 694     mfoot_fk: StringProperty(name="MFoot FK Name")
 695
 696     thigh_ik: StringProperty(name="Thigh IK Name")
 697     shin_ik:  StringProperty(name="Shin IK Name")
 698     foot_ik:  StringProperty(name="Foot IK Name")
 699     mfoot_ik: StringProperty(name="MFoot IK Name")
 700
 701     @classmethod
 702     def poll(cls, context):
 703         return (context.active_object != None and context.mode == 'POSE')
 704
 705     def execute(self, context):
 706         fk2ik_leg(context.active_object, fk=[self.thigh_fk, self.shin_fk, self.foot_fk, self.mfoot_fk], ik=[self.thigh_ik, self.shin_ik, self.foot_ik, self.mfoot_ik])
 707         return {'FINISHED'}
 708
 709
 710 class Rigify_Leg_IK2FK(bpy.types.Operator):
 711     """ Snaps an IK leg to an FK leg.
 712     """
 713     bl_idname = "pose.rigify_leg_ik2fk_" + rig_id
 714     bl_label = "Rigify Snap IK leg to FK"
 715     bl_options = {'UNDO', 'INTERNAL'}
 716
 717     thigh_fk: StringProperty(name="Thigh FK Name")
 718     shin_fk:  StringProperty(name="Shin FK Name")
 719     mfoot_fk: StringProperty(name="MFoot FK Name")
 720     foot_fk:  StringProperty(name="Foot FK Name", default="")
 721     thigh_ik: StringProperty(name="Thigh IK Name")
 722     shin_ik:  StringProperty(name="Shin IK Name")
 723     foot_ik:  StringProperty(name="Foot IK Name")
 724     footroll: StringProperty(name="Foot Roll Name")
 725     pole:     StringProperty(name="Pole IK Name")
 726     mfoot_ik: StringProperty(name="MFoot IK Name")
 727
 728     main_parent: StringProperty(name="Main Parent", default="")
 729
 730     @classmethod
 731     def poll(cls, context):
 732         return (context.active_object != None and context.mode == 'POSE')
 733
 734     def execute(self, context):
 735         ik2fk_leg(context.active_object, fk=[self.thigh_fk, self.shin_fk, self.mfoot_fk, self.foot_fk], ik=[self.thigh_ik, self.shin_ik, self.foot_ik, self.footroll, self.pole, self.mfoot_ik, self.main_parent])
 736         return {'FINISHED'}
 737 ''']
 738
 739 REGISTER_OP_OLD_POLE = ['Rigify_Rot2PoleSwitch']
 740
 741 UTILITIES_OP_OLD_POLE = ['''
 742 ###########################
 743 ## IK Rotation Pole Snap ##
 744 ###########################
 745
 746 class Rigify_Rot2PoleSwitch(bpy.types.Operator):
 747     bl_idname = "pose.rigify_rot2pole_" + rig_id
 748     bl_label = "Rotation - Pole toggle"
 749     bl_description = "Toggles IK chain between rotation and pole target"
 750
 751     bone_name: StringProperty(default='')
 752     limb_type: StringProperty(name="Limb Type")
 753     controls: StringProperty(name="Controls string")
 754     ik_ctrl: StringProperty(name="IK Controls string")
 755     fk_ctrl: StringProperty(name="FK Controls string")
 756     parent: StringProperty(name="Parent name")
 757     pole: StringProperty(name="Pole name")
 758
 759     def execute(self, context):
 760         rig = context.object
 761
 762         if self.bone_name:
 763             bpy.ops.pose.select_all(action='DESELECT')
 764             rig.pose.bones[self.bone_name].bone.select = True
 765
 766         rotPoleToggle(rig, self.limb_type, self.controls, self.ik_ctrl, self.fk_ctrl, self.parent, self.pole)
 767         return {'FINISHED'}
 768 ''']
 769
 770 REGISTER_RIG_OLD_ARM = REGISTER_OP_OLD_ARM_FKIK + REGISTER_OP_OLD_POLE
 771
 772 UTILITIES_RIG_OLD_ARM = [
 773     *UTILITIES_FUNC_COMMON_IKFK,
 774     *UTILITIES_FUNC_OLD_ARM_FKIK,
 775     *UTILITIES_FUNC_OLD_POLE,
 776     *UTILITIES_OP_OLD_ARM_FKIK,
 777     *UTILITIES_OP_OLD_POLE,
 778 ]
 779
 780 REGISTER_RIG_OLD_LEG = REGISTER_OP_OLD_LEG_FKIK + REGISTER_OP_OLD_POLE
 781
 782 UTILITIES_RIG_OLD_LEG = [
 783     *UTILITIES_FUNC_COMMON_IKFK,
 784     *UTILITIES_FUNC_OLD_LEG_FKIK,
 785     *UTILITIES_FUNC_OLD_POLE,
 786     *UTILITIES_OP_OLD_LEG_FKIK,
 787     *UTILITIES_OP_OLD_POLE,
 788 ]
 789
 790 ##############################
 791 ## Default set of utilities ##
 792 ##############################
 793
 794 UI_REGISTER = [
 795     'RigUI',
 796     'RigLayers',
 797 ]
 798
 799 UI_UTILITIES = [
 800 ]
 801
 802 UI_SLIDERS = '''
 803 ###################
 804 ## Rig UI Panels ##
 805 ###################
 806
 807 class RigUI(bpy.types.Panel):
 808     bl_space_type = 'VIEW_3D'
 809     bl_region_type = 'UI'
 810     bl_label = "Rig Main Properties"
 811     bl_idname = "VIEW3D_PT_rig_ui_" + rig_id
 812     bl_category = 'Item'
 813
 814     @classmethod
 815     def poll(self, context):
 816         if context.mode != 'POSE':
 817             return False
 818         try:
 819             return (context.active_object.data.get("rig_id") == rig_id)
 820         except (AttributeError, KeyError, TypeError):
 821             return False
 822
 823     def draw(self, context):
 824         layout = self.layout
 825         pose_bones = context.active_object.pose.bones
 826         try:
 827             selected_bones = set(bone.name for bone in context.selected_pose_bones)
 828             selected_bones.add(context.active_pose_bone.name)
 829         except (AttributeError, TypeError):
 830             return
 831
 832         def is_selected(names):
 833             # Returns whether any of the named bones are selected.
 834             if isinstance(names, list) or isinstance(names, set):
 835                 return not selected_bones.isdisjoint(names)
 836             elif names in selected_bones:
 837                 return True
 838             return False
 839
 840         num_rig_separators = [-1]
 841
 842         def emit_rig_separator():
 843             if num_rig_separators[0] >= 0:
 844                 layout.separator()
 845             num_rig_separators[0] += 1
 846 '''
 847
 848 UI_REGISTER_BAKE_SETTINGS = ['RigBakeSettings']
 849
 850 UI_BAKE_SETTINGS = '''
 851 class RigBakeSettings(bpy.types.Panel):
 852     bl_space_type = 'VIEW_3D'
 853     bl_region_type = 'UI'
 854     bl_label = "Rig Bake Settings"
 855     bl_idname = "VIEW3D_PT_rig_bake_settings_" + rig_id
 856     bl_category = 'Item'
 857
 858     @classmethod
 859     def poll(self, context):
 860         return RigUI.poll(context) and find_action(context.active_object) is not None
 861
 862     def draw(self, context):
 863         RigifyBakeKeyframesMixin.draw_common_bake_ui(context, self.layout)
 864 '''
 865
 866 def layers_ui(layers, layout):
 867     """ Turn a list of booleans + a list of names into a layer UI.
 868     """
 869
 870     code = '''
 871 class RigLayers(bpy.types.Panel):
 872     bl_space_type = 'VIEW_3D'
 873     bl_region_type = 'UI'
 874     bl_label = "Rig Layers"
 875     bl_idname = "VIEW3D_PT_rig_layers_" + rig_id
 876     bl_category = 'Item'
 877
 878     @classmethod
 879     def poll(self, context):
 880         try:
 881             return (context.active_object.data.get("rig_id") == rig_id)
 882         except (AttributeError, KeyError, TypeError):
 883             return False
 884
 885     def draw(self, context):
 886         layout = self.layout
 887         col = layout.column()
 888 '''
 889     rows = {}
 890     for i in range(28):
 891         if layers[i]:
 892             if layout[i][1] not in rows:
 893                 rows[layout[i][1]] = []
 894             rows[layout[i][1]] += [(layout[i][0], i)]
 895
 896     keys = list(rows.keys())
 897     keys.sort()
 898
 899     for key in keys:
 900         code += "\n        row = col.row()\n"
 901         i = 0
 902         for l in rows[key]:
 903             if i > 3:
 904                 code += "\n        row = col.row()\n"
 905                 i = 0
 906             code += "        row.prop(context.active_object.data, 'layers', index=%s, toggle=True, text='%s')\n" % (str(l[1]), l[0])
 907             i += 1
 908
 909     # Root layer
 910     code += "\n        row = col.row()"
 911     code += "\n        row.separator()"
 912     code += "\n        row = col.row()"
 913     code += "\n        row.separator()\n"
 914     code += "\n        row = col.row()\n"
 915     code += "        row.prop(context.active_object.data, 'layers', index=28, toggle=True, text='Root')\n"
 916
 917     return code
 918
 919
 920 def quote_parameters(positional, named):
 921     """Quote the given positional and named parameters as a code string."""
 922     positional_list = [ repr(v) for v in positional ]
 923     named_list = [ "%s=%r" % (k, v) for k, v in named.items() ]
 924     return ', '.join(positional_list + named_list)
 925
 926 def indent_lines(lines, indent=4):
 927     if indent > 0:
 928         prefix = ' ' * indent
 929         return [ prefix + line for line in lines ]
 930     else:
 931         return lines
 932
 933
 934 class PanelLayout(object):
 935     """Utility class that builds code for creating a layout."""
 936
 937     def __init__(self, parent, index=0):
 938         if isinstance(parent, PanelLayout):
 939             self.parent = parent
 940             self.script = parent.script
 941         else:
 942             self.parent = None
 943             self.script = parent
 944
 945         self.header = []
 946         self.items = []
 947         self.indent = 0
 948         self.index = index
 949         self.layout = self._get_layout_var(index)
 950         self.is_empty = True
 951
 952     @staticmethod
 953     def _get_layout_var(index):
 954         return 'layout' if index == 0 else 'group' + str(index)
 955
 956     def clear_empty(self):
 957         self.is_empty = False
 958
 959         if self.parent:
 960             self.parent.clear_empty()
 961
 962     def get_lines(self):
 963         lines = []
 964
 965         for item in self.items:
 966             if isinstance(item, PanelLayout):
 967                 lines += item.get_lines()
 968             else:
 969                 lines.append(item)
 970
 971         if len(lines) > 0:
 972             return self.wrap_lines(lines)
 973         else:
 974             return []
 975
 976     def wrap_lines(self, lines):
 977         return self.header + indent_lines(lines, self.indent)
 978
 979     def add_line(self, line):
 980         assert isinstance(line, str)
 981
 982         self.items.append(line)
 983
 984         if self.is_empty:
 985             self.clear_empty()
 986
 987     def use_bake_settings(self):
 988         """This panel contains operators that need the common Bake settings."""
 989         self.parent.use_bake_settings()
 990
 991     def custom_prop(self, bone_name, prop_name, **params):
 992         """Add a custom property input field to the panel."""
 993         param_str = quote_parameters([ rna_idprop_quote_path(prop_name) ], params)
 994         self.add_line(
 995             "%s.prop(pose_bones[%r], %s)" % (self.layout, bone_name, param_str)
 996         )
 997
 998     def operator(self, operator_name, *, properties=None, **params):
 999         """Add an operator call button to the panel."""
1000         name = operator_name.format_map(self.script.format_args)
1001         param_str = quote_parameters([ name ], params)
1002         call_str = "%s.operator(%s)" % (self.layout, param_str)
1003         if properties:
1004             self.add_line("props = " + call_str)
1005             for k, v in properties.items():
1006                 self.add_line("props.%s = %r" % (k,v))
1007         else:
1008             self.add_line(call_str)
1009
1010     def add_nested_layout(self, name, params):
1011         param_str = quote_parameters([], params)
1012         sub_panel = PanelLayout(self, self.index + 1)
1013         sub_panel.header.append('%s = %s.%s(%s)' % (sub_panel.layout, self.layout, name, param_str))
1014         self.items.append(sub_panel)
1015         return sub_panel
1016
1017     def row(self, **params):
1018         """Add a nested row layout to the panel."""
1019         return self.add_nested_layout('row', params)
1020
1021     def column(self, **params):
1022         """Add a nested column layout to the panel."""
1023         return self.add_nested_layout('column', params)
1024
1025     def split(self, **params):
1026         """Add a split layout to the panel."""
1027         return self.add_nested_layout('split', params)
1028
1029
1030 class BoneSetPanelLayout(PanelLayout):
1031     """Panel restricted to a certain set of bones."""
1032
1033     def __init__(self, rig_panel, bones):
1034         assert isinstance(bones, frozenset)
1035         super().__init__(rig_panel)
1036         self.bones = bones
1037         self.show_bake_settings = False
1038
1039     def clear_empty(self):
1040         self.parent.bones |= self.bones
1041
1042         super().clear_empty()
1043
1044     def wrap_lines(self, lines):
1045         if self.bones != self.parent.bones:
1046             header = ["if is_selected(%r):" % (set(self.bones))]
1047             return header + indent_lines(lines)
1048         else:
1049             return lines
1050
1051     def use_bake_settings(self):
1052         self.show_bake_settings = True
1053         if not self.script.use_bake_settings:
1054             self.script.use_bake_settings = True
1055             self.script.add_utilities(SCRIPT_UTILITIES_BAKE)
1056             self.script.register_classes(SCRIPT_REGISTER_BAKE)
1057
1058
1059 class RigPanelLayout(PanelLayout):
1060     """Panel owned by a certain rig."""
1061
1062     def __init__(self, script, rig):
1063         super().__init__(script)
1064         self.bones = set()
1065         self.subpanels = OrderedDict()
1066
1067     def wrap_lines(self, lines):
1068         header = [ "if is_selected(%r):" % (set(self.bones)) ]
1069         prefix = [ "emit_rig_separator()" ]
1070         return header + indent_lines(prefix + lines)
1071
1072     def panel_with_selected_check(self, control_names):
1073         selected_set = frozenset(control_names)
1074
1075         if selected_set in self.subpanels:
1076             return self.subpanels[selected_set]
1077         else:
1078             panel = BoneSetPanelLayout(self, selected_set)
1079             self.subpanels[selected_set] = panel
1080             self.items.append(panel)
1081             return panel
1082
1083
1084 class ScriptGenerator(base_generate.GeneratorPlugin):
1085     """Generator plugin that builds the python script attached to the rig."""
1086
1087     priority = -100
1088
1089     def __init__(self, generator):
1090         super().__init__(generator)
1091
1092         self.ui_scripts = []
1093         self.ui_imports = UI_IMPORTS.copy()
1094         self.ui_utilities = UI_UTILITIES.copy()
1095         self.ui_register = UI_REGISTER.copy()
1096         self.ui_register_drivers = []
1097         self.ui_register_props = []
1098
1099         self.ui_rig_panels = OrderedDict()
1100
1101         self.use_bake_settings = False
1102
1103     # Structured panel code generation
1104     def panel_with_selected_check(self, rig, control_names):
1105         """Add a panel section with restricted selection."""
1106         rig_key = id(rig)
1107
1108         if rig_key in self.ui_rig_panels:
1109             panel = self.ui_rig_panels[rig_key]
1110         else:
1111             panel = RigPanelLayout(self, rig)
1112             self.ui_rig_panels[rig_key] = panel
1113
1114         return panel.panel_with_selected_check(control_names)
1115
1116     # Raw output
1117     def add_panel_code(self, str_list):
1118         """Add raw code to the panel."""
1119         self.ui_scripts += str_list
1120
1121     def add_imports(self, str_list):
1122         self.ui_imports += str_list
1123
1124     def add_utilities(self, str_list):
1125         self.ui_utilities += str_list
1126
1127     def register_classes(self, str_list):
1128         self.ui_register += str_list
1129
1130     def register_driver_functions(self, str_list):
1131         self.ui_register_drivers += str_list
1132
1133     def register_property(self, name, definition):
1134         self.ui_register_props.append((name, definition))
1135
1136     def initialize(self):
1137         self.format_args = {
1138             'rig_id': self.generator.rig_id,
1139         }
1140
1141     def finalize(self):
1142         metarig = self.generator.metarig
1143         rig_id = self.generator.rig_id
1144
1145         vis_layers = self.obj.data.layers
1146
1147         # Ensure the collection of layer names exists
1148         for i in range(1 + len(metarig.data.rigify_layers), 29):
1149             metarig.data.rigify_layers.add()
1150
1151         # Create list of layer name/row pairs
1152         layer_layout = []
1153         for l in metarig.data.rigify_layers:
1154             layer_layout += [(l.name, l.row)]
1155
1156         # Generate the UI script
1157         script = metarig.data.rigify_rig_ui
1158
1159         if script:
1160             script.clear()
1161         else:
1162             script_name = self.generator.obj.name + "_ui.py"
1163             script = bpy.data.texts.new(script_name)
1164             metarig.data.rigify_rig_ui = script
1165
1166         for s in OrderedDict.fromkeys(self.ui_imports):
1167             script.write(s + "\n")
1168
1169         script.write(UI_BASE_UTILITIES % rig_id)
1170
1171         for s in OrderedDict.fromkeys(self.ui_utilities):
1172             script.write(s + "\n")
1173
1174         script.write(UI_SLIDERS)
1175
1176         for s in self.ui_scripts:
1177             script.write("\n        " + s.replace("\n", "\n        ") + "\n")
1178
1179         if len(self.ui_scripts) > 0:
1180             script.write("\n        num_rig_separators[0] = 0\n")
1181
1182         for panel in self.ui_rig_panels.values():
1183             lines = panel.get_lines()
1184             if len(lines) > 1:
1185                 script.write("\n        ".join([''] + lines) + "\n")
1186
1187         if self.use_bake_settings:
1188             self.ui_register = UI_REGISTER_BAKE_SETTINGS + self.ui_register
1189             script.write(UI_BAKE_SETTINGS)
1190
1191         script.write(layers_ui(vis_layers, layer_layout))
1192
1193         script.write("\ndef register():\n")
1194
1195         ui_register = OrderedDict.fromkeys(self.ui_register)
1196         for s in ui_register:
1197             script.write("    bpy.utils.register_class("+s+")\n")
1198
1199         ui_register_drivers = OrderedDict.fromkeys(self.ui_register_drivers)
1200         for s in ui_register_drivers:
1201             script.write("    bpy.app.driver_namespace['"+s+"'] = "+s+"\n")
1202
1203         ui_register_props = OrderedDict.fromkeys(self.ui_register_props)
1204         for s in ui_register_props:
1205             script.write("    bpy.types.%s = %s\n " % (*s,))
1206
1207         script.write("\ndef unregister():\n")
1208
1209         for s in ui_register_props:
1210             script.write("    del bpy.types.%s\n" % s[0])
1211
1212         for s in ui_register:
1213             script.write("    bpy.utils.unregister_class("+s+")\n")
1214
1215         for s in ui_register_drivers:
1216             script.write("    del bpy.app.driver_namespace['"+s+"']\n")
1217
1218         script.write("\nregister()\n")
1219         script.use_module = True
1220
1221         # Run UI script
1222         exec(script.as_string(), {})
1223
1224         # Attach the script to the rig
1225         self.obj['rig_ui'] = script