1
2
3import numpy as np
4import torch
5import torchaudio
6import librosa
7import librosa.display
8import matplotlib.pyplot as plt
9import soundfile as sf
10from PIL import Image
11
12
13# Step 1: Encode Audio to Mel-Spectrogram
14def encode_audio_to_mel_spectrogram(audio_file, n_mels=128):
15 """
16 Encode an audio file to a mel-spectrogram.
17
18 Parameters:
19 - audio_file: Path to the audio file.
20 - n_mels: Number of mel bands (default: 128).
21
22 Returns:
23 - mel_spectrogram_db: Mel-spectrogram in dB scale.
24 - sample_rate: Sample rate of the audio file.
25 """
26 y, sample_rate = librosa.load(audio_file, sr=None) # Load audio
27 mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sample_rate, n_mels=n_mels)
28 mel_spectrogram_db = librosa.power_to_db(mel_spectrogram, ref=np.max) # Convert to dB
29 return mel_spectrogram_db, sample_rate
30
31# Improved Step 2: Save Mel-Spectrogram as Image
32def save_mel_spectrogram_image(mel_spectrogram_db, sample_rate, output_image='mel_spectrogram.png', method='matplotlib', figsize=(10, 4), cmap='hot'):
33 """
34 Save the mel-spectrogram as an image using the specified method.
35
36 Parameters:
37 - mel_spectrogram_db: Mel-spectrogram in dB scale.
38 - sample_rate: Sample rate of the audio file.
39 - output_image: Path to save the image.
40 - method: Method for saving ('matplotlib' or 'custom').
41 - figsize: Size of the figure for matplotlib (default: (10, 4)).
42 - cmap: Colormap for the spectrogram (default: 'hot').
43 """
44 if method == 'matplotlib':
45 plt.figure(figsize=figsize)
46 librosa.display.specshow(mel_spectrogram_db, sr=sample_rate, x_axis='time', y_axis='mel', cmap=cmap)
47 plt.colorbar(format='%+2.0f dB')
48 plt.title('Mel-Spectrogram')
49 plt.savefig(output_image)
50 plt.close()
51 print(f"Mel-spectrogram image saved using matplotlib as '{output_image}'")
52
53 elif method == 'custom':
54 # Convert dB scale to linear scale for image generation
55 mel_spectrogram_linear = librosa.db_to_power(mel_spectrogram_db)
56 # Create an image from the mel-spectrogram
57 image = image_from_spectrogram(mel_spectrogram_linear[np.newaxis, ...]) # Add channel dimension
58 # Save the image
59 image.save(output_image)
60 print(f"Mel-spectrogram image saved using custom method as '{output_image}'")
61
62 else:
63 raise ValueError("Invalid method. Choose 'matplotlib' or 'custom'.")
64
65
66# Spectrogram conversion functions
67def image_from_spectrogram(spectrogram: np.ndarray, power: float = 0.25) -> Image.Image:
68 """
69 Compute a spectrogram image from a spectrogram magnitude array.
70
71 Args:
72 spectrogram: (channels, frequency, time)
73 power: A power curve to apply to the spectrogram to preserve contrast
74
75 Returns:
76 image: (frequency, time, channels)
77 """
78 # Rescale to 0-1
79 max_value = np.max(spectrogram)
80 data = spectrogram / max_value
81
82 # Apply the power curve
83 data = np.power(data, power)
84
85 # Rescale to 0-255 and invert
86 data = 255 - (data * 255).astype(np.uint8)
87
88 # Convert to a PIL image
89 if data.shape[0] == 1:
90 image = Image.fromarray(data[0], mode="L").convert("RGB")
91 elif data.shape[0] == 2:
92 data = np.array([np.zeros_like(data[0]), data[0], data[1]]).transpose(1, 2, 0)
93 image = Image.fromarray(data, mode="RGB")
94 else:
95 raise NotImplementedError(f"Unsupported number of channels: {data.shape[0]}")
96
97 # Flip Y
98 image = image.transpose(Image.FLIP_TOP_BOTTOM)
99 return image
100
101
102# Step 3: Extract Mel-Spectrogram from Image (Direct Pixel Manipulation)
103def extract_mel_spectrogram_from_image(image_path):
104 """
105 Extract a mel-spectrogram from a saved image using pixel manipulation.
106
107 Parameters:
108 - image_path: Path to the spectrogram image file.
109
110 Returns:
111 - mel_spectrogram_db: The extracted mel-spectrogram in dB scale.
112 """
113 img = Image.open(image_path).convert('L') # Open image and convert to grayscale
114 img_array = np.array(img) # Convert to NumPy array
115 mel_spectrogram_db = img_array / 255.0 * -80 # Scale to dB range
116 return mel_spectrogram_db
117
118# Alternative Spectrogram Extraction (IFFT Method)
119def extract_spectrogram_with_ifft(mel_spectrogram_db):
120 """
121 Extracts the audio signal from a mel-spectrogram using the inverse FFT method.
122
123 Parameters:
124 - mel_spectrogram_db: The mel-spectrogram in dB scale.
125
126 Returns:
127 - audio: The reconstructed audio signal.
128 """
129 # Convert dB mel-spectrogram back to linear scale
130 mel_spectrogram = librosa.db_to_power(mel_spectrogram_db)
131
132 # Inverse mel transformation to get the audio signal
133 # Using IFFT (simplified for demonstration; typically requires phase info)
134 audio = librosa.feature.inverse.mel_to_audio(mel_spectrogram)
135
136 return audio
137
138# Step 4: Decode Mel-Spectrogram with Griffin-Lim
139def decode_mel_spectrogram_to_audio(mel_spectrogram_db, sample_rate, output_audio='griffin_reconstructed_audio.wav'):
140 """
141 Decode a mel-spectrogram into audio using Griffin-Lim algorithm.
142
143 Parameters:
144 - mel_spectrogram_db: The mel-spectrogram in dB scale.
145 - sample_rate: The sample rate for the audio file.
146 - output_audio: Path to save the reconstructed audio file.
147 """
148 # Convert dB mel-spectrogram back to linear scale
149 mel_spectrogram = librosa.db_to_power(mel_spectrogram_db)
150 # Perform Griffin-Lim to reconstruct audio
151 audio = librosa.griffinlim(mel_spectrogram)
152 # Save the generated audio
153 sf.write(output_audio, audio, sample_rate)
154 print(f"Griffin-Lim reconstructed audio saved as '{output_audio}'")
155 return audio
156
157# Step 5: Load MelGAN Vocoder
158def load_melgan_vocoder():
159 """
160 Load a lightweight pre-trained MelGAN vocoder for decoding mel-spectrograms.
161 Returns a torch MelGAN vocoder model.
162 """
163 model = torchaudio.models.MelGAN() # Load MelGAN model
164 model.eval() # Ensure the model is in evaluation mode
165 return model
166
167# Step 6: Decode Mel-Spectrogram with MelGAN
168def decode_mel_spectrogram_with_melgan(mel_spectrogram_db, sample_rate, output_audio='melgan_reconstructed_audio.wav'):
169 """
170 Decode a mel-spectrogram into audio using MelGAN vocoder.
171
172 Parameters:
173 - mel_spectrogram_db: The mel-spectrogram in dB scale.
174 - sample_rate: The sample rate for the audio file.
175 - output_audio: Path to save the reconstructed audio file.
176
177 Returns:
178 - audio: The reconstructed audio signal.
179 """
180 # Convert dB mel-spectrogram back to linear scale
181 mel_spectrogram = librosa.db_to_power(mel_spectrogram_db)
182 # Convert numpy array to torch tensor and adjust the shape
183 mel_spectrogram_tensor = torch.tensor(mel_spectrogram).unsqueeze(0) # Shape: [1, mel_bins, time_frames]
184
185 # Load the MelGAN vocoder model
186 melgan = load_melgan_vocoder()
187
188 # Pass the mel-spectrogram through MelGAN to generate audio
189 with torch.no_grad():
190 audio = melgan(mel_spectrogram_tensor).squeeze().numpy() # Squeeze to remove batch dimension
191
192 # Save the generated audio
193 sf.write(output_audio, audio, sample_rate)
194 print(f"MelGAN reconstructed audio saved as '{output_audio}'")
195 return audio
196def audio_from_waveform(samples: np.ndarray, sample_rate: int, normalize: bool = False) -> pydub.AudioSegment:
197 """
198 Convert a numpy array of samples of a waveform to an audio segment.
199
200 Args:
201 samples: (channels, samples) array
202 sample_rate: Sample rate of the audio.
203 normalize: Flag to normalize volume.
204
205 Returns:
206 pydub.AudioSegment
207 """
208 # Normalize volume to fit in int16
209 if normalize:
210 samples *= np.iinfo(np.int16).max / np.max(np.abs(samples))
211
212 # Transpose and convert to int16
213 samples = samples.transpose(1, 0).astype(np.int16)
214
215 # Write to the bytes of a WAV file
216 wav_bytes = io.BytesIO()
217 wavfile.write(wav_bytes, sample_rate, samples)
218 wav_bytes.seek(0)
219
220 # Read into pydub
221 return pydub.AudioSegment.from_wav(wav_bytes)
222
223
224def apply_filters(segment: pydub.AudioSegment, compression: bool = False) -> pydub.AudioSegment:
225 """
226 Apply post-processing filters to the audio segment to compress it and keep at a -10 dBFS level.
227
228 Args:
229 segment: The audio segment to filter.
230 compression: Flag to apply dynamic range compression.
231
232 Returns:
233 pydub.AudioSegment
234 """
235 if compression:
236 segment = pydub.effects.normalize(segment, headroom=0.1)
237 segment = segment.apply_gain(-10 - segment.dBFS)
238 segment = pydub.effects.compress_dynamic_range(
239 segment,
240 threshold=-20.0,
241 ratio=4.0,
242 attack=5.0,
243 release=50.0,
244 )
245
246 # Apply gain to desired dB level and normalize again
247 desired_db = -12
248 segment = segment.apply_gain(desired_db - segment.dBFS)
249 return pydub.effects.normalize(segment, headroom=0.1)
250
251
252def stitch_segments(segments: Sequence[pydub.AudioSegment], crossfade_s: float) -> pydub.AudioSegment:
253 """
254 Stitch together a sequence of audio segments with a crossfade between each segment.
255
256 Args:
257 segments: Sequence of audio segments to stitch.
258 crossfade_s: Duration of crossfade in seconds.
259
260 Returns:
261 pydub.AudioSegment
262 """
263 crossfade_ms = int(crossfade_s * 1000)
264 combined_segment = segments[0]
265 for segment in segments[1:]:
266 combined_segment = combined_segment.append(segment, crossfade=crossfade_ms)
267 return combined_segment
268
269
270def overlay_segments(segments: Sequence[pydub.AudioSegment]) -> pydub.AudioSegment:
271 """
272 Overlay a sequence of audio segments on top of each other.
273
274 Args:
275 segments: Sequence of audio segments to overlay.
276
277 Returns:
278 pydub.AudioSegment
279 """
280 assert len(segments) > 0
281 output: pydub.AudioSegment = segments[0]
282 for segment in segments[1:]:
283 output = output.overlay(segment)
284 return output
285
286
287
288# Step 7: Full Pipeline for Audio Processing with Customization
289def mel_spectrogram_pipeline(audio_file, output_image='mel_spectrogram.png',
290 output_audio_griffin='griffin_reconstructed_audio.wav',
291 output_audio_melgan='melgan_reconstructed_audio.wav',
292 extraction_method='pixel', # 'pixel' or 'ifft'
293 decoding_method='griffin'): # 'griffin' or 'melgan'
294 """
295 Full pipeline to encode audio to mel-spectrogram, save it as an image, extract the spectrogram from the image,
296 and decode it back to audio using the selected methods.
297
298 Parameters:
299 - audio_file: Path to the audio file to be processed.
300 - output_image: Path to save the mel-spectrogram image (default: 'mel_spectrogram.png').
301 - output_audio_griffin: Path to save the Griffin-Lim reconstructed audio.
302 - output_audio_melgan: Path to save the MelGAN reconstructed audio.
303 - extraction_method: Method for extraction ('pixel' or 'ifft').
304 - decoding_method: Method for decoding ('griffin' or 'melgan').
305 """
306 # Step 1: Encode (Audio -> Mel-Spectrogram)
307 mel_spectrogram_db, sample_rate = encode_audio_to_mel_spectrogram(audio_file)
308
309 # Step 2: Convert Mel-Spectrogram to Image and save it
310 save_mel_spectrogram_image(mel_spectrogram_db, sample_rate, output_image)
311
312 # Step 3: Extract Mel-Spectrogram from the image based on chosen method
313 if extraction_method == 'pixel':
314 extracted_mel_spectrogram_db = extract_mel_spectrogram_from_image(output_image)
315 elif extraction_method == 'ifft':
316 extracted_mel_spectrogram_db = extract_spectrogram_with_ifft(mel_spectrogram_db)
317 else:
318 raise ValueError("Invalid extraction method. Choose 'pixel' or 'ifft'.")
319
320 # Step 4: Decode based on the chosen decoding method
321 if decoding_method == 'griffin':
322 decode_mel_spectrogram_to_audio(extracted_mel_spectrogram_db, sample_rate, output_audio_griffin)
323 elif decoding_method == 'melgan':
324 decode_mel_spectrogram_with_melgan(extracted_mel_spectrogram_db, sample_rate, output_audio_melgan)
325 else:
326 raise ValueError("Invalid decoding method. Choose 'griffin' or 'melgan'.")
327
328# Example usage
329if __name__ == "__main__":
330 audio_file_path = 'your_audio_file.wav' # Specify the path to your audio file here
331 mel_spectrogram_pipeline(
332 audio_file_path,
333 output_image='mel_spectrogram.png',
334 output_audio_griffin='griffin_reconstructed_audio.wav',
335 output_audio_melgan='melgan_reconstructed_audio.wav',
336 extraction_method='pixel', # Choose 'pixel' or 'ifft'
337 decoding_method='griffin' # Choose 'griffin' or 'melgan'
338 )
339
340
341
342