Skip to content

an inference lib for image/video restoration with VapourSynth support

License

Notifications You must be signed in to change notification settings

TensoRaws/ccrestoration

Repository files navigation

ccrestoration

codecov CI-test Release-pypi PyPI version GitHub

an inference lib for image/video restoration with VapourSynth support, compatible with many community models

Install

Make sure you have Python >= 3.9 and PyTorch >= 1.13 installed

pip install ccrestoration
  • Install VapourSynth (optional)

Start

cv2

a simple example to use the SISR (Single Image Super-Resolution) model to process an image (APISR)

import cv2
import numpy as np

from ccrestoration import AutoModel, ConfigType, SRBaseModel

model: SRBaseModel = AutoModel.from_pretrained(
    pretrained_model_name=ConfigType.RealESRGAN_APISR_RRDB_GAN_generator_2x,
)

img = cv2.imdecode(np.fromfile("test.jpg", dtype=np.uint8), cv2.IMREAD_COLOR)
img = model.inference_image(img)
cv2.imwrite("test_out.jpg", img)

VapourSynth

a simple example to use the VSR (Video Super-Resolution) model to process a video (AnimeSR)

import vapoursynth as vs
from vapoursynth import core

from ccrestoration import AutoModel, BaseModelInterface, ConfigType

model: BaseModelInterface = AutoModel.from_pretrained(
    pretrained_model_name=ConfigType.RealESRGAN_AnimeJaNai_HD_V3_Compact_2x,
    tile=None
)

clip = core.bs.VideoSource(source="s.mp4")
clip = core.resize.Bicubic(clip=clip, matrix_in_s="709", format=vs.RGBH)
clip = model.inference_video(clip)
clip = core.resize.Bicubic(clip=clip, matrix_s="709", format=vs.YUV420P16)
clip.set_output()

See more examples in the example directory, ccrestoration can register custom configurations and models to extend the functionality

Current Support

It still in development, the following models are supported:

Notice

  • All the architectures have been edited to normalize input and output, and automatic padding is applied. The input and output tensor shapes may differ from the original architectures. For SR models, the input and output are both 4D tensors in the shape of (b, c, h, w). For VSR models, the input and output are both 5D tensors in the shape of (b, l, c, h, w).

  • For VSR models with equal l in input and output (f1, f2, f3, f4 -> f1', f2', f3', f4'), you can directly extend from class VSRBaseModel. For VSR models that output only one frame (f-2, f-1, f0, f1, f2 -> f0'), you also need to set self.one_frame_out = True.

Reference

License

This project is licensed under the MIT - see the LICENSE file for details.