Added stereo to mono conversion

The model does not support stereo audio. Separate channel enhancement can be experimented on in the future.

functions.py

@@ -1,83 +1,92 @@
-import torchaudio
-from torch import hamming_window, log10, no_grad, exp
-
-
-def return_input(user_input):
-    if user_input is None:
-        return None
-    return user_input
-
-def load_audio(audio_path):
-
-    audio_tensor, sr = torchaudio.load(audio_path)
-    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 16000)
-    audio_tensor.type()
-    return audio_tensor
-
-def load_audio_numpy(audio_path):
-    audio_tensor, sr = torchaudio.load(audio_path)
-    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 16000)
-    audio_array = audio_tensor.numpy()
-    return (16000, audio_array.ravel())
-
-def audio_to_spectrogram(audio):
-    transform_fn = torchaudio.transforms.Spectrogram(n_fft=512, hop_length=512//4, power=None, window_fn=hamming_window)
-    spectrogram = transform_fn(audio)
-    return spectrogram
-
-def extract_magnitude_and_phase(spectrogram):
-    magnitude, phase = spectrogram.abs(), spectrogram.angle()
-    return magnitude, phase
-
-def amplitude_to_db(magnitude_spec):
-    max_amplitude = magnitude_spec.max()
-    db_spectrogram = torchaudio.functional.amplitude_to_DB(magnitude_spec, 20, 10e-10, log10(max_amplitude), 100.0)
-    return db_spectrogram, max_amplitude
-
-def min_max_scaling(spectrogram, scaler):
-    # Min-Max scaling (soundness of the math is questionable due to the use of each spectrograms' max value during decibel-scaling)
-    spectrogram = scaler.transform(spectrogram)
-    return spectrogram
-
-def inverse_min_max(spectrogram, scaler):
-    spectrogram = scaler.inverse_transform(spectrogram)
-    return spectrogram
-
-def db_to_amplitude(db_spectrogram, max_amplitude):
-    return max_amplitude * 10**(db_spectrogram/20)
-
-def reconstruct_complex_spectrogram(magnitude, phase):
-    return magnitude * exp(1j*phase)
-     
-def inverse_fft(spectrogram):
-    inverse_fn = torchaudio.transforms.InverseSpectrogram(n_fft=512, hop_length=512//4, window_fn=hamming_window)
-    return inverse_fn(spectrogram)
-
-def transform_audio(audio, scaler):
-    spectrogram = audio_to_spectrogram(audio)
-    magnitude, phase = extract_magnitude_and_phase(spectrogram)
-    db_spectrogram, max_amplitude = amplitude_to_db(magnitude)
-    db_spectrogram = min_max_scaling(db_spectrogram, scaler)
-    return db_spectrogram.unsqueeze(0), phase, max_amplitude
-
-def spectrogram_to_audio(db_spectrogram, scaler, phase, max_amplitude):
-    db_spectrogram = db_spectrogram.squeeze(0)
-    db_spectrogram = inverse_min_max(db_spectrogram, scaler)
-    spectrogram = db_to_amplitude(db_spectrogram, max_amplitude)
-    complex_spec = reconstruct_complex_spectrogram(spectrogram, phase)
-    audio = inverse_fft(complex_spec)
-    return audio
-
-def save_audio(audio):
-    torchaudio.save(r"enhanced_audio.wav", audio, 16000)
-    return r"enhanced_audio.wav"
-
-def predict(user_input, model, scaler):
-    audio = load_audio(user_input)
-    spectrogram, phase, max_amplitude = transform_audio(audio, scaler)
-    
-    with no_grad():
-        enhanced_spectrogram = model.forward(spectrogram)
-    enhanced_audio = spectrogram_to_audio(enhanced_spectrogram, scaler, phase, max_amplitude)
-    enhanced_audio_path = save_audio(enhanced_audio)
+import torchaudio
+from torch import mean as _mean
+from torch import hamming_window, log10, no_grad, exp
+
+
+def return_input(user_input):
+    if user_input is None:
+        return None
+    return user_input
+
+
+def stereo_to_mono_convertion(waveform):
+    if waveform.shape[0] > 1:
+        waveform = _mean(waveform, dim=0, keepdims=True)
+        return waveform
+    else:
+        return waveform
+        
+def load_audio(audio_path):
+
+    audio_tensor, sr = torchaudio.load(audio_path)
+    audio_tensor = stereo_to_mono_convertion(audio_tensor)
+    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 16000)
+    return audio_tensor
+
+def load_audio_numpy(audio_path):
+    audio_tensor, sr = torchaudio.load(audio_path)
+    audio_tensor = torchaudio.functional.resample(audio_tensor, sr, 16000)
+    audio_array = audio_tensor.numpy()
+    return (16000, audio_array.ravel())
+
+def audio_to_spectrogram(audio):
+    transform_fn = torchaudio.transforms.Spectrogram(n_fft=512, hop_length=512//4, power=None, window_fn=hamming_window)
+    spectrogram = transform_fn(audio)
+    return spectrogram
+
+def extract_magnitude_and_phase(spectrogram):
+    magnitude, phase = spectrogram.abs(), spectrogram.angle()
+    return magnitude, phase
+
+def amplitude_to_db(magnitude_spec):
+    max_amplitude = magnitude_spec.max()
+    db_spectrogram = torchaudio.functional.amplitude_to_DB(magnitude_spec, 20, 10e-10, log10(max_amplitude), 100.0)
+    return db_spectrogram, max_amplitude
+
+def min_max_scaling(spectrogram, scaler):
+    # Min-Max scaling (soundness of the math is questionable due to the use of each spectrograms' max value during decibel-scaling)
+    spectrogram = scaler.transform(spectrogram)
+    return spectrogram
+
+def inverse_min_max(spectrogram, scaler):
+    spectrogram = scaler.inverse_transform(spectrogram)
+    return spectrogram
+
+def db_to_amplitude(db_spectrogram, max_amplitude):
+    return max_amplitude * 10**(db_spectrogram/20)
+
+def reconstruct_complex_spectrogram(magnitude, phase):
+    return magnitude * exp(1j*phase)
+     
+def inverse_fft(spectrogram):
+    inverse_fn = torchaudio.transforms.InverseSpectrogram(n_fft=512, hop_length=512//4, window_fn=hamming_window)
+    return inverse_fn(spectrogram)
+
+def transform_audio(audio, scaler):
+    spectrogram = audio_to_spectrogram(audio)
+    magnitude, phase = extract_magnitude_and_phase(spectrogram)
+    db_spectrogram, max_amplitude = amplitude_to_db(magnitude)
+    db_spectrogram = min_max_scaling(db_spectrogram, scaler)
+    return db_spectrogram.unsqueeze(0), phase, max_amplitude
+
+def spectrogram_to_audio(db_spectrogram, scaler, phase, max_amplitude):
+    db_spectrogram = db_spectrogram.squeeze(0)
+    db_spectrogram = inverse_min_max(db_spectrogram, scaler)
+    spectrogram = db_to_amplitude(db_spectrogram, max_amplitude)
+    complex_spec = reconstruct_complex_spectrogram(spectrogram, phase)
+    audio = inverse_fft(complex_spec)
+    return audio
+
+def save_audio(audio):
+    torchaudio.save(r"enhanced_audio.wav", audio, 16000)
+    return r"enhanced_audio.wav"
+
+def predict(user_input, model, scaler):
+    audio = load_audio(user_input)
+    spectrogram, phase, max_amplitude = transform_audio(audio, scaler)
+    
+    with no_grad():
+        enhanced_spectrogram = model.forward(spectrogram)
+    enhanced_audio = spectrogram_to_audio(enhanced_spectrogram, scaler, phase, max_amplitude)
+    enhanced_audio_path = save_audio(enhanced_audio)
     return enhanced_audio_path