642 lines
19 KiB
Python
642 lines
19 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.console import PipelineProgressView, ProgressView
|
|
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
|
|
|
|
|
|
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 ffmpeg_colorspace_from_matrix(matrix: int | None) -> str | None:
|
|
if matrix == 1:
|
|
return "bt709"
|
|
if matrix in {5, 6}:
|
|
return "smpte170m"
|
|
if matrix in {9, 10}:
|
|
return "bt2020nc"
|
|
return None
|
|
|
|
|
|
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 encode_video_only_y4m(
|
|
*,
|
|
y4m_command: list[str],
|
|
ffmpeg: Path,
|
|
script: Path,
|
|
audio_source: Path | None,
|
|
output: Path,
|
|
options: VideoCodecOptions | None = None,
|
|
audio_options: AudioEncodeOptions | None = None,
|
|
audio_start: float = 0.0,
|
|
audio_duration: float | None = None,
|
|
audio_segments: list[tuple[float, float]] | None = None,
|
|
audio_tempo: float = 1.0,
|
|
audio_sample_rate: int | None = None,
|
|
source_has_audio: bool = False,
|
|
allow_audio_copy: bool = False,
|
|
pixel_format: str = "yuv420p",
|
|
colorspace: str | None = None,
|
|
frame_count: int | None = None,
|
|
progress_label: str | None = None,
|
|
) -> VideoEncodeResult:
|
|
if options is None:
|
|
options = VideoCodecOptions(codec="libx264", crf=16, preset="slow")
|
|
if audio_options is None:
|
|
audio_options = AudioEncodeOptions(mode="none")
|
|
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
command = ffmpeg_y4m_command(
|
|
ffmpeg=ffmpeg,
|
|
audio_source=audio_source,
|
|
output=output,
|
|
options=options,
|
|
audio_options=audio_options,
|
|
audio_start=audio_start,
|
|
audio_duration=audio_duration,
|
|
audio_segments=audio_segments,
|
|
audio_tempo=audio_tempo,
|
|
audio_sample_rate=audio_sample_rate,
|
|
source_has_audio=source_has_audio,
|
|
allow_audio_copy=allow_audio_copy,
|
|
pixel_format=pixel_format,
|
|
colorspace=colorspace,
|
|
)
|
|
run = run_y4m_encoder(
|
|
y4m_command=y4m_command,
|
|
script=script,
|
|
encoder_command=command,
|
|
frame_count=frame_count,
|
|
progress_label=progress_label or output.name,
|
|
read_progress=read_ffmpeg_progress,
|
|
)
|
|
if run.encoder_returncode != 0:
|
|
raise RuntimeError(
|
|
f"FFmpeg failed for {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 {script} with exit {run.renderer_returncode}:\n"
|
|
+ run.renderer_log
|
|
)
|
|
|
|
return VideoEncodeResult(
|
|
script=script,
|
|
output=output,
|
|
renderer_returncode=run.renderer_returncode,
|
|
ffmpeg_returncode=run.encoder_returncode,
|
|
encoder_log=run.encoder_log,
|
|
)
|
|
|
|
|
|
def ffmpeg_y4m_command(
|
|
*,
|
|
ffmpeg: Path,
|
|
audio_source: Path | None,
|
|
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,
|
|
) -> list[str]:
|
|
command = [
|
|
str(ffmpeg), "-y", "-hide_banner", "-loglevel", "info", "-nostats",
|
|
"-progress", "pipe:2", "-f", "yuv4mpegpipe", "-i", "pipe:0",
|
|
]
|
|
if audio_options.mode != "none" and source_has_audio:
|
|
if audio_source is None:
|
|
raise RuntimeError("An audio source is required when audio encoding is enabled")
|
|
add_audio_input_options(
|
|
command,
|
|
audio_options=audio_options,
|
|
audio_start=audio_start,
|
|
audio_duration=audio_duration,
|
|
allow_audio_copy=allow_audio_copy,
|
|
)
|
|
command.extend(["-i", str(audio_source)])
|
|
command.extend([
|
|
"-map", "0:v:0", "-c:v", options.codec, "-preset", options.preset,
|
|
"-crf", str(options.crf), "-pix_fmt", pixel_format,
|
|
])
|
|
add_video_profile_level(command, options)
|
|
add_video_threads(command, options)
|
|
if colorspace is not None:
|
|
command.extend(["-colorspace", colorspace])
|
|
add_audio_options(
|
|
command,
|
|
audio_options=audio_options,
|
|
audio_start=audio_start,
|
|
audio_duration=audio_duration,
|
|
audio_segments=audio_segments,
|
|
audio_tempo=audio_tempo,
|
|
audio_sample_rate=audio_sample_rate,
|
|
source_has_audio=source_has_audio,
|
|
allow_audio_copy=allow_audio_copy,
|
|
)
|
|
command.append(str(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 _avs_path(path: str) -> str:
|
|
return path.replace("\\", "/").replace('"', '\\"')
|
|
|
|
|
|
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":
|
|
command.extend(["-x265-params", f"pools={options.threads}"])
|
|
else:
|
|
command.extend(["-threads", str(options.threads)])
|
|
|
|
|
|
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)
|