def landmark_preprocess(noisy_path, dataset_path):
noisy_train_path = os.path.join(noisy_path, "train")
noisy_test_path = os.path.join(noisy_path, "test")
noisy_val_path = os.path.join(noisy_path, "validation")
train_path = os.path.join(dataset_path, "train")
test_path = os.path.join(dataset_path, "test")
val_path = os.path.join(dataset_path, "validation")
checkpath(dataset_path)
checkpath(train_path)
checkpath(test_path)
checkpath(val_path)
fa = FaceAlignment(LandmarksType._2D, device="cuda:1")
train_files = file_list(noisy_train_path, ".png")
for i in train_files:
landmarks = fa.get_landmarks_from_image(i)[0]
img = plot_landmarks((300, 300, 3), landmarks)
img.save(os.path.join(train_path, "lm" + i.split("y")[-1]))
test_files = file_list(noisy_test_path, ".png")
for i in test_files:
landmarks = fa.get_landmarks_from_image(i)[0]
img = plot_landmarks((300, 300, 3), landmarks)
img.save(os.path.join(test_path, "lm" + i.split("y")[-1]))
val_files = file_list(noisy_val_path, ".png")
print(noisy_val_path)
for i in val_files:
landmarks = fa.get_landmarks_from_image(i)[0]
img = plot_landmarks((300, 300, 3), landmarks)
img.save(os.path.join(val_path, "lm" + i.split("y")[-1]))
def end2end_mask_encoding(
image: np.ndarray,
face_aligment_class: face_alignment.FaceAlignment) -> dict:
landmarks = face_aligment_class.get_landmarks_from_image(image)
landmarks = np.floor(landmarks[0]).astype(np.int32)
target_points, s1, s2 = extract_target_points_and_characteristic(landmarks)
mask_rgba_crop, target_points = extract_polygon(image, target_points)
mask_rgba_crop, target_points = rotate_image_and_points(
mask_rgba_crop, s1, target_points)
res = {
's1': s1,
's2': s2,
'points': target_points.tolist(),
'base64_img': image_to_string(mask_rgba_crop)
}
return res
def eval(input_path, output_path, checkpoint_path, model, gpu):
input = Image.open(input_path)
input = input.convert("RGB")
w, h = input.size
w_, h_ = 128 * (w // 128), 128 * (h // 128)
fa = FaceAlignment(LandmarksType._2D, device="cuda:" + str(gpu))
landmarks = fa.get_landmarks_from_image(input_path)[0]
landmark_img = plot_landmarks(np.array(input), landmarks)
transform_forward = transforms.Compose([
transforms.Resize((w_, h_)),
transforms.CenterCrop((w_, h_)),
transforms.ToTensor()
])
transform_backward = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((w, h)),
transforms.CenterCrop((w, h)),
])
input = transform_forward(input)
landmark_img = transform_forward(landmark_img)
if model == "Pix2Pix":
NFNet = Pix2Pix()
else:
NFNet = ResResNet()
checkpoint = torch.load(checkpoint_path)
NFNet.load_state_dict(checkpoint['my_classifier'])
NFNet.to(gpu)
x = torch.cat((input, landmark_img), 0)
x = x.unsqueeze(0)
x = x.to(gpu)
output = NFNet(x)
output = output.to("cpu")
output = transform_backward(output[0])
output.save(output_path)
def end2end_mask_generation(image: np.ndarray,
masks_database: dict,
face_aligment_class: face_alignment.FaceAlignment,
input_face_landmarks: np.ndarray = None):
if input_face_landmarks is None:
face_landmarks = face_aligment_class.get_landmarks_from_image(image)
if len(face_landmarks) == 0:
raise RuntimeError('Can\'t find facial landmarks')
face_landmarks = np.floor(face_landmarks[0]).astype(np.int32)
else:
face_landmarks = input_face_landmarks
_, _, s2 = extract_target_points_and_characteristic(face_landmarks)
sampling_mask_data = extract_mask_from_base(masks_database, s2)
face_with_mask = apply_mask_to_image_with_face(image, sampling_mask_data,
face_landmarks)
return face_with_mask