465 lines
18 KiB
Python
465 lines
18 KiB
Python
#!/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()
|