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