Spaces:
Runtime error
Runtime error
import torch | |
def lengths_to_mask(lengths, max_len): | |
# max_len = max(lengths) | |
mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1) | |
return mask | |
def collate_tensors(batch): | |
dims = batch[0].dim() | |
max_size = [max([b.size(i) for b in batch]) for i in range(dims)] | |
size = (len(batch),) + tuple(max_size) | |
canvas = batch[0].new_zeros(size=size) | |
for i, b in enumerate(batch): | |
sub_tensor = canvas[i] | |
for d in range(dims): | |
sub_tensor = sub_tensor.narrow(d, 0, b.size(d)) | |
sub_tensor.add_(b) | |
return canvas | |
def collate(batch): | |
notnone_batches = [b for b in batch if b is not None] | |
databatch = [b['inp'] for b in notnone_batches] | |
if 'lengths' in notnone_batches[0]: | |
lenbatch = [b['lengths'] for b in notnone_batches] | |
else: | |
lenbatch = [len(b['inp'][0][0]) for b in notnone_batches] | |
databatchTensor = collate_tensors(databatch) | |
lenbatchTensor = torch.as_tensor(lenbatch) | |
maskbatchTensor = lengths_to_mask(lenbatchTensor, databatchTensor.shape[-1]).unsqueeze(1).unsqueeze(1) # unqueeze for broadcasting | |
motion = databatchTensor | |
cond = {'y': {'mask': maskbatchTensor, 'lengths': lenbatchTensor}} | |
if 'text' in notnone_batches[0]: | |
textbatch = [b['text'] for b in notnone_batches] | |
cond['y'].update({'text': textbatch}) | |
if 'tokens' in notnone_batches[0]: | |
textbatch = [b['tokens'] for b in notnone_batches] | |
cond['y'].update({'tokens': textbatch}) | |
if 'action' in notnone_batches[0]: | |
actionbatch = [b['action'] for b in notnone_batches] | |
cond['y'].update({'action': torch.as_tensor(actionbatch).unsqueeze(1)}) | |
# collate action textual names | |
if 'action_text' in notnone_batches[0]: | |
action_text = [b['action_text']for b in notnone_batches] | |
cond['y'].update({'action_text': action_text}) | |
# 'st_idx': b['st_idx'], | |
# 'ed_idx': b['ed_idx'], | |
if 'st_idx' in notnone_batches[0]: | |
st_idx_batch = [b['st_idx'] for b in notnone_batches] | |
cond['y'].update({'st_idx': torch.tensor(st_idx_batch, dtype=torch.long)}) | |
if 'ed_idx' in notnone_batches[0]: | |
st_idx_batch = [b['ed_idx'] for b in notnone_batches] | |
cond['y'].update({'ed_idx': torch.tensor(st_idx_batch, dtype=torch.long)}) | |
# 'pert_verts': b['pert_verts'], | |
# 'verts': b['verts'], | |
if 'pert_verts' in notnone_batches[0]: | |
pert_verts_batch = [b['pert_verts'] for b in notnone_batches] | |
cond['y'].update({'pert_verts': torch.stack(pert_verts_batch, dim=0)}) | |
if 'verts' in notnone_batches[0]: | |
verts_batch = [b['verts'] for b in notnone_batches] | |
cond['y'].update({'verts': torch.stack(verts_batch, dim=0)}) | |
# 'avg_joints': b['avg_joints'], | |
# 'std_joints': b['std_joints'], | |
# ### avg_joints, std_joints ### # | |
if 'avg_joints' in notnone_batches[0]: | |
verts_batch = [b['avg_joints'] for b in notnone_batches] | |
cond['y'].update({'avg_joints': torch.stack(verts_batch, dim=0)}) | |
if 'std_joints' in notnone_batches[0]: # | |
verts_batch = [b['std_joints'] for b in notnone_batches] | |
cond['y'].update({'std_joints': torch.stack(verts_batch, dim=0)}) | |
# 'object_id': b['object_id'], | |
# 'object_global_orient': b['object_global_orient'], | |
# 'object_transl': b['object_transl'], | |
if 'object_global_orient' in notnone_batches[0]: # nnbsz x nnframes x 3 --> orientations # | |
obj_global_orient_batch = [b['object_global_orient'] for b in notnone_batches] | |
cond['y'].update({'object_global_orient': torch.stack(obj_global_orient_batch, dim=0)}) | |
if 'object_transl' in notnone_batches[0]: # nnbsz x nnframes x 3 --> orientations # | |
obj_transl_batch = [b['object_transl'] for b in notnone_batches] | |
cond['y'].update({'object_transl': torch.stack(obj_transl_batch, dim=0)}) | |
if 'object_id' in notnone_batches[0]: # nnbsz x nnframes x 3 --> orientations # | |
# nn_bsz --> a one-D tensor here # | |
object_id_batch = [b['object_id'] for b in notnone_batches] | |
cond['y'].update({'object_id': torch.tensor(object_id_batch, dtype=torch.long)}) | |
# 'sampled_base_pts_nearest_obj_pc': b['sampled_base_pts_nearest_obj_pc'], | |
# 'sampled_base_pts_nearest_obj_vns': b['sampled_base_pts_nearest_obj_vns'], | |
return motion, cond | |
# an adapter to our collate func | |
def t2m_collate(batch): | |
# batch.sort(key=lambda x: x[3], reverse=True) # nf x nnjoints x 3 -> # | |
adapted_batch = [{ | |
'inp': torch.tensor(b[4].T).float().unsqueeze(1), | |
# 'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen] | |
'text': b[2], #b[0]['caption'] | |
'tokens': b[6], | |
'lengths': b[5], | |
} for b in batch] | |
return collate(adapted_batch) | |
# an adapter to our collate func | |
def motion_ours_collate(batch): | |
# batch.sort(key=lambda x: x[3], reverse=True) # nf x nnjoints x 3 -> # | |
adapted_batch = [{ | |
'inp': b[4].permute(1, 2, 0).contiguous(), | |
# 'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen] | |
'text': b[2], #b[0]['caption'] | |
'tokens': b[6], | |
'lengths': b[5], | |
} for b in batch] | |
return collate(adapted_batch) | |
# An adapter for ours-single-seq collect function # | |
def motion_ours_singe_seq_collate(batch): | |
# batch.sort(key=lambda x: x[3], reverse=True) # nf x nnjoints x 3 -> # | |
# 'pert_verts': pert_rhand_verts, | |
# 'verts': rhand_verts, | |
adapted_batch = [{ | |
'inp': b['motion'].permute(1, 2, 0).contiguous(), | |
# 'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen] | |
'text': b['caption'], #b[0]['caption'] | |
'tokens': b['tokens'], | |
'lengths': b['m_length'], | |
'st_idx': b['st_idx'], | |
'ed_idx': b['ed_idx'], | |
'pert_verts': b['pert_verts'], | |
'verts': b['verts'], | |
'avg_joints': b['avg_joints'], | |
'std_joints': b['std_joints'], | |
'object_id': b['object_id'], | |
'object_global_orient': b['object_global_orient'], | |
'object_transl': b['object_transl'], | |
} for b in batch] | |
return collate(adapted_batch) | |
# # An adapter for ours-single-seq collect function # | |
# def motion_ours_singe_seq_collate_ambient_objbase(batch): | |
# # batch.sort(key=lambda x: x[3], reverse=True) # nf x nnjoints x 3 -> # | |
# # 'pert_verts': pert_rhand_verts, | |
# # 'verts': rhand_verts, | |
# adapted_batch = [{ | |
# 'inp': b['motion'].permute(1, 2, 0).contiguous(), | |
# # 'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen] | |
# 'text': b['caption'], #b[0]['caption'] | |
# 'tokens': b['tokens'], | |
# 'lengths': b['m_length'], | |
# 'st_idx': b['st_idx'], | |
# 'ed_idx': b['ed_idx'], | |
# 'pert_verts': b['pert_verts'], | |
# 'verts': b['verts'], | |
# 'avg_joints': b['avg_joints'], | |
# 'std_joints': b['std_joints'], | |
# 'object_id': b['object_id'], | |
# 'object_global_orient': b['object_global_orient'], | |
# 'object_transl': b['object_transl'], | |
# 'sampled_base_pts_nearest_obj_pc': b['sampled_base_pts_nearest_obj_pc'], | |
# 'sampled_base_pts_nearest_obj_vns': b['sampled_base_pts_nearest_obj_vns'], | |
# } for b in batch] | |
# return collate(adapted_batch) | |
# An adapter for ours-single-seq collect function # | |
def motion_ours_obj_base_rel_dist_collate(batch): | |
# batch.sort(key=lambda x: x[3], reverse=True) # nf x nnjoints x 3 -> # | |
# 'pert_verts': pert_rhand_verts, | |
# 'verts': rhand_verts, | |
## motion ours, obj_base rel_dist, ... ## | |
batch_dict = {} | |
data_keys = batch[0].keys() | |
for k in data_keys: | |
if k == 'caption': | |
batched_caption = [b[k] for b in batch] | |
batch_dict.update({'text': batched_caption}) | |
elif k == 'lengths': | |
batched_length = [b[k] for b in batch] | |
max_length = max(batched_length) | |
batched_length = torch.tensor(batched_length, dtype=torch.long) | |
batch_dict.update({'lengths': batched_length}) | |
# maskbatchTensor = lengths_to_mask(lenbatchTensor, databatchTensor.shape[-1]).unsqueeze(1).unsqueeze(1) # unqueeze for broadcasting | |
masks = lengths_to_mask(batched_length, max_length) | |
batch_dict.update({'mask': masks}) | |
elif k in ["object_id", 'st_idx', 'ed_idx']: | |
batched_kdata = [b[k] for b in batch] | |
batched_kdata = torch.tensor(batched_kdata, dtype=torch.long) | |
batch_dict.update({k: batched_kdata}) | |
elif k in ["obj_verts", "obj_normals", "obj_faces"]: | |
batched_kdata = [b[k] for b in batch] | |
batch_dict.update({k: batched_kdata}) | |
else: | |
batched_kdata = [b[k] for b in batch] | |
batched_kdata = torch.stack(batched_kdata, dim=0) | |
batch_dict.update({k: batched_kdata}) | |
return batch_dict | |
# batch_dict = { | |
# k: [] | |
# } | |
# for b in batch: | |
# for k in : | |
# adapted_batch = [{ | |
# '' | |
# 'inp': b['motion'].permute(1, 2, 0).contiguous(), | |
# # 'inp': torch.tensor(b[4].T).float().unsqueeze(1), # [seqlen, J] -> [J, 1, seqlen] | |
# 'text': b['caption'], #b[0]['caption'] | |
# 'tokens': b['tokens'], | |
# 'lengths': b['m_length'], | |
# 'st_idx': b['st_idx'], | |
# 'ed_idx': b['ed_idx'], | |
# 'pert_verts': b['pert_verts'], | |
# 'verts': b['verts'], | |
# 'avg_joints': b['avg_joints'], | |
# 'std_joints': b['std_joints'], | |
# 'object_id': b['object_id'], | |
# 'object_global_orient': b['object_global_orient'], | |
# 'object_transl': b['object_transl'], | |
# } for b in batch] | |
# return collate(adapted_batch) | |