remove bl_operators/nla.py, move bake_action function into bpy_extras.anim_utils...
[blender.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-80 compliant>
20
21 __all__ = (
22     "bake_action",
23     )
24
25 import bpy
26
27
28 def bake_action(frame_start,
29                 frame_end,
30                 frame_step=1,
31                 only_selected=False,
32                 do_pose=True,
33                 do_object=True,
34                 do_constraint_clear=False,
35                 do_clean=False,
36                 action=None,
37                 ):
38
39     """
40     Return an image from the file path with options to search multiple paths
41     and return a placeholder if its not found.
42
43     :arg frame_start: First frame to bake.
44     :type frame_start: int
45     :arg frame_end: Last frame to bake.
46     :type frame_end: int
47     :arg frame_step: Frame step.
48     :type frame_step: int
49     :arg only_selected: Only bake selected data.
50     :type only_selected: bool
51     :arg do_pose: Bake pose channels.
52     :type do_pose: bool
53     :arg do_object: Bake objects.
54     :type do_object: bool
55     :arg do_constraint_clear: Remove constraints.
56     :type do_constraint_clear: bool
57     :arg do_clean: Remove redundant keyframes after baking.
58     :type do_clean: bool
59     :arg action: An action to bake the data into, or None for a new action
60        to be created.
61     :type action: :class:`bpy.types.Action` or None
62     
63     :return: an action or None
64     :rtype: :class:`bpy.types.Action`
65     """
66
67     # -------------------------------------------------------------------------
68     # Helper Functions
69
70     def pose_frame_info(obj):
71         from mathutils import Matrix
72
73         info = {}
74
75         pose = obj.pose
76
77         pose_items = pose.bones.items()
78
79         for name, pbone in pose_items:
80             binfo = {}
81             bone = pbone.bone
82
83             binfo["parent"] = getattr(bone.parent, "name", None)
84             binfo["bone"] = bone
85             binfo["pbone"] = pbone
86             binfo["matrix_local"] = bone.matrix_local.copy()
87             try:
88                 binfo["matrix_local_inv"] = binfo["matrix_local"].inverted()
89             except:
90                 binfo["matrix_local_inv"] = Matrix()
91
92             binfo["matrix"] = bone.matrix.copy()
93             binfo["matrix_pose"] = pbone.matrix.copy()
94             try:
95                 binfo["matrix_pose_inv"] = binfo["matrix_pose"].inverted()
96             except:
97                 binfo["matrix_pose_inv"] = Matrix()
98
99             info[name] = binfo
100
101         for name, pbone in pose_items:
102             binfo = info[name]
103             binfo_parent = binfo.get("parent", None)
104             if binfo_parent:
105                 binfo_parent = info[binfo_parent]
106
107             matrix = binfo["matrix_pose"]
108             rest_matrix = binfo["matrix_local"]
109
110             if binfo_parent:
111                 matrix = binfo_parent["matrix_pose_inv"] * matrix
112                 rest_matrix = binfo_parent["matrix_local_inv"] * rest_matrix
113
114             binfo["matrix_key"] = rest_matrix.inverted() * matrix
115
116         return info
117
118
119     def obj_frame_info(obj):
120         info = {}
121         # parent = obj.parent
122         info["matrix_key"] = obj.matrix_local.copy()
123         return info
124
125     # -------------------------------------------------------------------------
126     # Setup the Context
127
128     # TODO, pass data rather then grabbing from the context!
129     scene = bpy.context.scene
130     obj = bpy.context.object
131     pose = obj.pose
132     frame_back = scene.frame_current
133
134     if pose is None:
135         do_pose = False
136
137     if do_pose is None and do_object is None:
138         return None
139
140     pose_info = []
141     obj_info = []
142
143     frame_range = range(frame_start, frame_end + 1, frame_step)
144
145     # -------------------------------------------------------------------------
146     # Collect transformations
147
148     # could speed this up by applying steps here too...
149     for f in frame_range:
150         scene.frame_set(f)
151
152         if do_pose:
153             pose_info.append(pose_frame_info(obj))
154         if do_object:
155             obj_info.append(obj_frame_info(obj))
156
157         f += 1
158
159     # -------------------------------------------------------------------------
160     # Create action
161
162     # incase animation data hassnt been created
163     atd = obj.animation_data_create()
164     if action is None:
165         action = bpy.data.actions.new("Action")
166     atd.action = action
167
168     if do_pose:
169         pose_items = pose.bones.items()
170     else:
171         pose_items = []  # skip
172
173     # -------------------------------------------------------------------------
174     # Apply transformations to action
175
176     # pose
177     for name, pbone in (pose_items if do_pose else ()):
178         if only_selected and not pbone.bone.select:
179             continue
180
181         if do_constraint_clear:
182             while pbone.constraints:
183                 pbone.constraints.remove(pbone.constraints[0])
184
185         for f in frame_range:
186             matrix = pose_info[(f - frame_start) // frame_step][name]["matrix_key"]
187
188             # pbone.location = matrix.to_translation()
189             # pbone.rotation_quaternion = matrix.to_quaternion()
190             pbone.matrix_basis = matrix
191
192             pbone.keyframe_insert("location", -1, f, name)
193
194             rotation_mode = pbone.rotation_mode
195
196             if rotation_mode == 'QUATERNION':
197                 pbone.keyframe_insert("rotation_quaternion", -1, f, name)
198             elif rotation_mode == 'AXIS_ANGLE':
199                 pbone.keyframe_insert("rotation_axis_angle", -1, f, name)
200             else:  # euler, XYZ, ZXY etc
201                 pbone.keyframe_insert("rotation_euler", -1, f, name)
202
203             pbone.keyframe_insert("scale", -1, f, name)
204
205     # object. TODO. multiple objects
206     if do_object:
207         if do_constraint_clear:
208             while obj.constraints:
209                 obj.constraints.remove(obj.constraints[0])
210
211         for f in frame_range:
212             matrix = obj_info[(f - frame_start) // frame_step]["matrix_key"]
213             obj.matrix_local = matrix
214
215             obj.keyframe_insert("location", -1, f)
216
217             rotation_mode = obj.rotation_mode
218
219             if rotation_mode == 'QUATERNION':
220                 obj.keyframe_insert("rotation_quaternion", -1, f)
221             elif rotation_mode == 'AXIS_ANGLE':
222                 obj.keyframe_insert("rotation_axis_angle", -1, f)
223             else:  # euler, XYZ, ZXY etc
224                 obj.keyframe_insert("rotation_euler", -1, f)
225
226             obj.keyframe_insert("scale", -1, f)
227
228     scene.frame_set(frame_back)
229
230     # -------------------------------------------------------------------------
231     # Clean
232
233     if do_clean:
234         for fcu in action.fcurves:
235             keyframe_points = fcu.keyframe_points
236             i = 1
237             while i < len(fcu.keyframe_points) - 1:
238                 val_prev = keyframe_points[i - 1].co[1]
239                 val_next = keyframe_points[i + 1].co[1]
240                 val = keyframe_points[i].co[1]
241
242                 if abs(val - val_prev) + abs(val - val_next) < 0.0001:
243                     keyframe_points.remove(keyframe_points[i])
244                 else:
245                     i += 1
246
247     return action