Spaces:
Sleeping
Sleeping
| 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 |