Spaces:
Running
Running
import os | |
import onnxruntime as rt | |
import sys | |
import PIL | |
from PIL import Image, ImageOps, ImageFile | |
import numpy as np | |
from pathlib import Path | |
import collections | |
from typing import Union, List | |
import scipy.ndimage | |
import requests | |
MODEL_FILE = "ffhqu2vintage512_pix2pixHD_v1E11-inp2inst-simp.onnx" | |
so = rt.SessionOptions() | |
so.inter_op_num_threads = 4 | |
so.intra_op_num_threads = 4 | |
session = rt.InferenceSession(MODEL_FILE, sess_options=so) | |
input_name = session.get_inputs()[0].name | |
print("input_name = " + str(input_name)) | |
output_name = session.get_outputs()[0].name | |
print("output_name = " + str(output_name)) | |
import face_detection | |
def array_to_image(array_in): | |
array_in = np.squeeze(255*(array_in + 1)/2) | |
array_in = np.transpose(array_in, (1, 2, 0)) | |
im = Image.fromarray(array_in.astype(np.uint8)) | |
return im | |
def image_as_array(image_in): | |
im_array = np.array(image_in, np.float32) | |
im_array = (im_array/255)*2 - 1 | |
im_array = np.transpose(im_array, (2, 0, 1)) | |
im_array = np.expand_dims(im_array, 0) | |
return im_array | |
def find_aligned_face(image_in, size=512): | |
aligned_image, n_faces, quad = face_detection.align(image_in, face_index=0, output_size=size) | |
return aligned_image, n_faces, quad | |
def align_first_face(image_in, size=512): | |
aligned_image, n_faces, quad = find_aligned_face(image_in,size=size) | |
if n_faces == 0: | |
try: | |
image_in = ImageOps.exif_transpose(image_in) | |
except: | |
print("exif problem, not rotating") | |
image_in = image_in.resize((size, size)) | |
im_array = image_as_array(image_in) | |
else: | |
im_array = image_as_array(aligned_image) | |
return im_array | |
def img_concat_h(im1, im2): | |
dst = Image.new('RGB', (im1.width + im2.width, im1.height)) | |
dst.paste(im1, (0, 0)) | |
dst.paste(im2, (im1.width, 0)) | |
return dst | |
import gradio as gr | |
def face2vintage( | |
img: Image.Image, | |
size: int | |
) -> Image.Image: | |
aligned_img = align_first_face(img) | |
if aligned_img is None: | |
output=None | |
else: | |
output = session.run([output_name], {input_name: aligned_img})[0] | |
output = array_to_image(output) | |
aligned_img = array_to_image(aligned_img).resize((output.width, output.height)) | |
output = img_concat_h(aligned_img, output) | |
return output | |
def inference(img): | |
out = face2vintage(img, 512) | |
return out | |
title = "Vintage style Pix2PixHD" | |
description = "Style a face to look more \"Vintage\". Upload an image with a face, or click on one of the examples below. If a face could not be detected, an image will still be created." | |
article = "<hr><p style='text-align: center'>See the <a href='https://github.com/justinpinkney/pix2pixHD' target='_blank'>Github Repo</a></p><p style='text-align: center'>samples: <img src='https://hf.space/gradioiframe/Norod78/VintageStyle/file/Sample00001.jpg' alt='Sample00001'/><img src='https://hf.space/gradioiframe/Norod78/VintageStyle/file/Sample00002.jpg' alt='Sample00002'/><img src='https://hf.space/gradioiframe/Norod78/VintageStyle/file/Sample00003.jpg' alt='Sample00003'/><img src='https://hf.space/gradioiframe/Norod78/VintageStyle/file/Sample00004.jpg' alt='Sample00004'/><img src='https://hf.space/gradioiframe/Norod78/VintageStyle/file/Sample00005.jpg' alt='Sample00005'/></p><p>The \"Vintage Style\" Pix2PixHD model was trained by <a href='https://linktr.ee/Norod78' target='_blank'>Doron Adler</a></p>" | |
examples=[['Example00001.jpg'],['Example00002.jpg'],['Example00003.jpg'],['Example00004.jpg'],['Example00005.jpg'], ['Example00006.jpg']] | |
gr.Interface( | |
inference, | |
inputs=[gr.Image(type="pil", label="Input")], | |
outputs=[gr.Image(type="pil", label="Output")], | |
title=title, | |
description=description, | |
article=article, | |
examples=examples, | |
enable_queue=True, | |
allow_flagging=False | |
).launch() |