honney/movie_sync
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: background waveform/thumbnail loading with RMS envelope

Browse Source
This commit is contained in:
Honney OC committed 2026-09-15 18:38:59 +00:00
1 parent 4147310706
commit bc5bbb5e23
1 file changed
+164 -99
Download Patch File Download Diff File Expand all files Collapse all files
picker.py
+164 -99
View File
@@ -7,7 +7,7 @@ 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 os, sys, argparse, subprocess, json, math, signal, struct
import numpy as np
try:
@@ -22,7 +22,6 @@ 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
@@ -53,28 +52,45 @@ def probe_video_info(path: str) -> dict:
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
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:
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))
# This doesn't work well — use a different approach: extract raw PCM but in chunks
pass
except Exception:
pass
return thumbs
# 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,
@@ -94,6 +110,76 @@ def extract_audio_segment(path: str, start: float, duration: float, sr: int = 48
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__()
@@ -104,14 +190,15 @@ class PickerWindow(QtWidgets.QWidget):
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.worker = None
self.setWindowTitle(f"picker — {os.path.basename(file_a)}")
self.resize(1400, 700)
self._setup_ui()
self._load_waveforms()
self._start_loading()
def _setup_ui(self):
lay = QtWidgets.QVBoxLayout(self)
@@ -120,10 +207,12 @@ class PickerWindow(QtWidgets.QWidget):
# ── Video info ──
info = QtWidgets.QHBoxLayout()
lay.addLayout(info)
info.addWidget(QtWidgets.QLabel(
self.info_label = 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)"
))
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 ──
@@ -168,12 +257,16 @@ class PickerWindow(QtWidgets.QWidget):
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())
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)
@@ -181,36 +274,39 @@ class PickerWindow(QtWidgets.QWidget):
# ── Plots ──
# Row 0: Thumbnails
self.thumb_label = QtWidgets.QLabel("Thumbnails (loading...)")
self.thumb_label.setMinimumHeight(120)
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")
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")
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 (vertical span on both plots)
# 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))
self.curve_b = self.plot_b.plot(pen=pg.mkPen("r", width=1))
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
@@ -234,86 +330,52 @@ class PickerWindow(QtWidgets.QWidget):
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
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)
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)
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")
# 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)
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
))
os.unlink(tmp_path)
except Exception as e:
self.thumb_label.setText(f"Thumbnail error: {e}")
self.status_label.setText("Ready")
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
@@ -321,8 +383,6 @@ class PickerWindow(QtWidgets.QWidget):
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")
@@ -337,7 +397,6 @@ class PickerWindow(QtWidgets.QWidget):
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]")
@@ -360,6 +419,12 @@ class PickerWindow(QtWidgets.QWidget):
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")