Upload 2 files

inference_dnr_onnx.py

@@ -0,0 +1,194 @@
+#!/usr/bin/env python3
+"""
+US-DNR-003: Pure onnxruntime inference for TIGER-DnR (Dialog/Effect/Music separation).
+
+Uses only onnxruntime + audio I/O, no look2hear import.
+STFT/ISTFT performed in Python, separator network runs in ONNX.
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+import torchaudio
+import numpy as np
+import onnxruntime as ort
+
+
+def load_audio(audio_path, target_sr=44100):
+    """Load and preprocess audio to 44.1kHz."""
+    waveform, sr = torchaudio.load(audio_path)
+
+    # Resample if needed
+    if sr != target_sr:
+        resampler = torchaudio.transforms.Resample(sr, target_sr)
+        waveform = resampler(waveform)
+
+    # Convert to mono if stereo
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+
+    return waveform, target_sr
+
+
+def save_audio(audio_tensor, output_path, sample_rate=44100):
+    """Save audio tensor to file."""
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    torchaudio.save(output_path, audio_tensor, sample_rate)
+
+
+def onnx_separate(onnx_path, audio_tensor, win=2048, stride=512):
+    """
+    Separate audio using ONNX model.
+
+    Args:
+        onnx_path: Path to ONNX separator model
+        audio_tensor: [C, T] audio tensor
+        win: STFT window size
+        stride: STFT hop length
+
+    Returns:
+        Tuple of (dialog, effect, music) tensors, each [C, T]
+    """
+    # Create ONNX session
+    sess_options = ort.SessionOptions()
+    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+
+    # Try CUDA first, fallback to CPU
+    providers = []
+    if 'CUDAExecutionProvider' in ort.get_available_providers():
+        providers.append('CUDAExecutionProvider')
+        print(f"[inference_dnr_onnx] Using CUDAExecutionProvider")
+    else:
+        providers.append('CPUExecutionProvider')
+        print(f"[inference_dnr_onnx] Using CPUExecutionProvider")
+
+    session = ort.InferenceSession(onnx_path, sess_options, providers=providers)
+
+    # Ensure [C, T] shape
+    if audio_tensor.ndim == 1:
+        audio_tensor = audio_tensor.unsqueeze(0)
+
+    nch = audio_tensor.shape[0]
+    original_length = audio_tensor.shape[-1]
+    audio_flat = audio_tensor.view(-1)  # Flatten to [nch*T]
+
+    # Compute STFT
+    print(f"[inference_dnr_onnx] Computing STFT...")
+    window = torch.hann_window(win).type(audio_flat.dtype)
+    spec = torch.stft(
+        audio_flat,
+        n_fft=win,
+        hop_length=stride,
+        window=window,
+        return_complex=True
+    )  # [F, T_frames]
+
+    # Extract real and imaginary parts
+    spec_real = spec.real.unsqueeze(0).numpy()  # [1, F, T_frames]
+    spec_imag = spec.imag.unsqueeze(0).numpy()  # [1, F, T_frames]
+    print(f"[inference_dnr_onnx] STFT shape: {spec_real.shape}")
+
+    # Run ONNX inference
+    print(f"[inference_dnr_onnx] Running ONNX separator...")
+    outputs = session.run(
+        None,
+        {
+            'spec_real': spec_real.astype(np.float32),
+            'spec_imag': spec_imag.astype(np.float32)
+        }
+    )
+
+    # outputs: [dialog_real, dialog_imag, effect_real, effect_imag, music_real, music_imag]
+    dialog_real, dialog_imag, effect_real, effect_imag, music_real, music_imag = outputs
+
+    # Convert back to complex spectrograms
+    dialog_spec = torch.complex(
+        torch.from_numpy(dialog_real).squeeze(0),
+        torch.from_numpy(dialog_imag).squeeze(0)
+    )
+    effect_spec = torch.complex(
+        torch.from_numpy(effect_real).squeeze(0),
+        torch.from_numpy(effect_imag).squeeze(0)
+    )
+    music_spec = torch.complex(
+        torch.from_numpy(music_real).squeeze(0),
+        torch.from_numpy(music_imag).squeeze(0)
+    )
+
+    # ISTFT to get time-domain signals
+    print(f"[inference_dnr_onnx] Computing ISTFT...")
+    dialog = torch.istft(
+        dialog_spec,
+        n_fft=win,
+        hop_length=stride,
+        window=window,
+        length=original_length
+    )
+    effect = torch.istft(
+        effect_spec,
+        n_fft=win,
+        hop_length=stride,
+        window=window,
+        length=original_length
+    )
+    music = torch.istft(
+        music_spec,
+        n_fft=win,
+        hop_length=stride,
+        window=window,
+        length=original_length
+    )
+
+    # Reshape to [C, T]
+    dialog = dialog.view(nch, -1)
+    effect = effect.view(nch, -1)
+    music = music.view(nch, -1)
+
+    return dialog, effect, music
+
+
+def main():
+    parser = argparse.ArgumentParser(description="TIGER-DnR ONNX inference (no look2hear)")
+    parser.add_argument("--audio_path", default="test/test_mixture_466.wav", help="Input audio file")
+    parser.add_argument("--output_dir", default="separated_audio_dnr_onnx", help="Output directory")
+    parser.add_argument("--onnx_path", default="onnx/tiger_dnr_separator.onnx", help="ONNX model path")
+    args = parser.parse_args()
+
+    print(f"[inference_dnr_onnx] TIGER-DnR ONNX Inference")
+    print(f"[inference_dnr_onnx] Input: {args.audio_path}")
+    print(f"[inference_dnr_onnx] Output: {args.output_dir}")
+    print(f"[inference_dnr_onnx] Model: {args.onnx_path}")
+
+    # Check inputs
+    if not os.path.exists(args.audio_path):
+        print(f"[inference_dnr_onnx] ERROR: Audio file not found: {args.audio_path}")
+        sys.exit(1)
+
+    if not os.path.exists(args.onnx_path):
+        print(f"[inference_dnr_onnx] ERROR: ONNX model not found: {args.onnx_path}")
+        sys.exit(1)
+
+    # Load audio
+    print(f"[inference_dnr_onnx] Loading audio...")
+    audio, sr = load_audio(args.audio_path)
+    print(f"[inference_dnr_onnx] Audio shape: {audio.shape}, sample rate: {sr}")
+
+    # Separate
+    dialog, effect, music = onnx_separate(args.onnx_path, audio)
+
+    # Save outputs
+    print(f"[inference_dnr_onnx] Saving separated audio...")
+    save_audio(dialog, os.path.join(args.output_dir, "dialog.wav"), sr)
+    save_audio(effect, os.path.join(args.output_dir, "effect.wav"), sr)
+    save_audio(music, os.path.join(args.output_dir, "music.wav"), sr)
+
+    print(f"[inference_dnr_onnx] Saved dialog.wav")
+    print(f"[inference_dnr_onnx] Saved effect.wav")
+    print(f"[inference_dnr_onnx] Saved music.wav")
+    print(f"[inference_dnr_onnx] SUCCESS")
+
+
+if __name__ == "__main__":
+    main()

inference_onnx.py

@@ -0,0 +1,99 @@
+import argparse
+import os
+import sys
+
+import numpy as np
+import onnxruntime as ort
+import torch
+import torchaudio
+import torchaudio.transforms as T
+
+
+TARGET_SR = 16000
+CHUNK_LEN = TARGET_SR * 4  # must match dummy length in export_onnx.py
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Pure onnxruntime TIGER-speech inference.")
+    p.add_argument("--audio_path", default="test/mix.wav",
+                   help="Path to mixture wav.")
+    p.add_argument("--output_dir", default="separated_audio_onnx",
+                   help="Directory to save separated spkN.wav files.")
+    p.add_argument("--onnx_path", default="onnx/tiger_speech.onnx",
+                   help="Exported ONNX model (from export_onnx.py).")
+    return p.parse_args()
+
+
+def load_audio(audio_path):
+    waveform, original_sr = torchaudio.load(audio_path)
+    print(f"Loaded {audio_path}: sr={original_sr}, shape={tuple(waveform.shape)}")
+    if original_sr != TARGET_SR:
+        print(f"Resampling {original_sr} Hz -> {TARGET_SR} Hz")
+        waveform = T.Resample(orig_freq=original_sr, new_freq=TARGET_SR)(waveform)
+    if waveform.dim() == 1:
+        waveform = waveform.unsqueeze(0)
+    if waveform.shape[0] > 1:
+        print(f"Downmixing {waveform.shape[0]} channels -> mono")
+        waveform = waveform.mean(dim=0, keepdim=True)
+    return waveform  # [1, T]
+
+
+def build_session(onnx_path):
+    available = ort.get_available_providers()
+    if "CUDAExecutionProvider" in available:
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    else:
+        providers = ["CPUExecutionProvider"]
+    sess = ort.InferenceSession(onnx_path, providers=providers)
+    chosen = sess.get_providers()[0]
+    print(f"onnxruntime provider: {chosen}")
+    return sess
+
+
+def run_chunks(sess, mono_wave):
+    in_name = sess.get_inputs()[0].name
+    out_name = sess.get_outputs()[0].name
+    total = mono_wave.shape[-1]
+    outputs = []
+    for start in range(0, total, CHUNK_LEN):
+        end = min(start + CHUNK_LEN, total)
+        chunk = mono_wave[:, start:end]
+        pad = CHUNK_LEN - chunk.shape[-1]
+        if pad > 0:
+            chunk = torch.nn.functional.pad(chunk, (0, pad))
+        x = chunk.unsqueeze(0).contiguous().numpy().astype(np.float32)  # [1,1,CHUNK_LEN]
+        y = sess.run([out_name], {in_name: x})[0]  # [1, num_spk, CHUNK_LEN]
+        if pad > 0:
+            y = y[..., : CHUNK_LEN - pad]
+        outputs.append(y[0])  # [num_spk, chunk_len]
+    return np.concatenate(outputs, axis=-1)  # [num_spk, total]
+
+
+def main():
+    args = parse_args()
+
+    if not os.path.isfile(args.audio_path):
+        print(f"ERROR: audio not found: {args.audio_path}")
+        sys.exit(1)
+    if not os.path.isfile(args.onnx_path):
+        print(f"ERROR: onnx not found: {args.onnx_path}")
+        sys.exit(1)
+
+    waveform = load_audio(args.audio_path)  # [1, T]
+    print(f"Preprocessed shape: {tuple(waveform.shape)} (16 kHz mono)")
+
+    sess = build_session(args.onnx_path)
+    estimates = run_chunks(sess, waveform)  # [num_spk, T]
+    num_spk = estimates.shape[0]
+    print(f"Separation complete: num_spk={num_spk}, samples={estimates.shape[-1]}")
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    for i in range(num_spk):
+        out_path = os.path.join(args.output_dir, f"spk{i+1}.wav")
+        track = torch.from_numpy(estimates[i]).unsqueeze(0)
+        torchaudio.save(out_path, track, TARGET_SR)
+        print(f"Saved spk{i+1} -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()