pep8 cleanup and fix for keymap test operator from my own recent fix.
[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     def obj_frame_info(obj):
119         info = {}
120         # parent = obj.parent
121         info["matrix_key"] = obj.matrix_local.copy()
122         return info
123
124     # -------------------------------------------------------------------------
125     # Setup the Context
126
127     # TODO, pass data rather then grabbing from the context!
128     scene = bpy.context.scene
129     obj = bpy.context.object
130     pose = obj.pose
131     frame_back = scene.frame_current
132
133     if pose is None:
134         do_pose = False
135
136     if do_pose is None and do_object is None:
137         return None
138
139     pose_info = []
140     obj_info = []
141
142     frame_range = range(frame_start, frame_end + 1, frame_step)
143
144     # -------------------------------------------------------------------------
145     # Collect transformations
146
147     # could speed this up by applying steps here too...
148     for f in frame_range:
149         scene.frame_set(f)
150
151         if do_pose:
152             pose_info.append(pose_frame_info(obj))
153         if do_object:
154             obj_info.append(obj_frame_info(obj))
155
156         f += 1
157
158     # -------------------------------------------------------------------------
159     # Create action
160
161     # incase animation data hassnt been created
162     atd = obj.animation_data_create()
163     if action is None:
164         action = bpy.data.actions.new("Action")
165     atd.action = action
166
167     if do_pose:
168         pose_items = pose.bones.items()
169     else:
170         pose_items = []  # skip
171
172     # -------------------------------------------------------------------------
173     # Apply transformations to action
174
175     # pose
176     for name, pbone in (pose_items if do_pose else ()):
177         if only_selected and not pbone.bone.select:
178             continue
179
180         if do_constraint_clear:
181             while pbone.constraints:
182                 pbone.constraints.remove(pbone.constraints[0])
183
184         for f in frame_range:
185             f_step = (f - frame_start) // frame_step
186             matrix = pose_info[f_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