Spaces:
Running
Running
| 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_0521.png', | |
| 'faces/makeup/vFG756.png'], | |
| [, | |
| 'faces/no_makeup/xfsy_0068.png', | |
| 'faces/makeup/XMY-136.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() |