from enum import Enum
from collections.abc import Callable, Sequence
import os
import vapoursynth as vs
from vapoursynth import core
import vssource
from vstools import initialize_clip, depth, set_output, finalize_clip, PlanesT, KwargsT
import lvsfunc as lvs
from vodesfunc import RescaleBuilder, NNEDI_Doubler, DescaleTarget
import vskernels as vsk
from vsscale import Waifu2x, ArtCNN
from vsdenoise import nl_means, MVTools, MVToolsPresets
from vsdehalo import fine_dehalo
from vsdeband import F3kdb, masked_deband

core.max_cache_size = 12 * 1024

#region Types

class ImportMethod(Enum):
    DGDecNV = 0
    LSmashWorks = 1
    BestSource = 2
    FFMS2 = 3

class RescaleDetails:
    width: int | float = 1280
    height: int | float = 720
    doubler: NNEDI_Doubler | Waifu2x | ArtCNN = ArtCNN()
    kernel: vsk.Kernel = vsk.Bilinear
    border_handling: 0 | 1 = 0

    def __init__(self, width: int | float = 1280, height: int | float = 720, doubler: NNEDI_Doubler | Waifu2x | ArtCNN = ArtCNN(), kernel: vsk.Kernel = vsk.Bilinear, border_handling: 0 | 1 = 0):
        self.width = width
        self.height = height
        self.doubler = doubler
        self.kernel = kernel
        self.border_handling = border_handling

class Details:
    name: str | None
    path: str
    importer: ImportMethod
    rescale: RescaleDetails | None = None

    def __init__(self, path: str, importer: ImportMethod, name: str | None = None, rescale: RescaleDetails | None = None):
        self.name = name
        self.path = path
        self.importer = importer
        self.rescale = rescale

class Source:
    importer_or_video: vs.VideoNode | Callable[[str, ImportMethod], vs.VideoNode]
    process_methods: list[Callable[[vs.VideoNode], vs.VideoNode]] | None = None
    details: Details | None = None
    label: str | None = None

    def __init__(self, importer_or_video: vs.VideoNode | Callable[[str, ImportMethod], vs.VideoNode], process_methods: list[Callable[[vs.VideoNode], vs.VideoNode]] | None = None, details: Details | None = None, label: str | None = None):
        self.importer_or_video = importer_or_video if isinstance(importer_or_video, vs.VideoNode) else importer_or_video()
        self.process_methods = process_methods
        self.details = details
        self.label = label

#endregion Types

#region Processing Lambdas

