Spaces:
Running
Running
import gradio as gr | |
import tensorflow as tf | |
import numpy as np | |
import os, cv2, shutil | |
import torch | |
from basicsr.archs.srvgg_arch import SRVGGNetCompact | |
from gfpgan.utils import GFPGANer | |
from realesrgan.utils import RealESRGANer | |
from imageio.v2 import imread | |
from PIL import Image | |
os.system("pip freeze") | |
os.system("wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -P .") | |
os.system("wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth -P .") | |
pb = "weights/dmt.pb" | |
style_dim = 8 | |
img_size = 256 | |
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 = False | |
upsampler = RealESRGANer( | |
scale=4, | |
model_path=model_path, | |
model=model, | |
tile=0, | |
tile_pad=10, | |
pre_pad=0, | |
half=half | |
) | |
def preprocess(img): | |
return (img / 255. - 0.5) * 2 | |
def deprocess(img): | |
return (img + 1) / 2 | |
def load_image(path): | |
img = cv2.resize(imread(path), (img_size, img_size)) | |
img_ = np.expand_dims(preprocess(img), 0) | |
return img / 255., img_ | |
def inference(A, B): | |
with tf.Graph().as_default(): | |
output_graph_def = tf.compat.v1.GraphDef() | |
with open(pb, 'rb') as fr: | |
output_graph_def.ParseFromString(fr.read()) | |
tf.import_graph_def(output_graph_def, name='') | |
sess = tf.compat.v1.Session() | |
try: | |
sess.run(tf.compat.v1.global_variables_initializer()) | |
graph = tf.compat.v1.get_default_graph() | |
Xs = graph.get_tensor_by_name('decoder_1/g:0') | |
X = graph.get_tensor_by_name('X:0') | |
Y = graph.get_tensor_by_name('Y:0') | |
A_img, A_img_ = load_image(A) | |
B_img, B_img_ = load_image(B) | |
Xs_ = sess.run(Xs, feed_dict={X: A_img_, Y: B_img_}) | |
output = deprocess(Xs_)[0] | |
output = np.array(np.array(output) * 255, dtype=np.uint8) | |
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR) | |
face_enhancer = GFPGANer( | |
model_path='GFPGANv1.3.pth', | |
upscale=4, | |
arch='clean', | |
channel_multiplier=2, | |
bg_upsampler=upsampler | |
) | |
try: | |
_, _, output = face_enhancer.enhance(output, has_aligned=False, only_center_face=False, paste_back=True) | |
except RuntimeError as error: | |
print('Error:', error) | |
save_path = "output/result.png" | |
os.makedirs(os.path.dirname(save_path), exist_ok=True) | |
cv2.imwrite(save_path, output) | |
return output, save_path | |
finally: | |
sess.close() | |
def makeupTransfer(arr1, arr2): | |
print("-" * 8) | |
shutil.rmtree("input/", ignore_errors=True) | |
os.makedirs("input/", exist_ok=True) | |
output1 = cv2.cvtColor(arr1, cv2.COLOR_BGR2RGB) | |
output2 = cv2.cvtColor(arr2, cv2.COLOR_BGR2RGB) | |
cv2.imwrite("input/original.png", output1) | |
cv2.imwrite("input/ref.png", output2) | |
no_makeup = "input/original.png" | |
makeup = "input/ref.png" | |
output, save_path = inference(no_makeup, makeup) | |
output = cv2.imread(save_path) | |
return cv2.cvtColor(output, cv2.COLOR_BGR2RGB) | |
def clear_inputs(): | |
return None, None, None | |
no_makeup_examples = [ | |
['faces/no_makeup/xfsy_0226.png'], | |
['faces/no_makeup/XYUH-006.png'], | |
['faces/no_makeup/vSYYZ306.png'], | |
['faces/no_makeup/xfsy_0405.png'], | |
['faces/no_makeup/xfsy_0055.png'], | |
['faces/no_makeup/xfsy_0521.png'], | |
['faces/no_makeup/vSYYZ639.png'], | |
['faces/no_makeup/vSYYZ429.png'], | |
['faces/no_makeup/xfsy_0068.png'] | |
] | |
makeup_examples = [ | |
['faces/makeup/XMY-136.png'], | |
['faces/makeup/XMY-266.png'], | |
['faces/makeup/vRX916.png'], | |
['faces/makeup/vFG137.png'], | |
['faces/makeup/vFG756.png'], | |
['faces/makeup/XMY-074.png'], | |
['faces/makeup/vFG112.png'], | |
['faces/makeup/XMY-014.png'], | |
['faces/makeup/vFG56.png'] | |
] | |
with gr.Blocks() as app: | |
gr.Markdown("# Makeup Transfer Program") | |
gr.Markdown( | |
"""*DISCLAIMER*: Undesirable results may occur due to limitations in the algorithm | |
or the quality of the input image. To ensure the best possible output, please use high-resolution | |
and high-quality images. Low-quality, blurry, poorly lit images, or images that are not exclusively | |
of faces may lead to suboptimal results. Additionally, the algorithm is designed to work specifically | |
with face images in a 1:1 aspect ratio. Using non-face images or images with incorrect dimensions may | |
result in errors or inaccurate outputs. The developers of the program and the algorithm are not | |
responsible for inaccuracies caused by unsuitable input images, unexpected system behavior, or | |
deviations from the specified requirements.""" | |
) | |
gr.Markdown( | |
"""*NOTE*: Please press the "Clear" button to restart the program. This ensures that no | |
previous inputs and outputs are present.""" | |
) | |
with gr.Row(): | |
no_makeup_image = gr.Image(label="Upload Your Image", type='numpy') | |
makeup_image = gr.Image(label="Upload Target Makeup", type='numpy') | |
output_image = gr.Image(label="Output Image", type='numpy', format="png") | |
with gr.Row(): | |
clear_button = gr.ClearButton([no_makeup_image, makeup_image, output_image]) | |
transfer_button = gr.Button("Transfer Makeup", variant="primary") | |
transfer_button.click( | |
fn=makeupTransfer, | |
inputs=[no_makeup_image, makeup_image], | |
outputs=[output_image] | |
) | |
with gr.Row(): | |
gr.Examples(no_makeup_examples, inputs=[no_makeup_image], label="No Makeup Image Examples") | |
gr.Examples(makeup_examples, inputs=[makeup_image], label="Makeup Image Examples") | |
app.launch() |