185 lines
6.3 KiB
Python
185 lines
6.3 KiB
Python
import subprocess
|
|
import tempfile
|
|
import os
|
|
from pathlib import Path
|
|
from typing import List, Tuple, Optional
|
|
|
|
import numpy as np
|
|
from scipy.io import wavfile
|
|
|
|
class ConfidenceError(Exception):
|
|
"""Raised when an area is found but the energy difference is too low to be confident."""
|
|
pass
|
|
|
|
class MissingDetectionError(Exception):
|
|
"""Raised when the script fails to find both an intro and an outro."""
|
|
pass
|
|
|
|
def get_video_duration(file: Path) -> float:
|
|
"""Uses ffprobe to get the total duration of the video."""
|
|
cmd = [
|
|
"ffprobe",
|
|
"-v", "error",
|
|
"-show_entries", "format=duration",
|
|
"-of", "default=noprint_wrappers=1:nokey=1",
|
|
str(file)
|
|
]
|
|
result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
|
try:
|
|
return float(result.stdout.strip())
|
|
except ValueError:
|
|
raise RuntimeError(f"Could not determine duration for {file}")
|
|
|
|
def extract_audio_segment_tmp(
|
|
file: Path,
|
|
start_or_section: float | tuple[float, float],
|
|
end: Optional[float] = None,
|
|
audio_stream: Optional[int] = None
|
|
) -> Path:
|
|
"""Extracts audio using FFmpeg and saves it to a temporary WAV file."""
|
|
if isinstance(start_or_section, tuple):
|
|
if len(start_or_section) != 2:
|
|
raise ValueError("Tuple must strictly be (start, end)")
|
|
start = float(start_or_section[0])
|
|
duration = float(start_or_section[1])
|
|
elif isinstance(start_or_section, (int, float)):
|
|
if end is None:
|
|
raise ValueError("The 'end' parameter must be provided when using a separate start float.")
|
|
start = float(start_or_section)
|
|
end_time = float(end)
|
|
duration = end_time - start
|
|
else:
|
|
raise TypeError("Invalid type for start_or_section. Expected float or tuple.")
|
|
|
|
if duration <= 0:
|
|
raise ValueError("Duration must be above 0")
|
|
|
|
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)
|
|
]
|
|
|
|
if audio_stream is not None:
|
|
cmd += ["-map", f"0:a:{audio_stream}"]
|
|
else:
|
|
cmd += ["-map", "0:a:0"]
|
|
|
|
cmd += [
|
|
"-vn",
|
|
"-acodec", "pcm_s16le",
|
|
"-ar", "44100",
|
|
"-ac", "2",
|
|
str(tmp_path)
|
|
]
|
|
|
|
subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=True)
|
|
return tmp_path
|
|
|
|
def find_high_power_segment(wav_path: Path, expected_length: float, confidence_threshold: float = 1.0) -> float:
|
|
rate, data = wavfile.read(wav_path)
|
|
|
|
if len(data.shape) > 1:
|
|
data = data.mean(axis=1)
|
|
|
|
data = data.astype(np.float32) / (np.max(np.abs(data)) + 1e-9)
|
|
|
|
chunk_size = rate
|
|
num_chunks = len(data) // chunk_size
|
|
|
|
# Calculate RMS (Volume) and Standard Deviation (Variation) for every second
|
|
powers = np.zeros(num_chunks)
|
|
variations = np.zeros(num_chunks)
|
|
|
|
for i in range(num_chunks):
|
|
chunk = data[i * chunk_size : (i + 1) * chunk_size]
|
|
powers[i] = np.sqrt(np.mean(chunk**2))
|
|
variations[i] = np.std(chunk)
|
|
|
|
window_size = int(expected_length)
|
|
if len(powers) < window_size:
|
|
raise ValueError("Audio segment too short.")
|
|
|
|
# SCORING ALGORITHM:
|
|
# We want High Power and Low Variation (that "solid block" look)
|
|
# We invert variation so that a small variation results in a higher score.
|
|
scores = []
|
|
for i in range(len(powers) - window_size + 1):
|
|
win_power = np.mean(powers[i : i + window_size])
|
|
win_var = np.mean(variations[i : i + window_size])
|
|
|
|
# The score increases with power and decreases as the variation increases
|
|
# We add a small constant to win_var to avoid division by zero
|
|
score = win_power / (win_var + 0.05)
|
|
scores.append(score)
|
|
|
|
best_idx = np.argmax(scores)
|
|
|
|
# Confidence check: is our best score significantly better than the average score?
|
|
avg_score = np.mean(scores)
|
|
if scores[best_idx] < avg_score * confidence_threshold:
|
|
raise ConfidenceError("Detection unreliable: no distinct high-power constant block found.")
|
|
|
|
return float(best_idx)
|
|
|
|
def get_intro_outro(file: Path, length: float = 90.0) -> List[Tuple[float, float]]:
|
|
"""
|
|
Finds the intro and outro of a video file by analyzing audio power peaks.
|
|
Searches the first 8 minutes and the last 8 minutes.
|
|
"""
|
|
results = []
|
|
total_duration = get_video_duration(file)
|
|
search_window = 8 * 60.0 # 8 minutes
|
|
|
|
# --- 1. FIND INTRO ---
|
|
intro_search_dur = min(search_window, total_duration)
|
|
intro_wav = extract_audio_segment_tmp(file, 0.0, intro_search_dur)
|
|
|
|
try:
|
|
intro_start_rel = find_high_power_segment(intro_wav, length)
|
|
results.append((intro_start_rel, length))
|
|
except ConfidenceError:
|
|
intro_wav.unlink()
|
|
raise ConfidenceError("Could not confidently detect the Intro.")
|
|
finally:
|
|
if intro_wav.exists():
|
|
intro_wav.unlink() # Cleanup temp file
|
|
|
|
# --- 2. FIND OUTRO ---
|
|
outro_search_start = max(0.0, total_duration - search_window)
|
|
outro_search_dur = total_duration - outro_search_start
|
|
outro_wav = extract_audio_segment_tmp(file, outro_search_start, total_duration)
|
|
|
|
try:
|
|
outro_start_rel = find_high_power_segment(outro_wav, length)
|
|
# Add the 8-minute offset back to get the absolute timestamp in the video
|
|
outro_start_abs = outro_search_start + outro_start_rel
|
|
results.append((outro_start_abs, length))
|
|
except ConfidenceError:
|
|
outro_wav.unlink()
|
|
raise ConfidenceError("Could not confidently detect the Outro.")
|
|
finally:
|
|
if outro_wav.exists():
|
|
outro_wav.unlink() # Cleanup temp file
|
|
|
|
# Ensure we got both
|
|
if len(results) < 2:
|
|
raise MissingDetectionError("Failed to find both an intro and an outro.")
|
|
|
|
return results
|
|
|
|
# --- Example Usage ---
|
|
if __name__ == "__main__":
|
|
try:
|
|
# Requires: pip install numpy scipy
|
|
file_path = Path("/home/honney/mount/Media/tmp/serien/ready/Charlotte/Season 01/[NO CHAPTER]_E04.mkv")
|
|
if file_path.exists():
|
|
timestamps = get_intro_outro(file_path, length=90.0)
|
|
print("Found Intro and Outro [Start, Duration]:", timestamps)
|
|
except Exception as e:
|
|
print(f"Error: {e}") |