Bake Action: operate on selected objects
[blender-staging.git] / release / scripts / modules / bpy_extras / anim_utils.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 __all__ = (
22     "bake_action",
23     "bake_action_objects",
24
25     "bake_action_iter",
26     "bake_action_objects_iter",
27 )
28
29 import bpy
30
31
32 def bake_action(
33         obj,
34         *,
35         action, frames,
36         **kwargs,
37 ):
38     """
39     :arg obj: Object to bake.
40     :type obj: :class:`bpy.types.Object`
41     :arg action: An action to bake the data into, or None for a new action
42        to be created.
43     :type action: :class:`bpy.types.Action` or None
44     :arg frames: Frames to bake.
45     :type frames: iterable of int
46
47     :return: an action or None
48     :rtype: :class:`bpy.types.Action`
49     """
50     if not (do_pose or do_object):
51         return None
52
53     action, = bake_action_objects(
54         [(obj, action)],
55         frames,
56         **kwargs,
57     )
58     return action
59
60
61 def bake_action_objects(
62         object_action_pairs,
63         *,
64         frames,
65         **kwargs,
66 ):
67     """
68     A version of :func:`bake_action_objects_iter` that takes frames and returns the output.
69
70     :arg frames: Frames to bake.
71     :type frames: iterable of int
72
73     :return: A sequence of Action or None types (aligned with `object_action_pairs`)
74     :rtype: sequence of :class:`bpy.types.Action`
75     """
76     iter = bake_action_objects_iter(object_action_pairs, **kwargs)
77     iter.send(None)
78     for frame in frames:
79         iter.send(frame)
80     return iter.send(None)
81
82
83 def bake_action_objects_iter(
84         object_action_pairs,
85         **kwargs,
86 ):
87     """
88     An coroutine that bakes actions for multiple objects.
89
90     :arg object_action_pairs: Sequence of object action tuples,
91        action is the destination for the baked data. When None a new action will be created.
92     :type object_action_pairs: Sequence of (:class:`bpy.types.Object`, :class:`bpy.types.Action`)
93     """
94     scene = bpy.context.scene
95     frame_back = scene.frame_current
96     iter_all = tuple(
97         bake_action_iter(obj, action=action, **kwargs)
98         for (obj, action) in object_action_pairs
99     )
100     for iter in iter_all:
101         iter.send(None)
102     while True:
103         frame = yield None
104         if frame is None:
105             break
106         scene.frame_set(frame)
107         scene.update()
108         for iter in iter_all:
109             iter.send(frame)
110     scene.frame_set(frame_back)
111     yield tuple(iter.send(None) for iter in iter_all)
112
113
114 # XXX visual keying is actually always considered as True in this code...
115 def bake_action_iter(
116         obj,
117         *,
118         action,
119         only_selected=False,
120         do_pose=True,
121         do_object=True,
122         do_visual_keying=True,
123         do_constraint_clear=False,
124         do_parents_clear=False,
125         do_clean=False,
126 ):
127     """
128     An coroutine that bakes action for a single object.
129
130     :arg obj: Object to bake.
131     :type obj: :class:`bpy.types.Object`
132     :arg action: An action to bake the data into, or None for a new action
133        to be created.
134     :type action: :class:`bpy.types.Action` or None
135     :arg only_selected: Only bake selected bones.
136     :type only_selected: bool
137     :arg do_pose: Bake pose channels.
138     :type do_pose: bool
139     :arg do_object: Bake objects.
140     :type do_object: bool
141     :arg do_visual_keying: Use the final transformations for baking ('visual keying')
142     :type do_visual_keying: bool
143     :arg do_constraint_clear: Remove constraints after baking.
144     :type do_constraint_clear: bool
145     :arg do_parents_clear: Unparent after baking objects.
146     :type do_parents_clear: bool
147     :arg do_clean: Remove redundant keyframes after baking.
148     :type do_clean: bool
149
150     :return: an action or None
151     :rtype: :class:`bpy.types.Action`
152     """
153     # -------------------------------------------------------------------------
154     # Helper Functions and vars
155
156     def pose_frame_info(obj):
157         matrix = {}
158         for name, pbone in obj.pose.bones.items():
159             if do_visual_keying:
160                 # Get the final transform of the bone in its own local space...
161                 matrix[name] = obj.convert_space(pbone, pbone.matrix, 'POSE', 'LOCAL')
162             else:
163                 matrix[name] = pbone.matrix_basis.copy()
164         return matrix
165
166     if do_parents_clear:
167         if do_visual_keying:
168             def obj_frame_info(obj):
169                 return obj.matrix_world.copy()
170         else:
171             def obj_frame_info(obj):
172                 parent = obj.parent
173                 matrix = obj.matrix_basis
174                 if parent:
175                     return parent.matrix_world * matrix
176                 else:
177                     return matrix.copy()
178     else:
179         if do_visual_keying:
180             def obj_frame_info(obj):
181                 parent = obj.parent
182                 matrix = obj.matrix_world
183                 if parent:
184                     return parent.matrix_world.inverted_safe() * matrix
185                 else:
186                     return matrix.copy()
187         else:
188             def obj_frame_info(obj):
189                 return obj.matrix_basis.copy()
190
191     # -------------------------------------------------------------------------
192     # Setup the Context
193
194     if obj.pose is None:
195         do_pose = False
196
197     if not (do_pose or do_object):
198         raise Exception("Pose and object baking is disabled, no action needed")
199
200     pose_info = []
201     obj_info = []
202
203     options = {'INSERTKEY_NEEDED'}
204
205     # -------------------------------------------------------------------------
206     # Collect transformations
207
208     while True:
209         # Caller is responsible for setting the frame and updating the scene.
210         frame = yield None
211
212         # Signal we're done!
213         if frame is None:
214             break
215
216         if do_pose:
217             pose_info.append((frame, pose_frame_info(obj)))
218         if do_object:
219             obj_info.append((frame, obj_frame_info(obj)))
220
221     # -------------------------------------------------------------------------
222     # Clean (store initial data)
223     if do_clean and action is not None:
224         clean_orig_data = {fcu: {p.co[1] for p in fcu.keyframe_points} for fcu in action.fcurves}
225     else:
226         clean_orig_data = {}
227
228     # -------------------------------------------------------------------------
229     # Create action
230
231     # in case animation data hasn't been created
232     atd = obj.animation_data_create()
233     if action is None:
234         action = bpy.data.actions.new("Action")
235
236     # Leave tweak mode before trying to modify the action (T48397)
237     if atd.use_tweak_mode:
238         atd.use_tweak_mode = False
239
240     atd.action = action
241
242     # -------------------------------------------------------------------------
243     # Apply transformations to action
244
245     # pose
246     if do_pose:
247         for name, pbone in obj.pose.bones.items():
248             if only_selected and not pbone.bone.select:
249                 continue
250
251             if do_constraint_clear:
252                 while pbone.constraints:
253                     pbone.constraints.remove(pbone.constraints[0])
254
255             # create compatible eulers
256             euler_prev = None
257
258             for (f, matrix) in pose_info:
259                 pbone.matrix_basis = matrix[name].copy()
260
261                 pbone.keyframe_insert("location", -1, f, name, options)
262
263                 rotation_mode = pbone.rotation_mode
264                 if rotation_mode == 'QUATERNION':
265                     pbone.keyframe_insert("rotation_quaternion", -1, f, name, options)
266                 elif rotation_mode == 'AXIS_ANGLE':
267                     pbone.keyframe_insert("rotation_axis_angle", -1, f, name, options)
268                 else:  # euler, XYZ, ZXY etc
269                     if euler_prev is not None:
270                         euler = pbone.rotation_euler.copy()
271                         euler.make_compatible(euler_prev)
272                         pbone.rotation_euler = euler
273                         euler_prev = euler
274                         del euler
275                     else:
276                         euler_prev = pbone.rotation_euler.copy()
277                     pbone.keyframe_insert("rotation_euler", -1, f, name, options)
278
279                 pbone.keyframe_insert("scale", -1, f, name, options)
280
281     # object. TODO. multiple objects
282     if do_object:
283         if do_constraint_clear:
284             while obj.constraints:
285                 obj.constraints.remove(obj.constraints[0])
286
287         # create compatible eulers
288         euler_prev = None
289
290         for (f, matrix) in obj_info:
291             name = "Action Bake"  # XXX: placeholder
292             obj.matrix_basis = matrix
293
294             obj.keyframe_insert("location", -1, f, name, options)
295
296             rotation_mode = obj.rotation_mode
297             if rotation_mode == 'QUATERNION':
298                 obj.keyframe_insert("rotation_quaternion", -1, f, name, options)
299             elif rotation_mode == 'AXIS_ANGLE':
300                 obj.keyframe_insert("rotation_axis_angle", -1, f, name, options)
301             else:  # euler, XYZ, ZXY etc
302                 if euler_prev is not None:
303                     euler = obj.rotation_euler.copy()
304                     euler.make_compatible(euler_prev)
305                     obj.rotation_euler = euler
306                     euler_prev = euler
307                     del euler
308                 else:
309                     euler_prev = obj.rotation_euler.copy()
310                 obj.keyframe_insert("rotation_euler", -1, f, name, options)
311
312             obj.keyframe_insert("scale", -1, f, name, options)
313
314         if do_parents_clear:
315             obj.parent = None
316
317     # -------------------------------------------------------------------------
318     # Clean
319
320     if do_clean:
321         for fcu in action.fcurves:
322             fcu_orig_data = clean_orig_data.get(fcu, set())
323
324             keyframe_points = fcu.keyframe_points
325             i = 1
326             while i < len(keyframe_points) - 1:
327                 val = keyframe_points[i].co[1]
328
329                 if val in fcu_orig_data:
330                     i += 1
331                     continue
332
333                 val_prev = keyframe_points[i - 1].co[1]
334                 val_next = keyframe_points[i + 1].co[1]
335
336                 if abs(val - val_prev) + abs(val - val_next) < 0.0001:
337                     keyframe_points.remove(keyframe_points[i])
338                 else:
339                     i += 1
340
341     yield action