import math import tempfile import gradio import gradio.inputs import gradio.outputs import matplotlib.pyplot as plt import markdown import numpy as np import torch from df import config from df.enhance import enhance, init_df, load_audio, save_audio device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model, df, _ = init_df() model = model.to(device=device).eval() def mix_at_snr(clean, noise, snr, eps=1e-10): """Mix clean and noise signal at a given SNR. Args: clean: 1D Tensor with the clean signal to mix. noise: 1D Tensor of shape. snr: Signal to noise ratio. Returns: clean: 1D Tensor with gain changed according to the snr. noise: 1D Tensor with the combined noise channels. mix: 1D Tensor with added clean and noise signals. """ clean = torch.as_tensor(clean).mean(0, keepdim=True) noise = torch.as_tensor(noise).mean(0, keepdim=True) if noise.shape[1] < clean.shape[1]: noise = noise.repeat((1, int(math.ceil(clean.shape[1] / noise.shape[1])))) max_start = int(noise.shape[1] - clean.shape[1]) start = torch.randint(0, max_start, ()).item() noise = noise[:, start : start + clean.shape[1]] E_speech = torch.mean(clean.pow(2)) + eps E_noise = torch.mean(noise.pow(2)) K = torch.sqrt((E_noise / E_speech) * 10 ** (snr / 10) + eps) noise = noise / K mixture = clean + noise assert torch.isfinite(mixture).all() return clean, noise, mixture def mix_and_denoise(speech, speech_alt, noise, snr): print(speech, noise, snr) if noise is None: noise = "samples/dkitchen.wav" if speech is None or speech == "": speech = "samples/p232_013_clean.wav" if speech_alt is not None: speech = speech_alt print(speech, noise, snr) sr = config("sr", 48000, int, section="df") speech, _ = load_audio(speech, sr) noise, _ = load_audio(noise, sr) speech, noise, noisy = mix_at_snr(speech, noise, snr) enhanced = enhance(model, df, noisy) lim = torch.linspace(0.0, 1.0, int(sr * 0.15)).unsqueeze(0) lim = torch.cat((lim, torch.ones(1, enhanced.shape[1] - lim.shape[1])), dim=1) enhanced = enhanced * lim noisy_fn = tempfile.NamedTemporaryFile(suffix="noisy.wav", delete=False).name save_audio(noisy_fn, noisy, sr) enhanced_fn = tempfile.NamedTemporaryFile(suffix="enhanced.wav", delete=False).name save_audio(enhanced_fn, enhanced, sr) print("saved audios", noisy_fn, enhanced_fn) return ( noisy_fn, spec_figure(noisy, sr=sr), enhanced_fn, spec_figure(enhanced, sr=sr), ) def specshow( spec, ax=None, title=None, xlabel=None, ylabel=None, sr=48000, n_fft=None, hop=None, t=None, f=None, vmin=-100, vmax=0, xlim=None, ylim=None, cmap="viridis", ): """Plots a spectrogram of shape [F, T]""" spec_np = spec.cpu().numpy() if isinstance(spec, torch.Tensor) else spec if ax is not None: set_title = ax.set_title set_xlabel = ax.set_xlabel set_ylabel = ax.set_ylabel set_xlim = ax.set_xlim set_ylim = ax.set_ylim else: ax = plt set_title = plt.title set_xlabel = plt.xlabel set_ylabel = plt.ylabel set_xlim = plt.xlim set_ylim = plt.ylim if n_fft is None: if spec.shape[0] % 2 == 0: n_fft = spec.shape[0] * 2 else: n_fft = (spec.shape[0] - 1) * 2 hop = hop or n_fft // 4 if t is None: t = np.arange(0, spec_np.shape[-1]) * hop / sr if f is None: f = np.arange(0, spec_np.shape[0]) * sr // 2 / (n_fft // 2) / 1000 im = ax.pcolormesh( t, f, spec_np, rasterized=True, shading="auto", vmin=vmin, vmax=vmax, cmap=cmap ) if title is not None: set_title(title) if xlabel is not None: set_xlabel(xlabel) if ylabel is not None: set_ylabel(ylabel) if xlim is not None: set_xlim(xlim) if ylim is not None: set_ylim(ylim) return im def spec_figure( audio: torch.Tensor, figsize=(15, 5), colorbar=False, colorbar_format=None, figure=None, return_im=False, labels=True, **kwargs, ) -> plt.Figure: audio = torch.as_tensor(audio) if labels: kwargs.setdefault("xlabel", "Time [s]") kwargs.setdefault("ylabel", "Frequency [Hz]") n_fft = kwargs.setdefault("n_fft", 1024) hop = kwargs.setdefault("hop", 512) w = torch.hann_window(n_fft, device=audio.device) spec = torch.stft(audio, n_fft, hop, window=w, return_complex=False) spec = spec.div_(w.pow(2).sum()) spec = torch.view_as_complex(spec).abs().clamp_min(1e-12).log10().mul(10) kwargs.setdefault("vmax", max(0.0, spec.max().item())) if figure is None: figure = plt.figure(figsize=figsize) figure.set_tight_layout(True) if spec.dim() > 2: spec = spec.squeeze(0) im = specshow(spec, **kwargs) if colorbar: ckwargs = {} if "ax" in kwargs: if colorbar_format is None: if ( kwargs.get("vmin", None) is not None or kwargs.get("vmax", None) is not None ): colorbar_format = "%+2.0f dB" ckwargs = {"ax": kwargs["ax"]} plt.colorbar(im, format=colorbar_format, **ckwargs) if return_im: return im return figure inputs = [ gradio.inputs.Audio( source="microphone", type="filepath", optional=True, label="Record your own voice", ), gradio.inputs.Audio( source="upload", type="filepath", optional=True, label="Alternative: Upload speech sample", ), gradio.inputs.Audio( source="upload", type="filepath", optional=True, label="Upload noise sample" ), gradio.inputs.Slider(minimum=-10, maximum=40, step=5, default=10), # SNR ] examples = [ [ "samples/p232_013_clean.wav", "samples/p232_013_clean.wav", "samples/dkitchen.wav", 10, ], [ "samples/p232_013_clean.wav", "samples/p232_019_clean.wav", "samples/dliving.wav", 10, ], ] outputs = [ gradio.outputs.Audio(label="Noisy"), gradio.outputs.Image(type="plot"), gradio.outputs.Audio(label="Enhanced"), gradio.outputs.Image(type="plot"), ] description = ( "This demo denoises audio files using DeepFilterNet. Try it with your own voice!" ) iface = gradio.Interface( fn=mix_and_denoise, title="DeepFilterNet Demo", inputs=inputs, outputs=outputs, examples=examples, description=description, layout="horizontal", allow_flagging="never", article=markdown.markdown(open("usage.md").read()), ) iface.launch(cache_examples=False)