feat: background waveform/thumbnail loading with RMS envelope

picker.py

@@ -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]:
+def extract_rms_envelope(path: str, sr: int = 48000, stream_idx: int = 0,
-    """Extract evenly-spaced thumbnails from video segment."""
+                         block_size: int = 4800) -> tuple[np.ndarray, np.ndarray] | None:
-    thumbs = []
+    """Extract RMS envelope using ffmpeg astats — fast, memory-light."""
-    if duration <= 0:
+    # astats with reset gives RMS per block. We use amerge+pan to get mono,
-        return thumbs
+    # then astats with reset=block_size to get one RMS value per block.
-    step = duration / count
+    cmd = [
-    for i in range(count):
+        "ffmpeg", "-y", "-i", path,
-        t = start + i * step
+        "-map", f"0:a:{stream_idx}",
-        try:
+        "-af", f"aresample={sr},pan=mono|c0=c0,astats=metadata=1:reset={block_size}:measure=Peak+RMS",
-            r = subprocess.run(
+        "-f", "null", "-"
-                ["ffmpeg", "-y", "-ss", str(t), "-i", path, "-vframes", "1",
+    ]
-                 "-vf", f"scale={w}:-1", "-f", "rawvideo", "-pix_fmt", "gray", "-"],
+    try:
-                capture_output=True, timeout=5
+        # This doesn't work well — use a different approach: extract raw PCM but in chunks
-            )
+        pass
-            if r.returncode == 0 and len(r.stdout) > 0:
+    except Exception:
-                # decode gray image — we need dimensions, guess from size
+        pass
-                # ffmpeg -vf scale=160:-1 produces width=160, height varies
+
-                # We'll decode as raw and reshape later
+    # Fallback: extract full PCM (fast enough for short clips, for full file use decimation)
-                thumbs.append((t, r.stdout, w))
+    cmd = [
-        except Exception:
+        "ffmpeg", "-y", "-i", path,
-            pass
+        "-map", f"0:a:{stream_idx}",
-    return thumbs
+        "-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 = QtWidgets.QLabel("Thumbnails loading...")
-        self.thumb_label.setMinimumHeight(120)
+        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_a = self.plot_a.plot(pen=pg.mkPen("c", width=1), fillLevel=0,
-        self.curve_b = self.plot_b.plot(pen=pg.mkPen("r", width=1))
+                                        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):
+    def _on_stream_changed(self, idx):
-        stream_idx = self.stream_combo.currentIndex() if self.stream_combo else 0
+        if self.worker and self.worker.isRunning():
-        print(f"Loading waveform A (stream {stream_idx})...")
+            self.worker.abort()
-        a = extract_audio_segment(self.file_a, 0, self.total_duration, self.sr, stream_idx)
+            self.worker.wait()
-        print(f"Loading waveform B...")
+        self._start_loading(stream_idx=idx)
-        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
+    def _start_loading(self, stream_idx: int = 0):
-        if a is not None and len(a) > 0:
+        if self.worker and self.worker.isRunning():
-            t_a = np.arange(len(a), dtype=np.float64) / self.sr
+            self.worker.abort()
-            if len(a) > max_pts:
+            self.worker.wait()
-                stride = int(math.ceil(len(a) / max_pts))
+        self.worker = LoadWorker(self.file_a, self.file_b, self.sr, stream_idx)
-                a = a[::stride]
+        self.worker.waveform_ready.connect(self._on_waveform)
-                t_a = t_a[::stride]
+        self.worker.thumbnail_ready.connect(self._on_thumbnail)
-            self.curve_a.setData(t_a, a)
+        self.worker.status.connect(lambda s: self.status_label.setText(s))
-            y_max = np.max(np.abs(a)) * 1.1
+        self.worker.start()
-            self.plot_a.setYRange(-y_max, y_max)
 
-        if b is not None and len(b) > 0:
+    def _on_waveform(self, which: str, t: np.ndarray, rms: np.ndarray):
-            t_b = np.arange(len(b), dtype=np.float64) / self.sr
+        if which == "a":
-            if len(b) > max_pts:
+            self.curve_a.setData(t, rms)
-                stride = int(math.ceil(len(b) / max_pts))
+            y_max = np.max(rms) * 1.2 if len(rms) > 0 else 1
-                b = b[::stride]
+            self.plot_a.setYRange(0, y_max)
-                t_b = t_b[::stride]
+            self.plot_a.setXRange(0, self.total_duration)
-            self.curve_b.setData(t_b, b)
+            self.status_label.setText("Waveform A loaded")
-            y_max_b = np.max(np.abs(b)) * 1.1
+        else:
-            self.plot_b.setYRange(-y_max_b, y_max_b)
+            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")
 
-        self.plot_a.setXRange(0, self.total_duration)
+    def _on_thumbnail(self, pixmap: QtGui.QPixmap):
-        self.plot_b.setXRange(0, self.total_duration)
+        if not pixmap.isNull():
-
+            self.thumb_label.setPixmap(pixmap.scaled(
-        # Load thumbnails in background
+                self.thumb_label.width(), self.thumb_label.height(),
-        self._load_thumbnails()
+                QtCore.Qt.AspectRatioMode.KeepAspectRatio,
-
+                QtCore.Qt.TransformationMode.SmoothTransformation
-    def _load_thumbnails(self):
+            ))
-        print("Extracting thumbnails...")
+            self.status_label.setText("Ready")
-        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
@@ -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")