Spaces:
Runtime error
Runtime error
File size: 8,664 Bytes
b794328 ab48a0d b794328 f5a1e92 5236fda f5a1e92 5ed8759 b794328 f5a1e92 ab48a0d f5a1e92 ab48a0d 5ed8759 ab48a0d f5a1e92 b794328 6cb67f8 b794328 77307a6 b794328 77307a6 5236fda 77307a6 7fec513 77307a6 7fec513 77307a6 b794328 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
import os
import cv2
import gradio as gr
import torch
from basicsr.archs.srvgg_arch import SRVGGNetCompact
from gfpgan.utils import GFPGANer
from realesrgan.utils import RealESRGANer
os.system("pip freeze")
# download weights
if not os.path.exists('realesr-general-x4v3.pth'):
os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .")
if not os.path.exists('GFPGANv1.2.pth'):
os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.2.pth -P .")
if not os.path.exists('GFPGANv1.3.pth'):
os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .")
if not os.path.exists('GFPGANv1.4.pth'):
os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth -P .")
if not os.path.exists('RestoreFormer.pth'):
os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth -P .")
if not os.path.exists('CodeFormer.pth'):
os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/CodeFormer.pth -P .")
if not os.path.exists('HanamichiSakuragi.jpg'):
torch.hub.download_url_to_file(
'https://haoluobo.com/wp-content/uploads/2023/01/%E6%A8%B1%E6%9C%A8%E8%8A%B1%E9%81%93.jpg',
'HanamichiSakuragi.jpg')
torch.hub.download_url_to_file(
'https://haoluobo.com/wp-content/uploads/2023/01/%E6%9D%8E%E4%B8%96%E6%B0%91.jpg',
'LiShiming.jpg')
torch.hub.download_url_to_file(
'https://haoluobo.com/wp-content/uploads/2023/01/%E4%B9%BE%E9%9A%86.jpg',
'QianLong.jpg')
torch.hub.download_url_to_file(
'https://user-images.githubusercontent.com/17445847/187401133-8a3bf269-5b4d-4432-b2f0-6d26ee1d3307.png',
'10045.png')
# background enhancer with RealESRGAN
model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
model_path = 'realesr-general-x4v3.pth'
half = True if torch.cuda.is_available() else False
upsampler = RealESRGANer(scale=4, model_path=model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
os.makedirs('output', exist_ok=True)
# def inference(img, version, scale, weight):
def inference(img, version, scale, blur_face):
scale = int(scale)
blur_face = int(blur_face)
if blur_face % 2 != 1:
blur_face += 1
if blur_face < 3:
blur_face = 0
# weight /= 100
print(img, version, scale)
if scale > 4:
scale = 4 # avoid too large scale value
try:
extension = os.path.splitext(os.path.basename(str(img)))[1]
img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
if len(img.shape) == 3 and img.shape[2] == 4:
img_mode = 'RGBA'
elif len(img.shape) == 2: # for gray inputs
img_mode = None
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
else:
img_mode = None
h, w = img.shape[0:2]
if h < 300:
img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
if version == 'v1.2':
face_enhancer = GFPGANer(
model_path='GFPGANv1.2.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
elif version == 'v1.3':
face_enhancer = GFPGANer(
model_path='GFPGANv1.3.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
elif version == 'v1.4':
face_enhancer = GFPGANer(
model_path='GFPGANv1.4.pth', upscale=2, arch='clean', channel_multiplier=2, bg_upsampler=upsampler)
elif version == 'RestoreFormer':
face_enhancer = GFPGANer(
model_path='RestoreFormer.pth', upscale=2, arch='RestoreFormer', channel_multiplier=2, bg_upsampler=upsampler)
# elif version == 'CodeFormer':
# face_enhancer = GFPGANer(
# model_path='CodeFormer.pth', upscale=2, arch='CodeFormer', channel_multiplier=2, bg_upsampler=upsampler)
try:
# _, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True, weight=weight)
face_helper = face_enhancer.face_helper
align_warp_face = face_helper.align_warp_face
def new_align_warp_face(*args, **kwargs):
align_warp_face(*args, **kwargs) # save_cropped_path
face_helper.org_cropped_faces = face_helper.cropped_faces
if blur_face >= 3:
face_helper.cropped_faces = [cv2.GaussianBlur(e, (blur_face, blur_face), 0) for e in face_helper.cropped_faces]
print("find face count:", len(face_helper.cropped_faces))
face_helper.align_warp_face = new_align_warp_face
_, _, output = face_enhancer.enhance(img, has_aligned=False, only_center_face=False, paste_back=True)
except RuntimeError as error:
print('Error', error)
try:
if scale != 2:
interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
h, w = img.shape[0:2]
output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
except Exception as error:
print('wrong scale input.', error)
if img_mode == 'RGBA': # RGBA images should be saved in png format
extension = 'png'
else:
extension = 'jpg'
save_path = f'output/out.{extension}'
cv2.imwrite(save_path, output)
output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
return (
output,
save_path,
[cv2.cvtColor(e, cv2.COLOR_BGR2RGB) for e in face_enhancer.face_helper.org_cropped_faces],
[cv2.cvtColor(e, cv2.COLOR_BGR2RGB) for e in face_enhancer.face_helper.restored_faces]
)
except Exception as error:
print('global exception', error)
return None, None
title = "GFPGAN: Practical Face Restoration Algorithm"
description = r"""Gradio demo for <a href='https://github.com/TencentARC/GFPGAN' target='_blank'><b>GFPGAN: Towards Real-World Blind Face Restoration with Generative Facial Prior</b></a>.<br>
It can be used to restore your **old photos** or improve **AI-generated faces**.<br>
To use it, simply upload your image.<br>
If GFPGAN is helpful, please help to ⭐ the <a href='https://github.com/TencentARC/GFPGAN' target='_blank'>Github Repo</a> and recommend it to your friends 😊<br>
This demo was forked by [vicalloy](https://github.com/vicalloy), add `face blur` param to optimize painting face enhance.
"""
article = r"""
[](https://github.com/TencentARC/GFPGAN/releases)
[](https://github.com/TencentARC/GFPGAN)
[](https://arxiv.org/abs/2101.04061)
If you have any question, please email 📧 `[email protected]` or `[email protected]`.
"""
with gr.Blocks() as demo:
gr.Markdown("<center><h1>%s</h1></center>" % title)
gr.Markdown(description)
with gr.Row(equal_height=False):
with gr.Column():
file_path = gr.components.Image(type="filepath", label="Input")
version = gr.components.Radio(['v1.2', 'v1.3', 'v1.4', 'RestoreFormer'], type="value", value='v1.4', label='version')
rescaling_factor = gr.components.Radio(['1', '2', '4'], type="value", value='2', label='Rescaling factor')
blur_face = gr.Slider(label='Blur face', minimum=0, maximum=55, value=0, step=1)
submit = gr.Button("Submit")
with gr.Column():
output_img = gr.components.Image(type="numpy", label="Output (The whole image)")
download = gr.components.File(label="Download the output image")
with gr.Row():
with gr.Column():
input_faces = gr.Gallery(label="Input faces").style(height="auto")
with gr.Column():
output_faces = gr.Gallery(label="Output faces").style(height="auto")
gr.Examples([['HanamichiSakuragi.jpg', 'v1.4', 2, 31], ['LiShiming.jpg', 'v1.4', 2, 3], ['QianLong.jpg', 'v1.4', 2, 3],
['10045.png', 'v1.4', 2, 0]], [file_path, version, rescaling_factor, blur_face])
gr.Markdown(article)
submit.click(
inference,
inputs=[file_path, version, rescaling_factor, blur_face],
outputs=[output_img, download, input_faces, output_faces]
)
demo.queue(concurrency_count=4)
demo.launch()
|