luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

optimize smpl

Browse Source
This commit is contained in:
luni
2023-07-03 17:27:34 +02:00
parent 592b8df0f0
commit 25c0a09296
1 changed files with 27 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
smpl/smpl_torch.py
+25 -15
View File
@@ -10,7 +10,7 @@ from time import time
from torch.nn import Module
class SMPL(Module):
def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False):
def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False, batch_size=4):
super(SMPL, self).__init__()
self.J_regressor = []
@@ -56,6 +56,14 @@ class SMPL(Module):
self.shape = (1,27578,3)
self.device = device
self.batch_size = batch_size
self.with_zeros_ones = torch.tensor(
[[[0.0, 0.0, 0.0, 1.0]]], dtype=torch.float32
).expand(self.batch_size,-1,-1).to(self.device)
self.results_zeros = torch.zeros((self.batch_size, 24, 1), dtype=torch.float32).to(self.device)
self.J_ones = torch.ones((self.batch_size, 24, 1), dtype=torch.float32, device=self.device)
self.rest_shape_h_ones_full = torch.ones((self.batch_size, self.weights_all.shape[1], 1), dtype=torch.float32).to(self.device)
if orig_smpl:
self.to_keep = self.to_keep[self.to_keep < 6890]
@@ -67,6 +75,8 @@ class SMPL(Module):
self.weights = self.weights_all[:,self.to_keep,:]
self.shape = (1, 6890, 3)
self.rest_shape_h_ones = torch.ones((self.batch_size, self.to_keep.shape[0], 1), dtype=torch.float32).to(self.device)
@staticmethod
def rodrigues(r):
"""
@@ -91,7 +101,7 @@ class SMPL(Module):
m = torch.stack(
(z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick,
-r_hat[:, 0, 0], -r_hat[:, 0, 1], r_hat[:, 0, 0], z_stick), dim=1)
m = torch.reshape(m, (-1, 3, 3))
m = m.view((-1, 3, 3))
i_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) \
+ torch.zeros((theta_dim, 3, 3), dtype=torch.float32)).to(r.device)
A = r_hat.permute(0, 2, 1)
@@ -158,8 +168,8 @@ class SMPL(Module):
# v_shaped = torch.tensordot(betas, self.shapedirs, dims=([1], [2])) + self.v_template
v_shaped = []
v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + torch.reshape(self.v_template[0], shape=self.shape)) # female
v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + torch.reshape(self.v_template[1], shape=self.shape)) # male
v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[0], dims=[[1], [2]]) + self.v_template[0].view(self.shape)) # female
v_shaped.append(torch.tensordot(betas[:,:10], self.shapedirs[1], dims=[[1], [2]]) + self.v_template[1].view(self.shape)) # male
v_shaped = torch.stack(v_shaped, axis=-1)
# J = torch.matmul(self.J_regressor, v_shaped)
@@ -167,7 +177,7 @@ class SMPL(Module):
J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
else:
J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
J0 = J[:,0,:].view((-1, 1, 3))
J = J - J0
R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
@@ -177,12 +187,12 @@ class SMPL(Module):
R_cube = R_cube_big[:, 1:, :, :]
I_cube = (torch.eye(3, dtype=torch.float32).unsqueeze(dim=0) + \
torch.zeros((batch_num, R_cube.shape[1], 3, 3), dtype=torch.float32)).to(self.device)
lrotmin = (R_cube - I_cube).reshape(batch_num, -1, 1).squeeze(dim=2)
lrotmin = (R_cube - I_cube).view(batch_num, -1, 1).squeeze(dim=2)
v_posed = torch.stack([v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)] + torch.tensordot(lrotmin[i], self.posedirs[torch.greater(g, 0.5).to(torch.long)], dims=([0], [2])) for i, g in enumerate(betas[:,-1])])
results = []
results.append(
self.with_zeros(torch.cat((R_cube_big[:, 0], torch.reshape(J[:, 0, :], (-1, 3, 1))), dim=2))
self.with_zeros(torch.cat((R_cube_big[:, 0], J[:, 0, :].view((-1, 3, 1))), dim=2))
)
for i in range(1, 24):
results.append(
@@ -190,7 +200,7 @@ class SMPL(Module):
results[self.parent[i]],
self.with_zeros(
torch.cat(
(R_cube_big[:, i], torch.reshape(J[:, i, :] - J[:, self.parent[i], :], (-1, 3, 1))),
(R_cube_big[:, i], (J[:, i, :] - J[:, self.parent[i], :]).view((-1, 3, 1))),
dim=2
)
)
@@ -202,25 +212,25 @@ class SMPL(Module):
self.pack(
torch.matmul(
stacked,
torch.reshape(
torch.cat((J, torch.zeros((batch_num, 24, 1), dtype=torch.float32).to(self.device)), dim=2),
torch.cat((J, self.results_zeros), dim=2).view(
(batch_num, 24, 4, 1)
)
)
)
joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
joints = torch.matmul(results, torch.transpose(torch.concat([J, self.J_ones], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
weights = self.weights if verts is not None else self.weights_all
rest_shape_h_ones_local = self.rest_shape_h_ones if verts is not None else self.rest_shape_h_ones_full
T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
rest_shape_h = torch.cat(
(v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
(v_posed, rest_shape_h_ones_local), dim=2
)
v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
v = torch.matmul(T, rest_shape_h.view((batch_num, -1, 4, 1)))
v = v.view((batch_num, -1, 4))[:, :, :3]
result = v
root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
root_joint = joints[:,0,:].view((-1, 1, 3))
result = result - root_joint
joints = joints - root_joint