class Processors:
    @staticmethod
    def import_video(path: str, import_method: ImportMethod) -> vs.VideoNode:
        # Ensure source is an absolute path
        if not os.path.isabs(path):
            path = os.path.abspath(path)

        # Ensure source exists
        if not os.path.exists(path):
            raise ValueError(f'Source "{path}" not found!')

        match import_method:
            case ImportMethod.DGDecNV:
                return vssource.source(path)
            case ImportMethod.LSmashWorks:
                return core.lsmas.LWLibavSource(path)
            case ImportMethod.BestSource:
                return core.bs.VideoSource(path)
            case ImportMethod.FFMS2:
                return core.ffms2.Source(path)
            case _:
                raise ValueError(f"Invalid import method: {import_method}")

    @staticmethod
    def crop(left: int = 0, right: int = 0, top: int = 0, bottom: int = 0) -> Callable[[vs.VideoNode], vs.VideoNode]:
        """Returns a function that crops a VideoNode with the given dimensions."""
        def crop_function(video: vs.VideoNode) -> vs.VideoNode:
            return video.std.Crop(left=left, right=right, top=top, bottom=bottom)
        return crop_function

    @staticmethod
    def trim(start: int | None = None, end: int | None = None) -> Callable[[vs.VideoNode], vs.VideoNode]:
        """Returns a function that trims a VideoNode with the given dimensions."""
        def trim_function(video: vs.VideoNode) -> vs.VideoNode:
            return video.std.Trim(first=start, last=end)
        return trim_function

    @staticmethod
    def hermite(width: int | float, height: int | float, border_handling: 0 | 1 = 0) -> vs.VideoNode:
        """Scales a VideoNode with the given dimensions."""
        def scale_function(video: vs.VideoNode) -> vs.VideoNode:
            return vsk.Hermite.scale(video, width=width, height=height)
        return scale_function

    @staticmethod
    def rescale(doubler: NNEDI_Doubler | Waifu2x | ArtCNN, width: int | float, height: int | float, kernel: vsk.Kernel, match_centers_scaling: bool = False) -> Callable[[vs.VideoNode], vs.VideoNode]:
        def rescale_function(video: vs.VideoNode) -> vs.VideoNode:
            native_res = lvs.get_match_centers_scaling(video, target_width=width, target_height=height) if (match_centers_scaling) else KwargsT(width=width, height=height)
            print(f"Rescaling at {width}x{height} using {get_doubler_name(doubler)} and {get_kernel_name(kernel)} kernel")
            _builder, rescaled = RescaleBuilder(video).descale(kernel, **native_res).double(doubler).errormask().linemask().downscale(vsk.Hermite(linear=True)).final()
            return rescaled
        return rescale_function

    @staticmethod
    def denoise(
        strength: float | Sequence[float] = 1.2,
        temporal_radius: int | Sequence[int] = 2,
        search_radius: int | Sequence[int] = 2,
        similarity_radius: int | Sequence[int] = 4,
        planes: PlanesT = None
    ):
        """Returns a function that denoises a VideoNode with the given parameters."""
        def denoise_function(video: vs.VideoNode) -> vs.VideoNode:
            return nl_means(
                video,
                strength=strength,
                tr=temporal_radius,
                sr=search_radius,
                simr=similarity_radius,
                planes=planes
            )
        return denoise_function

    @staticmethod
    def dehalo(planes: PlanesT = [0,1,2]) -> Callable[[vs.VideoNode], vs.VideoNode]:
        def dehalo_function(video: vs.VideoNode) -> vs.VideoNode:
            return fine_dehalo(video, planes=planes)
        return dehalo_function

    @staticmethod
    def deband(
        radius: int = 16,
        threshold: float | list[float] = 96,
        grain: float | list[float] = [0.23, 0],
        sigma: float = 1.25,
        rxsigma: list[int] = [50, 220, 300],
        pf_sigma: float | None = 1.25,
        brz: tuple[float, float] = (0.038, 0.068),
        rg_mode: int = 0,
        planes: PlanesT = [0,1,2]
    ) -> Callable[[vs.VideoNode], vs.VideoNode]:
        def deband_function(video: vs.VideoNode) -> vs.VideoNode:
            return masked_deband(
                video,
                radius=radius,
                thr=threshold,
                grain=grain,
                sigma=sigma,
                rxsigma=rxsigma,
                pf_sigma=pf_sigma,
                brz=brz,
                rg_mode=rg_mode,
                planes=planes
            )
        return deband_function

    @staticmethod
    def dither(bit_depth: int = 10) -> Callable[[vs.VideoNode], vs.VideoNode]:
        def dither_function(video: vs.VideoNode) -> vs.VideoNode:
            return depth(
                video,
                bit_depth
            )
        return dither_function

#endregion Processing Lambdas

#region Utilities

def get_doubler_name(doubler: NNEDI_Doubler | Waifu2x | ArtCNN) -> str:
    """Returns a friendly string that represents the current DOUBLER."""
    if isinstance(doubler, Waifu2x):
        return 'Waifu2x'
    elif isinstance(doubler, NNEDI_Doubler):
        return 'NNEDI3'
    elif isinstance(doubler, ArtCNN.C16F64):
        return 'ArtCNN.C16F64'
    elif isinstance(doubler, ArtCNN.C4F32):
        return 'ArtCNN.C4F32'
    elif isinstance(doubler, ArtCNN):
        return 'ArtCNN'
    else:
        from warnings import warn
        warn(f'Invalid DOUBLER: {doubler}')
        return f'{doubler}'

def get_kernel_name(kernel: vsk.Kernel) -> str:
    """Returns a friendly string that represents the current KERNEL."""
    kernel_name_parts = f'{kernel.__class__}'.replace('<', '').replace('>', '').split('.')
    kernel_name = kernel_name_parts[len(kernel_name_parts) - 1].replace("'", '')

    return kernel_name

