| 27 | |
| 28 | |
| 29 | class CachedFolder(datasets.DatasetFolder): |
| 30 | def __init__( |
| 31 | self, |
| 32 | root: str, |
| 33 | ): |
| 34 | super().__init__( |
| 35 | root, |
| 36 | loader=None, |
| 37 | extensions=(".npz",), |
| 38 | ) |
| 39 | |
| 40 | def __getitem__(self, index: int): |
| 41 | """ |
| 42 | Args: |
| 43 | index (int): Index |
| 44 | |
| 45 | Returns: |
| 46 | tuple: (moments, target). |
| 47 | """ |
| 48 | path, target = self.samples[index] |
| 49 | |
| 50 | data = np.load(path) |
| 51 | if torch.rand(1) < 0.5: # randomly hflip |
| 52 | moments = data['moments'] |
| 53 | else: |
| 54 | moments = data['moments_flip'] |
| 55 | |
| 56 | return moments, target |