alphaatlas / vapoursynth-super-resolution-helper Goto Github PK

# super resolution using OpenCV
# note: the input image to the network is not cropped, which might triggers out-of-memory error.

import os

# Set OpenCL device in format `<Platform>:<CPU|GPU|ACCELERATOR|nothing=GPU/CPU>:<DeviceName or ID>`
# examples: 'AMD:GPU:', ':GPU:1', 'Intel:CPU:', 
# https://github.com/opencv/opencv/wiki/OpenCL-optimizations#opencv-opencl-configuration-options
os.environ['OPENCV_OPENCL_DEVICE'] = 'AMD:GPU:' # use GPU to accelerate processing


import vapoursynth as vs
import cv2
import mvsfunc as mvf
import muvsfunc_numpy as mufnp

core = vs.get_core()


# global params
src = core.std.BlankClip(width=640, height=360, length=1000, format=vs.RGBS) # can be RGB/YUV/GRAY

if 'GPU' in os.environ['OPENCV_OPENCL_DEVICE']:
    if cv2.ocl.haveOpenCL() and cv2.ocl.useOpenCL():
        backend = cv2.dnn.DNN_BACKEND_OPENCV
        target = cv2.dnn.DNN_TARGET_OPENCL # available on NVIDIA GPU since OpenCV 4.0.1, but only works on Intel GPU before OpenCV 3.4.2
    else:
        backend = cv.dnn.DNN_BACKEND_DEFAULT
        target = cv2.dnn.DNN_TARGET_CPU

###Change SR parameters here:
sr_args = dict(up_scale=2, is_rgb_model=True, pad=(7,7,7,7))

#Example caffe model.
# https://github.com/alterzero/DBPN-caffe
# https://drive.google.com/drive/folders/1ahbeoEHkjxoo4NV1wReOmpoRWbl448z-?usp=sharing
#sr_args = dict(prototxt=r'DBPN_mat_2x.prototxt',
        caffe_model=r'DBPN_2x.caffemodel', up_scale=2, is_rgb_model=True)


# internel functions
def channel_last(arr):
    """Convert a CHW array to HWC."""
    ndim = arr.ndim
    return arr.swapaxes(ndim - 3, ndim - 2).swapaxes(ndim - 2, ndim - 1)


def super_resolution_core(img, net, pad=None, crop=None):
    if pad is not None:
        img = cv2.copyMakeBorder(img, *pad, 1)

    blob = cv2.dnn.blobFromImage(img)

    net.setInput(blob, '')

    super_res = net.forward()

    if img.ndim == 2:
        if crop is not None:
            return super_res[0, 0, crop[0]:-crop[1], crop[2]:-crop[3]]
        else:
            return super_res[0, 0, :, :]
    else:
        # the output is BGR rather than RGB so channel reversal is needed
        if crop is not None:
            return channel_last(super_res[0, ::-1, crop[0]:-crop[1], crop[2]:-crop[3]])
        else:
            return channel_last(super_res[0, ::-1, :, :])


def run_super_resolution(clip, prototxt, caffe_model, up_scale=2, is_rgb_model=True, pad=None, crop=None, backend=None, target=None):
    """ Super-Resolution without color family hadling
    """
    ###Put the path(s) to your model here!
    ###See: https://docs.opencv.org/3.4/d6/d0f/group__dnn.html#ga3b34fe7a29494a6a4295c169a7d32422
    net = cv2.dnn.readNet("Path to model")

    if backend is not None:
        net.setPreferableBackend(backend)

    if target is not None:
        net.setPreferableTarget(target)

    if up_scale != 1:
        blank = core.std.BlankClip(clip, width=clip.width*up_scale, height=clip.height*up_scale)
        super_res = mufnp.numpy_process([blank, clip], super_resolution_core, net=net, 
            input_per_plane=(not is_rgb_model), output_per_plane=(not is_rgb_model), pad=pad, crop=crop, 
            omit_first_clip=True)
    else:
        super_res = mufnp.numpy_process(clip, super_resolution_core, net=net, 
            input_per_plane=(not is_rgb_model), output_per_plane=(not is_rgb_model), pad=pad, crop=crop)

    return super_res


