Added tests and more

check_final_file.py

@@ -0,0 +1,210 @@
+from pathlib import Path
+from typing import Tuple, List, Optional
+import json
+import subprocess
+import tempfile
+import numpy as np
+import soundfile as sf
+from scipy.signal import correlate
+
+def run(cmd):
+    """
+    Run a subprocess command.
+
+    Args:
+        cmd: Command and arguments as a list.
+
+    Raises:
+        subprocess.CalledProcessError: If the command exits with a non-zero status.
+    """
+    try:
+        result = subprocess.run(
+            cmd,
+            check=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True
+        )
+        return result
+    except subprocess.CalledProcessError as e:
+        print(e.stdout)
+        raise
+
+def _get_chapters(file: Path) -> List[Tuple[float, float]]:
+    """Return list of (start_time, duration) for each chapter."""
+    cmd = [
+        "ffprobe",
+        "-i", str(file),
+        "-print_format", "json",
+        "-show_chapters",
+        "-loglevel", "error"
+    ]
+
+    result = run(cmd)
+    data = json.loads(result.stdout)
+
+    chapters = []
+    for ch in data.get("chapters", []):
+        start = float(ch["start_time"])
+        end = float(ch["end_time"])
+        duration = end - start
+        chapters.append((start, duration))
+
+    return chapters
+
+
+def get_intro(file: Path, time: float) -> Tuple[float, float]:
+    """
+    Find the chapter whose duration is closest to `time`,
+    considering only chapters with duration between time/2 and time*2.
+    Returns (start_time, duration).
+    """
+    chapters = _get_chapters(file)
+
+    if not chapters:
+        raise ValueError("No chapters found")
+
+    # Filter by allowed duration range
+    min_len = time / 2
+    max_len = time * 2
+    filtered = [
+        (start, dur) for start, dur in chapters
+        if min_len <= dur <= max_len
+    ]
+
+    if not filtered:
+        raise ValueError("No chapters match duration constraints")
+
+    # Sort by closeness to target time
+    filtered.sort(key=lambda x: abs(x[1] - time))
+
+    # Take the two closest
+    closest_two = filtered[:2]
+
+    # Choose the better one (closest duration)
+    best = min(closest_two, key=lambda x: abs(x[1] - time))
+
+    return best
+
+def get_audio_streams(file: Path) -> List[int]:
+    """
+    Return a list of available audio stream indices (0-based for ffmpeg mapping).
+    
+    Example return:
+        [0, 1]  -> corresponds to 0:a:0 and 0:a:1
+    """
+    cmd = [
+        "ffprobe",
+        "-i", str(file),
+        "-select_streams", "a",
+        "-show_entries", "stream=index",
+        "-of", "json",
+        "-loglevel", "error"
+    ]
+
+    result = run(cmd)
+    data = json.loads(result.stdout)
+
+    # ffprobe gives absolute stream indices, but we want relative audio indices
+    streams = data.get("streams", [])
+
+    # Convert to 0-based audio stream numbering
+    return list(range(len(streams)))
+
+def extract_audio_segment(
+    file: Path,
+    intro: Tuple[float, float],
+    audio_stream: Optional[int] = None
+) -> Path:
+    """
+    Extract audio using (start, duration) from get_intro.
+
+    Parameters:
+        file: input video file
+        intro: (start, duration)
+        audio_stream: index of audio stream (0 = first audio track, etc.)
+
+    Returns:
+        Path to temporary audio file
+    """
+    start, duration = intro
+
+    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    tmp_path = Path(tmp.name)
+    tmp.close()
+
+    cmd = [
+        "ffmpeg",
+        "-y",
+        "-i", str(file),
+        "-ss", str(start),
+        "-t", str(duration),
+    ]
+
+    # Select specific audio stream if provided
+    if audio_stream is not None:
+        cmd += ["-map", f"0:a:{audio_stream}"]
+    else:
+        cmd += ["-map", "0:a:0"]  # default: first audio track
+
+    cmd += [
+        "-vn",
+        "-acodec", "pcm_s16le",
+        "-ar", "44100",
+        "-ac", "2",
+        str(tmp_path)
+    ]
+
+    run(cmd)
+    return tmp_path
+
+def extract_audio(file: Path, duration: float):
+    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+
+    run([
+        "ffmpeg",
+        "-y",
+        "-i", str(file),
+        "-map", "0:a:0",        # first audio stream
+        "-ac", "1",             # mono
+        "-ar", "44100",         # fixed sample rate
+        "-t", str(duration),    # limit duration
+        tmp.name
+    ])
+
+    data, sr = sf.read(tmp.name)
+    return data, sr
+
+def get_audio_offset(file1: Path, file2: Path, important_time: float):
+    sig1, sr1 = extract_audio(file1, important_time)
+    sig2, sr2 = extract_audio(file2, important_time)
+
+    if sr1 != sr2:
+        raise ValueError("Sample rates do not match")
+
+    # normalize signals
+    sig1 = sig1 - np.mean(sig1)
+    sig2 = sig2 - np.mean(sig2)
+
+    corr = correlate(sig1, sig2, mode="full")
+    lag = np.argmax(corr) - (len(sig2) - 1)
+
+    offset_seconds = lag / sr1
+    return offset_seconds
+
+def check_final_file(video: Path, fps: float) -> bool:
+    try:
+        intro_duration = 90
+        intro = get_intro(video, 90.0)
+        audio_streams = get_audio_streams(video)
+        extract_path1 = extract_audio_segment(video, intro, audio_streams[0])
+        extract_path2 = extract_audio_segment(video, intro, audio_streams[-1])
+        offset = get_audio_offset(extract_path1, extract_path2, intro[-1])
+        if np.abs(offset) < (1/fps):
+            return True
+        else:
+            print(f"NOT CORRECT!!!!!!!!!!!!!!!!!!!!!!!!: {offset}, {1/fps}")
+            return False
+    except:
+        print("Chapter detection not working")
+        return False