samplegen/audio_processor.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
import torch
import torchaudio
import numpy as np
import os
import soundfile as sf
import matplotlib.pyplot as plt

class AudioProcessor:
    def __init__(self):
        # Configuration
        self.duration = 4
        self.sample_rate = 16000
        self.hop_length = 502
        self.n_fft = 2048
        self.n_stft = int((self.n_fft//2) + 1)
        self.n_mels = 128

        self.max_freq = -float("inf")

        self.mel_module = torchaudio.transforms.MelSpectrogram(
            sample_rate=self.sample_rate,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            n_mels=self.n_mels
        )

        self.inv_mel_module = torchaudio.transforms.InverseMelScale(
            n_stft=self.n_stft,
            n_mels=self.n_mels,
            sample_rate=self.sample_rate
        )

        self.gl_module = torchaudio.transforms.GriffinLim(
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            length=self.duration*self.sample_rate
        )

    def signal_to_spectrogram(self, signal):
        spectrogram = self.mel_module(signal).unsqueeze(0)
        return spectrogram

    def spectrogram_to_signal(self, spectrogram):
        inv_mel = self.inv_mel_module(spectrogram)
        signal = self.gl_module(inv_mel).squeeze(0)
        return signal

    def batch_spectrograms_to_signal(self, spectrograms):
        inv_mels = self.inv_mel_module(spectrograms)
        signals = self.gl_module(inv_mels).squeeze(1)
        return signals

    def normalize(self, arr):
        return arr/self.max_freq

    def denormalize(self, arr):
        return arr*self.max_freq

    def save_spectrogram(self, spectrogram, name):
        save_path = os.path.join(self.spectrograms_dir, name + ".npy")
        np.save(save_path, spectrogram)
        return save_path

    def reconstruct(self, norm_spectrogram, name="reconstructed", save_dir="./generated"):
        # Only normalize values greater than 0.01 to limit noisiness
        norm_spectrogram[norm_spectrogram < 0.005] = 0

        spectrogram = self.denormalize(norm_spectrogram)
        signal = self.spectrogram_to_signal(spectrogram)
        save_path = os.path.join(save_dir, name+".wav")

        self.save_signal(signal, save_path)
        return signal

    def save_signal(self, signal, path):
        sf.write(path, signal, self.sample_rate)

    def display(self, spectrogram, name):
        # librosa.display.specshow(spectrogram, sr=self.sample_rate, hop_length=self.hop_length, x_axis='time', y_axis='mel')
        plt.figure(figsize=(10, 4))
        plt.imshow(spectrogram, aspect='auto', origin='lower')
        plt.colorbar(format='%+2.0f dB')
        plt.title(name)
        plt.xlabel('Time')
        plt.ylabel('Frequency')
        plt.tight_layout()
        plt.show()