tweaking ¶

def auto_balance(
    clip: vs.VideoNode,
    target_max: SupportsFloat | None = None,
    relative_sat: float = 1.0,
    range_in: ColorRange = ColorRange.LIMITED,
    frame_overrides: Override | Sequence[Override] = [],
    ref: vs.VideoNode | None = None,
    radius: int = 1,
    delta_thr: float = 0.4,
    min_thr: float = 1.0,
    max_thr: float = 5.0,
    min_thr_tr: float = 1.0,
    max_thr_tr: float = 5.0,
    balance_mode: BalanceMode = BalanceMode.UNDIMMING,
    weight_mode: BalanceWeightMode = BalanceWeightMode.MEAN,
    prop: bool = False,
) -> vs.VideoNode:
    import numpy as np

    ref_clip = fallback(ref, clip)

    assert check_variable(clip, auto_balance)
    assert check_variable(ref_clip, auto_balance)

    if ref_clip.format.sample_type is vs.FLOAT:
        raise CustomValueError(auto_balance, "Float auto_balance not implemented yet!")

    zero = scale_value(16, 8, ref_clip, range_in, scale_offsets=True)

    target = float(
        fallback(
            target_max, scale_value(235, input_depth=8, output_depth=ref_clip, range_in=range_in, scale_offsets=True)
        )
    )

    if weight_mode == BalanceWeightMode.NONE:
        raise CustomValueError(auto_balance, "Global weight mode can't be NONE!")

    ref_stats = ref_clip.std.PlaneStats()

    over_mapped = list[tuple[range, float, BalanceWeightMode]]()

    if frame_overrides:
        frame_overrides = [frame_overrides] if isinstance(frame_overrides, Override) else list(frame_overrides)

        over_frames, over_conts, over_int_modes = list(zip(*frame_overrides))

        oframes_ranges = [range(start, stop + 1) for start, stop in normalize_ranges(clip, list(over_frames))]

        over_mapped = list(zip(oframes_ranges, over_conts, over_int_modes))

    clipfrange = range(0, clip.num_frames)

    def _weighted(x: float, y: float, z: float) -> float:
        return max(1e-6, x - z) / max(1e-6, y - z)

    nobalanceclip = clip.std.SetFrameProps(AutoBalance=False) if prop else clip

    def _autobalance(n: int, f: Sequence[vs.VideoFrame]) -> vs.VideoNode:
        override: tuple[range, float, BalanceWeightMode] | None = next((x for x in over_mapped if n in x[0]), None)

        psvalues: Any = np.asarray(
            [_weighted(target, get_prop(frame.props, "PlaneStatsMax", int), zero) for frame in f]
        )

        middle_idx = psvalues.size // 2

        mean_value = np.mean(psvalues)

        if not override and not (mean_value >= min_thr_tr and mean_value <= max_thr_tr):
            return nobalanceclip

        curr_value = psvalues[middle_idx]

        if not override and not (curr_value >= min_thr and curr_value <= max_thr):
            return nobalanceclip

        if balance_mode == BalanceMode.UNDIMMING:
            psvalues[psvalues < 1.0] = 1.0
        elif balance_mode == BalanceMode.DIMMING:
            psvalues[psvalues > 1.0] = 1.0

        psvalues[(abs(psvalues - curr_value) > delta_thr)] = curr_value

        def _get_cont(mode: BalanceWeightMode, frange: range) -> Any:
            if mode == BalanceWeightMode.INTERPOLATE:
                if radius < 1:
                    raise CustomValueError(auto_balance, "Radius has to be >= 1 with BalanceWeightMode.INTERPOLATE!")

                weight = (n - (frange.start - 1)) / (frange.stop - (frange.start - 1))

                weighted_prev = psvalues[middle_idx - 1] * (1 - weight)
                weighted_next = psvalues[middle_idx + 1] * weight

                return weighted_prev + weighted_next

            if mode == BalanceWeightMode.MEDIAN:
                return np.median(psvalues)

            if mode == BalanceWeightMode.MEAN:
                return psvalues.mean()

            if mode == BalanceWeightMode.MAX:
                return psvalues.max()

            if mode == BalanceWeightMode.MIN:
                return psvalues.min()

            return psvalues[middle_idx]

        if override:
            frange, cont, override_mode = override

            if override_mode == BalanceWeightMode.NONE:
                return nobalanceclip

            if cont is not None:
                psvalues[
                    max(0, middle_idx - (n - frange.start)) : min(len(psvalues), middle_idx + (frange.stop - n))
                ] = cont

            if override_mode != weight_mode:
                cont = _get_cont(override_mode, frange)
        else:
            cont = _get_cont(weight_mode, clipfrange)

        sat = (cont - 1) * relative_sat + 1

        fix = tweak_clip(clip, cont, sat, range_in=range_in)

        if prop:
            return fix.std.SetFrameProps(AutoBalance=True, AutoBalanceCont=cont, AutoBalanceSat=sat)

        return fix

    stats_clips = [
        *(ref_stats[0] * i + ref_stats[:-i] for i in range(1, radius + 1)),
        ref_stats,
        *(ref_stats[i:] + ref_stats[-1] * i for i in range(1, radius + 1)),
    ]

    return clip.std.FrameEval(_autobalance, stats_clips, clip)

