feat: add speech noise separation

experiments/bs-roformer-vocal-separation/speech_enhance.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import argparse
+import subprocess
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+from clearvoice import ClearVoice
+
+
+MODEL = "MossFormer2_SE_48K"
+SAMPLE_RATE = 48_000
+
+
+def convert_input(source: Path, target: Path) -> None:
+    completed = subprocess.run(
+        [
+            "ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
+            "-i", str(source), "-vn", "-ar", str(SAMPLE_RATE),
+            "-c:a", "pcm_f32le", str(target),
+        ],
+        capture_output=True,
+        text=True,
+        timeout=1800,
+    )
+    if completed.returncode:
+        raise RuntimeError(completed.stderr[-4000:] or "ffmpeg input conversion failed")
+
+
+def enhance(source: Path, speech_path: Path, noise_path: Path) -> None:
+    with tempfile.TemporaryDirectory(prefix="clearvoice-") as temp_dir:
+        converted = Path(temp_dir) / "input-48k.wav"
+        convert_input(source, converted)
+        audio, sample_rate = sf.read(converted, dtype="float32", always_2d=True)
+        if sample_rate != SAMPLE_RATE:
+            raise RuntimeError(f"unexpected sample rate: {sample_rate}")
+
+        model = ClearVoice(task="speech_enhancement", model_names=[MODEL])
+        # Use ClearVoice's file-I/O path so recordings longer than its 20-second
+        # one-pass window are segmented correctly. Run each channel separately
+        # because the enhancement network itself is mono, then restore stereo.
+        channels = []
+        for channel_index in range(audio.shape[1]):
+            channel_path = Path(temp_dir) / f"channel-{channel_index}.wav"
+            sf.write(channel_path, audio[:, channel_index], SAMPLE_RATE, subtype="FLOAT")
+            result = np.asarray(model(str(channel_path), False), dtype=np.float32).squeeze()
+            if result.ndim != 1:
+                raise RuntimeError(f"unexpected ClearVoice output shape: {result.shape}")
+            channels.append(result)
+        enhanced = np.column_stack(channels)
+
+        length = min(len(audio), len(enhanced))
+        original = audio[:length]
+        speech = enhanced[:length]
+        noise = original - speech
+
+        # FLAC does not support floating-point samples. PCM_24 retains ample
+        # headroom and avoids the invalid FLOAT/FLAC combination in libsndfile.
+        sf.write(speech_path, np.clip(speech, -1.0, 1.0), SAMPLE_RATE, format="FLAC", subtype="PCM_24")
+        sf.write(noise_path, np.clip(noise, -1.0, 1.0), SAMPLE_RATE, format="FLAC", subtype="PCM_24")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("input", nargs="?", type=Path)
+    parser.add_argument("speech", nargs="?", type=Path)
+    parser.add_argument("noise", nargs="?", type=Path)
+    parser.add_argument("--download-only", action="store_true")
+    args = parser.parse_args()
+    if args.download_only:
+        ClearVoice(task="speech_enhancement", model_names=[MODEL])
+        return
+    if not all((args.input, args.speech, args.noise)):
+        parser.error("input, speech and noise output paths are required")
+    enhance(args.input, args.speech, args.noise)
+
+
+if __name__ == "__main__":
+    main()