makeup-transfer / app.py
rashid996958's picture
Update app.py
84a7577 verified
import tensorflow as tf
import numpy as np
from imageio.v2 import imread
import os, glob, cv2, shutil
from super_image import EdsrModel, ImageLoader
from PIL import Image
import gradio as gr
pb = 'dmt.pb'
style_dim = 8
img_size=256
model_scale = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)
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()
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)
image = Image.fromarray(output)
inputs = ImageLoader.load_image(image)
preds = model_scale(inputs)
ImageLoader.save_image(preds, 'output/scaled_2x.png')
# return output
def makeupTransfer(arr1,arr2):
print("-"*8)
shutil.rmtree("input/")
os.makedirs("input/")
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"
inference(no_makeup, makeup)
img = cv2.imread("output/scaled_2x.png")
return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# return tr_out
examples = [['faces/no_makeup/xfsy_0226.png', 'faces/makeup/XMY-136.png'],
['faces/no_makeup/XYUH-006.png', 'faces/makeup/XMY-266.png'],
['faces/no_makeup/vSYYZ306.png', 'faces/makeup/vRX916.png'],
['faces/no_makeup/xfsy_0405.png', 'faces/makeup/vFG137.png'],
['faces/no_makeup/xfsy_0055.png', 'faces/makeup/vFG756.png'],
['faces/no_makeup/xfsy_0521.png', 'faces/makeup/XMY-074.png'],
['faces/no_makeup/vSYYZ639.png', 'faces/makeup/vFG112.png'],
['faces/no_makeup/vSYYZ429.png', 'faces/makeup/XMY-014.png'],
['faces/no_makeup/xfsy_0068.png', 'faces/makeup/vFG56.png']]
app = gr.Interface(fn=makeupTransfer,
inputs=[gr.Image(label="Reference Image",type='numpy'),
gr.Image(label="Makeup Image",type='numpy')],
outputs=gr.Image(label="Makeup Transfer Image",type='numpy'),
title="MakeUp Transfer APP",
examples=examples,
cache_examples=True
)
app.launch(share=True)