From ad38400ab9bde18d3793a639f8406f707a48ecb4 Mon Sep 17 00:00:00 2001 From: arch1t3cht Date: Tue, 7 Feb 2023 08:57:31 +0100 Subject: [PATCH] vapoursynth: Improve default scripts and add utility functions Add a utility library that wraps LWLibavSource and can parse its .lwi file to obtain timecodes and keyframes. It also contains a function to generate and save keyframes using WWXD or Scxvid. Update the default scripts to use these functions. --- .gitignore | 1 + automation/meson.build | 4 + automation/vapoursynth/aegisub_vs.py | 201 ++++++++++++++++++ .../win_installer/fragment_automation.iss | 2 + src/libresrc/default_config.json | 4 +- src/libresrc/osx/default_config.json | 4 +- 6 files changed, 212 insertions(+), 4 deletions(-) create mode 100644 automation/vapoursynth/aegisub_vs.py diff --git a/.gitignore b/.gitignore index f0468aa12..ce9cab360 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,6 @@ /tests/data +automation/vapoursynth/__pycache__ packages/desktop/aegisub.desktop packages/desktop/aegisub.desktop.template src/libresrc/bitmap.cpp diff --git a/automation/meson.build b/automation/meson.build index cfa428d57..dc3319658 100644 --- a/automation/meson.build +++ b/automation/meson.build @@ -41,3 +41,7 @@ install_data( 'include/aegisub/unicode.moon', 'include/aegisub/util.moon', install_dir: automation_dir / 'include' / 'aegisub') + +install_data( + 'vapoursynth/aegisub_vs.py', + install_dir: automation_dir / 'vapoursynth') diff --git a/automation/vapoursynth/aegisub_vs.py b/automation/vapoursynth/aegisub_vs.py new file mode 100644 index 000000000..510f5b2ac --- /dev/null +++ b/automation/vapoursynth/aegisub_vs.py @@ -0,0 +1,201 @@ +""" +Utility functions for loading video files into Aegisub using the Vapoursynth +video provider. + +When encountering a file whose file extension is not .py or .vpy, the +Vapoursynth audio and video providers will execute the respective default +script set in Aegisub's configuration, with the following string variables set: +- filename: The path to the file that's being opened. +- __aegi_data, __aegi_dictionary, __aegi_local, __aegi_script, __aegi_temp, __aegi_user: + The values of ?data, ?dictionary, etc. respectively. +- __aegi_vscache: The path to a directory where the Vapoursynth script can + store cache files. This directory is cleaned by Aegisub when it gets too + large (as defined by Aegisub's configuration). + +The provider reads the video from the script's 0-th output node. By default, +the video is assumed to be CFR. The script can pass further information to +Aegisub using the following variables: + - __aegi_timecodes: List[int] | str: The timecodes for the video, or the + path to a timecodes file. + - __aegi_keyframes: List[int] | str: List of frame numbers to load as + keyframes, or the path to a keyframes file. + - __aegi_hasaudio: int: If nonzero, Aegisub will try to load an audio track + from the same file. + +This module provides some utility functions to obtain timecodes, keyframes, and +other data. +""" +import os +import os.path +import re +from collections import deque +from typing import Any, Dict, List, Tuple + +import vapoursynth as vs +core = vs.core + +def make_lwi_cache_filename(filename: str) -> str: + """ + Given a path to a video, will return a file name like the one LWLibavSource + would use for a .lwi file. + """ + max_len = 254 + extension = ".lwi" + + if len(filename) + len(extension) > max_len: + filename = filename[-(max_len + len(extension)):] + + return "".join(("_" if c in "/\\:" else c) for c in filename) + extension + + +def make_keyframes_filename(filename: str) -> str: + """ + Given a path `path/to/file.mkv`, will return the path + `path/to/file_keyframes.txt`. + """ + extlen = filename[::-1].find(".") + 1 + return filename[:len(filename) - extlen] + "_keyframes.txt" + + +lwindex_re1 = re.compile(r"Index=(?P-?[0-9]+),POS=(?P-?[0-9]+),PTS=(?P-?[0-9]+),DTS=(?P-?[0-9]+),EDI=(?P-?[0-9]+)") +lwindex_re2 = re.compile(r"Key=(?P-?[0-9]+),Pic=(?P-?[0-9]+),POC=(?P-?[0-9]+),Repeat=(?P-?[0-9]+),Field=(?P-?[0-9]+)") +streaminfo_re = re.compile(r"Codec=(?P[0-9]+),TimeBase=(?P[0-9\/]+),Width=(?P[0-9]+),Height=(?P[0-9]+),Format=(?P[0-9a-zA-Z]+),ColorSpace=(?P[0-9]+)") + +class LWIndexFrame: + pts: int + key: int + + def __init__(self, raw: list[str]): + match1 = lwindex_re1.match(raw[0]) + match2 = lwindex_re2.match(raw[1]) + if not match1 or not match2: + raise ValueError("Invalid lwindex format") + self.pts = int(match1.group("PTS")) + self.key = int(match2.group("Key")) + + def __int__(self) -> int: + return self.pts + + +def info_from_lwindex(indexfile: str) -> Dict[str, List[int]]: + """ + Given a path to an .lwi file, will return a dictionary containing + information about the video, with the keys + - timcodes: The timecodes. + - keyframes: Array of frame numbers of keyframes. + """ + with open(indexfile, encoding="latin1") as f: + index = f.read().splitlines() + + indexstart, indexend = index.index("") + 1, index.index("") + frames = [LWIndexFrame(index[i:i+2]) for i in range(indexstart, indexend, 2)] + frames.sort(key=int) + + streaminfo = streaminfo_re.match(index[indexstart - 2]) + if not streaminfo: + raise ValueError("Invalid lwindex format") + + timebase_num, timebase_den = [int(i) for i in streaminfo.group("TimeBase").split("/")] + + return { + "timecodes": [(f.pts * 1000 * timebase_num) // timebase_den for f in frames], + "keyframes": [i for i, f in enumerate(frames) if f.key], + } + + +def wrap_lwlibavsource(filename: str, cachedir: str, **kwargs: Any) -> Tuple[vs.VideoNode, Dict[str, List[int]]]: + """ + Given a path to a video file and a directory to store index files in + (usually __aegi_vscache), will open the video with LWLibavSource and read + the generated .lwi file to obtain the timecodes and keyframes. + Additional keyword arguments are passed on to LWLibavSource. + """ + try: + os.mkdir(cachedir) + except FileExistsError: + pass + cachefile = os.path.join(cachedir, make_lwi_cache_filename(filename)) + + if not hasattr(core, "lsmas"): + raise vs.Error("To use Aegisub's LWLibavSource wrapper, the `lsmas` plugin for VapourSynth must be installed") + + if b"-Dcachedir" not in core.lsmas.Version()["config"]: # type: ignore + raise vs.Error("To use Aegisub's LWLibavSource wrapper, the `lsmas` plugin must support the `cachedir` option for LWLibavSource.") + + clip = core.lsmas.LWLibavSource(source=filename, cachefile=cachefile, **kwargs) + + return clip, info_from_lwindex(cachefile) + + +def make_keyframes(clip: vs.VideoNode, use_scxvid: bool = False, + resize_h: int = 360, resize_format: int = vs.YUV420P8, + **kwargs: Any) -> List[int]: + """ + Generates a list of keyframes from a clip, using either WWXD or Scxvid. + Will be slightly more efficient with the `akarin` plugin installed. + + :param clip: Clip to process. + :param use_scxvid: Whether to use Scxvid. If False, the function uses WWXD. + :param resize_h: Height to resize the clip to before processing. + :param resize_format: Format to convert the clip to before processing. + + The remaining keyword arguments are passed on to the respective filter. + """ + + clip = core.resize.Bilinear(clip, width=resize_h * clip.width // clip.height, height=resize_h, format=resize_format); + try: + clip = core.scxvid.Scxvid(clip, **kwargs) if use_scxvid else core.wwxd.WWXD(clip, **kwargs) + except AttributeError: + raise vs.Error("To use the keyframe generation, the `{}` plugin for VapourSynth must be installed" + .format("scxvid" if use_scxvid else "wwxd")) + + keyframes = {} + def _cb(n: int, f: vs.VideoFrame) -> vs.VideoFrame: + keyframes[n] = f.props._SceneChangePrev if use_scxvid else f.props.Scenechange # type: ignore + return f + + deque(clip.std.ModifyFrame(clip, _cb).frames(close=True), 0) + return [n for n in range(clip.num_frames) if keyframes[n]] + + +def save_keyframes(filename: str, keyframes: List[int]): + """ + Saves a list of keyframes in Aegisub's keyframe format v1 to a file with + the given filename. + """ + with open(filename, "w") as f: + f.write("# keyframe format v1\n") + f.write("fps 0\n") + f.write("".join(f"{n}\n" for n in keyframes)) + + +def get_keyframes(filename: str, clip: vs.VideoNode, **kwargs: Any) -> str: + """ + When not already present, creates a keyframe file for the given clip next + to the given filename using WWXD or Scxvid (see the make_keyframes docstring). + Additional keyword arguments are passed on to make_keyframes. + """ + kffilename = make_keyframes_filename(filename) + + if not os.path.exists(kffilename): + keyframes = make_keyframes(clip, **kwargs) + save_keyframes(kffilename, keyframes) + + return kffilename + + +def check_audio(filename: str, **kwargs: Any) -> bool: + """ + Checks whether the given file has an audio track by trying to open it with + BestAudioSource. Requires the `bas` plugin to return correct results, but + won't crash if it's not installed. + Additional keyword arguments are passed on to BestAudioSource. + """ + try: + vs.core.bas.Source(source=filename, **kwargs) + return True + except AttributeError: + pass + except vs.Error: + pass + return False diff --git a/packages/win_installer/fragment_automation.iss b/packages/win_installer/fragment_automation.iss index 6843a3388..fc52a96a9 100644 --- a/packages/win_installer/fragment_automation.iss +++ b/packages/win_installer/fragment_automation.iss @@ -30,6 +30,8 @@ DestDir: {app}\automation\include; Source: {#SOURCE_ROOT}\automation\include\uni DestDir: {app}\automation\include; Source: {#SOURCE_ROOT}\automation\include\utils.lua; Flags: ignoreversion overwritereadonly uninsremovereadonly; Attribs: readonly; Components: main DestDir: {app}\automation\include; Source: {#SOURCE_ROOT}\automation\include\utils-auto4.lua; Flags: ignoreversion overwritereadonly uninsremovereadonly; Attribs: readonly; Components: main +DestDir: {app}\automation\vapoursynth; Source: {#SOURCE_ROOT}\automation\vapoursynth\aegisub_vs.py; Flags: ignoreversion overwritereadonly uninsremovereadonly; Attribs: readonly; Components: main + #ifdef DEPCTRL ; DepCtrl DestDir: {userappdata}\Aegisub\automation\include\l0; Source: {#DEPS_DIR}\DependencyControl\modules\*; Flags: ignoreversion recursesubdirs createallsubdirs; Components: macros\modules\depctrl diff --git a/src/libresrc/default_config.json b/src/libresrc/default_config.json index dbdc303a3..5daae04bd 100644 --- a/src/libresrc/default_config.json +++ b/src/libresrc/default_config.json @@ -338,7 +338,7 @@ "Aegisub Cache" : true }, "VapourSynth" : { - "Default Script" : "import vapoursynth as vs\nvs.core.bas.Source(source=filename).set_output()" + "Default Script" : "# This default script will load an audio file using BestAudioSource.\n# It requires the `bas` plugin.\n\nimport vapoursynth as vs\ntry:\n vs.core.bas.Source(source=filename).set_output()\nexcept AttributeError:\n raise vs.Error(\"To use Aegisub's default audio loader, the `bas` plugin for VapourSynth must be installed\")" } }, "Avisynth" : { @@ -373,7 +373,7 @@ "Seek Preroll" : 12 }, "VapourSynth" : { - "Default Script" : "import vapoursynth as vs\nvs.core.lsmas.LWLibavSource(source=filename).set_output()" + "Default Script" : "# This default script will load a video file using LWLibavSource.\n# It requires the `lsmas` plugin.\n# See ?data/automation/vapoursynth/aegisub_vs.py for more information.\n\nimport aegisub_vs as a\nimport vapoursynth as vs\n\nclip, videoinfo = a.wrap_lwlibavsource(filename, __aegi_vscache)\nclip.set_output()\n__aegi_timecodes = videoinfo[\"timecodes\"]\n__aegi_keyframes = videoinfo[\"keyframes\"]\n# Uncomment this to automatically generate keyframes at scene changes.\n#__aegi_keyframes = a.get_keyframes(filename, clip)\n\n# Check if the file has an audio track. This requires the `bas` plugin.\n__aegi_hasaudio = 1 if a.check_audio(filename) else 0" } } }, diff --git a/src/libresrc/osx/default_config.json b/src/libresrc/osx/default_config.json index 412f13f50..c819f2338 100644 --- a/src/libresrc/osx/default_config.json +++ b/src/libresrc/osx/default_config.json @@ -338,7 +338,7 @@ "Aegisub Cache" : true }, "VapourSynth" : { - "Default Script" : "import vapoursynth as vs\nvs.core.bas.Source(source=filename).set_output()" + "Default Script" : "# This default script will load an audio file using BestAudioSource.\n# It requires the `bas` plugin.\n\nimport vapoursynth as vs\ntry:\n vs.core.bas.Source(source=filename).set_output()\nexcept AttributeError:\n raise vs.Error(\"To use Aegisub's default audio loader, the `bas` plugin for VapourSynth must be installed\")" } }, "Avisynth" : { @@ -373,7 +373,7 @@ "Seek Preroll" : 12 }, "VapourSynth" : { - "Default Script" : "import vapoursynth as vs\nvs.core.lsmas.LWLibavSource(source=filename).set_output()" + "Default Script" : "# This default script will load a video file using LWLibavSource.\n# It requires the `lsmas` plugin.\n# See ?data/automation/vapoursynth/aegisub_vs.py for more information.\n\nimport aegisub_vs as a\nimport vapoursynth as vs\n\nclip, videoinfo = a.wrap_lwlibavsource(filename, __aegi_vscache)\nclip.set_output()\n__aegi_timecodes = videoinfo[\"timecodes\"]\n__aegi_keyframes = videoinfo[\"keyframes\"]\n# Uncomment this to automatically generate keyframes at scene changes.\n#__aegi_keyframes = a.get_keyframes(filename, clip)\n\n# Check if the file has an audio track. This requires the `bas` plugin.\n__aegi_hasaudio = 1 if a.check_audio(filename) else 0" } } },