Merging r39312 through r39329 from trunk into soc-2011-tomato
[blender.git] / release / scripts / startup / bl_operators / nla.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-80 compliant>
20
21 import bpy
22 from bpy.types import Operator
23
24
25 def pose_frame_info(obj):
26     from mathutils import Matrix
27
28     info = {}
29
30     pose = obj.pose
31
32     pose_items = pose.bones.items()
33
34     for name, pbone in pose_items:
35         binfo = {}
36         bone = pbone.bone
37
38         binfo["parent"] = getattr(bone.parent, "name", None)
39         binfo["bone"] = bone
40         binfo["pbone"] = pbone
41         binfo["matrix_local"] = bone.matrix_local.copy()
42         try:
43             binfo["matrix_local_inv"] = binfo["matrix_local"].inverted()
44         except:
45             binfo["matrix_local_inv"] = Matrix()
46
47         binfo["matrix"] = bone.matrix.copy()
48         binfo["matrix_pose"] = pbone.matrix.copy()
49         try:
50             binfo["matrix_pose_inv"] = binfo["matrix_pose"].inverted()
51         except:
52             binfo["matrix_pose_inv"] = Matrix()
53
54         info[name] = binfo
55
56     for name, pbone in pose_items:
57         binfo = info[name]
58         binfo_parent = binfo.get("parent", None)
59         if binfo_parent:
60             binfo_parent = info[binfo_parent]
61
62         matrix = binfo["matrix_pose"]
63         rest_matrix = binfo["matrix_local"]
64
65         if binfo_parent:
66             matrix = binfo_parent["matrix_pose_inv"] * matrix
67             rest_matrix = binfo_parent["matrix_local_inv"] * rest_matrix
68
69         binfo["matrix_key"] = rest_matrix.inverted() * matrix
70
71     return info
72
73
74 def obj_frame_info(obj):
75     info = {}
76     # parent = obj.parent
77     info["matrix_key"] = obj.matrix_local.copy()
78     return info
79
80
81 def bake(frame_start,
82          frame_end, step=1,
83          only_selected=False,
84          do_pose=True,
85          do_object=True,
86          do_constraint_clear=False,
87          ):
88
89     scene = bpy.context.scene
90     obj = bpy.context.object
91     pose = obj.pose
92     frame_back = scene.frame_current
93
94     if pose is None:
95         do_pose = False
96
97     if do_pose is None and do_object is None:
98         return None
99
100     pose_info = []
101     obj_info = []
102
103     frame_range = range(frame_start, frame_end + 1, step)
104
105     # -------------------------------------------------------------------------
106     # Collect transformations
107
108     # could speed this up by applying steps here too...
109     for f in frame_range:
110         scene.frame_set(f)
111
112         if do_pose:
113             pose_info.append(pose_frame_info(obj))
114         if do_object:
115             obj_info.append(obj_frame_info(obj))
116
117         f += 1
118
119     # -------------------------------------------------------------------------
120     # Create action
121
122     # incase animation data hassnt been created
123     atd = obj.animation_data_create()
124     action = bpy.data.actions.new("Action")
125     atd.action = action
126
127     if do_pose:
128         pose_items = pose.bones.items()
129     else:
130         pose_items = []  # skip
131
132     # -------------------------------------------------------------------------
133     # Apply transformations to action
134
135     # pose
136     for name, pbone in (pose_items if do_pose else ()):
137         if only_selected and not pbone.bone.select:
138             continue
139
140         if do_constraint_clear:
141             while pbone.constraints:
142                 pbone.constraints.remove(pbone.constraints[0])
143
144         for f in frame_range:
145             matrix = pose_info[(f - frame_start) // step][name]["matrix_key"]
146
147             # pbone.location = matrix.to_translation()
148             # pbone.rotation_quaternion = matrix.to_quaternion()
149             pbone.matrix_basis = matrix
150
151             pbone.keyframe_insert("location", -1, f, name)
152
153             rotation_mode = pbone.rotation_mode
154
155             if rotation_mode == 'QUATERNION':
156                 pbone.keyframe_insert("rotation_quaternion", -1, f, name)
157             elif rotation_mode == 'AXIS_ANGLE':
158                 pbone.keyframe_insert("rotation_axis_angle", -1, f, name)
159             else:  # euler, XYZ, ZXY etc
160                 pbone.keyframe_insert("rotation_euler", -1, f, name)
161
162             pbone.keyframe_insert("scale", -1, f, name)
163
164     # object. TODO. multiple objects
165     if do_object:
166         if do_constraint_clear:
167             while obj.constraints:
168                 obj.constraints.remove(obj.constraints[0])
169
170         for f in frame_range:
171             matrix = obj_info[(f - frame_start) // step]["matrix_key"]
172             obj.matrix_local = matrix
173
174             obj.keyframe_insert("location", -1, f)
175
176             rotation_mode = obj.rotation_mode
177
178             if rotation_mode == 'QUATERNION':
179                 obj.keyframe_insert("rotation_quaternion", -1, f)
180             elif rotation_mode == 'AXIS_ANGLE':
181                 obj.keyframe_insert("rotation_axis_angle", -1, f)
182             else:  # euler, XYZ, ZXY etc
183                 obj.keyframe_insert("rotation_euler", -1, f)
184
185             obj.keyframe_insert("scale", -1, f)
186
187     scene.frame_set(frame_back)
188
189     return action
190
191
192 from bpy.props import IntProperty, BoolProperty, EnumProperty
193
194
195 class BakeAction(Operator):
196     '''Bake animation to an Action'''
197     bl_idname = "nla.bake"
198     bl_label = "Bake Action"
199     bl_options = {'REGISTER', 'UNDO'}
200
201     frame_start = IntProperty(name="Start Frame",
202             description="Start frame for baking",
203             default=1, min=1, max=300000)
204     frame_end = IntProperty(name="End Frame",
205             description="End frame for baking",
206             default=250, min=1, max=300000)
207     step = IntProperty(name="Frame Step",
208             description="Frame Step",
209             default=1, min=1, max=120)
210     only_selected = BoolProperty(name="Only Selected",
211             default=True)
212     clear_consraints = BoolProperty(name="Clear Constraints",
213             default=False)
214     bake_types = EnumProperty(
215             name="Bake Data",
216             options={'ENUM_FLAG'},
217             items=(('POSE', "Pose", ""),
218                    ('OBJECT', "Object", ""),
219                    ),
220             default={'POSE'},
221             )
222
223     def execute(self, context):
224
225         action = bake(self.frame_start,
226                       self.frame_end,
227                       self.step,
228                       self.only_selected,
229                       'POSE' in self.bake_types,
230                       'OBJECT' in self.bake_types,
231                       self.clear_consraints,
232                       )
233
234         if action is None:
235             self.report({'INFO'}, "Nothing to bake")
236             return {'CANCELLED'}
237
238         # basic cleanup, could move elsewhere
239         for fcu in action.fcurves:
240             keyframe_points = fcu.keyframe_points
241             i = 1
242             while i < len(fcu.keyframe_points) - 1:
243                 val_prev = keyframe_points[i - 1].co[1]
244                 val_next = keyframe_points[i + 1].co[1]
245                 val = keyframe_points[i].co[1]
246
247                 if abs(val - val_prev) + abs(val - val_next) < 0.0001:
248                     keyframe_points.remove(keyframe_points[i])
249                 else:
250                     i += 1
251
252         return {'FINISHED'}
253
254     def invoke(self, context, event):
255         wm = context.window_manager
256         return wm.invoke_props_dialog(self)