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