use set as a suffix (matches operators)
[blender.git] / release / scripts / op / 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 compliant>
20
21 import bpy
22
23
24 def pose_info():
25     from mathutils import Matrix
26
27     info = {}
28
29     obj = bpy.context.object
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"].copy().invert()
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"].copy().invert()
51         except:
52             binfo["matrix_pose_inv"] = Matrix()
53
54         print(binfo["matrix_pose"])
55         info[name] = binfo
56
57     for name, pbone in pose_items:
58         binfo = info[name]
59         binfo_parent = binfo.get("parent", None)
60         if binfo_parent:
61             binfo_parent = info[binfo_parent]
62
63         matrix = binfo["matrix_pose"]
64         rest_matrix = binfo["matrix_local"]
65
66         if binfo_parent:
67             matrix = binfo_parent["matrix_pose_inv"] * matrix
68             rest_matrix = binfo_parent["matrix_local_inv"] * rest_matrix
69
70         matrix = rest_matrix.copy().invert() * matrix
71
72         binfo["matrix_key"] = matrix.copy()
73
74     return info
75
76
77 def bake(frame_start, frame_end, step=1, only_selected=False):
78     # import nla; reload(nla); nla.bake()
79
80     scene = bpy.context.scene
81     obj = bpy.context.object
82     pose = obj.pose
83
84     info_ls = []
85
86     frame_range = range(frame_start, frame_end + 1, step)
87
88     # could spped this up by applying steps here too...
89     for f in frame_range:
90         scene.frame_set(f)
91
92         info = pose_info()
93         info_ls.append(info)
94         f += 1
95
96     action = bpy.data.actions.new("Action")
97
98     bpy.context.object.animation_data.action = action
99
100     pose_items = pose.bones.items()
101
102     for name, pbone in pose_items:
103         if only_selected and not pbone.select:
104             continue
105
106         for f in frame_range:
107             matrix = info_ls[int((f - frame_start) / step)][name]["matrix_key"]
108
109             #pbone.location = matrix.translation_part()
110             #pbone.rotation_quaternion = matrix.to_quat()
111             pbone.matrix_local = [f for v in matrix for f in v]
112
113             pbone.keyframe_insert("location", -1, f, name)
114
115             rotation_mode = pbone.rotation_mode
116
117             if rotation_mode == 'QUATERNION':
118                 pbone.keyframe_insert("rotation_quaternion", -1, f, name)
119             elif rotation_mode == 'AXIS_ANGLE':
120                 pbone.keyframe_insert("rotation_axis_angle", -1, f, name)
121             else: # euler, XYZ, ZXY etc
122                 pbone.keyframe_insert("rotation_euler", -1, f, name)
123
124             pbone.keyframe_insert("scale", -1, f, name)
125
126     return action
127
128
129 from bpy.props import *
130
131
132 class BakeAction(bpy.types.Operator):
133     '''Bake animation to an Action'''
134     bl_idname = "nla.bake"
135     bl_label = "Bake Action"
136     bl_options = {'REGISTER', 'UNDO'}
137
138     frame_start = IntProperty(name="Start Frame",
139             description="Start frame for baking",
140             default=1, min=1, max=300000)
141     frame_end = IntProperty(name="End Frame",
142             description="End frame for baking",
143             default=250, min=1, max=300000)
144     step = IntProperty(name="Frame Step",
145             description="Frame Step",
146             default=1, min=1, max=120)
147     only_selected = BoolProperty(name="Only Selected",
148             default=True)
149
150     def execute(self, context):
151         props = self.properties
152
153         action = bake(props.frame_start, props.frame_end, props.step, props.show_only_selected)
154
155         # basic cleanup, could move elsewhere
156         for fcu in action.fcurves:
157             keyframe_points = fcu.keyframe_points
158             i = 1
159             while i < len(fcu.keyframe_points) - 1:
160                 val_prev = keyframe_points[i - 1].co[1]
161                 val_next = keyframe_points[i + 1].co[1]
162                 val = keyframe_points[i].co[1]
163
164                 if abs(val - val_prev) + abs(val - val_next) < 0.0001:
165                     keyframe_points.remove(keyframe_points[i])
166                 else:
167                     i += 1
168
169         return {'FINISHED'}
170
171     def invoke(self, context, event):
172         wm = context.window_manager
173         return wm.invoke_props_dialog(self)
174
175
176 #def menu_func(self, context):
177 #    self.layout.operator(BakeAction.bl_idname, text="Bake Armature Action")
178
179
180 def register():
181     pass
182     # bpy.types.INFO_MT_mesh_add.append(menu_func)
183
184
185 def unregister():
186     pass
187     # bpy.types.INFO_MT_mesh_add.remove(menu_func)
188
189 if __name__ == "__main__":
190     register()