def fit_SMPLD(scans, smpl_pkl, gender='male', save_path=None, scale_file=None):
# Get SMPL faces
sp = SmplPaths(gender=gender)
smpl_faces = sp.get_faces()
th_faces = torch.tensor(smpl_faces.astype('float32'), dtype=torch.long).cuda()
# Batch size
batch_sz = len(scans)
# Init SMPL
pose, betas, trans = [], [], []
for spkl in smpl_pkl:
smpl_dict = pkl.load(open(spkl, 'rb'))
p, b, t = smpl_dict['pose'], smpl_dict['betas'], smpl_dict['trans']
pose.append(p)
if len(b) == 10:
temp = np.zeros((300,))
temp[:10] = b
b = temp.astype('float32')
betas.append(b)
trans.append(t)
pose, betas, trans = np.array(pose), np.array(betas), np.array(trans)
betas, pose, trans = torch.tensor(betas), torch.tensor(pose), torch.tensor(trans)
smpl = th_batch_SMPL(batch_sz, betas, pose, trans, faces=th_faces).cuda()
verts, _, _, _ = smpl()
init_smpl_meshes = [tm.from_tensors(vertices=v.clone().detach(),
faces=smpl.faces) for v in verts]
# Load scans
th_scan_meshes = []
for scan in scans:
print('scan path ...', scan)
temp = Mesh(filename=scan)
th_scan = tm.from_tensors(torch.tensor(temp.v.astype('float32'), requires_grad=False, device=DEVICE),
torch.tensor(temp.f.astype('int32'), requires_grad=False, device=DEVICE).long())
th_scan_meshes.append(th_scan)
if scale_file is not None:
for n, sc in enumerate(scale_file):
dat = np.load(sc, allow_pickle=True)
th_scan_meshes[n].vertices += torch.tensor(dat[1]).to(DEVICE)
th_scan_meshes[n].vertices *= torch.tensor(dat[0]).to(DEVICE)
# Optimize
optimize_offsets(th_scan_meshes, smpl, init_smpl_meshes, 5, 10)
print('Done')
verts, _, _, _ = smpl()
th_smpl_meshes = [tm.from_tensors(vertices=v,
faces=smpl.faces) for v in verts]
if save_path is not None:
if not exists(save_path):
os.makedirs(save_path)
names = [split(s)[1] for s in scans]
# Save meshes
save_meshes(th_smpl_meshes, [join(save_path, n.replace('.ply', '_smpld.obj')) for n in names])
save_meshes(th_scan_meshes, [join(save_path, n) for n in names])
# Save params
for p, b, t, d, n in zip(smpl.pose.cpu().detach().numpy(), smpl.betas.cpu().detach().numpy(),
smpl.trans.cpu().detach().numpy(), smpl.offsets.cpu().detach().numpy(), names):
smpl_dict = {'pose': p, 'betas': b, 'trans': t, 'offsets': d}
pkl.dump(smpl_dict, open(join(save_path, n.replace('.ply', '_smpld.pkl')), 'wb'))
return smpl.pose.cpu().detach().numpy(), smpl.betas.cpu().detach().numpy(), \
smpl.trans.cpu().detach().numpy(), smpl.offsets.cpu().detach().numpy()
def fit_SMPLD(scans, smpl_pkl=None, gender='male', save_path=None, display=False):
# Get SMPL faces
sp = SmplPaths(gender=gender)
smpl_faces = sp.get_faces()
th_faces = torch.tensor(smpl_faces.astype('float32'), dtype=torch.long).cuda()
# Batch size
batch_sz = len(scans)
# Init SMPL
if smpl_pkl is None or smpl_pkl[0] is None:
print('SMPL not specified, fitting SMPL now')
pose, betas, trans = fit_SMPL(scans, None, gender, save_path, display)
else:
pose, betas, trans = [], [], []
for spkl in smpl_pkl:
smpl_dict = pkl.load(open(spkl, 'rb'), encoding='latin-1')
p, b, t = smpl_dict['pose'], smpl_dict['betas'], smpl_dict['trans']
pose.append(p)
if len(b) == 10:
temp = np.zeros((300,))
temp[:10] = b
b = temp.astype('float32')
betas.append(b)
trans.append(t)
pose, betas, trans = np.array(pose), np.array(betas), np.array(trans)
betas, pose, trans = torch.tensor(betas), torch.tensor(pose), torch.tensor(trans)
smpl = th_batch_SMPL(batch_sz, betas, pose, trans, faces=th_faces).cuda()
verts, _, _, _ = smpl()
init_smpl_meshes = [tm.from_tensors(vertices=v.clone().detach(),
faces=smpl.faces) for v in verts]
# Load scans
th_scan_meshes = []
for scan in scans:
th_scan = tm.from_obj(scan)
if save_path is not None:
th_scan.save_mesh(join(save_path, split(scan)[1]))
th_scan.vertices = th_scan.vertices.cuda()
th_scan.faces = th_scan.faces.cuda()
th_scan.vertices.requires_grad = False
th_scan_meshes.append(th_scan)
# Optimize
optimize_offsets(th_scan_meshes, smpl, init_smpl_meshes, 5, 10)
print('Done')
verts, _, _, _ = smpl()
th_smpl_meshes = [tm.from_tensors(vertices=v,
faces=smpl.faces) for v in verts]
if save_path is not None:
if not exists(save_path):
os.makedirs(save_path)
names = [split(s)[1] for s in scans]
# Save meshes
save_meshes(th_smpl_meshes, [join(save_path, n.replace('.obj', '_smpld.obj')) for n in names])
save_meshes(th_scan_meshes, [join(save_path, n) for n in names])
# Save params
for p, b, t, d, n in zip(smpl.pose.cpu().detach().numpy(), smpl.betas.cpu().detach().numpy(),
smpl.trans.cpu().detach().numpy(), smpl.offsets.cpu().detach().numpy(), names):
smpl_dict = {'pose': p, 'betas': b, 'trans': t, 'offsets': d}
pkl.dump(smpl_dict, open(join(save_path, n.replace('.obj', '_smpld.pkl')), 'wb'))
return smpl.pose.cpu().detach().numpy(), smpl.betas.cpu().detach().numpy(), \
smpl.trans.cpu().detach().numpy(), smpl.offsets.cpu().detach().numpy()
def fit_SMPL(scans, scan_labels, gender='male', save_path=None, scale_file=None, display=None):
"""
:param save_path:
:param scans: list of scan paths
:param pose_files:
:return:
"""
# Get SMPL faces
sp = SmplPaths(gender=gender)
smpl_faces = sp.get_faces()
th_faces = torch.tensor(smpl_faces.astype('float32'), dtype=torch.long).to(DEVICE)
# Load SMPL parts
part_labels = pkl.load(open('/BS/bharat-3/work/IPNet/assets/smpl_parts_dense.pkl', 'rb'))
labels = np.zeros((6890,), dtype='int32')
for n, k in enumerate(part_labels):
labels[part_labels[k]] = n
labels = torch.tensor(labels).unsqueeze(0).to(DEVICE)
# Load scan parts
scan_part_labels = []
for sc_l in scan_labels:
temp = torch.tensor(np.load(sc_l).astype('int32')).to(DEVICE)
scan_part_labels.append(temp)
# Batch size
batch_sz = len(scans)
# Set optimization hyper parameters
iterations, pose_iterations, steps_per_iter, pose_steps_per_iter = 3, 2, 30, 30
prior = get_prior(gender=gender, precomputed=True)
pose_init = torch.zeros((batch_sz, 72))
pose_init[:, 3:] = prior.mean
betas, pose, trans = torch.zeros((batch_sz, 300)), pose_init, torch.zeros((batch_sz, 3))
# Init SMPL, pose with mean smpl pose, as in ch.registration
smpl = th_batch_SMPL(batch_sz, betas, pose, trans, faces=th_faces).to(DEVICE)
smpl_part_labels = torch.cat([labels] * batch_sz, axis=0)
th_scan_meshes, centers = [], []
for scan in scans:
print('scan path ...', scan)
temp = Mesh(filename=scan)
th_scan = tm.from_tensors(torch.tensor(temp.v.astype('float32'), requires_grad=False, device=DEVICE),
torch.tensor(temp.f.astype('int32'), requires_grad=False, device=DEVICE).long())
th_scan_meshes.append(th_scan)
if scale_file is not None:
for n, sc in enumerate(scale_file):
dat = np.load(sc, allow_pickle=True)
th_scan_meshes[n].vertices += torch.tensor(dat[1]).to(DEVICE)
th_scan_meshes[n].vertices *= torch.tensor(dat[0]).to(DEVICE)
# Optimize pose first
optimize_pose_only(th_scan_meshes, smpl, pose_iterations, pose_steps_per_iter, scan_part_labels, smpl_part_labels,
display=None if display is None else 0)
# Optimize pose and shape
optimize_pose_shape(th_scan_meshes, smpl, iterations, steps_per_iter, scan_part_labels, smpl_part_labels,
display=None if display is None else 0)
verts, _, _, _ = smpl()
th_smpl_meshes = [tm.from_tensors(vertices=v, faces=smpl.faces) for v in verts]
if save_path is not None:
if not exists(save_path):
os.makedirs(save_path)
names = [split(s)[1] for s in scans]
# Save meshes
save_meshes(th_smpl_meshes, [join(save_path, n.replace('.ply', '_smpl.obj')) for n in names])
save_meshes(th_scan_meshes, [join(save_path, n) for n in names])
# Save params
for p, b, t, n in zip(smpl.pose.cpu().detach().numpy(), smpl.betas.cpu().detach().numpy(),
smpl.trans.cpu().detach().numpy(), names):
smpl_dict = {'pose': p, 'betas': b, 'trans': t}
pkl.dump(smpl_dict, open(join(save_path, n.replace('.ply', '_smpl.pkl')), 'wb'))
return smpl.pose.cpu().detach().numpy(), smpl.betas.cpu().detach().numpy(), smpl.trans.cpu().detach().numpy()