Merge-Anime-Lanuages/check_final_file.py

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