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()

# https://en.wikipedia.org/wiki/Unsharp_masking
# https://stackoverflow.com/a/55590133/1495606
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()