Files

502 lines
18 KiB
Python

from __future__ import annotations
import os
import shutil
import sys
import time
from datetime import timedelta, timezone
from pathlib import Path
from statistics import median
from tools.video_formatting import (
format_frame_count,
format_framerate,
format_seconds,
join_parts,
video_parts,
)
from tools.video_inputs import VideoInput
from tools.video_probe import VideoProbe
WORKSPACE_DIR = Path("video_workspace")
SOURCE_DIR = ".source"
HELPER_SCRIPT_NAME = "avisynth_helpers.avs"
def workspace_root(script_path: Path) -> Path:
return script_path.resolve().parent / WORKSPACE_DIR
def visible_script_path(root: Path, item: VideoInput) -> Path:
return root / item.collapsed_relative.with_name(
f"{item.collapsed_relative.name}.avs"
)
def hidden_source_script_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".source.avs")
def frame_timestamp_values_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".timestamps.txt")
def frame_time_values_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".frame-times.txt")
def frame_time_label_values_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".frame-time-labels.txt")
def frame_duration_values_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".durations.txt")
def frame_duration_label_values_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".duration-labels.txt")
def validation_script_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".validate.avs")
def validation_csv_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".frames.csv")
def editable_audio_script_path(root: Path, item: VideoInput) -> Path:
visible = visible_script_path(root, item)
return visible.with_name(f"{visible.stem}_audio.avs")
def rendered_audio_wav_path(root: Path, item: VideoInput) -> Path:
return _source_sidecar_path(root, item, ".audio.wav")
def _source_sidecar_path(root: Path, item: VideoInput, suffix: str) -> Path:
return root / SOURCE_DIR / item.collapsed_relative.with_name(
f"{item.collapsed_relative.name}{suffix}"
)
def existing_visible_scripts(root: Path, items: list[VideoInput]) -> list[Path]:
return [path for path in (visible_script_path(root, item) for item in items) if path.exists()]
def write_workspace(
*,
root: Path,
items: list[VideoInput],
probes_by_path: dict[Path, VideoProbe],
overwrite_visible: bool | set[Path],
ffms2_plugin_path: Path | None,
) -> tuple[list[Path], list[Path]]:
written: list[Path] = []
skipped_visible: list[Path] = []
source_root = root / SOURCE_DIR
_clear_source_root(source_root)
helper_script = helper_script_path()
for source_id, item in enumerate(items):
probe = probes_by_path[item.path.resolve()]
source_script = hidden_source_script_path(root, item)
timestamp_values = frame_timestamp_values_path(root, item)
frame_time_values = frame_time_values_path(root, item)
frame_time_label_values = frame_time_label_values_path(root, item)
duration_values = frame_duration_values_path(root, item)
duration_label_values = frame_duration_label_values_path(root, item)
validation_script = validation_script_path(root, item)
visible_script = visible_script_path(root, item)
source_script.parent.mkdir(parents=True, exist_ok=True)
visible_script.parent.mkdir(parents=True, exist_ok=True)
timestamp_values.write_text(_timestamp_values_file(probe), encoding="utf-8")
frame_time_values.write_text(_frame_time_values_file(probe), encoding="utf-8")
frame_time_label_values.write_text(
_frame_time_label_values_file(probe), encoding="utf-8"
)
duration_values.write_text(_duration_values_file(probe), encoding="utf-8")
duration_label_values.write_text(_duration_label_values_file(probe), encoding="utf-8")
source_script.write_text(
_hidden_source_script(
item,
probe,
source_id,
source_script,
helper_script,
ffms2_plugin_path,
timestamp_values,
frame_time_values,
frame_time_label_values,
duration_values,
duration_label_values,
),
encoding="utf-8",
)
validation_script.write_text(
_validation_script(visible_script, validation_script),
encoding="utf-8",
)
written.extend([
timestamp_values,
frame_time_values,
frame_time_label_values,
duration_values,
duration_label_values,
source_script,
validation_script,
])
if visible_script.exists() and not _should_overwrite_visible(visible_script, overwrite_visible):
skipped_visible.append(visible_script)
else:
visible_script.write_text(
_visible_script(
visible_script,
source_script,
item,
probe,
),
encoding="utf-8",
)
written.append(visible_script)
return written, skipped_visible
def _clear_source_root(source_root: Path) -> None:
if not source_root.exists():
return
for attempt in range(5):
try:
shutil.rmtree(source_root)
return
except OSError as exc:
if attempt == 4:
raise RuntimeError(
f"Could not reset {source_root}. Close any program using files in "
"video_workspace/.source and try again."
) from exc
time.sleep(0.25)
def _should_overwrite_visible(path: Path, overwrite_visible: bool | set[Path]) -> bool:
if isinstance(overwrite_visible, bool):
return overwrite_visible
return path.resolve() in overwrite_visible
def _visible_script(
visible_script: Path,
source_script: Path,
item: VideoInput,
probe: VideoProbe,
) -> str:
import_path = _avs_path(os.path.relpath(source_script, visible_script.parent))
audio_script_name = f"{visible_script.stem}_audio.avs"
timing_kind = probe.timing.timing_kind.upper()
final_clip = _default_final_clip(probe)
windows_32bit_note = (
"# End this file with a commented #32bit marker to use 32-bit AviSynth.\n"
if sys.platform == "win32"
else ""
)
return (
"# Don't touch.\n"
f'Import("{import_path}")\n'
"\n"
f"# Source: {item.path}\n"
f"# Video: {_format_video_summary(probe)}\n"
f"# Duration: {format_seconds(probe.duration_seconds)}, {_format_frame_count(probe)}\n"
f"# Framerate: {format_framerate(timing_kind, probe.avg_frame_rate)}\n"
"\n"
"# Trimmed frames also trim audio according to the timestamps of the trimmed frames.\n"
"# MBT_Drop marks frames for dropping without changing other frame timings.\n"
"# Frame properties must be preserved.\n"
"\n"
"# Bundled helper functions:\n"
"# MBT_Drop(last, int start, int \"end\"), MBT_Undrop(...)\n"
"# MBT_DropEvery(last, int cycle, int offset0, ...), MBT_UndropEvery(...)\n"
"# Similar syntax as Trim/SelectEvery. end defaults to -1, meaning one frame.\n"
"# Examples: MBT_Drop(last, 100), MBT_Drop(last, 100, -20),\n"
"# MBT_DropEvery(last, 4, 1, 3)\n"
"# MBT_Info(last, int \"size\", int \"align\")\n"
"# Overlay source frame, video time, absolute time, duration, and drop marker.\n"
"# Resize(clip v, int \"w\", int \"h\", float \"dar\", float \"xcrop\", float \"ycrop\", int \"pos\", bool \"linear\", string \"input_matrix\", string \"output_matrix\")\n"
"# Resize/crop. linear auto-enables for large downscales unless set explicitly.\n"
"# MBT_CorrectMatrix(last)\n"
"# Convert to the resolution-expected matrix using _Matrix as input metadata.\n"
"# DeShake(clip c, int \"w\", int \"h\", float \"dar\", float \"crop\", float \"xcrop\", float \"ycrop\", int \"pos\", bool \"zoom\", bool \"rot\", float \"freq\")\n"
"# Stabilize through MVTools + DePan; those plugins must be loaded separately.\n"
"# Cropf(...)\n"
"# Crop using fractional dimensions.\n"
"# RotateCrop(clip c, float angle, float \"dar\")\n"
"# Rotate with manyPlus Turn, then crop the largest background-free rectangle.\n"
"# FixContrast(...)\n"
"# Remap luma levels to limited-range video luma.\n"
"\n"
"# Optional audio edits:\n"
f"# Copy this script to {audio_script_name} next to this file.\n"
"# Edit that copy for audio fades/channel changes/etc. Its video output is ignored.\n"
"# Its audio output will be trimmed according to this script's frame edits.\n"
f"{windows_32bit_note}"
"\n"
"# User edits below this line. The imported source clip will already be in \"last\".\n"
"\n"
f"{final_clip}\n"
)
def _hidden_source_script(
item: VideoInput,
probe: VideoProbe,
source_id: int,
source_script: Path,
helper_script: Path,
ffms2_plugin_path: Path | None,
timestamp_values: Path,
frame_time_values: Path,
frame_time_label_values: Path,
duration_values: Path,
duration_label_values: Path,
) -> str:
source_path = _avs_path(str(item.path.resolve()))
import_path = _avs_path(os.path.relpath(helper_script, source_script.parent))
cache_path = _avs_path(
str(source_script.with_name(f"{item.collapsed_relative.name}.ffindex").resolve())
)
audio_cache_path = _avs_path(
str(source_script.with_name(f"{item.collapsed_relative.name}.audio.ffindex").resolve())
)
timestamp_values_path = _avs_path(str(timestamp_values.resolve()))
frame_time_values_path = _avs_path(str(frame_time_values.resolve()))
frame_time_label_values_path = _avs_path(str(frame_time_label_values.resolve()))
duration_values_path = _avs_path(str(duration_values.resolve()))
duration_label_values_path = _avs_path(str(duration_label_values.resolve()))
if ffms2_plugin_path is None:
plugin_setup = "# FFMS2 is expected from Avisynth autoload or a previously loaded plugin.\n"
else:
plugin_setup = f'LoadPlugin("{_avs_path(str(ffms2_plugin_path.resolve()))}")\n'
audio_setup = (
"try {\n"
" mbt_video_only = mbt_video_only\n"
"} catch (err_msg) {\n"
" mbt_video_only = false\n"
"}\n"
f'source = mbt_video_only ? video : MBT_AudioSource(video, "{source_path}", "{audio_cache_path}")\n'
if probe.audio_stream_count > 0
else "source = video\n"
)
validation_blank_setup = (
"try {\n"
" mbt_validation_blank = mbt_validation_blank\n"
"} catch (err_msg) {\n"
" mbt_validation_blank = false\n"
"}\n"
"video = mbt_validation_blank ? BlankClip(video, color=$000000) : video\n"
)
return (
"# Generated by media-batch-tools. This file is regenerated by video_encode.py.\n"
"# Do not put manual edits here; edit the matching script in video_workspace/.\n"
f"# Source: {source_path}\n"
f"# Source id: {source_id}\n"
f"# Detected timing: {probe.timing.timing_kind.upper()}\n"
"# Requires FFMS2 to provide FFVideoSource.\n"
"\n"
f"{plugin_setup}"
f'Import("{import_path}")\n'
"\n"
f'video = FFVideoSource("{source_path}", cachefile="{cache_path}")\n'
f"{validation_blank_setup}"
f"{audio_setup}"
f"last = MBT_MarkSource(source, {source_id}, {_matrix_code_for_source(probe)})\n"
f'last = ConditionalReader(last, "{timestamp_values_path}", "mbt_absolute_timestamp", false)\n'
f'last = ConditionalReader(last, "{frame_time_values_path}", "mbt_frame_time_seconds", false)\n'
f'last = ConditionalReader(last, "{frame_time_label_values_path}", "mbt_relative_timestamp", false)\n'
f'last = ConditionalReader(last, "{duration_values_path}", "mbt_frame_duration_seconds", false)\n'
f'last = ConditionalReader(last, "{duration_label_values_path}", "mbt_frame_duration", false)\n'
"last\n"
)
def _validation_script(
visible_script: Path,
validation_script: Path,
) -> str:
import_path = _avs_path(os.path.relpath(visible_script, validation_script.parent))
return (
"# Generated by media-batch-tools. This file is regenerated by video_encode.py.\n"
"# The media-batch-tools runner reads frame identity properties from this clip.\n"
f'Import("{import_path}")\n'
"last\n"
)
def _timestamp_values_file(probe: VideoProbe) -> str:
lines = [
"# Generated by media-batch-tools for AviSynth ConditionalReader.",
"Type string",
"Default",
]
if probe.metadata_end is None or probe.duration_seconds is None:
return "\n".join(lines) + "\n"
source_end = probe.metadata_end
if source_end.tzinfo is None:
source_end = source_end.replace(tzinfo=timezone.utc)
source_begin = source_end - timedelta(seconds=probe.duration_seconds)
for frame, timestamp in enumerate(probe.frame_timestamps):
value = (source_begin + timedelta(seconds=timestamp)).astimezone(timezone.utc)
lines.append(f"{frame} {value.isoformat().replace('+00:00', 'Z')}")
return "\n".join(lines) + "\n"
def _frame_time_values_file(probe: VideoProbe) -> str:
lines = [
"# Generated by media-batch-tools for AviSynth ConditionalReader.",
"Type float",
"Default 0.0",
]
for frame, timestamp in enumerate(probe.frame_timestamps):
lines.append(f"{frame} {timestamp:.9f}")
return "\n".join(lines) + "\n"
def _frame_time_label_values_file(probe: VideoProbe) -> str:
lines = [
"# Generated by media-batch-tools for AviSynth ConditionalReader.",
"Type string",
"Default 0.000000",
]
for frame, timestamp in enumerate(probe.frame_timestamps):
lines.append(f"{frame} {_format_relative_timestamp(timestamp)}")
return "\n".join(lines) + "\n"
def _duration_values_file(probe: VideoProbe) -> str:
lines = [
"# Generated by media-batch-tools for AviSynth ConditionalReader.",
"Type float",
"Default 0.0",
]
for frame, duration in enumerate(_source_frame_durations(probe)):
lines.append(f"{frame} {duration:.9f}")
return "\n".join(lines) + "\n"
def _duration_label_values_file(probe: VideoProbe) -> str:
lines = [
"# Generated by media-batch-tools for AviSynth ConditionalReader.",
"Type string",
"Default",
]
for frame, duration in enumerate(_source_frame_durations(probe)):
lines.append(f"{frame} {_format_duration_label(duration)}")
return "\n".join(lines) + "\n"
def _format_duration_label(duration: float) -> str:
if duration <= 0:
return f"{duration:.6f} (1/inf)"
reciprocal = 1.0 / duration
return f"{duration:.6f} (1/{reciprocal:.2g})"
def _format_relative_timestamp(seconds: float) -> str:
if seconds < 60.0:
return f"{seconds:.6f}"
total = int(seconds)
fraction = seconds - total
fractional = f"{fraction:.6f}"[1:]
if total < 3600:
minutes = total // 60
remaining_seconds = total % 60
return f"{minutes}:{remaining_seconds:02d}{fractional} ({seconds:.6f})"
hours = total // 3600
minutes = (total % 3600) // 60
remaining_seconds = total % 60
return f"{hours}:{minutes:02d}:{remaining_seconds:02d}{fractional} ({seconds:.6f})"
def _source_frame_durations(probe: VideoProbe) -> list[float]:
timestamps = probe.frame_timestamps
if not timestamps:
return []
durations: list[float] = []
for current, next_timestamp in zip(timestamps, timestamps[1:]):
durations.append(max(0.0, next_timestamp - current))
fallback = median(durations) if durations else 0.0
if probe.duration_seconds is not None and len(timestamps) > 1:
last_duration = probe.duration_seconds - timestamps[-1]
if last_duration <= 0:
last_duration = fallback
else:
last_duration = fallback
durations.append(max(0.0, last_duration))
return durations
def _format_video_summary(probe: VideoProbe) -> str:
parts = video_parts(
width=probe.width,
height=probe.height,
pixel_format=probe.pixel_format,
bit_depth=probe.bit_depth,
)
parts.append(_format_matrix(probe))
return join_parts(parts)
def _format_matrix(probe: VideoProbe) -> str:
expected = probe.expected_color_matrix or "unknown expected"
if probe.color_matrix is not None:
if probe.color_space:
return f"{probe.color_matrix} ({probe.color_space}), expected {expected}"
return f"{probe.color_matrix}, expected {expected}"
if probe.color_space:
return f"{probe.color_space}, expected {expected}"
return f"unknown, expected {expected}"
def _matrix_code_for_source(probe: VideoProbe) -> int:
return _matrix_code(probe.color_matrix or probe.expected_color_matrix)
def _matrix_code(matrix: str | None) -> int:
if matrix == "Rec709":
return 1
if matrix == "Rec2020":
return 9
if matrix == "Rec601":
return 6
return 2
def _default_final_clip(probe: VideoProbe) -> str:
source_matrix = probe.color_matrix
expected_matrix = probe.expected_color_matrix
if source_matrix is None or expected_matrix is None or source_matrix == expected_matrix:
return "last"
return "MBT_CorrectMatrix(last) # source metadata differs from resolution default"
def _format_frame_count(probe: VideoProbe) -> str:
return format_frame_count(probe.stream_frame_count, len(probe.frame_timestamps))
def _avs_path(path: str) -> str:
return path.replace("\\", "/").replace('"', '\\"')
def helper_script_path() -> Path:
return Path(__file__).resolve().with_name(HELPER_SCRIPT_NAME)