Added tests and more
This commit is contained in:
Hannes committed
1 parent
295177be6f
commit
a7a7abf720
9 files changed
+1327
No files matched your search
@@ -0,0 +1,298 @@
|
||||
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")
|
||||
Reference in new issue
Block a user