Spaces:
Running
Running
| 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 | |
| ) |