From ffcdf94fc40b245fd53fed75765c8a546693c6c2 Mon Sep 17 00:00:00 2001 From: louis Date: Sun, 30 May 2021 19:36:17 -0400 Subject: [PATCH] yt_common: chapters and audio mixins --- yt_common/yt_common/__init__.py | 2 +- yt_common/yt_common/automation.py | 48 ++++++++++------------- yt_common/yt_common/chapters.py | 63 +++++++++++++++++++++++++++++++ yt_common/yt_common/config.py | 53 ++++++++++++++++++++++++-- yt_common/yt_common/source.py | 5 ++- 5 files changed, 136 insertions(+), 35 deletions(-) create mode 100644 yt_common/yt_common/chapters.py diff --git a/yt_common/yt_common/__init__.py b/yt_common/yt_common/__init__.py index 6b683c2..57a5874 100644 --- a/yt_common/yt_common/__init__.py +++ b/yt_common/yt_common/__init__.py @@ -1 +1 @@ -from . import antialiasing, automation, config, data, deband, logging, source # noqa: F401 +from . import antialiasing, automation, chapters, config, data, deband, logging, scale, source # noqa: F401 diff --git a/yt_common/yt_common/automation.py b/yt_common/yt_common/automation.py index 53113f1..f78616d 100644 --- a/yt_common/yt_common/automation.py +++ b/yt_common/yt_common/automation.py @@ -10,8 +10,9 @@ import subprocess from lvsfunc.render import clip_async_render -from typing import Any, BinaryIO, Callable, List, Optional, Sequence, Tuple, cast +from typing import Any, BinaryIO, Callable, List, Optional, Sequence, Set, Tuple, cast +from .chapters import Chapter, make_chapters from .config import Config from .logging import log from .source import FileSource @@ -115,45 +116,33 @@ class Encoder(): class AudioGetter(): config: Config src: FileSource - cleanup: List[str] + cleanup: Set[str] def __init__(self, config: Config, src: FileSource) -> None: self.config = config self.src = src - self.cleanup = [] + self.cleanup = set() 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, trims[0].trim or (0, None), streams=0)[0] - self.cleanup.append(audio_cut) + self.cleanup.add(audio_cut) if ftrim: audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source())[0] - self.cleanup.append(audio_cut) + self.cleanup.add(audio_cut) + + audio_cut = self.config.encode_audio(audio_cut) + self.cleanup.add(audio_cut) return audio_cut - def encode_audio(self, path: str, codec_args: List[str]) -> str: - ffmpeg_args = [ - "ffmpeg", - "-hide_banner", "-loglevel", "panic", - "-i", path, - "-y", - "-map", "0:a", - ] + codec_args + [AUDIO_ENCODE] - print("+ " + " ".join(ffmpeg_args)) - subprocess.call(ffmpeg_args) - - self.cleanup.append(AUDIO_ENCODE) - - return AUDIO_ENCODE - def do_cleanup(self) -> None: for f in self.cleanup: os.remove(f) - self.cleanup = [] + self.cleanup.clear() class SelfRunner(): @@ -175,7 +164,7 @@ class SelfRunner(): 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: + chapters: Optional[List[Chapter]] = None) -> None: self.config = config self.src = source self.video_clean = False @@ -230,11 +219,11 @@ class SelfRunner(): if start >= end: raise ValueError("Start frame must be before end frame!") - out_name = f"{self.config.title.lower()}_{self.config.desc}_{self.suffix}" + out_name = f"{self.config.title.lower().replace(' ', '_')}_{self.config.desc}_{self.suffix}" if args.audio_only: out_name += ".mka" - self._do_audio(start, audio_end, audio_codec, out_name=out_name) + self._do_audio(start, audio_end, out_name=out_name) self.audio.do_cleanup() log.success("--- AUDIO ENCODE COMPLETE ---") return @@ -269,7 +258,11 @@ class SelfRunner(): self.timecodes = [round(float(1e9*f*(1/self.clip.fps)))/1e9 for f in range(0, self.clip.num_frames + 1)] \ if self.clip.fps_den != 0 and len(self.timecodes) == 0 else self.timecodes - self._do_audio(start, audio_end, audio_codec) + self._do_audio(start, audio_end) + + if chapters: + log.status("--- GENERATING CHAPTERS ---") + make_chapters(chapters, self.timecodes, f"{self.config.desc}.xml") try: log.status("--- MUXING FILE ---") @@ -289,15 +282,12 @@ class SelfRunner(): log.success("--- ENCODE COMPLETE ---") - def _do_audio(self, start: int, end: int, codec: Optional[List[str]], out_name: Optional[str] = None) -> None: + 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) log.status("--- TRIMMING AUDIO ---") self.audio_file = self.audio.trim_audio((start, end)) - if codec: - log.status("--- TRANSCODING AUDIO ---") - self.audio_file = self.audio.encode_audio(self.audio_file, codec) if out_name: shutil.copy(self.audio_file, out_name) self.audio_file = out_name diff --git a/yt_common/yt_common/chapters.py b/yt_common/yt_common/chapters.py new file mode 100644 index 0000000..3d6a8e4 --- /dev/null +++ b/yt_common/yt_common/chapters.py @@ -0,0 +1,63 @@ +from lxml import etree +from random import getrandbits + +from typing import Dict, List, NamedTuple, Set + +LANGMAP: Dict[str, str] = { + "eng": "en", + "und": "und", + "jpn": "ja", +} + + +class Chapter(NamedTuple): + title: str + frame: int + lang: str = "eng" + + +class RandMan: + used: Set[int] + + def __init__(self) -> None: + self.used = set() + + def get_rand(self, bits: int = 64) -> str: + r = getrandbits(bits) + while r in self.used: + r = getrandbits(bits) + self.used.add(r) + return str(r) + + +def timecode_to_timestamp(stamp: float) -> str: + m = int(stamp // 60) + stamp %= 60 + h = int(m // 60) + m %= 60 + return f"{h:02d}:{m:02d}:{stamp:06.3f}000000" + + +def make_chapters(chapters: List[Chapter], timecodes: List[float], outfile: str) -> None: + chapters.sort(key=lambda c: c.frame) + rand = RandMan() + + root = etree.Element("Chapters") + ed = etree.SubElement(root, "EditionEntry") + etree.SubElement(ed, "EditionUID").text = rand.get_rand() + + for i, c in enumerate(chapters): + start = timecode_to_timestamp(timecodes[c.frame]) + end = timecode_to_timestamp(timecodes[chapters[i+1].frame]) if i < len(chapters) - 1 else None + atom = etree.SubElement(ed, "ChapterAtom") + etree.SubElement(atom, "ChapterTimeStart").text = start + if end is not None: + etree.SubElement(atom, "ChapterTimeEnd").text = end + disp = etree.SubElement(atom, "ChapterDisplay") + etree.SubElement(disp, "ChapterString").text = c.title + etree.SubElement(disp, "ChapLanguageIETF").text = LANGMAP[c.lang] + etree.SubElement(disp, "ChapterLanguage").text = c.lang + etree.SubElement(atom, "ChapterUID").text = rand.get_rand() + + with open(outfile, "wb") as o: + o.write(etree.tostring(root, encoding="utf-8", xml_declaration=True, pretty_print=True)) diff --git a/yt_common/yt_common/config.py b/yt_common/yt_common/config.py index 0369f3d..d741eec 100644 --- a/yt_common/yt_common/config.py +++ b/yt_common/yt_common/config.py @@ -1,4 +1,8 @@ -from typing import Union +import os +import subprocess + +from abc import ABC, abstractmethod +from typing import List, Union class Config(): @@ -16,5 +20,48 @@ class Config(): self.datapath = datapath def format_filename(self, filename: str) -> str: - return filename.format(epnum=self.desc, title=self.title, - title_long=self.title_long, resolution=self.resolution) + fname = filename.format(epnum=self.desc, title=self.title, + title_long=self.title_long, resolution=self.resolution) + return os.path.join(f"../{self.desc}/", fname) + + def encode_audio(self, afile: str) -> str: + return afile # default: passthrough + + +class AudioEncoder(ABC): + @abstractmethod + def encode_audio(self, afile: str) -> str: + pass + + +class FFAudio(AudioEncoder): + def encode_audio(self, afile: str) -> str: + ffmpeg_args = [ + "ffmpeg", + "-hide_banner", "-loglevel", "panic", + "-i", afile, + "-y", + "-map", "0:a", + ] + self.codec_args() + ["_ffaudio_encode.mka"] + print("+ " + " ".join(ffmpeg_args)) + subprocess.call(ffmpeg_args) + return "_ffaudio_encode.mka" + + @abstractmethod + def codec_args(self) -> List[str]: + pass + + +class OpusMixin(FFAudio): + def codec_args(self) -> List[str]: + return ["-c:a", "libopus", "-b:a", "192k", "-sample_fmt", "s16"] + + +class FdkAacMixin(FFAudio): + def codec_args(self) -> List[str]: + return ["-c:a", "libfdk_aac", "-b:a", "256k", "-sample_fmt", "s16"] + + +class FlacMixin(FFAudio): + def codec_args(self) -> List[str]: + return ["-c:a", "flac"] diff --git a/yt_common/yt_common/source.py b/yt_common/yt_common/source.py index 6c1adcc..4cdc974 100644 --- a/yt_common/yt_common/source.py +++ b/yt_common/yt_common/source.py @@ -78,8 +78,9 @@ class FileSource(ABC): 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 __init__(self, src: Union[str, List[str], FileTrim, List[FileTrim]]) -> None: + srcl = src if isinstance(src, list) else [src] + self.src = [FileTrim(s, (None, None)) if isinstance(s, str) else s for s in srcl] def get_audio(self) -> List[FileTrim]: return self.src