Spaces:
Running on Zero
Running on Zero
File size: 4,321 Bytes
2680bd5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | import torch
def lengths_to_mask(lengths, max_len) -> torch.Tensor:
'''
lengths: [B]
'''
mask = torch.arange(max_len, device=lengths.device).expand(len(lengths), max_len) < lengths.unsqueeze(1) # (B, max_len)
return mask
def collate_tensors(batch) -> torch.Tensor: # 合并不同长度的 tensor 为一个 batch
dims = batch[0].dim()
max_size = [max([b.size(i) for b in batch]) for i in range(dims)] # (dims,)
size = (len(batch),) + tuple(max_size) # (B, max_size[0], max_size[1], ...)
canvas = batch[0].new_zeros(size=size) # (B, max_size[0], max_size[1], ...)
for i, b in enumerate(batch):
subtensor = canvas[i]
for d in range(dims):
sub_tensor = subtensor.narrow(d, 0, b.size(d)) # 从一个 Tensor 中提取一个子区域，并返回一个视图（与原来的 Tensor 共享底层存储）
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] # TODO: to be checked!!!
databatchTensor = collate_tensors(databatch)
lenbatchTensor = torch.as_tensor(lenbatch)
maskbatchTensor = lengths_to_mask(lenbatchTensor, databatchTensor.shape[-1]).unsqueeze(1).unsqueeze(1) # (B, 1, 1, max_seq_len)
motion = databatchTensor
cond = dict(
y=dict(
mask=maskbatchTensor, # (B, 1, 1, max_seq_len)
lengths=lenbatchTensor # (B)
)
)
if 'text' in notnone_batches[0]:
textbatch = [b['text'] for b in notnone_batches]
cond['y'].update({"text" : textbatch})
if 'contact_label' in notnone_batches[0]:
contact_label_batch = [b['contact_label'] for b in notnone_batches]
cond['y'].update({"contact_label" : collate_tensors(contact_label_batch)})
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 = torch.tensor([b['action'] for b in notnone_batches]).long()
cond['y'].update({"action" : actionbatch})
if 'action_text' in notnone_batches[0]:
action_text = [b['action_text'] for b in notnone_batches]
cond['y'].update({"action_text" : action_text})
if 'prefix' in notnone_batches[0]:
cond['y'].update(dict(
prefix=collate_tensors([
b['prefix']
for b in notnone_batches
])
))
return motion, cond
def motion_action_collate(batch):
adapted_batch = [dict(
inp=torch.tensor(b[0]).float(), # (T, 2J, 3)
lengths=b[1],
**(dict(
action=b[2]
) if len(b)>2 else dict())
) for b in batch]
return collate(adapted_batch)
def motion_text_collate(batch):
adapted_batch = [
dict(
inp=torch.tensor(b[0]).float(), # (T, 2J, 3)
lengths=b[1],
text=b[2]
)
for b in batch
]
return collate(adapted_batch)
def motion_text_treble_collate(batch):
"""
Collate function for three-part text annotations.
Each text is a dict with keys: 'left', 'right', 'two_hands_relation'
"""
adapted_batch = [
dict(
inp=torch.tensor(b[0]).float(), # (T, 2J, 3)
lengths=b[1],
text=b[2] # b[2] is a dict: {'left': str, 'right': str, 'two_hands_relation': str}
)
for b in batch
]
if len(batch[0]) == 4:
for i in range(len(adapted_batch)):
adapted_batch[i]['contact_label'] = torch.tensor(batch[i][3]) # (T, N_contacts)
motion, cond = collate(adapted_batch)
# Reorganize text from list of dicts to dict of lists
if 'text' in cond['y']:
text_batch = cond['y']['text']
if isinstance(text_batch[0], dict):
cond['y']['text'] = {
'left': [t['left'] for t in text_batch],
'right': [t['right'] for t in text_batch],
'two_hands_relation': [t['two_hands_relation'] for t in text_batch]
}
return motion, cond
|