| import os |
| import PIL |
| import cv2 |
| import pickle |
| import argparse |
| import numpy as np |
| import face_alignment |
| import matplotlib.pyplot as plt |
| import matplotlib.patches as patches |
| from matplotlib.path import Path |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description="Plot facial landmarks from an image.") |
| parser.add_argument( |
| "--image_path", |
| type=str, |
| default=None, |
| help="Path to the image file." |
| ) |
| parser.add_argument("--size", type=int, default=512) |
| parser.add_argument("--crop", action="store_true", help="Crop around the face image.") |
| parser.add_argument( |
| "--output_dir", |
| type=str, |
| default="output/landmarks/", |
| help="Folder to save landmark images." |
| ) |
| args = parser.parse_args() |
|
|
| return args |
|
|
| def get_patch(landmarks, color='lime', closed=False): |
| contour = landmarks |
| ops = [Path.MOVETO] + [Path.LINETO]*(len(contour)-1) |
| facecolor = (0, 0, 0, 0) |
| if closed: |
| contour.append(contour[0]) |
| ops.append(Path.CLOSEPOLY) |
| facecolor = color |
| path = Path(contour, ops) |
| return patches.PathPatch(path, facecolor=facecolor, edgecolor=color, lw=4) |
|
|
| def bbox_from_landmarks(landmarks): |
| landmarks_x, landmarks_y = zip(*landmarks) |
|
|
| x_min, x_max = min(landmarks_x), max(landmarks_x) |
| y_min, y_max = min(landmarks_y), max(landmarks_y) |
| width = x_max - x_min |
| height = y_max - y_min |
|
|
| |
| x_min -= 25 |
| y_min -= 25 |
| width += 50 |
| height += 50 |
| bbox = (x_min, y_min, width, height) |
| return bbox |
|
|
| def plot_landmarks(landmarks, crop=False, size=512): |
| if crop: |
| (x_min, y_min, width, height) = bbox_from_landmarks(landmarks) |
| |
| landmarks_np = np.array(landmarks) |
| landmarks_np[:, 0] = (landmarks_np[:, 0] - x_min) * size / width |
| landmarks_np[:, 1] = (landmarks_np[:, 1] - y_min) * size / height |
| landmarks = landmarks_np.tolist() |
| |
| dpi = 72 |
| fig, ax = plt.subplots(1, figsize=[size/dpi, size/dpi], tight_layout={'pad':0}) |
| fig.set_dpi(dpi) |
|
|
| black = np.zeros((size, size, 3)) |
| ax.imshow(black) |
|
|
| face_patch = get_patch(landmarks[0:17]) |
| l_eyebrow = get_patch(landmarks[17:22], color='yellow') |
| r_eyebrow = get_patch(landmarks[22:27], color='yellow') |
| nose_v = get_patch(landmarks[27:31], color='orange') |
| nose_h = get_patch(landmarks[31:36], color='orange') |
| l_eye = get_patch(landmarks[36:42], color='magenta', closed=True) |
| r_eye = get_patch(landmarks[42:48], color='magenta', closed=True) |
| outer_lips = get_patch(landmarks[48:60], color='cyan', closed=True) |
| inner_lips = get_patch(landmarks[60:68], color='blue', closed=True) |
|
|
| ax.add_patch(face_patch) |
| ax.add_patch(l_eyebrow) |
| ax.add_patch(r_eyebrow) |
| ax.add_patch(nose_v) |
| ax.add_patch(nose_h) |
| ax.add_patch(l_eye) |
| ax.add_patch(r_eye) |
| ax.add_patch(outer_lips) |
| ax.add_patch(inner_lips) |
|
|
| plt.axis('off') |
|
|
| fig.canvas.draw() |
| buffer, (width, height) = fig.canvas.print_to_buffer() |
| assert width == height |
| assert width == size |
|
|
| buffer = np.frombuffer(buffer, np.uint8).reshape((height, width, 4)) |
| buffer = buffer[:, :, 0:3] |
| plt.close(fig) |
| return PIL.Image.fromarray(buffer) |
|
|
| def get_landmarks(image): |
| fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, flip_input=False, face_detector='sfd') |
| faces = fa.get_landmarks_from_image(image) |
| if faces is None or len(faces) == 0: |
| return None |
| landmarks = faces[0] |
| return landmarks |
|
|
| def save_landmarks(args): |
| os.makedirs(args.output_dir, exist_ok=True) |
| |
| image_name = os.path.basename(args.image_path) |
| image = cv2.imread(args.image_path) |
| image = cv2.resize(image, (args.size, args.size)) |
| landmarks = get_landmarks(image) |
| if landmarks is None: |
| print(f'No faces found in {image_name}') |
| return |
| |
| filename = f'{args.output_dir}/{image_name}' |
| if args.crop: |
| landmarks_cropped_image = plot_landmarks(landmarks.tolist(), crop=True, size=args.size) |
| landmarks_cropped_image.save(filename) |
| else: |
| landmarks_image = plot_landmarks(landmarks.tolist(), size=args.size) |
| landmarks_image.save(filename) |
| print(f'Landmark saved in {filename}') |
|
|
| if __name__ == '__main__': |
| args = parse_args() |
| save_landmarks(args) |
|
|
