extract_subs_v1.vpy

SourceFile = str(Source)
dir_ = str(OutputDir)
DimensionCropBox = [int(width), int(height)]
HeightCropBox = int(CropBox_y)
HeightCropBoxAlt= int(CropBoxAlt_y)
Supersampling= int(Supersampling)
ExpandRatio = int(ExpandRatio)
ModeU = Resampler
ThresholdI = int(WhiteThresh)
ThresholdO = int(BlackThresh)
ThresholdSCD = float(DetectionThresh)


import vapoursynth as vs
if ModeU=='nnedi3':
    import edi_rpow2 as edi
import havsfunc as haf
import functools
import numpy as np
import cv2 as cv
import mvsfunc as mvs
import sys
import os

from multiprocessing.dummy import Pool as ThreadPool
from pathlib import Path, PurePath

core = vs.core

def minimum(x,y):
    return min(x,y)


displayname = PurePath(SourceFile).name

if not Path(Path(dir_).joinpath(displayname)).exists():
    displayname_def_path = Path(Path(dir_).joinpath(displayname, "default"))
    dpname_def_abspath = Path(displayname_def_path).resolve()
    Path(dpname_def_abspath).mkdir(parents=True, exist_ok=True)
    if HeightCropBoxAlt != -1:
        Path(Path(Path(dir_).joinpath(displayname, "alt"))).mkdir(parents=True, exist_ok=True)
    
Clip = core.ffms2.Source(
    source=SourceFile
    )

if type(ThresholdI) is int and type(ThresholdO) is int:
    Clip = core.std.ShufflePlanes(clips=Clip, planes=0, colorfamily=vs.GRAY)

# Upscale
if Supersampling < 0:
    if Clip.width/Clip.height>16/9:    
        TargetRes = 1920
        CurrentRes = Clip.width    
    else:
        TargetRes = 1080
        CurrentRes = Clip.height
    if ModeU == "nnedi3":
        Ss = TargetRes/CurrentRes/1.125    
    else:
        Ss = TargetRes/CurrentRes
elif Supersampling==0:
    Ss = 1
else:
    Ss = Supersampling

if ModeU == "nnedi3" and Ss != 1:
    if Ss-int(Ss)>0:
        Ss = int(Ss/2)*2+2
    else:
        Ss = int(Ss/2)*2
    if Supersampling<0:
        Ssbis = TargetRes/(CurrentRes*Ss)
    else:
        Ssbis = Supersampling/Ss


def Resizing(clip, width, height, height2):
    # clip = core.std.Crop(clip=clip, left=0, top=300, right=0, bottom=0)
    clip = core.std.CropAbs(
        clip=clip, width=width, height=height, left=int((clip.width-width)/2), top=clip.height-height2
    )
    if Ss != 1:
        if ModeU == "nnedi3" or ModeU == "waifu2x":
            if ModeU == "nnedi3":
                clip = edi.nnedi3_rpow2(clip=clip, rfactor=Ss)
            else :
                clip = core.fmtc.bitdepth(clip=clip, bits=32)
                clip = core.w2xc.Waifu2x(clip=clip, scale=Ss)
                if Ssbis!=1:
                    clip = core.fmtc.bitdepth(clip=clip, bits=16)
                else :
                    clip = core.fmtc.bitdepth(clip=clip, bits=8)
            if Ssbis!=1:
                clip = core.fmtc.resample(clip=clip, scale=Ssbis, kernel="sinc", taps=2)
                clip = core.fmtc.bitdepth(clip=clip, bits=8)
        else:
            clip = core.fmtc.resample(clip=clip, scale=Ss, kernel="sinc", taps=2)
            clip = core.fmtc.bitdepth(clip=clip, bits=8)
    elif clip.format.bits_per_sample!=8:
        clip = core.fmtc.bitdepth(clip=clip, bits=8)
    return clip
    

