| import random |
| import codecs as cs |
| import numpy as np |
| from torch.utils import data |
| from rich.progress import track |
| from os.path import join as pjoin |
| from .dataset_m import MotionDataset |
| from .dataset_t2m import Text2MotionDataset |
|
|
|
|
| class MotionDatasetVQ(Text2MotionDataset): |
| def __init__( |
| self, |
| data_root, |
| split, |
| mean, |
| std, |
| max_motion_length, |
| min_motion_length, |
| win_size, |
| unit_length=4, |
| fps=20, |
| tmpFile=True, |
| tiny=False, |
| debug=False, |
| **kwargs, |
| ): |
| super().__init__(data_root, split, mean, std, max_motion_length, |
| min_motion_length, unit_length, fps, tmpFile, tiny, |
| debug, **kwargs) |
|
|
| |
| self.window_size = win_size |
| name_list = list(self.name_list) |
| for name in self.name_list: |
| motion = self.data_dict[name]["motion"] |
| if motion.shape[0] < self.window_size: |
| name_list.remove(name) |
| self.data_dict.pop(name) |
| self.name_list = name_list |
|
|
| def __len__(self): |
| return len(self.name_list) |
|
|
| def __getitem__(self, item): |
| idx = self.pointer + item |
| data = self.data_dict[self.name_list[idx]] |
| motion, length = data["motion"], data["length"] |
|
|
| idx = random.randint(0, motion.shape[0] - self.window_size) |
| motion = motion[idx:idx + self.window_size] |
| motion = (motion - self.mean) / self.std |
|
|
| return None, motion, length, None, None, None, None, |
|
|
