makeup-transfer / app.py
tryontech's picture
Update app.py
fa59a68 verified
raw
history blame
5.92 kB
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()