Create paintingface

paintingface

@@ -0,0 +1,110 @@
+import os
+os.system("pip install dlib")
+import sys
+import face_detection
+from PIL import Image, ImageOps, ImageFile
+import numpy as np
+import cv2 as cv
+import torch
+import gradio as gr
+
+torch.set_grad_enabled(False)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", device=device).eval()
+model2 = torch.hub.load("AK391/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1",  device=device).eval()
+face2paint = torch.hub.load("bryandlee/animegan2-pytorch:main", "face2paint", device=device)
+image_format = "png" #@param ["jpeg", "png"]
+
+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
+
+# detect a face
+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
+
+# clip the face, return array
+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
+
+def paintface(img: Image.Image,size: int) -> Image.Image:
+    aligned_img = align_first_face(img,size)
+    if aligned_img is None:
+        output=None
+    else:
+        im_in = array_to_image(aligned_img).convert("RGB")
+        im_out1 = face2paint(model, im_in, side_by_side=False)
+        im_out2 = face2paint(model2, im_in, side_by_side=False)
+
+        output = img_concat_h(im_out1, im_out2)         
+    return output
+
+def generate(img):
+    out = paintface(img, 400)
+    return out