File size: 1,688 Bytes
da855ff |
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 |
from torch.utils.data import DataLoader
from data_loaders.tensors import collate as all_collate
from data_loaders.tensors import t2m_collate
def get_dataset_class(name):
if name == "amass":
from .amass import AMASS
return AMASS
elif name == "uestc":
from .a2m.uestc import UESTC
return UESTC
elif name == "humanact12":
from .a2m.humanact12poses import HumanAct12Poses
return HumanAct12Poses
elif name == "humanml":
from data_loaders.humanml.data.dataset import HumanML3D
return HumanML3D
elif name == "kit":
from data_loaders.humanml.data.dataset import KIT
return KIT
else:
raise ValueError(f'Unsupported dataset name [{name}]')
def get_collate_fn(name, hml_mode='train'):
if hml_mode == 'gt':
from data_loaders.humanml.data.dataset import collate_fn as t2m_eval_collate
return t2m_eval_collate
if name in ["humanml", "kit"]:
return t2m_collate
else:
return all_collate
def get_dataset(name, num_frames, split='train', hml_mode='train'):
DATA = get_dataset_class(name)
if name in ["humanml", "kit"]:
dataset = DATA(split=split, num_frames=num_frames, mode=hml_mode)
else:
dataset = DATA(split=split, num_frames=num_frames)
return dataset
def get_dataset_loader(name, batch_size, num_frames, split='train', hml_mode='train'):
dataset = get_dataset(name, num_frames, split, hml_mode)
collate = get_collate_fn(name, hml_mode)
loader = DataLoader(
dataset, batch_size=batch_size, shuffle=True,
num_workers=8, drop_last=True, collate_fn=collate
)
return loader |