Spaces:
Runtime error
Runtime error
| import cv2 | |
| from basicsr.archs.rrdbnet_arch import RRDBNet | |
| from web.log import logger | |
| from realesrgan import RealESRGANer | |
| from gfpgan import GFPGANer | |
| import numpy as np | |
| # restorer | |
| def create_real_esrganer(model_path, num_block, scale=4): | |
| upsampler = RealESRGANer( | |
| scale=scale, | |
| model_path=model_path, | |
| dni_weight=None, | |
| model=RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=num_block, num_grow_ch=32, scale=scale), | |
| tile=0, | |
| tile_pad=10, | |
| pre_pad=0, | |
| half=False, | |
| gpu_id=None) | |
| face_enhancer = GFPGANer( | |
| model_path='weights/GFPGANv1.4.pth', | |
| upscale=scale, | |
| arch='clean', | |
| channel_multiplier=2, | |
| bg_upsampler=upsampler) | |
| return upsampler, face_enhancer | |
| def convert_img(path: str, origin: str, output_name: str, face_enhance_open:bool, amine:bool): | |
| scale = 4 | |
| if amine: | |
| upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus_anime_6B.pth', 6, scale) | |
| else: | |
| upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus.pth', 23, scale) | |
| img = cv2.imread(f"{path}/{origin}", cv2.IMREAD_UNCHANGED) | |
| logger.info("convert_img img shape: %s", img.shape) | |
| try: | |
| if face_enhance_open: | |
| _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True) | |
| else: | |
| output, _ = upsampler.enhance(img, outscale=scale) | |
| except Exception as error: | |
| logger.info('Error %s', error) | |
| logger.info('If you encounter CUDA out of memory, try to set --tile with a smaller number.') | |
| else: | |
| cv2.imwrite(f"{path}/{output_name}", output) | |
| def convert_img_row(data: bytes, output_name: str, face_enhance_open:bool, amine:bool): | |
| if amine: | |
| upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus_anime_6B.pth', 6) | |
| else: | |
| upsampler, face_enhancer = create_real_esrganer('weights/RealESRGAN_x4plus.pth', 23) | |
| img = cv2.imdecode(np.fromstring(data, np.uint8), cv2.IMREAD_UNCHANGED) | |
| try: | |
| if face_enhance_open: | |
| _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True) | |
| else: | |
| output, _ = upsampler.enhance(img, outscale=4) | |
| except RuntimeError as error: | |
| print('Error', error) | |
| print('If you encounter CUDA out of memory, try to set --tile with a smaller number.') | |
| else: | |
| return cv2.imencode(f'.{output_name}', output)[1].tobytes() | |