VapourSynth-Super-Resolution-Helper icon indicating copy to clipboard operation
VapourSynth-Super-Resolution-Helper copied to clipboard

AMD/Intel GPU Support

Open AlphaAtlas opened this issue 6 years ago • 0 comments

AMD GPUs (and Intel IGPs) can theoretically run a variety of pretrained models with OpenCV. WolframRhodium wrote some example VapourSynth code just for that: https://github.com/WolframRhodium/muvsfunc/blob/master/Collections/examples/super_resolution_opencv.vpy

Here's a slightly modified version that ran a Caffe model on a lowly AMD 6620G IGP (with a tiny PNG texture):

# 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()

However, I can't seem to load any pretrained PyTorch models, getting the same error as this user when I try: https://answers.opencv.org/question/220422/cv2dnnreadnetfromtorchpretrained_model_path-throws-error/

Does OpenCV even support PyTorch models? It seemingly recognizes and tries to load the .pth file, but I haven't found any specific documentation on PyTorch support. I probably need to open an issue in the OpenCV GitHub repo or on their own support site if I can't resolve this problem.

AlphaAtlas avatar Nov 10 '19 01:11 AlphaAtlas