| import copy |
| from typing import Callable, Optional |
|
|
| import numpy as np |
| from omegaconf import DictConfig |
|
|
| import torch |
|
|
| from .base import BaseDataModule |
| from .humanml.dataset import Text2MotionDataset, MotionDataset |
| from .humanml.scripts.motion_process import recover_from_ric |
|
|
|
|
| |
| dataset_map = {'humanml3d': (263, 22), 'kit': (251, 21)} |
|
|
|
|
| class DataModule(BaseDataModule): |
|
|
| def __init__(self, |
| name: str, |
| cfg: DictConfig, |
| motion_only: bool, |
| collate_fn: Optional[Callable] = None, |
| **kwargs) -> None: |
| super().__init__(collate_fn=collate_fn) |
| self.cfg = cfg |
| self.name = name |
| self.nfeats, self.njoints = dataset_map[name] |
| self.hparams = copy.deepcopy({**kwargs, 'njoints': self.njoints}) |
| self.Dataset = MotionDataset if motion_only else Text2MotionDataset |
| sample_overrides = {"tiny": True, "progress_bar": False} |
| self._sample_set = self.get_sample_set(overrides=sample_overrides) |
|
|
| def denorm_spatial(self, hint: torch.Tensor) -> torch.Tensor: |
| raw_mean = torch.tensor(self._sample_set.raw_mean).to(hint) |
| raw_std = torch.tensor(self._sample_set.raw_std).to(hint) |
| hint = hint * raw_std + raw_mean |
| return hint |
|
|
| def norm_spatial(self, hint: torch.Tensor) -> torch.Tensor: |
| raw_mean = torch.tensor(self._sample_set.raw_mean).to(hint) |
| raw_std = torch.tensor(self._sample_set.raw_std).to(hint) |
| hint = (hint - raw_mean) / raw_std |
| return hint |
|
|
| def feats2joints(self, features: torch.Tensor) -> torch.Tensor: |
| mean = torch.tensor(self.hparams['mean']).to(features) |
| std = torch.tensor(self.hparams['std']).to(features) |
| features = features * std + mean |
| return recover_from_ric(features, self.njoints) |
|
|
| def renorm4t2m(self, features: torch.Tensor) -> torch.Tensor: |
| |
| ori_mean = torch.tensor(self.hparams['mean']).to(features) |
| ori_std = torch.tensor(self.hparams['std']).to(features) |
| eval_mean = torch.tensor(self.hparams['mean_eval']).to(features) |
| eval_std = torch.tensor(self.hparams['std_eval']).to(features) |
| features = features * ori_std + ori_mean |
| features = (features - eval_mean) / eval_std |
| return features |
|
|
| def mm_mode(self, mm_on: bool = True) -> None: |
| if mm_on: |
| self.is_mm = True |
| self.name_list = self.test_dataset.name_list |
| self.mm_list = np.random.choice(self.name_list, |
| self.cfg.TEST.MM_NUM_SAMPLES, |
| replace=False) |
| self.test_dataset.name_list = self.mm_list |
| else: |
| self.is_mm = False |
| self.test_dataset.name_list = self.name_list |
|
|
