File size: 4,515 Bytes
332190f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
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) # Transparent fill color, if open
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
# Give it a little room; I think it works anyway
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)
# print(x_min, y_min, width, height)
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()
# Precisely control output image size
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)
|