face-swap / utils /utils.py
felixrosberg's picture
with private models
69c590e
raw
history blame
14.2 kB
import json
from tensorflow.keras.models import model_from_json
from networks.layers import AdaIN, AdaptiveAttention
import tensorflow as tf
import numpy as np
import cv2
import math
from skimage import transform as trans
from scipy.signal import convolve2d
from skimage.color import rgb2yuv, yuv2rgb
from PIL import Image
def save_model_internal(model, path, name, num):
json_model = model.to_json()
with open(path + name + '.json', "w") as json_file:
json_file.write(json_model)
model.save_weights(path + name + '_' + str(num) + '.h5')
def load_model_internal(path, name, num):
with open(path + name + '.json', 'r') as json_file:
model_dict = json_file.read()
mod = model_from_json(model_dict, custom_objects={'AdaIN': AdaIN, 'AdaptiveAttention': AdaptiveAttention})
mod.load_weights(path + name + '_' + str(num) + '.h5')
return mod
def save_training_meta(state_dict, path, num):
with open(path + str(num) + '.json', 'w') as json_file:
json.dump(state_dict, json_file, indent=2)
def load_training_meta(path, num):
with open(path + str(num) + '.json', 'r') as json_file:
state_dict = json.load(json_file)
return state_dict
def log_info(sw, results_dict, iteration):
with sw.as_default():
for key in results_dict.keys():
tf.summary.scalar(key, results_dict[key], step=iteration)
src1 = np.array([[51.642, 50.115], [57.617, 49.990], [35.740, 69.007],
[51.157, 89.050], [57.025, 89.702]],
dtype=np.float32)
# <--left
src2 = np.array([[45.031, 50.118], [65.568, 50.872], [39.677, 68.111],
[45.177, 86.190], [64.246, 86.758]],
dtype=np.float32)
# ---frontal
src3 = np.array([[39.730, 51.138], [72.270, 51.138], [56.000, 68.493],
[42.463, 87.010], [69.537, 87.010]],
dtype=np.float32)
# -->right
src4 = np.array([[46.845, 50.872], [67.382, 50.118], [72.737, 68.111],
[48.167, 86.758], [67.236, 86.190]],
dtype=np.float32)
# -->right profile
src5 = np.array([[54.796, 49.990], [60.771, 50.115], [76.673, 69.007],
[55.388, 89.702], [61.257, 89.050]],
dtype=np.float32)
src = np.array([src1, src2, src3, src4, src5])
src_map = {112: src, 224: src * 2}
# Left eye, right eye, nose, left mouth, right mouth
arcface_src = np.array(
[[38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
[41.5493, 92.3655], [70.7299, 92.2041]],
dtype=np.float32)
arcface_src = np.expand_dims(arcface_src, axis=0)
def extract_face(img, bb, absolute_center, mode='arcface', extention_rate=0.05, debug=False):
"""Extract face from image given a bounding box"""
# bbox
x1, y1, x2, y2 = bb + 60
adjusted_absolute_center = (absolute_center[0] + 60, absolute_center[1] + 60)
if debug:
print(bb + 60)
x1, y1, x2, y2 = bb
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 3)
cv2.circle(img, absolute_center, 1, (255, 0, 255), 2)
Image.fromarray(img).show()
x1, y1, x2, y2 = bb + 60
# Pad image in case face is out of frame
padded_img = np.zeros(shape=(248, 248, 3), dtype=np.uint8)
padded_img[60:-60, 60:-60, :] = img
if debug:
cv2.rectangle(padded_img, (x1, y1), (x2, y2), (0, 255, 255), 3)
cv2.circle(padded_img, adjusted_absolute_center, 1, (255, 255, 255), 2)
Image.fromarray(padded_img).show()
y_len = abs(y1 - y2)
x_len = abs(x1 - x2)
new_len = (y_len + x_len) // 2
extension = int(new_len * extention_rate)
x_adjust = (x_len - new_len) // 2
y_adjust = (y_len - new_len) // 2
x_1_adjusted = x1 + x_adjust - extension
x_2_adjusted = x2 - x_adjust + extension
if mode == 'arcface':
y_1_adjusted = y1 - extension
y_2_adjusted = y2 - 2 * y_adjust + extension
else:
y_1_adjusted = y1 + 2 * y_adjust - extension
y_2_adjusted = y2 + extension
move_x = adjusted_absolute_center[0] - (x_1_adjusted + x_2_adjusted) // 2
move_y = adjusted_absolute_center[1] - (y_1_adjusted + y_2_adjusted) // 2
x_1_adjusted = x_1_adjusted + move_x
x_2_adjusted = x_2_adjusted + move_x
y_1_adjusted = y_1_adjusted + move_y
y_2_adjusted = y_2_adjusted + move_y
# print(y_1_adjusted, y_2_adjusted, x_1_adjusted, x_2_adjusted)
return padded_img[y_1_adjusted:y_2_adjusted, x_1_adjusted:x_2_adjusted]
def distance(a, b):
return np.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2)
def euclidean_distance(a, b):
x1 = a[0]; y1 = a[1]
x2 = b[0]; y2 = b[1]
return np.sqrt(((x2 - x1) * (x2 - x1)) + ((y2 - y1) * (y2 - y1)))
def align_face(img, landmarks, debug=False):
nose, right_eye, left_eye = landmarks
left_eye_x = left_eye[0]
left_eye_y = left_eye[1]
right_eye_x = right_eye[0]
right_eye_y = right_eye[1]
center_eye = ((left_eye[0] + right_eye[0]) // 2, (left_eye[1] + right_eye[1]) // 2)
if left_eye_y < right_eye_y:
point_3rd = (right_eye_x, left_eye_y)
direction = -1
else:
point_3rd = (left_eye_x, right_eye_y)
direction = 1
if debug:
cv2.circle(img, point_3rd, 1, (255, 0, 0), 1)
cv2.circle(img, center_eye, 1, (255, 0, 0), 1)
cv2.line(img, right_eye, left_eye, (0, 0, 0), 1)
cv2.line(img, left_eye, point_3rd, (0, 0, 0), 1)
cv2.line(img, right_eye, point_3rd, (0, 0, 0), 1)
a = euclidean_distance(left_eye, point_3rd)
b = euclidean_distance(right_eye, left_eye)
c = euclidean_distance(right_eye, point_3rd)
cos_a = (b * b + c * c - a * a) / (2 * b * c)
angle = np.arccos(cos_a)
angle = (angle * 180) / np.pi
if direction == -1:
angle = 90 - angle
ang = math.radians(direction * angle)
else:
ang = math.radians(direction * angle)
angle = 0 - angle
M = cv2.getRotationMatrix2D((64, 64), angle, 1)
new_img = cv2.warpAffine(img, M, (128, 128),
flags=cv2.INTER_CUBIC)
rotated_nose = (int((nose[0] - 64) * np.cos(ang) - (nose[1] - 64) * np.sin(ang) + 64),
int((nose[0] - 64) * np.sin(ang) + (nose[1] - 64) * np.cos(ang) + 64))
rotated_center_eye = (int((center_eye[0] - 64) * np.cos(ang) - (center_eye[1] - 64) * np.sin(ang) + 64),
int((center_eye[0] - 64) * np.sin(ang) + (center_eye[1] - 64) * np.cos(ang) + 64))
abolute_center = (rotated_center_eye[0], (rotated_nose[1] + rotated_center_eye[1]) // 2)
if debug:
cv2.circle(new_img, rotated_nose, 1, (0, 0, 255), 1)
cv2.circle(new_img, rotated_center_eye, 1, (0, 0, 255), 1)
cv2.circle(new_img, abolute_center, 1, (0, 0, 255), 1)
return new_img, abolute_center
def estimate_norm(lmk, image_size=112, mode='arcface', shrink_factor=1.0):
assert lmk.shape == (5, 2)
tform = trans.SimilarityTransform()
lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
min_M = []
min_index = []
min_error = float('inf')
src_factor = image_size / 112
if mode == 'arcface':
src = arcface_src * shrink_factor + (1 - shrink_factor) * 56
src = src * src_factor
else:
src = src_map[image_size] * src_factor
for i in np.arange(src.shape[0]):
tform.estimate(lmk, src[i])
M = tform.params[0:2, :]
results = np.dot(M, lmk_tran.T)
results = results.T
error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1)))
# print(error)
if error < min_error:
min_error = error
min_M = M
min_index = i
return min_M, min_index
def inverse_estimate_norm(lmk, t_lmk, image_size=112, mode='arcface', shrink_factor=1.0):
assert lmk.shape == (5, 2)
tform = trans.SimilarityTransform()
lmk_tran = np.insert(lmk, 2, values=np.ones(5), axis=1)
min_M = []
min_index = []
min_error = float('inf')
src_factor = image_size / 112
if mode == 'arcface':
src = arcface_src * shrink_factor + (1 - shrink_factor) * 56
src = src * src_factor
else:
src = src_map[image_size] * src_factor
for i in np.arange(src.shape[0]):
tform.estimate(t_lmk, lmk)
M = tform.params[0:2, :]
results = np.dot(M, lmk_tran.T)
results = results.T
error = np.sum(np.sqrt(np.sum((results - src[i])**2, axis=1)))
# print(error)
if error < min_error:
min_error = error
min_M = M
min_index = i
return min_M, min_index
def norm_crop(img, landmark, image_size=112, mode='arcface', shrink_factor=1.0):
"""
Align and crop the image based of the facial landmarks in the image. The alignment is done with
a similarity transformation based of source coordinates.
:param img: Image to transform.
:param landmark: Five landmark coordinates in the image.
:param image_size: Desired output size after transformation.
:param mode: 'arcface' aligns the face for the use of Arcface facial recognition model. Useful for
both facial recognition tasks and face swapping tasks.
:param shrink_factor: Shrink factor that shrinks the source landmark coordinates. This will include more border
information around the face. Useful when you want to include more background information when performing face swaps.
The lower the shrink factor the more of the face is included. Default value 1.0 will align the image to be ready
for the Arcface recognition model, but usually omits part of the chin. Value of 0.0 would transform all source points
to the middle of the image, probably rendering the alignment procedure useless.
If you process the image with a shrink factor of 0.85 and then want to extract the identity embedding with arcface,
you simply do a central crop of factor 0.85 to yield same cropped result as using shrink factor 1.0. This will
reduce the resolution, the recommendation is to processed images to output resolutions higher than 112 is using
Arcface. This will make sure no information is lost by resampling the image after central crop.
:return: Returns the transformed image.
"""
M, pose_index = estimate_norm(landmark, image_size, mode, shrink_factor=shrink_factor)
warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
return warped
def transform_landmark_points(M, points):
lmk_tran = np.insert(points, 2, values=np.ones(5), axis=1)
transformed_lmk = np.dot(M, lmk_tran.T)
transformed_lmk = transformed_lmk.T
return transformed_lmk
def multi_convolver(image, kernel, iterations):
if kernel == "Sharpen":
kernel = np.array([[0, -1, 0],
[-1, 5, -1],
[0, -1, 0]])
elif kernel == "Unsharp_mask":
kernel = np.array([[1, 4, 6, 4, 1],
[4, 16, 24, 16, 1],
[6, 24, -476, 24, 1],
[4, 16, 24, 16, 1],
[1, 4, 6, 4, 1]]) * (-1 / 256)
elif kernel == "Blur":
kernel = (1 / 16.0) * np.array([[1., 2., 1.],
[2., 4., 2.],
[1., 2., 1.]])
for i in range(iterations):
image = convolve2d(image, kernel, 'same', boundary='fill', fillvalue = 0)
return image
def convolve_rgb(image, kernel, iterations=1):
img_yuv = rgb2yuv(image)
img_yuv[:, :, 0] = multi_convolver(img_yuv[:, :, 0], kernel,
iterations)
final_image = yuv2rgb(img_yuv)
return final_image.astype('float32')
def generate_mask_from_landmarks(lms, im_size):
blend_mask_lm = np.zeros(shape=(im_size, im_size, 3), dtype='float32')
# EYES
blend_mask_lm = cv2.circle(blend_mask_lm,
(int(lms[0][0]), int(lms[0][1])), 12, (255, 255, 255), 30)
blend_mask_lm = cv2.circle(blend_mask_lm,
(int(lms[1][0]), int(lms[1][1])), 12, (255, 255, 255), 30)
blend_mask_lm = cv2.circle(blend_mask_lm,
(int((lms[0][0] + lms[1][0]) / 2), int((lms[0][1] + lms[1][1]) / 2)),
16, (255, 255, 255), 65)
# NOSE
blend_mask_lm = cv2.circle(blend_mask_lm,
(int(lms[2][0]), int(lms[2][1])), 5, (255, 255, 255), 5)
blend_mask_lm = cv2.circle(blend_mask_lm,
(int((lms[0][0] + lms[1][0]) / 2), int(lms[2][1])), 16, (255, 255, 255), 100)
# MOUTH
blend_mask_lm = cv2.circle(blend_mask_lm,
(int(lms[3][0]), int(lms[3][1])), 6, (255, 255, 255), 30)
blend_mask_lm = cv2.circle(blend_mask_lm,
(int(lms[4][0]), int(lms[4][1])), 6, (255, 255, 255), 30)
blend_mask_lm = cv2.circle(blend_mask_lm,
(int((lms[3][0] + lms[4][0]) / 2), int((lms[3][1] + lms[4][1]) / 2)),
16, (255, 255, 255), 40)
return blend_mask_lm
def display_distance_text(im, distance, lms, im_w, im_h, scale=2):
blended_insert = cv2.putText(im, str(distance)[:4],
(int(lms[4] * im_w * 0.5), int(lms[5] * im_h * 0.8)),
cv2.FONT_HERSHEY_SIMPLEX, scale * 0.5, (0.08, 0.16, 0.08), int(scale * 2))
blended_insert = cv2.putText(blended_insert, str(distance)[:4],
(int(lms[4] * im_w * 0.5), int(lms[5] * im_h * 0.8)),
cv2.FONT_HERSHEY_SIMPLEX, scale* 0.5, (0.3, 0.7, 0.32), int(scale * 1))
return blended_insert
def get_lm(annotation, im_w, im_h):
lm_align = np.array([[annotation[4] * im_w, annotation[5] * im_h],
[annotation[6] * im_w, annotation[7] * im_h],
[annotation[8] * im_w, annotation[9] * im_h],
[annotation[10] * im_w, annotation[11] * im_h],
[annotation[12] * im_w, annotation[13] * im_h]],
dtype=np.float32)
return lm_align