From 2aa6d03810a0915615a2d6c007dc2582d6676b3f Mon Sep 17 00:00:00 2001 From: Honney OC Date: Mon, 14 Sep 2026 02:02:26 +0000 Subject: [PATCH] feat: improve audio stream selection and caching --- apply_sync.py | 23 ++++++++-------------- extract.py | 34 ++++++++++++++++++++++++++++++++ verify_sync.py | 49 +++++++++++++++++++++++++++++++++++++---------- waveform_cache.py | 9 +++++---- 4 files changed, 86 insertions(+), 29 deletions(-) diff --git a/apply_sync.py b/apply_sync.py index 5f31964..8ff6337 100644 --- a/apply_sync.py +++ b/apply_sync.py @@ -75,7 +75,8 @@ def _codec_args(codec: str, out_idx: int) -> list[str]: def build_ffmpeg_cmd(file_a: str, file_b: str, lag_sec: float, out_path: str) -> list[str]: - """Build ffmpeg command to mux A video + A audio + delayed B audio streams.""" + """Build ffmpeg command to mux A video + all A audio + all delayed B audio streams.""" + a_streams = probe_audio_streams(file_a) b_streams = probe_audio_streams(file_b) lag_ms = lag_sec * 1000 @@ -99,22 +100,22 @@ def build_ffmpeg_cmd(file_a: str, file_b: str, lag_sec: float, out_path: str) -> if filters: cmd += ["-filter_complex", ";".join(filters)] - # Map: video from A, all audio from A, all delayed B audio + # Map: video from A, ALL audio from A, all delayed B audio cmd += ["-map", "0:v", "-map", "0:a"] for i in range(len(b_streams)): cmd += ["-map", f"[d{i}]"] # Codecs: copy video + A's audio, re-encode B's filtered audio with same codec - n_a_audio = len(probe_audio_streams(file_a)) + n_a = len(a_streams) cmd += ["-c:v", "copy"] - for i in range(n_a_audio): + for i in range(n_a): cmd += [f"-c:a:{i}", "copy"] for i, s in enumerate(b_streams): - cmd += _codec_args(s["codec"], n_a_audio + i) + cmd += _codec_args(s["codec"], n_a + i) - # Language metadata for B audio streams (they become the last N audio streams in output) + # Language metadata for B audio streams (they start after A's streams) for i, s in enumerate(b_streams): - out_audio_idx = n_a_audio + i # 0-based index in output audio streams + out_audio_idx = n_a + i if set_german and i == 0: cmd += [f"-metadata:s:a:{out_audio_idx}", "language=ger"] elif s["language"] not in ("und", "unknown", ""): @@ -183,14 +184,6 @@ def main(): print(f"\nSynced output: {out_path}") - # Auto-launch verify_sync in full mode - verify_script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "verify_sync.py") - if os.path.exists(verify_script): - verify_cmd = [sys.executable, verify_script, "--config", args.config or "config.toml", - "--lag", f"{args.lag:.6f}", "--mode", "full"] - print(f"Launching verify: {' '.join(verify_cmd)}") - subprocess.Popen(verify_cmd) - sys.exit(0) diff --git a/extract.py b/extract.py index 803d81d..f9fb745 100644 --- a/extract.py +++ b/extract.py @@ -410,6 +410,8 @@ _stems = _derive_stems(INPUT_FILES) # ─── Prepare output folder ─── os.makedirs(OUT_DIR, exist_ok=True) +_audio_stream_selections = [] # track selected stream indices for config + for _file_idx, INPUT_FILE in enumerate(INPUT_FILES): # ─── Derive output filenames ─── @@ -430,6 +432,7 @@ for _file_idx, INPUT_FILE in enumerate(INPUT_FILES): print(f"[audio] {stem}: selected stream {audio_stream}") else: print(f"[audio] {stem}: no audio selected, skipping audio extraction") + _audio_stream_selections.append(audio_stream) # ─── Extract video-only (duration s, no audio, CPU FFV1 re-encode) ─── # Only apply fps filter if probed fps differs from target (avoid needless re-timing) @@ -489,3 +492,34 @@ for _file_idx, INPUT_FILE in enumerate(INPUT_FILES): print(f"-" * 50) print(f"All done. Files are in:", OUT_DIR) + +# ─── Save selected audio stream indices to config.toml ─── +if len(_audio_stream_selections) >= 2 and CONFIG_PATH and os.path.exists(CONFIG_PATH): + try: + with open(CONFIG_PATH, "r") as f: + cfg_text = f.read() + # Update or add audio_stream_a and audio_stream_b in [extract] section + import re + def _cfg_set(key: str, val) -> str: + nonlocal cfg_text + pat = re.compile(rf'(^{key}\s*=\s*).*$') + replacement = f'{key} = {val}' + if pat.search(cfg_text): + cfg_text = pat.sub(replacement, cfg_text, count=1) + else: + # Insert before the next section header or at end of [extract] + ins_pat = re.compile(r'(\[audio\])') + if ins_pat.search(cfg_text): + cfg_text = ins_pat.sub(f'{replacement}\n\n\\1', cfg_text, count=1) + else: + cfg_text += f"\n{replacement}\n" + _cfg_set("audio_stream_a", _audio_stream_selections[0] + if _audio_stream_selections[0] is not None else "null") + _cfg_set("audio_stream_b", _audio_stream_selections[1] + if _audio_stream_selections[1] is not None else "null") + with open(CONFIG_PATH, "w") as f: + f.write(cfg_text) + print(f"[config] saved audio_stream_a={_audio_stream_selections[0]}, " + f"audio_stream_b={_audio_stream_selections[1]} → {CONFIG_PATH}") + except Exception as e: + print(f"[config] warn: could not save stream indices: {e}", file=sys.stderr) diff --git a/verify_sync.py b/verify_sync.py index 1263053..c0a42fc 100644 --- a/verify_sync.py +++ b/verify_sync.py @@ -297,7 +297,8 @@ class VerifyWindow(QtWidgets.QMainWindow): def __init__(self, video_path: str, wav_a_path: str, wav_b_path: str, sr: int = 48000, lag_sec: float = 0.0, mode: str = "short", out_dir: str = "output", - temp_dir: str | None = None, volume: int = 70): + temp_dir: str | None = None, volume: int = 70, + cache_key_a: str | None = None, cache_key_b: str | None = None): super().__init__() self.video_path = video_path self.wav_a_path = wav_a_path @@ -325,9 +326,9 @@ class VerifyWindow(QtWidgets.QMainWindow): # Load audio print(f"Loading audio A: {wav_a_path}") - self.a_orig, _ = cached_load_wav(wav_a_path, target_sr=sr, _loader=load_wav) + self.a_orig, _ = cached_load_wav(wav_a_path, target_sr=sr, _loader=load_wav, _cache_key=cache_key_a) print(f"Loading audio B: {wav_b_path}") - self.b_orig, _ = cached_load_wav(wav_b_path, target_sr=sr, _loader=load_wav) + self.b_orig, _ = cached_load_wav(wav_b_path, target_sr=sr, _loader=load_wav, _cache_key=cache_key_b) # Apply lag alignment self.a_aligned, self.b_aligned = apply_lag(self.a_orig, self.b_orig, self.lag_samples) @@ -688,18 +689,44 @@ def main(): print("ERROR: synced MKV needs at least 2 audio streams (A + B)", file=sys.stderr) sys.exit(1) + # Determine which streams to extract from synced MKV + # A's streams are first, B's delayed streams follow + video_a = ext.get("file_a", "") + a_streams_in = probe_audio_streams(video_a) if os.path.exists(video_a) else [] + n_a = len(a_streams_in) + sel_a = ext.get("audio_stream_a", 0) or 0 + sel_b = ext.get("audio_stream_b", 0) or 0 + sel_a = int(sel_a) + sel_b = int(sel_b) + # Output indices: A's selected stream stays at sel_a, B's becomes n_a + sel_b + out_idx_a = sel_a + out_idx_b = n_a + sel_b + print(f"A has {n_a} streams, selected [{sel_a}] -> output [{out_idx_a}]") + print(f"B selected [{sel_b}] -> output [{out_idx_b}]") + + if out_idx_a >= len(streams) or out_idx_b >= len(streams): + print(f"ERROR: computed stream indices [{out_idx_a}, {out_idx_b}] out of range " + f"(file has {len(streams)} streams)", file=sys.stderr) + sys.exit(1) + temp_dir = tempfile.mkdtemp(prefix="verify_sync_") wav_a = os.path.join(temp_dir, "audio_a.wav") wav_b = os.path.join(temp_dir, "audio_b.wav") - print(f"Extracting audio stream 0 (A) to {wav_a}...") - if not extract_audio_from_mkv(synced, 0, wav_a, sr): - print("ERROR: failed to extract audio stream 0", file=sys.stderr) + # Stable cache keys based on source MKV + stream index (not temp file path) + import hashlib + _mkv_mtime = os.path.getmtime(synced) + cache_key_a = hashlib.sha256(f"{synced}|{out_idx_a}|{sr}|{_mkv_mtime:.6f}".encode()).hexdigest()[:16] + cache_key_b = hashlib.sha256(f"{synced}|{out_idx_b}|{sr}|{_mkv_mtime:.6f}".encode()).hexdigest()[:16] + + print(f"Extracting audio stream {out_idx_a} (A) to {wav_a}...") + if not extract_audio_from_mkv(synced, out_idx_a, wav_a, sr): + print(f"ERROR: failed to extract audio stream {out_idx_a}", file=sys.stderr) sys.exit(1) - print(f"Extracting audio stream 1 (B) to {wav_b}...") - if not extract_audio_from_mkv(synced, 1, wav_b, sr): - print("ERROR: failed to extract audio stream 1", file=sys.stderr) + print(f"Extracting audio stream {out_idx_b} (B) to {wav_b}...") + if not extract_audio_from_mkv(synced, out_idx_b, wav_b, sr): + print(f"ERROR: failed to extract audio stream {out_idx_b}", file=sys.stderr) sys.exit(1) video_path = synced @@ -737,6 +764,8 @@ def main(): else: print(f"WARNING: extracted video not found ({extracted_video}), using original", file=sys.stderr) temp_dir = None + cache_key_a = None + cache_key_b = None print(f"Mode: {args.mode}") print(f"Video: {video_path}") @@ -750,7 +779,7 @@ def main(): w = VerifyWindow(video_path, wav_a, wav_b, sr=sr, lag_sec=lag_sec, mode=args.mode, out_dir=out_dir, temp_dir=temp_dir, - volume=volume) + volume=volume, cache_key_a=cache_key_a, cache_key_b=cache_key_b) w.show() sys.exit(app.exec()) diff --git a/waveform_cache.py b/waveform_cache.py index bcd56c7..f707841 100644 --- a/waveform_cache.py +++ b/waveform_cache.py @@ -11,7 +11,7 @@ import numpy as np CACHE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "cache", "waveforms") -def _cache_key(path: str, target_sr: int, mono: bool) -> str: +def _cache_key_path(path: str, target_sr: int, mono: bool) -> str: real = os.path.realpath(path) mtime = os.path.getmtime(real) raw = f"{real}|{target_sr}|{mono}|{mtime:.6f}" @@ -19,9 +19,10 @@ def _cache_key(path: str, target_sr: int, mono: bool) -> str: def cached_load_wav(path: str, target_sr: int = 48000, mono: bool = True, - _loader=None) -> tuple[np.ndarray, int]: - """Load WAV with disk cache. Falls back to _loader(data, sr) if cache miss.""" - key = _cache_key(path, target_sr, mono) + _loader=None, _cache_key: str | None = None) -> tuple[np.ndarray, int]: + """Load WAV with disk cache. Falls back to _loader(data, sr) if cache miss. + _cache_key: override cache key (e.g. source_mkv + stream_idx) for temp files.""" + key = _cache_key if _cache_key else _cache_key_path(path, target_sr, mono) npy_path = os.path.join(CACHE_DIR, f"{key}.npy") sr_path = os.path.join(CACHE_DIR, f"{key}.sr")