class Pose2MeshModel(nn.Module):
def __init__(self):
super(Pose2MeshModel, self).__init__()
if torch.cuda.is_available():
self.reg = AcosRegressor(hidden_dim=256).cuda()
self.smpl = SMPLModel(device=torch.device('cuda'),
model_path = './model_24_joints.pkl',
simplify=True
)
else:
self.reg = AcosRegressorRegressor(hidden_dim=256).cpu()
self.smpl = SMPLModel(device=torch.device('cpu'),
model_path = './model_24_joints.pkl',
simplify=True
)
ckpt_path = './checkpoints_0303_24_joints'
state_dict = torch.load('%s/regressor_040.pth' % (ckpt_path))
self.reg.load_state_dict(state_dict)
def forward(self, input):
trans = torch.zeros((input.shape[0], 3), device=input.device)
betas = torch.zeros((input.shape[0], 10), device=input.device)
thetas = self.reg(input)
print('Estimated theta:\n', thetas.detach().cpu().numpy())
mesh, joints = self.smpl(betas, thetas, trans)
return mesh, joints
def evaluate(self, input, save_dir):
mesh, joints = self.forward(input)
self.smpl.write_obj(mesh[0].detach().cpu().numpy(), save_dir)
np.savetxt('recon_pose.xyz', joints[0].detach().cpu().numpy().reshape(24,3), delimiter=' ')
def run_test():
if platform == 'linux':
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
data_type = torch.float32
device = torch.device('cuda')
pose_size = 72
beta_size = 10
np.random.seed(9608)
model = SMPLModel(
device=device,
model_path='./model_lsp.pkl',
data_type=data_type,
)
dataset = Human36MDataset(model, max_item=100, calc_mesh=True)
# generate mesh, align with 14 point ground truth
case_num = 10
data = dataset[:case_num]
meshes = data['meshes']
input = data['lsp_joints']
target_2d = data['gt2d']
target_3d = data['gt3d']
transforms = map_3d_to_2d(input, target_2d, target_3d)
# Important: mesh should be centered at the origin!
deformed_meshes = transforms(meshes)
mesh_3d = deformed_meshes.detach().cpu().numpy()
# visualize(data['imagename'], mesh_3d[:,:,:2].astype(np.int),
# target_2d.detach().cpu().numpy().astype(np.int))
for i, mesh in enumerate(mesh_3d):
model.write_obj(mesh,
'_test_cache/real_mesh_{}.obj'.format(i)) # weird.
UV_position_map, UV_scatter, rgbs_backup = get_UV(mesh, 300)
# write colorized coordinates to ply
write_ply('_test_cache/colored_mesh_{}.ply'.format(i), mesh,
rgbs_backup)
out = np.concatenate((UV_position_map, UV_scatter), axis=1)
imsave('_test_cache/UV_position_map_{}.png'.format(i), out)
resampled_mesh = resample(UV_position_map)
model.write_obj(resampled_mesh,
'_test_cache/recon_mesh_{}.obj'.format(i))
def create_meshes(UV_label_root=None, uv_prefix='smpl_fbx_template'):
if platform == 'linux':
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
data_type = torch.float32
device = torch.device('cuda')
pose_size = 72
beta_size = 10
np.random.seed(9608)
model = SMPLModel(
device=device,
model_path='./model_lsp.pkl',
data_type=data_type,
)
dataset = Human36MWashedDataset(
model,
calc_mesh=True,
root_dir='/home/wzeng/mydata/h3.6m/images_washed')
generator = UV_Map_Generator(UV_height=256,
UV_pickle=uv_prefix + '.pickle')
# create root folder for UV labels
if UV_label_root is None:
UV_label_root = dataset.root_dir.replace('images_washed',
'meshes_washed')
if not os.path.isdir(UV_label_root):
os.makedirs(UV_label_root)
subs = [
sub for sub in os.listdir(dataset.root_dir)
if os.path.isdir(dataset.root_dir + '/' + sub)
]
for sub in subs:
os.makedirs(UV_label_root + '/' + sub)
else:
print('{} folder exists, process terminated...'.format(UV_label_root))
# return
# generate mesh, align with 14 point ground truth
batch_size = 64
total_batch_num = dataset.length // batch_size + 1
_loop = tqdm(range(total_batch_num), ncols=80)
for batch_id in _loop:
data = dataset[batch_id * batch_size:(batch_id + 1) * batch_size]
meshes = data['meshes']
input = data['lsp_joints']
target_2d = data['gt2d']
target_3d = data['gt3d']
imagename = [UV_label_root + str for str in data['imagename']]
transforms = map_3d_to_2d(input, target_2d, target_3d)
# Important: mesh should be centered at the origin!
deformed_meshes = transforms(meshes)
mesh_3d = deformed_meshes.detach().cpu().numpy()
test_folder = UV_label_root
'''
test_folder = UV_label_root
visualize(test_folder, data['imagename'], mesh_3d[:,:,:2].astype(np.int),
target_2d.detach().cpu().numpy().astype(np.int), dataset.root_dir)
'''
s = time()
for name, mesh in zip(imagename, mesh_3d):
# this approach didn't normalize the coordinates
model.write_obj(mesh, name[:-4] + '.obj')
'''
model = SMPLModel(device=device, model_path = './model_24_joints.pkl',
simplify=True)
if not os.path.isdir('joint2pose_result'):
os.makedirs('joint2pose_result')
loss_op = nn.L1Loss()
betas = torch.zeros((1, beta_size), dtype=torch.float64, device=device)
trans = torch.zeros((1, 3), dtype=torch.float64, device=device)
for i in range(10):
print('Test case %d:' % (i+1))
real_pose = torch.from_numpy((np.random.rand(1, pose_size) - 0.5) * 1)\
.type(torch.float64).to(device)
real_result, real_joints = model(betas, real_pose, trans)
model.write_obj(real_result[0].detach().cpu().numpy(), 'joint2pose_result/real_mesh_{}.obj'.format(i))
# Initialize a pose from zero and optimize it
test_pose = torch.zeros((1, pose_size), dtype=torch.float64, device=device, requires_grad=True)
optimizer = SGD(iter([test_pose]), lr=0.0005, momentum=0.2)
#optimizer = Adam(iter([test_pose]), lr=0.0001, betas=(0.5,0.999))
s = time()
prev_loss = None
for step in range(2000):
_, test_joints = model(betas, test_pose, trans)
loss = loss_op(test_joints, real_joints)
if step % 50 == 0:
print('Step {:03d}: loss: {:10.6f}'.format(step, loss.data.item()))
cur_loss = loss.data.item()
def run_test():
if platform == 'linux':
os.environ['CUDA_VISIBLE_DEVICES'] = '2'
data_type = torch.float32
device = torch.device('cuda')
pose_size = 72
beta_size = 10
np.random.seed(9608)
model = SMPLModel(
device=device,
model_path='./model_lsp.pkl',
data_type=data_type,
)
dataset = Human36MDataset(model, max_item=100, calc_mesh=True)
# generate mesh, align with 14 point ground truth
case_num = 10
data = dataset[:case_num]
meshes = data['meshes']
input = data['lsp_joints']
target_2d = data['gt2d']
target_3d = data['gt3d']
transforms = map_3d_to_2d(input, target_2d, target_3d)
# Important: mesh should be centered at the origin!
deformed_meshes = transforms(meshes)
mesh_3d = deformed_meshes.detach().cpu().numpy()
file_prefix = 'smpl_fbx_template'
generator = UV_Map_Generator(UV_height=256,
UV_pickle=file_prefix + '.pickle')
test_folder = '_test_smpl_fbx'
if not os.path.isdir(test_folder):
os.makedirs(test_folder)
visualize(test_folder, data['imagename'], mesh_3d[:, :, :2].astype(np.int),
target_2d.detach().cpu().numpy().astype(np.int))
s = time()
for i, mesh in enumerate(mesh_3d):
model.write_obj(mesh, '{}/real_mesh_{}.obj'.format(test_folder,
i)) # weird.
UV_position_map, verts_backup = \
generator.get_UV_map(mesh)
# write colorized coordinates to ply
UV_scatter, _, _ = generator.render_point_cloud(verts=mesh)
generator.write_ply('{}/colored_mesh_{}.ply'.format(test_folder, i),
mesh)
out = np.concatenate((UV_position_map, UV_scatter), axis=1)
imsave('{}/UV_position_map_{}.png'.format(test_folder, i), out)
resampled_mesh = generator.resample(UV_position_map)
model.write_obj(resampled_mesh,
'{}/recon_mesh_{}.obj'.format(test_folder, i))
print('{} cases for {}s'.format(case_num, time() - s))