def RGBBinarize(clip, threshold):
    R = core.std.ShufflePlanes(clips=clip, planes=0, colorfamily=vs.GRAY)
    G = core.std.ShufflePlanes(clips=clip, planes=1, colorfamily=vs.GRAY)
    B = core.std.ShufflePlanes(clips=clip, planes=2, colorfamily=vs.GRAY)
    for i in range(0,int(len(threshold)/3)):
        i = i*3
        x_ = "x " + str(threshold[i])
        y_ = " >= y " + str(threshold[i+1])
        z_ = " >= or z " + str(threshold[i+2])
        RGB = core.std.Expr(
            clips=[R,G,B],
            expr=[x_ + y_ + z_ + " >= or 255 0 ?"]
        )
        if i == 0:
            clipfin = RGB
        else:
            clipfin = core.std.Merge(clipfin,RGB)
    clipfin = core.std.Binarize(clip=clipfin, threshold=1)
    return clipfin


def Cleaning(clip, blank, e):
    if type(ThresholdI) is list or type(ThresholdO) is list:
        clipRGB = core.fmtc.resample(clip=clip, css="444")
        clipRGB = core.fmtc.matrix(clip=clipRGB, mat="709", col_fam=vs.RGB)
        clipRGB = core.fmtc.bitdepth(clip=clipRGB, bits=8)
        
    if type(ThresholdI) is int and type(ThresholdO) is int:
        White_Raw = core.std.Binarize(clip=clip, threshold=ThresholdI)
        Bright_Raw = core.std.Binarize(clip=clip, threshold=ThresholdO)
    elif type(ThresholdI) is int and type(ThresholdO) is list:
        White_Raw = core.std.ShufflePlanes(clips=clip, planes=0, colorfamily=vs.GRAY)
        White_Raw = core.std.Binarize(clip=White_Raw, threshold=ThresholdI)
        Bright_Raw = RGBBinarize(clipRGB, ThresholdO)
    elif type(ThresholdI) is list and type(ThresholdO) is int:
        White_Raw = RGBBinarize(clipRGB, ThresholdI)
        Bright_Raw = core.std.ShufflePlanes(clips=clip, planes=0, colorfamily=vs.GRAY)
        Bright_Raw = core.std.Binarize(clip=Bright_Raw, threshold=ThresholdO)
    else:
        White_Raw = RGBBinarize(clipRGB, ThresholdI)
        Bright_Raw = RGBBinarize(clipRGB, ThresholdO)

    Bright_Out = core.std.Lut2(clipa=Bright_Raw, clipb=Rect, function=minimum)

    Bright_Not = core.misc.Hysteresis(clipa=Bright_Out, clipb=Bright_Raw)
    Bright_Not = core.std.Invert(Bright_Not)

    White_Txt = core.std.MaskedMerge(blank,White_Raw, Bright_Not)

    White_Lb = haf.mt_inpand_multi(src=White_Txt, sw=int(e), sh=int(e), mode="ellipse")
    White_Lb = haf.mt_expand_multi(src=White_Lb, sw=int(e), sh=int(e), mode="ellipse")

    White_Ub = haf.mt_inpand_multi(src=White_Txt, sw=int(5*e), sh=int(5*e), mode="ellipse")
    White_Ub = haf.mt_expand_multi(src=White_Ub, sw=int(3*e), sh=int(3*e), mode="ellipse")
    White_Ub = core.std.Invert(White_Ub)

    White = core.std.MaskedMerge(Blank, White_Lb, White_Ub)
    White = core.misc.Hysteresis(clipa=White, clipb=White_Txt)

    ClipCleaning = core.std.MaskedMerge(Blank, White_Raw, White)
    ClipCleaning = core.std.Median(clip=ClipCleaning)
    
    return ClipCleaning


def SceneLog(n, f, clip, dir_, zero_pad):
    if (
        f[0].props._SceneChangePrev == 1 or f[0].props._SceneChangeNext == 1
    ) and f[0].props.PlaneStatsMax > 1.:
        nfn = ""
        if f[0].props._SceneChangePrev == 1:
            frame = f[1]
            v = cv.merge([
                np.array(frame[i], copy = False) for i in reversed(range(frame.format.num_planes))
            ])
            nfn = "%0*d.png" % (zero_pad, n)
            cv.imwrite(Path(dir_).joinpath(nfn), v)
        with open(Path(dir_).joinpath("SceneChanges.csv"), "a") as ofile:
            ofile.write("%0*d,%d,%d,\"%s\"\n" % (
                zero_pad, n, f[0].props._SceneChangePrev, f[0].props._SceneChangeNext, nfn
            ))
    return clip


