298 lines
7.7 KiB
Python
298 lines
7.7 KiB
Python
from pathlib import Path
|
|
from typing import Tuple, List, Optional
|
|
from scipy.signal import correlate
|
|
import soundfile as sf
|
|
import numpy as np
|
|
import subprocess
|
|
import json
|
|
import tempfile
|
|
import os
|
|
|
|
|
|
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 probe_video(path: Path):
|
|
"""Get video FPS, duration, codec, and stream start time."""
|
|
import json
|
|
cmd = [
|
|
"ffprobe",
|
|
"-v", "error",
|
|
"-print_format", "json",
|
|
"-show_streams",
|
|
"-show_format",
|
|
path
|
|
]
|
|
data = json.loads(run(cmd).stdout)
|
|
|
|
duration = float(data["format"]["duration"])
|
|
for s in data["streams"]:
|
|
if s["codec_type"] == "video":
|
|
num, den = s["r_frame_rate"].split("/")
|
|
fps = float(num) / float(den)
|
|
codec = s["codec_name"]
|
|
start_time = float(s.get("start_time", 0))
|
|
return fps, duration, codec, start_time
|
|
|
|
raise RuntimeError("No video stream found")
|
|
|
|
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 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 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 play_audio_segment(file: Path, intro: Tuple[float, float]) -> None:
|
|
"""
|
|
Extract and play audio from (start, duration).
|
|
"""
|
|
audio_file = extract_audio_segment(file, intro)
|
|
|
|
try:
|
|
run(
|
|
["ffplay", "-nodisp", "-autoexit", str(audio_file)]
|
|
)
|
|
finally:
|
|
os.remove(audio_file)
|
|
|
|
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 main(video: Path) -> bool:
|
|
|
|
intro_duration = 90
|
|
print("finding intro")
|
|
intro = get_intro(video, 90.0)
|
|
print("geting audio streams")
|
|
audio_streams = get_audio_streams(video)
|
|
print("extract first intro")
|
|
extract_path1 = extract_audio_segment(video, intro, audio_streams[0])
|
|
print("extract second intro")
|
|
extract_path2 = extract_audio_segment(video, intro, audio_streams[-1])
|
|
|
|
# # or directly play:
|
|
# time = (0, 20)
|
|
# play_audio_segment(extract_path1, time)
|
|
# play_audio_segment(extract_path2, time)
|
|
|
|
fps, _, _, _ = probe_video(video)
|
|
|
|
print("Calculate")
|
|
|
|
offset = get_audio_offset(extract_path1, extract_path2, intro[-1]-10)
|
|
|
|
if np.abs(offset) < (1/fps):
|
|
print(f"no offset: {offset}")
|
|
return True
|
|
else:
|
|
print(f"NOT CORRECT!!!!!!!!!!!!!!!!!!!!!!!!: {offset}, {1/fps}")
|
|
return False
|
|
|
|
extensions=("mp4", "mkv", "avi", "mov", "wmv", "flv", "webm", "lrv", "gif")
|
|
|
|
def check_video_ext(video: Path) -> bool:
|
|
if video.suffix.strip(".").lower() in extensions:
|
|
return True
|
|
return False
|
|
|
|
if __name__ == '__main__':
|
|
import sys
|
|
# sys.argv.remove(sys.argv[0])
|
|
if len(sys.argv) == 2:
|
|
print("One file")
|
|
if os.path.isfile(sys.argv[1]):
|
|
main(Path(sys.argv[1]))
|
|
else: raise
|
|
elif len(sys.argv) > 2:
|
|
sys.argv.remove(sys.argv[0])
|
|
error_files = []
|
|
for str_path in sys.argv:
|
|
video = Path(str_path)
|
|
if (check_video_ext(video)):
|
|
print(f"{video}")
|
|
working = main(video)
|
|
if not working:
|
|
error_files.append(video)
|
|
else: print(f"Not a compatabile Video: {video}")
|
|
if error_files != []:
|
|
print(f"These files are broken:")
|
|
for file in error_files:
|
|
print(file)
|
|
else:
|
|
print(f"Use this script like this:\npython detect_chapter.py /path/to/file") |