Spaces:
Sleeping
Sleeping
File size: 3,529 Bytes
b03a8f2 |
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 |
import os
import numpy as np
import torch
import smplx
from tqdm import tqdm
root_dir = './BEAT2/beat_english_v2.0.0/smplxflame_30'
output_dir = "./BEAT2/beat_english_v2.0.0/footcontact"
os.makedirs(output_dir, exist_ok=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = smplx_model = smplx.create(
"./emage_evaltools/smplx_models/",
model_type='smplx',
gender='NEUTRAL_2020',
use_face_contour=False,
num_betas=300,
num_expression_coeffs=100,
ext='npz',
use_pca=False,
).eval().to(device)
max_length = 128
for data_file in tqdm(os.listdir(root_dir)):
if not data_file.endswith(".npz"):
continue
data = np.load(os.path.join(root_dir, data_file), allow_pickle=True)
betas = data["betas"]
poses = data["poses"]
trans = data["trans"]
if "expressions" in data:
exps = data["expressions"]
else:
exps = np.zeros((poses.shape[0], 100))
n, c = poses.shape
betas = betas.reshape(1, 300)
betas = np.tile(betas, (n, 1))
betas = torch.from_numpy(betas).float().to(device)
poses = torch.from_numpy(poses.reshape(n, c)).float().to(device)
exps = torch.from_numpy(exps.reshape(n, 100)).float().to(device)
trans = torch.from_numpy(trans.reshape(n, 3)).float().to(device)
s, r = n//max_length, n%max_length
all_tensor = []
for i in range(s):
with torch.no_grad():
joints = model(
betas=betas[i*max_length:(i+1)*max_length],
transl=trans[i*max_length:(i+1)*max_length],
expression=exps[i*max_length:(i+1)*max_length],
jaw_pose=poses[i*max_length:(i+1)*max_length,66:69],
global_orient=poses[i*max_length:(i+1)*max_length,:3],
body_pose=poses[i*max_length:(i+1)*max_length,3:66],
left_hand_pose=poses[i*max_length:(i+1)*max_length,75:120],
right_hand_pose=poses[i*max_length:(i+1)*max_length,120:165],
leye_pose=poses[i*max_length:(i+1)*max_length,69:72],
reye_pose=poses[i*max_length:(i+1)*max_length,72:75],
return_joints=True
)['joints'][:, (7,8,10,11), :].reshape(max_length, 4, 3).cpu()
all_tensor.append(joints)
if r != 0:
with torch.no_grad():
joints = model(
betas=betas[s*max_length:s*max_length+r],
transl=trans[s*max_length:s*max_length+r],
expression=exps[s*max_length:s*max_length+r],
jaw_pose=poses[s*max_length:s*max_length+r,66:69],
global_orient=poses[s*max_length:s*max_length+r,:3],
body_pose=poses[s*max_length:s*max_length+r,3:66],
left_hand_pose=poses[s*max_length:s*max_length+r,75:120],
right_hand_pose=poses[s*max_length:s*max_length+r,120:165],
leye_pose=poses[s*max_length:s*max_length+r,69:72],
reye_pose=poses[s*max_length:s*max_length+r,72:75],
return_joints=True
)['joints'][:, (7,8,10,11), :].reshape(r, 4, 3).cpu()
all_tensor.append(joints)
joints = torch.cat(all_tensor, axis=0)
feetv = torch.zeros(joints.shape[1], joints.shape[0])
joints = joints.permute(1, 0, 2)
feetv[:, :-1] = (joints[:, 1:] - joints[:, :-1]).norm(dim=-1)
contacts = (feetv < 0.01).numpy().astype(float)
contacts = contacts.transpose(1, 0)
np.save(os.path.join(output_dir, data_file.replace(".npz", ".npy")), contacts)
|