From 2569dd0bd53f48497ef56f9e972d017c37558105 Mon Sep 17 00:00:00 2001 From: LightArrowsEXE Date: Tue, 27 Aug 2024 16:28:39 +0200 Subject: [PATCH] Further fixes, add preset/custom list support --- vsdeinterlace/wobbly/base.py | 32 ++++++-- vsdeinterlace/wobbly/info.py | 146 ++++++++++++++++++++++++++++++--- vsdeinterlace/wobbly/types.py | 37 ++++++++- vsdeinterlace/wobbly/wibbly.py | 2 +- vsdeinterlace/wobbly/wobbly.py | 83 ++++++++++++++++--- 5 files changed, 272 insertions(+), 28 deletions(-) diff --git a/vsdeinterlace/wobbly/base.py b/vsdeinterlace/wobbly/base.py index 4785f55..8290287 100644 --- a/vsdeinterlace/wobbly/base.py +++ b/vsdeinterlace/wobbly/base.py @@ -1,12 +1,14 @@ from typing import Sequence from vstools import (CustomValueError, DependencyNotFoundError, FieldBased, - FuncExceptT, VSFunction, core, replace_ranges, vs) + FieldBasedT, FuncExceptT, VSFunction, core, + replace_ranges, vs) from vsdeinterlace.combing import fix_interlaced_fades +from vsdeinterlace.wobbly.info import CustomList from .info import FreezeFrame, InterlacedFade, OrphanField -from .types import Match +from .types import CustomPostFiltering, Match class _WobblyProcessBase: @@ -24,16 +26,18 @@ def _check_plugin_installed(self, plugin: str, func_except: FuncExceptT | None = def _apply_fieldmatches( self, clip: vs.VideoNode, matches: Sequence[Match], + tff: FieldBasedT | None = None, func_except: FuncExceptT | None = None ) -> vs.VideoNode: """Apply fieldmatches to a clip.""" self._check_plugin_installed('fh', func_except) - match_clips = dict[str, vs.VideoNode]() + tff = FieldBased.from_param_or_video(tff, clip) - for match in set(matches): - match_clips |= {match: clip.std.SetFrameProps(wobbly_match=match)} + clip = clip.fh.FieldHint(None, tff, ''.join(matches)) + + match_clips = {match: clip.std.SetFrameProps(wobbly_match=match) for match in set(matches)} return clip.std.FrameEval(lambda n: match_clips.get(matches[n])) @@ -119,3 +123,21 @@ def _apply_interlaced_fades( fix_interlaced_fades(clip, colors=0, planes=0, func=func).std.SetFrameProps(wobbly_fif=True), [f.framenum for f in ifades] ) + + def _apply_custom_list( + self, clip: vs.VideoNode, custom_list: list[CustomList], pos: CustomPostFiltering + ) -> vs.VideoNode: + """Apply a list of custom functions to a clip based on the pos.""" + + for custom in custom_list: + if custom.position == pos: + custom_clip = custom.preset.apply_preset(clip) + + custom_clip = custom_clip.std.SetFrameProps( + wobbly_custom_list_name=custom.name, + wobbly_custom_list_position=str(pos) + ) + + clip = replace_ranges(clip, custom_clip, custom.frames) + + return clip diff --git a/vsdeinterlace/wobbly/info.py b/vsdeinterlace/wobbly/info.py index edb9918..efab763 100644 --- a/vsdeinterlace/wobbly/info.py +++ b/vsdeinterlace/wobbly/info.py @@ -3,11 +3,12 @@ from typing import Any, Callable, NamedTuple from vstools import (CustomValueError, FieldBased, FieldBasedT, - FileNotExistsError, FrameRangeN, FuncExceptT, - SceneChangeMode, SPath, SPathLike, VSCoreProxy, core, vs) + FileNotExistsError, FrameRangeN, FrameRangesN, + FuncExceptT, SceneChangeMode, SPath, SPathLike, + VSCoreProxy, core, vs) from .exceptions import InvalidMatchError -from .types import Match, SectionPreset +from .types import CustomPostFiltering, Match, SectionPreset __all__: list[str] = [ "WobblyMeta", "WobblyVideo", @@ -79,13 +80,37 @@ def source(self, func_except: FuncExceptT | None = None, **kwargs: Any) -> vs.Vi src = self.source_filter(sfile.to_str(), **kwargs) # type:ignore[operator] - if src.fps != self.framerate: - src = src.std.AssumeFPS(src, self.framerate.numerator, self.framerate.denominator) + src = self.set_framerate(src) + + return self.trim(src) + + def set_framerate(self, clip: vs.VideoNode) -> vs.VideoNode: + """ + Set the framerate of a clip to the base framerate. + + :param clip: The clip to set the framerate of. + + :return: The clip with the framerate set. + """ + + if clip.fps != self.framerate: + return clip.std.AssumeFPS(clip, self.framerate.numerator, self.framerate.denominator) + + return clip + + def trim(self, clip: vs.VideoNode) -> vs.VideoNode: + """ + Apply trims to a clip. + + :param clip: The clip to apply the trims to. + + :return: The trimmed clip. + """ if not self.trims: - return src + return clip - return core.std.Splice([src.std.Trim(s, e) for s, e in self.trims]) + return core.std.Splice([clip.std.Trim(s, e) for s, e in self.trims]) # TODO: refactor this @@ -268,7 +293,7 @@ def as_frame_range(self) -> FrameRangeN: return (self.start_frame, self.end_frame) -@dataclass +@dataclass(unsafe_hash=True) class FreezeFrame(_HoldsStartEndFrames): """Frame ranges to freeze.""" @@ -291,7 +316,7 @@ def __post_init__(self) -> None: raise CustomValueError("Frame number must be greater than or equal to 0!") -@dataclass +@dataclass(unsafe_hash=True) class InterlacedFade(_HoldsFrameNum): """Information about interlaced fades.""" @@ -307,7 +332,7 @@ def __post_init__(self) -> None: super().__post_init__() -@dataclass +@dataclass(unsafe_hash=True) class OrphanField(_HoldsFrameNum): """Information about the orphan fields.""" @@ -329,3 +354,104 @@ def deinterlace_order(self) -> FieldBased: """The fieldorder to deinterlace in to properly deinterlace the orphan field.""" return FieldBased.TFF if self.match in ('n', 'p') else FieldBased.BFF + + +@dataclass(unsafe_hash=True) +class Preset: + """A filtering preset.""" + + name: str + """The section the preset applies to.""" + + contents: str + """The preset to apply to the section.""" + + def __str__(self) -> str: + return f"Preset({self.name=}, {self.contents=})" + + def __post_init__(self) -> None: + clip = core.std.BlankClip() + + local_namespace = { + "clip": clip, + "core": core + } + + try: + exec(self.contents, {}, local_namespace) + except Exception as e: + raise CustomValueError( + f"Invalid preset contents ({self.contents=})! Original error: {e}", Preset + ) from e + + def apply_preset(self, clip: vs.VideoNode) -> vs.VideoNode: + """ + Apply the preset to a clip. + + :param clip: The clip to apply the preset to. + + :return: The clip with the preset applied. + """ + + local_namespace = {"clip": clip} + + try: + exec(self.contents, {}, local_namespace) + + clip = local_namespace.get("clip", clip) + except Exception as e: + raise CustomValueError(f"Could not apply preset ({self.contents=})!", self.apply_preset) from e + + return clip + + +@dataclass +class CustomList: + """Custom filtering applied to a given frame range.""" + + name: str + """The name of the custom list.""" + + preset: Preset + """The preset used for the custom list.""" + + position: CustomPostFiltering + """The position to apply the custom filter.""" + + frames: FrameRangesN = field(default_factory=lambda: []) + """The frame ranges to apply the custom filter to.""" + + def __init__( + self, name: str, + preset: Preset | str, + position: CustomPostFiltering | str, + frames: list[list[int]] | None = None, + presets: set[Preset] = set(), + ) -> None: + self.name = name + + if not preset: + raise CustomValueError("A preset must be set!", CustomList) + + if not position: + raise CustomValueError("A position must be set!", CustomList) + + self.position = CustomPostFiltering.from_str(position) if isinstance(position, str) else position + self.preset = self._get_preset_from_presets(preset, presets) if isinstance(preset, str) else preset + + self.frames = [] + + for frame_range in frames or []: + self.frames.append(tuple(frame_range)) # type:ignore + + def __str__(self) -> str: + return f"CustomList({self.name=}, {self.preset=}, {self.position=}, {self.frames=})" + + def _get_preset_from_presets(self, name: str, presets: set[Preset]) -> Preset: + """Get the preset from the list of presets.""" + + for preset in presets: + if preset.name == name: + return preset + + raise CustomValueError(f"Could not find the preset in the list of presets ({name=})!", CustomList) diff --git a/vsdeinterlace/wobbly/types.py b/vsdeinterlace/wobbly/types.py index f08394a..bcecc82 100644 --- a/vsdeinterlace/wobbly/types.py +++ b/vsdeinterlace/wobbly/types.py @@ -1,10 +1,14 @@ +from __future__ import annotations + +from enum import Enum from typing import Callable, Literal, TypeAlias -from vstools import vs +from vstools import NotFoundEnumValue, vs __all__ = [ "Match", "OrphanMatch", - "SectionPreset" + "SectionPreset", + "CustomPostFiltering", ] @@ -16,3 +20,32 @@ SectionPreset = Callable[[vs.VideoNode], vs.VideoNode] """A callable preset applied to a section.""" + + +class CustomPostFiltering(Enum): + """When to perform custom filtering.""" + + SOURCE = -1 + """Apply the custom filter after the source clip is loaded.""" + + FIELD_MATCH = 0 + """Apply the custom filter after the field match is applied.""" + + DECIMATE = 1 + """Apply the custom filter after the decimation is applied.""" + + @classmethod + def from_str(cls, value: str) -> CustomPostFiltering: + """Convert a string to a CustomPosition.""" + + norm_val = value.upper().replace('POST', '').strip().replace(' ', '_') + + try: + return cls[norm_val] + except KeyError: + raise NotFoundEnumValue( + "Could not find a matching CustomPostFiltering value!", cls.from_str, value + ) + + def __str__(self) -> str: + return self.name.replace('_', ' ').title() diff --git a/vsdeinterlace/wobbly/wibbly.py b/vsdeinterlace/wobbly/wibbly.py index 5d7d53f..be87c92 100644 --- a/vsdeinterlace/wobbly/wibbly.py +++ b/vsdeinterlace/wobbly/wibbly.py @@ -66,7 +66,7 @@ def _get_clip(self, display: bool = False) -> vs.VideoNode: odd_avg = separated[1::2].std.PlaneStats() if hasattr(core, 'akarin'): - wclip = core.akarin.PropExpr( # type:ignore + wclip = core.akarin.PropExpr( [wclip, even_avg, odd_avg], lambda: {'WibblyFieldDiff': 'y.PlaneStatsAverage z.PlaneStatsAverage - abs'} ) diff --git a/vsdeinterlace/wobbly/wobbly.py b/vsdeinterlace/wobbly/wobbly.py index 38a9433..816a632 100644 --- a/vsdeinterlace/wobbly/wobbly.py +++ b/vsdeinterlace/wobbly/wobbly.py @@ -12,12 +12,14 @@ FunctionUtil, Keyframes, SPath, SPathLike, Timecodes, UnsupportedFieldBasedError, VSFunction, core, vs) +from vsdeinterlace.wobbly.info import Preset + from .base import _WobblyProcessBase from .exceptions import InvalidCycleError, InvalidMatchError -from .info import (FreezeFrame, InterlacedFade, OrphanField, Section, - VDecParams, VfmParams, WobblyMeta, WobblyVideo, +from .info import (CustomList, FreezeFrame, InterlacedFade, OrphanField, + Section, VDecParams, VfmParams, WobblyMeta, WobblyVideo, _HoldsFrameNum, _HoldsStartEndFrames) -from .types import Match, OrphanMatch +from .types import CustomPostFiltering, Match, OrphanMatch __all__: list[str] = [ "WobblyParsed", @@ -59,7 +61,7 @@ class WobblyParsed(_WobblyProcessBase): decimations: set[int] = field(default_factory=set) """A set of frames to decimate.""" - sections: set[Section] = field(default_factory=set) + sections: list[Section] = field(default_factory=list) """A set of Section objects representing the scenes of a video.""" interlaced_fades: set[InterlacedFade] = field(default_factory=set) @@ -68,6 +70,12 @@ class WobblyParsed(_WobblyProcessBase): freeze_frames: set[FreezeFrame] = field(default_factory=set) """A list of FreezeFrame objects representing ranges to freeze, and which frames to replace them with.""" + presets: set[Preset] = field(default_factory=set) + """A set of Presets used in the wobbly file.""" + + custom_lists: list[CustomList] = field(default_factory=list) + """A set of CustomLists used in the wobbly file.""" + def __init__(self, file_path: SPathLike, func_except: FuncExceptT | None = None) -> None: """ Parse a wobbly file and return a WobblyParsed object. @@ -122,6 +130,8 @@ def __init__(self, file_path: SPathLike, func_except: FuncExceptT | None = None) self._set_interlaced_fades() self._set_freeze_frames() self._set_orphan_frames() + self._set_presets() + self._set_custom_lists() # Further sanitizing where necessary. self._remove_ifades_from_combed() @@ -158,6 +168,9 @@ def apply( if clip is None: clip = self.work_clip.source(func_except) + else: + clip = self.work_clip.trim(clip) + clip = self.work_clip.set_framerate(clip) func = FunctionUtil(clip, func_except, None, vs.YUV) @@ -165,8 +178,15 @@ def apply( orphan_proc = self._get_orphans_to_process(orphan_handling, func.func) - if matches_to_proc := self._force_c_match(self.matches, orphan_proc): # type:ignore[arg-type] - wclip = self._apply_fieldmatches(wclip, matches_to_proc, func.func) + if self.custom_lists: + wclip = self._apply_custom_list(wclip, self.custom_lists, CustomPostFiltering.SOURCE) + + if self.matches: + matches_to_proc = self._force_c_match(self.matches, orphan_proc) # type:ignore[arg-type] + wclip = self._apply_fieldmatches(wclip, matches_to_proc, self.field_order, func.func) + + if self.custom_lists: + wclip = self._apply_custom_list(wclip, self.custom_lists, CustomPostFiltering.FIELD_MATCH) if self.freeze_frames: wclip = self._apply_freezeframes(wclip, self.freeze_frames, func.func) @@ -186,6 +206,9 @@ def apply( if self.decimations: wclip = wclip.std.DeleteFrames(list(self.decimations)) + if self.custom_lists: + wclip = self._apply_custom_list(wclip, self.custom_lists, CustomPostFiltering.DECIMATE) + wclip = FieldBased.PROGRESSIVE.apply(wclip) return func.return_clip(wclip) @@ -193,7 +216,13 @@ def apply( def remove_sections(self) -> None: """Remove all sections from the wobbly data except for the first one.""" - self.sections = {next(iter(self.sections))} + self.sections = [self.sections[0]] + + def remove_presets(self) -> None: + """Remove all presets from the wobbly data.""" + + for section in self.sections: + section.presets = [] def _check_wob_path(self, file_path: SPathLike) -> SPath: """Check the wob file path and return an SPath object.""" @@ -240,14 +269,14 @@ def _set_sections(self) -> None: sections_data: list[dict[str, Any]] = self._get_val("sections", [{}]) - self.sections = { + self.sections = [ Section( section.get("start", 0), sections_data[i + 1].get("start", 0) - 1 if i < len(sections_data) - 1 else len(self.matches) - 1, section.get("presets", []) ) for i, section in enumerate(sections_data) - } + ] def _set_interlaced_fades(self) -> None: """Set the interlaced fades attribute.""" @@ -273,6 +302,8 @@ def _set_freeze_frames(self) -> None: def _set_orphan_frames(self) -> None: """Set the orphan frames attribute.""" + self.orphan_frames = set() + try: for section in self.sections: frame_num = section.start_frame @@ -292,6 +323,35 @@ def _set_orphan_frames(self) -> None: except IndexError as e: raise CustomIndexError(" ".join([str(e), f"(frame: {frame_num})"]), self._func) + def _set_presets(self) -> None: + """Set the presets attribute.""" + + preset_data: list[dict[str, str]] = self._get_val("presets", []) + + self.presets = { + Preset( + preset.get("name", "fallback name"), + preset.get("contents", "raise RuntimeError('No contents provided!')"), + ) + for preset in preset_data + } + + def _set_custom_lists(self) -> None: + """Set the custom lists attribute.""" + + custom_list_data: list[dict[str, Any]] = self._get_val("custom lists", []) + + self.custom_lists = [ + CustomList( + custom_list.get("name", "fallback name"), + custom_list.get("preset", ""), + custom_list.get("position", ""), + custom_list.get("frames", []), + self.presets + ) + for custom_list in custom_list_data + ] + def _remove_ifades_from_combed(self) -> None: """Remove interlaced fades from the combed frames attribute.""" @@ -389,6 +449,7 @@ def parse_wobbly( wobbly_file: SPathLike, clip: vs.VideoNode | None = None, tff: FieldBasedT | None = None, + orphan_deinterlacing_function: VSFunction = core.resize.Bob, orphan_handling: bool | OrphanMatch | Sequence[OrphanMatch] = False, ) -> vs.VideoNode: """ @@ -401,4 +462,6 @@ def parse_wobbly( wob_file = SPath(wobbly_file) - return WobblyParsed(wob_file, parse_wobbly).apply(clip, tff, orphan_handling, parse_wobbly) + return WobblyParsed(wob_file, parse_wobbly).apply( + clip, tff, orphan_handling, orphan_deinterlacing_function, parse_wobbly + )