def get_import_method_name(import_method: ImportMethod) -> str:
    """Returns a friendly string that represents the current IMPORT_METHOD."""
    if import_method is ImportMethod.DGDecNV:
        return 'DGDecNV'
    elif import_method is ImportMethod.LSmashWorks:
        return 'LSmashWorks'
    elif import_method is ImportMethod.BestSource:
        return 'BestSource'
    elif import_method is ImportMethod.FFMS2:
        return 'FFMS2'
    else:
        raise ValueError(f"Invalid IMPORT_METHOD: {import_method}")

def process_source(source: Source) -> vs.VideoNode:
    opening_part = '' if source.details.path is None else f'Opening {source.details.name if source.details.name is not None else ''}({os.path.basename(source.details.path)}) with {get_import_method_name(source.details.importer)}'
    rescaling_part = '' if source.details.rescale is None else f' and rescaling at {source.details.rescale.width}x{source.details.rescale.height} using {get_doubler_name(source.details.rescale.doubler)} and {get_kernel_name(source.details.rescale.kernel)} kernel'

    print(f'{opening_part}{rescaling_part}')

    if source.process_methods is not None:
        for method in source.process_methods:
            source.importer_or_video = method(source.importer_or_video)

    return finalize_clip(source.importer_or_video)

def generate_label(details: Details) -> str:
    return f'{os.path.basename(details.path) if details.name is None else details.name}{'' if details.rescale is None else f': Rescale[{details.rescale.height}p,{get_kernel_name(details.rescale.kernel)}]'}'

#endregion Utilities

#region Setup

default_path = os.environ.get('SOURCE') or 'C:/Golden.Time.S01E01.Spring.Time.1080p.BluRay.Remux.Dual-Audio.FLAC2.0.H.264-CRUCiBLE.mkv'
vid_1 = Processors.import_video(path=default_path, import_method=ImportMethod.DGDecNV)

SOURCES: list[Source] = [
    Source(
        importer_or_video=vid_1,
        details=Details(
            path=default_path,
            importer=ImportMethod.DGDecNV,
            name='JPBD',
        ),
    ),
    Source(
        importer_or_video=vid_1,
        process_methods=[
            lambda video: initialize_clip(video),
            Processors.dehalo([0, 1, 2]),
            Processors.dither(),
        ],
        details=Details(
            path=default_path,
            importer=ImportMethod.DGDecNV,
            name='JPBD + Dehalo',
        ),
    ),
    Source(
        importer_or_video=vid_1,
        process_methods=[
            lambda video: initialize_clip(video),
            Processors.rescale(
                doubler=ArtCNN.C4F32(),
                width=vid_1.width * (918 / vid_1.height),
                height=918,
                kernel=vsk.Bilinear(),
            ),
            # Processors.denoise(
            #     strength=0.2,
            #     temporal_radius=2,
            #     search_radius=[3, 2, 2],
            #     planes=[0, 1, 2]
            # ),
            # Processors.dehalo([0, 1, 2]),
            # Processors.deband(
            #     grain=0,
            #     threshold=64,
            #     planes=[0, 1, 2]
            # ),
            Processors.dither(),
        ],
        details=Details(
            path=default_path,
            importer=ImportMethod.DGDecNV,
            rescale=RescaleDetails(
                width=vid_1.width * (918 / vid_1.height),
                height=918,
                doubler=ArtCNN.C4F32(),
                kernel=vsk.Bilinear(),
                # border_handling=1,
            ),
            name='JPBD',
        ),
    ),
    # Source(
    #     importer_or_video=Processors.import_video(
    #         path='C:/Ex/SomeOtherS01E01.mkv',
    #         import_method=ImportMethod.LSmashWorks
    #     ),
    #     details=Details(
    #         path='C:/Ex/SomeOtherS01E01.mkv',
    #         importer=ImportMethod.LSmashWorks,
    #         name='[Some Other Releaser] HEVC 1kbps',
    #     ),
    # ),
]

#endregion Setup

#region Main

# Ensure at least one source is set
if len(SOURCES) == 0:
    raise ValueError('No sources set! Make sure to set at least one source before running this script.')

# Map sources into processed VideoNodes
videos = [process_source(source) for source in SOURCES]

# For each video and its SOURCE set output using source.label or source.path
for video, source in zip(videos, SOURCES):
    label = source.label if source.label is not None else generate_label(source.details)

    set_output(video, label)

#endregion Main
Edit
Pub: 22 May 2025 11:43 UTC
Views: 7