| import numpy as np |
| import random, torch, joblib |
| from torch.utils.data import Dataset |
| from mmengine.registry import DATASETS, FUNCTIONS |
| import time, os |
|
|
| FPS = 30 |
| @DATASETS.register_module() |
| class AnyMoMotionDataset(Dataset): |
| ''' |
| data format: |
| dict( |
| motion_272: np.ndarray (T, 272) |
| ) |
| ''' |
|
|
| def __init__(self, data_path_file='', window_size=-1, unit_length=4, mean_path='', std_path='', |
| min_motion_length=60, max_motion_length=300, interval=1): |
| super().__init__() |
| self.min_motion_length = min_motion_length |
| self.max_motion_length = max_motion_length |
| self.window_size = window_size |
| self.unit_length = unit_length |
|
|
| if isinstance(data_path_file, list): |
| self.data_paths = [] |
| for path in data_path_file: |
| with open(path, 'r') as f: |
| self.data_paths.extend([line.strip() for line in f.readlines()]) |
|
|
| elif isinstance(data_path_file, str): |
| with open(data_path_file, 'r') as f: |
| self.data_paths = [line.strip() for line in f.readlines()] |
| else: |
| raise NotImplementedError |
| |
| self.data_paths = self.data_paths[::interval] |
| self.motion_mean = torch.from_numpy(np.load(mean_path).astype(np.float32)) |
| self.motion_std = torch.from_numpy(np.load(std_path).astype(np.float32)) |
|
|
| def __len__(self): |
| return len(self.data_paths) |
| |
| def __getitem__(self, index): |
| data_infos = joblib.load(self.data_paths[index]) |
| data = [] |
| for data_info in data_infos: |
| |
| motion = torch.as_tensor(data_info['motion_272'], dtype=torch.float32) |
|
|
| |
| T = motion.shape[0] |
| if self.window_size > 0: |
| assert T >= self.window_size |
| idx = 0 if T == self.window_size else random.randint(0, T - self.window_size) |
| motion = motion[idx:idx+self.window_size] |
| motion_length = self.window_size |
| else: |
| motion_length = (T // self.unit_length) * self.unit_length |
| if motion_length > self.max_motion_length: |
| motion_length = self.max_motion_length |
| idx = random.randint(0, T - motion_length) |
| motion = motion[idx:idx+motion_length] |
| |
| motion = (motion - self.motion_mean) / self.motion_std |
| pad_motion = torch.zeros(self.max_motion_length - motion_length, motion.shape[1]).to(motion) |
| motion = torch.cat([motion, pad_motion], dim=0) |
| data.append(dict( |
| motion=motion, |
| motion_length=motion_length, |
| mean=self.motion_mean, |
| std=self.motion_std, |
| )) |
|
|
| return data |
|
|
| @FUNCTIONS.register_module() |
| def anymo_motion_collate_fn(batch): |
| new_batch = [] |
| for item in batch: |
| new_batch.extend(item) |
| motion = [item['motion'] for item in new_batch] |
| motion_length = torch.as_tensor([item['motion_length'] for item in new_batch]) |
| motion = torch.stack(motion) |
|
|
| mean = new_batch[0]['mean'] |
| std = new_batch[0]['std'] |
| return dict( |
| motion=motion, |
| motion_length=motion_length, |
|
|
| mean=mean, |
| std=std, |
| ) |
|
|
|
|
| @DATASETS.register_module() |
| class MotionDataset(Dataset): |
| ''' |
| data format: |
| dict( |
| motion_272: np.ndarray (T, 272) |
| ) |
| ''' |
|
|
| def __init__(self, split_file='', window_size=-1, unit_length=4, mean_path='', std_path='', |
| min_motion_length=60, max_motion_length=300, repeat_time=1, interval=1): |
| super().__init__() |
| self.min_motion_length = min_motion_length |
| self.max_motion_length = max_motion_length |
| self.window_size = window_size |
| self.unit_length = unit_length |
| self.repeat_time = repeat_time |
| st = time.time() |
| print('load data...') |
| data_infos = joblib.load(split_file) |
| print(f'load data with {time.time() - st}s') |
| self.data_infos = [] |
| for data_info in data_infos: |
| if data_info['motion_272'].shape[0] > self.min_motion_length: |
| self.data_infos.append(data_info) |
| self.data_infos = self.data_infos[::interval] |
| self.motion_mean = torch.from_numpy(np.load(mean_path).astype(np.float32)) |
| self.motion_std = torch.from_numpy(np.load(std_path).astype(np.float32)) |
|
|
| def __len__(self): |
| return len(self.data_infos) * self.repeat_time |
| |
| def __getitem__(self, index): |
| |
| data_info = self.data_infos[index % len(self.data_infos)] |
| |
| motion = torch.as_tensor(data_info['motion_272'], dtype=torch.float32) |
|
|
| |
| T = motion.shape[0] |
| if self.window_size > 0: |
| assert T >= self.window_size |
| idx = 0 if T == self.window_size else random.randint(0, T - self.window_size) |
| motion = motion[idx:idx+self.window_size] |
| motion_length = self.window_size |
| else: |
| motion_length = (T // self.unit_length) * self.unit_length |
| if motion_length > self.max_motion_length: |
| motion_length = self.max_motion_length |
| idx = random.randint(0, T - motion_length) |
| motion = motion[idx:idx+motion_length] |
| |
| motion = (motion - self.motion_mean) / self.motion_std |
| pad_motion = torch.zeros(self.max_motion_length - motion_length, motion.shape[1]).to(motion) |
| motion = torch.cat([motion, pad_motion], dim=0) |
|
|
| return dict( |
| motion=motion, |
| motion_length=motion_length, |
| mean=self.motion_mean, |
| std=self.motion_std, |
| ) |
|
|
| @FUNCTIONS.register_module() |
| def motion_collate_fn(batch): |
| motion = [item['motion'] for item in batch] |
| motion_length = torch.as_tensor([item['motion_length'] for item in batch]) |
| motion = torch.stack(motion) |
|
|
| mean = batch[0]['mean'] |
| std = batch[0]['std'] |
| return dict( |
| motion=motion, |
| motion_length=motion_length, |
|
|
| mean=mean, |
| std=std, |
| ) |
|
|