def tweak_clip(
    clip: vs.VideoNode,
    cont: float = 1.0,
    sat: float = 1.0,
    bright: float = 0.0,
    hue: float = 0.0,
    relative_sat: float | None = None,
    range_in: ColorRange | None = None,
    range_out: ColorRange | None = None,
    clamp: bool = True,
    pre: vs.VideoNode | Callable[[vs.VideoNode], vs.VideoNode] | None = None,
) -> vs.VideoNode:
    assert clip.format

    bits = get_depth(clip)

    range_in = fallback(range_in, ColorRange(clip))
    range_out = fallback(range_out, range_in)

    sv_args_out = dict[str, Any](
        input_depth=8, output_depth=bits, range_in=range_in, range_out=range_out, scale_offsets=True
    )

    luma_min = scale_value(16, **sv_args_out)
    chroma_min = scale_value(16, **sv_args_out, chroma=True)

    luma_max = scale_value(235, **sv_args_out)
    chroma_max = scale_value(240, **sv_args_out, chroma=True)

    chroma_center = get_neutral_value(clip)

    if relative_sat is not None:
        sat = cont if cont == 1.0 or relative_sat == 1.0 else (cont - 1.0) * relative_sat + 1.0

    cont = max(cont, 0.0)
    sat = max(sat, 0.0)

    if (hue == bright == 0.0) and (sat == cont == 1.0):
        return clip

    pre_clip = pre(clip) if callable(pre) else fallback(pre, clip)

    clips = [pre_clip]

    yexpr = list[Any](["x"])
    cexpr = list[Any](["x"])

    if (hue != 0.0 or sat != 1.0) and clip.format.color_family != vs.GRAY:
        hue *= pi / degrees(pi)

        hue_sin, hue_cos = sin(hue), cos(hue)

        normalize = [chroma_center, ExprOp.SUB]

        cexpr.extend([normalize, hue_cos, sat, ExprOp.MUL * 2])

        if hue != 0:
            clips += [pre_clip.std.ShufflePlanes([0, 2, 1], vs.YUV)]
            cexpr.extend(["y", normalize, hue_sin, sat, ExprOp.MUL * 2, ExprOp.ADD])

        cexpr.extend([chroma_center, ExprOp.ADD])

        if clamp and range_out:
            cexpr.extend(StrList([chroma_min, ExprOp.MAX, chroma_max, ExprOp.MIN]))

    if bright != 0 or cont != 1:
        if luma_min > 0:
            yexpr.extend([luma_min, ExprOp.SUB])

        if cont != 1:
            yexpr.extend([cont, ExprOp.MUL])

        if (luma_min + bright) != 0:
            yexpr.extend([luma_min, bright, ExprOp.ADD * 2])

        if clamp and range_out:
            yexpr.extend(StrList([luma_min, ExprOp.MAX, luma_max, ExprOp.MIN]))

    tclip = norm_expr(clips, (yexpr, cexpr))

    return tclip