82 lines
3.1 KiB
Python
82 lines
3.1 KiB
Python
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()
|