diff --git a/Makefile b/Makefile index a2cc20d..f8124b3 100644 --- a/Makefile +++ b/Makefile @@ -14,7 +14,7 @@ else PIP_INSTALL = $(PIP) install endif -.PHONY: default install install-cpu install-rocm install-cuda extract extract-loud ast sync verify clean help +.PHONY: default install install-cpu install-rocm install-cuda extract extract-loud picker ast sync verify clean help default: ast @@ -26,6 +26,7 @@ help: @echo " make install-cuda # force CUDA torch" @echo " make extract # run extract.py (uses config.toml if present)" @echo " make extract-loud # extract with loud-only gating (keep explosions/transients)" + @echo " make picker # visual picker: select start/duration before extraction" @echo " make ast # run sync_audio_test.py (audio_sync_test, PyQtGraph GUI)" @echo " make sync LAG=0.042 # run apply_sync.py to mux synced output (LAG in seconds)" @echo " make verify LAG=0.042 # verify sync: video + dual waveform player" @@ -86,6 +87,9 @@ extract: extract-loud: @if [ -x $(PY) ]; then $(PY) extract.py loud; else python3 extract.py loud; fi +picker: + @if [ -x $(PY) ]; then QT_QPA_PLATFORM=xcb QT_STYLE_OVERRIDE=Fusion $(PY) picker.py --config $(CONFIG); else QT_QPA_PLATFORM=xcb QT_STYLE_OVERRIDE=Fusion python3 picker.py --config $(CONFIG); fi + sync: @if [ -x $(PY) ]; then PYTHONPATH=. $(PY) apply_sync.py --config $(CONFIG) --lag $(LAG); else PYTHONPATH=. python3 apply_sync.py --config $(CONFIG) --lag $(LAG); fi diff --git a/picker.py b/picker.py new file mode 100644 index 0000000..cc6ba8f --- /dev/null +++ b/picker.py @@ -0,0 +1,399 @@ +#!/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 +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) + +# Reuse dark theme from verify_sync +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_thumbnails(path: str, start: float, duration: float, count: int = 20, w: int = 160) -> list[np.ndarray]: + """Extract evenly-spaced thumbnails from video segment.""" + thumbs = [] + if duration <= 0: + return thumbs + step = duration / count + for i in range(count): + t = start + i * step + try: + r = subprocess.run( + ["ffmpeg", "-y", "-ss", str(t), "-i", path, "-vframes", "1", + "-vf", f"scale={w}:-1", "-f", "rawvideo", "-pix_fmt", "gray", "-"], + capture_output=True, timeout=5 + ) + if r.returncode == 0 and len(r.stdout) > 0: + # decode gray image — we need dimensions, guess from size + # ffmpeg -vf scale=160:-1 produces width=160, height varies + # We'll decode as raw and reshape later + thumbs.append((t, r.stdout, w)) + except Exception: + pass + return thumbs + + +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 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", []) + + # Current selection + self.start = max(0.0, min(start, self.total_duration)) + self.duration = max(1.0, duration) + + self.setWindowTitle(f"picker — {os.path.basename(file_a)}") + self.resize(1400, 700) + self._setup_ui() + self._load_waveforms() + + def _setup_ui(self): + lay = QtWidgets.QVBoxLayout(self) + lay.setSpacing(4) + + # ── Video info ── + info = QtWidgets.QHBoxLayout() + lay.addLayout(info) + info.addWidget(QtWidgets.QLabel( + f"Video: {self.video_info.get('width', '?')}x{self.video_info.get('height', '?')} | " + f"{self.total_duration:.1f}s | {len(self.audio_streams)} audio stream(s)" + )) + 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(lambda _: self._load_waveforms()) + else: + self.stream_combo = None + + sel.addStretch() + + 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(120) + 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") + 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") + 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 (vertical span on both plots) + 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)) + self.curve_b = self.plot_b.plot(pen=pg.mkPen("r", width=1)) + + 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 _load_waveforms(self): + stream_idx = self.stream_combo.currentIndex() if self.stream_combo else 0 + print(f"Loading waveform A (stream {stream_idx})...") + a = extract_audio_segment(self.file_a, 0, self.total_duration, self.sr, stream_idx) + print(f"Loading waveform B...") + b = extract_audio_segment(self.file_b, 0, self.total_duration, self.sr, 0) if self.file_b and os.path.exists(self.file_b) else None + + max_pts = 200_000 + if a is not None and len(a) > 0: + t_a = np.arange(len(a), dtype=np.float64) / self.sr + if len(a) > max_pts: + stride = int(math.ceil(len(a) / max_pts)) + a = a[::stride] + t_a = t_a[::stride] + self.curve_a.setData(t_a, a) + y_max = np.max(np.abs(a)) * 1.1 + self.plot_a.setYRange(-y_max, y_max) + + if b is not None and len(b) > 0: + t_b = np.arange(len(b), dtype=np.float64) / self.sr + if len(b) > max_pts: + stride = int(math.ceil(len(b) / max_pts)) + b = b[::stride] + t_b = t_b[::stride] + self.curve_b.setData(t_b, b) + y_max_b = np.max(np.abs(b)) * 1.1 + self.plot_b.setYRange(-y_max_b, y_max_b) + + self.plot_a.setXRange(0, self.total_duration) + self.plot_b.setXRange(0, self.total_duration) + + # Load thumbnails in background + self._load_thumbnails() + + def _load_thumbnails(self): + print("Extracting thumbnails...") + thumbs = extract_thumbnails(self.file_a, self.start, self.duration, count=20, w=160) + if not thumbs: + self.thumb_label.setText("Thumbnails unavailable") + return + + # Stitch thumbnails horizontally into a pixmap + imgs = [] + for t, raw, w in thumbs: + # Decode raw gray — need height. We'll use ffmpeg to produce png instead + pass + + # Simpler: extract a sprite sheet with ffmpeg + step = self.duration / 20 if self.duration > 0 else 1 + try: + r = subprocess.run( + ["ffmpeg", "-y", "-ss", str(self.start), "-i", self.file_a, "-t", str(self.duration), + "-vf", f"fps=1/{max(0.1, step):.2f},scale=160:-1,tile=20x1"], + capture_output=True, timeout=30 + ) + if r.returncode == 0: + # Write to temp file and load as pixmap + import tempfile + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: + tmp.write(r.stdout) + tmp_path = tmp.name + pixmap = QtGui.QPixmap(tmp_path) + self.thumb_label.setPixmap(pixmap.scaled( + self.thumb_label.width(), self.thumb_label.height(), + QtCore.Qt.AspectRatioMode.KeepAspectRatio, + QtCore.Qt.TransformationMode.SmoothTransformation + )) + os.unlink(tmp_path) + except Exception as e: + self.thumb_label.setText(f"Thumbnail error: {e}") + + def _save_config(self): + config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.toml") + # Read existing config + cfg = {} + if os.path.exists(config_path): + with open(config_path, "rb") as f: + cfg = tomllib.load(f) + + # Update extract section + 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) + + # Write back + import io + 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') + + # Copy audio and sync sections from existing if present + 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 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()