923ca92a16297d804f505838d2389c36b5acb18f
[blender-staging.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 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"].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         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.inverted() * 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     scene = bpy.context.scene
79     obj = bpy.context.object
80     pose = obj.pose
81
82     info_ls = []
83
84     frame_range = range(frame_start, frame_end + 1, step)
85
86     # could spped this up by applying steps here too...
87     for f in frame_range:
88         scene.frame_set(f)
89
90         info = pose_info()
91         info_ls.append(info)
92         f += 1
93
94     action = bpy.data.actions.new("Action")
95
96     bpy.context.object.animation_data.action = action
97
98     pose_items = pose.bones.items()
99
100     for name, pbone in pose_items:
101         if only_selected and not pbone.bone.select:
102             continue
103
104         for f in frame_range:
105             matrix = info_ls[int((f - frame_start) / step)][name]["matrix_key"]
106
107             #pbone.location = matrix.to_translation()
108             #pbone.rotation_quaternion = matrix.to_quaternion()
109             pbone.matrix_basis = matrix
110
111             pbone.keyframe_insert("location", -1, f, name)
112
113             rotation_mode = pbone.rotation_mode
114
115             if rotation_mode == 'QUATERNION':
116                 pbone.keyframe_insert("rotation_quaternion", -1, f, name)
117             elif rotation_mode == 'AXIS_ANGLE':
118                 pbone.keyframe_insert("rotation_axis_angle", -1, f, name)
119             else:  # euler, XYZ, ZXY etc
120                 pbone.keyframe_insert("rotation_euler", -1, f, name)
121
122             pbone.keyframe_insert("scale", -1, f, name)
123
124     return action
125
126
127 from bpy.props import IntProperty, BoolProperty
128
129
130 class BakeAction(bpy.types.Operator):
131     '''Bake animation to an Action'''
132     bl_idname = "nla.bake"
133     bl_label = "Bake Action"
134     bl_options = {'REGISTER', 'UNDO'}
135
136     frame_start = IntProperty(name="Start Frame",
137             description="Start frame for baking",
138             default=1, min=1, max=300000)
139     frame_end = IntProperty(name="End Frame",
140             description="End frame for baking",
141             default=250, min=1, max=300000)
142     step = IntProperty(name="Frame Step",
143             description="Frame Step",
144             default=1, min=1, max=120)
145     only_selected = BoolProperty(name="Only Selected",
146             default=True)
147
148     def execute(self, context):
149
150         action = bake(self.frame_start, self.frame_end, self.step, self.only_selected)
151
152         # basic cleanup, could move elsewhere
153         for fcu in action.fcurves:
154             keyframe_points = fcu.keyframe_points
155             i = 1
156             while i < len(fcu.keyframe_points) - 1:
157                 val_prev = keyframe_points[i - 1].co[1]
158                 val_next = keyframe_points[i + 1].co[1]
159                 val = keyframe_points[i].co[1]
160
161                 if abs(val - val_prev) + abs(val - val_next) < 0.0001:
162                     keyframe_points.remove(keyframe_points[i])
163                 else:
164                     i += 1
165
166         return {'FINISHED'}
167
168     def invoke(self, context, event):
169         wm = context.window_manager
170         return wm.invoke_props_dialog(self)