File size: 3,071 Bytes
0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f 8e07454 0e3fe1f d07f96e 0e3fe1f e6c483f 0e3fe1f 8e07454 d07f96e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import os
import onnxruntime as rt
import sys
from PIL import Image, ImageOps
import numpy as np
from pathlib import Path
import collections
from typing import Union, List
import scipy.ndimage
import requests
import gradio as gr
import face_detection # Ensure this is the adjusted face_detection.py
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))
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
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."
article = "<hr><p style='text-align: center'>See the <a href='https://github.com/justinpinkney/pix2pixHD' target='_blank'>Github Repo</a></p>"
examples = [['Example00001.jpg'], ['Example00002.jpg'], ['Example00003.jpg'], ['Example00004.jpg'], ['Example00005.jpg']]
demo = gr.Interface(
inference,
inputs=[gr.Image(type="pil", label="Input")],
outputs=[gr.Image(type="pil", label="Output")],
title=title,
description=description,
article=article,
)
demo.queue()
demo.launch() |