zombie.py

# import sys
# sys.path.append("pix2pixHD")

from PIL import Image
import numpy as np
import align
import cv2

from faceparsing2 import get_face_mask

# sys.path.append(r"..\face-parsing.PyTorch")
# import inferface

# def get_model():
#     with open("opt.pkl","rb") as f:
#         opt = pickle.load(f)
    
#     return create_model(opt)

def normalized_tensor(pil):
    transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5),(0.5, 0.5, 0.5))])
    return transform(pil)

def tensor2im(tensor):
    return Image.fromarray((255 * ((tensor.cpu().numpy() * .5) + .5)).astype(np.uint8).transpose((1,2,0)))

def normalized_array(image):
    np_array = np.array(image)
    return np.expand_dims(((np_array / 255 - .5)/.5).astype(np.float32).transpose((2,0,1)),0)

def array2im(np_array):
    return Image.fromarray((255 * ((np_array * .5) + .5)).astype(np.uint8).transpose((1,2,0)))

def square_human_2_zombie(image,model):
    tensor = normalized_tensor(image)
    generated = model(tensor.unsqueeze(0))
    return tensor2im(generated.data[0])

def square_human_2_zombie_onnx(image,model):
    np_array = normalized_array(image)
    ort_inputs = {"input": np_array}
    ort_outs = model.run(None, ort_inputs)
    return array2im(ort_outs[0][0])

def human_2_zombie_onnx(rgb_image,model,output_size=512):
    square,_ = align.align(rgb_image,enable_padding=False,output_size=output_size)
    if square is None:
        return None
    return square_human_2_zombie_onnx(square,model)

def human_2_zombie(rgb_image,model,output_size=512):
    square,_ = align.align(rgb_image,enable_padding=False,output_size=output_size)
    if square is None:
        return None
    return square_human_2_zombie(square,model)

def human_path_2_zombie(path):
    rgb_image = ImageOpen(path).convert("RGB")
    return human_2_zombie(rgb_image)

def transition(path,model):
    rgb_image = ImageOpen(path).convert("RGB")
    square,_ = align.align(rgb_image,enable_padding=False)
    if square is None:
        return None
    zombie = square_human_2_zombie(square,model)
    return get_concat_h(square,zombie)

def transition_onnx(rgb_image,model):
    square,_ = align.align(rgb_image,enable_padding=False,model_type="mediapipe")
    # if square is None:
    #     return None

    #Take the largest square in the upper left corner
    #And resize it to 512x512
    # side = min(rgb_image.width,rgb_image.height)
    # square = rgb_image.crop((0,0,side,side))
    # square = square.resize((512,512))
    
    zombie = square_human_2_zombie_onnx(square,model)
    
    return get_concat_h(square,zombie)

def get_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 get_concat_multiple_h(*ims):
    to_return = ims[0]
    for i in range(1,len(ims)):
        to_return = get_concat_h(to_return,ims[i])
        
    return to_return

def ImageOpen(filepath):
    try:
        image=Image.open(filepath)

        for orientation in ExifTags.TAGS.keys():
            if ExifTags.TAGS[orientation]=='Orientation':
                break

        exif=dict(image._getexif().items())

        if exif[orientation] == 3:
            image=image.rotate(180, expand=True)
        elif exif[orientation] == 6:
            image=image.rotate(270, expand=True)
        elif exif[orientation] == 8:
            image=image.rotate(90, expand=True)

        image = image.convert('RGB')

        return image
    except (AttributeError, KeyError, IndexError):
        # cases: image don't have getexif
        return Image.open(filepath)

def do_face(aligned, box, im_array, ort_session):
        z = square_human_2_zombie_onnx(aligned,ort_session)

        t = cv2.getAffineTransform(np.array([[0,0],[0,511],[511,511]],dtype="float32"),box[:3,:].astype("float32"))

        h,w,c = im_array.shape
        face_in_place = cv2.warpAffine(np.array(z),t,(w,h))
        
        face_classes = get_face_mask(aligned)
        face_mask = np.logical_and(face_classes < 14, face_classes > 0)
        
        imagemask = cv2.warpAffine(face_mask.astype("uint8"),t,(w,h))

        imagemask = imagemask.astype("uint8")

        cv2.copyTo(face_in_place,imagemask,im_array)
        
def make_faces_zombie(path, facenet, ort_session):
    im = Image.open(path)
    im = im.convert(mode="RGB")
    im_array = np.array(im)

    return make_faces_zombie_from_array(im_array, facenet, ort_session)

def make_faces_zombie_from_array(im_array_rgb, facenet, ort_session):
    im_array_rgb = np.copy(im_array_rgb)

    faces = align.aligns(Image.fromarray(im_array_rgb),enable_padding=True,output_size=512,model_type="mediapipe")
    if faces is None:
        faces = []
    
    for aligned,box in faces:
        do_face(aligned,box,im_array_rgb,ort_session)

    return Image.fromarray(im_array_rgb)