218 lines
7.9 KiB
Python
218 lines
7.9 KiB
Python
import os
|
||
import subprocess
|
||
import glob
|
||
import sys
|
||
import json
|
||
import shlex
|
||
|
||
# ---------- ffprobe helpers ----------
|
||
|
||
def _run_probe(args):
|
||
return json.loads(
|
||
subprocess.run(args, capture_output=True, text=True, check=True).stdout
|
||
)
|
||
|
||
def get_media_info(file_path):
|
||
"""
|
||
Detect container-agnostic media properties required to reproduce the same
|
||
encoding format for an intro clip that can be safely appended.
|
||
"""
|
||
# Video stream (include codec, pix_fmt, color info, fps as *string* rational)
|
||
vprobe = _run_probe([
|
||
"ffprobe","-v","error",
|
||
"-select_streams","v:0",
|
||
"-show_entries","stream=codec_name,pix_fmt,width,height,r_frame_rate,color_range,color_space,color_transfer,color_primaries",
|
||
"-of","json", file_path
|
||
])
|
||
v = vprobe["streams"][0]
|
||
width = int(v["width"])
|
||
height = int(v["height"])
|
||
fps_rational = v.get("r_frame_rate","24000/1001") # keep rational string for exact match
|
||
vcodec = v.get("codec_name","h264")
|
||
pix_fmt = v.get("pix_fmt","yuv420p")
|
||
color_range = v.get("color_range") # "tv" or "pc" typically
|
||
color_space = v.get("color_space") # e.g. "bt709", "bt2020nc"
|
||
color_trc = v.get("color_transfer") # e.g. "smpte2084", "bt709"
|
||
color_primaries = v.get("color_primaries") # e.g. "bt2020", "bt709"
|
||
|
||
# Audio stream (channels, sample_rate, codec, layout if present)
|
||
# Some inputs have no audio; handle gracefully.
|
||
try:
|
||
aprobe = _run_probe([
|
||
"ffprobe","-v","error",
|
||
"-select_streams","a:0",
|
||
"-show_entries","stream=codec_name,channels,sample_rate,channel_layout",
|
||
"-of","json", file_path
|
||
])
|
||
a = aprobe["streams"][0]
|
||
achannels = int(a.get("channels", 2))
|
||
arate = int(a.get("sample_rate", 48000))
|
||
acodec = a.get("codec_name", "aac")
|
||
alayout = a.get("channel_layout") # may be None
|
||
has_audio = True
|
||
except Exception:
|
||
achannels = 0
|
||
arate = 0
|
||
acodec = None
|
||
alayout = None
|
||
has_audio = False
|
||
|
||
return {
|
||
"width": width, "height": height, "fps_rational": fps_rational,
|
||
"vcodec": vcodec, "pix_fmt": pix_fmt,
|
||
"color_range": color_range, "color_space": color_space,
|
||
"color_trc": color_trc, "color_primaries": color_primaries,
|
||
"has_audio": has_audio, "achannels": achannels,
|
||
"arate": arate, "acodec": acodec, "alayout": alayout
|
||
}
|
||
|
||
# ---------- encoder selection ----------
|
||
|
||
VIDEO_ENCODER_MAP = {
|
||
"av1": ("libaom-av1", ["-crf","30","-b:v","0","-cpu-used","6"]),
|
||
"hevc": ("libx265", ["-crf","20","-preset","medium"]),
|
||
"h264": ("libx264", ["-crf","20","-preset","medium"]),
|
||
"vp9": ("libvpx-vp9", ["-crf","32","-b:v","0","-row-mt","1"]),
|
||
"mpeg2video": ("mpeg2video", []),
|
||
"mpeg4": ("mpeg4", []),
|
||
# fall back handled below
|
||
}
|
||
|
||
AUDIO_ENCODER_MAP = {
|
||
"aac": ("aac", ["-b:a","192k"]),
|
||
"opus": ("libopus", ["-b:a","160k"]),
|
||
"ac3": ("ac3", ["-b:a","384k"]),
|
||
"eac3": ("eac3", ["-b:a","384k"]),
|
||
"flac": ("flac", []),
|
||
"vorbis":("libvorbis", ["-q:a","4"]),
|
||
"mp3": ("libmp3lame", ["-b:a","192k"]),
|
||
"pcm_s16le": ("pcm_s16le", []),
|
||
# more can be added as needed
|
||
}
|
||
|
||
def pick_video_encoder(vcodec):
|
||
enc, extra = VIDEO_ENCODER_MAP.get(vcodec, ("libx264", ["-crf","20","-preset","medium"]))
|
||
return enc, extra
|
||
|
||
def pick_audio_encoder(acodec):
|
||
if acodec is None:
|
||
return None, [] # no audio in input
|
||
enc, extra = AUDIO_ENCODER_MAP.get(acodec, ("aac", ["-b:a","192k"]))
|
||
return enc, extra
|
||
|
||
# ---------- intro generation ----------
|
||
|
||
def build_intro_cmd(info, duration, intro_file):
|
||
"""
|
||
Build an ffmpeg command that creates a silent (black) intro with *matching*
|
||
A/V format to the input.
|
||
"""
|
||
w, h = info["width"], info["height"]
|
||
fps = info["fps_rational"] # keep as rational string
|
||
venc, venc_extra = pick_video_encoder(info["vcodec"])
|
||
pix_fmt = info["pix_fmt"]
|
||
|
||
cmd = [
|
||
"ffmpeg",
|
||
"-f","lavfi","-i", f"color=black:s={w}x{h}:r={fps}:d={duration}",
|
||
]
|
||
|
||
|
||
# Audio input (only if input has audio; otherwise omit to keep track counts aligned)
|
||
if info["has_audio"]:
|
||
ch = info["achannels"]
|
||
sr = info["arate"]
|
||
layout = info["alayout"] or ( "stereo" if ch == 2 else f"{ch}c" )
|
||
cmd += ["-f","lavfi","-i", f"anullsrc=channel_layout={layout}:sample_rate={sr}"]
|
||
else:
|
||
# still add a null audio to ensure mkv track counts match? no — input had no audio.
|
||
pass
|
||
|
||
# Map streams explicitly
|
||
# If there is audio in the input, produce both v+a; else only v.
|
||
# (FFmpeg will auto-create stream 0:v:0 and 1:a:0 if two inputs present.)
|
||
# Codec settings
|
||
cmd += ["-c:v", venc, "-pix_fmt", pix_fmt]
|
||
cmd += venc_extra
|
||
|
||
# Preserve basic color tags when present (helps SDR/HDR tagging)
|
||
if info["color_primaries"]:
|
||
cmd += ["-color_primaries", info["color_primaries"]]
|
||
if info["color_trc"]:
|
||
cmd += ["-color_trc", info["color_trc"]]
|
||
if info["color_space"]:
|
||
cmd += ["-colorspace", info["color_space"]]
|
||
if info["color_range"]:
|
||
cmd += ["-color_range", info["color_range"]] # "pc" or "tv"
|
||
|
||
if info["has_audio"]:
|
||
aenc, aenc_extra = pick_audio_encoder(info["acodec"])
|
||
cmd += ["-c:a", aenc] + aenc_extra
|
||
# Ensure channels & rate match exactly
|
||
cmd += ["-ac", str(info["achannels"]), "-ar", str(info["arate"])]
|
||
|
||
# Keep duration tight and avoid stuck encodes
|
||
cmd += ["-shortest", "-y", intro_file]
|
||
return cmd
|
||
|
||
# ---------- main pipeline ----------
|
||
|
||
def extend_with_intro(input_file, output_folder, duration=0.5):
|
||
base_name = os.path.basename(input_file)
|
||
intro_file = os.path.join(output_folder, f"{base_name}_intro.mkv")
|
||
output_file = os.path.join(output_folder, base_name)
|
||
|
||
info = get_media_info(input_file)
|
||
|
||
# Step 1: generate intro clip
|
||
ffmpeg_intro = build_intro_cmd(info, duration, intro_file)
|
||
print("FFmpeg intro command:\n " + " ".join(shlex.quote(x) for x in ffmpeg_intro))
|
||
subprocess.run(ffmpeg_intro, check=True)
|
||
|
||
# Step 2: mkvmerge – intro + main (VA) + main (subs+attachments+chapters)
|
||
mkvmerge_cmd = [
|
||
"mkvmerge",
|
||
"-o", output_file,
|
||
"--append-to", "1:0:0:0,1:1:0:1",
|
||
intro_file,
|
||
"+",
|
||
"--no-subtitles", "--no-chapters", "--no-attachments", # VA only from main
|
||
input_file,
|
||
"--no-video", "--no-audio", # subs+attachments+chapters only
|
||
input_file,
|
||
]
|
||
print("mkvmerge command:\n " + " ".join(shlex.quote(x) for x in mkvmerge_cmd))
|
||
subprocess.run(mkvmerge_cmd, check=True)
|
||
|
||
os.remove(intro_file)
|
||
print(f"Extended {base_name} -> {output_file}")
|
||
|
||
def process_folder(input_folder, duration=0.5):
|
||
output_folder = os.path.join(input_folder, "output")
|
||
os.makedirs(output_folder, exist_ok=True)
|
||
mkv_files = glob.glob(os.path.join(input_folder, "*.mkv"))
|
||
if not mkv_files:
|
||
print("No MKV files found in the specified folder.")
|
||
return
|
||
|
||
for mkv_file in mkv_files:
|
||
try:
|
||
extend_with_intro(mkv_file, output_folder, duration=duration)
|
||
except subprocess.CalledProcessError as e:
|
||
# Surface stderr if present to help debug codec/param mismatches
|
||
print(f"Error processing {mkv_file}: {getattr(e, 'stderr', e)}")
|
||
|
||
if __name__ == "__main__":
|
||
if len(sys.argv) < 2 or len(sys.argv) > 3:
|
||
print("Usage: python extend_video.py /path/to/folder [duration_seconds]")
|
||
sys.exit(1)
|
||
|
||
input_folder = sys.argv[1]
|
||
if not os.path.isdir(input_folder):
|
||
print(f"Error: '{input_folder}' is not a valid folder path.")
|
||
sys.exit(1)
|
||
|
||
duration = float(sys.argv[2]) if len(sys.argv) == 3 else 1
|
||
print(f"Extending videos in folder: {input_folder} with intro duration: {duration} seconds\n")
|
||
process_folder(input_folder, duration=duration)
|