| from torch.utils.data import DataLoader |
| from data_loaders.tensors import collate as all_collate |
| from data_loaders.tensors import t2m_collate, t2m_prefix_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 |
| elif name == "preprocessed_posterior": |
| from data_loaders.preprocessed_posterior_loader import PreprocessedPosteriorDataset |
| return PreprocessedPosteriorDataset |
| else: |
| raise ValueError(f'Unsupported dataset name [{name}]') |
|
|
| def get_collate_fn(name, hml_mode='train', pred_len=0, batch_size=1): |
| 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"]: |
| if pred_len > 0: |
| return lambda x: t2m_prefix_collate(x, pred_len=pred_len) |
| return lambda x: t2m_collate(x, batch_size) |
| elif name == "preprocessed_posterior": |
| from data_loaders.preprocessed_posterior_loader import collate_preprocessed_posterior |
| return collate_preprocessed_posterior |
| else: |
| return all_collate |
|
|
|
|
| def get_dataset(name, num_frames, split='train', hml_mode='train', abs_path='.', fixed_len=0, |
| device=None, autoregressive=False, cache_path=None, posterior_dir=None, |
| max_samples=None, resample=False): |
| DATA = get_dataset_class(name) |
| if name in ["humanml", "kit"]: |
| dataset = DATA(split=split, num_frames=num_frames, mode=hml_mode, abs_path=abs_path, fixed_len=fixed_len, |
| device=device, autoregressive=autoregressive) |
| elif name == "preprocessed_posterior": |
| dataset = DATA(posterior_dir=posterior_dir, max_samples=max_samples, resample=resample) |
| 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', fixed_len=0, pred_len=0, |
| device=None, autoregressive=False, posterior_dir=None, max_samples=None, resample=False, |
| num_workers=8, shuffle=True): |
| """ |
| Get a data loader. |
| |
| Args: |
| name: dataset name; supports 'humanml', 'kit', 'preprocessed_posterior', etc. |
| batch_size: batch size |
| num_frames: number of frames |
| split: dataset split (train/val/test) |
| hml_mode: HumanML mode |
| fixed_len: fixed length |
| pred_len: prediction length |
| device: device |
| autoregressive: whether to use autoregressive mode |
| posterior_dir: directory path for the preprocessed_posterior dataset |
| max_samples: maximum number of samples (for preprocessed_posterior) |
| resample: whether to resample (for preprocessed_posterior) |
| num_workers: number of data loading workers |
| shuffle: whether to shuffle data |
| |
| Returns: |
| DataLoader object |
| """ |
| dataset = get_dataset(name, num_frames, split=split, hml_mode=hml_mode, fixed_len=fixed_len, |
| device=device, autoregressive=autoregressive, posterior_dir=posterior_dir, |
| max_samples=max_samples, resample=resample) |
| |
| collate = get_collate_fn(name, hml_mode, pred_len, batch_size) |
|
|
| loader = DataLoader( |
| dataset, batch_size=batch_size, shuffle=shuffle, |
| num_workers=num_workers, drop_last=True, collate_fn=collate |
| ) |
|
|
| return loader |
|
|
|
|
| def get_preprocessed_posterior_loader(posterior_dir, batch_size=32, shuffle=True, |
| num_workers=8, max_samples=None, resample=False, |
| mean=None, std=None, drop_last=True): |
| """ |
| Convenience function for creating a preprocessed posterior data loader. |
| |
| Args: |
| posterior_dir: directory containing preprocessed posterior files |
| batch_size: batch size |
| shuffle: whether to shuffle |
| num_workers: number of data loading workers |
| max_samples: maximum number of samples |
| resample: whether to resample the latent while loading |
| mean: mean used for normalization |
| std: standard deviation used for normalization |
| drop_last: whether to drop the last incomplete batch |
| |
| Returns: |
| DataLoader object with a dataset attribute |
| """ |
| from data_loaders.preprocessed_posterior_loader import get_preprocessed_posterior_loader as _get_loader |
| return _get_loader( |
| posterior_dir=posterior_dir, |
| batch_size=batch_size, |
| shuffle=shuffle, |
| num_workers=num_workers, |
| max_samples=max_samples, |
| resample=resample, |
| mean=mean, |
| std=std, |
| drop_last=drop_last |
| ) |
