pixel_shifter.py

import vapoursynth as vs
import numpy as np
from vsutil import depth

core = vs.core

def to_rgbs(clip: vs.VideoNode, matrix: int = 1) -> vs.VideoNode:
    clip = depth(clip, 32).std.SetFrameProp('_Matrix', intval=matrix)
    clip = core.resize.Bicubic(clip, format=vs.RGBS)
    return clip

def to_yuvps(clip: vs.VideoNode, matrix: int = 1) -> vs.VideoNode:
    return core.resize.Bicubic(clip, format=vs.YUV420P16, matrix=matrix)

def vsclip_ndarray(clip):
    assert isinstance(clip, vs.VideoNode)

    num = clip.num_frames
    planes = clip.format.num_planes
    dn = []
    for n, f in enumerate(clip.frames(close=True)):
        d = []
        for p in range(planes):
            arr = np.array(f[p], copy=False)
            d.append(arr.reshape([1] + list(arr.shape) + [1]))
        dn.append(np.concatenate(d, axis=3))
    dn = np.concatenate(dn, axis=0)
    
    return dn

def float32_vsclip(s, clip=None):
    assert isinstance(s, np.ndarray)

    if len(s.shape) <= 3: s = s.reshape([1] + list(s.shape))
    num = s.shape[-4]
    height = s.shape[-3]
    width = s.shape[-2]
    planes = s.shape[-1]
    if clip is None:
        clip = core.std.BlankClip(None, width, height, vs.RGBS if planes == 3 else vs.GRAYS, num)

    def convert_func(n, f):
        fout = f.copy()
        for p in range(planes):
            d = np.array(fout[p], copy=False)
            np.copyto(d, s[n, :, :, p])
            del d
        return fout
    return core.std.ModifyFrame(clip, clip, convert_func)


clip = key
src = core.ffms2.Source(clip)
# src = src[0:10]
# src.set_output(0)
src = to_rgbs(src)

arr = vsclip_ndarray(src)
# print(type(arr))

# np.savez_compressed('data.npz', arr)
for frame in range(0, src.num_frames):
    for i in range(20, 1070):
        arr[frame][i][:] = arr[frame][i+1][:]
     
arr1 = float32_vsclip(arr)
arr1 = arr1.std.AssumeFPS(fpsnum=24000, fpsden=1001)
print(type(arr1))
# arr1.set_output(2)
final = depth( 
              to_yuvps(arr1.resize.Spline36(range_in_s="limited", range_s="full")), 10
              )

final.set_output()