Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add all useful fmtc dithertypes and set default to void #85

Merged
merged 10 commits into from
Sep 20, 2023
100 changes: 92 additions & 8 deletions vstools/functions/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@


class DitherType(CustomStrEnum):
"""Enum for `zimg_dither_type_e`."""
"""Enum for `zimg_dither_type_e` and fmtc `dmode`."""

AUTO = 'auto'
"""Choose automatically."""
Expand All @@ -54,6 +54,78 @@ class DitherType(CustomStrEnum):
ERROR_DIFFUSION = 'error_diffusion'
"""Floyd-Steinberg error diffusion."""

ERROR_DIFFUSION_FMTC = 'error_diffusion_fmtc'
"""
Floyd-Steinberg error diffusion.
Modified for serpentine scan (avoids worm artefacts).
"""

SIERRA_2_4A = 'sierra_2_4a'
"""
Another type of error diffusion.
Quick and excellent quality, similar to Floyd-Steinberg.
"""

STUCKI = 'stucki'
"""
Another error diffusion kernel.
Preserves delicate edges better but distorts gradients.
"""

ATKINSON = 'atkinson'
"""
Another error diffusion kernel.
Generates distinct patterns but keeps clean the flat areas (noise modulation).
"""

OSTROMOUKHOV = 'ostromoukhov'
"""
Another error diffusion kernel.
Slow, available only for integer input at the moment. Avoids usual F-S artefacts.
"""

VOID = 'void'
"""A way to generate blue-noise dither and has a much better visual aspect than ordered dithering."""

QUASIRANDOM = 'quasirandom'
"""
Dither using quasirandom sequences.
Good intermediated between Void and cluster and error diffusion algorithms.
"""

def apply(
self, clip: vs.VideoNode, fmt_out: vs.VideoFormat, range_in: ColorRange | None, range_out: ColorRange | None
) -> vs.VideoNode:
from ..utils import get_video_format

assert self != DitherType.AUTO, CustomValueError("Cannot apply AUTO.", self.__class__)

fmt = get_video_format(clip)

if not self.is_fmtc:
return clip.resize.Point(
format=fmt_out.id, dither_type=self.value.lower(),
range_in=range_in and range_in.value_zimg, range=range_out and range_out.value_zimg
)

if fmt.sample_type is vs.FLOAT:
if self == DitherType.OSTROMOUKHOV:
raise CustomValueError("Ostromoukhov can't be used for float input.", self.__class__)

# Workaround because fmtc doesn't support FLOAT 16 input
if fmt.bits_per_sample < 32:
clip = clip.resize.Point(format=fmt.replace(bits_per_sample=32).id, dither_type='none')

return clip.fmtc.bitdepth(
dmode=_dither_fmtc_types.get(self), bits=fmt_out.bits_per_sample,
fulls=None if not range_in else range_in == ColorRange.FULL,
fulld=None if not range_out else range_out == ColorRange.FULL
)

@property
def is_fmtc(self) -> bool:
return self in _dither_fmtc_types

