Update app.py

app.py

@@ -6,7 +6,6 @@ import time
 import gradio as gr
 from typing import List, Optional, Tuple, Union
 
-import gradio as gr
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -42,19 +41,6 @@ NOISES = {
 
 
 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]:
@@ -77,21 +63,17 @@ def mix_at_snr(clean, noise, snr, eps=1e-10):
     return clean, noise, mixture
 
 
-def load_audio_gradio(
-    audio_or_file: Union[None, str, Tuple[int, np.ndarray]], sr: int
-) -> Optional[Tuple[Tensor, AudioMetaData]]:
+def load_audio_gradio(audio_or_file: Union[None, str, Tuple[int, np.ndarray]], sr: int) -> Optional[Tuple[Tensor, AudioMetaData]]:
     if audio_or_file is None:
         return None
     if isinstance(audio_or_file, str):
         if audio_or_file.lower() == "none":
             return None
-        # First try default format
         audio, meta = load_audio(audio_or_file, sr)
     else:
         meta = AudioMetaData(-1, -1, -1, -1, "")
         assert isinstance(audio_or_file, (tuple, list))
         meta.sample_rate, audio_np = audio_or_file
-        # Gradio documentation says, the shape is [samples, 2], but apparently sometimes its not.
         audio_np = audio_np.reshape(audio_np.shape[0], -1).T
         if audio_np.dtype == np.int16:
             audio_np = (audio_np / (1 << 15)).astype(np.float32)
