52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
import os
|
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '5'
|
|
|
|
import numpy as np
|
|
from smpl_tf import SMPL as SMPL_tf
|
|
from smpl_torch import SMPL as SMPL_torch
|
|
from utils import plot_mesh, plot_mesh_2
|
|
from time import time
|
|
|
|
import torch
|
|
import tensorflow as tf
|
|
|
|
if __name__ == '__main__':
|
|
device = torch.device('cuda')
|
|
|
|
smpl_tf : SMPL_tf = SMPL_tf('model_files')
|
|
smpl_torch : SMPL_torch = SMPL_torch(device=device, path='model_files')
|
|
|
|
|
|
pose = ((np.random.rand(4, 72) - 0.5) * 0.4).astype(np.float32)
|
|
betas = ((np.random.rand(4, 10) - 0.5) * 0.06).astype(np.float32)
|
|
betas = np.concatenate((betas, np.array([0,1,0,1], dtype=np.float32)[:,None]), axis=1)
|
|
|
|
# pose = np.zeros((4, 72), dtype=np.float32)
|
|
# betas = np.zeros((4, 11), dtype=np.float32)
|
|
|
|
pose_torch = torch.tensor(pose, device=device)
|
|
betas_torch = torch.tensor(betas, device=device)
|
|
|
|
result_torch, joints_torch, _, _ = smpl_torch(betas_torch, pose_torch)
|
|
s = time()
|
|
result_torch, joints_torch, v_shaped_1, v_posed_1 = smpl_torch(betas_torch, pose_torch)
|
|
print(time() - s)
|
|
R_1 = smpl_torch.rodrigues(pose_torch.view(-1, 1, 3)).reshape(4, -1, 3, 3)
|
|
|
|
# plot_mesh(result[0].cpu().numpy(), model.f).show()
|
|
|
|
R = smpl_tf.run_rodrigues(pose)
|
|
result_tf, joints_torch, _, _ = smpl_tf.run(R, betas)
|
|
s = time()
|
|
R = smpl_tf.run_rodrigues(pose)
|
|
result_tf, joints_torch, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
|
|
print(time() - s)
|
|
|
|
result_torch = result_torch.cpu().numpy()
|
|
result_tf = result_tf.numpy()
|
|
diff = np.abs(result_torch - result_tf).sum(-1).sum(-1)
|
|
|
|
plot_mesh_2(result_torch[0], result_tf[0], smpl_tf.f).show()
|
|
plot_mesh_2(result_torch[1], result_tf[1], smpl_tf.f).show()
|
|
pass |