A custom implementation for OSX pipeline.
| 33 | from detrsmpl.data.datasets.pipelines.transforms import Normalize |
| 34 | |
| 35 | class Cache(): |
| 36 | """A custom implementation for OSX pipeline.""" |
| 37 | def __init__(self, load_path=None): |
| 38 | if load_path is not None: |
| 39 | self.load(load_path) |
| 40 | |
| 41 | def load(self, load_path): |
| 42 | self.load_path = load_path |
| 43 | self.cache = np.load(load_path, allow_pickle=True) |
| 44 | self.data_len = self.cache['data_len'] |
| 45 | self.data_strategy = self.cache['data_strategy'] |
| 46 | assert self.data_len == len(self.cache) - 2 # data_len, data_strategy |
| 47 | self.cache = None |
| 48 | |
| 49 | @classmethod |
| 50 | def save(cls, save_path, data_list, data_strategy): |
| 51 | assert save_path is not None, 'save_path is None' |
| 52 | data_len = len(data_list) |
| 53 | cache = {} |
| 54 | for i, data in enumerate(data_list): |
| 55 | cache[str(i)] = data |
| 56 | assert len(cache) == data_len |
| 57 | # update meta |
| 58 | cache.update({'data_len': data_len, 'data_strategy': data_strategy}) |
| 59 | # import pdb; pdb.set_trace() |
| 60 | np.savez_compressed(save_path, **cache) |
| 61 | print(f'Cache saved to {save_path}.') |
| 62 | |
| 63 | # def shuffle(self): |
| 64 | # random.shuffle(self.mapping) |
| 65 | |
| 66 | def __len__(self): |
| 67 | return self.data_len |
| 68 | |
| 69 | def __getitem__(self, idx): |
| 70 | if self.cache is None: |
| 71 | self.cache = np.load(self.load_path, allow_pickle=True) |
| 72 | # mapped_idx = self.mapping[idx] |
| 73 | # cache_data = self.cache[str(mapped_idx)] |
| 74 | # print(self.cache.files) |
| 75 | cache_data = self.cache[str(idx)] |
| 76 | data = cache_data.item() |
| 77 | return data |
| 78 | |
| 79 | |
| 80 | class HumanDataset(torch.utils.data.Dataset): |