upscaler/realesrgan_upscaler.py

from basicsr.archs.rrdbnet_arch import RRDBNet
from .real_esrgan.realesrgan import RealESRGANer

import cv2
import numpy as np
from PIL import Image


class RealESRGAN:
    def __init__(self, checkpoint_path):

        self.netscale = 4

        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)

        self.upsampler = RealESRGANer(
            scale=self.netscale,
            model_path=checkpoint_path,
            dni_weight=None,
            model=model,
            tile=0,
            tile_pad=10,
            pre_pad=0,
            half=True)
    
    def upscale(self, pil_image, scale_factor=3):
        cv2_img = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
        op, _ = self.upsampler.enhance(cv2_img, outscale=scale_factor)
        pil_image_fin = Image.fromarray(cv2.cvtColor(op, cv2.COLOR_BGR2RGB))
        return pil_image_fin