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