feat: add config-driven extraction and audio sync tooling
This commit is contained in:
Honney OC committed
1 parent
476b2543d1
commit
3a18a8b2c2
5 files changed
+798
-10
No files matched your search
@@ -0,0 +1,415 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PyQtGraph GUI to visualize audio cross-correlation and interactively tune lag.
|
||||
|
||||
Features (per user request):
|
||||
- uv venv via Makefile, torch auto GPU (ROCm/CUDA) with cpu fallback
|
||||
- config.toml for file_a/file_b, duration/start (extract duration is source of truth)
|
||||
- PyQtGraph waveforms (decimated for 4M samples) + xcorr plot
|
||||
- QSlider + QDoubleSpinBox numeric entry for lag (seconds + samples)
|
||||
- Volume slider (sounddevice optional, playback with lag applied, silence padded for delay)
|
||||
|
||||
Run: make ast or .venv/bin/python sync_audio_test.py --config config.toml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
import os, sys, math, argparse, pathlib
|
||||
import numpy as np
|
||||
|
||||
# allow tomllib (py3.11+) with tomli fallback
|
||||
try:
|
||||
import tomllib
|
||||
except ModuleNotFoundError:
|
||||
import tomli as tomllib
|
||||
|
||||
try:
|
||||
import pyqtgraph as pg
|
||||
from PyQt6 import QtWidgets, QtCore
|
||||
HAS_PG = True
|
||||
except ImportError as e:
|
||||
print("ERROR: pyqtgraph/PyQt6 not installed — run make install", file=sys.stderr)
|
||||
print(e, file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
try:
|
||||
import soundfile as sf
|
||||
HAS_SF = True
|
||||
except ImportError:
|
||||
HAS_SF = False
|
||||
|
||||
try:
|
||||
import sounddevice as sd
|
||||
HAS_SD = True
|
||||
except ImportError:
|
||||
HAS_SD = False
|
||||
|
||||
from sync_audio import load_wav, correlate, device_info, lag_to_seconds
|
||||
|
||||
|
||||
def load_config(path: str) -> dict:
|
||||
if not os.path.exists(path):
|
||||
return {}
|
||||
with open(path, "rb") as f:
|
||||
return tomllib.load(f)
|
||||
|
||||
|
||||
def decimate_for_plot(arr: np.ndarray, max_pts: int = 120_000) -> tuple[np.ndarray, int]:
|
||||
if len(arr) <= max_pts:
|
||||
return arr, 1
|
||||
stride = int(math.ceil(len(arr) / max_pts))
|
||||
return arr[::stride], stride
|
||||
|
||||
|
||||
def find_wavs_from_config(cfg: dict) -> tuple[str, str, int, float, float]:
|
||||
audio_cfg = cfg.get("audio", {})
|
||||
extract_cfg = cfg.get("extract", {})
|
||||
sync_cfg = cfg.get("sync", {})
|
||||
sr = int(audio_cfg.get("sample_rate", 48000))
|
||||
# priority: explicit wav paths else derived from extracted dir
|
||||
a_wav = audio_cfg.get("file_a_wav")
|
||||
b_wav = audio_cfg.get("file_b_wav")
|
||||
# fallback to extract files
|
||||
out_dir = extract_cfg.get("out_dir", "extracted")
|
||||
file_a = extract_cfg.get("file_a")
|
||||
file_b = extract_cfg.get("file_b")
|
||||
if not a_wav and file_a:
|
||||
stem = os.path.splitext(os.path.basename(file_a))[0]
|
||||
cand = os.path.join(out_dir, f"{stem}_audio.wav")
|
||||
if os.path.exists(cand):
|
||||
a_wav = cand
|
||||
else:
|
||||
a_wav = cand # keep for error msg
|
||||
if not b_wav and file_b:
|
||||
stem = os.path.splitext(os.path.basename(file_b))[0]
|
||||
cand = os.path.join(out_dir, f"{stem}_audio.wav")
|
||||
if os.path.exists(cand):
|
||||
b_wav = cand
|
||||
else:
|
||||
b_wav = cand
|
||||
# also allow sync-specific override
|
||||
duration = extract_cfg.get("duration", 90)
|
||||
start = float(extract_cfg.get("start", 0.0))
|
||||
# duration_sec comment: only duration needed is off the extract — we respect that
|
||||
return a_wav, b_wav, sr, float(start), float(duration)
|
||||
|
||||
|
||||
class SyncWindow(QtWidgets.QMainWindow):
|
||||
def __init__(self, a_path: str, b_path: str, sr: int = 48000, method: str = "phat", device: str = "auto"):
|
||||
super().__init__()
|
||||
self.setWindowTitle(f"movie_sync — audio sync test — {device_info()}")
|
||||
self.resize(1400, 900)
|
||||
self.a_path = a_path
|
||||
self.b_path = b_path
|
||||
self.sr = sr
|
||||
self.method = method
|
||||
self.device_pref = device
|
||||
self._updating = False
|
||||
|
||||
# load
|
||||
print(f"Loading {a_path}")
|
||||
self.a, self.sr_a = load_wav(a_path, target_sr=sr)
|
||||
print(f"Loading {b_path}")
|
||||
self.b, self.sr_b = load_wav(b_path, target_sr=sr)
|
||||
# keep originals for lag apply
|
||||
self.a_orig = self.a.copy()
|
||||
self.b_orig = self.b.copy()
|
||||
|
||||
# initial correlate
|
||||
self.lag_samples, self.peak, self.corr = correlate(self.a, self.b, sr=sr, method=method, device=device)
|
||||
self.lag_sec = self.lag_samples / sr
|
||||
print(f"initial lag {self.lag_samples} samples = {self.lag_sec:.6f}s peak={self.peak:.4f}")
|
||||
|
||||
# decimated views
|
||||
self.a_plot, self.stride_a = decimate_for_plot(self.a)
|
||||
self.b_plot, self.stride_b = decimate_for_plot(self.b)
|
||||
self.t_a = np.arange(len(self.a_plot)) * self.stride_a / sr
|
||||
self.t_b = np.arange(len(self.b_plot)) * self.stride_b / sr
|
||||
|
||||
# xcorr display — corr is already fftshifted
|
||||
N = len(self.corr)
|
||||
# lags in seconds for display
|
||||
# raw corr was shifted: center at N//2, range [- (len(b)-1) .. len(a)-1]
|
||||
# for display we map indices to lag
|
||||
total = len(self.a) + len(self.b) - 1
|
||||
n_pow2 = N
|
||||
# shifted lags: -N//2 .. N//2-1
|
||||
self.corr_lags = np.arange(N) - N//2
|
||||
self.corr_lags_sec = self.corr_lags / sr
|
||||
|
||||
central = QtWidgets.QWidget()
|
||||
self.setCentralWidget(central)
|
||||
lay = QtWidgets.QVBoxLayout(central)
|
||||
|
||||
# info bar
|
||||
info = QtWidgets.QLabel(f"A: {os.path.basename(a_path)} ({len(self.a)/sr:.1f}s) B: {os.path.basename(b_path)} ({len(self.b)/sr:.1f}s) sr={sr} method={method} {device_info()}")
|
||||
info.setWordWrap(True)
|
||||
lay.addWidget(info)
|
||||
|
||||
# plots
|
||||
self.gw = pg.GraphicsLayoutWidget()
|
||||
lay.addWidget(self.gw, stretch=1)
|
||||
|
||||
# waveform A
|
||||
self.p1 = self.gw.addPlot(row=0, col=0, title="Waveform A (reference)")
|
||||
self.p1.showGrid(x=True, y=True)
|
||||
self.p1.setLabel("bottom", "Time", "s")
|
||||
self.curve_a = self.p1.plot(self.t_a, self.a_plot, pen=pg.mkPen("#4cc9f0", width=1))
|
||||
self.p1.setMouseEnabled(x=True, y=False)
|
||||
|
||||
# waveform B with lag
|
||||
self.p2 = self.gw.addPlot(row=1, col=0, title=f"Waveform B (shifted by lag)")
|
||||
self.p2.showGrid(x=True, y=True)
|
||||
self.p2.setLabel("bottom", "Time", "s")
|
||||
self.curve_b = self.p2.plot(self.t_b, self.b_plot, pen=pg.mkPen("#f72585", width=1))
|
||||
self.p2.setMouseEnabled(x=True, y=False)
|
||||
# link x
|
||||
self.p2.setXLink(self.p1)
|
||||
|
||||
# xcorr
|
||||
self.p3 = self.gw.addPlot(row=2, col=0, title="Cross-correlation (GCC-PHAT)" if method=="phat" else "Cross-correlation")
|
||||
self.p3.showGrid(x=True, y=True)
|
||||
self.p3.setLabel("bottom", "Lag", "s")
|
||||
# decimate corr for perf if huge
|
||||
corr_dec, corr_stride = decimate_for_plot(self.corr, max_pts=200_000)
|
||||
corr_lags_dec = self.corr_lags_sec[::corr_stride] if corr_stride>1 else self.corr_lags_sec
|
||||
self.curve_corr = self.p3.plot(corr_lags_dec, np.abs(corr_dec), pen=pg.mkPen("#ffd166", width=1))
|
||||
self.peak_line = pg.InfiniteLine(pos=self.lag_sec, angle=90, pen=pg.mkPen("#06d6a0", width=2, style=QtCore.Qt.PenStyle.DashLine))
|
||||
self.p3.addItem(self.peak_line)
|
||||
self.p3.setXRange(self.lag_sec-2, self.lag_sec+2)
|
||||
|
||||
# controls
|
||||
ctr = QtWidgets.QHBoxLayout()
|
||||
lay.addLayout(ctr)
|
||||
|
||||
ctr.addWidget(QtWidgets.QLabel("Lag (s):"))
|
||||
self.slider = QtWidgets.QSlider(QtCore.Qt.Orientation.Horizontal)
|
||||
self.slider.setRange(-10000, 10000) # will map to seconds via scale
|
||||
# range +/-10s in ms
|
||||
self.slider.setMinimum(int(-10 * 1000))
|
||||
self.slider.setMaximum(int(10 * 1000))
|
||||
self.slider.setValue(int(self.lag_sec * 1000))
|
||||
self.slider.setTickInterval(100)
|
||||
ctr.addWidget(self.slider, stretch=1)
|
||||
|
||||
self.spin_sec = QtWidgets.QDoubleSpinBox()
|
||||
self.spin_sec.setRange(-60, 60)
|
||||
self.spin_sec.setDecimals(6)
|
||||
self.spin_sec.setSingleStep(0.001)
|
||||
self.spin_sec.setValue(self.lag_sec)
|
||||
self.spin_sec.setSuffix(" s")
|
||||
self.spin_sec.setKeyboardTracking(False)
|
||||
ctr.addWidget(self.spin_sec)
|
||||
|
||||
self.spin_samp = QtWidgets.QSpinBox()
|
||||
self.spin_samp.setRange(-10_000_000, 10_000_000)
|
||||
self.spin_samp.setSingleStep(1)
|
||||
self.spin_samp.setValue(self.lag_samples)
|
||||
self.spin_samp.setSuffix(" samp")
|
||||
ctr.addWidget(self.spin_samp)
|
||||
|
||||
btn_auto = QtWidgets.QPushButton("Auto peak")
|
||||
ctr.addWidget(btn_auto)
|
||||
|
||||
btn_swap = QtWidgets.QPushButton("Swap A/B")
|
||||
ctr.addWidget(btn_swap)
|
||||
|
||||
# volume + playback
|
||||
ctr2 = QtWidgets.QHBoxLayout()
|
||||
lay.addLayout(ctr2)
|
||||
ctr2.addWidget(QtWidgets.QLabel("Volume:"))
|
||||
self.vol = QtWidgets.QSlider(QtCore.Qt.Orientation.Horizontal)
|
||||
self.vol.setRange(0, 100)
|
||||
self.vol.setValue(70)
|
||||
self.vol.setMaximumWidth(200)
|
||||
ctr2.addWidget(self.vol)
|
||||
self.vol_label = QtWidgets.QLabel("70%")
|
||||
ctr2.addWidget(self.vol_label)
|
||||
|
||||
self.btn_play_a = QtWidgets.QPushButton("Play A (2s)")
|
||||
self.btn_play_b = QtWidgets.QPushButton("Play B+lag (2s)")
|
||||
self.btn_play_mix = QtWidgets.QPushButton("Play Mix")
|
||||
self.btn_stop = QtWidgets.QPushButton("Stop")
|
||||
for b in [self.btn_play_a, self.btn_play_b, self.btn_play_mix, self.btn_stop]:
|
||||
ctr2.addWidget(b)
|
||||
if not HAS_SD:
|
||||
for b in [self.btn_play_a, self.btn_play_b, self.btn_play_mix, self.btn_stop]:
|
||||
b.setEnabled(False)
|
||||
b.setToolTip("sounddevice not installed — run make install (optional)")
|
||||
ctr2.addWidget(QtWidgets.QLabel("(sounddevice missing — playback disabled)"))
|
||||
|
||||
self.ffmpeg_label = QtWidgets.QLabel()
|
||||
self.ffmpeg_label.setWordWrap(True)
|
||||
self.ffmpeg_label.setTextInteractionFlags(QtCore.Qt.TextInteractionFlag.TextSelectableByMouse)
|
||||
lay.addWidget(self.ffmpeg_label)
|
||||
|
||||
# signals
|
||||
self.slider.valueChanged.connect(self.on_slider)
|
||||
self.spin_sec.valueChanged.connect(self.on_spin_sec)
|
||||
self.spin_samp.valueChanged.connect(self.on_spin_samp)
|
||||
self.vol.valueChanged.connect(lambda v: self.vol_label.setText(f"{v}%"))
|
||||
btn_auto.clicked.connect(self.on_auto)
|
||||
btn_swap.clicked.connect(self.on_swap)
|
||||
self.btn_play_a.clicked.connect(lambda: self.play("a"))
|
||||
self.btn_play_b.clicked.connect(lambda: self.play("b"))
|
||||
self.btn_play_mix.clicked.connect(lambda: self.play("mix"))
|
||||
self.btn_stop.clicked.connect(self.stop_play)
|
||||
|
||||
self.update_labels()
|
||||
|
||||
def _set_lag(self, lag_sec: float, source: str = ""):
|
||||
if self._updating:
|
||||
return
|
||||
self._updating = True
|
||||
lag_sec = float(np.clip(lag_sec, -60, 60))
|
||||
lag_samp = int(round(lag_sec * self.sr))
|
||||
self.lag_sec = lag_sec
|
||||
self.lag_samples = lag_samp
|
||||
# update widgets
|
||||
self.slider.blockSignals(True); self.slider.setValue(int(lag_sec*1000)); self.slider.blockSignals(False)
|
||||
self.spin_sec.blockSignals(True); self.spin_sec.setValue(lag_sec); self.spin_sec.blockSignals(False)
|
||||
self.spin_samp.blockSignals(True); self.spin_samp.setValue(lag_samp); self.spin_samp.blockSignals(False)
|
||||
self.peak_line.setPos(lag_sec)
|
||||
# shift B curve visually by lag (pad with silence not needed for visualize, just shift x)
|
||||
# we keep t_b but offset x
|
||||
shift = lag_sec
|
||||
# pyqtgraph setPos shifts items; we recreate x
|
||||
# Instead shift data: plot t_b + lag
|
||||
self.curve_b.setData(self.t_b + shift, self.b_plot)
|
||||
self.update_labels()
|
||||
self._updating = False
|
||||
|
||||
def update_labels(self):
|
||||
# ffmpeg adelay for mux: add silence for delay (no sound for build-in delay)
|
||||
# If lag>0: B delayed vs A => to mux B onto A video, delay B audio by lag (pad silence)
|
||||
# adelay expects ms per channel, e.g. adelay=1234|1234
|
||||
lag_ms = self.lag_sec * 1000
|
||||
if lag_ms >= 0:
|
||||
ffcmd = f'ffmpeg -i video.mkv -i audio.wav -filter_complex "[1:a]adelay={lag_ms:.0f}:all=1[ad]" -map 0:v -map "[ad]" -c:v copy -c:a pcm_s16le out.mkv'
|
||||
note = f"B delayed by {self.lag_sec:.6f}s ({self.lag_samples} samp) → pad B start with silence {lag_ms:.0f} ms"
|
||||
else:
|
||||
# B earlier => pad A (or trim B). For adding B to A file, we need to delay video or pad A? Simplest pad A? But we control only audio: delay B negative means B earlier, so we delay A or use -itsoffset
|
||||
note = f"B earlier by {abs(self.lag_sec):.6f}s → pad A / use -itsoffset {abs(lag_ms):.0f}ms or trim B by {abs(self.lag_samples)} samp"
|
||||
ffcmd = f'# B early: ffmpeg -itsoffset {abs(lag_ms)/1000:.6f} -i audio.wav -i video.mkv ... OR pad video audio with silence'
|
||||
self.ffmpeg_label.setText(f"{note}\n{ffcmd}")
|
||||
|
||||
# slots
|
||||
def on_slider(self, v):
|
||||
self._set_lag(v/1000.0, "slider")
|
||||
|
||||
def on_spin_sec(self, v):
|
||||
self._set_lag(float(v), "sec")
|
||||
|
||||
def on_spin_samp(self, v):
|
||||
self._set_lag(int(v)/self.sr, "samp")
|
||||
|
||||
def on_auto(self):
|
||||
# recompute peak
|
||||
lag, peak, corr = correlate(self.a_orig, self.b_orig, sr=self.sr, method=self.method, device=self.device_pref)
|
||||
self.corr = corr
|
||||
# update corr curve? re-decimate
|
||||
self._set_lag(lag / self.sr, "auto")
|
||||
print(f"auto lag {lag} samp {lag/self.sr:.6f}s peak={peak:.4f}")
|
||||
|
||||
def on_swap(self):
|
||||
self.a_orig, self.b_orig = self.b_orig, self.a_orig
|
||||
self.a, self.b = self.a_orig.copy(), self.b_orig.copy()
|
||||
self.a_plot, _ = decimate_for_plot(self.a)
|
||||
self.b_plot, _ = decimate_for_plot(self.b)
|
||||
self.curve_a.setData(self.t_a, self.a_plot)
|
||||
# recalc
|
||||
lag, _, corr = correlate(self.a, self.b, sr=self.sr, method=self.method, device=self.device_pref)
|
||||
self.corr = corr
|
||||
self._set_lag(lag/self.sr, "swap")
|
||||
|
||||
def play(self, which: str):
|
||||
if not HAS_SD:
|
||||
return
|
||||
vol = self.vol.value()/100.0
|
||||
sr = self.sr
|
||||
# play 2 seconds from 10s mark to hear sync
|
||||
start_samp = int(10 * sr)
|
||||
dur_samp = int(2 * sr)
|
||||
if start_samp + dur_samp > len(self.a_orig):
|
||||
start_samp = 0
|
||||
a_seg = self.a_orig[start_samp:start_samp+dur_samp] * vol
|
||||
# apply lag with silence padding (no sound for delay) — Option A
|
||||
lag = self.lag_samples
|
||||
# For playback mix, we need to align: if B delayed (lag>0), pad B front with silence
|
||||
# For preview we shift B segment accordingly
|
||||
b_start = start_samp
|
||||
if lag > 0:
|
||||
# B delayed => B segment should be taken earlier to align? For playback we pad B with lag silence, so effective B at time t is original at t-lag
|
||||
idx = b_start - lag
|
||||
if idx < 0:
|
||||
b_seg = np.concatenate([np.zeros(-idx, dtype=np.float32), self.b_orig[0:b_start + dur_samp - (-idx) if b_start+dur_samp - (-idx)>0 else 0]])
|
||||
# slice to dur
|
||||
b_seg = b_seg[:dur_samp]
|
||||
if len(b_seg) < dur_samp:
|
||||
b_seg = np.concatenate([b_seg, np.zeros(dur_samp-len(b_seg), dtype=np.float32)])
|
||||
else:
|
||||
b_seg = self.b_orig[idx:idx+dur_samp]
|
||||
elif lag < 0:
|
||||
idx = b_start - lag # lag negative => plus
|
||||
b_seg = self.b_orig[idx:idx+dur_samp] if idx+dur_samp < len(self.b_orig) else np.zeros(dur_samp, dtype=np.float32)
|
||||
if len(b_seg) < dur_samp:
|
||||
b_seg = np.concatenate([b_seg, np.zeros(dur_samp-len(b_seg), dtype=np.float32)])
|
||||
else:
|
||||
b_seg = self.b_orig[b_start:b_start+dur_samp]
|
||||
b_seg = b_seg * vol
|
||||
try:
|
||||
sd.stop()
|
||||
if which == "a":
|
||||
sd.play(a_seg, samplerate=sr)
|
||||
elif which == "b":
|
||||
sd.play(b_seg, samplerate=sr)
|
||||
elif which == "mix":
|
||||
mix = (a_seg + b_seg) / 2.0
|
||||
sd.play(mix, samplerate=sr)
|
||||
except Exception as e:
|
||||
print(f"playback failed {e}", file=sys.stderr)
|
||||
|
||||
def stop_play(self):
|
||||
if HAS_SD:
|
||||
sd.stop()
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--config", default="config.toml")
|
||||
ap.add_argument("--a", default=None)
|
||||
ap.add_argument("--b", default=None)
|
||||
ap.add_argument("--device", default=None, choices=["auto","cpu","cuda"])
|
||||
ap.add_argument("--method", default=None, choices=["phat","standard"])
|
||||
args = ap.parse_args()
|
||||
|
||||
cfg = load_config(args.config)
|
||||
cfg_a, cfg_b, sr, _, _ = find_wavs_from_config(cfg)
|
||||
a_path = args.a or cfg_a
|
||||
b_path = args.b or cfg_b
|
||||
method = args.method or cfg.get("sync", {}).get("method", "phat")
|
||||
device = args.device or cfg.get("sync", {}).get("use_gpu", "auto")
|
||||
# map toml "use_gpu" to device
|
||||
if device == "auto":
|
||||
device = "auto"
|
||||
elif device == "cpu":
|
||||
device = "cpu"
|
||||
elif device in ("cuda", "gpu"):
|
||||
device = "cuda"
|
||||
|
||||
if not a_path or not b_path:
|
||||
print(f"ERROR: no input wavs. Set file_a/file_b in {args.config} or --a/--b", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
if not os.path.exists(a_path):
|
||||
print(f"ERROR: A not found {a_path}", file=sys.stderr); sys.exit(1)
|
||||
if not os.path.exists(b_path):
|
||||
print(f"ERROR: B not found {b_path}", file=sys.stderr); sys.exit(1)
|
||||
|
||||
app = QtWidgets.QApplication(sys.argv)
|
||||
pg.setConfigOptions(antialias=True, background="#121212", foreground="#eee")
|
||||
w = SyncWindow(a_path, b_path, sr=sr, method=method, device=device)
|
||||
w.show()
|
||||
sys.exit(app.exec())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user