| 17 | |
| 18 | class PeDataset(BaseDataset): |
| 19 | def __init__(self, prefix, shuffle=False): |
| 20 | super().__init__(shuffle) |
| 21 | self.data_dir = hparams['binary_data_dir'] |
| 22 | self.prefix = prefix |
| 23 | self.hparams = hparams |
| 24 | self.sizes = np.load(f'{self.data_dir}/{self.prefix}_lengths.npy') |
| 25 | self.indexed_ds = None |
| 26 | |
| 27 | # pitch stats |
| 28 | f0_stats_fn = f'{self.data_dir}/train_f0s_mean_std.npy' |
| 29 | if os.path.exists(f0_stats_fn): |
| 30 | hparams['f0_mean'], hparams['f0_std'] = self.f0_mean, self.f0_std = np.load(f0_stats_fn) |
| 31 | hparams['f0_mean'] = float(hparams['f0_mean']) |
| 32 | hparams['f0_std'] = float(hparams['f0_std']) |
| 33 | else: |
| 34 | hparams['f0_mean'], hparams['f0_std'] = self.f0_mean, self.f0_std = None, None |
| 35 | |
| 36 | if prefix == 'test': |
| 37 | if hparams['num_test_samples'] > 0: |
| 38 | self.avail_idxs = list(range(hparams['num_test_samples'])) + hparams['test_ids'] |
| 39 | self.sizes = [self.sizes[i] for i in self.avail_idxs] |
| 40 | |
| 41 | def _get_item(self, index): |
| 42 | if hasattr(self, 'avail_idxs') and self.avail_idxs is not None: |