| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
| import torch.nn.functional as F |
|
|
| |
|
|
| num_frequencies = None |
|
|
| |
|
|
|
|
| class CharacterDataset(Dataset): |
| def __init__( |
| self, |
| name: str, |
| dataset_dir: str, |
| standardize: bool, |
| num_feats: int, |
| num_cams: int, |
| sequential: bool, |
| num_frequencies: int, |
| min_freq: int, |
| max_freq: int, |
| load_vertices: bool, |
| **kwargs, |
| ): |
| super().__init__() |
| self.modality = "char" |
| self.name = name |
| self.dataset_dir = Path(dataset_dir) |
| self.traj_dir = self.dataset_dir / "traj" |
| self.data_dir = self.dataset_dir / self.name |
| self.vert_dir = self.dataset_dir / "vert_raw" |
| self.center_dir = self.dataset_dir / "char_raw" |
|
|
| self.filenames = None |
| self.standardize = standardize |
| if self.standardize: |
| mean_std = kwargs["standardization"] |
| self.norm_mean = torch.Tensor(mean_std["norm_mean_h"])[:, None] |
| self.norm_std = torch.Tensor(mean_std["norm_std_h"])[:, None] |
| self.velocity = mean_std["velocity"] |
|
|
| self.num_cams = num_cams |
| self.num_feats = num_feats |
| self.sequential = sequential |
| self.num_frequencies = num_frequencies |
| self.min_freq = min_freq |
| self.max_freq = max_freq |
|
|
| self.load_vertices = load_vertices |
|
|
| def __len__(self): |
| return len(self.filenames) |
|
|
| def __getitem__(self, index): |
| filename = self.filenames[index] |
|
|
| char_filename = filename + ".npy" |
| char_path = self.data_dir / char_filename |
|
|
| raw_char_feature = torch.from_numpy(np.load((char_path))).to(torch.float32) |
| padding_size = self.num_cams - raw_char_feature.shape[0] |
| padded_raw_char_feature = F.pad( |
| raw_char_feature, (0, 0, 0, padding_size) |
| ).permute(1, 0) |
|
|
| center_path = self.center_dir / char_filename |
| center_offset = torch.from_numpy(np.load(center_path)[0]).to(torch.float32) |
| if self.load_vertices: |
| vert_path = self.vert_dir / char_filename |
| raw_verts = np.load(vert_path, allow_pickle=True)[()] |
| if raw_verts["vertices"] is None: |
| num_frames = raw_char_feature.shape[0] |
| verts = torch.zeros((num_frames, 6890, 3), dtype=torch.float32) |
| padded_verts = torch.zeros( |
| (self.num_cams, 6890, 3), dtype=torch.float32 |
| ) |
| faces = torch.zeros((13776, 3), dtype=torch.int16) |
| else: |
| verts = torch.from_numpy(raw_verts["vertices"]).to(torch.float32) |
| verts -= center_offset |
| padded_verts = F.pad(verts, (0, 0, 0, 0, 0, padding_size)) |
| faces = torch.from_numpy(raw_verts["faces"]).to(torch.int16) |
|
|
| char_feature = raw_char_feature.clone() |
| if self.velocity: |
| velocity = char_feature[1:].clone() - char_feature[:-1].clone() |
| char_feature = torch.cat([raw_char_feature[0][None], velocity]) |
|
|
| if self.standardize: |
| |
| if len(self.norm_mean) == 6: |
| char_feature[0] -= self.norm_mean[:3, 0].to(raw_char_feature.device) |
| char_feature[0] /= self.norm_std[:3, 0].to(raw_char_feature.device) |
| char_feature[1:] -= self.norm_mean[3:, 0].to(raw_char_feature.device) |
| char_feature[1:] /= self.norm_std[3:, 0].to(raw_char_feature.device) |
| |
| else: |
| char_feature -= self.norm_mean[:, 0].to(raw_char_feature.device) |
| char_feature /= self.norm_std[:, 0].to(raw_char_feature.device) |
| padded_char_feature = F.pad( |
| char_feature, |
| (0, 0, 0, self.num_cams - char_feature.shape[0]), |
| ) |
|
|
| if self.sequential: |
| padded_char_feature = padded_char_feature.permute(1, 0) |
| else: |
| padded_char_feature = padded_char_feature.reshape(-1) |
|
|
| raw_feats = {"char_raw_feat": padded_raw_char_feature} |
| if self.load_vertices: |
| raw_feats["char_vertices"] = padded_verts |
| raw_feats["char_faces"] = faces |
|
|
| return char_filename, padded_char_feature, raw_feats |
|
|