@overload
@staticmethod
def should_dither(
Expand Down Expand Up @@ -154,6 +226,17 @@ def should_dither(
return in_range == ColorRange.FULL and (in_bits, out_bits) != (8, 16)


_dither_fmtc_types: dict[DitherType, int] = {
DitherType.SIERRA_2_4A: 3,
DitherType.STUCKI: 4,
DitherType.ATKINSON: 5,
DitherType.ERROR_DIFFUSION_FMTC: 6,
DitherType.OSTROMOUKHOV: 7,
DitherType.VOID: 8,
DitherType.QUASIRANDOM: 9,
}


@disallow_variable_format
def depth(
clip: vs.VideoNode, bitdepth: VideoFormatT | HoldsVideoFormatT | int | None = None, /,
Expand Down Expand Up @@ -183,7 +266,7 @@ def depth(
:param sample_type: Desired sample type of output clip. Allows overriding default float/integer behavior.
Accepts ``vapoursynth.SampleType`` enums ``vapoursynth.INTEGER`` and ``vapoursynth.FLOAT``
or their values, ``0`` and ``1`` respectively.
:param range_in: Input pixel range (defaults to input `clip`'s range).
:param range_in: Input pixel range (defaults to input `clip`'s range).
:param range_out: Output pixel range (defaults to input `clip`'s range).
:param dither_type: Dithering algorithm. Allows overriding default dithering behavior. See :py:class:`Dither`.

Expand All @@ -200,7 +283,7 @@ def depth(
from .funcs import fallback

in_fmt = get_video_format(clip)
out_fmt = get_video_format(fallback(bitdepth, clip), sample_type=sample_type) # type: ignore
out_fmt = get_video_format(fallback(bitdepth, clip), sample_type=sample_type)

range_out = ColorRange.from_param(range_out)
range_in = ColorRange.from_param(range_in)
Expand All @@ -217,16 +300,17 @@ def depth(
if dither_type is DitherType.AUTO:
should_dither = DitherType.should_dither(in_fmt, out_fmt, range_in, range_out)

dither_type = DitherType.ERROR_DIFFUSION if should_dither else DitherType.NONE
if hasattr(clip, "fmtc"):
dither_type = DitherType.VOID
else:
dither_type = DitherType.ERROR_DIFFUSION if out_fmt.bits_per_sample == 8 else DitherType.ORDERED
dither_type = dither_type if should_dither else DitherType.NONE

new_format = in_fmt.replace(
bits_per_sample=out_fmt.bits_per_sample, sample_type=out_fmt.sample_type
)

return clip.resize.Point(
format=new_format.id, range_in=range_in and range_in.value_zimg, range=range_out and range_out.value_zimg,
dither_type=dither_type
)
return dither_type.apply(clip, new_format, range_in, range_out)


_f2c_cache = WeakValueDictionary[int, vs.VideoNode]()
Expand Down
33 changes: 20 additions & 13 deletions vstools/utils/clips.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
PrimariesT, PropEnum, Transfer, TransferT
)
from ..exceptions import CustomValueError, InvalidColorFamilyError
from ..functions import check_variable, depth, fallback, get_y, join
from ..functions import check_variable, depth, fallback, get_y, join, DitherType
from ..types import F_VD, FuncExceptT, HoldsVideoFormatT, P
from . import vs_proxy as vs
from .info import get_depth, get_video_format, get_w
Expand All @@ -26,14 +26,16 @@


def finalize_clip(
clip: vs.VideoNode, bits: int | None = 10, clamp_tv_range: bool = True, *, func: FuncExceptT | None = None
clip: vs.VideoNode, bits: int | None = 10, clamp_tv_range: bool = True,
dither_type: DitherType = DitherType.AUTO, *, func: FuncExceptT | None = None
) -> vs.VideoNode:
"""
Finalize a clip for output to the encoder.

:param clip: Clip to output.
:param bits: Output bits.
:param clamp_tv_range: Whether to clamp to tv range.
:param dither_type: Dithering used for the bitdepth conversion.
:param func: Optional function this was called from.

:return: Dithered down and optionally clamped clip.
Expand All @@ -42,7 +44,7 @@ def finalize_clip(
assert check_variable(clip, func or finalize_clip)

if bits:
clip = depth(clip, bits)
clip = depth(clip, bits, dither_type=dither_type)
else:
bits = get_depth(clip)

Expand All @@ -65,35 +67,35 @@ def finalize_clip(
@overload
def finalize_output(
function: None = None, /, *, bits: int | None = 10,
clamp_tv_range: bool = True, func: FuncExceptT | None = None
clamp_tv_range: bool = True, dither_type: DitherType = DitherType.AUTO, func: FuncExceptT | None = None
) -> Callable[[F_VD], F_VD] | F_VD:
...


@overload
def finalize_output(
function: F_VD, /, *, bits: int | None = 10,
clamp_tv_range: bool = True, func: FuncExceptT | None = None
clamp_tv_range: bool = True, dither_type: DitherType = DitherType.AUTO, func: FuncExceptT | None = None
) -> F_VD:
...


def finalize_output(
function: F_VD | None = None, /, *, bits: int | None = 10,
clamp_tv_range: bool = True, func: FuncExceptT | None = None
clamp_tv_range: bool = True, dither_type: DitherType = DitherType.AUTO, func: FuncExceptT | None = None
) -> Callable[[F_VD], F_VD] | F_VD:
"""Decorator implementation of finalize_clip."""

if function is None:
return cast(
Callable[[F_VD], F_VD],
partial(finalize_output, bits=bits, clamp_tv_range=clamp_tv_range, func=func)
partial(finalize_output, bits=bits, clamp_tv_range=clamp_tv_range, dither_type=dither_type, func=func)
)

@wraps(function)
def _wrapper(*args: Any, **kwargs: Any) -> vs.VideoNode:
assert function
return finalize_clip(function(*args, **kwargs), bits, clamp_tv_range, func=func)
return finalize_clip(function(*args, **kwargs), bits, clamp_tv_range, dither_type, func=func)

return cast(F_VD, _wrapper)

Expand All @@ -106,7 +108,8 @@ def initialize_clip(
chroma_location: ChromaLocationT | None = None,
color_range: ColorRangeT | None = None,
field_based: FieldBasedT | None = None,
strict: bool = False, *, func: FuncExceptT | None = None
strict: bool = False,
dither_type: DitherType = DitherType.AUTO, *, func: FuncExceptT | None = None
) -> vs.VideoNode:
"""
Initialize a clip with default props.
Expand All @@ -126,6 +129,7 @@ def initialize_clip(
:param field_based: FieldBased prop to set. If None, tries to get the FieldBased from existing props.
:param strict: Whether to be strict about existing properties.
If True, throws an exception if certain frame properties are not found.
:param dither_type: Dithering used for the bitdepth conversion.
:param func: Optional function this was called from.

:return: Clip with relevant frame properties set, and optionally dithered up to 16 bits.
Expand All @@ -150,7 +154,7 @@ def initialize_clip(
if bits is None:
return clip

return depth(clip, bits)
return depth(clip, bits, dither_type=dither_type)


@overload
Expand All @@ -162,6 +166,7 @@ def initialize_input(
chroma_location: ChromaLocationT | None = None,
color_range: ColorRangeT | None = None,
field_based: FieldBasedT | None = None,
dither_type: DitherType = DitherType.AUTO,
func: FuncExceptT | None = None
) -> Callable[[F_VD], F_VD]:
...
Expand All @@ -176,7 +181,8 @@ def initialize_input(
chroma_location: ChromaLocationT | None = None,
color_range: ColorRangeT | None = None,
field_based: FieldBasedT | None = None,
strict: bool = False, func: FuncExceptT | None = None
strict: bool = False,
dither_type: DitherType = DitherType.AUTO, func: FuncExceptT | None = None
) -> F_VD:
...

Expand All @@ -189,7 +195,8 @@ def initialize_input(
chroma_location: ChromaLocationT | None = None,
color_range: ColorRangeT | None = None,
field_based: FieldBasedT | None = None,
strict: bool = False, func: FuncExceptT | None = None
strict: bool = False,
dither_type: DitherType = DitherType.AUTO, func: FuncExceptT | None = None
) -> Callable[[F_VD], F_VD] | F_VD:
"""
Decorator implementation of ``initialize_clip``
Expand All @@ -199,7 +206,7 @@ def initialize_input(
bits=bits,
matrix=matrix, transfer=transfer, primaries=primaries,
chroma_location=chroma_location, color_range=color_range,
field_based=field_based, strict=strict, func=func
field_based=field_based, strict=strict, dither_type=dither_type, func=func
)

if function is None:
Expand Down