def super_resolution(clip, up_scale=2, is_rgb_model=True, pad=None, crop=None, backend=None, target=None, pre_upscale=False, upscale_uv=False, merge_residual=False):
    """ Super-Resolution with color family hadling

    The color space of the output depends on the algorithm
    """

    isGray = clip.format.color_family == vs.GRAY
    isRGB = clip.format.color_family == vs.RGB

    if is_rgb_model and not isRGB:
        clip = mvf.ToRGB(clip, depth=32)

    elif not is_rgb_model:
        if isRGB:
            clip = mvf.ToYUV(clip, depth=32)

        if not isGray and not upscale_uv: # isYUV/RGB and only upscale Y
            clip = mvf.GetPlane(clip)

    clip = mvf.Depth(clip, depth=32)

    if pre_upscale:
        clip = core.resize.Bicubic(clip, clip.width*up_scale, clip.height*up_scale, filter_param_a=0, filter_param_b=0.5)
        up_scale = 1

    super_res = run_super_resolution(clip, prototxt=prototxt, caffe_model=caffe_model, 
        up_scale=up_scale, is_rgb_model=is_rgb_model, pad=pad, crop=crop, backend=backend, target=target)

    if merge_residual:
        low_res = core.resize.Bicubic(clip, super_res.width, super_res.height, filter_param_a=0, filter_param_b=0.5)
        super_res = core.std.Expr([super_res, low_res], ['x y +'])

    return super_res

sr = super_resolution(src, **sr_args, backend=backend, target=target)

sr.set_output()

Existing VapourSynth Fatpack installation found in this directory. Please delete
 it or move this installer somewhere else!
If you just want to update everything, please use the updater batch file

Press ENTER to continue...

C:\VapourSynth64Portable\VapourSynth64>python.exe -m pip install vsgan
Collecting vsgan
  Using cached vsgan-1.0.8-py3-none-any.whl (9.9 kB)
Collecting vapoursynth
  Using cached VapourSynth-49.zip (388 kB)
    ERROR: Command errored out with exit status 1:
     command: 'C:\VapourSynth64Portable\VapourSynth64\python.exe' -c 'import sys, setuptools, tokenize; sys.argv[0] = '"'"'C:\\Users\\STREAM~1\\AppData\\Local\\Temp\\pip-install-b1yauu0s\\vapoursynth\\setup.py'"'"'; __file__='"'"'C:\\Users\\STREAM~1\\AppData\\Local\\Temp\\pip-install-b1yauu0s\\vapoursynth\\setup.py'"'"';f=getattr(tokenize, '"'"'open'"'"', open)(__file__);code=f.read().replace('"'"'\r\n'"'"', '"'"'\n'"'"');f.close();exec(compile(code, __file__, '"'"'exec'"'"'))' egg_info --egg-base 'C:\Users\STREAM~1\AppData\Local\Temp\pip-pip-egg-info-drpuisxm'
         cwd: C:\Users\STREAM~1\AppData\Local\Temp\pip-install-b1yauu0s\vapoursynth\
    Complete output (14 lines):
    Traceback (most recent call last):
      File "C:\Users\STREAM~1\AppData\Local\Temp\pip-install-b1yauu0s\vapoursynth\setup.py", line 60, in <module>
        dll_path = query(winreg.HKEY_LOCAL_MACHINE, REGISTRY_PATH, REGISTRY_KEY)
      File "C:\Users\STREAM~1\AppData\Local\Temp\pip-install-b1yauu0s\vapoursynth\setup.py", line 36, in query
        reg_key = winreg.OpenKey(hkey, path, 0, winreg.KEY_READ)
    FileNotFoundError: [WinError 2] The system cannot find the file specified

    During handling of the above exception, another exception occurred:

    Traceback (most recent call last):
      File "<string>", line 1, in <module>
      File "C:\Users\STREAM~1\AppData\Local\Temp\pip-install-b1yauu0s\vapoursynth\setup.py", line 63, in <module>
        raise OSError("Couldn't detect vapoursynth installation path")
    OSError: Couldn't detect vapoursynth installation path
    ----------------------------------------
ERROR: Command errored out with exit status 1: python setup.py egg_info Check the logs for full command output.

Traceback (most recent call last):
File "src\cython\vapoursynth.pyx", line 1946, in vapoursynth.vpy_evaluateScript
File "src\cython\vapoursynth.pyx", line 1947, in vapoursynth.vpy_evaluateScript
File "D:/VapourSynth64Portable/VapourSynthScripts/ProcessVideoAuto.vpy", line 62, in 
clip = muf.super_resolution(clip, **manual_sr_args)
File "D:\VapourSynth64Portable\Scripts\muvsfunc.py", line 4561, in super_resolution
super_res = inference(clip, device_id[0])
File "D:\VapourSynth64Portable\Scripts\muvsfunc.py", line 4533, in inference
ctx=2 if dev_id >= 0 else 1, dev_id=max(dev_id, 0))
File "src\cython\vapoursynth.pyx", line 1825, in vapoursynth.Function.__call__
vapoursynth.Error: Predict: Function does not take argument(s) named padding