WeAreSyncing/clap_detection.py

"""Sharp audio onset (clap/snap/beep) detection, shared by the GUI and
scripts/detect_claps.py's CLI. Ported from that standalone script so both
paths use one implementation instead of drifting apart.
"""
import os
import subprocess
import tempfile
import warnings
from dataclasses import dataclass
from typing import List, Optional

try:
    import librosa
except ImportError:
    librosa = None


class ClapDetectionUnavailable(RuntimeError):
    pass


@dataclass(frozen=True)
class ClapEvent:
    time_sec: float
    strength: float


def _probe_audio_start_time(input_path: str, stream_specifier: str) -> Optional[float]:
    result = subprocess.run(
        [
            "ffprobe", "-v", "error", "-select_streams", stream_specifier,
            "-show_entries", "stream=start_time",
            "-of", "default=noprint_wrappers=1:nokey=1",
            input_path,
        ],
        capture_output=True, text=True,
    )
    try:
        return float(result.stdout.strip())
    except ValueError:
        # Some containers/codecs legitimately report "N/A" (genuinely no
        # start offset to speak of), but this could also mean ffprobe just
        # couldn't express a real nonzero offset for this container -- we
        # can't tell those apart here, so surface it instead of silently
        # assuming 0, which would leave onset times subtly, silently wrong.
        warnings.warn(
            f"could not determine audio start_time for {input_path!r} "
            f"(ffprobe returned {result.stdout.strip()!r}); assuming 0.0, "
            "which may be wrong for this file's container/codec."
        )
        return None


def _extract_audio_wav(input_path: str, audio_stream_index: Optional[int] = None) -> str:
    # The audio stream's own start_time is exactly how a container encodes
    # "this track starts playing N seconds later than the others" -- i.e.
    # it's the very thing this whole tool exists to fix. ffmpeg's default
    # `-vn` extraction drops it and renormalizes the output to start at 0,
    # which would silently make onset detection blind to that offset.
    # Compensate by inserting the missing lead-in as real silence, so times
    # measured from this WAV's t=0 still match the file's true timeline.
    #
    # audio_stream_index (ffprobe/ffmpeg global stream index, e.g. mpv's
    # own `ff-index` for the track it's currently playing) makes stream
    # selection explicit and -- critically -- lets the caller guarantee
    # this matches whichever audio track is actually audible in the
    # player. Without it, ffmpeg's own default-stream heuristic picks
    # *a* audio track, which isn't guaranteed to be the same one on
    # multi-audio-track files (extra language dubs, a separate mic/system
    # audio track from a screen recorder, etc.) -- something our own
    # single-track synthetic test fixtures never exercise, but real-world
    # files often have.
    stream_specifier = f"{audio_stream_index}" if audio_stream_index is not None else "a:0"
    audio_start = _probe_audio_start_time(input_path, stream_specifier) or 0.0

    fd, wav_path = tempfile.mkstemp(suffix=".wav")
    os.close(fd)
    cmd = ["ffmpeg", "-y", "-i", input_path, "-vn"]
    cmd += ["-map", f"0:{audio_stream_index}" if audio_stream_index is not None else "a:0"]
    # aresample=async=1 keeps the decoded audio conformed to its own
    # internal PTS, filling any gap with silence instead of silently
    # closing it up. Long recordings (stream captures, OBS-style long
    # sessions) can have a genuine multi-second discontinuity mid-file --
    # e.g. a brief audio-device dropout -- from a real capture glitch, not
    # anything to do with sync between tracks. Plain -vn extraction just
    # concatenates packets and ignores the gap, silently shortening the
    # output and shifting every onset after that point earlier by however
    # long the gap was. Confirmed on a real ~110min capture with exactly
    # one such gap (+1.621s at ~15.5min in): without this filter the
    # extracted WAV came out 1.6s short of the source's true duration;
    # with it, the duration matches exactly. This is a different failure
    # mode from the leading start_time offset below (which aresample alone
    # does NOT fix -- it re-zeros to its own first frame, same as plain
    # extraction), so both are needed together, not either/or.
    af_chain = ["aresample=async=1"]
    if audio_start > 0.001:
        delay_ms = int(round(audio_start * 1000))
        af_chain.append(f"adelay={delay_ms}|{delay_ms}")
    cmd += ["-af", ",".join(af_chain)]
    cmd += ["-acodec", "pcm_s16le", "-ar", "22050", "-ac", "1", wav_path]

    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        os.unlink(wav_path)
        raise RuntimeError(f"ffmpeg audio extraction failed: {result.stderr}")
    return wav_path


def detect_claps(
    input_path: str,
    start_offset: float = 0.0,
    threshold: float = 0.7,
    min_interval: float = 0.3,
    max_results: Optional[int] = None,
    audio_stream_index: Optional[int] = None,
) -> List[ClapEvent]:
    """Detect sharp audio onsets in input_path. `start_offset`: only onsets
    at or after this absolute timestamp are returned (for a from-cursor
    search). `max_results` caps how many are returned, for a "find the next
    one" search instead of a full scan. `audio_stream_index` (an ffmpeg
    global stream index) pins detection to a specific audio track on
    multi-audio-track files -- pass the ff-index of whichever track mpv is
    actually playing (see main.py) so detection can't silently analyze a
    different track than the one being listened to.

    The whole file is always decoded and analyzed regardless of
    start_offset -- an earlier version tried to skip ahead with ffmpeg's
    `-ss` for a from-cursor scan, but combining that with the start_time
    compensation above (both are timestamp-origin shifts on the same file)
    was too easy to get subtly wrong, and getting a "which timestamp is
    correct" bug wrong silently is the one thing this tool cannot afford.
    Filtering the results in Python after a full decode is slower but
    unambiguous. Audio-only onset detection is fast enough in practice that
    this hasn't been a real cost.
    """
    if librosa is None:
        raise ClapDetectionUnavailable(
            "librosa is not installed. Install with: pip install librosa numpy soundfile"
        )

    wav_path = _extract_audio_wav(input_path, audio_stream_index=audio_stream_index)
    try:
        y, sr = librosa.load(wav_path, sr=22050)
        onset_env = librosa.onset.onset_strength(y=y, sr=sr)
        onset_frames = librosa.onset.onset_detect(
            onset_envelope=onset_env,
            sr=sr,
            units="frames",
            backtrack=False,
            delta=threshold,
            wait=int(min_interval * sr / 512),
        )
        onset_times = librosa.frames_to_time(onset_frames, sr=sr)
        strengths = onset_env[onset_frames]

        all_events: List[ClapEvent] = []
        last_time = -min_interval
        for t, s in zip(onset_times, strengths):
            t = float(t)
            if t - last_time < min_interval:
                continue
            all_events.append(ClapEvent(time_sec=t, strength=float(s)))
            last_time = t

        events = [e for e in all_events if e.time_sec >= start_offset]
        if max_results is not None:
            events = events[:max_results]
        return events
    finally:
        os.unlink(wav_path)