Spaces:
Runtime error
Runtime error
| import tempfile | |
| from df.utils import resample | |
| 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 | |
| import math | |
| 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() | |
| print("start:", start, clean.shape) | |
| 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 | |
| print("mixture:", mixture.shape) | |
| assert torch.isfinite(mixture).all() | |
| max_m = mixture.abs().max() | |
| if max_m > 1: | |
| print(f"Clipping detected during mixing. Reducing gain by {1/max_m}") | |
| clean, noise, mixture = clean / max_m, noise / max_m, mixture / max_m | |
| return clean, noise, mixture | |
| def mix_and_denoise(speech_rec, speech_upl, noise, snr): | |
| sr = config("sr", 48000, int, section="df") | |
| print(speech_rec, speech_upl, noise, snr) | |
| if noise is None: | |
| noise = "samples/dkitchen.wav" | |
| sp_kwargs = {} | |
| if speech_rec is None or "none" in speech_rec: | |
| speech = "samples/p232_013_clean.wav" | |
| if speech_upl is not None and "none" not in speech_upl: | |
| print("using speech_upl") | |
| speech = speech_upl | |
| else: | |
| speech = speech_rec | |
| sp_kwargs = {"frame_offset": 4800} | |
| speech, meta = load_audio(speech, sr, **sp_kwargs) | |
| print(f"Loaded speech with shape {speech.shape}") | |
| noise, _ = load_audio(noise, sr) | |
| if meta.sample_rate != sr: | |
| # Low pass filter by resampling | |
| noise = resample(resample(noise, sr, meta.sample_rate), meta.sample_rate, sr) | |
| print(f"Loaded noise with shape {noise.shape}") | |
| speech, noise, noisy = mix_at_snr(speech, noise, snr) | |
| print("Start denoising audio") | |
| enhanced = enhance(model, df, noisy) | |
| print("Denoising finished") | |
| 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 | |
| if meta.sample_rate != sr: | |
| enhanced = resample(enhanced, sr, meta.sample_rate) | |
| noisy = resample(noisy, sr, meta.sample_rate) | |
| 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="inferno", | |
| ): | |
| """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 = [ | |
| [ | |
| "none", | |
| "samples/p232_013_clean.wav", | |
| "samples/dkitchen.wav", | |
| 10, | |
| ], | |
| [ | |
| "none", | |
| "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) | |