fix [#28413] bpy.ops.nla.bake can't bake from frame 0
[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          action=None,
88          ):
89
90     scene = bpy.context.scene
91     obj = bpy.context.object
92     pose = obj.pose
93     frame_back = scene.frame_current
94
95     if pose is None:
96         do_pose = False
97
98     if do_pose is None and do_object is None:
99         return None
100
101     pose_info = []
102     obj_info = []
103
104     frame_range = range(frame_start, frame_end + 1, step)
105
106     # -------------------------------------------------------------------------
107     # Collect transformations
108
109     # could speed this up by applying steps here too...
110     for f in frame_range:
111         scene.frame_set(f)
112
113         if do_pose:
114             pose_info.append(pose_frame_info(obj))
115         if do_object:
116             obj_info.append(obj_frame_info(obj))
117
118         f += 1
119
120     # -------------------------------------------------------------------------
121     # Create action
122
123     # incase animation data hassnt been created
124     atd = obj.animation_data_create()
125     if action is None:
126         action = bpy.data.actions.new("Action")
127     atd.action = action
128
129     if do_pose:
130         pose_items = pose.bones.items()
131     else:
132         pose_items = []  # skip
133
134     # -------------------------------------------------------------------------
135     # Apply transformations to action
136
137     # pose
138     for name, pbone in (pose_items if do_pose else ()):
139         if only_selected and not pbone.bone.select:
140             continue
141
142         if do_constraint_clear:
143             while pbone.constraints:
144                 pbone.constraints.remove(pbone.constraints[0])
145
146         for f in frame_range:
147             matrix = pose_info[(f - frame_start) // step][name]["matrix_key"]
148
149             # pbone.location = matrix.to_translation()
150             # pbone.rotation_quaternion = matrix.to_quaternion()
151             pbone.matrix_basis = matrix
152
153             pbone.keyframe_insert("location", -1, f, name)
154
155             rotation_mode = pbone.rotation_mode
156
157             if rotation_mode == 'QUATERNION':
158                 pbone.keyframe_insert("rotation_quaternion", -1, f, name)
159             elif rotation_mode == 'AXIS_ANGLE':
160                 pbone.keyframe_insert("rotation_axis_angle", -1, f, name)
161             else:  # euler, XYZ, ZXY etc
162                 pbone.keyframe_insert("rotation_euler", -1, f, name)
163
164             pbone.keyframe_insert("scale", -1, f, name)
165
166     # object. TODO. multiple objects
167     if do_object:
168         if do_constraint_clear:
169             while obj.constraints:
170                 obj.constraints.remove(obj.constraints[0])
171
172         for f in frame_range:
173             matrix = obj_info[(f - frame_start) // step]["matrix_key"]
174             obj.matrix_local = matrix
175
176             obj.keyframe_insert("location", -1, f)
177
178             rotation_mode = obj.rotation_mode
179
180             if rotation_mode == 'QUATERNION':
181                 obj.keyframe_insert("rotation_quaternion", -1, f)
182             elif rotation_mode == 'AXIS_ANGLE':
183                 obj.keyframe_insert("rotation_axis_angle", -1, f)
184             else:  # euler, XYZ, ZXY etc
185                 obj.keyframe_insert("rotation_euler", -1, f)
186
187             obj.keyframe_insert("scale", -1, f)
188
189     scene.frame_set(frame_back)
190
191     return action
192
193
194 from bpy.props import IntProperty, BoolProperty, EnumProperty
195
196
197 class BakeAction(Operator):
198     '''Bake animation to an Action'''
199     bl_idname = "nla.bake"
200     bl_label = "Bake Action"
201     bl_options = {'REGISTER', 'UNDO'}
202
203     frame_start = IntProperty(
204             name="Start Frame",
205             description="Start frame for baking",
206             min=0, max=300000,
207             default=1,
208             )
209     frame_end = IntProperty(
210             name="End Frame",
211             description="End frame for baking",
212             min=1, max=300000,
213             default=250,
214             )
215     step = IntProperty(
216             name="Frame Step",
217             description="Frame Step",
218             min=1, max=120,
219             default=1,
220             )
221     only_selected = BoolProperty(
222             name="Only Selected",
223             default=True,
224             )
225     clear_consraints = BoolProperty(
226             name="Clear Constraints",
227             default=False,
228             )
229     bake_types = EnumProperty(
230             name="Bake Data",
231             options={'ENUM_FLAG'},
232             items=(('POSE', "Pose", ""),
233                    ('OBJECT', "Object", ""),
234                    ),
235             default={'POSE'},
236             )
237
238     def execute(self, context):
239
240         action = bake(self.frame_start,
241                       self.frame_end,
242                       self.step,
243                       self.only_selected,
244                       'POSE' in self.bake_types,
245                       'OBJECT' in self.bake_types,
246                       self.clear_consraints,
247                       )
248
249         if action is None:
250             self.report({'INFO'}, "Nothing to bake")
251             return {'CANCELLED'}
252
253         # basic cleanup, could move elsewhere
254         for fcu in action.fcurves:
255             keyframe_points = fcu.keyframe_points
256             i = 1
257             while i < len(fcu.keyframe_points) - 1:
258                 val_prev = keyframe_points[i - 1].co[1]
259                 val_next = keyframe_points[i + 1].co[1]
260                 val = keyframe_points[i].co[1]
261
262                 if abs(val - val_prev) + abs(val - val_next) < 0.0001:
263                     keyframe_points.remove(keyframe_points[i])
264                 else:
265                     i += 1
266
267         return {'FINISHED'}
268
269     def invoke(self, context, event):
270         wm = context.window_manager
271         return wm.invoke_props_dialog(self)