-
Notifications
You must be signed in to change notification settings - Fork 1
/
ik_baker.py
237 lines (201 loc) · 10.8 KB
/
ik_baker.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
import math
import bpy
from functools import lru_cache
from collections import defaultdict
import time
START_IDX = 0
END_IDX = 1
def get_ik_targets(obj):
"""
Returns a list of all the IK targets of the given object.
Note that this should be ordered so that the root bones are first, then all the other ik targets are last.
"""
for ik_target in obj.ik_targets:
if ik_target.target_bone == "b__ROOT__":
yield ik_target
for ik_target in obj.ik_targets:
if ik_target.target_bone == "b__ROOT__Adjust":
yield ik_target
for ik_target in obj.ik_targets:
if "__subroot__" in ik_target.target_bone:
yield ik_target
for ik_target in obj.ik_targets:
if ik_target.target_bone != "b__ROOT__" and ik_target.target_bone != "b__ROOT__Adjust" and "__subroot__" not in\
ik_target.target_bone:
yield ik_target
class s4animtool_OT_bakeik(bpy.types.Operator):
"""Bake the IK weights"""
bl_idname = "s4animtools.bakeik"
bl_label = "Bake ik"
bl_options = {"REGISTER", "UNDO"}
def get_keyframes(self, obj, data_path):
keyframes = []
anim = obj.animation_data
if anim is not None and anim.action is not None:
for fcu in anim.action.fcurves:
if fcu.data_path == data_path:
for keyframe in fcu.keyframe_points:
x, y = keyframe.co
if x not in keyframes:
keyframes.append((math.ceil(x)))
return keyframes
def set_interpolation_keyframes(self, obj, bone):
"""
IK Ranges are always a set of two keyframes, start and end.
What this does is find the previous and next ik weights, and set the ik
weight to be zero on those for a smooth blend.
"""
frames = {}
for weight_idx in range(11):
keyframes = self.get_keyframes(obj, f'pose.bones["{bone.name}"].ik_weight_{weight_idx}')
if len(keyframes) % 2 != 0 and len(keyframes) != 1:
raise Exception(f"IK Ranges on ik weights aren't even or a single frame... somehow. {len(keyframes)}")
else:
if len(keyframes) == 1:
continue
if len(keyframes) == 0:
continue
for i in range(0, len(keyframes), 2):
frames[(bone.name, weight_idx, math.floor(i/2))] = (keyframes[i], keyframes[i+1])
#print(len(frames[(bone.name,weight_idx)]))
frames_ordered = sorted(frames,key=lambda k: frames[k][0])
#TODO what is this pls rewrite
for idx, item in enumerate(frames_ordered):
bone, weight_idx, range_idx = item
if idx == 0:
pre_start_frame = -1
else:
pre_start_frame = frames[frames_ordered[idx - 1]][END_IDX]
# TODO ugly hack
if pre_start_frame == 0:
pre_start_frame = -1
"""
Set post end frame to beyond the end of a clip if the frame
it's requesting is beyond the current clip range. +1 offset too for some reason
"""
if idx + 1 >= len(frames_ordered):
post_end_frame = frames[frames_ordered[idx]][END_IDX] + 1
else:
"""
Get the next ik target and set the weight to be zero on this ik target weight.
"""
next_ik_target_range = frames_ordered[idx + 1]
post_end_frame = frames[next_ik_target_range][START_IDX]
data_path = f'pose.bones["{bone}"].ik_weight_{weight_idx}'
fc = obj.animation_data.action.fcurves.find(data_path)
fc.keyframe_points.insert(pre_start_frame, 0)
fc.keyframe_points.insert(post_end_frame, 0)
fc.update()
def execute(self, context):
t1 = time.time()
obj = context.object
bones_to_interpolate = []
positions = defaultdict(list)
quaternions = defaultdict(list)
if obj.ik_idx >= 0 and obj.ik_targets:
s4animtool_OT_bakeik.remove_IK(obj)
for idx, item in enumerate(get_ik_targets(obj)):
chain_bone = obj.pose.bones[item.chain_bone]
chain_idx = item.chain_idx
target_rig = bpy.data.objects[item.target_obj]
target_bone = target_rig.pose.bones[item.target_bone]
quaternions[(chain_bone, target_bone, target_rig)] = defaultdict(list)
positions[(chain_bone, target_bone, target_rig)] = defaultdict(list)
if obj.ik_idx >= 0 and obj.ik_targets:
for i in range(context.scene.frame_start, context.scene.frame_end, 1):
bpy.context.scene.frame_set(i)
for idx, item in enumerate(get_ik_targets(obj)):
chain_bone, target_bone, target_rig = self.get_bonedata(item, obj)
src_matrix = target_rig.matrix_world @ target_bone.matrix
target_matrix = obj.matrix_world @ chain_bone.matrix
matrix_data = src_matrix.inverted() @ target_matrix
rotation_data = matrix_data.to_quaternion()
translation_data = matrix_data.to_translation()
if i == 0:
print(translation_data, rotation_data)
for axis_idx, value in enumerate(translation_data):
positions[(chain_bone, target_bone, target_rig)][axis_idx].extend([i, value])
for axis_idx, value in enumerate(rotation_data):
quaternions[(chain_bone, target_bone, target_rig)][axis_idx].extend([i, value])
for potential_ik_bone in obj.pose.bones:
for idx in range(11):
for pos_idx in range(3):
data_path = f'pose.bones["{potential_ik_bone.name}"].ik_pos_{idx}'
old_fc = obj.animation_data.action.fcurves.find(data_path, index=pos_idx)
if old_fc is not None:
obj.animation_data.action.fcurves.remove(old_fc)
for rot_idx in range(4):
data_path = f'pose.bones["{potential_ik_bone.name}"].ik_rot_{idx}'
old_fc = obj.animation_data.action.fcurves.find(data_path, index=rot_idx)
if old_fc is not None:
obj.animation_data.action.fcurves.remove(old_fc)
data_path = f'pose.bones["{potential_ik_bone.name}"].ik_weight_{idx}'
s4animtool_OT_bakeik.find_and_remove(data_path, obj)
current_bone_idx = defaultdict(int)
for idx, item in enumerate(get_ik_targets(obj)):
chain_bone, target_bone, target_rig = self.get_bonedata(item, obj)
for pos_idx in range(3):
data_path = f'pose.bones["{chain_bone.name}"].ik_pos_{current_bone_idx[chain_bone]}'
fc = obj.animation_data.action.fcurves.new(data_path, index=pos_idx)
#print(positions[(chain_bone, target_bone, target_rig)][pos_idx])
fc.keyframe_points.add(round(len(positions[(chain_bone, target_bone, target_rig)][pos_idx]) / 2))
fc.keyframe_points.foreach_set("co", positions[(chain_bone, target_bone, target_rig)][pos_idx])
fc.update()
for rot_idx in range(4):
data_path = f'pose.bones["{chain_bone.name}"].ik_rot_{current_bone_idx[chain_bone]}'
fc = obj.animation_data.action.fcurves.new(data_path, index=rot_idx)
fc.keyframe_points.add(round(len(quaternions[(chain_bone, target_bone, target_rig)][rot_idx]) / 2))
fc.keyframe_points.foreach_set("co", quaternions[(chain_bone, target_bone, target_rig)][rot_idx])
fc.update()
current_bone_idx[chain_bone] += 1
current_bone_idx = defaultdict(int)
if obj.ik_idx >= 0 and obj.ik_targets:
for idx, item in enumerate(get_ik_targets(obj)):
chain_bone = obj.pose.bones[item.chain_bone]
chain_idx = item.chain_idx
target_rig = bpy.data.objects[item.target_obj]
data_path = f'pose.bones["{chain_bone.name}"].ik_weight_{current_bone_idx[chain_bone]}'
s4animtool_OT_bakeik.find_and_remove(data_path, obj)
fc = obj.animation_data.action.fcurves.new(data_path)
for range_value in item.ranges:
if range_value.start_time == range_value.end_time:
if range_value.start_time == 0:
fc.keyframe_points.insert(-1, 0)
fc.keyframe_points.insert(0, 0)
else:
fc.keyframe_points.insert(range_value.start_time, 1)
fc.keyframe_points.insert(range_value.end_time, 1)
fc.update()
if chain_bone not in bones_to_interpolate:
bones_to_interpolate.append(chain_bone)
current_bone_idx[chain_bone] += 1
for bone in bones_to_interpolate:
self.set_interpolation_keyframes(obj, bone)
bpy.context.view_layer.objects.active = obj
t2 = time.time()
print(f"Took {t2-t1} seconds for ik baking")
return {'FINISHED'}
def create_ik_target_name(self, chain_bone, obj, target_bone, target_rig, chain_idx):
ik_target_empty_name = "{} IK {} {} {} {}".format(obj.name.strip(), chain_idx, chain_bone.name,
target_rig.name, target_bone.name)
return ik_target_empty_name
def get_bonedata(self, item, obj):
chain_bone = obj.pose.bones[item.chain_bone]
target_rig = bpy.data.objects[item.target_obj]
target_bone = target_rig.pose.bones[item.target_bone]
return chain_bone, target_bone, target_rig
@staticmethod
def remove_IK(obj):
for bone in obj.pose.bones:
for weight_idx in range(0, 9):
print(weight_idx)
s4animtool_OT_bakeik.find_and_remove(f'pose.bones["{bone.name}"].ik_weight_{weight_idx}', obj)
for i in range(3):
s4animtool_OT_bakeik.find_and_remove(f'pose.bones["{bone.name}"].ik_pos_{weight_idx}', obj, index=i)
for i in range(4):
s4animtool_OT_bakeik.find_and_remove(f'pose.bones["{bone.name}"].ik_rot_{weight_idx}', obj, index=i)
@staticmethod
def find_and_remove(data_path, obj, index=0):
old_fc = obj.animation_data.action.fcurves.find(data_path, index=index)
if old_fc is not None:
obj.animation_data.action.fcurves.remove(old_fc)