From 0a65dbd75d35ef7903876db1a45a39dc33567948 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Mon, 6 Feb 2023 15:53:34 +0000 Subject: [PATCH] smaller error but still not same --- smpl_tf.py | 21 +++++++++++---------- smpl_torch.py | 29 +++++++++++++++------------- test.py | 52 +++++++++++++++++++++++++++++++++++++++++++++++++++ utils.py | 40 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 23 deletions(-) diff --git a/smpl_tf.py b/smpl_tf.py index 3ec5af6..9091a39 100644 --- a/smpl_tf.py +++ b/smpl_tf.py @@ -84,12 +84,13 @@ class SMPL: self.to_keep = sio.loadmat(os.path.join(path, 'keep_remove.mat'))['keep'].reshape(-1) self.weights_all = self.weights - if use_one_hot: - self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0) - self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1]) - else: - self.weights = tf.gather(self.weights, self.to_keep, axis=1) + # if use_one_hot: + # self.weights = tf.gather(self.weights_one_hot, self.to_keep, axis=0) + # self.weights = tf.tile(self.weights[tf.newaxis,:,:], [2,1,1]) + # else: + # self.weights = tf.gather(self.weights, self.to_keep, axis=1) + self.shape = (1,27578,3) if orig_smpl: @@ -189,21 +190,21 @@ class SMPL: v = tf.matmul(T, tf.reshape(rest_shape_h, (-1, 4, 1))) v = tf.reshape(v, (-1, 4))[:, :3] - return v, tf.squeeze(J), v_posed + return v, tf.squeeze(J), v_posed, lrotmin gender = tf.cast(tf.greater_equal(betas[:, 10] , 0.5), tf.int32) - T_out, J_out, v_posed = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32)) + T_out, J_out, v_posed,lrotmin = tf.map_fn(single_run, (R, v_shaped, gender), (tf.float32, tf.float32, tf.float32, tf.float32)) J0 = tf.reshape(J_out[:,0,:], (-1, 1, 3)) J_out = tf.subtract(J_out, J0) T_out = tf.subtract(T_out, J0) - v_posed = tf.subtract(v_posed, J0) + # v_posed = tf.subtract(v_posed, J0) if remove: T_out = tf.gather(T_out, self.to_keep, axis=1) - v_posed = tf.gather(v_posed, self.to_keep, axis=1) + # v_posed = tf.gather(v_posed, self.to_keep, axis=1) - return T_out, J_out, v_posed + return T_out, J_out, lrotmin, v_posed def run_with_weights(self, R, betas, v_cloth, weights=None, filter=None): diff --git a/smpl_torch.py b/smpl_torch.py index e7fc7b9..1299632 100644 --- a/smpl_torch.py +++ b/smpl_torch.py @@ -10,10 +10,10 @@ from time import time from torch.nn import Module from utils import plot_mesh -class SMPLModel(Module): +class SMPL(Module): def __init__(self, device=None, path='', orig_smpl=False, is_skirt=False): - super(SMPLModel, self).__init__() + super(SMPL, self).__init__() self.J_regressor = [] self.v_template = [] self.shapedirs = [] @@ -80,11 +80,14 @@ class SMPLModel(Module): ------- Rotation matrix of shape [batch_size * angle_num, 3, 3]. """ - eps = r.clone().normal_(std=1e-8) - theta = torch.norm(r + eps, dim=(1, 2), keepdim=True) # dim cannot be tuple + # eps = r.clone().normal_(std=1e-8) + # theta = torch.norm(r + eps, dim=(1, 2), keepdim=True) # dim cannot be tuple + theta = torch.norm(r, dim=(1, 2), keepdim=True) + theta = torch.maximum(theta, torch.ones_like(theta) * np.finfo(np.float32).eps) + theta_dim = theta.shape[0] r_hat = r / theta - cos = torch.cos(theta) + # cos = torch.cos(theta) z_stick = torch.zeros(theta_dim, dtype=torch.float32).to(r.device) m = torch.stack( (z_stick, -r_hat[:, 0, 2], r_hat[:, 0, 1], r_hat[:, 0, 2], z_stick, @@ -93,8 +96,9 @@ class SMPLModel(Module): 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) - dot = torch.matmul(A, r_hat) - R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m + # dot = torch.matmul(A, r_hat) + # R = cos * i_cube + (1 - cos) * dot + torch.sin(theta) * m + R = i_cube + m * torch.sin(theta) + torch.matmul(m, m) * (1 - torch.cos(theta)) return R @staticmethod @@ -138,7 +142,7 @@ class SMPLModel(Module): for f in self.faces + 1: fp.write('f %d %d %d\n' % (f[0], f[1], f[2])) - def forward(self, betas, pose, trans, remove=True): + def forward(self, betas, pose, remove=True): """ Construct a compute graph that takes in parameters and outputs a tensor as @@ -211,7 +215,7 @@ class SMPLModel(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 + torch.reshape(trans, (batch_num, 1, 3)) + result = v # estimate 3D joint locations # print(result.shape) # print(self.joint_regressor.shape) @@ -224,7 +228,7 @@ class SMPLModel(Module): if remove: result = result[:,self.to_keep,:] - return result, joints + return result, joints, lrotmin, v_posed def test_gpu(gpu_id=[0], path=''): @@ -239,7 +243,7 @@ def test_gpu(gpu_id=[0], path=''): beta_size = 10 np.random.seed(9608) - model = SMPLModel(device=device, path=path) + 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) @@ -247,8 +251,7 @@ def test_gpu(gpu_id=[0], path=''): .type(torch.float32).to(device) betas = torch.concat((betas, torch.zeros((4,1), device=device)), axis=1) s = time() - trans = torch.from_numpy(np.zeros((4, 3))).type(torch.float32).to(device) - result, joints = model(betas, pose, trans) + result, joints = model(betas, pose) print(time() - s) plot_mesh(result[0].cpu().numpy(), model.f).show() diff --git a/test.py b/test.py index e69de29..9657e2f 100644 --- a/test.py +++ b/test.py @@ -0,0 +1,52 @@ +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 \ No newline at end of file diff --git a/utils.py b/utils.py index 08346c7..f5f0564 100644 --- a/utils.py +++ b/utils.py @@ -1,4 +1,5 @@ import plotly.graph_objects as go +from plotly.subplots import make_subplots def plot_mesh(verts, faces): @@ -25,4 +26,43 @@ def plot_mesh(verts, faces): fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0)) + return fig + + +def plot_mesh_2(verts_1, verts_2, faces): + # fig = make_subplots(rows=1, cols=2, start_cell="top-left", specs = [ [{ 'type': 'scene' } for _ in range(2)] ]) + fig = go.Figure() + + fig.add_trace(go.Mesh3d( + x=verts_1[:,0], + y=verts_1[:,2], + z=verts_1[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text="", + ),) + + fig.add_trace(go.Mesh3d( + x=verts_2[:,0], + y=verts_2[:,2], + z=verts_2[:,1], + i = faces[:,0], + j = faces[:,1], + k = faces[:,2], + text="", + ),) + + fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1), + # scene2_aspectmode='manual', scene2_aspectratio=dict(x=1, y=1, z=1) + ) + helper = dict(xaxis = dict(range=[-1,1], ), + yaxis = dict(range=[-1,1], ), + zaxis = dict(range=[-1,1], ) ) + fig.update_layout(scene = helper, + # scene2 = helper + ) + + fig.update_layout(margin=dict(l=0, r=0, b=0, t=0, pad=0)) + return fig \ No newline at end of file