yt_common: steal zastin's scary debanding mask

yt_common/yt_common/__init__.py

@@ -1 +1 @@
-from . import antialiasing, config, automation, logging, source  # noqa: F401
+from . import antialiasing, automation, config, deband, logging, source  # noqa: F401

yt_common/yt_common/automation.py

@@ -12,7 +12,7 @@ from typing import Any, BinaryIO, Callable, List, Optional, Sequence, Union, cas
 
 from .config import Config
 from .logging import log
-from .source import AMAZON_FILENAME, ER_FILENAME, SUBSPLS_FILENAME, FUNI_INTRO, glob_crc
+from .source import AMAZON_FILENAME, ER_FILENAME, SUBSPLS_FILENAME, FUNI_INTRO, glob_filename
 
 core = vs.core
 
@@ -141,17 +141,14 @@ class AudioGetter():
             log.success("Found Amazon audio")
             return
 
-        # as of Ep4 SubsPlease is using new funi 128kbps aac while erai has 256kbps still
         try:
-            if os.path.isfile(self.config.format_filename(ER_FILENAME)):
-                self.audio_file = self.config.format_filename(ER_FILENAME)
-                self.video_src = core.ffms2.Source(self.audio_file)
-            elif os.path.isfile(glob_crc(self.config.format_filename(SUBSPLS_FILENAME))):
-                self.audio_file = glob_crc(self.config.format_filename(SUBSPLS_FILENAME))
-                self.video_src = core.ffms2.Source(self.audio_file)
-                log.warn("Using SubsPlease, audio may be worse than Erai-Raws")
-            else:
-                raise FileNotFoundError()
+            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
@@ -203,6 +200,8 @@ class SelfRunner():
     encoder: Encoder
     audio: AudioGetter
 
+    profile: str
+
     def __init__(self, config: Config, final_filter: Callable[[], vs.VideoNode],
                  workraw_filter: Optional[Callable[[], vs.VideoNode]] = None) -> None:
         self.config = config
@@ -211,35 +210,30 @@ class SelfRunner():
 
         parser = argparse.ArgumentParser(description=f"Encode {self.config.title} Episode {self.config.epnum:02d}")
         if workraw_filter:
-            parser.add_argument("-w", "--workraw", help="Encode workraw, fast x264", action="store_true")
-        parser.add_argument("-s", "--start", nargs='?', type=int, help="Start encode at frame START")
-        parser.add_argument("-e", "--end", nargs='?', type=int, help="Stop encode at frame END (inclusive)")
+            parser.add_argument("-w", "--workraw", help="Encode workraw, fast x264.", action="store_true")
+        parser.add_argument("-s", "--start", nargs='?', type=int, help="Start encode at frame START.")
+        parser.add_argument("-e", "--end", nargs='?', type=int, help="Stop encode at frame END (inclusive).")
         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("-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")
-        parser.add_argument("-c", "--comparison", help="Output a comparison between workraw and final",
+        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. \
+                            Overrides SUFFIX when set. Defaults to \"workraw\" when WORKRAW is set, else \"final.\"")
+        parser.add_argument("-c", "--comparison", help="Output a comparison between workraw and final. \
+                            Will search for the output file to include in comparison, if present.",
                             action="store_true")
         args = parser.parse_args()
 
-        if args.comparison and workraw_filter:
-            log.status("Generating comparison...")
-            gencomp(10, "comp", src=workraw_filter(), final=final_filter())
-            log.status("Comparison generated.")
-            return
-
         self.workraw = args.workraw if workraw_filter else False
+        self.profile = "workraw" if self.workraw else "final"
+        self.profile = args.profile or self.profile
         self.suffix = args.suffix if args.suffix is not None else "workraw" if self.workraw else "premux"
+        self.suffix = args.profile or self.suffix
 
         self.clip = workraw_filter() if workraw_filter and self.workraw else final_filter()
 
-        basename = "workraw-settings" if self.workraw else "final-settings"
-        settings_path = os.path.join(self.config.datapath, basename)
-
-        if not os.path.isfile(settings_path):
-            raise FileNotFoundError(f"Failed to find {settings_path}!")
-
         start = args.start if args.start is not None else 0
         if args.end is not None:
             if args.end < 0:
@@ -260,8 +254,27 @@ class SelfRunner():
         if start >= end:
             raise ValueError("Start frame must be before end frame!")
 
