|
""" |
|
This file is used for deploying hugging face demo: |
|
https://huggingface.co/spaces/sczhou/CodeFormer |
|
""" |
|
|
|
import sys |
|
sys.path.append('CodeFormer') |
|
import os |
|
import cv2 |
|
import torch |
|
import torch.nn.functional as F |
|
import gradio as gr |
|
|
|
from torchvision.transforms.functional import normalize |
|
|
|
from basicsr.utils import imwrite, img2tensor, tensor2img |
|
from basicsr.utils.download_util import load_file_from_url |
|
from facelib.utils.face_restoration_helper import FaceRestoreHelper |
|
from basicsr.archs.rrdbnet_arch import RRDBNet |
|
from basicsr.utils.realesrgan_utils import RealESRGANer |
|
from facelib.utils.misc import is_gray |
|
|
|
from basicsr.utils.registry import ARCH_REGISTRY |
|
|
|
|
|
os.system("pip freeze") |
|
|
|
pretrain_model_url = { |
|
'codeformer': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth', |
|
'detection': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/detection_Resnet50_Final.pth', |
|
'parsing': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth', |
|
'realesrgan': 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth' |
|
} |
|
|
|
if not os.path.exists('CodeFormer/weights/CodeFormer/codeformer.pth'): |
|
load_file_from_url(url=pretrain_model_url['codeformer'], model_dir='CodeFormer/weights/CodeFormer', progress=True, file_name=None) |
|
if not os.path.exists('CodeFormer/weights/facelib/detection_Resnet50_Final.pth'): |
|
load_file_from_url(url=pretrain_model_url['detection'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None) |
|
if not os.path.exists('CodeFormer/weights/facelib/parsing_parsenet.pth'): |
|
load_file_from_url(url=pretrain_model_url['parsing'], model_dir='CodeFormer/weights/facelib', progress=True, file_name=None) |
|
if not os.path.exists('CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth'): |
|
load_file_from_url(url=pretrain_model_url['realesrgan'], model_dir='CodeFormer/weights/realesrgan', progress=True, file_name=None) |
|
|
|
|
|
torch.hub.download_url_to_file( |
|
'https://replicate.com/api/models/sczhou/codeformer/files/fa3fe3d1-76b0-4ca8-ac0d-0a925cb0ff54/06.png', |
|
'01.png') |
|
torch.hub.download_url_to_file( |
|
'https://replicate.com/api/models/sczhou/codeformer/files/a1daba8e-af14-4b00-86a4-69cec9619b53/04.jpg', |
|
'02.jpg') |
|
torch.hub.download_url_to_file( |
|
'https://replicate.com/api/models/sczhou/codeformer/files/542d64f9-1712-4de7-85f7-3863009a7c3d/03.jpg', |
|
'03.jpg') |
|
torch.hub.download_url_to_file( |
|
'https://replicate.com/api/models/sczhou/codeformer/files/a11098b0-a18a-4c02-a19a-9a7045d68426/010.jpg', |
|
'04.jpg') |
|
torch.hub.download_url_to_file( |
|
'https://replicate.com/api/models/sczhou/codeformer/files/7cf19c2c-e0cf-4712-9af8-cf5bdbb8d0ee/012.jpg', |
|
'05.jpg') |
|
torch.hub.download_url_to_file( |
|
'https://raw.githubusercontent.com/sczhou/CodeFormer/master/inputs/cropped_faces/0729.png', |
|
'06.png') |
|
|
|
def imread(img_path): |
|
img = cv2.imread(img_path) |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
return img |
|
|
|
|
|
def set_realesrgan(): |
|
half = True if torch.cuda.is_available() else False |
|
model = RRDBNet( |
|
num_in_ch=3, |
|
num_out_ch=3, |
|
num_feat=64, |
|
num_block=23, |
|
num_grow_ch=32, |
|
scale=2, |
|
) |
|
upsampler = RealESRGANer( |
|
scale=2, |
|
model_path="CodeFormer/weights/realesrgan/RealESRGAN_x2plus.pth", |
|
model=model, |
|
tile=400, |
|
tile_pad=40, |
|
pre_pad=0, |
|
half=half, |
|
) |
|
return upsampler |
|
|
|
upsampler = set_realesrgan() |
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
codeformer_net = ARCH_REGISTRY.get("CodeFormer")( |
|
dim_embd=512, |
|
codebook_size=1024, |
|
n_head=8, |
|
n_layers=9, |
|
connect_list=["32", "64", "128", "256"], |
|
).to(device) |
|
ckpt_path = "CodeFormer/weights/CodeFormer/codeformer.pth" |
|
checkpoint = torch.load(ckpt_path)["params_ema"] |
|
codeformer_net.load_state_dict(checkpoint) |
|
codeformer_net.eval() |
|
|
|
os.makedirs('output', exist_ok=True) |
|
|
|
def inference(image, face_align, background_enhance, face_upsample, upscale, codeformer_fidelity): |
|
"""Run a single prediction on the model""" |
|
try: |
|
|
|
only_center_face = False |
|
draw_box = False |
|
detection_model = "retinaface_resnet50" |
|
|
|
print('Inp:', image, background_enhance, face_upsample, upscale, codeformer_fidelity) |
|
face_align = face_align if face_align is not None else True |
|
background_enhance = background_enhance if background_enhance is not None else True |
|
face_upsample = face_upsample if face_upsample is not None else True |
|
upscale = upscale if (upscale is not None and upscale > 0) else 2 |
|
|
|
has_aligned = not face_align |
|
upscale = 1 if has_aligned else upscale |
|
|
|
img = cv2.imread(str(image), cv2.IMREAD_COLOR) |
|
print('\timage size:', img.shape) |
|
|
|
upscale = int(upscale) |
|
if upscale > 4: |
|
upscale = 4 |
|
if upscale > 2 and max(img.shape[:2])>1000: |
|
upscale = 2 |
|
if max(img.shape[:2]) > 1500: |
|
upscale = 1 |
|
background_enhance = False |
|
face_upsample = False |
|
|
|
face_helper = FaceRestoreHelper( |
|
upscale, |
|
face_size=512, |
|
crop_ratio=(1, 1), |
|
det_model=detection_model, |
|
save_ext="png", |
|
use_parse=True, |
|
device=device, |
|
) |
|
bg_upsampler = upsampler if background_enhance else None |
|
face_upsampler = upsampler if face_upsample else None |
|
|
|
if has_aligned: |
|
|
|
img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR) |
|
face_helper.is_gray = is_gray(img, threshold=5) |
|
if face_helper.is_gray: |
|
print('\tgrayscale input: True') |
|
face_helper.cropped_faces = [img] |
|
else: |
|
face_helper.read_image(img) |
|
|
|
num_det_faces = face_helper.get_face_landmarks_5( |
|
only_center_face=only_center_face, resize=640, eye_dist_threshold=5 |
|
) |
|
print(f'\tdetect {num_det_faces} faces') |
|
|
|
face_helper.align_warp_face() |
|
|
|
|
|
for idx, cropped_face in enumerate(face_helper.cropped_faces): |
|
|
|
cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True) |
|
normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True) |
|
cropped_face_t = cropped_face_t.unsqueeze(0).to(device) |
|
|
|
try: |
|
with torch.no_grad(): |
|
output = codeformer_net( |
|
cropped_face_t, w=codeformer_fidelity, adain=True |
|
)[0] |
|
restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1)) |
|
del output |
|
torch.cuda.empty_cache() |
|
except RuntimeError as error: |
|
print(f"Failed inference for CodeFormer: {error}") |
|
restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1)) |
|
|
|
restored_face = restored_face.astype("uint8") |
|
face_helper.add_restored_face(restored_face, cropped_face) |
|
|
|
|
|
if not has_aligned: |
|
|
|
if bg_upsampler is not None: |
|
|
|
bg_img = bg_upsampler.enhance(img, outscale=upscale)[0] |
|
else: |
|
bg_img = None |
|
face_helper.get_inverse_affine(None) |
|
|
|
if face_upsample and face_upsampler is not None: |
|
restored_img = face_helper.paste_faces_to_input_image( |
|
upsample_img=bg_img, |
|
draw_box=draw_box, |
|
face_upsampler=face_upsampler, |
|
) |
|
else: |
|
restored_img = face_helper.paste_faces_to_input_image( |
|
upsample_img=bg_img, draw_box=draw_box |
|
) |
|
else: |
|
restored_img = restored_face |
|
|
|
|
|
save_path = f'output/out.png' |
|
imwrite(restored_img, str(save_path)) |
|
|
|
restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB) |
|
return restored_img |
|
except Exception as error: |
|
print('Global exception', error) |
|
return None, None |
|
|
|
|
|
title = "CodeFormer: Robust Face Restoration and Enhancement Network" |
|
|
|
description = r"""<center><img src='https://user-images.githubusercontent.com/14334509/189166076-94bb2cac-4f4e-40fb-a69f-66709e3d98f5.png' alt='CodeFormer logo'></center> |
|
<br> |
|
<b>Official Gradio demo</b> for <a href='https://github.com/sczhou/CodeFormer' target='_blank'><b>Towards Robust Blind Face Restoration with Codebook Lookup Transformer (NeurIPS 2022)</b></a><br> |
|
🔥 CodeFormer is a robust face restoration algorithm for old photos or AI-generated faces.<br> |
|
🤗 Try CodeFormer for improved stable-diffusion generation!<br> |
|
""" |
|
|
|
article = r""" |
|
If CodeFormer is helpful, please help to ⭐ the <a href='https://github.com/sczhou/CodeFormer' target='_blank'>Github Repo</a>. Thanks! |
|
[](https://github.com/sczhou/CodeFormer) |
|
|
|
--- |
|
|
|
📝 **Citation** |
|
|
|
If our work is useful for your research, please consider citing: |
|
```bibtex |
|
@inproceedings{zhou2022codeformer, |
|
author = {Zhou, Shangchen and Chan, Kelvin C.K. and Li, Chongyi and Loy, Chen Change}, |
|
title = {Towards Robust Blind Face Restoration with Codebook Lookup TransFormer}, |
|
booktitle = {NeurIPS}, |
|
year = {2022} |
|
} |
|
``` |
|
|
|
📋 **License** |
|
|
|
This project is licensed under <a rel="license" href="https://github.com/sczhou/CodeFormer/blob/master/LICENSE">S-Lab License 1.0</a>. |
|
Redistribution and use for non-commercial purposes should follow this license. |
|
|
|
📧 **Contact** |
|
|
|
If you have any questions, please feel free to reach me out at <b>[email protected]</b>. |
|
|
|
🤗 **Find Me:** |
|
<style type="text/css"> |
|
td { |
|
padding-right: 0px !important; |
|
} |
|
|
|
.gradio-container-4-37-2 .prose table, .gradio-container-4-37-2 .prose tr, .gradio-container-4-37-2 .prose td, .gradio-container-4-37-2 .prose th { |
|
border: 0px solid #ffffff; |
|
border-bottom: 0px solid #ffffff; |
|
} |
|
|
|
</style> |
|
|
|
<table> |
|
<tr> |
|
<td><a href="https://github.com/sczhou"><img style="margin:-0.8em 0 2em 0" src="https://img.shields.io/github/followers/sczhou?style=social" alt="Github Follow"></a></td> |
|
<td><a href="https://twitter.com/ShangchenZhou"><img style="margin:-0.8em 0 2em 0" src="https://img.shields.io/twitter/follow/ShangchenZhou?label=%40ShangchenZhou&style=social" alt="Twitter Follow"></a></td> |
|
</tr> |
|
</table> |
|
|
|
<center><img src='https://api.infinitescript.com/badgen/count?name=sczhou/CodeFormer<ext=Visitors&color=6dc9aa' alt='visitors'></center> |
|
""" |
|
|
|
demo = gr.Interface( |
|
inference, [ |
|
gr.Image(type="filepath", label="Input"), |
|
gr.Checkbox(value=True, label="Pre_Face_Align"), |
|
gr.Checkbox(value=True, label="Background_Enhance"), |
|
gr.Checkbox(value=True, label="Face_Upsample"), |
|
gr.Number(value=2, label="Rescaling_Factor (up to 4)"), |
|
gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity (0 for better quality, 1 for better identity)') |
|
], [ |
|
gr.Image(type="numpy", label="Output") |
|
], |
|
title=title, |
|
description=description, |
|
article=article, |
|
examples=[ |
|
['01.png', True, True, True, 2, 0.7], |
|
['02.jpg', True, True, True, 2, 0.7], |
|
['03.jpg', True, True, True, 2, 0.7], |
|
['04.jpg', True, True, True, 2, 0.1], |
|
['05.jpg', True, True, True, 2, 0.1], |
|
['06.png', False, True, True, 1, 0.5] |
|
], |
|
concurrency_limit=2 |
|
) |
|
|
|
DEBUG = os.getenv('DEBUG') == '1' |
|
|
|
demo.launch(debug=DEBUG, share=True) |
