added autorig neck type
[blender.git] / release / scripts / modules / rigify / neck.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., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
16 #
17 # ##### END GPL LICENSE BLOCK #####
18
19 import bpy
20 from rigify import bone_class_instance, copy_bone_simple
21 from rna_prop_ui import rna_idprop_ui_prop_get
22
23
24 def main(obj, orig_bone_name):
25     from Mathutils import Vector
26     
27     arm = obj.data
28
29     # Initialize container classes for convenience
30     mt = bone_class_instance(obj, ["body", "head"]) # meta
31     mt.head = orig_bone_name
32     mt.update()
33     mt.body = mt.head_e.parent.name
34     mt.update()
35
36     # child chain of the 'head'
37     children = mt.head_e.children
38     if len(children) != 1:
39         print("expected the head to have only 1 child.")
40
41     child = children[0]
42     neck_chain = [child] + child.children_recursive_basename
43     neck_chain = [child.name for child in neck_chain]
44     
45     mt_chain = bone_class_instance(obj, [("neck_%.2d" % (i + 1)) for i in range(len(neck_chain))]) # 99 bones enough eh?
46     for i, child_name in enumerate(neck_chain):
47         setattr(mt_chain, ("neck_%.2d" % (i + 1)), child_name)
48     mt_chain.update()
49
50     neck_chain_basename = mt_chain.neck_01_e.basename
51     neck_chain_segment_length = mt_chain.neck_01_e.length
52     
53     ex = bone_class_instance(obj, ["body", "head", "head_hinge","neck_socket"]) # hinge & extras
54     
55     # Add the head hinge at the bodys location, becomes the parent of the original head
56     
57     
58     # Copy the head bone and offset
59     ex.head_e = copy_bone_simple(arm, mt.head, "MCH_%s" % mt.head, parent=True)
60     ex.head = ex.head_e.name
61     # offset
62     head_length = ex.head_e.length
63     ex.head_e.head.y += head_length / 2.0
64     ex.head_e.tail.y += head_length / 2.0
65     
66     # Yes, use the body bone but call it a head hinge
67     ex.head_hinge_e = copy_bone_simple(arm, mt.body, "MCH_%s_hinge" % mt.head, parent=True)
68     ex.head_hinge = ex.head_hinge_e.name
69     ex.head_hinge_e.head.y += head_length / 4.0
70     ex.head_hinge_e.tail.y += head_length / 4.0
71
72     # reparent the head, assume its not connected
73     mt.head_e.parent = ex.head_hinge_e
74     
75     # Insert the neck socket, the head copys this loation
76     ex.neck_socket_e = arm.edit_bones.new("MCH-%s_socked" % neck_chain_basename)
77     ex.neck_socket = ex.neck_socket_e.name
78     ex.neck_socket_e.parent = mt.body_e
79     ex.neck_socket_e.head = mt.head_e.head
80     ex.neck_socket_e.tail = mt.head_e.head - Vector(0.0, neck_chain_segment_length / 2.0, 0.0)
81     ex.neck_socket_e.roll = 0.0
82     
83     # offset the head, not really needed since it has a copyloc constraint
84     mt.head_e.head.y += head_length / 4.0
85     mt.head_e.tail.y += head_length / 4.0
86
87     for i in range(len(neck_chain)):
88         neck_e = getattr(mt_chain, "neck_%.2d_e" % (i + 1))
89         
90         # dont store parent names, re-reference as each chain bones parent.
91         neck_e_parent = arm.edit_bones.new("MCH-rot_%s" % neck_e.name)
92         neck_e_parent.head = neck_e.head
93         neck_e_parent.tail = neck_e.head + Vector(0.0, 0.0, neck_chain_segment_length / 2.0)
94         neck_e_parent.roll = 0.0
95         
96         orig_parent = neck_e.parent
97         neck_e.connected = False
98         neck_e.parent = neck_e_parent
99         neck_e_parent.connected = False
100         
101         if i == 0:
102             neck_e_parent.parent = mt.body_e
103         else:
104             neck_e_parent.parent = orig_parent
105     
106     
107     bpy.ops.object.mode_set(mode='OBJECT')
108     
109     mt.update()
110     mt_chain.update()
111     ex.update()
112     
113     # Simple one off constraints, no drivers
114     con = mt.head_p.constraints.new('COPY_LOCATION')
115     con.target = obj
116     con.subtarget = ex.neck_socket
117     
118     con = ex.head_p.constraints.new('COPY_ROTATION')
119     con.target = obj
120     con.subtarget = mt.head
121     
122     # driven hinge
123     prop = rna_idprop_ui_prop_get(mt.head_p, "hinge", create=True)
124     mt.head_p["hinge"] = 0.0
125     prop["soft_min"] = 0.0
126     prop["soft_max"] = 1.0
127     
128     con = ex.head_hinge_p.constraints.new('COPY_ROTATION')
129     con.name = "hinge"
130     con.target = obj
131     con.subtarget = mt.body
132     
133     # add driver
134     hinge_driver_path = mt.head_p.path_to_id() + '["hinge"]'
135     
136     fcurve = con.driver_add("influence", 0)
137     driver = fcurve.driver
138     tar = driver.targets.new()
139     driver.type = 'AVERAGE'
140     tar.name = "var"
141     tar.id_type = 'OBJECT'
142     tar.id = obj
143     tar.rna_path = hinge_driver_path
144     
145     #mod = fcurve_driver.modifiers.new('GENERATOR')
146     mod = fcurve.modifiers[0]
147     mod.poly_order = 1
148     mod.coefficients[0] = 1.0
149     mod.coefficients[1] = -1.0
150     
151     head_driver_path = mt.head_p.path_to_id()
152     
153     # b01/max(0.001,b01+b02+b03+b04+b05)
154     target_names = [("b%.2d" % (i + 1)) for i in range(len(neck_chain))]
155     expression_suffix = "/max(0.001,%s)" % "+".join(target_names)
156     
157     
158     for i in range(len(neck_chain)):
159         neck_p = getattr(mt_chain, "neck_%.2d_p" % (i + 1))
160         neck_p.lock_location = True, True, True
161         neck_p.lock_location = True, True, True
162         neck_p.lock_rotations_4d = True
163         
164         # Add bend prop
165         prop_name = "bend_%.2d" % (i + 1)
166         prop = rna_idprop_ui_prop_get(mt.head_p, prop_name, create=True)
167         mt.head_p[prop_name] = 1.0
168         prop["soft_min"] = 0.0
169         prop["soft_max"] = 1.0
170         
171         # add parent constraint
172         neck_p_parent = neck_p.parent
173         
174         # add constraint
175         con = neck_p_parent.constraints.new('COPY_ROTATION')
176         con.name = "Copy Rotation"
177         con.target = obj
178         con.subtarget = ex.head
179         con.owner_space = 'LOCAL'
180         con.target_space = 'LOCAL'
181         
182         fcurve = con.driver_add("influence", 0)
183         driver = fcurve.driver
184         driver.type = 'SCRIPTED'
185         # b01/max(0.001,b01+b02+b03+b04+b05)
186         driver.expression = target_names[i] + expression_suffix
187         fcurve.modifiers.remove(0) # grr dont need a modifier
188
189         for j in range(len(neck_chain)):
190             tar = driver.targets.new()
191             tar.name = target_names[j]
192             tar.id_type = 'OBJECT'
193             tar.id = obj
194             tar.rna_path = head_driver_path + ('["bend_%.2d"]' % (j + 1))