xingt/ImageUpscaler

1import cv22from basicsr.archs.rrdbnet_arch import RRDBNet3from web.log import logger4from realesrgan import RealESRGANer5from gfpgan import GFPGANer6import numpy as np7 8# restorer9def create_real_esrganer(model_path, num_block, scale=4):10    upsampler = RealESRGANer(11        scale=scale,12        model_path=model_path,13        dni_weight=None,14        model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=scale),15        tile=0,16        tile_pad=10,17        pre_pad=0,18        half=False,19        gpu_id=None)20 21    face_enhancer = GFPGANer(22        model_path='weights/GFPGANv1.4.pth',23        upscale=scale,24        arch='clean',25        channel_multiplier=2,26        bg_upsampler=upsampler)27    return upsampler, face_enhancer28 29def convert_img(path: str, origin: str, output_name: str, face_enhance_open:bool, amine:bool):30    scale = 431    if amine:32        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus_anime_6B.pth', 6, scale)33    else:34        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus.pth', 23, scale)35    img = cv2.imread(f"{path}/{origin}", cv2.IMREAD_UNCHANGED)36    logger.info("convert_img img shape: %s", img.shape)37    try:38        if face_enhance_open:39            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)40        else:41            output, _ = upsampler.enhance(img, outscale=scale)42    except Exception as error:43        logger.info('Error %s', error)44        logger.info('If you encounter CUDA out of memory, try to set --tile with a smaller number.')45    else:46        cv2.imwrite(f"{path}/{output_name}", output)47 48 49def convert_img_row(data: bytes, output_name: str, face_enhance_open:bool, amine:bool):50    if amine:51        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus_anime_6B.pth', 6)52    else:53        upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus.pth', 23)54    img = cv2.imdecode(np.fromstring(data, np.uint8), cv2.IMREAD_UNCHANGED)55    try:56        if face_enhance_open:57            _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)58        else:59            output, _ = upsampler.enhance(img, outscale=4)60    except RuntimeError as error:61        print('Error', error)62        print('If you encounter CUDA out of memory, try to set --tile with a smaller number.')63    else:64        return cv2.imencode(f'.{output_name}', output)[1].tobytes()65