descreen vapoursynth

descreen.py

from typing import Any

import cv2
import numpy as np
import vapoursynth as vs
from numpy.typing import NDArray


src = core.imwri.Read(r'Scans_assault_loli\02.png')
src.set_output(0)


def normalize(h: int, w: int) -> NDArray[Any]:
    x = np.arange(w)
    y = np.arange(h)
    cx = np.abs(x - w // 2) ** 0.5
    cy = np.abs(y - h // 2) ** 0.5
    energy = cx[None, :] + cy[:, None]
    return np.maximum(energy * energy, 0.01)


def ellipse(w: int, h: int) -> NDArray[Any]:
    offset = (w + h) / 2. / (w * h)
    y, x = np.ogrid[-h: h + 1., -w: w + 1.]
    return np.array((x / w)**2 + (y / h)**2 - offset <= 1, np.uint8)


def frame_to_array(f: vs.VideoFrame) -> NDArray[Any]:
    # return np.dstack([f.get_read_array(i) for i in range(f.format.num_planes - 1, -1, -1)])
    return np.dstack([f.get_read_array(i) for i in range(f.format.num_planes)])


def array_to_frame(array: NDArray[Any], frame: vs.VideoFrame) -> vs.VideoFrame:
    f = frame.copy()
    # for i in range(f.format.num_planes - 1, -1, -1):
    for i in range(f.format.num_planes):
        np.copyto(
            np.asarray(f.get_write_array(i)),
            array[:, :, i],
            # casting="unsafe"
        )

    return f


def descreen(clip: vs.VideoNode, threshold: int = 92, radius: int = 6, middle: int = 4) -> vs.VideoNode:

    clip = clip.resize.Bicubic(format=vs.RGB24).std.ShufflePlanes([2, 1, 0], vs.RGB)

    rows, cols = clip.height, clip.width
    coefs = normalize(rows, cols)

    ew, eh = cols // (middle * 2), rows // (middle * 2)
    pw, ph = (cols - ew * 2) // 2, (rows - eh * 2) // 2

    mid = np.pad(ellipse(ew, eh), ((ph, rows - ph - eh * 2 - 1), (pw, cols - pw - ew * 2 - 1)), 'constant')

    def _process(n: int, f: vs.VideoFrame) -> vs.VideoFrame:
        array_frame = np.array(frame_to_array(f).transpose(2, 0, 1), np.float32)

        for i in range(3):
            fftimg = cv2.dft(array_frame[i], flags=18)
            fftimg = np.fft.fftshift(fftimg)

            spectrum = 20 * np.log(cv2.magnitude(fftimg[:, :, 0], fftimg[:, :, 1]) * coefs)

            _, thresh = cv2.threshold(np.float32(np.maximum(0, spectrum)), threshold, 255, cv2.THRESH_BINARY)
            thresh *= 1 - mid
            thresh = cv2.dilate(thresh, ellipse(radius, radius))
            thresh = cv2.GaussianBlur(thresh, (0, 0), radius / 3., 0, 0, cv2.BORDER_REPLICATE)
            thresh = 1 - thresh / 255

            img_back = fftimg * np.repeat(thresh[..., None], 2, axis=2)
            img_back = np.fft.ifftshift(img_back)
            img_back = cv2.idft(img_back)

            array_frame[i] = cv2.magnitude(img_back[:, :, 0], img_back[:, :, 1])

        return array_to_frame(
            np.round(array_frame).clip(0, 255).transpose(1, 2, 0).astype(np.uint8),
            f
        )

    return core.std.ModifyFrame(clip, clip, _process).std.ShufflePlanes([2, 1, 0], vs.RGB)


srcdec = descreen(src)
srcdec.set_output(1)