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.avisynth_paths import avs_path 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, avisynth_primaries_from_color_primaries, avisynth_range_from_color_range, avisynth_transfer_from_color_transfer, hdr_kind, ) 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(clip c, int start, int \"end\"), MBT_Undrop(...)\n" "# MBT_DropEvery(clip c, 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(clip c, int \"size\", int \"align\")\n" "# Overlay source frame, video time, absolute time, duration, and drop marker.\n" "# MBT_AbsoluteTimeStretch(clip c, float factor)\n" "# Map source timestamps to real time, anchored to the metadata end timestamp.\n" "# Example: MBT_AbsoluteTimeStretch(last, 1.0/8.0) for 8x slow motion.\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. HDR linear resizing requires HDRTools; linear=false preserves HDR unchanged.\n" "# MBT_CorrectMatrix(clip c, string \"input_matrix\", string \"output_matrix\")\n" "# Correct SDR matrix metadata. HDR color systems are preserved.\n" "# MBT_ToSDR(clip c, float \"hdr_peak\", float \"sdr_peak\", string \"method\", float \"exposure\", bool \"chroma_correction\", string \"output_matrix\")\n" "# HDR to SDR. MBT_ToSDR(10000, 125, \"MPC\") uses a 10,000 cd/m² PQ source\n" "# peak and a 125 cd/m² SDR display target. Methods: MPC (default), Hable, Mobius,\n" "# Reinhard, ACES, BT2446C. exposure overrides the method's default curve scale; chroma_correction\n" "# applies only to BT2446C.\n" "# MPC applies the Hable curve to XYZ luminance, then preserves chromaticity and compresses\n" "# target-gamut chroma. Defaults: PQ hdr_peak=10000, HLG hdr_peak=1000, sdr_peak=125,\n" "# chroma_correction=false, output_matrix=Rec709/Rec601 by output size.\n" "# HDR10+ metadata is dynamic and is not used by HDRTools; this is a static tone map.\n" "# Add LoadPlugin(\"path/to/HDRTools.dll\") below when using HDRTools.\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(clip c, float l, float t, float w, float h)\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(clip c, int \"low\", int \"high\")\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' f'global mbt_color_matrix = "{probe.color_matrix or probe.expected_color_matrix or "Rec709"}"\n' f"global mbt_color_primaries = {_primaries_code_for_source(probe)}\n" f"global mbt_color_transfer = {_transfer_code_for_source(probe)}\n" f"global mbt_color_range = {_range_code_for_source(probe)}\n" f"global mbt_hdr_peak = {_hdr_peak_for_source(probe):.6f}\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)}, {_primaries_code_for_source(probe)}, {_transfer_code_for_source(probe)}, {_range_code_for_source(probe)}, {_hdr_peak_for_source(probe):.6f}, {_source_end_epoch_seconds(probe):.9f}, {(probe.duration_seconds or 0.0):.9f})\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 = MBT_MarkTiming(last)\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 _source_end_epoch_seconds(probe: VideoProbe) -> float: if probe.metadata_end is None: return 0.0 source_end = probe.metadata_end if source_end.tzinfo is None: source_end = source_end.replace(tzinfo=timezone.utc) return source_end.timestamp() 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)) detected_hdr = hdr_kind(probe.color_transfer, dynamic=probe.hdr10plus) if detected_hdr is not None: parts.append(detected_hdr) 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 _primaries_code_for_source(probe: VideoProbe) -> int: return avisynth_primaries_from_color_primaries(probe.color_primaries) def _transfer_code_for_source(probe: VideoProbe) -> int: return avisynth_transfer_from_color_transfer(probe.color_transfer) def _range_code_for_source(probe: VideoProbe) -> int: return avisynth_range_from_color_range(probe.color_range) def _hdr_peak_for_source(probe: VideoProbe) -> float: # Zero means unknown. MBT_ToSDR then uses HDRTools' own HDR-mode default # instead of pretending that an absent mastering display is 1000 cd/m². return probe.hdr_peak_luminance or 0.0 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: if hdr_kind(probe.color_transfer) is not None: return "last" 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 helper_script_path() -> Path: return Path(__file__).resolve().with_name(HELPER_SCRIPT_NAME)