HeightCropBox = HeightCropBox + DimensionCropBox[1]
if HeightCropBoxAlt >= 0:
    HeightCropBoxAlt = HeightCropBoxAlt + DimensionCropBox[1]

ClipResized = Resizing(
    Clip, DimensionCropBox[0], DimensionCropBox[1], HeightCropBox
)

Blackclip = core.std.BlankClip(
    width=int(ClipResized.width-20), height=int(ClipResized.height-20), format=vs.GRAY8, color=0
)
Rect = core.std.AddBorders(
    clip=Blackclip, left=10, right=10, top=10, bottom=10, color=255
)
Blank = core.std.BlankClip(ClipResized, format=vs.GRAY8)

ClipCleaned = Cleaning(ClipResized,Blank, ExpandRatio)

with open(f"{dir_}/{displayname}/default/SceneChanges.csv", "w") as ofile:
    ofile.write("[Video Informations]\nfps={0:.2f}\nframe_count={}\n\n[Scene Informations]\n".format(
        Clip.fps,
        Clip.num_frames
    ))
    ofile.write("frame,is_start,is_end,subimage\n")
    
ClipCleanedSC = core.std.CropAbs(
    clip=ClipCleaned,
    width=int(ClipCleaned.width/2.7), height=int(ClipCleaned.height/2.7),
    left=int(ClipCleaned.width*(1-1/2.7)/2), top=int(ClipCleaned.height/2)
)
ClipCleanedSC = core.misc.SCDetect(clip=ClipCleanedSC, threshold=ThresholdSCD)

ClipCleanedSC = core.std.PlaneStats(ClipCleanedSC)

zero_pad = len(str(ClipCleaned.num_frames))
ClipCleaned = core.std.FrameEval(
    ClipCleaned,
    functools.partial(
        SceneLog, clip=ClipCleaned,
        dir=Path(dir_).joinpath(displayname, "default"), zero_pad=zero_pad
    ),
    prop_src=[ClipCleanedSC, ClipCleaned]
)


if HeightCropBoxAlt >= 0:
    ClipResizedAlt = Resizing(
        Clip, DimensionCropBox[0], DimensionCropBox[1], HeightCropBoxAlt
    )
    
    ClipCleanedAlt = Cleaning(ClipResizedAlt, Blank, ExpandRatio)
    
    with open(f"{dir_}/{displayname}/alt/SceneChanges.csv", "w") as ofile:
        ofile.write("[Video Informations]\nfps={0:.2f}\nframe_count={}\n\n[Scene Informations]\n".format(
            Clip.fps,
            Clip.num_frames
        ))
        ofile.write("frame,is_start,is_end,subimage\n")
    ClipCleanedAltSC = core.std.CropAbs(
        clip=ClipCleanedAlt,
        width=int(ClipCleanedAlt.width/2.7), height=int(ClipCleanedAlt.height/2.7),
        left=int(ClipCleanedAlt.width*(1-1/2.7)/2), top=int(ClipCleanedAlt.height*(1/2-1/2.7))
    )
    ClipCleanedAltSC = core.misc.SCDetect(clip =ClipCleanedAltSC, threshold=ThresholdSCD)
    ClipCleanedAltSC = core.std.PlaneStats(ClipCleanedAltSC)
    ClipCleanedAlt = core.std.FrameEval(
        ClipCleanedAlt,
        functools.partial(
            SceneLog, clip=ClipCleanedAlt,
            dir=Path(dir_).joinpath(displayname, "alt"), zero_pad=zero_pad
        ),
        prop_src=[ClipCleanedAltSC, ClipCleanedAlt]
    )
    
    Clip = core.std.StackVertical([ClipCleanedAlt, ClipCleaned])

else:
    Clip = ClipCleaned
    
Clip.set_output()