| import os | |
| import glob | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| class HumanML3DDataset(Dataset): | |
| def __init__(self, data_dir='./new_joint_vecs', window_size=100): | |
| self.data_dir = str(data_dir) | |
| self.window_size = window_size | |
| self.file_paths = glob.glob(os.path.join(self.data_dir, '*.npy')) | |
| data_dir_path = Path(self.data_dir).resolve() | |
| base_dir = data_dir_path.parent | |
| self.mean = torch.from_numpy(np.load(base_dir / 'Mean.npy')).float() | |
| self.std = torch.from_numpy(np.load(base_dir / 'Std.npy')).float() | |
| print(f"Dataset initialized: found {len(self.file_paths)} files.") | |
| def __len__(self): | |
| return len(self.file_paths) | |
| def __getitem__(self, idx): | |
| data = np.load(self.file_paths[idx]) | |
| tensor_data = torch.from_numpy(data).float() | |
| tensor_data = (tensor_data - self.mean) / self.std | |
| t_len, _ = tensor_data.shape | |
| if t_len > self.window_size: | |
| start = torch.randint(0, t_len - self.window_size + 1, (1,)).item() | |
| tensor_data = tensor_data[start : start + self.window_size, :] | |
| elif t_len < self.window_size: | |
| padding_size = self.window_size - t_len | |
| last_pose = tensor_data[-1].unsqueeze(0) | |
| padding = last_pose.repeat(padding_size, 1) | |
| tensor_data = torch.cat([tensor_data, padding], dim=0) | |
| tensor_data = tensor_data.permute(1, 0) | |
| return tensor_data | |