def get_tpose_smpl():
sp = SmplPaths(gender='neutral')
smpl = sp.get_smpl()
smpl.trans[:] = 0
smpl.pose[:] = 0
return smpl
def save_output(self, names, pose_, betas_, trans_, corr_, vcs, save_name,
epoch, it):
from psbody.mesh import Mesh
from lib.smpl_paths import SmplPaths
sp = SmplPaths(gender='male')
smpl = sp.get_smpl()
for nam, p, b, t, c, vc in zip(names, pose_, betas_, trans_, corr_,
vcs):
name = split(nam)[1]
smpl.pose[:] = p
smpl.betas[:10] = b
smpl.trans[:] = t
# save registration
Mesh(smpl.r, smpl.f).write_ply(
join(self.exp_path, save_name + '_ep_{}'.format(epoch),
name + '_{}_reg.ply'.format(it)))
# save raw correspondences
temp = Mesh(c, [])
temp.set_vertex_colors(vc)
temp.write_ply(
join(self.exp_path, save_name + '_ep_{}'.format(epoch),
name + '_{}_craw.ply'.format(it)))
# save SMPL params
with open(
join(self.exp_path, save_name + '_ep_{}'.format(epoch),
name + '_{}_reg.pkl'.format(it)), 'wb') as f:
pkl.dump({'pose': p, 'betas': b, 'trans': t}, f)
def __init__(self,
mode,
batch_sz,
data_path=DATA_PATH,
split_file='assets/data_split_01.pkl',
num_workers=12,
augment=False,
naked=False):
self.mode = mode
self.path = data_path
with open(split_file, "rb") as f:
self.split = pkl.load(f)
self.data = self.split[mode]
self.batch_size = batch_sz
self.num_workers = num_workers
self.augment = augment
self.naked = naked
sp = SmplPaths(gender='male')
self.ref_smpl = sp.get_smpl()
self.vt, self.ft = sp.get_vt_ft()
# Load smpl part labels
with open('assets/smpl_parts_dense.pkl', 'rb') as f:
dat = pkl.load(f, encoding='latin-1')
self.smpl_parts = np.zeros((6890, 1))
for n, k in enumerate(dat):
self.smpl_parts[dat[k]] = n
def get_prior(gender='male', precomputed=False):
if precomputed:
prior = Prior(sm=None)
return prior['Generic']
else:
from lib.smpl_paths import SmplPaths
dp = SmplPaths(gender=gender)
if gender == 'neutral':
dp_prior = SmplPaths(gender='male')
else:
dp_prior = dp
prior = Prior(dp_prior.get_smpl())
return prior['Generic']
def __init__(self,
model,
device,
train_loader,
val_loader,
exp_name,
opt_dict={},
optimizer='Adam',
checkpoint_number=-1,
train_supervised=False):
"""
:param model: correspondence prediction network
:param device: cuda or cpu
:param train_loader:
:param val_loader:
:param exp_name:
:param opt_dict: dict containing optimization specific parameteres
:param optimizer:
:param checkpoint_number: load a specific checkpoint, -1 => load latest
:param train_supervised:
"""
self.model = model.to(device)
self.device = device
self.opt_dict = self.parse_opt_dict(opt_dict)
self.optimizer_type = optimizer
self.optimizer = self.init_optimizer(optimizer,
self.model.parameters(),
learning_rate=0.001)
self.train_data_loader = train_loader
self.val_data_loader = val_loader
self.train_supervised = train_supervised
self.checkpoint_number = checkpoint_number
# Load vsmpl
self.vsmpl = VolumetricSMPL(
'/BS/bharat-2/work/LearntRegistration/test_data/volumetric_smpl_function_64',
device, 'male')
sp = SmplPaths(gender='male')
self.ref_smpl = sp.get_smpl()
self.template_points = torch.tensor(trimesh.Trimesh(
vertices=self.ref_smpl.r,
faces=self.ref_smpl.f).sample(NUM_POINTS).astype('float32'),
requires_grad=False).unsqueeze(0)
self.pose_prior = get_prior('male', precomputed=True)
# Load smpl part labels
with open(
'/BS/bharat-2/work/LearntRegistration/test_data/smpl_parts_dense.pkl',
'rb') as f:
dat = pkl.load(f, encoding='latin-1')
self.smpl_parts = np.zeros((6890, 1))
for n, k in enumerate(dat):
self.smpl_parts[dat[k]] = n
self.exp_path = join(os.path.dirname(__file__),
'../experiments/{}'.format(exp_name))
self.checkpoint_path = join(self.exp_path,
'checkpoints/'.format(exp_name))
if not os.path.exists(self.checkpoint_path):
print(self.checkpoint_path)
os.makedirs(self.checkpoint_path)
self.writer = SummaryWriter(
join(self.exp_path, 'summary'.format(exp_name)))
self.val_min = None