nightfury commited on
Commit
8da77a8
1 Parent(s): 49a60a3

Create spectro.py

Browse files
Files changed (1) hide show
  1. spectro.py +185 -0
spectro.py ADDED
@@ -0,0 +1,185 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Audio processing tools to convert between spectrogram images and waveforms.
3
+ """
4
+ import io
5
+ import typing as T
6
+
7
+ import numpy as np
8
+ from PIL import Image
9
+ import pydub
10
+ from scipy.io import wavfile
11
+ import torch
12
+ import torchaudio
13
+
14
+
15
+ def wav_bytes_from_spectrogram_image(image: Image.Image) -> T.Tuple[io.BytesIO, float]:
16
+ """
17
+ Reconstruct a WAV audio clip from a spectrogram image. Also returns the duration in seconds.
18
+ """
19
+
20
+ max_volume = 50
21
+ power_for_image = 0.25
22
+ Sxx = spectrogram_from_image(image, max_volume=max_volume, power_for_image=power_for_image)
23
+
24
+ sample_rate = 44100 # [Hz]
25
+ clip_duration_ms = 5000 # [ms]
26
+
27
+ bins_per_image = 512
28
+ n_mels = 512
29
+
30
+ # FFT parameters
31
+ window_duration_ms = 100 # [ms]
32
+ padded_duration_ms = 400 # [ms]
33
+ step_size_ms = 10 # [ms]
34
+
35
+ # Derived parameters
36
+ num_samples = int(image.width / float(bins_per_image) * clip_duration_ms) * sample_rate
37
+ n_fft = int(padded_duration_ms / 1000.0 * sample_rate)
38
+ hop_length = int(step_size_ms / 1000.0 * sample_rate)
39
+ win_length = int(window_duration_ms / 1000.0 * sample_rate)
40
+
41
+ samples = waveform_from_spectrogram(
42
+ Sxx=Sxx,
43
+ n_fft=n_fft,
44
+ hop_length=hop_length,
45
+ win_length=win_length,
46
+ num_samples=num_samples,
47
+ sample_rate=sample_rate,
48
+ mel_scale=True,
49
+ n_mels=n_mels,
50
+ max_mel_iters=200,
51
+ num_griffin_lim_iters=32,
52
+ )
53
+
54
+ wav_bytes = io.BytesIO()
55
+ wavfile.write(wav_bytes, sample_rate, samples.astype(np.int16))
56
+ wav_bytes.seek(0)
57
+
58
+ duration_s = float(len(samples)) / sample_rate
59
+
60
+ return wav_bytes, duration_s
61
+
62
+
63
+ def spectrogram_from_image(
64
+ image: Image.Image, max_volume: float = 50, power_for_image: float = 0.25
65
+ ) -> np.ndarray:
66
+ """
67
+ Compute a spectrogram magnitude array from a spectrogram image.
68
+
69
+ TODO(hayk): Add image_from_spectrogram and call this out as the reverse.
70
+ """
71
+ # Convert to a numpy array of floats
72
+ data = np.array(image).astype(np.float32)
73
+
74
+ # Flip Y take a single channel
75
+ data = data[::-1, :, 0]
76
+
77
+ # Invert
78
+ data = 255 - data
79
+
80
+ # Rescale to max volume
81
+ data = data * max_volume / 255
82
+
83
+ # Reverse the power curve
84
+ data = np.power(data, 1 / power_for_image)
85
+
86
+ return data
87
+
88
+
89
+ def spectrogram_from_waveform(
90
+ waveform: np.ndarray,
91
+ sample_rate: int,
92
+ n_fft: int,
93
+ hop_length: int,
94
+ win_length: int,
95
+ mel_scale: bool = True,
96
+ n_mels: int = 512,
97
+ ) -> np.ndarray:
98
+ """
99
+ Compute a spectrogram from a waveform.
100
+ """
101
+
102
+ spectrogram_func = torchaudio.transforms.Spectrogram(
103
+ n_fft=n_fft,
104
+ power=None,
105
+ hop_length=hop_length,
106
+ win_length=win_length,
107
+ )
108
+
109
+ waveform_tensor = torch.from_numpy(waveform.astype(np.float32)).reshape(1, -1)
110
+ Sxx_complex = spectrogram_func(waveform_tensor).numpy()[0]
111
+
112
+ Sxx_mag = np.abs(Sxx_complex)
113
+
114
+ if mel_scale:
115
+ mel_scaler = torchaudio.transforms.MelScale(
116
+ n_mels=n_mels,
117
+ sample_rate=sample_rate,
118
+ f_min=0,
119
+ f_max=10000,
120
+ n_stft=n_fft // 2 + 1,
121
+ norm=None,
122
+ mel_scale="htk",
123
+ )
124
+
125
+ Sxx_mag = mel_scaler(torch.from_numpy(Sxx_mag)).numpy()
126
+
127
+ return Sxx_mag
128
+
129
+
130
+ def waveform_from_spectrogram(
131
+ Sxx: np.ndarray,
132
+ n_fft: int,
133
+ hop_length: int,
134
+ win_length: int,
135
+ num_samples: int,
136
+ sample_rate: int,
137
+ mel_scale: bool = True,
138
+ n_mels: int = 512,
139
+ max_mel_iters: int = 200,
140
+ num_griffin_lim_iters: int = 32,
141
+ device: str = "cuda:0",
142
+ ) -> np.ndarray:
143
+ """
144
+ Reconstruct a waveform from a spectrogram.
145
+
146
+ This is an approximate inverse of spectrogram_from_waveform, using the Griffin-Lim algorithm
147
+ to approximate the phase.
148
+ """
149
+ Sxx_torch = torch.from_numpy(Sxx).to(device)
150
+
151
+ # TODO(hayk): Make this a class that caches the two things
152
+
153
+ if mel_scale:
154
+ mel_inv_scaler = torchaudio.transforms.InverseMelScale(
155
+ n_mels=n_mels,
156
+ sample_rate=sample_rate,
157
+ f_min=0,
158
+ f_max=10000,
159
+ n_stft=n_fft // 2 + 1,
160
+ norm=None,
161
+ mel_scale="htk",
162
+ max_iter=max_mel_iters,
163
+ ).to(device)
164
+
165
+ Sxx_torch = mel_inv_scaler(Sxx_torch)
166
+
167
+ griffin_lim = torchaudio.transforms.GriffinLim(
168
+ n_fft=n_fft,
169
+ win_length=win_length,
170
+ hop_length=hop_length,
171
+ power=1.0,
172
+ n_iter=num_griffin_lim_iters,
173
+ ).to(device)
174
+
175
+ waveform = griffin_lim(Sxx_torch).cpu().numpy()
176
+
177
+ return waveform
178
+
179
+
180
+ def mp3_bytes_from_wav_bytes(wav_bytes: io.BytesIO) -> io.BytesIO:
181
+ mp3_bytes = io.BytesIO()
182
+ sound = pydub.AudioSegment.from_wav(wav_bytes)
183
+ sound.export(mp3_bytes, format="mp3")
184
+ mp3_bytes.seek(0)
185
+ return mp3_bytes