From be1a0d8ac7088b4cb142dcea384b7e6a070d8bcc Mon Sep 17 00:00:00 2001 From: louis Date: Tue, 11 May 2021 23:58:39 -0400 Subject: [PATCH] yt_common: automation: audio handling updates --- .gitignore | 1 + Vivy/01/01.vpy | 4 +- Vivy/02/02.vpy | 4 +- Vivy/03/03.vpy | 4 +- Vivy/04/04.vpy | 6 +- Vivy/05/05.vpy | 4 +- Vivy/06/06.vpy | 4 +- Vivy/07/07.vpy | 4 +- Vivy/setup.py | 2 +- Vivy/vivy_common/filter.py | 4 +- yt_common/setup.py | 2 +- yt_common/yt_common/__init__.py | 2 +- yt_common/yt_common/antialiasing.py | 4 +- yt_common/yt_common/automation.py | 107 +++++++----------- yt_common/yt_common/data.py | 3 + .../yt_common}/shaders/.gitignore | 0 yt_common/yt_common/source.py | 82 ++++++++++++-- 17 files changed, 133 insertions(+), 104 deletions(-) create mode 100644 yt_common/yt_common/data.py rename {Vivy/vivy_common => yt_common/yt_common}/shaders/.gitignore (100%) diff --git a/.gitignore b/.gitignore index 08e83df..10b47ab 100644 --- a/.gitignore +++ b/.gitignore @@ -22,3 +22,4 @@ ffmpeg2pass*.log **/results/ **/__pycache__/ **/*.egg-info/ +**/bdmv/ diff --git a/Vivy/01/01.vpy b/Vivy/01/01.vpy index d18783c..2fdfd1a 100644 --- a/Vivy/01/01.vpy +++ b/Vivy/01/01.vpy @@ -50,7 +50,7 @@ LETTERBOX: List[Range] = [ def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() wakas = [waka[:37301] + core.std.BlankClip(waka, length=6) + waka[37301:] for waka in wakas] waka = wakas[0] waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE) @@ -71,6 +71,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/02/02.vpy b/Vivy/02/02.vpy index ac8efe9..47b620a 100644 --- a/Vivy/02/02.vpy +++ b/Vivy/02/02.vpy @@ -58,7 +58,7 @@ NOAA: List[Range] = PIXELSHIT def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() wakas = [waka[:32344] + waka[32349:] for waka in wakas] waka = wakas[0] waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE) @@ -78,6 +78,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/03/03.vpy b/Vivy/03/03.vpy index 9fb919a..6458ffa 100644 --- a/Vivy/03/03.vpy +++ b/Vivy/03/03.vpy @@ -42,7 +42,7 @@ LETTERBOX: List[Range] = [(0, 432)] def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() waka = wakas[0] waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE) src = bounded_dehardsub(waka, ref, SIGNS_RU) @@ -62,6 +62,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/04/04.vpy b/Vivy/04/04.vpy index f2020a9..f2634a5 100644 --- a/Vivy/04/04.vpy +++ b/Vivy/04/04.vpy @@ -14,8 +14,6 @@ from typing import List import os core = vs.core -core.num_threads = 16 - EPNUM: int = int(os.path.basename(os.path.splitext(__file__)[0])) CONFIG: VivyConfig = VivyConfig(EPNUM) @@ -39,7 +37,7 @@ LETTERBOX: List[Range] = [(0, 1151)] def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() wakas = [w[:33665] + core.std.BlankClip(w, length=21) + w[33665:] for w in wakas] waka = wakas[0] waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE) @@ -62,6 +60,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/05/05.vpy b/Vivy/05/05.vpy index 6ccdbe0..fba378e 100644 --- a/Vivy/05/05.vpy +++ b/Vivy/05/05.vpy @@ -50,7 +50,7 @@ LETTERBOX: List[Range] = [(0, 1153)] def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() waka = wakas[0] waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE) src = bounded_dehardsub(waka, ref, SIGNS_RU, wakas) @@ -71,6 +71,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/06/06.vpy b/Vivy/06/06.vpy index 3cd12d0..8d60b2e 100644 --- a/Vivy/06/06.vpy +++ b/Vivy/06/06.vpy @@ -110,7 +110,7 @@ LETTERBOX: List[Range] = [(0, 2150), (8791, 10693), (13427, 15153), (27878, 2800 def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() wakas = [w[:33669] + w.std.BlankClip(length=17) + w[33669:] for w in wakas] ref = ref.resize.Point(range_in=CRange.FULL, range=CRange.LIMITED) if SOURCE.ref_is_funi else ref waka = wakas[0] @@ -134,6 +134,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/07/07.vpy b/Vivy/07/07.vpy index 1267645..925eec3 100644 --- a/Vivy/07/07.vpy +++ b/Vivy/07/07.vpy @@ -91,7 +91,7 @@ LETTERBOX: List[Range] = [] def filter_basic() -> vs.VideoNode: - wakas, ref = SOURCE.source() + wakas, ref = SOURCE.dhs_source() waka = wakas[0] waka, wakas = waka_replace(waka, wakas[1:], WAKA_REPLACE) src = bounded_dehardsub(waka, ref, SIGNS_RU, wakas) @@ -116,6 +116,6 @@ def filter() -> vs.VideoNode: if __name__ == "__main__": - SelfRunner(CONFIG, filter, filter_basic) + SelfRunner(CONFIG, SOURCE, filter, filter_basic) else: filter() diff --git a/Vivy/setup.py b/Vivy/setup.py index 31a003a..f7320f0 100755 --- a/Vivy/setup.py +++ b/Vivy/setup.py @@ -19,7 +19,7 @@ setuptools.setup( "Operating System :: OS Independent", ], package_data={ - 'vivy_common': ['py.typed', 'workraw-settings', 'final-settings', 'shaders/FSRCNNX_x2_56-16-4-1.glsl'], + 'vivy_common': ['py.typed', 'workraw-settings', 'final-settings'], }, python_requires='>=3.8', ) diff --git a/Vivy/vivy_common/filter.py b/Vivy/vivy_common/filter.py index 08fd861..1a377ae 100644 --- a/Vivy/vivy_common/filter.py +++ b/Vivy/vivy_common/filter.py @@ -12,16 +12,14 @@ from lvsfunc.types import Range from typing import List, Optional, Union from yt_common import antialiasing +from yt_common.data import FSRCNNX from yt_common.denoise import bm3d from yt_common.deband import morpho_mask -import os import vsutil core = vs.core -FSRCNNX = os.path.join(os.path.dirname(__file__), "shaders/FSRCNNX_x2_56-16-4-1.glsl") - def fsrcnnx_rescale(src: vs.VideoNode, noscale: Optional[List[Range]] = None, kernel: Optional[lvf.kernels.Kernel] = None) -> vs.VideoNode: diff --git a/yt_common/setup.py b/yt_common/setup.py index a0cb328..8935a11 100755 --- a/yt_common/setup.py +++ b/yt_common/setup.py @@ -19,7 +19,7 @@ setuptools.setup( "Operating System :: OS Independent", ], package_data={ - 'yt_common': ['py.typed'], + 'yt_common': ['py.typed', 'shaders/FSRCNNX_x2_56-16-4-1.glsl'], }, python_requires='>=3.8', ) diff --git a/yt_common/yt_common/__init__.py b/yt_common/yt_common/__init__.py index a6a8178..99c657c 100644 --- a/yt_common/yt_common/__init__.py +++ b/yt_common/yt_common/__init__.py @@ -1 +1 @@ -from . import antialiasing, automation, config, deband, denoise, logging, source # noqa: F401 +from . import antialiasing, automation, config, data, deband, denoise, logging, source # noqa: F401 diff --git a/yt_common/yt_common/antialiasing.py b/yt_common/yt_common/antialiasing.py index 19845c6..38cf9bc 100644 --- a/yt_common/yt_common/antialiasing.py +++ b/yt_common/yt_common/antialiasing.py @@ -35,9 +35,9 @@ def combine_mask(clip: vs.VideoNode, weak: Union[Range, List[Range], None] = Non def sraa_clamp(clip: vs.VideoNode, mask: Optional[vs.VideoNode] = None, - strength: float = 2, opencl: bool = False, + strength: float = 2, opencl: bool = True, postprocess: Optional[Callable[[vs.VideoNode], vs.VideoNode]] = None) -> vs.VideoNode: - sraa = upscaled_sraa(clip, rfactor=1.3, opencl=opencl, downscaler=Bicubic(b=0, c=1/2).scale) + sraa = upscaled_sraa(clip, rfactor=1.3, nnedi3cl=opencl, downscaler=Bicubic(b=0, c=1/2).scale) sraa = postprocess(sraa) if postprocess else sraa clamp = clamp_aa(clip, nnedi3(clip, opencl=opencl), sraa, strength=strength) return core.std.MaskedMerge(clip, clamp, mask, planes=0) if mask else clamp diff --git a/yt_common/yt_common/automation.py b/yt_common/yt_common/automation.py index 86c79a7..bb083db 100644 --- a/yt_common/yt_common/automation.py +++ b/yt_common/yt_common/automation.py @@ -8,16 +8,15 @@ import shutil import string import subprocess -from typing import Any, BinaryIO, Callable, List, Optional, Sequence, Union, cast +from typing import Any, BinaryIO, Callable, List, Optional, Sequence, cast from .config import Config from .logging import log -from .source import AMAZON_FILENAME_CBR, AMAZON_FILENAME_VBR, ER_FILENAME, SUBSPLS_FILENAME, FUNI_INTRO, glob_filename +from .source import FileSource core = vs.core -AUDIO_OVERRIDE: str = "audio.mka" -AUDIO_CUT: str = "_audiogetter_cut.mka" +AUDIO_ENCODE: str = "_audiogetter_encode.mka" def bin_to_plat(binary: str) -> str: @@ -109,6 +108,7 @@ class AudioGetter(): TODO: really should modularize this a bit instead of assuming amazon->funi """ config: Config + src: FileSource audio_file: str audio_start: int @@ -116,77 +116,41 @@ class AudioGetter(): cleanup: List[str] - def __init__(self, config: Config, override: Optional[str] = None) -> None: + def __init__(self, config: Config, src: FileSource) -> None: self.config = config + self.src = src self.audio_start = 0 self.video_src = None self.cleanup = [] - if override is not None: - if os.path.isfile(override): - self.audio_file = override - else: - raise FileNotFoundError(f"Audio file {override} not found!") + def trim_audio(self, ftrim: Optional[acsuite.types.Trim] = None) -> str: + trims = self.src.get_audio() + if not trims or len(trims) > 1: + raise NotImplementedError("Please implement multifile trimming") + audio_cut = acsuite.eztrim(trims[0].path, cast(acsuite.types.Trim, trims[0].trim))[0] + self.cleanup.append(audio_cut) - # drop "audio.m4a" into the folder and it'll get used - if os.path.isfile(AUDIO_OVERRIDE): - self.audio_file = AUDIO_OVERRIDE - return + if ftrim: + audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source())[0] + self.cleanup.append(audio_cut) - # look for amazon first - if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_CBR)): - self.audio_file = self.config.format_filename(AMAZON_FILENAME_CBR) - self.video_src = core.ffms2.Source(self.audio_file) - log.success("Found Amazon audio") - return + return audio_cut - if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_VBR)): - self.audio_file = self.config.format_filename(AMAZON_FILENAME_VBR) - self.video_src = core.ffms2.Source(self.audio_file) - log.success("Found Amazon audio") - return + def encode_audio(self, path: str, args: List[str]) -> str: + ffmpeg_args = [ + "ffmpeg", + "-hide_banner", "-loglevel", "panic", + "-i", path, + "-y", + "-map", "0:a", + ] + args + [AUDIO_ENCODE] + print("+ " + " ".join(ffmpeg_args)) + subprocess.call(ffmpeg_args) - try: - self.audio_file = glob_filename(self.config.format_filename(SUBSPLS_FILENAME)) - self.video_src = core.ffms2.Source(self.audio_file) - except FileNotFoundError: - pass - try: - self.audio_file = glob_filename(self.config.format_filename(ER_FILENAME)) - self.video_src = core.ffms2.Source(self.audio_file) - except FileNotFoundError: - log.error("Could not find audio") - raise + self.cleanup.append(AUDIO_ENCODE) - self.audio_start = FUNI_INTRO - log.warn("No Amazon audio, falling back to Funi") - - def trim_audio(self, src: vs.VideoNode, - trims: Union[acsuite.Trim, List[acsuite.Trim], None] = None) -> str: - if isinstance(trims, tuple): - trims = [trims] - - if trims is None or len(trims) == 0: - if self.audio_start == 0: - trims = [(None, None)] - else: - trims = [(self.audio_start, None)] - else: - if self.audio_start != 0: - trims = [(s+self.audio_start if s is not None and s >= 0 else s, - e+self.audio_start if e is not None and e > 0 else e) - for s, e in trims] - - if os.path.isfile(AUDIO_CUT): - os.remove(AUDIO_CUT) - - acsuite.eztrim(self.video_src if self.video_src else src, trims, - self.audio_file, AUDIO_CUT, quiet=True) - - self.cleanup.append(AUDIO_CUT) - - return AUDIO_CUT + return AUDIO_ENCODE def do_cleanup(self) -> None: for f in self.cleanup: @@ -196,6 +160,7 @@ class AudioGetter(): class SelfRunner(): config: Config + src: FileSource clip: vs.VideoNode workraw: bool @@ -208,9 +173,11 @@ class SelfRunner(): profile: str - def __init__(self, config: Config, final_filter: Callable[[], vs.VideoNode], - workraw_filter: Optional[Callable[[], vs.VideoNode]] = None) -> None: + def __init__(self, config: Config, source: FileSource, final_filter: Callable[[], vs.VideoNode], + workraw_filter: Optional[Callable[[], vs.VideoNode]] = None, + audio_codec: Optional[List[str]] = None) -> None: self.config = config + self.src = source self.video_clean = False self.audio_clean = False @@ -222,7 +189,6 @@ class SelfRunner(): parser.add_argument("-k", "--keep", help="Keep raw video", action="store_true") parser.add_argument("-b", "--encoder", type=str, help="Override detected encoder binary.") parser.add_argument("-f", "--force", help="Overwrite existing intermediaries.", action="store_true") - parser.add_argument("-a", "--audio", type=str, help="Force audio file") parser.add_argument("-x", "--suffix", type=str, help="Change the suffix of the mux. \ Will be overridden by PROFILE if set.") parser.add_argument("-p", "--profile", type=str, help="Set the encoder profile. \ @@ -283,10 +249,13 @@ class SelfRunner(): start, end) log.status("--- LOOKING FOR AUDIO ---") - self.audio = AudioGetter(self.config, args.audio) + self.audio = AudioGetter(self.config, self.src) log.status("--- TRIMMING AUDIO ---") - self.audio_file = self.audio.trim_audio(self.clip, (start, end)) + self.audio_file = self.audio.trim_audio((start, end)) + if audio_codec: + log.status("--- TRANSCODING AUDIO ---") + self.audio_file = self.audio.encode_audio(self.audio_file, audio_codec) try: log.status("--- MUXING FILE ---") diff --git a/yt_common/yt_common/data.py b/yt_common/yt_common/data.py new file mode 100644 index 0000000..6c8ffb3 --- /dev/null +++ b/yt_common/yt_common/data.py @@ -0,0 +1,3 @@ +import os + +FSRCNNX: str = os.path.join(os.path.dirname(__file__), "shaders/FSRCNNX_x2_56-16-4-1.glsl") diff --git a/Vivy/vivy_common/shaders/.gitignore b/yt_common/yt_common/shaders/.gitignore similarity index 100% rename from Vivy/vivy_common/shaders/.gitignore rename to yt_common/yt_common/shaders/.gitignore diff --git a/yt_common/yt_common/source.py b/yt_common/yt_common/source.py index e4a7c1c..172d2dd 100644 --- a/yt_common/yt_common/source.py +++ b/yt_common/yt_common/source.py @@ -1,15 +1,16 @@ import vapoursynth as vs -import vsutil +from acsuite.types import Trim +from lvsfunc.misc import replace_ranges from lvsfunc.types import Range -import lvsfunc as lvf +from vsutil import depth import glob import os from abc import ABC, abstractmethod -from typing import Any, List, Tuple +from typing import Any, List, NamedTuple, Optional, Tuple, Union from .config import Config from .logging import log @@ -23,6 +24,23 @@ AMAZON_FILENAME_CBR: str = "{title_long} - {epnum:02d} (Amazon Prime CBR {resolu AMAZON_FILENAME_VBR: str = "{title_long} - {epnum:02d} (Amazon Prime VBR {resolution}p).mkv" +class FileTrim(NamedTuple): + path: str + trim: Optional[Trim] + + def apply_trim(self, clip: vs.VideoNode) -> vs.VideoNode: + if self.trim is None: + return clip + s, e = self.trim + if s is None and e is None: + return clip + if s is None: + return clip[:e] + if e is None: + return clip[s:] + return clip[s:e] + + def waka_replace(src: vs.VideoNode, wakas: List[vs.VideoNode], ranges: List[List[Range]] ) -> Tuple[vs.VideoNode, List[vs.VideoNode]]: if len(wakas) == 0: @@ -30,8 +48,8 @@ def waka_replace(src: vs.VideoNode, wakas: List[vs.VideoNode], ranges: List[List new_wakas = [] for waka, r in zip(wakas, ranges): tmp = src - src = lvf.misc.replace_ranges(src, waka, r) - new_wakas.append(lvf.misc.replace_ranges(waka, tmp, r)) + src = replace_ranges(src, waka, r) + new_wakas.append(replace_ranges(waka, tmp, r)) return src, new_wakas @@ -43,7 +61,34 @@ def glob_filename(pattern: str) -> str: return res[0] -class DehardsubFileFinder(ABC): +class FileSource(ABC): + def _open(self, path: str) -> vs.VideoNode: + return depth(core.lsmas.LWLibavSource(path), 16) if path.lower().endswith(".m2ts") \ + else depth(core.ffms2.Source(path), 16) + + @abstractmethod + def get_audio(self) -> List[FileTrim]: + pass + + @abstractmethod + def source(self) -> vs.VideoNode: + pass + + +class SimpleSource(FileSource): + src: List[FileTrim] + + def __init__(self, src: Union[FileTrim, List[FileTrim]]) -> None: + self.src = src if isinstance(src, list) else [src] + + def get_audio(self) -> List[FileTrim]: + return self.src + + def source(self) -> vs.VideoNode: + return depth(core.std.Splice([s.apply_trim(self._open(s.path)) for s in self.src]), 16) + + +class DehardsubFileFinder(FileSource): config: Config def __init__(self, config: Config) -> None: @@ -57,14 +102,17 @@ class DehardsubFileFinder(ABC): def get_ref(self) -> vs.VideoNode: pass - def source(self) -> Tuple[List[vs.VideoNode], vs.VideoNode]: + def source(self) -> vs.VideoNode: + return self.get_ref() + + def dhs_source(self) -> Tuple[List[vs.VideoNode], vs.VideoNode]: wakas: List[vs.VideoNode] = [] for f in [self.config.format_filename(f) for f in self.get_waka_filenames()]: if not os.path.isfile(f): log.warn("Missing a waka!") continue - wakas.append(vsutil.depth(core.ffms2.Source(f), 16)) - ref = vsutil.depth(self.get_ref(), 16) + wakas.append(self._open(f)) + ref = self.get_ref() return wakas, ref @@ -75,12 +123,24 @@ class FunimationSource(DehardsubFileFinder): self.ref_is_funi = False super().__init__(*args, **kwargs) + def get_audio(self) -> List[FileTrim]: + if self.ref_is_funi: + return [FileTrim(self.get_funi_filename(), (FUNI_INTRO, None))] + + if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_CBR)): + return [FileTrim(self.config.format_filename(AMAZON_FILENAME_CBR), None)] + + if os.path.isfile(self.config.format_filename(AMAZON_FILENAME_VBR)): + return [FileTrim(self.config.format_filename(AMAZON_FILENAME_VBR), None)] + + raise FileNotFoundError("Failed to find audio that should exist!") + def get_amazon(self) -> vs.VideoNode: if not os.path.isfile(self.config.format_filename(AMAZON_FILENAME_CBR)): log.warn("Amazon not found, falling back to Funimation") raise FileNotFoundError() log.success("Found Amazon video") - return core.ffms2.Source(self.config.format_filename(AMAZON_FILENAME_CBR)) + return self._open(self.config.format_filename(AMAZON_FILENAME_CBR)) def get_funi_filename(self) -> str: try: @@ -95,7 +155,7 @@ class FunimationSource(DehardsubFileFinder): raise def get_funi(self) -> vs.VideoNode: - return core.ffms2.Source(self.get_funi_filename())[FUNI_INTRO:] + return self._open(self.get_funi_filename())[FUNI_INTRO:] def get_ref(self) -> vs.VideoNode: try: