def facenet_distance(img1, img2, percentage=False, embedding=False):
"""
Outputs distance between 0 and 2 with 0 is identical and 2 is most different
:param embedding: bool - return embedding
:param percentage: bool - return percentage
:param img1: img1 tensor
:param img2: img2 tensor
:return: dist between 0 (identical) and 2 (most dissimilar), + embedding/percentange if True
"""
img1 = mtcnn(Image.fromarray(tensor2im(img1)))
img2 = mtcnn(Image.fromarray(tensor2im(img2)))
if img1 is None or img2 is None:
dist = 2
embedding_list = [torch.zeros(512), torch.zeros(512)]
else:
img_stack = torch.stack([img1, img2]).to(device)
embedding_list = resnet(img_stack) # .detach().cpu()
embedding_dist = embedding_list[0] - embedding_list[1]
dist = embedding_dist.norm().item()
if percentage:
return dist, ((2 - dist) * 50)
elif embedding:
return dist, embedding_list
else:
return dist
def run_segmentation_model(model, opt, root_folder, warp=False):
print("\n---- Running Face2Mask Model ----\n")
dataset = create_dataset(opt)
for i, data in tqdm(enumerate(dataset)):
if i >= opt.num_test:
break
# run model and get processed image
model.set_input(data)
model.test()
# get mask
visuals_mask = model.get_current_visuals()
mask = tensor2im(visuals_mask['fake_B'])
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2RGB)
mask = clean_mask(mask)
mask = Image.fromarray(mask)
mask = make_color_transparent(mask, (0, 0, 0), tolerance=50)
# open original
img_path = model.get_image_paths()[0]
img_name = img_path.split("/")[-1][:-4]
image = Image.open(img_path)
# add background to mask
overlayed_mask = overlay_two_images(image, mask)
# save file
overlayed_mask.save(f'{root_folder}/{img_name}.png')
print("SAVED GENERATED MASKS\n")
def run_blending_model(model,
opt,
root_folder,
result_folder,
blend_back=False):
print("\n---- Running Face2Face Bleding Model ---- \n")
if blend_back:
with open(f'{root_folder}/{generated_folder}/coordinates.txt',
'r') as f:
reinsert_rect = json.load(f)
Path(f"{root_folder}/{blended_reinserted_folder}/{result_folder}"
).mkdir(parents=True, exist_ok=True)
dataset = create_dataset(opt)
for i, data in tqdm(enumerate(dataset)):
if i >= opt.num_test:
break
# test model
model.set_input(data)
model.test()
# get visuals
visuals_face = model.get_current_visuals()
img_path = model.get_image_paths()[0]
file_name = img_path.split('/')[-1]
# Tensor to image
generated_image = tensor2im(visuals_face['fake_B'])
generated_image = cv2.cvtColor(generated_image, cv2.COLOR_RGB2BGR)
# get key and file name
key_name = file_name[:-4]
raw_file_name = "_".join(file_name.split("_")[1:])
# get original image
out = _get_maybe_modified_file(raw_file_name, root_folder)
# save generated image
cv2.imwrite(
f"{root_folder}/{blended_folder}/{result_folder}/{file_name}",
generated_image)
if blend_back:
# get face bounding box
x, y, x_end, y_end = reinsert_rect[key_name]["reinsert_range"]
shape = out.shape
# reinsert blended image
generated_image = cv2.resize(
generated_image, (x_end - x + min(shape[1] - x_end, 0),
y_end - y + min(shape[0] - y_end, 0)))
out[y:y_end, x:x_end, :] = generated_image
# save
cv2.imwrite(
f"{root_folder}/{blended_reinserted_folder}/{result_folder}/{file_name}",
out)
print("SAVED BLENDED FACE\n")
def run_generation_model(model,
opt,
root_folder,
smoothEdge=20,
min_width=30,
margin=0):
print("\n---- Running Mask2Face Model ---- \n")
with open(f'{root_folder}/{test_folder}/coordinates.txt', 'r') as f:
face_extraction_data = json.load(f)
# logging
num_too_small = 0
num_replaced = 0
insert_data = {}
dataset = create_dataset(opt)
for i, data in tqdm(enumerate(dataset)):
if i >= opt.num_test:
break
# test model
model.set_input(data)
model.test()
# get visuals
visuals_face = model.get_current_visuals()
img_path = model.get_image_paths()[0]
key_name = img_path.split('/')[-1][:-4]
idx = key_name.index("_")
raw_file_name = key_name[idx + 1:] + ".png"
# Tensor to image
generated_image = tensor2im(visuals_face['fake_B'])
generated_image = cv2.cvtColor(generated_image, cv2.COLOR_RGB2BGR)
# get face bounding box
try:
x, y, w, h = face_extraction_data[key_name]["rect_cv"]
except KeyError:
# this handles files that were not deleted and are not longer in the test set
print(
f"{key_name} key not found in dict face_extraction_data! Skipping."
)
continue
# only continue if width of face that got reconstructed is larger that 50 pixels
if w > min_width:
out = _get_maybe_modified_file(raw_file_name, root_folder)
# save generated image
cv2.imwrite(f"{root_folder}/{generated_folder}/{key_name}.png",
generated_image)
# load keypoint data
keypoints = face_extraction_data[key_name]["keypoints"]
alignment_params = face_extraction_data[key_name][
"alignment_params"]
out, capture, out_margin, capture_margin, reinsert_range = reinsert_aligned_into_image(
generated_image,
out,
alignment_params,
keypoints,
smoothEdge=smoothEdge,
margin=margin,
clean_merge=True)
insert_data[key_name] = {'reinsert_range': reinsert_range}
cv2.imwrite(
f"{root_folder}/{generated_reinserted_folder}/{raw_file_name}",
out)
cv2.imwrite(f"{root_folder}/{to_blend_folder}/{key_name}.png",
out_margin)
num_replaced += 1
else:
num_too_small += 1
with open(f"{root_folder}/{generated_folder}/coordinates.txt",
'w') as outfile:
json.dump(insert_data, outfile)
print(
f"In total {num_replaced} faces replaced. "
f"{num_too_small} incedents of faces with width < {min_width} were NOT replaced. \n"
)
print("SAVED GENERATED FACE\n")
def align_fake(self, margin=40, alignUnaligned=True):
# get params
desiredLeftEye = [
float(self.alignment_params["desiredLeftEye"][0]),
float(self.alignment_params["desiredLeftEye"][1])
]
rotation_point = self.alignment_params["eyesCenter"]
angle = -self.alignment_params["angle"]
h, w = self.fake_B.shape[2:]
# get original positions
m1 = round(w * 0.5)
m2 = round(desiredLeftEye[1] * w)
# define the scale factor
scale = 1 / self.alignment_params["scale"]
width = int(self.alignment_params["shape"][0])
long_edge_size = width / abs(np.cos(np.deg2rad(angle)))
w_original = int(scale * long_edge_size)
h_original = int(scale * long_edge_size)
# get offset
tX = w_original * 0.5
tY = h_original * desiredLeftEye[1]
# get rotation center
center = torch.ones(1, 2)
center[..., 0] = m1
center[..., 1] = m2
# compute the transformation matrix
M: torch.tensor = kornia.get_rotation_matrix2d(center, angle,
scale).to(self.device)
M[0, 0, 2] += (tX - m1)
M[0, 1, 2] += (tY - m2)
# get insertion point
x_start = int(rotation_point[0] - (0.5 * w_original))
y_start = int(rotation_point[1] - (desiredLeftEye[1] * h_original))
# _, _, h_tensor, w_tensor = self.real_B_unaligned_full.shape
# # # # # # # # # # # # # # # # # # ## # # # # # # # ## # # # # ## # # # # # # # # # ## # #
# get safe margin
h_size_tensor, w_size_tensor = self.real_B_unaligned_full.shape[2:]
margin = max(
min(
y_start - max(0, y_start - margin),
x_start - max(0, x_start - margin),
min(y_start + h_original + margin, h_size_tensor) - y_start -
h_original,
min(x_start + w_original + margin, w_size_tensor) - x_start -
w_original,
), 0)
# get face + margin unaligned space
self.real_B_aligned_margin = self.real_B_unaligned_full[:, :, y_start -
margin:
y_start +
h_original +
margin,
x_start -
margin:
x_start +
w_original +
margin]
# invert matrix
M_inverse = kornia.invert_affine_transform(M)
# update output size to fit the 256 + scaled margin
old_size = self.real_B_aligned_margin.shape[2]
new_size = old_size + 2 * round(float(margin * scale))
_, _, h_tensor, w_tensor = self.real_B_aligned_margin.shape
self.real_B_aligned_margin = kornia.warp_affine(
self.real_B_aligned_margin, M_inverse, dsize=(new_size, new_size))
# padding_mode="reflection")
self.fake_B_aligned_margin = self.real_B_aligned_margin.clone(
).requires_grad_(True)
# update margin as we now scale the image!
# update start point
start = round(float(margin * scale * new_size / old_size))
print(start)
# point = torch.tensor([0, 0, 1], dtype=torch.float)
# M_ = M_inverse[0].clone().detach()
# M_ = torch.cat((M_, torch.tensor([[0, 0, 1]], dtype=torch.float)))
#
# M_n = M[0].clone().detach()
# M_n = torch.cat((M_n, torch.tensor([[0, 0, 1]], dtype=torch.float)))
#
# start_tensor = torch.matmul(torch.matmul(point, M_) + margin, M_n)
# print(start_tensor)
# start_y, start_x = round(float(start_tensor[0])), round(float(start_tensor[1]))
# reinsert into tensor
self.fake_B_aligned_margin[0, :, start:start + 256,
start:start + 256] = self.real_B
Image.fromarray(tensor2im(self.real_B_aligned_margin)).save(
"/home/mo/datasets/ff_aligned_unaligned/real.png")
Image.fromarray(tensor2im(self.fake_B_aligned_margin)).save(
"/home/mo/datasets/ff_aligned_unaligned/fake.png")
exit()
# # # # # # # # # # # # # # # # # # ## # # # # # # # ## # # # # ## # # # # # # # # # ## # #
if not alignUnaligned:
# Now apply the transformation to original image
# clone fake
fake_B_clone = self.fake_B.clone().requires_grad_(True)
# apply warp
fake_B_warped: torch.tensor = kornia.warp_affine(
fake_B_clone, M, dsize=(h_original, w_original))
# make sure warping does not exceed real_B_unaligned_full dimensions
if y_start < 0:
fake_B_warped = fake_B_warped[:, :, abs(y_start):h_original, :]
h_original += y_start
y_start = 0
if x_start < 0:
fake_B_warped = fake_B_warped[:, :, :, abs(x_start):w_original]
w_original += x_start
x_start = 0
if y_start + h_original > h_tensor:
h_original -= (y_start + h_original - h_tensor)
fake_B_warped = fake_B_warped[:, :, 0:h_original, :]
if x_start + w_original > w_tensor:
w_original -= (x_start + w_original - w_tensor)
fake_B_warped = fake_B_warped[:, :, :, 0:w_original]
# create mask that is true where fake_B_warped is 0
# This is the background that is not filled with image after the transformation
mask = ((fake_B_warped[0][0] == 0) & (fake_B_warped[0][1] == 0) &
(fake_B_warped[0][2] == 0))
# fill fake_B_filled where mask = False with self.real_B_unaligned_full
fake_B_filled = torch.where(
mask,
self.real_B_unaligned_full[:, :, y_start:y_start + h_original,
x_start:x_start + w_original],
fake_B_warped)
# reinsert into tensor
self.fake_B_unaligned = self.real_B_unaligned_full.clone(
).requires_grad_(True)
mask = torch.zeros_like(self.fake_B_unaligned, dtype=torch.bool)
mask[0, :, y_start:y_start + h_original,
x_start:x_start + w_original] = True
self.fake_B_unaligned = self.fake_B_unaligned.masked_scatter(
mask, fake_B_filled)
# cutout tensor
h_size_tensor, w_size_tensor = self.real_B_unaligned_full.shape[2:]
margin = max(
min(
y_start - max(0, y_start - margin),
x_start - max(0, x_start - margin),
min(y_start + h_original + margin, h_size_tensor) -
y_start - h_original,
min(x_start + w_original + margin, w_size_tensor) -
x_start - w_original,
), 0)
self.fake_B_unaligned = self.fake_B_unaligned[:, :, y_start -
margin:y_start +
h_original + margin,
x_start -
margin:x_start +
w_original + margin]
self.real_B_unaligned = self.real_B_unaligned_full[:, :, y_start -
margin:y_start +
h_original +
margin,
x_start -
margin:x_start +
w_original +
margin]