超解像GAN Real-ESRGAN をオブジェクト化して動かしてみる【Windows11】

import argparse
import cv2
import glob
import os
import time
from basicsr.archs.rrdbnet_arch import RRDBNet

from realesrgan import RealESRGANer
#from realesrgan.archs.srvgg_arch import SRVGGNetCompact


#GAN
class esrgan():
    def __init__(self, scale):
        """Inference demo for Real-ESRGAN.
        """
        # ==================== setting argument ====================== #
        parser = argparse.ArgumentParser()
        parser.add_argument('-s', '--outscale', type=float, default=scale, help='The final upsampling scale of the image')
        parser.add_argument('--suffix', type=str, default='out', help='Suffix of the restored image')
        parser.add_argument('-t', '--tile', type=int, default=0, help='Tile size, 0 for no tile during testing')
        parser.add_argument('--tile_pad', type=int, default=10, help='Tile padding')
        parser.add_argument('--pre_pad', type=int, default=0, help='Pre padding size at each border')
        parser.add_argument('--face_enhance', action='store_true', help='Use GFPGAN to enhance face')
        parser.add_argument('--half', action='store_true', help='Use half precision during inference')
        parser.add_argument(
            '--alpha_upsampler',
            type=str,
            default='realesrgan',
            help='The upsampler for the alpha channels. Options: realesrgan | bicubic')
        parser.add_argument(
            '--ext',
            type=str,
            default='auto',
            help='Image extension. Options: auto | jpg | png, auto means using the same extension as inputs')
        args = parser.parse_args()

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

        # determine model paths
        model_path = os.path.join('./weights/', "RealESRGAN_x4plus" + '.pth')
        # ================================ setting GAN ================================ #
        # restorer
        upsampler = RealESRGANer(
            scale=netscale,
            model_path=model_path,
            model=model,
            tile=args.tile,
            tile_pad=args.tile_pad,
            pre_pad=args.pre_pad,
            half=args.half
        )
        # ================================ init parameters =========================== #
        self.args = args
        self.upsampler = upsampler

    # action gan for webcam
    def gan_webcam(self, img_cv2):
            img = img_cv2
            try:
                output, _ = self.upsampler.enhance(img, outscale=self.args.outscale)
                return output, True
            except RuntimeError as error:
                print('Error', error)
                print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
                return img, False


if __name__ == '__main__':
    # init
    GAN = esrgan(4)
    start = time.time()

    # read cv2_image
    img = cv2.imread("./input.png")

    # gan
    output, ret = GAN.gan_webcam(img)

    # ret = True is successs
    if ret:
        print(time.time() - start)
        cv2.imshow("test", output)
        cv2.imwrite("output.png", output)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

if __name__ == '__main__':
    # init
    GAN = esrgan(4)

    # read cv2_image
    img = cv2.imread("./input.png")

    # gan
    output, ret = GAN.single_gan(img)

    # ret = True is success
    if ret:
        cv2.imshow("test", output)
        cv2.imwrite("./output.png", output)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

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

# determine model paths
model_path = os.path.join('./weights/', "RealESRGAN_x4plus" + '.pth')

# ================================ setting GAN ================================ #
# restorer
upsampler = RealESRGANer(
    scale=netscale,
    model_path=model_path,
    model=model,
    tile=args.tile,
    tile_pad=args.tile_pad,
    pre_pad=args.pre_pad,
    half=args.half
)
# ================================ init parameters =========================== #
self.args = args
self.upsampler = upsampler

# action gan for webcam
def gan_webcam(self, img_cv2):
    img = img_cv2
    try:
        output, _ = self.upsampler.enhance(img, outscale=self.args.outscale)
        return output, True
    except RuntimeError as error:
        print('Error', error)
        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')
        return img, False