Python Binding for realcugan-ncnn-vulkan with PyBind11
RealCUGAN is a Generative Adversarial Network (GAN) based model for image super-resolution (SR). This wrapper provides an easy-to-use interface for running the pre-trained RealCUGAN model.
System | Status | CPU (32bit) | CPU (64bit) | GPU (32bit) | GPU (64bit) |
---|---|---|---|---|---|
Linux (Clang) | — | ✅ | — | ✅ | |
Linux (GCC) | — | ✅ | — | ✅ | |
Windows | — | ✅ | — | ✅ | |
MacOS | — | ✅ | — | ✅ | |
MacOS (ARM) | — | ✅ | — | ✅ |
To use this package, simply install it via pip:
pip install realcugan-ncnn-py
For Linux user:
apt install -y libomp5 libvulkan-dev
Then, import the Realcugan class from the package:
from realcugan_ncnn_py import Realcugan
To initialize the model:
realcugan = Realcugan(gpuid: int = 0, tta_mode: bool = False, num_threads: int = 1, noise: int = -1, scale: int = 2, tilesize: int = 0, syncgap: int = 3, model: str = "models-se")
# model can be "models-se" or "models-pro" or "models-nose"
# or an absolute path to the models' directory
Here, gpuid specifies the GPU device to use (-1 means use CPU), tta_mode enables test-time augmentation, num_threads sets the number of threads for processing, noise specifies the level of noise to apply to the image (-1 to 3), scale is the scaling factor for super-resolution (1 to 4), tilesize specifies the tile size for processing (0 or >= 32), syncgap is the sync gap mode, and model specifies the name of the pre-trained model to use.
Once the model is initialized, you can use the upscale method to super-resolve your images:
from PIL import Image
realcugan = Realcugan(gpuid=0, scale=2, noise=3)
with Image.open("input.jpg") as image:
image = realcugan.process_pil(image)
image.save("output.jpg", quality=95)
import cv2
realcugan = Realcugan(gpuid=0, scale=2, noise=3)
image = cv2.imdecode(np.fromfile("input.jpg", dtype=np.uint8), cv2.IMREAD_COLOR)
image = realcugan.process_cv2(image)
cv2.imencode(".jpg", image)[1].tofile("output_cv2.jpg")
import subprocess as sp
# your ffmpeg parameters
command_out = [FFMPEG_BIN,........]
command_in = [FFMPEG_BIN,........]
pipe_out = sp.Popen(command_out, stdout=sp.PIPE, bufsize=10 ** 8)
pipe_in = sp.Popen(command_in, stdin=sp.PIPE)
realcugan = Realcugan(gpuid=0, scale=2, noise=3)
while True:
raw_image = pipe_out.stdout.read(src_width * src_height * 3)
if not raw_image:
break
raw_image = realcugan.process_bytes(raw_image, src_width, src_height, 3)
pipe_in.stdin.write(raw_image)
The project just only been tested in Ubuntu 18+ and Debian 9+ environments on Linux, so if the project does not work on your system, please try building it.
The following references were used in the development of this project:
nihui/realcugan-ncnn-vulkan - This project was the main inspiration for our work. It provided the core implementation of the Real-CUGAN algorithm using the ncnn and Vulkan libraries.
Real-CUGAN - Real-CUGAN is an AI super resolution model for anime images, trained in a million scale anime dataset, using the same architecture as Waifu2x-CUNet.
media2x/realcugan-ncnn-vulkan-python - This project was used as a reference for implementing the wrapper. Special thanks to the original author for sharing the code.
ncnn - ncnn is a high-performance neural network inference framework developed by Tencent AI Lab.
This project is licensed under the BSD 3-Clause - see the LICENSE file for details.