samplegen/data.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
import tensorflow_datasets as tfds
import torch
import torch.nn.functional as F
import numpy as np
import random
from torch.utils.data import Dataset
from audio_processor import AudioProcessor

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

        if split == 'train_subset':
            ds = tfds.load('nsynth', split='train', shuffle_files=True)
            examples = tfds.as_numpy(ds)
            examples = list(examples)
            random.shuffle(examples)
            examples = examples[:8192]  # Limiting to 8,192 random samples (The original dataset is like 300k samples, would take too long to train on)
        elif split in ['test', 'valid']:
            ds = tfds.load('nsynth', split=split, shuffle_files=True)
            examples = tfds.as_numpy(ds)
            examples = list(examples)
        else:
            raise ValueError(f"Unknown split: {split}")

        self.ds = []

        for example in examples:
            id = bytearray(example['id'])
            if len(id) < 32:
                id += b'\0' * (32 - len(id))
            elif len(id) > 32:
                id = id[:32]
            id = np.array(id)
            id = torch.tensor(id).float()

            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({
                'id': id,
                'audio': audio,
                'spectrogram': spectrogram,
                'family': family_vec,
                'instrument': instrument_vec,
                'source': source_vec,
                'note': note_vec,
                'qualities': qualities_vec
            })

    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 get_split(split):
    audio_processor = AudioProcessor()
    ds = NSynthDataset(audio_processor, split=split)
    ds.normalize()

    print(f"Loaded split: {split}")
    return ds

def get_splits(splits = ['test']):
    audio_processor = AudioProcessor()
    datasets = []
    for split in splits:
        datasets.append(NSynthDataset(audio_processor, split=split))

    for ds in datasets:
        ds.normalize()

    print(f"Loaded splits: {', '.join(splits)}")
    return datasets