makeup-transfer / app.py
rashid996958's picture
tf version fix and example added
359aac9
raw
history blame
2.99 kB
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)
print(preds.shape)
ImageLoader.save_image(preds, 'output/scaled_2x.png')
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)
return cv2.cvtColor(cv2.imread("output/scaled_2x.png"), cv2.COLOR_BGR2RGB)
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
)
app.launch()