samplegen/miniset.py at main · samjwebster/samplegen · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import tensorflow_datasets as tfds
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset
from audio_processor import AudioProcessor
import pickle as pkl

class NSynthDataset(Dataset):
    def __init__(self, audio_processor, split='test'):
        self.audio_processor = audio_processor

        num_families = 11
        num_instruments = 1006
        num_sources = 3

        max_pitch = 128
        max_velocity = 127

        ds = tfds.load('nsynth', split=split, shuffle_files=True)
        self.ds = []

        counter = 0
        for example in tfds.as_numpy(ds):
            audio = torch.from_numpy(example['audio'])
            spectrogram = self.audio_processor.signal_to_spectrogram(audio)
            if spectrogram.max() > self.audio_processor.max_freq:
                self.audio_processor.max_freq = spectrogram.max()

            family_vec = F.one_hot(torch.tensor(example["instrument"]["family"]), num_families).float()
            instrument_vec = F.one_hot(torch.tensor(example["instrument"]["label"]), num_instruments).float()
            source_vec = F.one_hot(torch.tensor(example["instrument"]["source"]), num_sources).float()

            note_vec = torch.cat([
                torch.tensor([example["pitch"]/max_pitch]),
                torch.tensor([example["velocity"]/max_velocity])
            ]).float()
            qualities_vec = torch.tensor([1 if value else 0 for value in example['qualities'].values()]).float()

            self.ds.append({
                'audio': audio,
                'spectrogram': spectrogram,
                'family': family_vec,
                'instrument': instrument_vec,
                'source': source_vec,
                'note': note_vec,
                'qualities': qualities_vec
            })
            counter += 1
            if counter == 50: # DIFFERENCE - We only get the first 50 samples!
                break

        # numpy_ds = tfds.as_numpy(ds)
        # arr_ds = [item for item in numpy_ds]
        # arr_ds = arr_ds[:50] # DIFFERENCE - We only get the first 50 samples!

        # # For every signal, compute Mel spectrogram
        # for item in arr_ds:
        #     item['audio'] = torch.from_numpy(item['audio'])
        #     spectrogram = self.audio_processor.signal_to_spectrogram(item['audio'])
        #     item['spectrogram'] = spectrogram
        #     if spectrogram.max() > self.audio_processor.max_freq:
        #         self.audio_processor.max_freq = spectrogram.max()

        # self.ds = arr_ds

    def normalize(self):
        # Now that we know the max frequency, we can normalize the spectrograms
        for item in self.ds:
            spectrogram_normalized = self.audio_processor.normalize(item['spectrogram'])
            item['spectrogram_normalized'] = spectrogram_normalized

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        return self.ds[idx]

def save_splits():
    splits = ['test', 'valid']
    audio_processor = AudioProcessor()
    datasets = {}
    for split in splits:
        datasets[split] = NSynthDataset(audio_processor, split=split)

    for split in splits:
        datasets[split].normalize()

    with open("./minisets/train.pkl", "wb") as f:
        pkl.dump(datasets['test'], f)
    with open("./minisets/validate.pkl", "wb") as f:
        pkl.dump(datasets['valid'], f)

def load_splits():
    with open("./minisets/train.pkl", "rb") as f:
        train = pkl.load(f)
    with open("./minisets/validate.pkl", "rb") as f:
        validate = pkl.load(f)

    return train, validate

if __name__ == "__main__":
    save_splits()