yt_common: chapters and audio mixins

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

yt_common/yt_common/antialiasing.py

@@ -1,6 +1,6 @@
 import vapoursynth as vs
 
-import vsutil
+import math
 
 from havsfunc import SMDegrain
 from lvsfunc.aa import nnedi3, clamp_aa, upscaled_sraa
@@ -8,7 +8,9 @@ from lvsfunc.kernels import Bicubic
 from lvsfunc.misc import replace_ranges, scale_thresh
 from lvsfunc.types import Range
 
-from typing import Callable, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from .scale import nnedi3_double
 
 
 core = vs.core
@@ -16,8 +18,9 @@ core = vs.core
 
 def mask_weak(clip: vs.VideoNode) -> vs.VideoNode:
     # use bilateral to smear the noise as much as possible without destroying lines
-    pre = vsutil.get_y(clip).bilateral.Bilateral(sigmaS=5, sigmaR=0.75)
-    # frei-chen edge detection
+    pre = clip.std.ShufflePlanes(planes=0, colorfamily=vs.GRAY) \
+        .bilateral.Bilateral(sigmaS=5, sigmaR=0.75)
+    # g41f's broken frei-chen edge detection
     gx = core.std.Convolution(pre, [-7, 0, 7, -10, 0, 10, -7, 0, 7], divisor=7, saturate=False)
     gy = core.std.Convolution(pre, [-7, -10, -7, 0, 0, 0, 7, 10, 7], divisor=7, saturate=False)
     return core.std.Expr([gx, gy], 'x dup * y dup * + sqrt').std.Binarize(scale_thresh(0.25, clip)) \
@@ -25,8 +28,9 @@ def mask_weak(clip: vs.VideoNode) -> vs.VideoNode:
 
 
 def mask_strong(clip: vs.VideoNode) -> vs.VideoNode:
-    mask: vs.VideoNode = SMDegrain(vsutil.get_y(clip), tr=3, thSAD=500, RefineMotion=True, prefilter=4) \
-        .std.Prewitt().std.Binarize(scale_thresh(0.25, clip)).std.Maximum().std.Convolution([1]*9)
+    pre: vs.VideoNode = SMDegrain(clip.std.ShufflePlanes(planes=0, colorfamily=vs.GRAY),
+                                  tr=3, thSAD=500, RefineMotion=True, prefilter=4)
+    mask = pre.std.Prewitt().std.Binarize(scale_thresh(0.25, clip)).std.Maximum().std.Convolution([1]*9)
     return mask
 
 
@@ -34,10 +38,56 @@ def combine_mask(clip: vs.VideoNode, weak: Union[Range, List[Range], None] = Non
     return replace_ranges(mask_strong(clip), mask_weak(clip), weak or [])
 
 
-def sraa_clamp(clip: vs.VideoNode, mask: Optional[vs.VideoNode] = None,
+def sraa_clamp(clip: vs.VideoNode, mask: Optional[Callable[[vs.VideoNode], vs.VideoNode]] = None,
                strength: float = 2, opencl: bool = False,
                postprocess: Optional[Callable[[vs.VideoNode], vs.VideoNode]] = None) -> vs.VideoNode:
     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
+    return core.std.MaskedMerge(clip, clamp, mask(sraa), planes=0) if mask else clamp
+
+
+def _zastin_default_mask(clip: vs.VideoNode) -> vs.VideoNode:
+    return clip.std.ShufflePlanes(planes=0, colorfamily=vs.GRAY) \
+        .std.Prewitt() \
+        .std.Binarize(scale_thresh(0.25, clip)) \
+        .std.Maximum() \
+        .std.BoxBlur(0, 1, 1, 1, 1)
+
+
+def zastinAA(clip: vs.VideoNode, rfactor: float = 2.0,
+             maskfun: Callable[[vs.VideoNode], vs.VideoNode] = _zastin_default_mask,
+             **eedi3override: Any) -> vs.VideoNode:
+    eedi3args: Dict[str, Any] = {'field': 0, 'alpha': 0.125, 'beta': 0.25, 'gamma': 40,
+                                 'nrad': 2, 'mdis': 20, 'vcheck': 2, 'vthresh0': 12,
+                                 'vthresh1': 24, 'vthresh2': 4}
+    eedi3args.update(eedi3override)
+    y = clip.std.ShufflePlanes(planes=0, colorfamily=vs.GRAY)
+
+    def eedi3s(clip: vs.VideoNode, mclip: Optional[vs.VideoNode] = None) -> vs.VideoNode:
+        out = clip.eedi3m.EEDI3(mclip=mclip, sclip=clip, **eedi3args)
+        if mclip is not None:
+            return core.std.Expr([clip, out, mclip], 'z y x ?')
+        return out
+
+    def resize_mclip(mclip: vs.VideoNode, w: Optional[int] = None, h: Optional[int] = None) -> vs.VideoNode:
+        iw = mclip.width
+        ih = mclip.height
+        ow = w if w is not None else iw
+        oh = h if h is not None else ih
+        if (ow > iw and ow/iw != ow//iw) or (oh > ih and oh/ih != oh//ih):
+            mclip = mclip.resize.Point(iw * math.ceil(ow / iw), ih * math.ceil(oh / ih))
+        return mclip.fmtc.resample(ow, oh, kernel='box', fulls=1, fulld=1)
+
+    aaw = round(y.width * rfactor) & ~1
+    aah = round(y.height * rfactor) & ~1
+    mask = maskfun(y)
+    mclip = resize_mclip(mask, aaw, aah)
+
+    aa = y.std.Transpose()
+    aa = nnedi3_double(aa)
+    aa = aa.resize.Spline16(aah, aaw)
+    aa = eedi3s(aa, mclip=mclip.std.Transpose()).std.Transpose()
+    aa = eedi3s(aa, mclip=mclip).resize.Spline16(y.width, y.height)
+
+    return core.std.ShufflePlanes([core.std.MaskedMerge(y, aa, mask), clip], planes=[0, 1, 2], colorfamily=vs.YUV)

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))

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"]

yt_common/yt_common/scale.py

@@ -0,0 +1,12 @@
+import vapoursynth as vs
+
+from typing import Any, Dict
+
+
+def nnedi3_double(clip: vs.VideoNode) -> vs.VideoNode:
+    nnargs: Dict[str, Any] = dict(field=0, dh=True, nsize=4, nns=4, qual=2, pscrn=2)
+    nn = clip.std.Transpose() \
+        .nnedi3.nnedi3(**nnargs) \
+        .std.Transpose() \
+        .nnedi3.nnedi3(**nnargs)
+    return nn.resize.Bicubic(src_top=0.5, src_left=0.5)

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