@@ -109,7 +91,7 @@ def demo_fn(speech_upl: str, noise_type: str, snr: int, mic_input: Optional[str]
     snr = int(snr)
     noise_fn = NOISES[noise_type]
     meta = AudioMetaData(-1, -1, -1, -1, "")
-    max_s = 10  # limit to 10 seconds
+    max_s = 10
     if speech_upl is not None:
         sample, meta = load_audio(speech_upl, sr)
         max_len = max_s * sr
@@ -120,13 +102,11 @@ def demo_fn(speech_upl: str, noise_type: str, snr: int, mic_input: Optional[str]
         sample, meta = load_audio("samples/p232_013_clean.wav", sr)
         sample = sample[..., : max_s * sr]
     if sample.dim() > 1 and sample.shape[0] > 1:
-        assert (
-            sample.shape[1] > sample.shape[0]
-        ), f"Expecting channels first, but got {sample.shape}"
+        assert sample.shape[1] > sample.shape[0], f"Expecting channels first, but got {sample.shape}"
         sample = sample.mean(dim=0, keepdim=True)
     logger.info(f"Loaded sample with shape {sample.shape}")
     if noise_fn is not None:
-        noise, _ = load_audio(noise_fn, sr)  # type: ignore
+        noise, _ = load_audio(noise_fn, sr)
         logger.info(f"Loaded noise with shape {noise.shape}")
         _, _, sample = mix_at_snr(sample, noise, snr)
     logger.info("Start denoising audio")
@@ -155,24 +135,47 @@ def demo_fn(speech_upl: str, noise_type: str, snr: int, mic_input: Optional[str]
     return noisy_wav, noisy_im, enhanced_wav, enh_im
 
 
-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]"""
+def denoise_api(audio_file_path: str, noise_type: str = "None", snr: int = 10):
+    sr = config("sr", 48000, int, section="df")
+    sample, meta = load_audio(audio_file_path, sr)
+    noise = None if noise_type == "None" else load_audio(NOISES[noise_type], sr)[0]
+    _, _, sample_mix = mix_at_snr(sample, noise, snr) if noise is not None else (sample, None, sample)
+    enhanced = enhance(model, df, sample_mix)
+    noisy_wav = tempfile.NamedTemporaryFile(suffix="noisy.wav", delete=False).name
+    enhanced_wav = tempfile.NamedTemporaryFile(suffix="enhanced.wav", delete=False).name
+    save_audio(noisy_wav, sample_mix, sr)
+    save_audio(enhanced_wav, enhanced, sr)
+    return {
+        "enhanced_audio": enhanced_wav,
+        "snr": snr,
+        "status": "done"
+    }
+
+
+def toggle(choice):
+    if choice == "mic":
+        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
+    else:
+        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
+
+
+def cleanup_tmp(filter: List[str] = [], hours_keep=2):
+    filter.append("p232")
+    logger.info(f"Filter: {filter}")
+    if os.path.exists("/tmp"):
+        for f in glob.glob("/tmp/*"):
+            print(f"Got file {f}")
+            is_old = (time.time() - os.path.getmtime(f)) / 3600 > hours_keep
+            filtered = any(filt in f for filt in filter if filt is not None)
+            if is_old and not filtered:
+                try:
+                    os.remove(f)
+                    logger.info(f"Removed file {f}")
+                except Exception as e:
+                    logger.warning(f"failed to remove file {f}: {e}")
+
+
+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"):
     spec_np = spec.cpu().numpy() if isinstance(spec, torch.Tensor) else spec
     if ax is not None:
         set_title = ax.set_title
@@ -188,18 +191,11 @@ def specshow(
         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
+        n_fft = spec.shape[0] * 2 if spec.shape[0] % 2 == 0 else (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
-    )
+    t = np.arange(0, spec_np.shape[-1]) * hop / sr if t is None else t
+    f = np.arange(0, spec_np.shape[0]) * sr // 2 / (n_fft // 2) / 1000 if f is None else f
+    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:
@@ -213,15 +209,7 @@ def specshow(
     return im
 
 
-def spec_im(
-    audio: torch.Tensor,
-    figsize=(15, 5),
-    colorbar=False,
-    colorbar_format=None,
-    figure=None,
-    labels=True,
-    **kwargs,
-) -> Image:
+def spec_im(audio: torch.Tensor, figsize=(15, 5), colorbar=False, colorbar_format=None, figure=None, labels=True, **kwargs) -> Image:
     audio = torch.as_tensor(audio)
     if labels:
         kwargs.setdefault("xlabel", "Time [s]")
@@ -233,7 +221,6 @@ def spec_im(
     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)
@@ -252,85 +239,41 @@ def spec_im(
     return Image.frombytes("RGB", figure.canvas.get_width_height(), figure.canvas.tostring_rgb())
 
 
-def cleanup_tmp(filter: List[str] = [], hours_keep=2):
-    filter.append("p232")
-    logger.info(f"Filter: {filter}")
-    # Cleanup some old wav files
-    if os.path.exists("/tmp"):
-        for f in glob.glob("/tmp/*"):
-            print(f"Got file {f}")
-            is_old = (time.time() - os.path.getmtime(f)) / 3600 > hours_keep
-            filtered = any(filt in f for filt in filter if filt is not None)
-            if is_old and not filtered:
-                try:
-                    os.remove(f)
-                    logger.info(f"Removed file {f}")
-                except Exception as e:
-                    logger.warning(f"failed to remove file {f}: {e}")
-
-
-def toggle(choice):
-    if choice == "mic":
-        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
-    else:
-        return gr.update(visible=False, value=None), gr.update(visible=True, value=None)
-
-
 with gr.Blocks() as demo:
     with gr.Row():
-        gr.Markdown(
-            """
+        gr.Markdown("""
             ## DeepFilterNet2 Demo\
-
             This demo denoises audio files using DeepFilterNet. Try it with your own voice!
-            """
-        )
+        """)
     with gr.Row():
         with gr.Column():
-            radio = gr.Radio(
-                ["mic", "file"], value="file", label="How would you like to upload your audio?"
-            )
+            radio = gr.Radio(["mic", "file"], value="file", label="How would you like to upload your audio?")
             mic_input = gr.Mic(label="Input", type="filepath", visible=False)
             audio_file = gr.Audio(type="filepath", label="Input", visible=True)
             inputs = [
                 audio_file,
-                gr.Dropdown(
-                    label="Add background noise",
-                    choices=list(NOISES.keys()),
-                    value="None",
-                ),
-                gr.Dropdown(
-                    label="Noise Level (SNR)",
-                    choices=["-5", "0", "10", "20"],
-                    value="10",
-                ),
+                gr.Dropdown(label="Add background noise", choices=list(NOISES.keys()), value="None"),
+                gr.Dropdown(label="Noise Level (SNR)", choices=["-5", "0", "10", "20"], value="10"),
                 mic_input,
             ]
             btn = gr.Button("Generate")
         with gr.Column():
             outputs = [
-                # gr.Video(type="filepath", label="Noisy audio"),
                 gr.Audio(type="filepath", label="Noisy audio"),
                 gr.Image(label="Noisy spectrogram"),
-                # gr.Video(type="filepath", label="Enhanced audio"),
                 gr.Audio(type="filepath", label="Enhanced audio"),
                 gr.Image(label="Enhanced spectrogram"),
             ]
     btn.click(fn=demo_fn, inputs=inputs, outputs=outputs, api_name='denoise')
     radio.change(toggle, radio, [mic_input, audio_file])
-    gr.Examples(
-        [
-            ["./samples/p232_013_clean.wav", "Kitchen", "10"],
-            ["./samples/p232_013_clean.wav", "Cafe", "10"],
-            ["./samples/p232_019_clean.wav", "Cafe", "10"],
-            ["./samples/p232_019_clean.wav", "River", "10"],
-        ],
-        fn=demo_fn,
-        inputs=inputs,
-        outputs=outputs,
-        cache_examples=True,
-    ),
+    gr.Examples([
+        ["./samples/p232_013_clean.wav", "Kitchen", "10"],
+        ["./samples/p232_013_clean.wav", "Cafe", "10"],
+        ["./samples/p232_019_clean.wav", "Cafe", "10"],
+        ["./samples/p232_019_clean.wav", "River", "10"],
+    ], fn=demo_fn, inputs=inputs, outputs=outputs, cache_examples=True)
     gr.Markdown(open("usage.md").read())
 
 cleanup_tmp()
 demo.launch()
+