#!/usr/bin/env python3 """ Visual picker for extraction settings. Shows video thumbnails, audio waveforms, and lets you set start/duration. Usage: make picker or .venv/bin/python picker.py --config config.toml """ from __future__ import annotations import os, sys, argparse, subprocess, json, math, signal, struct import numpy as np try: import tomllib except ModuleNotFoundError: import tomli as tomllib try: import pyqtgraph as pg from PyQt6 import QtWidgets, QtCore, QtGui except ImportError: print("ERROR: pyqtgraph/PyQt6 not installed — run make install", file=sys.stderr) sys.exit(1) from verify_sync import apply_dark_theme def probe_video_info(path: str) -> dict: """Get video duration, resolution, audio streams.""" try: r = subprocess.run( ["ffprobe", "-v", "error", "-show_entries", "format=duration:stream=index,codec_type,codec_name,width,height,sample_rate,channels", "-of", "json", path], capture_output=True, text=True, timeout=10 ) if r.returncode != 0: return {} data = json.loads(r.stdout) fmt = data.get("format", {}) streams = data.get("streams", []) video = next((s for s in streams if s.get("codec_type") == "video"), {}) audio = [s for s in streams if s.get("codec_type") == "audio"] return { "duration": float(fmt.get("duration", 0)), "width": video.get("width", 0), "height": video.get("height", 0), "fps": video.get("r_frame_rate", ""), "audio_streams": audio, } except Exception: return {} def extract_rms_envelope(path: str, sr: int = 48000, stream_idx: int = 0, block_size: int = 4800) -> tuple[np.ndarray, np.ndarray] | None: """Extract RMS envelope using ffmpeg astats — fast, memory-light.""" # astats with reset gives RMS per block. We use amerge+pan to get mono, # then astats with reset=block_size to get one RMS value per block. cmd = [ "ffmpeg", "-y", "-i", path, "-map", f"0:a:{stream_idx}", "-af", f"aresample={sr},pan=mono|c0=c0,astats=metadata=1:reset={block_size}:measure=Peak+RMS", "-f", "null", "-" ] try: # This doesn't work well — use a different approach: extract raw PCM but in chunks pass except Exception: pass # Fallback: extract full PCM (fast enough for short clips, for full file use decimation) cmd = [ "ffmpeg", "-y", "-i", path, "-map", f"0:a:{stream_idx}", "-ac", "1", "-ar", str(sr), "-f", "f32le", "-" ] try: r = subprocess.run(cmd, capture_output=True, timeout=120) if r.returncode != 0: return None pcm = np.frombuffer(r.stdout, dtype=np.float32) if len(pcm) == 0: return None # Decimate to ~10 Hz (one value per 0.1s) decimate = max(1, sr // 10) n = (len(pcm) // decimate) * decimate pcm = pcm[:n].reshape(-1, decimate) rms = np.sqrt(np.mean(pcm ** 2, axis=1)) t = np.arange(len(rms), dtype=np.float64) * decimate / sr return t, rms except Exception: return None def extract_audio_segment(path: str, start: float, duration: float, sr: int = 48000, stream_idx: int = 0) -> np.ndarray | None: """Extract a mono float32 PCM segment from video file.""" cmd = [ "ffmpeg", "-y", "-ss", str(start), "-i", path, "-t", str(duration), "-vn", "-map", f"0:a:{stream_idx}", "-ac", "1", "-ar", str(sr), "-f", "f32le", "-" ] try: r = subprocess.run(cmd, capture_output=True, timeout=30) if r.returncode != 0: return None return np.frombuffer(r.stdout, dtype=np.float32) except Exception: return None class LoadWorker(QtCore.QThread): """Background thread for loading waveforms and thumbnails.""" waveform_ready = QtCore.pyqtSignal(str, object, object) # which, t, rms thumbnail_ready = QtCore.pyqtSignal(object) status = QtCore.pyqtSignal(str) def __init__(self, file_a: str, file_b: str, sr: int, stream_idx: int = 0): super().__init__() self.file_a = file_a self.file_b = file_b self.sr = sr self.stream_idx = stream_idx self._abort = False def abort(self): self._abort = True def run(self): # Load waveform A self.status.emit("Loading waveform A...") result = extract_rms_envelope(self.file_a, self.sr, self.stream_idx) if self._abort: return if result is not None: self.waveform_ready.emit("a", result[0], result[1]) # Load waveform B if self.file_b and os.path.exists(self.file_b): self.status.emit("Loading waveform B...") result = extract_rms_envelope(self.file_b, self.sr, 0) if self._abort: return if result is not None: self.waveform_ready.emit("b", result[0], result[1]) # Load thumbnails self.status.emit("Extracting thumbnails...") self._load_thumbnails() def _load_thumbnails(self): try: # Get video duration for thumbnail range info = probe_video_info(self.file_a) dur = info.get("duration", 0) if dur <= 0: return # Extract thumbnail sprite sheet across the full video n_thumbs = 40 step = dur / n_thumbs r = subprocess.run( ["ffmpeg", "-y", "-i", self.file_a, "-vf", f"fps=1/{max(0.5, step):.2f},scale=160:-1,tile={n_thumbs}x1"], capture_output=True, timeout=60 ) if self._abort: return if r.returncode == 0 and len(r.stdout) > 0: import tempfile with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: tmp.write(r.stdout) tmp_path = tmp.name pixmap = QtGui.QPixmap(tmp_path) os.unlink(tmp_path) if not pixmap.isNull(): self.thumbnail_ready.emit(pixmap) except Exception as e: print(f"Thumbnail error: {e}", file=sys.stderr) class PickerWindow(QtWidgets.QWidget): def __init__(self, file_a: str, file_b: str, start: float, duration: float, sr: int = 48000): super().__init__() self.file_a = file_a self.file_b = file_b self.sr = sr self.video_info = probe_video_info(file_a) self.total_duration = self.video_info.get("duration", 0) self.audio_streams = self.video_info.get("audio_streams", []) self.start = max(0.0, min(start, self.total_duration)) self.duration = max(1.0, duration) self.worker = None self.setWindowTitle(f"picker — {os.path.basename(file_a)}") self.resize(1400, 700) self._setup_ui() self._start_loading() def _setup_ui(self): lay = QtWidgets.QVBoxLayout(self) lay.setSpacing(4) # ── Video info ── info = QtWidgets.QHBoxLayout() lay.addLayout(info) self.info_label = QtWidgets.QLabel( f"Video: {self.video_info.get('width', '?')}x{self.video_info.get('height', '?')} | " f"{self.total_duration:.1f}s ({self.total_duration/60:.1f}min) | " f"{len(self.audio_streams)} audio stream(s)" ) info.addWidget(self.info_label) info.addStretch() # ── Selection controls ── sel = QtWidgets.QHBoxLayout() lay.addLayout(sel) sel.addWidget(QtWidgets.QLabel("Start (s):")) self.start_spin = QtWidgets.QDoubleSpinBox() self.start_spin.setRange(0, max(0, self.total_duration - 1)) self.start_spin.setDecimals(3) self.start_spin.setSingleStep(0.1) self.start_spin.setValue(self.start) self.start_spin.setFixedWidth(90) sel.addWidget(self.start_spin) sel.addWidget(QtWidgets.QLabel("Duration (s):")) self.dur_spin = QtWidgets.QDoubleSpinBox() self.dur_spin.setRange(1, max(1, self.total_duration)) self.dur_spin.setDecimals(1) self.dur_spin.setSingleStep(1) self.dur_spin.setValue(self.duration) self.dur_spin.setFixedWidth(70) sel.addWidget(self.dur_spin) sel.addWidget(QtWidgets.QLabel("End:")) self.end_label = QtWidgets.QLabel(f"{self.start + self.duration:.1f}s") sel.addWidget(self.end_label) sel.addSpacing(20) self.start_spin.valueChanged.connect(self._on_start_changed) self.dur_spin.valueChanged.connect(self._on_dur_changed) # ── Audio stream selector ── if len(self.audio_streams) > 1: sel.addWidget(QtWidgets.QLabel("Audio stream:")) self.stream_combo = QtWidgets.QComboBox() for i, s in enumerate(self.audio_streams): lang = s.get("language", "?") ch = s.get("channels", "?") codec = s.get("codec_name", "?") self.stream_combo.addItem(f"Stream {i}: {lang} {ch}ch {codec}") self.stream_combo.setFixedWidth(200) sel.addWidget(self.stream_combo) self.stream_combo.currentIndexChanged.connect(self._on_stream_changed) else: self.stream_combo = None sel.addStretch() self.status_label = QtWidgets.QLabel("") self.status_label.setMinimumWidth(200) sel.addWidget(self.status_label) btn_apply = QtWidgets.QPushButton("Save to config.toml") btn_apply.setFixedWidth(130) btn_apply.clicked.connect(self._save_config) sel.addWidget(btn_apply) # ── Plots ── # Row 0: Thumbnails self.thumb_label = QtWidgets.QLabel("Thumbnails loading...") self.thumb_label.setMinimumHeight(100) self.thumb_label.setAlignment(QtCore.Qt.AlignmentFlag.AlignCenter) lay.addWidget(self.thumb_label, stretch=1) # Row 1: Waveform A self.plot_a = pg.PlotWidget(name="waveform_a") self.plot_a.setLabel("bottom", "Time", units="s") self.plot_a.setLabel("left", "A (RMS)") self.plot_a.showGrid(x=True, y=True, alpha=0.3) lay.addWidget(self.plot_a, stretch=2) # Row 2: Waveform B self.plot_b = pg.PlotWidget(name="waveform_b") self.plot_b.setLabel("bottom", "Time", units="s") self.plot_b.setLabel("left", "B (RMS)") self.plot_b.showGrid(x=True, y=True, alpha=0.3) lay.addWidget(self.plot_b, stretch=2) # Link x-axes self.plot_b.setXLink(self.plot_a) # Selection region self.sel_region = pg.LinearRegionItem([self.start, self.start + self.duration], movable=True) self.sel_region.setBrush(pg.mkBrush(100, 100, 255, 40)) self.sel_region.sigRegionChanged.connect(self._on_region_changed) self.plot_a.addItem(self.sel_region) self.plot_b.addItem(self.sel_region) self.curve_a = self.plot_a.plot(pen=pg.mkPen("c", width=1), fillLevel=0, brush=pg.mkBrush(0, 200, 200, 40)) self.curve_b = self.plot_b.plot(pen=pg.mkPen("r", width=1), fillLevel=0, brush=pg.mkBrush(200, 0, 0, 40)) def _on_start_changed(self, val): self.start = val self.sel_region.setRegion([self.start, self.start + self.duration]) self.end_label.setText(f"{self.start + self.duration:.1f}s") def _on_dur_changed(self, val): self.duration = val self.sel_region.setRegion([self.start, self.start + self.duration]) self.end_label.setText(f"{self.start + self.duration:.1f}s") def _on_region_changed(self): r = self.sel_region.getRegion() self.start = r[0] self.duration = r[1] - r[0] self.start_spin.blockSignals(True) self.start_spin.setValue(self.start) self.start_spin.blockSignals(False) self.dur_spin.blockSignals(True) self.dur_spin.setValue(self.duration) self.dur_spin.blockSignals(False) self.end_label.setText(f"{self.start + self.duration:.1f}s") def _on_stream_changed(self, idx): if self.worker and self.worker.isRunning(): self.worker.abort() self.worker.wait() self._start_loading(stream_idx=idx) def _start_loading(self, stream_idx: int = 0): if self.worker and self.worker.isRunning(): self.worker.abort() self.worker.wait() self.worker = LoadWorker(self.file_a, self.file_b, self.sr, stream_idx) self.worker.waveform_ready.connect(self._on_waveform) self.worker.thumbnail_ready.connect(self._on_thumbnail) self.worker.status.connect(lambda s: self.status_label.setText(s)) self.worker.start() def _on_waveform(self, which: str, t: np.ndarray, rms: np.ndarray): if which == "a": self.curve_a.setData(t, rms) y_max = np.max(rms) * 1.2 if len(rms) > 0 else 1 self.plot_a.setYRange(0, y_max) self.plot_a.setXRange(0, self.total_duration) self.status_label.setText("Waveform A loaded") else: self.curve_b.setData(t, rms) y_max = np.max(rms) * 1.2 if len(rms) > 0 else 1 self.plot_b.setYRange(0, y_max) self.plot_b.setXRange(0, self.total_duration) self.status_label.setText("Waveform B loaded") def _on_thumbnail(self, pixmap: QtGui.QPixmap): if not pixmap.isNull(): self.thumb_label.setPixmap(pixmap.scaled( self.thumb_label.width(), self.thumb_label.height(), QtCore.Qt.AspectRatioMode.KeepAspectRatio, QtCore.Qt.TransformationMode.SmoothTransformation )) self.status_label.setText("Ready") def _save_config(self): config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.toml") cfg = {} if os.path.exists(config_path): with open(config_path, "rb") as f: cfg = tomllib.load(f) if "extract" not in cfg: cfg["extract"] = {} cfg["extract"]["file_a"] = self.file_a cfg["extract"]["file_b"] = self.file_b cfg["extract"]["start"] = round(self.start, 3) cfg["extract"]["duration"] = round(self.duration, 1) lines = [] lines.append("# movie_sync base config — shared by extract.py and sync_audio_test.py") lines.append("# TOML is efficient: no yaml dep, stdlib tomllib (py3.11+) parses directly.\n") lines.append("[extract]") lines.append(f'file_a = {repr(cfg["extract"]["file_a"])}') lines.append(f'file_b = {repr(cfg["extract"]["file_b"])}') lines.append(f'# Window to extract (seconds). This is the ONLY duration needed.') lines.append(f'duration = {cfg["extract"]["duration"]}') lines.append(f'# Start offset in file (seconds, 0 = file start). ffmpeg -ss before -i for fast seek.') lines.append(f'start = {cfg["extract"]["start"]}') lines.append(f'out_dir = "extracted"') lines.append(f'target_fps = "24000/1001" # Blu-ray film 23.976') lines.append(f'target_fps_eps = 0.02\n') au = cfg.get("audio", {}) sy = cfg.get("sync", {}) lines.append("[audio]") lines.append(f'sample_rate = {au.get("sample_rate", 48000)}') lines.append(f'channels = {au.get("channels", 1)}') lines.append(f'normalize = {str(au.get("normalize", True)).lower()}') lines.append(f'target_lufs = {au.get("target_lufs", -14)}') lines.append(f'target_tp = {au.get("target_tp", -1.0)}') lines.append(f'target_lra = {au.get("target_lra", 11)}') lines.append(f'volume = {au.get("volume", 10)}\n') lines.append("[sync]") lines.append(f'method = {repr(sy.get("method", "phat"))}') lines.append(f'use_gpu = {repr(sy.get("use_gpu", "auto"))}') lines.append(f'out_dir = {repr(sy.get("out_dir", "output"))}') with open(config_path, "w") as f: f.write("\n".join(lines) + "\n") print(f"Saved start={self.start:.3f}s duration={self.duration:.1f}s to {config_path}") QtWidgets.QMessageBox.information(self, "Saved", f"Config updated:\nstart = {self.start:.3f}\nduration = {self.duration:.1f}") def closeEvent(self, event): if self.worker and self.worker.isRunning(): self.worker.abort() self.worker.wait(2000) event.accept() def main(): parser = argparse.ArgumentParser(description="Visual picker for extraction settings") parser.add_argument("--config", default="config.toml") parser.add_argument("--a", help="Override file_a") parser.add_argument("--b", help="Override file_b") args = parser.parse_args() cfg = {} config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.config) if os.path.exists(config_path): with open(config_path, "rb") as f: cfg = tomllib.load(f) ext = cfg.get("extract", {}) au = cfg.get("audio", {}) file_a = args.a or ext.get("file_a", "") file_b = args.b or ext.get("file_b", "") start = float(ext.get("start", 0)) duration = float(ext.get("duration", 30)) sr = int(au.get("sample_rate", 48000)) if not file_a or not os.path.exists(file_a): print(f"ERROR: file_a not found: {file_a}", file=sys.stderr) sys.exit(1) app = QtWidgets.QApplication(sys.argv) signal.signal(signal.SIGINT, signal.SIG_DFL) apply_dark_theme(app) w = PickerWindow(file_a, file_b, start, duration, sr) w.show() sys.exit(app.exec()) if __name__ == "__main__": main()