From ef1c61d966a8dd304e65f411827c4299a75a9462 Mon Sep 17 00:00:00 2001 From: louis Date: Sat, 5 Jun 2021 22:44:37 -0400 Subject: [PATCH] yt_common: split up automation.py --- .../S2 [BD]/01/01.vpy | 3 +- .../S2 [BD]/02/02.vpy | 1 + .../S2 [BD]/03/03.vpy | 3 +- .../S2 [BD]/04/04.vpy | 1 + .../S2 [BD]/05/05.vpy | 1 + .../S2 [BD]/06/06.vpy | 3 +- .../S2 [BD]/07/07.vpy | 3 +- .../S2 [BD]/08/08.vpy | 3 +- .../S2 [BD]/09/09.vpy | 1 + .../S2 [BD]/10/10.vpy | 1 + .../S2 [BD]/11/11.vpy | 3 +- .../S2 [BD]/12/12.vpy | 3 +- .../S2 [BD]/nced/nced.vpy | 1 + .../S2 [BD]/ncop/ncop.vpy | 1 + .../S2 [BD]/recap/recap.vpy | 3 +- yt_common/yt_common/__init__.py | 3 +- yt_common/yt_common/audio.py | 69 ++++++- yt_common/yt_common/automation.py | 189 +----------------- yt_common/yt_common/util.py | 8 + yt_common/yt_common/video.py | 109 ++++++++++ 20 files changed, 218 insertions(+), 191 deletions(-) create mode 100644 yt_common/yt_common/video.py diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/01/01.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/01/01.vpy index 3865466..2ec072c 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/01/01.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/01/01.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/02/02.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/02/02.vpy index dcd2ecf..a3d2f8d 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/02/02.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/02/02.vpy @@ -6,6 +6,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.misc import replace_ranges from lvsfunc.types import Range diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/03/03.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/03/03.vpy index 0f4e565..c9e6c58 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/03/03.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/03/03.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.misc import replace_ranges from lvsfunc.types import Range diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/04/04.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/04/04.vpy index 43fcbb9..bd7661b 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/04/04.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/04/04.vpy @@ -6,6 +6,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/05/05.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/05/05.vpy index 8d4ee79..b608c0b 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/05/05.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/05/05.vpy @@ -6,6 +6,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/06/06.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/06/06.vpy index dc4941f..08ef85e 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/06/06.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/06/06.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/07/07.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/07/07.vpy index 9400b26..9dbca5c 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/07/07.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/07/07.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/08/08.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/08/08.vpy index f14f49c..57722c3 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/08/08.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/08/08.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/09/09.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/09/09.vpy index acdcc10..de63c8a 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/09/09.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/09/09.vpy @@ -6,6 +6,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/10/10.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/10/10.vpy index 84e9c4f..3c9df69 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/10/10.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/10/10.vpy @@ -6,6 +6,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.chapters import Chapter, Edition from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/11/11.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/11/11.vpy index 8726b1c..7618445 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/11/11.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/11/11.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter, Edition from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/12/12.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/12/12.vpy index 56a6105..5bb6e5f 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/12/12.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/12/12.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/nced/nced.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/nced/nced.vpy index 0d7da1c..b046f42 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/nced/nced.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/nced/nced.vpy @@ -5,6 +5,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/ncop/ncop.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/ncop/ncop.vpy index b94e4cb..b4cf5be 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/ncop/ncop.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/ncop/ncop.vpy @@ -5,6 +5,7 @@ from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deban from yt_common.automation import SelfRunner from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/Tensei Shitara Slime Datta Ken/S2 [BD]/recap/recap.vpy b/Tensei Shitara Slime Datta Ken/S2 [BD]/recap/recap.vpy index 6ec5847..1501c3c 100644 --- a/Tensei Shitara Slime Datta Ken/S2 [BD]/recap/recap.vpy +++ b/Tensei Shitara Slime Datta Ken/S2 [BD]/recap/recap.vpy @@ -3,9 +3,10 @@ import vapoursynth as vs from tensura_common import (TenSuraS2Config, TenSuraS2BDSource, antialias, deband, denoise, descale, edgefix, finalize, megurumono_scenefilter, regrain) -from yt_common.automation import SelfRunner, Zone +from yt_common.automation import SelfRunner from yt_common.chapters import Chapter from yt_common.source import FileTrim +from yt_common.video import Zone from lvsfunc.mask import BoundingBox from lvsfunc.misc import replace_ranges diff --git a/yt_common/yt_common/__init__.py b/yt_common/yt_common/__init__.py index 5082dbb..b278d9a 100644 --- a/yt_common/yt_common/__init__.py +++ b/yt_common/yt_common/__init__.py @@ -1 +1,2 @@ -from . import antialiasing, audio, automation, chapters, config, data, deband, logging, scale, source # noqa: F401 +from . import (antialiasing, audio, automation, chapters, config, data, deband, # noqa: F401 + logging, scale, source, video) diff --git a/yt_common/yt_common/audio.py b/yt_common/yt_common/audio.py index ee25b1b..9bbf0fc 100644 --- a/yt_common/yt_common/audio.py +++ b/yt_common/yt_common/audio.py @@ -1,9 +1,19 @@ +import acsuite +import os + from abc import ABC, abstractmethod from subprocess import call -from typing import List, NamedTuple +from typing import TYPE_CHECKING, List, Set, NamedTuple, Optional +from .config import Config from .util import get_temp_filename +if TYPE_CHECKING: + from .source import FileSource + + +AUDIO_PFX: str = "_audio_temp_" + class AudioEncoder(ABC): @abstractmethod @@ -63,3 +73,60 @@ class CodecFlac(FFAudio): class AudioStream(NamedTuple): stream_index: int # zero-indexed, ignores video streams codec: AudioEncoder + name: str = "" + language: str = "jpn" + + +class AudioTrimmer(): + config: Config + src: "FileSource" + cleanup: Set[str] + + def __init__(self, config: Config, src: "FileSource") -> None: + self.config = config + self.src = src + self.cleanup = set() + + def trim_audio(self, ftrim: Optional[acsuite.types.Trim] = None) -> str: + streams = sorted(self.src.audio_streams(), key=lambda s: s.stream_index) + if len(streams) == 0: + return "" + trims = self.src.audio_src() + ffmpeg = acsuite.ffmpeg.FFmpegAudio() + + tlist: List[str] = [] + for t in trims: + audio_cut = acsuite.eztrim(t.path, t.trim or (0, None), ref_clip=self.src.audio_ref(), + outfile=get_temp_filename(prefix=AUDIO_PFX+"cut_", suffix=".mka"), + streams=[s.stream_index for s in streams])[0] + self.cleanup.add(audio_cut) + tlist.append(audio_cut) + + if len(tlist) > 1: + audio_cut = ffmpeg.concat(*tlist) + self.cleanup.add(audio_cut) + + if ftrim: + audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source(), + outfile=get_temp_filename(prefix=AUDIO_PFX+"fcut_", suffix=".mka"))[0] + self.cleanup.add(audio_cut) + + if len(streams) > 1: + splits = [get_temp_filename(prefix=AUDIO_PFX+"split_", suffix=".mka") for _ in range(0, len(streams))] + ffmpeg.split(audio_cut, splits) + self.cleanup |= set(splits) + encode = [streams[i].codec.encode_audio(f) for i, f in enumerate(splits)] + self.cleanup |= set(encode) + audio_cut = get_temp_filename(prefix=AUDIO_PFX+"join_", suffix=".mka") + ffmpeg.join(audio_cut, *encode) + else: + audio_cut = streams[0].codec.encode_audio(audio_cut) + + self.cleanup.add(audio_cut) + + return audio_cut + + def do_cleanup(self) -> None: + for f in self.cleanup: + os.remove(f) + self.cleanup.clear() diff --git a/yt_common/yt_common/automation.py b/yt_common/yt_common/automation.py index 3e226da..f56e1ee 100644 --- a/yt_common/yt_common/automation.py +++ b/yt_common/yt_common/automation.py @@ -1,197 +1,24 @@ import vapoursynth as vs -import acsuite import argparse import os import random import shutil -import string import subprocess -from lvsfunc.render import clip_async_render, find_scene_changes +from lvsfunc.render import find_scene_changes -from typing import Any, BinaryIO, Callable, List, NamedTuple, Optional, Sequence, Set, Tuple, cast +from typing import Callable, List, Optional +from .audio import AudioTrimmer from .chapters import Chapter, Edition, make_chapters, make_qpfile from .config import Config from .logging import log -from .util import get_temp_filename from .source import FileSource +from .video import VideoEncoder, Zone core = vs.core -AUDIO_PFX: str = "_audiogetter_temp_" - - -def bin_to_plat(binary: str) -> str: - if os.name == "nt": - return binary if binary.lower().endswith(".exe") else f"{binary}.exe" - else: - return binary if not binary.lower().endswith(".exe") else binary[:-len(".exe")] - - -def forward_signal(signum: int, frame: Any, process: Any) -> None: - log.warn("Forwarding SIGINT") - process.send_signal(signum) - - -class Zone(NamedTuple): - r: Tuple[int, int] - b: float - - -class Encoder(): - clip: vs.VideoNode - - binary: str - params: Sequence[str] - force: bool - - out_template: str - - cleanup: List[str] - - def __init__(self, settings_path: str, binary: Optional[str] = None, force: bool = False) -> None: - self.binary = binary if binary is not None else "" - self.force = force - self.cleanup = [] - - self._get_encoder_settings(settings_path) - - def encode(self, clip: vs.VideoNode, filename: str, start: int = 0, end: int = 0, - zones: Optional[List[Zone]] = None, qpfile: Optional[str] = None, - timecode_file: Optional[str] = None, want_timecodes: bool = False) -> Tuple[str, List[float]]: - end = end if end != 0 else clip.num_frames - want_timecodes = True if timecode_file else want_timecodes - - outfile = self.out_template.format(filename=filename) - - if os.path.isfile(outfile) and not self.force: - log.warn("Existing output detected, skipping encode!") - return outfile, [] - - params: List[str] = [] - for p in self.params: - if p == "$ZONES": - if zones: - zones.sort(key=lambda z: z.r[0]) - params.append("--zones") - zargs: List[str] = [] - for z in zones: - if z.r[0] - start >= 0 and z.r[0] < end: - s = z.r[0] - start - e = z.r[1] - start - e = e if e < end - start else end - start - 1 - zargs.append(f"{s},{e},b={z.b}") - params.append("/".join(zargs)) - elif p == "$QPFILE": - if qpfile: - params += ["--qpfile", qpfile] - else: - params.append(p.format(frames=end-start, filename=filename, qpfile="qpfile.txt")) - - log.status("--- RUNNING ENCODE ---") - - print("+ " + " ".join([self.binary] + list(params))) - - process = subprocess.Popen([self.binary] + list(params), stdin=subprocess.PIPE) - - # i want the encoder to handle any ctrl-c so it exits properly - # forward_to_proc = functools.partial(forward_signal, process=process) - # signal.signal(signal.SIGINT, forward_to_proc) - # turns out this didn't work out the way i had hoped - - # use the python renderer only if we need timecodes because it's slower - timecodes: List[float] = [] - if want_timecodes: - timecode_io = open(timecode_file, "w") if timecode_file else None - timecodes = clip_async_render(clip[start:end], cast(BinaryIO, process.stdin), timecode_io) - else: - clip[start:end].output(cast(BinaryIO, process.stdin), y4m=True) - process.communicate() - - # vapoursynth should handle this itself but just in case - if process.returncode != 0: - log.error("--- ENCODE FAILED ---") - raise BrokenPipeError(f"Pipe to {self.binary} broken") - - log.success("--- ENCODE FINISHED ---") - self.cleanup.append(outfile) - return outfile, timecodes - - def _get_encoder_settings(self, settings_path: str) -> None: - with open(settings_path, "r") as settings: - keys = " ".join([line.strip() for line in settings if not line.strip().startswith("#")]).split(" ") - - # verify that the settings contain an output file template - outputs = [k for k in keys[1:] if any([name == "filename" for _, name, _, _ in string.Formatter().parse(k)])] - if not outputs or len(outputs) > 1: - raise Exception("Failed to find unambiguous output file for encoder!") - self.out_template = outputs[0] - - self.binary = bin_to_plat(keys[0]) if not self.binary else self.binary - self.params = keys[1:] - - def do_cleanup(self) -> None: - for f in self.cleanup: - os.remove(f) - self.cleanup = [] - - -class AudioGetter(): - config: Config - src: FileSource - cleanup: Set[str] - - def __init__(self, config: Config, src: FileSource) -> None: - self.config = config - self.src = src - self.cleanup = set() - - def trim_audio(self, ftrim: Optional[acsuite.types.Trim] = None) -> str: - streams = sorted(self.src.audio_streams(), key=lambda s: s.stream_index) - if len(streams) == 0: - return "" - trims = self.src.audio_src() - ffmpeg = acsuite.ffmpeg.FFmpegAudio() - - tlist: List[str] = [] - for t in trims: - audio_cut = acsuite.eztrim(t.path, t.trim or (0, None), ref_clip=self.src.audio_ref(), - outfile=get_temp_filename(prefix=AUDIO_PFX+"cut_", suffix=".mka"), - streams=[s.stream_index for s in streams])[0] - self.cleanup.add(audio_cut) - tlist.append(audio_cut) - - if len(tlist) > 1: - audio_cut = ffmpeg.concat(*tlist) - self.cleanup.add(audio_cut) - - if ftrim: - audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source(), - outfile=get_temp_filename(prefix=AUDIO_PFX+"fcut_", suffix=".mka"))[0] - self.cleanup.add(audio_cut) - - if len(streams) > 1: - splits = [get_temp_filename(prefix=AUDIO_PFX+"split_", suffix=".mka") for _ in range(0, len(streams))] - ffmpeg.split(audio_cut, splits) - self.cleanup |= set(splits) - encode = [streams[i].codec.encode_audio(f) for i, f in enumerate(splits)] - self.cleanup |= set(encode) - audio_cut = get_temp_filename(prefix=AUDIO_PFX+"join_", suffix=".mka") - ffmpeg.join(audio_cut, *encode) - else: - audio_cut = streams[0].codec.encode_audio(audio_cut) - - self.cleanup.add(audio_cut) - - return audio_cut - - def do_cleanup(self) -> None: - for f in self.cleanup: - os.remove(f) - self.cleanup.clear() - class SelfRunner(): config: Config @@ -206,8 +33,8 @@ class SelfRunner(): timecode_file: Optional[str] qpfile: Optional[str] - encoder: Encoder - audio: AudioGetter + encoder: VideoEncoder + audio: AudioTrimmer profile: str @@ -313,7 +140,7 @@ class SelfRunner(): if not os.path.isfile(settings_path): raise FileNotFoundError(f"Failed to find {settings_path}!") - self.encoder = Encoder(settings_path, args.encoder, args.force) + self.encoder = VideoEncoder(settings_path, args.encoder, args.force) # we only want to generate timecodes if vfr, otherwise we can just calculate them self.timecode_file = f"{self.config.desc}_{self.suffix}_{start}_{end}_timecodes.txt" \ if self.clip.fps_den == 0 else None @@ -351,7 +178,7 @@ class SelfRunner(): def _do_audio(self, start: int, end: int, out_name: Optional[str] = None) -> None: log.status("--- LOOKING FOR AUDIO ---") - self.audio = AudioGetter(self.config, self.src) + self.audio = AudioTrimmer(self.config, self.src) log.status("--- TRIMMING AUDIO ---") self.audio_file = self.audio.trim_audio((start, end)) diff --git a/yt_common/yt_common/util.py b/yt_common/yt_common/util.py index d808066..0c966cd 100644 --- a/yt_common/yt_common/util.py +++ b/yt_common/yt_common/util.py @@ -1,5 +1,13 @@ +import os import tempfile def get_temp_filename(prefix: str = "", suffix: str = "") -> str: return f"{prefix}{next(tempfile._get_candidate_names())}{suffix}" # type: ignore + + +def bin_to_plat(binary: str) -> str: + if os.name == "nt": + return binary if binary.lower().endswith(".exe") else f"{binary}.exe" + else: + return binary if not binary.lower().endswith(".exe") else binary[:-len(".exe")] diff --git a/yt_common/yt_common/video.py b/yt_common/yt_common/video.py new file mode 100644 index 0000000..f2ddb01 --- /dev/null +++ b/yt_common/yt_common/video.py @@ -0,0 +1,109 @@ +import vapoursynth as vs +import os +import string +import subprocess + +from lvsfunc.render import clip_async_render + +from typing import BinaryIO, List, NamedTuple, Optional, Sequence, Tuple, cast + +from .logging import log +from .util import bin_to_plat + + +class Zone(NamedTuple): + r: Tuple[int, int] + b: float + + +class VideoEncoder(): + clip: vs.VideoNode + + binary: str + params: Sequence[str] + force: bool + + out_template: str + + cleanup: List[str] + + def __init__(self, settings_path: str, binary: Optional[str] = None, force: bool = False) -> None: + self.binary = binary if binary is not None else "" + self.force = force + self.cleanup = [] + + self._get_encoder_settings(settings_path) + + def encode(self, clip: vs.VideoNode, filename: str, start: int = 0, end: int = 0, + zones: Optional[List[Zone]] = None, qpfile: Optional[str] = None, + timecode_file: Optional[str] = None, want_timecodes: bool = False) -> Tuple[str, List[float]]: + end = end if end != 0 else clip.num_frames + want_timecodes = True if timecode_file else want_timecodes + + outfile = self.out_template.format(filename=filename) + + if os.path.isfile(outfile) and not self.force: + log.warn("Existing output detected, skipping encode!") + return outfile, [] + + params: List[str] = [] + for p in self.params: + if p == "$ZONES": + if zones: + zones.sort(key=lambda z: z.r[0]) + params.append("--zones") + zargs: List[str] = [] + for z in zones: + if z.r[0] - start >= 0 and z.r[0] < end: + s = z.r[0] - start + e = z.r[1] - start + e = e if e < end - start else end - start - 1 + zargs.append(f"{s},{e},b={z.b}") + params.append("/".join(zargs)) + elif p == "$QPFILE": + if qpfile: + params += ["--qpfile", qpfile] + else: + params.append(p.format(frames=end-start, filename=filename, qpfile="qpfile.txt")) + + log.status("--- RUNNING ENCODE ---") + + print("+ " + " ".join([self.binary] + list(params))) + + process = subprocess.Popen([self.binary] + list(params), stdin=subprocess.PIPE) + + # use the python renderer only if we need timecodes because it's slower + timecodes: List[float] = [] + if want_timecodes: + timecode_io = open(timecode_file, "w") if timecode_file else None + timecodes = clip_async_render(clip[start:end], cast(BinaryIO, process.stdin), timecode_io) + else: + clip[start:end].output(cast(BinaryIO, process.stdin), y4m=True) + process.communicate() + + # vapoursynth should handle this itself but just in case + if process.returncode != 0: + log.error("--- ENCODE FAILED ---") + raise BrokenPipeError(f"Pipe to {self.binary} broken") + + log.success("--- ENCODE FINISHED ---") + self.cleanup.append(outfile) + return outfile, timecodes + + def _get_encoder_settings(self, settings_path: str) -> None: + with open(settings_path, "r") as settings: + keys = " ".join([line.strip() for line in settings if not line.strip().startswith("#")]).split(" ") + + # verify that the settings contain an output file template + outputs = [k for k in keys[1:] if any([name == "filename" for _, name, _, _ in string.Formatter().parse(k)])] + if not outputs or len(outputs) > 1: + raise Exception("Failed to find unambiguous output file for encoder!") + self.out_template = outputs[0] + + self.binary = bin_to_plat(keys[0]) if not self.binary else self.binary + self.params = keys[1:] + + def do_cleanup(self) -> None: + for f in self.cleanup: + os.remove(f) + self.cleanup = []