Spaces:
Running
on
Zero
Running
on
Zero
import os | |
import torch | |
from gfpgan import GFPGANer | |
from tqdm import tqdm | |
from .videoio import load_video_to_cv2 | |
import cv2 | |
class GeneratorWithLen(object): | |
""" From https://stackoverflow.com/a/7460929 """ | |
def __init__(self, gen, length): | |
self.gen = gen | |
self.length = length | |
def __len__(self): | |
return self.length | |
def __iter__(self): | |
return self.gen | |
def enhancer_list(images, method='gfpgan', bg_upsampler='realesrgan'): | |
gen = enhancer_generator_no_len(images, method=method, bg_upsampler=bg_upsampler) | |
return list(gen) | |
def enhancer_generator_with_len(images, method='gfpgan', bg_upsampler='realesrgan'): | |
""" Provide a generator with a __len__ method so that it can passed to functions that | |
call len()""" | |
if os.path.isfile(images): # handle video to images | |
# TODO: Create a generator version of load_video_to_cv2 | |
images = load_video_to_cv2(images) | |
gen = enhancer_generator_no_len(images, method=method, bg_upsampler=bg_upsampler) | |
gen_with_len = GeneratorWithLen(gen, len(images)) | |
return gen_with_len | |
def enhancer_generator_no_len(images, method='gfpgan', bg_upsampler='realesrgan'): | |
""" Provide a generator function so that all of the enhanced images don't need | |
to be stored in memory at the same time. This can save tons of RAM compared to | |
the enhancer function. """ | |
print('face enhancer....') | |
if not isinstance(images, list) and os.path.isfile(images): # handle video to images | |
images = load_video_to_cv2(images) | |
# ------------------------ set up GFPGAN restorer ------------------------ | |
if method == 'gfpgan': | |
arch = 'clean' | |
channel_multiplier = 2 | |
model_name = 'GFPGANv1.4' | |
url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth' | |
elif method == 'RestoreFormer': | |
arch = 'RestoreFormer' | |
channel_multiplier = 2 | |
model_name = 'RestoreFormer' | |
url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth' | |
elif method == 'codeformer': # TODO: | |
arch = 'CodeFormer' | |
channel_multiplier = 2 | |
model_name = 'CodeFormer' | |
url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth' | |
else: | |
raise ValueError(f'Wrong model version {method}.') | |
# ------------------------ set up background upsampler ------------------------ | |
if bg_upsampler == 'realesrgan': | |
if not torch.cuda.is_available(): # CPU | |
import warnings | |
warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. ' | |
'If you really want to use it, please modify the corresponding codes.') | |
bg_upsampler = None | |
else: | |
from basicsr.archs.rrdbnet_arch import RRDBNet | |
from realesrgan import RealESRGANer | |
model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) | |
bg_upsampler = RealESRGANer( | |
scale=2, | |
model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth', | |
model=model, | |
tile=400, | |
tile_pad=10, | |
pre_pad=0, | |
half=True) # need to set False in CPU mode | |
else: | |
bg_upsampler = None | |
# determine model paths | |
model_path = os.path.join('gfpgan/weights', model_name + '.pth') | |
if not os.path.isfile(model_path): | |
model_path = os.path.join('checkpoints', model_name + '.pth') | |
if not os.path.isfile(model_path): | |
# download pre-trained models from url | |
model_path = url | |
restorer = GFPGANer( | |
model_path=model_path, | |
upscale=2, | |
arch=arch, | |
channel_multiplier=channel_multiplier, | |
bg_upsampler=bg_upsampler) | |
# ------------------------ restore ------------------------ | |
for idx in tqdm(range(len(images)), 'Face Enhancer:'): | |
img = cv2.cvtColor(images[idx], cv2.COLOR_RGB2BGR) | |
# restore faces and background if necessary | |
cropped_faces, restored_faces, r_img = restorer.enhance( | |
img, | |
has_aligned=False, | |
only_center_face=False, | |
paste_back=True) | |
r_img = cv2.cvtColor(r_img, cv2.COLOR_BGR2RGB) | |
yield r_img | |