def validation_step(self, batch, batch_idx):
drec = self(batch['pose_body'].view(-1, 63))
loss = self._compute_loss(batch, drec)
val_loss = loss['unweighted_loss']['loss_total']
if self.renderer is not None and self.global_rank == 0 and batch_idx % 500 == 0 and np.random.rand(
) > 0.5:
out_fname = makepath(
self.work_dir,
'renders/vald_rec_E{:03d}_It{:04d}_val_loss_{:.2f}.png'.format(
self.current_epoch, batch_idx, val_loss.item()),
isfile=True)
self.renderer([batch, drec], out_fname=out_fname)
dgen = self.vp_model.sample_poses(
self.vp_ps.logging.num_bodies_to_display)
out_fname = makepath(self.work_dir,
'renders/vald_gen_E{:03d}_I{:04d}.png'.format(
self.current_epoch, batch_idx),
isfile=True)
self.renderer([dgen], out_fname=out_fname)
progress_bar = {'v2v': val_loss}
return {
'val_loss': c2c(val_loss),
'progress_bar': progress_bar,
'log': progress_bar
}
def smpl_params2ply(bm, out_dir, pose_body, pose_hand = None, trans=None, betas=None, root_orient=None):
'''
:param bm: pytorch body model with batch_size 1
:param pose_body: can be a single list of pose parameters, or a list of list of pose parameters:
:param trans: Nx3
:param betas: Nxnum_betas
:return
dumps are all parameter as gltf objects
'''
faces = c2c(bm.f)
makepath(out_dir)
for fIdx in range(0, len(pose_body)):
bm.pose_body.data[0,:] = bm.pose_body.new(pose_body[fIdx].reshape(1,-1))
if pose_hand is not None: bm.pose_hand.data[0,:] = bm.pose_hand.new(pose_hand[fIdx].reshape(1,-1))
if trans is not None: bm.trans.data[0,:] = bm.trans.new(trans[fIdx].reshape(1,-1))
if betas is not None: bm.betas.data[0,:len(betas[fIdx])] = bm.betas.new(betas[fIdx])
if root_orient is not None: bm.root_orient.data[0,:] = bm.root_orient.new(root_orient[fIdx])
v = c2c(bm.forward().v)[0]
mesh = trimesh.base.Trimesh(v, faces)
mesh.export(os.path.join(out_dir, '%03d.ply' % fIdx))
def dump_amass2pytroch(datasets, amass_dir, out_posepath, logger = None, rnd_seed = 100, keep_rate = 0.01):
'''
Select random number of frames from central 80 percent of each mocap sequence
Save individual data features like pose and shape per frame in pytorch pt files
test set will have the extra field for original markers
:param datasets: the name of the dataset
:param amass_dir: directory of downloaded amass npz files. should be in this structure: path/datasets/subjects/*_poses.npz
:param out_posepath: the path for final pose.pt file
:param logger: an instance of human_body_prior.tools.omni_tools.log2file
:param rnd_seed:
:return: Number of datapoints dumped using out_poseth address pattern
'''
import glob
np.random.seed(rnd_seed)
makepath(out_posepath, isfile=True)
if logger is None:
starttime = datetime.now().replace(microsecond=0)
log_name = datetime.strftime(starttime, '%Y%m%d_%H%M')
logger = log2file(out_posepath.replace('pose.pt', '%s.log' % (log_name)))
logger('Creating pytorch dataset at %s' % out_posepath)
data_pose = []
data_dmpl = []
data_betas = []
data_gender = []
data_trans = []
for ds_name in datasets:
npz_fnames = glob.glob(os.path.join(amass_dir, ds_name, '*/*_poses.npz'))
logger('randomly selecting data points from %s.' % (ds_name))
for npz_fname in tqdm(npz_fnames):
try:
cdata = np.load(npz_fname)
except:
logger('Could not read %s! skipping..'%npz_fname)
continue
N = len(cdata['poses'])
cdata_ids = np.random.choice(list(range(int(0.1*N), int(0.9*N),1)), int(keep_rate*0.8*N), replace=False)#removing first and last 10% of the data to avoid repetitive initial poses
if len(cdata_ids)<1: continue
data_pose.extend(cdata['poses'][cdata_ids].astype(np.float32))
data_dmpl.extend(cdata['dmpls'][cdata_ids].astype(np.float32))
data_trans.extend(cdata['trans'][cdata_ids].astype(np.float32))
data_betas.extend(np.repeat(cdata['betas'][np.newaxis].astype(np.float32), repeats=len(cdata_ids), axis=0))
data_gender.extend([gdr2num[str(cdata['gender'].astype(np.str))] for _ in cdata_ids])
assert len(data_pose) != 0
torch.save(torch.tensor(np.asarray(data_pose, np.float32)), out_posepath)
torch.save(torch.tensor(np.asarray(data_dmpl, np.float32)), out_posepath.replace('pose.pt', 'dmpl.pt'))
torch.save(torch.tensor(np.asarray(data_betas, np.float32)), out_posepath.replace('pose.pt', 'betas.pt'))
torch.save(torch.tensor(np.asarray(data_trans, np.float32)), out_posepath.replace('pose.pt', 'trans.pt'))
torch.save(torch.tensor(np.asarray(data_gender, np.int32)), out_posepath.replace('pose.pt', 'gender.pt'))
return len(data_pose)
def registration2markers(registration_dir, out_marker_dir):
np.random.seed(100)
m2b_distance = 0.0095
genders = {
'50002': 'male',
'50004': 'female',
'50007': 'male',
'50009': 'male',
'50020': 'female',
'50021': 'female',
'50022': 'female',
'50025': 'female',
'50026': 'male',
'50027': 'male'
}
with open('./ssm_all_marker_placements.json') as f:
all_marker_placements = json.load(f)
all_mrks_keys = list(all_marker_placements.keys())
for dfaust_subject in genders.keys():
subject_reg_pkls = glob.glob(
os.path.join(registration_dir, dfaust_subject, '*.pkl'))
chosen_k = all_mrks_keys[np.random.choice(len(all_marker_placements))]
chosen_marker_set = all_marker_placements[chosen_k]
print('chose %s markerset for dfaust subject %s' %
(chosen_k, dfaust_subject))
for reg_pkl in subject_reg_pkls:
with open(reg_pkl, 'rb') as f:
data = pickle.load(f, encoding='latin-1')
marker_data = np.zeros([len(data['v']), len(chosen_marker_set), 3])
cur_m2b_distance = m2b_distance + abs(
np.random.normal(0, m2b_distance / 3., size=[3
])) # Noise in 3D
for fIdx in range(0, len(data['v'])):
vertices = rotate_mesh(data['v'][fIdx].copy(), 90)
vn = c2c(
compute_vertex_normal(torch.Tensor(vertices),
torch.Tensor(data['f'])))
for mrk_id, vid in enumerate(chosen_marker_set.values()):
marker_data[
fIdx,
mrk_id] = vertices[vid] + cur_m2b_distance * vn[vid]
outpath = makepath(os.path.join(out_marker_dir, dfaust_subject,
os.path.basename(reg_pkl)),
isfile=True)
np.savez(
outpath, **{
'markers': marker_data,
'labels': list(chosen_marker_set.keys()),
'frame_rate': 60,
'gender': genders[dfaust_subject]
})
def render_once(body_parms,
body_colors=[colors['grey'], colors['brown-light']],
out_fname=None):
'''
:param body_parms: list of dictionaries of body parameters.
:param body_colors: list of np arrays of color rgb values
:param movie_outpath: a mp4 path
:return:
'''
if out_fname is not None: makepath(out_fname, isfile=True)
assert len(body_parms) <= len(body_colors), ValueError(
'Not enough colors provided for #{} body_parms'.format(
len(body_parms)))
bs = body_parms[0]['pose_body'].shape[0]
body_ids = np.random.choice(bs, num_bodies_to_display)
body_evals = [
c2c(
bm(root_orient=v['root_orient'].view(bs, -1) if 'root_orient'
in v else torch.zeros(bs, 3).type_as(v['pose_body']),
pose_body=v['pose_body'].contiguous().view(bs, -1)).v)
for v in body_parms
]
num_verts = body_evals[0].shape[1]
render_meshes = []
for bId in body_ids:
concat_cur_meshes = None
for body, body_color in zip(body_evals, body_colors):
cur_body_mesh = Mesh(body[bId],
faces,
vertex_colors=np.ones([num_verts, 3]) *
body_color)
concat_cur_meshes = cur_body_mesh if concat_cur_meshes is None else mesh_cat(
concat_cur_meshes, cur_body_mesh)
render_meshes.append(concat_cur_meshes)
img = renderer(render_meshes)
if out_fname is not None: imagearray2file(img, out_fname, fps=10)
return
def perform_training(self, num_epochs=None, message=None):
starttime = datetime.now().replace(microsecond=0)
if num_epochs is None: num_epochs = self.ps.num_epochs
self.logger(
'Started Training at %s for %d epochs' % (datetime.strftime(starttime, '%Y-%m-%d_%H:%M:%S'), num_epochs))
vis_bm = BodyModel(self.ps.bm_path, 'smplh', num_betas=16).to(self.comp_device)
prev_lr = np.inf
scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=int(num_epochs // 3), gamma=0.5)
for epoch_num in range(1, num_epochs + 1):
scheduler.step()
cur_lr = self.optimizer.param_groups[0]['lr']
if cur_lr != prev_lr:
self.logger('--- Optimizer learning rate changed from %.2e to %.2e ---' % (prev_lr, cur_lr))
prev_lr = cur_lr
self.epochs_completed += 1
train_loss_dict = self.train()
eval_loss_dict = self.evaluate()
with torch.no_grad():
eval_msg = VPoserTrainer.creat_loss_message(eval_loss_dict, expr_code=self.ps.expr_code,
epoch_num=self.epochs_completed, it=len(self.ds_val),
try_num=self.try_num, mode='evald')
if eval_loss_dict['loss_total'] < self.best_loss_total:
self.ps.best_model_fname = makepath(os.path.join(self.ps.work_dir, 'snapshots', 'TR%02d_E%03d.pt' % (
self.try_num, self.epochs_completed)), isfile=True)
self.logger(eval_msg + ' ** ')
self.best_loss_total = eval_loss_dict['loss_total']
torch.save(self.vposer_model.module.state_dict() if isinstance(self.vposer_model, torch.nn.DataParallel) else self.vposer_model.state_dict(), self.ps.best_model_fname)
imgname = '[%s]_TR%02d_E%03d.png' % (self.ps.expr_code, self.try_num, self.epochs_completed)
imgpath = os.path.join(self.ps.work_dir, 'images', imgname)
try:
VPoserTrainer.vis_results(self.vis_dorig, self.vposer_model, bm=vis_bm, imgpath=imgpath)
except:
print('The visualization failed.')
else:
self.logger(eval_msg)
self.swriter.add_scalars('total_loss/scalars', {'train_loss_total': train_loss_dict['loss_total'],
'evald_loss_total': eval_loss_dict['loss_total'], },
self.epochs_completed)
# if early_stopping(eval_loss_dict['loss_total']):
# self.logger("Early stopping at epoch %d"%self.epochs_completed)
# break
endtime = datetime.now().replace(microsecond=0)
self.logger('Finished Training at %s\n' % (datetime.strftime(endtime, '%Y-%m-%d_%H:%M:%S')))
self.logger(
'Training done in %s! Best val total loss achieved: %.2e\n' % (endtime - starttime, self.best_loss_total))
self.logger('Best model path: %s\n' % self.ps.best_model_fname)
def train_vposer_once(_config):
resume_training_if_possible = True
model = VPoserTrainer(_config)
model.vp_ps.logging.expr_msg = create_expr_message(model.vp_ps)
# model.text_logger(model.vp_ps.logging.expr_msg.replace(". ", '.\n'))
dump_config(model.vp_ps,
osp.join(model.work_dir, '{}.yaml'.format(model.expr_id)))
logger = TensorBoardLogger(model.work_dir, name='tensorboard')
lr_monitor = LearningRateMonitor()
snapshots_dir = osp.join(model.work_dir, 'snapshots')
checkpoint_callback = ModelCheckpoint(
dirpath=makepath(snapshots_dir, isfile=True),
filename="%s_{epoch:02d}_{val_loss:.2f}" % model.expr_id,
save_top_k=1,
verbose=True,
monitor='val_loss',
mode='min',
)
early_stop_callback = EarlyStopping(
**model.vp_ps.train_parms.early_stopping)
resume_from_checkpoint = None
if resume_training_if_possible:
available_ckpts = sorted(glob.glob(osp.join(snapshots_dir, '*.ckpt')),
key=os.path.getmtime)
if len(available_ckpts) > 0:
resume_from_checkpoint = available_ckpts[-1]
model.text_logger(
'Resuming the training from {}'.format(resume_from_checkpoint))
trainer = pl.Trainer(
gpus=1,
weights_summary='top',
distributed_backend='ddp',
# replace_sampler_ddp=False,
# accumulate_grad_batches=4,
# profiler=False,
# overfit_batches=0.05,
# fast_dev_run = True,
# limit_train_batches=0.02,
# limit_val_batches=0.02,
# num_sanity_val_steps=2,
plugins=[DDPPlugin(find_unused_parameters=False)],
callbacks=[lr_monitor, early_stop_callback, checkpoint_callback],
max_epochs=model.vp_ps.train_parms.num_epochs,
logger=logger,
resume_from_checkpoint=resume_from_checkpoint)
trainer.fit(model)
def prepare_data(self):
'''' Similar to standard AMASS dataset preparation pipeline:
Donwload npz file, corresponding to body data from https://amass.is.tue.mpg.de/ and place them under amass_dir
'''
self.text_logger = log2file(makepath(self.work_dir,
'{}.log'.format(self.expr_id),
isfile=True),
prefix=self._log_prefix)
prepare_vposer_datasets(self.dataset_dir,
self.vp_ps.data_parms.amass_splits,
self.vp_ps.data_parms.amass_dir,
logger=self.text_logger)
def dump_vposer_samples(bm, pose_body, out_imgpath=False, save_ply=False):
'''
:param bm: the BodyModel instance
:param pose_body: Nx63 will pose the body
:param out_imgpath: the final png path
:param save_ply: if True will dump as ply files
:return:
'''
view_angles = [0, 90, -90]
imw, imh = 400, 400
mv = MeshViewer(width=imw, height=imh, use_offscreen=True)
images = np.zeros([len(view_angles), len(pose_body), 1, imw, imh, 3])
for cId in range(0, len(pose_body)):
bm.pose_body.data[:] = bm.pose_body.new(pose_body[cId].reshape(-1))
body_mesh = trimesh.Trimesh(vertices=c2c(bm().v[0]),
faces=c2c(bm.f),
vertex_colors=np.tile(
colors['grey'], (6890, 1)))
for rId, angle in enumerate(view_angles):
apply_mesh_tranfsormations_(
[body_mesh],
trimesh.transformations.rotation_matrix(
np.radians(angle), (0, 1, 0)))
mv.set_meshes([body_mesh], group_name='static')
images[rId, cId, 0] = mv.render()
apply_mesh_tranfsormations_(
[body_mesh],
trimesh.transformations.rotation_matrix(
np.radians(-angle), (0, 1, 0)))
if out_imgpath:
imagearray2file(images, out_imgpath)
np.savez(out_imgpath.replace('.png', '.npz'), pose=pose_body)
if save_ply:
im_id = os.path.basename(out_imgpath).split('.')[0]
out_dir = makepath(
os.path.join(os.path.dirname(out_imgpath), '%s_ply' % im_id))
smpl_params2ply(bm, out_dir=out_dir, pose_body=pose_body)
print('Saved image: %s' % out_imgpath)
return images
def extract_weights_asnumpy(exp_id, vp_model=False):
from human_body_prior.tools.omni_tools import makepath
from human_body_prior.tools.omni_tools import copy2cpu as c2c
vposer_pt, vposer_ps = load_vposer(exp_id, vp_model=vp_model)
save_wt_dir = makepath(os.path.join(vposer_ps.work_dir, 'weights_npy'))
weights = {}
for var_name, var in vposer_pt.named_parameters():
weights[var_name] = c2c(var)
np.savez(os.path.join(save_wt_dir, 'vposerWeights.npz'), **weights)
print(('Dumped weights as numpy arrays to %s' % save_wt_dir))
return vposer_ps, weights
def sample_vposer(expr_dir, bm, num_samples=5, vp_model='snapshot'):
from human_body_prior.tools.omni_tools import id_generator, makepath
from human_body_prior.tools.model_loader import load_vposer
from human_body_prior.tools.omni_tools import copy2cpu
vposer_pt, ps = load_vposer(expr_dir, vp_model=vp_model)
sampled_pose_body = copy2cpu(vposer_pt.sample_poses(num_poses=num_samples))
out_dir = makepath(
os.path.join(ps.work_dir, 'evaluations', 'pose_generation'))
out_imgpath = os.path.join(out_dir, '%s.png' % id_generator(6))
dump_vposer_samples(bm, sampled_pose_body, out_imgpath)
print('Dumped samples at %s' % out_dir)
return sampled_pose_body
def on_train_start(self):
if self.global_rank != 0: return
self.train_starttime = dt.now().replace(microsecond=0)
######## make a backup of vposer
git_repo_dir = os.path.abspath(__file__).split('/')
git_repo_dir = '/'.join(
git_repo_dir[:git_repo_dir.index('human_body_prior') + 1])
starttime = dt.strftime(self.train_starttime, '%Y_%m_%d_%H_%M_%S')
archive_path = makepath(self.work_dir,
'code',
'vposer_{}.tar.gz'.format(starttime),
isfile=True)
cmd = 'cd %s && git ls-files -z | xargs -0 tar -czf %s' % (
git_repo_dir, archive_path)
os.system(cmd)
########
self.text_logger(
'Created a git archive backup at {}'.format(archive_path))
dump_config(self.vp_ps,
osp.join(self.work_dir, '{}.yaml'.format(self.expr_id)))
def __init__(self, _config):
super(VPoserTrainer, self).__init__()
_support_data_dir = get_support_data_dir()
vp_ps = load_config(**_config)
make_deterministic(vp_ps.general.rnd_seed)
self.expr_id = vp_ps.general.expr_id
self.dataset_id = vp_ps.general.dataset_id
self.work_dir = vp_ps.logging.work_dir = makepath(
vp_ps.general.work_basedir, self.expr_id)
self.dataset_dir = vp_ps.logging.dataset_dir = osp.join(
vp_ps.general.dataset_basedir, vp_ps.general.dataset_id)
self._log_prefix = '[{}]'.format(self.expr_id)
self.text_logger = log2file(prefix=self._log_prefix)
self.seq_len = vp_ps.data_parms.num_timeseq_frames
self.vp_model = VPoser(vp_ps)
with torch.no_grad():
self.bm_train = BodyModel(vp_ps.body_model.bm_fname)
if vp_ps.logging.render_during_training:
self.renderer = vposer_trainer_renderer(
self.bm_train, vp_ps.logging.num_bodies_to_display)
else:
self.renderer = None
self.example_input_array = {
'pose_body': torch.ones(vp_ps.train_parms.batch_size, 63),
}
self.vp_ps = vp_ps
import os
from human_body_prior.tools.omni_tools import makepath, log2file
from human_body_prior.data.prepare_data import prepare_vposer_datasets
from datetime import datetime
expr_code = datetime.now().strftime("%d %m %Y %H:%M:%S")
amass_dir = r'/content/drive/My Drive/LAZAR/AMASS'
vposer_datadir = makepath('prepared/%s' % (expr_code))
logger = log2file(os.path.join(vposer_datadir, '%s.log' % (expr_code)))
logger('[%s] Preparing data for training VPoser.'%expr_code)
amass_splits = {
'vald': ['HumanEva', 'MPIHDM05', 'SFU', 'MPImosh'],
'test': ['Transitions_mocap', 'SSMsynced'],
'train': ['CMU', 'MPILimits', 'TotalCapture', 'EyesJapanDataset', 'KIT', 'BMLrub', 'EKUT', 'TCDhandMocap', 'ACCAD']
}
amass_splits['train'] = list(set(amass_splits['train']).difference(set(amass_splits['test'] + amass_splits['vald'])))
prepare_vposer_datasets(amass_splits, amass_dir, vposer_datadir, logger=logger)
def __init__(self, work_dir, ps):
from tensorboardX import SummaryWriter
from human_body_prior.data.dataloader import VPoserDS
self.pt_dtype = torch.float64 if ps.fp_precision == '64' else torch.float32
torch.manual_seed(ps.seed)
ps.work_dir = makepath(work_dir, isfile=False)
logger = log2file(os.path.join(work_dir, '%s.log' % ps.expr_code))
summary_logdir = os.path.join(work_dir, 'summaries')
self.swriter = SummaryWriter(log_dir=summary_logdir)
logger('tensorboard --logdir=%s' % summary_logdir)
logger('Torch Version: %s\n' % torch.__version__)
shutil.copy2(os.path.realpath(__file__), work_dir)
use_cuda = torch.cuda.is_available()
if use_cuda: torch.cuda.empty_cache()
self.comp_device = torch.device(
"cuda:%d" % ps.cuda_id if torch.cuda.is_available() else "cpu")
logger('%d CUDAs available!' % torch.cuda.device_count())
gpu_brand = torch.cuda.get_device_name(
ps.cuda_id) if use_cuda else None
logger('Training with %s [%s]' %
(self.comp_device,
gpu_brand) if use_cuda else 'Training on CPU!!!')
logger('Base dataset_dir is %s' % ps.dataset_dir)
kwargs = {'num_workers': ps.n_workers}
ds_train = VPoserDS(dataset_dir=os.path.join(ps.dataset_dir, 'train'))
self.ds_train = DataLoader(ds_train,
batch_size=ps.batch_size,
shuffle=True,
drop_last=True,
**kwargs)
ds_val = VPoserDS(dataset_dir=os.path.join(ps.dataset_dir, 'vald'))
self.ds_val = DataLoader(ds_val,
batch_size=ps.batch_size,
shuffle=True,
drop_last=True,
**kwargs)
ds_test = VPoserDS(dataset_dir=os.path.join(ps.dataset_dir, 'test'))
self.ds_test = DataLoader(ds_test,
batch_size=ps.batch_size,
shuffle=True,
drop_last=True,
**kwargs)
logger('Train dataset size %.2f M' %
(len(self.ds_train.dataset) * 1e-6))
logger('Validation dataset size %d' % len(self.ds_val.dataset))
logger('Test dataset size %d' % len(self.ds_test.dataset))
ps.data_shape = list(ds_val[0]['pose_aa'].shape)
self.vposer_model = VPoser(num_neurons=ps.num_neurons,
latentD=ps.latentD,
data_shape=ps.data_shape,
use_cont_repr=ps.use_cont_repr)
if ps.use_multigpu:
self.vposer_model = nn.DataParallel(self.vposer_model)
self.vposer_model.to(self.comp_device)
varlist = [var[1] for var in self.vposer_model.named_parameters()]
params_count = sum(p.numel() for p in varlist if p.requires_grad)
logger('Total Trainable Parameters Count is %2.2f M.' %
((params_count) * 1e-6))
self.optimizer = optim.Adam(varlist,
lr=ps.base_lr,
weight_decay=ps.reg_coef)
self.logger = logger
self.best_loss_total = np.inf
self.try_num = ps.try_num
self.epochs_completed = 0
self.ps = ps
if ps.best_model_fname is not None:
if isinstance(self.vposer_model, torch.nn.DataParallel):
self.vposer_model.module.load_state_dict(
torch.load(ps.best_model_fname,
map_location=self.comp_device))
else:
self.vposer_model.load_state_dict(
torch.load(ps.best_model_fname,
map_location=self.comp_device))
logger('Restored model from %s' % ps.best_model_fname)
chose_ids = np.random.choice(list(range(len(ds_val))),
size=ps.num_bodies_to_display,
replace=False,
p=None)
data_all = {}
for id in chose_ids:
for k, v in ds_val[id].items():
if k in data_all.keys():
data_all[k] = torch.cat([data_all[k], v[np.newaxis]],
dim=0)
else:
data_all[k] = v[np.newaxis]
self.vis_dorig = {
k: data_all[k].to(self.comp_device)
for k in data_all.keys()
}
self.bm = BodyModel(self.ps.bm_path,
'smplh',
batch_size=self.ps.batch_size,
use_posedirs=True).to(self.comp_device)
def dump_amass2pytroch(datasets,
amass_dir,
out_dir,
split_name,
logger=None,
rnd_seed=100):
'''
Select random number of frames from central 80 percent of each mocap sequence
:param datasets:
:param amass_dir:
:param out_dir:
:param split_name:
:param logger
:param rnd_seed:
:return:
'''
import glob
from tqdm import tqdm
assert split_name in ['train', 'vald', 'test']
np.random.seed(rnd_seed)
makepath(out_dir, isfile=False)
if logger is None:
starttime = datetime.now().replace(microsecond=0)
log_name = datetime.strftime(starttime, '%Y%m%d_%H%M')
logger = log2file(os.path.join(out_dir, '%s.log' % (log_name)))
logger('Creating pytorch dataset at %s' % out_dir)
if split_name in ['vald', 'test']:
keep_rate = 0.3 # this should be fixed for vald and test datasets
elif split_name == 'train':
keep_rate = 0.3 # 30 percent, which would give you around 3.5 M training data points
data_pose = []
data_betas = []
data_gender = []
data_trans = []
data_markers = []
for ds_name in datasets:
npz_fnames = glob.glob(os.path.join(amass_dir, ds_name, '*/*.npz'))
logger('randomly selecting data points from %s.' % (ds_name))
for npz_fname in tqdm(npz_fnames):
cdata = np.load(npz_fname)
N = len(cdata['poses'])
# skip first and last frames to avoid initial standard poses, e.g. T pose
cdata_ids = np.random.choice(list(
range(int(0.1 * N), int(0.9 * N), 1)),
int(keep_rate * 0.8 * N),
replace=False)
if len(cdata_ids) < 1: continue
data_pose.extend(cdata['poses'][cdata_ids].astype(np.float32))
data_trans.extend(cdata['trans'][cdata_ids].astype(np.float32))
data_betas.extend(
np.repeat(cdata['betas'][np.newaxis].astype(np.float32),
repeats=len(cdata_ids),
axis=0))
data_gender.extend([{
'male': -1,
'neutral': 0,
'female': 1
}[str(cdata['gender'].astype(np.str))] for _ in cdata_ids])
if split_name == 'test':
data_markers.extend(
np.repeat(cdata['betas'][np.newaxis].astype(np.float32),
repeats=len(cdata_ids),
axis=0))
outdir = makepath(os.path.join(out_dir, split_name))
assert len(data_pose) != 0
outpath = os.path.join(outdir, 'pose.pt')
torch.save(torch.tensor(np.asarray(data_pose, np.float32)), outpath)
outpath = os.path.join(outdir, 'betas.pt')
torch.save(torch.tensor(np.asarray(data_betas, np.float32)), outpath)
outpath = os.path.join(outdir, 'trans.pt')
torch.save(torch.tensor(np.asarray(data_trans, np.float32)), outpath)
outpath = os.path.join(outdir, 'gender.pt')
torch.save(torch.tensor(np.asarray(data_gender, np.int32)), outpath)
logger('Len. split %s %d' % (split_name, len(data_pose)))
def prepare_vposer_datasets(amass_splits,
amass_dir,
vposer_datadir,
logger=None):
if logger is None:
starttime = datetime.now().replace(microsecond=0)
log_name = datetime.strftime(starttime, '%Y%m%d_%H%M')
logger = log2file(os.path.join(vposer_datadir, '%s.log' % (log_name)))
logger('Creating pytorch dataset at %s' % vposer_datadir)
stageI_outdir = os.path.join(vposer_datadir, 'stage_I')
shutil.copy2(sys.argv[0],
os.path.join(vposer_datadir, os.path.basename(sys.argv[0])))
logger('Stage I: Fetch data from AMASS npz files')
for split_name, datasets in amass_splits.items():
if os.path.exists(os.path.join(stageI_outdir, split_name, 'pose.pt')):
continue
dump_amass2pytroch(datasets,
amass_dir,
stageI_outdir,
split_name=split_name,
logger=logger)
logger(
'Stage II: augment data by noise and save into h5 files to be used in a cross framework scenario.'
)
## Writing to h5 files is also convinient since appending to files is possible
from torch.utils.data import DataLoader
import tables as pytables
from tqdm import tqdm
class AMASS_ROW(pytables.IsDescription):
gender = pytables.Int16Col(1) # 1-character String
pose = pytables.Float32Col(52 * 3) # float (single-precision)
pose_matrot = pytables.Float32Col(52 * 9) # float (single-precision)
betas = pytables.Float32Col(16) # float (single-precision)
trans = pytables.Float32Col(3) # float (single-precision)
stageII_outdir = makepath(os.path.join(vposer_datadir, 'stage_II'))
batch_size = 256
max_num_epochs = 1 # how much augmentation we would get
for split_name in amass_splits.keys():
h5_outpath = os.path.join(stageII_outdir, '%s.h5' % split_name)
if os.path.exists(h5_outpath): continue
ds = AMASS_Augment(dataset_dir=os.path.join(stageI_outdir, split_name))
logger('%s has %d data points!' % (split_name, len(ds)))
dataloader = DataLoader(ds,
batch_size=batch_size,
shuffle=False,
num_workers=32,
drop_last=False)
with pytables.open_file(h5_outpath, mode="w") as h5file:
table = h5file.create_table('/', 'data', AMASS_ROW)
for epoch_num in range(max_num_epochs):
for bId, bData in tqdm(enumerate(dataloader)):
for i in range(len(bData['trans'])):
for k in bData.keys():
table.row[k] = c2c(bData[k][i])
table.row.append()
table.flush()
logger('Stage III: dump every thing as a final thing to pt files')
# we would like to use pt files because their interface could run in multiple threads
stageIII_outdir = makepath(os.path.join(vposer_datadir, 'stage_III'))
for split_name in amass_splits.keys():
h5_filepath = os.path.join(stageII_outdir, '%s.h5' % split_name)
if not os.path.exists(h5_filepath): continue
with pytables.open_file(h5_filepath, mode="r") as h5file:
data = h5file.get_node('/data')
data_dict = {k: [] for k in data.colnames}
for id in range(len(data)):
cdata = data[id]
for k in data_dict.keys():
data_dict[k].append(cdata[k])
for k, v in data_dict.items():
outfname = makepath(os.path.join(stageIII_outdir, split_name,
'%s.pt' % k),
isfile=True)
if os.path.exists(outfname): continue
torch.save(torch.from_numpy(np.asarray(v)), outfname)
logger('Dumped final pytorch dataset at %s' % stageIII_outdir)
def imagearray2file(img_array, outpath=None, fps=30):
'''
:param nparray: RxCxTxwidthxheightx3
:param outpath: the directory where T images will be dumped for each time point in range T
:param fps: fps of the gif file
:return:
it will return an image list with length T
if outpath is given as a png file, an image will be saved for each t in T.
if outpath is given as a gif file, an animated image with T frames will be created.
'''
import cv2
from human_body_prior.tools.omni_tools import makepath
makepath(outpath, isfile=True)
if not isinstance(img_array, np.ndarray) or img_array.ndim < 6:
raise ValueError(
'img_array should be a numpy array of shape RxCxTxwidthxheightx3')
R, C, T, img_h, img_w, img_c = img_array.shape
out_images = []
for tIdx in range(T):
row_images = []
for rIdx in range(R):
col_images = []
for cIdx in range(C):
col_images.append(img_array[rIdx, cIdx, tIdx])
row_images.append(np.hstack(col_images))
t_image = np.vstack(row_images)
out_images.append(t_image)
if outpath is not None:
if '.png' in outpath:
for tIdx in range(T):
if T > 1:
cur_outpath = outpath.replace('.png', '_%03d.png' % tIdx)
else:
cur_outpath = outpath
cv2.imwrite(cur_outpath, out_images[tIdx])
while not os.path.exists(cur_outpath):
continue # wait until the snapshot is written to the disk
elif '.gif' in outpath:
import imageio
with imageio.get_writer(outpath, mode='I', fps=fps) as writer:
for tIdx in range(T):
img = out_images[tIdx].astype(np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
writer.append_data(img)
elif '.avi' in outpath:
fourcc = cv2.VideoWriter_fourcc(*'XVID')
video = cv2.VideoWriter(outpath, fourcc, fps, (img_w, img_h), True)
for tIdx in range(T):
img = out_images[tIdx].astype(np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
video.write(img)
video.release()
cv2.destroyAllWindows()
elif '.mp4' in outpath:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
video = cv2.VideoWriter(outpath, fourcc, fps, (img_w, img_h), True)
for tIdx in range(T):
img = out_images[tIdx].astype(np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
video.write(img)
video.release()
cv2.destroyAllWindows()
return out_images
if __name__ == '__main__':
# ['CMU', 'Transitions_mocap', 'MPI_Limits', 'SSM_synced', 'TotalCapture', 'Eyes_Japan_Dataset', 'MPI_mosh', 'MPI_HDM05', 'HumanEva', 'ACCAD', 'EKUT', 'SFU', 'KIT', 'H36M', 'TCD_handMocap', 'BML']
msg = ''' Using standard AMASS dataset preparation pipeline:
0) Donwload all npz files from https://amass.is.tue.mpg.de/
1) Convert npz files to pytorch readable pt files.
2) Either use these files directly or augment them in parallel and write into h5 files
3)[optional] If you have augmented your data, dump augmented results into final pt files and use with your dataloader'''
expr_code = 'VXX_SVXX_TXX' #VERSION_SUBVERSION_TRY
amass_dir = 'PATH_TO_DOWNLOADED_NPZFILES/*/*_poses.npz'
work_dir = makepath('WHERE_YOU_WANT_YOUR_FILE_TO_BE_DUMPED/%s' %
(expr_code))
logger = log2file(os.path.join(work_dir, '%s.log' % (expr_code)))
logger('[%s] AMASS Data Preparation Began.' % expr_code)
logger(msg)
amass_splits = {
'vald': ['HumanEva', 'MPI_HDM05', 'SFU', 'MPI_mosh'],
'test': ['Transitions_mocap', 'SSM_synced'],
'train': [
'CMU', 'MPI_Limits', 'TotalCapture', 'Eyes_Japan_Dataset', 'KIT',
'BML', 'EKUT', 'TCD_handMocap'
] #ACCAD
}
amass_splits['train'] = list(
set(amass_splits['train']).difference(
def dump_amass2pytroch(datasets,
amass_dir,
out_posepath,
logger=None,
rnd_seed=100,
keep_rate=0.01):
'''
Select random number of frames from central 80 percent of each mocap sequence
Save individual data features like pose and shape per frame in pytorch pt files
test set will have the extra field for original markers
:param datasets: the name of the dataset
:param amass_dir: directory of downloaded amass npz files. should be in this structure: path/datasets/subjects/*_poses.npz
:param out_posepath: the path for final pose.pt file
:param logger: an instance of human_body_prior.tools.omni_tools.log2file
:param rnd_seed:
:return: Number of datapoints dumped using out_poseth address pattern
'''
import glob
np.random.seed(rnd_seed)
makepath(out_posepath, isfile=True)
if logger is None:
starttime = datetime.now().replace(microsecond=0)
log_name = datetime.strftime(starttime, '%Y%m%d_%H%M')
logger = log2file(
out_posepath.replace('pose.pt', '%s.log' % (log_name)))
logger('Creating pytorch dataset at %s' % out_posepath)
data_pose = []
data_betas = []
data_gender = []
data_trans = []
data_idx = []
data_frame = []
data_tightness = []
data_outfit = []
for ds_name in datasets:
npz_fnames = glob.glob(os.path.join(amass_dir, ds_name, '*/info.mat'))
logger('randomly selecting data points from %s.' % (ds_name))
for npz_fname in tqdm(npz_fnames):
try:
cdata = loadInfo(npz_fname)
cdata['idx'] = int(npz_fname.split("/")[-2])
except:
logger('Could not read %s! skipping..' % npz_fname)
continue
cdata['poses'] = cdata['poses'].T
cdata['trans'] = cdata['trans'].T
outfit_arr = np.zeros(len(outfit_types))
for key in cdata['outfit'].keys():
outfit_arr[outfit_types.index(key)] = fabric_types.index(
cdata['outfit'][key]['fabric']) + 1
if len(cdata['poses'].shape) < 2: continue
N = len(cdata['poses'])
cdata_ids = np.arange(N)
np.random.shuffle(cdata_ids)
if len(cdata_ids) < 1: continue
# try:
data_frame.extend(np.array(cdata_ids).astype(np.int32))
data_idx.extend(np.array([cdata['idx'] for _ in cdata_ids]))
data_pose.extend(cdata['poses'][cdata_ids].astype(np.float32))
data_trans.extend(cdata['trans'][cdata_ids].astype(np.float32))
data_betas.extend(
np.repeat(cdata['shape'][np.newaxis].astype(np.float32),
repeats=len(cdata_ids),
axis=0))
data_gender.extend(np.array([cdata['gender'] for _ in cdata_ids]))
data_tightness.extend(
np.repeat(cdata['tightness'][np.newaxis].astype(np.float32),
repeats=len(cdata_ids),
axis=0))
data_outfit.extend(
np.repeat(outfit_arr[np.newaxis].astype(np.int32),
repeats=len(cdata_ids),
axis=0))
# except:
# print(N, cdata['poses'].shape)
assert len(data_pose) != 0
torch.save(torch.tensor(np.asarray(data_pose, np.float32)), out_posepath)
torch.save(torch.tensor(np.asarray(data_betas, np.float32)),
out_posepath.replace('pose.pt', 'betas.pt'))
torch.save(torch.tensor(np.asarray(data_trans, np.float32)),
out_posepath.replace('pose.pt', 'trans.pt'))
torch.save(torch.tensor(np.asarray(data_gender, np.int32)),
out_posepath.replace('pose.pt', 'gender.pt'))
torch.save(torch.tensor(np.asarray(data_frame, np.int32)),
out_posepath.replace('pose.pt', 'frame.pt'))
torch.save(torch.tensor(np.asarray(data_idx, np.int32)),
out_posepath.replace('pose.pt', 'idx.pt'))
torch.save(torch.tensor(np.asarray(data_tightness, np.float32)),
out_posepath.replace('pose.pt', 'tightness.pt'))
torch.save(torch.tensor(np.asarray(data_outfit, np.int32)),
out_posepath.replace('pose.pt', 'outfit.pt'))
return len(data_pose)
def prepare_amass(amass_splits,
amass_dir,
work_dir,
logger=None,
betas_range=None,
betas_limit=None,
frame_len=None,
max_len=None,
downsample_rate=None):
if logger is None:
starttime = datetime.now().replace(microsecond=0)
log_name = datetime.strftime(starttime, '%Y%m%d_%H%M')
logger = log2file(os.path.join(work_dir, '%s.log' % (log_name)))
logger('Creating pytorch dataset at %s' % work_dir)
stageI_outdir = os.path.join(work_dir, 'stage_I')
shutil.copy2(sys.argv[0],
os.path.join(work_dir, os.path.basename(sys.argv[0])))
logger('Stage I: Fetch data from AMASS npz files')
# split mode - split a single dataset into train/vald/test with specified proportions
# e.g.
# amass_splits = {
# 'dataset': 'HumanEva',
# 'splits': (.85, .05, .1) # train, vald, test
# }
if 'splits' in amass_splits.keys():
import numbers
from functools import reduce
splits = amass_splits['splits']
_amass_splits = {}
assert [isinstance(s, numbers.Number) for s in splits] == [
True, True, True
], "amass_splits['splits'] must be (number, number, number)"
assert reduce(
lambda x, y: x + y, splits
) <= 1., "sum of amass_splits['splits'] must equal or less than 1.0"
for split_idx, split_name in enumerate(('train', 'vald', 'test')):
# if there is a zero-split, skip through the dataset creation
if split_idx > 0 and splits[split_idx] == 0: continue
final_splits = (0., 1.)
outpath = makepath(os.path.join(stageI_outdir, split_name,
'pose.pt'),
isfile=True)
# reconstruct amass_splits as normal mode for stage II and III
_amass_splits[split_name] = amass_splits['dataset']
if os.path.exists(outpath): continue
if split_name is 'train': final_splits = (0., splits[0])
elif split_name is 'vald':
final_splits = (splits[0], splits[0] + splits[1])
else:
final_splits = (splits[0] + splits[1],
splits[0] + splits[1] + splits[2])
if frame_len:
downsample_amass2pytroch(amass_splits['dataset'],
amass_dir,
outpath,
logger=logger,
betas_range=betas_range,
betas_limit=betas_limit,
splits=final_splits,
frame_len=frame_len,
max_len=max_len,
downsample_rate=downsample_rate)
else:
dump_amass2pytroch(amass_splits['dataset'],
amass_dir,
outpath,
logger=logger,
betas_range=betas_range,
betas_limit=betas_limit,
splits=final_splits,
max_len=max_len)
# assigin the reconstructed amass_splits back after stage I compeletion
amass_splits = _amass_splits
# normal mode - using different datasets as train/vald/test
# e.g.
# amass_splits = {
# 'vald': ['HumanEva'],
# 'test': ['SSM_synced'],
# 'train': ['CMU']
# }
else:
for split_name, datasets in amass_splits.items():
outpath = makepath(os.path.join(stageI_outdir, split_name,
'pose.pt'),
isfile=True)
if os.path.exists(outpath): continue
if frame_len:
downsample_amass2pytroch(datasets,
amass_dir,
outpath,
logger=logger,
betas_range=betas_range,
betas_limit=betas_limit,
frame_len=frame_len,
max_len=max_len,
downsample_rate=downsample_rate)
else:
dump_amass2pytroch(datasets,
amass_dir,
outpath,
logger=logger,
betas_range=betas_range,
betas_limit=betas_limit,
max_len=max_len)
logger(
'Stage II: augment the data and save into h5 files to be used in a cross framework scenario.'
)
class AMASS_ROW(pytables.IsDescription):
fid = pytables.Int16Col(1) # 1-character String
fname = pytables.Int32Col(1) # 1-character String
gender = pytables.Int16Col(1) # 1-character String
pose = pytables.Float32Col(52 * 3) # float (single-precision)
dmpl = pytables.Float32Col(8) # float (single-precision)
pose_matrot = pytables.Float32Col(52 * 9) # float (single-precision)
betas = pytables.Float32Col(16) # float (single-precision)
trans = pytables.Float32Col(3) # float (single-precision)
stageII_outdir = makepath(os.path.join(work_dir, 'stage_II'))
batch_size = 256
max_num_epochs = 1 # how much augmentation we would get
for split_name in amass_splits.keys():
h5_outpath = os.path.join(stageII_outdir, '%s.h5' % split_name)
if os.path.exists(h5_outpath): continue
ds = AMASS_Augment(dataset_dir=os.path.join(stageI_outdir, split_name))
logger('%s has %d data points!' % (split_name, len(ds)))
dataloader = DataLoader(ds,
batch_size=batch_size,
shuffle=False,
num_workers=32,
drop_last=False)
with pytables.open_file(h5_outpath, mode="w") as h5file:
table = h5file.create_table('/', 'data', AMASS_ROW)
for epoch_num in range(max_num_epochs):
for bId, bData in tqdm(enumerate(dataloader)):
for i in range(len(bData['trans'])):
for k in bData.keys():
table.row[k] = c2c(bData[k][i])
table.row.append()
table.flush()
logger(
'\nStage III: dump every data field for all the splits as final pytorch pt files'
)
# we would like to use pt files because their interface could run in multiple threads
stageIII_outdir = makepath(os.path.join(work_dir, 'stage_III'))
for split_name in amass_splits.keys():
h5_filepath = os.path.join(stageII_outdir, '%s.h5' % split_name)
if not os.path.exists(h5_filepath): continue
with pytables.open_file(h5_filepath, mode="r") as h5file:
data = h5file.get_node('/data')
data_dict = {k: [] for k in data.colnames}
for id in range(len(data)):
cdata = data[id]
for k in data_dict.keys():
data_dict[k].append(cdata[k])
for k, v in data_dict.items():
outfname = makepath(os.path.join(stageIII_outdir, split_name,
'%s.pt' % k),
isfile=True)
if os.path.exists(outfname): continue
torch.save(torch.from_numpy(np.asarray(v)), outfname)
logger('Dumped final pytorch dataset at %s' % stageIII_outdir)
def dump_amass2pytroch(datasets,
amass_dir,
out_posepath,
logger=None,
betas_range=None,
betas_limit=None,
splits=None,
rnd_seed=100,
keep_rate=0.01,
max_len=None):
'''
Select random number of frames from central 80 percent of each mocap sequence
Save individual data features like pose and shape per frame in pytorch pt files
test set will have the extra field for original markers
:param datasets: the name of the dataset
:param amass_dir: directory of downloaded amass npz files. should be in this structure: path/datasets/subjects/*_poses.npz
:param out_posepath: the path for final pose.pt file
:param logger: an instance of human_body_prior.tools.omni_tools.log2file
:param betas_range: variance of each beta
:param betas_limit: betas variance ranging from -betas_limit to betas_limit. only works with integer betas_range
:param splits: (splits_start, splits_end), e.g. (.85, .90) means splits 5% of the dataset starts from 85%
:param rnd_seed: random seed
:param max_len: max frame allowed
:return: Number of datapoints dumped using out_poseth address pattern
'''
import glob
np.random.seed(rnd_seed)
makepath(out_posepath, isfile=True)
if logger is None:
starttime = datetime.now().replace(microsecond=0)
log_name = datetime.strftime(starttime, '%Y%m%d_%H%M')
logger = log2file(
out_posepath.replace('pose.pt', '%s.log' % (log_name)))
logger('Creating pytorch dataset at %s' % out_posepath)
data_pose = []
data_dmpl = []
data_betas = []
data_gender = []
data_trans = []
data_fname = []
data_fid = []
for ds_name in datasets:
npz_fnames = glob.glob(
os.path.join(amass_dir, ds_name, '*/*_poses.npz'))
if splits:
logger(
f'randomly selecting {"%.1f" % ((splits[1] - splits[0]) * 100)}% data points from {ds_name}.'
)
else:
logger(f'randomly selecting data points from {ds_name}.')
for dir_id, npz_fname in enumerate(tqdm(npz_fnames)):
try:
cdata = np.load(npz_fname)
except:
logger('Could not read %s! skipping..' % npz_fname)
continue
N = len(cdata['poses'])
# fname = abs(hash(npz_fname.split('/')[-1].split('.')[0])) % (10 ** 8) # hash filename to a unique positive 8-digits integer
fname = dir_id * 1000
cdata_ids = np.random.choice(
list(range(int(0.1 * N), int(0.9 * N), 1)),
int(keep_rate * 0.8 * N),
replace=False
) # removing first and last 10% of the data to avoid repetitive initial poses
if len(cdata_ids) < 1: continue
if 'int' in str(type(betas_range)) or 'numpy.ndarray' in str(
type(betas_range)):
if 'int' in str(type(betas_range)):
if betas_range == 0:
data_pose.extend(cdata['poses'][cdata_ids].astype(
np.float32))
data_dmpl.extend(cdata['dmpls'][cdata_ids].astype(
np.float32))
data_trans.extend(cdata['trans'][cdata_ids].astype(
np.float32))
data_betas.extend(
np.repeat(cdata['betas'][np.newaxis].astype(
np.float32),
repeats=len(cdata_ids),
axis=0))
data_gender.extend([
gdr2num[str(cdata['gender'].astype(np.str))]
for _ in cdata_ids
])
data_fname.extend([fname for _ in cdata_ids])
data_fid.extend([i for i, _ in enumerate(cdata_ids)])
else:
assert betas_range % 2 == 0, ValueError(
'betas_range should be multiple to 2')
if betas_limit is None: betas_limit = 2.
# if `betas_range` is an integer,
# sample the number of betas1 and betas2
# that varience from -2. to 2. as follows:
beta1, beta2 = cdata['betas'][0], cdata['betas'][1]
# left range, right range
beta1_lr, beta1_rr = max(0., betas_limit + beta1), max(
0., betas_limit - beta1)
beta2_lr, beta2_rr = max(0., betas_limit + beta2), max(
0., betas_limit - beta2)
# left range percentage, right range percentage
beta1_lp, beta1_rp = beta1_lr / (
beta1_lr + beta1_rr), beta1_rr / (beta1_lr +
beta1_rr)
beta2_lp, beta2_rp = beta2_lr / (
beta2_lr + beta2_rr), beta2_rr / (beta2_lr +
beta2_rr)
# left range sample number
beta1_ln, beta2_ln = int(betas_range * beta1_lp), int(
betas_range * beta2_lp)
# do sampling for beta1 range
beta1_range = betas_range_sample(
betas_range, beta1, beta1_ln, betas_limit)
beta2_range = betas_range_sample(
betas_range, beta2, beta2_ln, betas_limit)
# reconstruct beatas_range as numpy.ndarray
betas_range = np.zeros(
(betas_range, len(cdata['betas'])))
betas_range[:, 0] = beta1_range
betas_range[:, 1] = beta2_range
for i, beta_delta in enumerate(betas_range):
cdata_betas = np.array(cdata['betas']).astype(np.float32)
data_pose.extend(cdata['poses'][cdata_ids].astype(
np.float32))
data_dmpl.extend(cdata['dmpls'][cdata_ids].astype(
np.float32))
data_trans.extend(cdata['trans'][cdata_ids].astype(
np.float32))
data_betas.extend(
np.repeat((cdata_betas +
beta_delta)[np.newaxis].astype(np.float32),
repeats=len(cdata_ids),
axis=0))
data_gender.extend([
gdr2num[str(cdata['gender'].astype(np.str))]
for _ in cdata_ids
])
data_fname.extend([fname + i for _ in cdata_ids])
data_fid.extend([ii for ii, _ in enumerate(cdata_ids)])
else:
data_pose.extend(cdata['poses'][cdata_ids].astype(np.float32))
data_dmpl.extend(cdata['dmpls'][cdata_ids].astype(np.float32))
data_trans.extend(cdata['trans'][cdata_ids].astype(np.float32))
data_betas.extend(
np.repeat(cdata['betas'][np.newaxis].astype(np.float32),
repeats=len(cdata_ids),
axis=0))
data_gender.extend([
gdr2num[str(cdata['gender'].astype(np.str))]
for _ in cdata_ids
])
data_fname.extend([fname for _ in cdata_ids])
data_fid.extend([i for i, _ in enumerate(cdata_ids)])
assert len(data_pose) != 0
assert len(data_pose) == len(data_dmpl) == len(data_betas) == len(
data_trans) == len(data_gender) == len(data_fname) == len(data_fid)
if splits:
import math
# split data
split_start = math.floor(len(data_pose) * splits[0])
split_end = math.floor(len(data_pose) * splits[1])
data_pose = data_pose[split_start:split_end]
data_dmpl = data_dmpl[split_start:split_end]
data_betas = data_betas[split_start:split_end]
data_trans = data_trans[split_start:split_end]
data_gender = data_gender[split_start:split_end]
data_fname = data_fname[split_start:split_end]
data_fid = data_fid[split_start:split_end]
assert len(data_pose) > 0
logger(
f'data length: {len(data_pose)}, parsing from proportion ({"%.1f" % splits[0]}, {"%.1f" % splits[1]}) to index ({split_start}, {split_end})\n\n'
)
torch.save(torch.tensor(np.asarray(data_pose, np.float32)), out_posepath)
torch.save(torch.tensor(np.asarray(data_dmpl, np.float32)),
out_posepath.replace('pose.pt', 'dmpl.pt'))
torch.save(torch.tensor(np.asarray(data_betas, np.float32)),
out_posepath.replace('pose.pt', 'betas.pt'))
torch.save(torch.tensor(np.asarray(data_trans, np.float32)),
out_posepath.replace('pose.pt', 'trans.pt'))
torch.save(torch.tensor(np.asarray(data_gender, np.int32)),
out_posepath.replace('pose.pt', 'gender.pt'))
torch.save(torch.tensor(np.asarray(data_fname, np.int32)),
out_posepath.replace('pose.pt', 'fname.pt'))
torch.save(torch.tensor(np.asarray(data_fid, np.int32)),
out_posepath.replace('pose.pt', 'fid.pt'))
return len(data_pose)
def prepare_vposer_datasets(vposer_dataset_dir,
amass_splits,
amass_dir,
logger=None):
if dataset_exists(vposer_dataset_dir):
if logger is not None:
logger('VPoser dataset already exists at {}'.format(
vposer_dataset_dir))
return
ds_logger = log2file(makepath(vposer_dataset_dir,
'dataset.log',
isfile=True),
write2file_only=True)
logger = ds_logger if logger is None else logger_sequencer(
[ds_logger, logger])
logger('Creating pytorch dataset at %s' % vposer_dataset_dir)
logger('Using AMASS body parameters from {}'.format(amass_dir))
shutil.copy2(__file__, vposer_dataset_dir)
# class AMASS_ROW(pytables.IsDescription):
#
# # gender = pytables.Int16Col(1) # 1-character String
# root_orient = pytables.Float32Col(3) # float (single-precision)
# pose_body = pytables.Float32Col(21 * 3) # float (single-precision)
# # pose_hand = pytables.Float32Col(2 * 15 * 3) # float (single-precision)
#
# # betas = pytables.Float32Col(16) # float (single-precision)
# # trans = pytables.Float32Col(3) # float (single-precision)
def fetch_from_amass(ds_names):
keep_rate = 0.3
npz_fnames = []
for ds_name in ds_names:
mosh_stageII_fnames = glob.glob(
osp.join(amass_dir, ds_name, '*/*_poses.npz'))
npz_fnames.extend(mosh_stageII_fnames)
logger('Found {} sequences from {}.'.format(
len(mosh_stageII_fnames), ds_name))
for npz_fname in npz_fnames:
cdata = np.load(npz_fname)
N = len(cdata['poses'])
# skip first and last frames to avoid initial standard poses, e.g. T pose
cdata_ids = np.random.choice(list(
range(int(0.1 * N), int(0.9 * N), 1)),
int(keep_rate * 0.8 * N),
replace=False)
if len(cdata_ids) < 1: continue
fullpose = cdata['poses'][cdata_ids].astype(np.float32)
yield {
'pose_body': fullpose[:, 3:66],
'root_orient': fullpose[:, :3]
}
for split_name, ds_names in amass_splits.items():
if dataset_exists(vposer_dataset_dir, split_names=[split_name]):
continue
logger('Preparing VPoser data for split {}'.format(split_name))
data_fields = {}
for data in fetch_from_amass(ds_names):
for k in data.keys():
if k not in data_fields: data_fields[k] = []
data_fields[k].append(data[k])
for k, v in data_fields.items():
outpath = makepath(vposer_dataset_dir,
split_name,
'{}.pt'.format(k),
isfile=True)
v = np.concatenate(v)
torch.save(torch.tensor(v), outpath)
logger('{} datapoints dumped for split {}. ds_meta_pklpath: {}'.format(
len(v), split_name, osp.join(vposer_dataset_dir, split_name)))
Configer(**{
'amass_splits': amass_splits.toDict(),
'amass_dir': amass_dir,
}).dump_settings(makepath(vposer_dataset_dir, 'settings.ini', isfile=True))
logger('Dumped final pytorch dataset at %s' % vposer_dataset_dir)
