|
import os |
|
os.system("pip install dlib") |
|
import sys |
|
import face_detection |
|
import PIL |
|
from PIL import Image, ImageOps, ImageFile |
|
import numpy as np |
|
import cv2 as cv |
|
import torch |
|
|
|
torch.set_grad_enabled(False) |
|
model = torch.jit.load('u2net_bce_itr_25000_train_3.856416_tar_0.547567-400x_360x.jit.pt') |
|
model.eval() |
|
|
|
|
|
|
|
def unsharp_mask(image, kernel_size=(5, 5), sigma=1.0, amount=2.0, threshold=0): |
|
"""Return a sharpened version of the image, using an unsharp mask.""" |
|
blurred = cv.GaussianBlur(image, kernel_size, sigma) |
|
sharpened = float(amount + 1) * image - float(amount) * blurred |
|
sharpened = np.maximum(sharpened, np.zeros(sharpened.shape)) |
|
sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape)) |
|
sharpened = sharpened.round().astype(np.uint8) |
|
if threshold > 0: |
|
low_contrast_mask = np.absolute(image - blurred) < threshold |
|
np.copyto(sharpened, image, where=low_contrast_mask) |
|
return sharpened |
|
|
|
def normPRED(d): |
|
ma = np.max(d) |
|
mi = np.min(d) |
|
|
|
dn = (d-mi)/(ma-mi) |
|
|
|
return dn |
|
|
|
def array_to_np(array_in): |
|
array_in = normPRED(array_in) |
|
array_in = np.squeeze(255.0*(array_in)) |
|
array_in = np.transpose(array_in, (1, 2, 0)) |
|
return array_in |
|
|
|
def array_to_image(array_in): |
|
array_in = normPRED(array_in) |
|
array_in = np.squeeze(255.0*(array_in)) |
|
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): |
|
image_in = np.array(image_in, np.float32) |
|
tmpImg = np.zeros((image_in.shape[0],image_in.shape[1],3)) |
|
image_in = image_in/np.max(image_in) |
|
if image_in.shape[2]==1: |
|
tmpImg[:,:,0] = (image_in[:,:,0]-0.485)/0.229 |
|
tmpImg[:,:,1] = (image_in[:,:,0]-0.485)/0.229 |
|
tmpImg[:,:,2] = (image_in[:,:,0]-0.485)/0.229 |
|
else: |
|
tmpImg[:,:,0] = (image_in[:,:,0]-0.485)/0.229 |
|
tmpImg[:,:,1] = (image_in[:,:,1]-0.456)/0.224 |
|
tmpImg[:,:,2] = (image_in[:,:,2]-0.406)/0.225 |
|
|
|
tmpImg = tmpImg.transpose((2, 0, 1)) |
|
image_out = np.expand_dims(tmpImg, 0) |
|
return image_out |
|
|
|
def find_aligned_face(image_in, size=400): |
|
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=400): |
|
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 face2doll( |
|
img: Image.Image, |
|
size: int |
|
) -> Image.Image: |
|
|
|
aligned_img = align_first_face(img) |
|
if aligned_img is None: |
|
output=None |
|
else: |
|
input = torch.Tensor(aligned_img) |
|
results = model(input) |
|
doll_np_image = array_to_np(results[1].detach().numpy()) |
|
doll_image = unsharp_mask(doll_np_image) |
|
doll_image = Image.fromarray(doll_image) |
|
|
|
output = img_concat_h(array_to_image(aligned_img), doll_image) |
|
del results |
|
|
|
return output |
|
|
|
def inference(img): |
|
out = face2doll(img, 400) |
|
return out |
|
|
|
|
|
title = "Face2Doll U2Net" |
|
description = "Style transfer a face into one of a \"Doll\". 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. Faces with glasses on, seem not to yield good results." |
|
article = "<hr><p style='text-align: center'>See the <a href='https://github.com/Norod/U-2-Net-StyleTransfer' target='_blank'>Github Repo</a></p><p style='text-align: center'>samples: <img src='https://hf.space/gradioiframe/Norod78/Face2Doll/file/Sample00001.jpg' alt='Sample00001'/><img src='https://hf.space/gradioiframe/Norod78/Face2Doll/file/Sample00002.jpg' alt='Sample00002'/><img src='https://hf.space/gradioiframe/Norod78/Face2Doll/file/Sample00003.jpg' alt='Sample00003'/><img src='https://hf.space/gradioiframe/Norod78/Face2Doll/file/Sample00004.jpg' alt='Sample00004'/><img src='https://hf.space/gradioiframe/Norod78/Face2Doll/file/Sample00005.jpg' alt='Sample00005'/></p><p>The \"Face2Doll (U2Net)\" 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, |
|
gr.inputs.Image(type="pil", label="Input"), |
|
gr.outputs.Image(type="pil", label="Output"), |
|
title=title, |
|
description=description, |
|
article=article, |
|
examples=examples, |
|
enable_queue=True, |
|
allow_flagging=False |
|
).launch() |
|
|
