|
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) |
|
|
