From d5ef998e3e26e45363e1ad96fe1605a16412b434 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 31 Mar 2023 09:35:31 +0000 Subject: [PATCH] finalize --- smpl_tf.py | 6 +++--- smpl_torch.py | 52 +++++---------------------------------------------- 2 files changed, 8 insertions(+), 50 deletions(-) diff --git a/smpl_tf.py b/smpl_tf.py index 9503b82..350e28f 100644 --- a/smpl_tf.py +++ b/smpl_tf.py @@ -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 diff --git a/smpl_torch.py b/smpl_torch.py index 1299632..22219cc 100644 --- a/smpl_torch.py +++ b/smpl_torch.py @@ -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,11 +211,7 @@ 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 joints = joints - 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') \ No newline at end of file + return result, joints