Spaces:
Running
Running
| import torch | |
| def foot_contact_by_height(pos): | |
| eps = 0.25 | |
| return (-eps < pos[..., 1]) * (pos[..., 1] < eps) | |
| def velocity(pos, padding=False): | |
| velo = pos[1:, ...] - pos[:-1, ...] | |
| velo_norm = torch.norm(velo, dim=-1) | |
| if padding: | |
| pad = torch.zeros_like(velo_norm[:1, :]) | |
| velo_norm = torch.cat([pad, velo_norm], dim=0) | |
| return velo_norm | |
| def foot_contact(pos, ref_height=1., threshold=0.018): | |
| velo_norm = velocity(pos) | |
| contact = velo_norm < threshold | |
| contact = contact.int() | |
| padding = torch.zeros_like(contact) | |
| contact = torch.cat([padding[:1, :], contact], dim=0) | |
| return contact | |
| def alpha(t): | |
| return 2.0 * t * t * t - 3.0 * t * t + 1 | |
| def lerp(a, l, r): | |
| return (1 - a) * l + a * r | |
| def constrain_from_contact(contact, glb, fid='TBD', L=5): | |
| """ | |
| :param contact: contact label | |
| :param glb: original global position | |
| :param fid: joint id to fix, corresponding to the order in contact | |
| :param L: frame to look forward/backward | |
| :return: | |
| """ | |
| T = glb.shape[0] | |
| for i, fidx in enumerate(fid): # fidx: index of the foot joint | |
| fixed = contact[:, i] # [T] | |
| s = 0 | |
| while s < T: | |
| while s < T and fixed[s] == 0: | |
| s += 1 | |
| if s >= T: | |
| break | |
| t = s | |
| avg = glb[t, fidx].clone() | |
| while t + 1 < T and fixed[t + 1] == 1: | |
| t += 1 | |
| avg += glb[t, fidx].clone() | |
| avg /= (t - s + 1) | |
| for j in range(s, t + 1): | |
| glb[j, fidx] = avg.clone() | |
| s = t + 1 | |
| for s in range(T): | |
| if fixed[s] == 1: | |
| continue | |
| l, r = None, None | |
| consl, consr = False, False | |
| for k in range(L): | |
| if s - k - 1 < 0: | |
| break | |
| if fixed[s - k - 1]: | |
| l = s - k - 1 | |
| consl = True | |
| break | |
| for k in range(L): | |
| if s + k + 1 >= T: | |
| break | |
| if fixed[s + k + 1]: | |
| r = s + k + 1 | |
| consr = True | |
| break | |
| if not consl and not consr: | |
| continue | |
| if consl and consr: | |
| litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)), | |
| glb[s, fidx], glb[l, fidx]) | |
| ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)), | |
| glb[s, fidx], glb[r, fidx]) | |
| itp = lerp(alpha(1.0 * (s - l + 1) / (r - l + 1)), | |
| ritp, litp) | |
| glb[s, fidx] = itp.clone() | |
| continue | |
| if consl: | |
| litp = lerp(alpha(1.0 * (s - l + 1) / (L + 1)), | |
| glb[s, fidx], glb[l, fidx]) | |
| glb[s, fidx] = litp.clone() | |
| continue | |
| if consr: | |
| ritp = lerp(alpha(1.0 * (r - s + 1) / (L + 1)), | |
| glb[s, fidx], glb[r, fidx]) | |
| glb[s, fidx] = ritp.clone() | |
| return glb | |