Files
media-batch-tools/tools/video_encode_output.py
T

731 lines
23 KiB
Python

from __future__ import annotations
import subprocess
import sys
import re
from contextlib import contextmanager
from collections.abc import Iterator
from dataclasses import dataclass
from pathlib import Path
from threading import Thread
from typing import Callable
from tools.avisynth_paths import avs_path
from tools.console import PipelineProgressView, ProgressView
from tools.video_color import ColorMetadata
from tools.video_formatting import chroma_family
ENCODER_PRESETS = (
"ultrafast",
"superfast",
"veryfast",
"faster",
"fast",
"medium",
"slow",
"slower",
"veryslow",
"placebo",
)
@dataclass(frozen=True)
class VideoCodecOptions:
codec: str
crf: int
preset: str
extension: str = ".mp4"
profile: str | None = None
level: str | None = None
threads: int | None = None
@dataclass(frozen=True)
class AudioEncodeOptions:
mode: str
bitrate: str = "192k"
preserve_pitch: bool = True
@dataclass(frozen=True)
class VideoEncodeRequest:
y4m_command: list[str]
ffmpeg: Path
script: Path
audio_source: Path
output: Path
options: VideoCodecOptions
audio_options: AudioEncodeOptions
audio_start: float
audio_duration: float | None
audio_segments: list[tuple[float, float]] | None
audio_tempo: float
audio_sample_rate: int | None
source_has_audio: bool
allow_audio_copy: bool
pixel_format: str
colorspace: str | None
color_metadata: ColorMetadata | None
static_hdr: tuple[str | None, int | None, int | None, bool] | None
dynamic_hdr10plus: Path | None
frame_count: int | None
progress_label: str
def default_audio_bitrate(mode: str) -> str:
return {"aac": "192k", "opus": "128k"}.get(mode, "")
@dataclass(frozen=True)
class VideoEncodeResult:
script: Path
output: Path
renderer_returncode: int
ffmpeg_returncode: int
encoder_log: str = ""
@dataclass(frozen=True)
class Y4MEncodeRun:
renderer_returncode: int
encoder_returncode: int
encoder_log: str
renderer_log: str
@dataclass
class RendererStderrReader:
process: subprocess.Popen
progress: PipelineProgressView | None
lines: list[str]
thread: Thread
def finish(self) -> str:
self.thread.join()
if self.progress is not None:
self.progress.finish()
return "".join(self.lines)
def choose_video_pixel_format(
ffmpeg: Path,
codec: str,
bit_depth: int,
chroma: str | None = None,
) -> str:
family = chroma_family(chroma)
supported = supported_pixel_formats(ffmpeg, codec, family)
preserving_formats = [
(depth, pix_fmt)
for depth, pix_fmt in supported
if depth >= bit_depth
]
if preserving_formats:
return min(preserving_formats, key=lambda item: item[0])[1]
raise RuntimeError(
f"{codec} does not report a yuv{family} pixel format that preserves "
f"the {bit_depth}-bit Avisynth output"
)
def pixel_format_bit_depth(pixel_format: str) -> int:
for depth in (16, 14, 12, 10):
if f"p{depth}" in pixel_format:
return depth
return 8
def supported_pixel_formats(ffmpeg: Path, codec: str, family: str) -> list[tuple[int, str]]:
completed = subprocess.run(
[str(ffmpeg), "-hide_banner", "-h", f"encoder={codec}"],
capture_output=True,
text=True,
)
if completed.returncode != 0:
return [(8, "yuv420p")]
supported: list[tuple[int, str]] = []
for line in completed.stdout.splitlines() + completed.stderr.splitlines():
if "Supported pixel formats:" not in line:
continue
formats = line.split(":", 1)[1].split()
base = f"yuv{family}p"
if base in formats:
supported.append((8, base))
for depth in (10, 12, 14, 16):
pix_fmt = f"{base}{depth}le"
if pix_fmt in formats:
supported.append((depth, pix_fmt))
break
if supported:
return supported
return [(8, f"yuv{family}p")] if family == "420" else []
def add_x265_color_metadata(command: list[str], color_metadata: ColorMetadata) -> None:
value = color_metadata.x265_params()
if value:
add_codec_params(command, "-x265-params", value)
def x264_static_hdr_options(
static_hdr: tuple[str | None, int | None, int | None, bool] | None,
) -> list[str]:
if static_hdr is None:
return []
mastering_display, max_content, max_average, _ = static_hdr
options: list[str] = []
if mastering_display:
options.extend(["--mastering-display", mastering_display])
if max_content is not None:
options.extend(["--cll", f"{max_content},{max_average or 0}"])
return options
def render_audio_wav(
*,
avs_runner: Path,
script: Path,
output: Path,
) -> None:
output.parent.mkdir(parents=True, exist_ok=True)
with output.open("wb") as wav_file:
completed = subprocess.run(
[str(avs_runner), "--wav", str(script)],
stdout=wav_file,
stderr=subprocess.PIPE,
text=False,
)
if completed.returncode != 0:
raise RuntimeError(
f"Avisynth runner failed to render audio for {script} with exit "
f"{completed.returncode}:\n"
+ completed.stderr.decode(errors="replace")
)
def remux_video(
*,
ffmpeg: Path,
source: Path,
audio_source: Path | None,
output: Path,
audio_options: AudioEncodeOptions,
source_has_audio: bool,
frame_count: int | None = None,
progress_label: str | None = None,
) -> VideoEncodeResult:
"""Copy the source video stream while optionally copying or encoding audio."""
output.parent.mkdir(parents=True, exist_ok=True)
command = [
str(ffmpeg), "-y", "-hide_banner", "-loglevel", "info", "-nostats",
"-progress", "pipe:2", "-i", str(source), "-map", "0:v:0", "-c:v", "copy",
]
if audio_options.mode == "none" or not source_has_audio:
command.append("-an")
else:
source_index = 0
if audio_source is not None and audio_source.resolve() != source.resolve():
command.extend(["-i", str(audio_source)])
source_index = 1
command.extend(["-map", f"{source_index}:a:0"])
if audio_options.mode == "copy":
command.extend(["-c:a", "copy"])
else:
add_audio_encoder_options(command, audio_options)
if output.suffix.lower() == ".mp4":
command.extend(["-movflags", "+write_colr"])
command.append(str(output))
process = subprocess.Popen(
command,
stdout=subprocess.DEVNULL,
stderr=subprocess.PIPE,
text=True,
)
encoder_log = read_ffmpeg_progress(
process,
total_frames=frame_count,
label=f"Remuxing {progress_label or output.name}",
)
if process.returncode != 0:
raise RuntimeError(
f"FFmpeg remux failed for {source} with exit {process.returncode}:\n{encoder_log}"
)
return VideoEncodeResult(
script=source,
output=output,
renderer_returncode=0,
ffmpeg_returncode=process.returncode,
encoder_log=encoder_log,
)
def encode_video_only_y4m(request: VideoEncodeRequest) -> VideoEncodeResult:
request.output.parent.mkdir(parents=True, exist_ok=True)
command = ffmpeg_y4m_command(request)
run = run_y4m_encoder(
y4m_command=request.y4m_command,
script=request.script,
encoder_command=command,
frame_count=request.frame_count,
progress_label=request.progress_label,
read_progress=read_ffmpeg_progress,
)
if run.encoder_returncode != 0:
raise RuntimeError(
f"FFmpeg failed for {request.script} with exit {run.encoder_returncode}:\n"
+ run.encoder_log
+ (
"\nAvisynth runner also exited with "
f"{run.renderer_returncode}:\n"
+ run.renderer_log
if run.renderer_returncode != 0
else ""
)
)
if run.renderer_returncode != 0:
raise RuntimeError(
f"Avisynth runner failed for {request.script} with exit {run.renderer_returncode}:\n"
+ run.renderer_log
)
return VideoEncodeResult(
script=request.script,
output=request.output,
renderer_returncode=run.renderer_returncode,
ffmpeg_returncode=run.encoder_returncode,
encoder_log=run.encoder_log,
)
def ffmpeg_y4m_command(request: VideoEncodeRequest) -> list[str]:
command = [
str(request.ffmpeg), "-y", "-hide_banner", "-loglevel", "info", "-nostats",
"-progress", "pipe:2", "-f", "yuv4mpegpipe", "-i", "pipe:0",
]
if request.audio_options.mode != "none" and request.source_has_audio:
add_audio_input_options(
command,
audio_options=request.audio_options,
audio_start=request.audio_start,
audio_duration=request.audio_duration,
allow_audio_copy=request.allow_audio_copy,
)
command.extend(["-i", str(request.audio_source)])
command.extend([
"-map", "0:v:0", "-c:v", request.options.codec, "-preset", request.options.preset,
"-crf", str(request.options.crf), "-pix_fmt", request.pixel_format,
])
add_video_profile_level(command, request.options)
add_video_threads(command, request.options)
if request.options.codec == "libx265" and request.color_metadata:
add_x265_color_metadata(command, request.color_metadata)
add_static_hdr_metadata(command, request.options, request.static_hdr)
add_dynamic_hdr10plus_metadata(command, request.options, request.dynamic_hdr10plus)
if request.color_metadata is not None:
command.extend(request.color_metadata.ffmpeg_args())
elif request.colorspace is not None:
command.extend(["-colorspace", request.colorspace])
add_audio_options(
command,
audio_options=request.audio_options,
audio_start=request.audio_start,
audio_duration=request.audio_duration,
audio_segments=request.audio_segments,
audio_tempo=request.audio_tempo,
audio_sample_rate=request.audio_sample_rate,
source_has_audio=request.source_has_audio,
allow_audio_copy=request.allow_audio_copy,
)
if request.output.suffix.lower() == ".mp4" and request.color_metadata and request.color_metadata.ffmpeg_args():
command.extend(["-movflags", "+write_colr"])
command.append(str(request.output))
return command
def run_y4m_encoder(
*,
y4m_command: list[str],
script: Path,
encoder_command: list[str],
frame_count: int | None,
progress_label: str,
read_progress: Callable[..., str],
) -> Y4MEncodeRun:
with video_only_script(script) as render_script:
renderer_process = subprocess.Popen(
[*y4m_command, str(render_script)], stdout=subprocess.PIPE, stderr=subprocess.PIPE
)
assert renderer_process.stdout is not None
renderer_reader = start_renderer_stderr_reader(
renderer_process,
rendering_label=f"Rendering {progress_label}",
encoding_label=f"Encoding {progress_label}",
encoding_total=frame_count,
)
encoder_process = subprocess.Popen(
encoder_command,
stdin=renderer_process.stdout,
stdout=subprocess.DEVNULL,
stderr=subprocess.PIPE,
text=True,
)
renderer_process.stdout.close()
encoder_log = read_progress(
encoder_process,
total_frames=frame_count,
label=f"Encoding {progress_label}",
pipeline_progress=renderer_reader.progress,
)
renderer_log = renderer_reader.finish()
renderer_returncode = renderer_process.wait()
return Y4MEncodeRun(
renderer_returncode=renderer_returncode,
encoder_returncode=encoder_process.returncode,
encoder_log=encoder_log,
renderer_log=renderer_log,
)
@contextmanager
def video_only_script(script: Path) -> Iterator[Path]:
wrapper = script.with_name(f".{script.name}.video-only.avs")
wrapper.write_text(
'mbt_video_only = true\n'
f'Import("{avs_path(script.name)}")\n'
"last\n",
encoding="utf-8",
)
try:
yield wrapper
finally:
wrapper.unlink(missing_ok=True)
def add_audio_input_options(
command: list[str],
*,
audio_options: AudioEncodeOptions,
audio_start: float,
audio_duration: float | None,
allow_audio_copy: bool,
) -> None:
if audio_options.mode != "copy":
return
if not allow_audio_copy:
return
if audio_start > 0.001:
command.extend(["-ss", f"{audio_start:.6f}"])
if audio_duration is not None:
command.extend(["-t", f"{audio_duration:.6f}"])
def add_audio_options(
command: list[str],
*,
audio_options: AudioEncodeOptions,
audio_start: float,
audio_duration: float | None,
audio_segments: list[tuple[float, float]] | None = None,
audio_tempo: float,
audio_sample_rate: int | None,
source_has_audio: bool,
allow_audio_copy: bool,
) -> None:
if audio_options.mode == "none" or not source_has_audio:
command.append("-an")
elif audio_options.mode in {"aac", "opus", "flac"}:
if audio_segments is not None:
add_segmented_encoded_audio(
command,
audio_segments=audio_segments,
audio_options=audio_options,
audio_tempo=audio_tempo,
audio_sample_rate=audio_sample_rate,
)
return
if audio_duration is None:
raise RuntimeError("audio_duration is required for AAC audio")
audio_filters = [
f"atrim=start={audio_start:.6f}:duration={audio_duration:.6f}",
"asetpts=PTS-STARTPTS",
]
audio_filters.extend(
speed_change_filters(audio_tempo, audio_options, audio_sample_rate)
)
command.extend(["-map", "1:a:0", "-af", ",".join(audio_filters)])
add_audio_encoder_options(command, audio_options)
elif audio_options.mode == "copy":
if not allow_audio_copy:
raise RuntimeError("Audio copy is only supported for whole-file or simple beginning/end trims.")
command.extend(["-map", "1:a:0", "-c:a", "copy"])
else:
raise RuntimeError(f"Unsupported audio mode: {audio_options.mode}")
def add_video_profile_level(command: list[str], options: VideoCodecOptions) -> None:
if options.profile:
command.extend(["-profile:v", options.profile])
if options.level:
command.extend(["-level:v", options.level])
def add_video_threads(command: list[str], options: VideoCodecOptions) -> None:
if options.threads is None:
return
if options.codec == "libx265":
add_codec_params(command, "-x265-params", f"pools={options.threads}")
else:
command.extend(["-threads", str(options.threads)])
def add_static_hdr_metadata(
command: list[str],
options: VideoCodecOptions,
static_hdr: tuple[str | None, int | None, int | None, bool] | None,
) -> None:
if static_hdr is None:
return
mastering_display, max_content, max_average, is_pq = static_hdr
params: list[str] = ["hdr10=1"] if options.codec == "libx265" and is_pq else []
if mastering_display:
params.append(f"master-display={mastering_display}")
if max_content is not None:
params.append(f"max-cll={max_content},{max_average or 0}")
if not params:
return
option = "-x265-params" if options.codec == "libx265" else "-x264-params"
add_codec_params(command, option, ":".join(params))
def add_dynamic_hdr10plus_metadata(
command: list[str],
options: VideoCodecOptions,
metadata_json: Path | None,
) -> None:
if metadata_json is None:
return
if options.codec != "libx265":
raise RuntimeError("HDR10+ re-encoding requires x265")
# FFmpeg passes x265 options as a colon-separated string. Escaping keeps
# the Windows drive separator inside the JSON path value.
json_path = x265_parameter_path(metadata_json.resolve())
add_codec_params(
command,
"-x265-params",
f"dhdr10-info={json_path}:dhdr10-opt=1",
)
def x265_parameter_path(path: Path) -> str:
return str(path).replace("\\", "/").replace(":", r"\:")
def add_codec_params(command: list[str], option: str, value: str) -> None:
if option in command:
index = command.index(option) + 1
command[index] = f"{command[index]}:{value}"
else:
command.extend([option, value])
def print_encoder_statistics(encoder_log: str) -> None:
statistics = encoder_statistics_lines(encoder_log)
if not statistics:
return
print(" encoder stats:")
for line in statistics:
print(f" {line}")
def encoder_statistics_lines(encoder_log: str) -> list[str]:
pattern = re.compile(
r"(frame [IPB]:.*|consecutive B-frames:.*|ref [PB] L[01]:.*)",
re.IGNORECASE,
)
statistics: list[str] = []
for line in encoder_log.splitlines():
match = pattern.search(line)
if match:
statistics.append(match.group(1))
return statistics
def start_renderer_stderr_reader(
process: subprocess.Popen,
*,
rendering_label: str,
encoding_label: str,
encoding_total: int | None,
) -> RendererStderrReader:
assert process.stderr is not None
progress = (
PipelineProgressView(
rendering_label=rendering_label,
encoding_label=encoding_label,
encoding_total=encoding_total,
)
if encoding_total and sys.stdout.isatty()
else None
)
lines: list[str] = []
def consume() -> None:
for raw_line in process.stderr:
line = raw_line.decode(errors="replace")
if progress is not None and line.startswith("mbt_render,"):
if _update_renderer_progress(progress, line):
continue
lines.append(line)
thread = Thread(target=consume, daemon=True)
thread.start()
return RendererStderrReader(process, progress, lines, thread)
def _update_renderer_progress(progress: PipelineProgressView, line: str) -> bool:
parts = line.strip().split(",")
if len(parts) != 6:
return False
try:
processed = int(parts[1])
total = int(parts[2])
elapsed = float(parts[4])
eta = float(parts[5])
except ValueError:
return False
progress.update_rendering(processed, total, elapsed=elapsed, eta=eta)
return True
def add_segmented_encoded_audio(
command: list[str],
*,
audio_segments: list[tuple[float, float]],
audio_options: AudioEncodeOptions,
audio_tempo: float,
audio_sample_rate: int | None,
) -> None:
if not audio_segments:
raise RuntimeError("audio_segments must not be empty")
filters: list[str] = []
labels: list[str] = []
for index, (start, duration) in enumerate(audio_segments):
label = f"a{index}"
filters.append(
f"[1:a]atrim=start={start:.6f}:duration={duration:.6f},"
f"asetpts=PTS-STARTPTS[{label}]"
)
labels.append(f"[{label}]")
if len(labels) == 1:
current_label = labels[0]
else:
filters.append(
"".join(labels) + f"concat=n={len(labels)}:v=0:a=1[acut]"
)
current_label = "[acut]"
speed_filters = speed_change_filters(audio_tempo, audio_options, audio_sample_rate)
if speed_filters:
filters.append(f"{current_label}{','.join(speed_filters)}[aout]")
elif current_label != "[aout]":
filters.append(f"{current_label}anull[aout]")
command.extend(["-filter_complex", ";".join(filters), "-map", "[aout]"])
add_audio_encoder_options(command, audio_options)
def add_audio_encoder_options(
command: list[str],
audio_options: AudioEncodeOptions,
) -> None:
if audio_options.mode == "aac":
command.extend(["-c:a", "aac", "-b:a", audio_options.bitrate])
elif audio_options.mode == "opus":
command.extend(["-c:a", "libopus", "-b:a", audio_options.bitrate])
elif audio_options.mode == "flac":
command.extend(["-c:a", "flac"])
else:
raise RuntimeError(f"Unsupported encoded audio mode: {audio_options.mode}")
def atempo_filters(tempo: float) -> list[str]:
if tempo <= 0:
raise RuntimeError(f"Invalid audio tempo factor: {tempo}")
if abs(tempo - 1.0) < 0.0001:
return []
filters: list[str] = []
remaining = tempo
while remaining < 0.5:
filters.append("atempo=0.5")
remaining /= 0.5
while remaining > 2.0:
filters.append("atempo=2.0")
remaining /= 2.0
filters.append(f"atempo={remaining:.8f}")
return filters
def speed_change_filters(
tempo: float,
audio_options: AudioEncodeOptions,
audio_sample_rate: int | None,
) -> list[str]:
if abs(tempo - 1.0) < 0.0001:
return []
if audio_options.preserve_pitch:
return atempo_filters(tempo)
if audio_sample_rate is None:
raise RuntimeError(
"Audio sample rate is required for speed changes without pitch preservation"
)
changed_rate = max(1, round(audio_sample_rate * tempo))
return [f"asetrate={changed_rate}", f"aresample={audio_sample_rate}"]
def read_ffmpeg_progress(
process: subprocess.Popen,
*,
total_frames: int | None,
label: str,
pipeline_progress: PipelineProgressView | None = None,
) -> str:
assert process.stderr is not None
errors: list[str] = []
progress = (
ProgressView(
total_frames or 0,
label,
embedded_percent=True,
show_rate=True,
)
if total_frames and sys.stdout.isatty() and pipeline_progress is None
else None
)
current_frame = 0
if progress is not None:
progress.update(0)
for line in process.stderr:
text = line.rstrip("\r\n")
if text.startswith("frame="):
try:
current_frame = int(text.split("=", 1)[1].strip())
except ValueError:
continue
if progress is not None:
progress.update(current_frame)
elif pipeline_progress is not None:
pipeline_progress.update_encoding(current_frame)
elif text.startswith(("fps=", "stream_", "bitrate=", "total_size=", "out_time_", "dup_frames=", "drop_frames=", "speed=", "progress=")):
continue
elif text:
errors.append(text)
process.wait()
if progress is not None:
progress.finish(keep=True)
return "\n".join(errors)