+        out_name = f"{self.config.title.lower()}_{self.config.epnum:02d}_{self.suffix}.mkv"
+
+        if args.comparison and workraw_filter:
+            log.status("Generating comparison...")
+            if os.path.isfile(out_name):
+                pmx = core.ffms2.Source(out_name)
+                if pmx.num_frames == self.clip.num_frames:
+                    pmx = pmx[start:end]
+                gencomp(10, "comp", src=workraw_filter()[start:end], final=final_filter()[start:end], encode=pmx)
+            else:
+                gencomp(10, "comp", src=workraw_filter()[start:end], final=final_filter()[start:end])
+            log.status("Comparison generated.")
+            return
+
+        settings_path = os.path.join(self.config.datapath, f"{self.profile}-settings")
+        if not os.path.isfile(settings_path):
+            raise FileNotFoundError(f"Failed to find {settings_path}!")
+
         self.encoder = Encoder(self.config.epnum, settings_path, args.encoder, args.force)
-        self.video_file = self.encoder.encode(self.clip, f"{self.config.epnum:02d}_{start}_{end}", start, end)
+        self.video_file = self.encoder.encode(self.clip, f"{self.config.epnum:02d}_{self.suffix}_{start}_{end}",
+                                              start, end)
 
         log.status("--- LOOKING FOR AUDIO ---")
         self.audio = AudioGetter(self.config, args.audio)
@@ -271,8 +284,7 @@ class SelfRunner():
 
         try:
             log.status("--- MUXING FILE ---")
-            if self._mux(f"{self.config.title.lower()}_{self.config.epnum:02d}_{self.suffix}.mkv",
-                         start == 0 and end == self.clip.num_frames) != 0:
+            if self._mux(out_name, start == 0 and end == self.clip.num_frames) != 0:
                 raise Exception("mkvmerge failed")
         except Exception:
             log.error("--- MUXING FAILED ---")
