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

finalize

Browse Source
This commit is contained in:
luni
2023-03-31 09:35:31 +00:00
parent a27f5caf0c
commit d5ef998e3e
2 changed files with 8 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
smpl_tf.py
+3 -3
View File
@@ -184,8 +184,8 @@ class SMPL:
stacked = tf.stack(results, axis=0)
results = stacked - self.pack(tf.matmul(stacked, tf.reshape(tf.concat((J, tf.zeros((24, 1), dtype=tf.float32)), axis=1), (24, 4, 1))))
# J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
# J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
J = tf.matmul(results, tf.expand_dims(tf.concat([J, tf.ones([24, 1], tf.float32)], axis=1), axis=1), transpose_b=True)
J = tf.gather(J, tf.constant([0, 1, 2], tf.int32), axis=1)
T = tf.tensordot(weights[g], results, axes=((1), (0)))
@@ -193,7 +193,7 @@ class SMPL:
v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1)))
v = tf.reshape(v, (-1, 4))[:, :3]
J = tf.transpose(tf.tensordot(v, self.J_regressor[g], axes=([0], [1])), [1,0])
# J = tf.transpose(tf.tensordot(v, self.J_regressor[g], axes=([0], [1])), [1,0])
return v, tf.squeeze(J), v_posed, lrotmin
smpl_torch.py
+4 -46
View File
@@ -134,13 +134,6 @@ class SMPL(Module):
ret = torch.cat((zeros43, x), dim=3)
return ret
def write_obj(self, verts, file_name):
with open(file_name, 'w') as fp:
for v in verts:
fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
for f in self.faces + 1:
fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))
def forward(self, betas, pose, remove=True):
@@ -208,7 +201,9 @@ class SMPL(Module):
)
)
)
# Restart from here
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]
T = torch.stack([torch.tensordot(results[i], self.weights[g.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
@@ -216,10 +211,6 @@ class SMPL(Module):
v = torch.matmul(T, torch.reshape(rest_shape_h, (batch_num, -1, 4, 1)))
v = torch.reshape(v, (batch_num, -1, 4))[:, :, :3]
result = v
# estimate 3D joint locations
# print(result.shape)
# print(self.joint_regressor.shape)
joints = torch.stack([torch.tensordot(result[i], self.J_regressor[g.to(torch.long)], dims=([0], [1])) for i, g in enumerate(betas[:,-1])]).transpose(1, 2)
root_joint = torch.reshape(joints[:,0,:], (-1, 1, 3))
result = result - root_joint
@@ -228,37 +219,4 @@ class SMPL(Module):
if remove:
result = result[:,self.to_keep,:]
return result, joints, lrotmin, v_posed
def test_gpu(gpu_id=[0], path=''):
if len(gpu_id) > 0 and torch.cuda.is_available():
os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id[0])
device = torch.device('cuda')
else:
device = torch.device('cpu')
#print(device)
pose_size = 72
beta_size = 10
np.random.seed(9608)
model = SMPL(device=device, path=path)
for i in range(5):
pose = torch.from_numpy((np.random.rand(4, pose_size) - 0.5) * 0.4)\
.type(torch.float32).to(device)
betas = torch.from_numpy((np.random.rand(4, beta_size) - 0.5) * 0.06) \
.type(torch.float32).to(device)
betas = torch.concat((betas, torch.zeros((4,1), device=device)), axis=1)
s = time()
result, joints = model(betas, pose)
print(time() - s)
plot_mesh(result[0].cpu().numpy(), model.f).show()
# outmesh_path = './smpl_torch_{}.obj'
# for i in range(result.shape[0]):
# model.write_obj(result[i], outmesh_path.format(i))
if __name__ == '__main__':
test_gpu([5], path='model_files')
return result, joints