218 lines
8.2 KiB
Python
218 lines
8.2 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Audio cross-correlation for movie_sync — torch GPU auto (ROCm/CUDA) with CPU fallback.
|
|
|
|
Content-to-content sync (Option A): ignores container start_time/edit-list mux delay.
|
|
If audio is delayed in container, we pad with silence on mux instead of shifting presentation.
|
|
Extracted wavs are pcm_f32le 48k mono (extract.py:41), so correlation is phase-accurate.
|
|
|
|
Usage:
|
|
from sync_audio import load_wav, correlate, find_lag
|
|
a, sr = load_wav("extracted/Drachenzaehmen_audio.wav")
|
|
lag_samples, peak, corr = correlate(a, b, sr=48000, method="phat", device="auto")
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
import os
|
|
import sys
|
|
import math
|
|
from typing import Tuple, Literal
|
|
|
|
import numpy as np
|
|
|
|
try:
|
|
import soundfile as sf # preferred for f32 wav
|
|
HAS_SF = True
|
|
except ImportError:
|
|
HAS_SF = False
|
|
|
|
try:
|
|
import torch
|
|
HAS_TORCH = True
|
|
except ImportError:
|
|
HAS_TORCH = False
|
|
|
|
|
|
def _get_device(pref: str = "auto") -> str:
|
|
if pref == "cpu":
|
|
return "cpu"
|
|
if pref in ("cuda", "auto") and HAS_TORCH and torch.cuda.is_available():
|
|
return "cuda"
|
|
return "cpu"
|
|
|
|
|
|
def load_wav(path: str, target_sr: int = 48000, mono: bool = True) -> Tuple[np.ndarray, int]:
|
|
"""Load wav as float32 mono. Resample if needed (cheap linear warn)."""
|
|
if not os.path.exists(path):
|
|
raise FileNotFoundError(path)
|
|
if HAS_SF:
|
|
data, sr = sf.read(path, always_2d=False, dtype="float32")
|
|
else:
|
|
# fallback scipy
|
|
from scipy.io import wavfile
|
|
sr, data = wavfile.read(path)
|
|
if data.dtype == np.int16:
|
|
data = data.astype(np.float32) / 32768.0
|
|
elif data.dtype == np.int32:
|
|
data = data.astype(np.float32) / 2147483648.0
|
|
elif data.dtype != np.float32:
|
|
data = data.astype(np.float32)
|
|
if data.ndim > 1:
|
|
if mono:
|
|
data = data.mean(axis=1).astype(np.float32)
|
|
if sr != target_sr:
|
|
# simple resample via scipy if available, else warn
|
|
try:
|
|
from scipy.signal import resample_poly
|
|
import math as _m
|
|
g = _m.gcd(sr, target_sr)
|
|
data = resample_poly(data, target_sr // g, sr // g).astype(np.float32)
|
|
sr = target_sr
|
|
except Exception:
|
|
print(f"WARNING: {path} sr={sr} != {target_sr} and resample failed, using as-is", file=sys.stderr)
|
|
return data.astype(np.float32), sr
|
|
|
|
|
|
def _next_pow2(n: int) -> int:
|
|
return 1 << (n - 1).bit_length()
|
|
|
|
|
|
def _normalize(a: np.ndarray) -> np.ndarray:
|
|
a = a - a.mean()
|
|
s = a.std()
|
|
return a / s if s > 1e-9 else a
|
|
|
|
|
|
def correlate(
|
|
a: np.ndarray,
|
|
b: np.ndarray,
|
|
sr: int = 48000,
|
|
method: Literal["phat", "standard"] = "phat",
|
|
device: str = "auto",
|
|
normalize: bool = True,
|
|
) -> Tuple[int, float, np.ndarray]:
|
|
"""
|
|
Return (lag_samples, peak_value, corr_array).
|
|
lag >0 means b delayed vs a (b starts later), need to shift b earlier by lag.
|
|
Convention: corr = ifft(FFT(a) * conj(FFT(b))). peak at lag = argmax.
|
|
GPU via torch if available, else numpy.
|
|
"""
|
|
if normalize:
|
|
a = _normalize(a)
|
|
b = _normalize(b)
|
|
|
|
n = len(a) + len(b) - 1
|
|
N = _next_pow2(n)
|
|
|
|
dev = _get_device(device)
|
|
|
|
if dev == "cuda" and HAS_TORCH:
|
|
# torch path — ROCm and CUDA both via cuda device
|
|
try:
|
|
ta = torch.from_numpy(a).to("cuda", dtype=torch.float32)
|
|
tb = torch.from_numpy(b).to("cuda", dtype=torch.float32)
|
|
Fa = torch.fft.rfft(ta, n=N)
|
|
Fb = torch.fft.rfft(tb, n=N)
|
|
R = Fa * torch.conj(Fb)
|
|
if method == "phat":
|
|
R = R / (torch.abs(R) + 1e-10)
|
|
corr_t = torch.fft.irfft(R, n=N)
|
|
# fftshift to center 0 lag
|
|
corr_t = torch.roll(corr_t, shifts=N // 2)
|
|
corr = corr_t.cpu().numpy()
|
|
# keep only valid length n centered
|
|
start = N // 2 - (len(b) - 1)
|
|
# Actually after roll, we want slice; simpler: use raw without roll and map lag later
|
|
# Recompute without roll for correct indexing, keep roll for display but fix lag calc
|
|
# For lag calc we need not shift: peak at N - lag etc. So we use shifted version + offset mapping
|
|
# To avoid confusion, return shifted corr and map lag accordingly
|
|
peak_idx = int(np.argmax(np.abs(corr)))
|
|
# lag = peak_idx - (len(b)-1) - (N//2 - (len(b)-1))? Let's instead use unshifted for lag and shifted for display
|
|
# Fallback: compute lag from unshifted
|
|
corr_raw = torch.fft.irfft(R, n=N).cpu().numpy()
|
|
peak_raw = int(np.argmax(np.abs(corr_raw)))
|
|
# raw layout: lag 0 at 0, positive lag (b delayed) wraps to end
|
|
if peak_raw > N // 2:
|
|
lag = peak_raw - N
|
|
else:
|
|
lag = peak_raw
|
|
peak_val = float(np.abs(corr_raw[peak_raw]))
|
|
# return shifted display corr + lag
|
|
return lag, peak_val, corr
|
|
except Exception as e:
|
|
print(f"torch cuda failed ({e}), falling back cpu", file=sys.stderr)
|
|
dev = "cpu"
|
|
|
|
# numpy/scipy cpu path
|
|
Fa = np.fft.rfft(a, n=N)
|
|
Fb = np.fft.rfft(b, n=N)
|
|
R = Fa * np.conj(Fb)
|
|
if method == "phat":
|
|
R = R / (np.abs(R) + 1e-10)
|
|
corr_raw = np.fft.irfft(R, n=N)
|
|
peak_raw = int(np.argmax(np.abs(corr_raw)))
|
|
if peak_raw > N // 2:
|
|
lag = peak_raw - N
|
|
else:
|
|
lag = peak_raw
|
|
peak_val = float(np.abs(corr_raw[peak_raw]))
|
|
# shifted for display centered at 0
|
|
corr_shifted = np.fft.fftshift(corr_raw)
|
|
# but we return raw shifted centered? Provide shifted for plotting
|
|
# For consistency with torch branch we return shifted display array
|
|
return lag, peak_val, corr_shifted
|
|
|
|
|
|
def lag_to_seconds(lag_samples: int, sr: int = 48000) -> float:
|
|
return lag_samples / float(sr)
|
|
|
|
|
|
def apply_lag_with_silence(wav: np.ndarray, lag_samples: int) -> np.ndarray:
|
|
"""
|
|
Apply lag by padding silence at start (no sound for delay).
|
|
lag >0: b delayed vs a => pad b at start with lag silence to align to a, or trim a.
|
|
For mux: if lag>0 (b starts later), we need to delay a by lag or提前 b.
|
|
This helper pads with zeros (silence) at start; negative lag pads the other direction.
|
|
Caller decides which file to pad.
|
|
"""
|
|
if lag_samples == 0:
|
|
return wav
|
|
if lag_samples > 0:
|
|
# b delayed => pad b front with lag silence to align to a
|
|
return np.concatenate([np.zeros(lag_samples, dtype=np.float32), wav])
|
|
else:
|
|
# b earlier => trim or pad a; this pads a-side: when applying to b, negative lag means b earlier, so we trim? Provide padded version for mux: pad with silence then trim alignment
|
|
# For generic helper: pad with silence at start if lag negative we pad other file, but this is applied to lagging file only
|
|
return np.concatenate([np.zeros(-lag_samples, dtype=np.float32), wav])
|
|
|
|
|
|
# torch device info for Makefile / GUI
|
|
def device_info() -> str:
|
|
if not HAS_TORCH:
|
|
return "torch not installed (cpu numpy)"
|
|
if torch.cuda.is_available():
|
|
try:
|
|
name = torch.cuda.get_device_name(0)
|
|
except Exception:
|
|
name = "cuda"
|
|
return f"torch {torch.__version__} + {name} (cuda/ROCm)"
|
|
return f"torch {torch.__version__} (cpu only)"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse, pathlib
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("a", help="wav A")
|
|
ap.add_argument("b", help="wav B")
|
|
ap.add_argument("--phat", action="store_true", help="use GCC-PHAT")
|
|
ap.add_argument("--device", default="auto", choices=["auto","cpu","cuda"])
|
|
ap.add_argument("--sr", type=int, default=48000)
|
|
args = ap.parse_args()
|
|
a, sa = load_wav(args.a, target_sr=args.sr)
|
|
b, sb = load_wav(args.b, target_sr=args.sr)
|
|
method = "phat" if args.phat else "standard"
|
|
lag, peak, _ = correlate(a, b, sr=args.sr, method=method, device=args.device)
|
|
print(f"device: {device_info()} method={method}")
|
|
print(f"lag_samples={lag} lag_sec={lag/args.sr:.6f}s peak={peak:.4f}")
|
|
print(f"apply: ffmpeg -i video.mkv -i audio.wav -filter_complex \"[1:a]adelay={max(0, lag/args.sr*1000):.0f}|{max(0, lag/args.sr*1000):.0f}[ad]\" -map 0:v -map \"[ad]\" out.mkv (pad silence for delay)")
|