ccd_new

ccd_new.py

from typing import cast

import lvsfunc
import vapoursynth as vs
from vsutil import EXPR_VARS, join, split

core = vs.core


def add_expr(n: int) -> str:
    return 'x y + ' + ' + '.join(EXPR_VARS[i] for i in range(2, n)) + ' +'


def ccd(clip: vs.VideoNode, threshold: float) -> vs.VideoNode:
    assert clip.format
    bits = clip.format.bits_per_sample
    is_float = clip.format.sample_type == vs.FLOAT
    peak = 1.0 if is_float else (1 << bits) - 1
    threshold /= peak
    # threshold = threshold ** 2 / 195075.0

    rgb = clip.resize.Bicubic(format=vs.RGBS)

    pre1 = rgb.resize.Point(
        clip.width+24, clip.height+24,
        src_left=-12, src_top=-12,
        src_width=src.width+24, src_height=src.height+24
    )
    pre2 = rgb.resize.Point(
        rgb.width+24, rgb.height+24,
        src_width=rgb.width+24, src_height=rgb.height+24
    )
    pre_planes = split(pre1)


    shift_planes_clips = [
        split(pre2.resize.Point(src_left=-x, src_top=-y))
        for x in range(0, 25, 8) for y in range(0, 25, 8)
    ]
    denoise_clips = [
        core.std.Expr(pre_planes + shift_planes, f'x a - dup * y b - dup * + z c - dup * + sqrt {threshold} <')
        for shift_planes in shift_planes_clips
    ]

    cond_planes_clips = [
        join([core.std.Expr([splane, dclip], 'y 0 > x 0 ?') for splane in splanes])
        for dclip, splanes in zip(denoise_clips, shift_planes_clips)
    ]

    denoise = core.std.Expr(denoise_clips, add_expr(len(denoise_clips)) + ' 1 +')
    denoise = join([denoise] * 3)

    n_op = len(cond_planes_clips) + 1
    avg = core.std.Expr([pre1] + cond_planes_clips + [denoise], add_expr(n_op) + f' {EXPR_VARS[n_op]} /')
    avg = avg.resize.Bicubic(format=clip.format.id, dither_type='error_diffusion', matrix=cast(int, clip.get_frame(0).props['_Matrix']))
    avg = avg.std.Crop(12, 12, 12, 12)


    assert avg.format
    return core.std.ShufflePlanes([clip, avg], [0, 1, 2], avg.format.color_family)


src = core.lsmas.LWLibavSource(
    r'G:\Encodages\ReZero [EN COURS]\REZERO_2ND_DISC5\BDMV\STREAM\00003.m2ts'
)
src = src.std.SetFrameProp('_Matrix', intval=1)
src.set_output(0)

lvsfunc.comparison.stack_planes(src).set_output(2)



denoisedd = ccd(src, 10)
denoisedd.text.Text('denoised').set_output(1)
lvsfunc.comparison.stack_planes(denoisedd).text.Text('denoised').set_output(3)