@@ -316,7 +328,7 @@ def gencomp(num: int = 10, path: str = "comp", matrix: str = "709", **clips: vs.
     if len(lens) != 1:
         raise ValueError("gencomp: 'Clips must be equal length!'")
 
-    frames = random.sample(range(lens.pop()), num)
+    frames = sorted(random.sample(range(lens.pop()), num))
 
     if os.path.exists(path):
         shutil.rmtree(path)

yt_common/yt_common/deband.py

@@ -0,0 +1,95 @@
+import vapoursynth as vs
+
+from enum import Enum
+from functools import partial
+from typing import Any, Callable, List, Optional, Tuple, Union
+
+import vsutil
+
+
+core = vs.core
+
+
+class XxpandModes(Enum):
+    SQUARE = 1
+    ELLIPSE = 2
+    LOSANGE = 3
+
+
+def mt_xxpand_multi(clip: vs.VideoNode, sw: int = 1, sh: Optional[int] = None,
+                    mode: XxpandModes = XxpandModes.SQUARE,
+                    planes: Union[int, List[int], None] = None, start: int = 0,
+                    m__imum: Callable[..., vs.VideoNode] = core.std.Maximum, **params: Any) -> List[vs.VideoNode]:
+    """
+    blame zastin for this
+    """
+    sh = vsutil.fallback(sh, sw)
+    assert clip.format is not None
+    planes = list(range(clip.format.num_planes)) if planes is None else [planes] if isinstance(planes, int) else planes
+
+    if mode == XxpandModes.ELLIPSE:
+        coordinates = [[1]*8, [0, 1, 0, 1, 1, 0, 1, 0], [0, 1, 0, 1, 1, 0, 1, 0]]
+    elif mode == XxpandModes.LOSANGE:
+        coordinates = [[0, 1, 0, 1, 1, 0, 1, 0]] * 3
+    else:
+        coordinates = [[1]*8] * 3
+
+    clips = [clip]
+
+    end = min(sw, sh) + start
+
+    for x in range(start, end):
+        clips += [m__imum(clips[-1], coordinates=coordinates[x % 3], planes=planes, **params)]
+
+    for x in range(end, end + sw - sh):
+        clips += [m__imum(clips[-1], coordinates=[0, 0, 0, 1, 1, 0, 0, 0], planes=planes, **params)]
+
+    for x in range(end, end + sh - sw):
+        clips += [m__imum(clips[-1], coordinates=[0, 1, 0, 0, 0, 0, 1, 0], planes=planes, **params)]
+
+    return clips
+
+
+def morpho_mask(clip: vs.VideoNode) -> Tuple[vs.VideoNode, vs.VideoNode]:
+    """
+    blame zastin for this
+    """
+    y = vsutil.get_y(clip)
+    assert y.format is not None
+
+    maxm = partial(mt_xxpand_multi, m__imum=core.std.Maximum)
+    minm = partial(mt_xxpand_multi, m__imum=core.std.Minimum)
+
+    ymax = maxm(y, sw=40, mode='ellipse')
+    ymin = minm(y, sw=40, mode='ellipse')
+
+    thr = round(0.0125 * ((1 << y.format.bits_per_sample) - 1))
+    ypw0 = y.std.Prewitt()
+    ypw = ypw0.std.Binarize(thr).rgvs.RemoveGrain(11)
+
+    rad = 3
+    thr = round(0.0098 * ((1 << y.format.bits_per_sample) - 1))
+    yrangesml = core.std.Expr([ymax[3], ymin[3]], 'x y - abs')
+    yrangesml = yrangesml.std.Binarize(thr).std.BoxBlur(0, 2, 1, 2, 1)
+
+    rad = 16
+    thr = round(0.0156 * ((1 << y.format.bits_per_sample) - 1))
+    yrangebig0 = core.std.Expr([ymax[rad], ymin[rad]], 'x y - abs')
+    yrangebig = yrangebig0.std.Binarize(thr)
+    yrangebig = minm(yrangebig,
+                     sw=rad * 3 // 4,
+                     threshold=(1 << y.format.bits_per_sample)//((rad * 3 // 4) + 1),
+                     mode='ellipse')[-1]
+    yrangebig = yrangebig.std.BoxBlur(0, rad//4, 1, rad//4, 1)
+
+    rad = 30
+    thr = round(0.0039 * ((1 << y.format.bits_per_sample) - 1))
+    ymph = core.std.Expr([y, maxm(ymin[rad], sw=rad, mode='ellipse')[rad],
+                          minm(ymax[rad], sw=rad, mode='ellipse')[rad]], 'x y - z x - max')
+    ymph = ymph.std.Binarize(round(0.00586 * ((1 << y.format.bits_per_sample) - 1)))
+    ymph = ymph.std.Minimum().std.Maximum()
+    ymph = ymph.std.BoxBlur(0, 4, 1, 4, 1)
+
+    grad_mask = core.std.Expr([ymph, yrangesml, ypw], expr="x y z max max")
+
+    return grad_mask, yrangebig

yt_common/yt_common/source.py

@@ -16,8 +16,8 @@ from .logging import log
 
 core = vs.core
 
-SUBSPLS_FILENAME: str = "[SubsPlease] {title_long} - {epnum:02d} ({resolution}p) [$CRC].mkv"
-ER_FILENAME: str = "[Erai-raws] {title_long} - {epnum:02d} [v0][{resolution}p].mkv"
+SUBSPLS_FILENAME: str = "[SubsPlease] {title_long} - {epnum:02d} ({resolution}p) [$GLOB].mkv"
+ER_FILENAME: str = "[Erai-raws] {title_long} - {epnum:02d} [v0][{resolution}p]$GLOB.mkv"
 FUNI_INTRO: int = 289
 AMAZON_FILENAME: str = "{title_long} - {epnum:02d} (Amazon Prime CBR {resolution}p).mkv"
 
@@ -37,8 +37,8 @@ def waka_replace(src: vs.VideoNode, wakas: List[vs.VideoNode], ranges: List[List
     return src, new_wakas
 
 
-def glob_crc(pattern: str) -> str:
-    res = glob.glob(glob.escape(pattern).replace("$CRC", "*"))
+def glob_filename(pattern: str) -> str:
+    res = glob.glob(glob.escape(pattern).replace("$GLOB", "*"))
     if len(res) == 0:
         raise FileNotFoundError(f"File matching \"{pattern}\" not found!")
     return res[0]
@@ -74,11 +74,16 @@ class FunimationSource(DehardsubFileFinder):
         return core.ffms2.Source(self.config.format_filename(AMAZON_FILENAME))
 
     def get_funi_filename(self) -> str:
-        if os.path.isfile(self.config.format_filename(ER_FILENAME)):
-            return self.config.format_filename(ER_FILENAME)
+        try:
+            return glob_filename(self.config.format_filename(ER_FILENAME))
+        except FileNotFoundError:
+            pass
 
-        log.error("Erai-raws not found, falling back to SubsPlease")
-        return glob_crc(self.config.format_filename(SUBSPLS_FILENAME))
+        try:
+            return glob_filename(self.config.format_filename(SUBSPLS_FILENAME))
+        except FileNotFoundError:
+            log.error("Could not find funimation video")
+            raise
 
     def get_funi(self) -> vs.VideoNode:
         return core.ffms2.Source(self.get_funi_filename())[FUNI_INTRO:]