diff --git a/src/falign/align.py b/src/falign/align.py new file mode 100644 index 0000000..2da1e3e --- /dev/null +++ b/src/falign/align.py @@ -0,0 +1,80 @@ +import numpy +import cv2 + + +SLICE_EYE_LEFT = slice(36, 42) +SLICE_EYE_RIGHT = slice(42, 48) +IDX_MOUTH_LEFT = 48 +IDX_MOUTH_RIGHT = 54 + + +def normalize_vector(v: numpy.ndarray): + norm = numpy.linalg.norm(v) + 1e-8 + return v / norm + + +def rotate_90(v: numpy.ndarray): + return numpy.array([-v[1], v[0]]) + + +def calculate_alignment_quadrangle(landmarks: numpy.ndarray, margin_ratio: float) -> numpy.ndarray: + """ + Calculate the quadrangle for alignment based on facial landmarks. + + Args: + landmarks: The landmarks array of shape (68, 2). + Returns: + The quadrangle array of shape (4, 2). + """ + + # compute auxiliary vectors + eye_left = landmarks[SLICE_EYE_LEFT].mean(axis=0) + eye_right = landmarks[SLICE_EYE_RIGHT].mean(axis=0) + mouth_left = landmarks[IDX_MOUTH_LEFT] + mouth_right = landmarks[IDX_MOUTH_RIGHT] + + # compute average positions + eye_avg = (eye_left + eye_right) * 0.5 + mouth_avg = (mouth_left + mouth_right) * 0.5 + + # compute direction vectors + eye_to_eye = eye_right - eye_left + eye_to_mouth = mouth_avg - eye_avg + + # set the quadrangle + x_axis = normalize_vector(eye_to_eye - rotate_90(eye_to_mouth)) + x_axis *= max(numpy.linalg.norm(eye_to_eye) * 2.0, numpy.linalg.norm(eye_to_mouth) * 1.8) + y_axis = rotate_90(x_axis) + center = eye_avg + eye_to_mouth * 0.1 + + # add margin by expanding the quadrangle + margin_factor = 1.0 + margin_ratio + x_axis_with_margin = x_axis * margin_factor + y_axis_with_margin = y_axis * margin_factor + + return numpy.stack([ + center - x_axis_with_margin - y_axis_with_margin, + center - x_axis_with_margin + y_axis_with_margin, + center + x_axis_with_margin + y_axis_with_margin, + center + x_axis_with_margin - y_axis_with_margin, + ]) + + +def align( + image: numpy.ndarray, landmarks: numpy.ndarray, + height: int, width: int, +) -> numpy.ndarray: + landmarks = landmarks.astype(numpy.float32) # (68, 2) + + quadrangle_target = numpy.array([ + [0, 0], + [0, height-1], + [width-1, height-1], + [width-1, 0] + ], dtype=numpy.float32) + + quadrangle = calculate_alignment_quadrangle(landmarks, margin_ratio=0.1).astype(numpy.float32) + transform = cv2.getPerspectiveTransform(quadrangle, quadrangle_target) + aligned = cv2.warpPerspective(image, transform, (width, height), borderMode=cv2.BORDER_REFLECT) + + return aligned diff --git a/src/falign/landmarks.py b/src/falign/landmarks.py index 84bdd1b..79da15f 100644 --- a/src/falign/landmarks.py +++ b/src/falign/landmarks.py @@ -1,21 +1,23 @@ -from numpy.typing import NDArray - from pathlib import Path +import numpy from skimage.io import imread import torch import face_alignment -def _get_face_landmarks(image: NDArray, face_aligner: face_alignment.FaceAlignment): - landmarks = face_aligner.get_landmarks(image) +def get_face_landmarks(image: numpy.ndarray, face_aligner: face_alignment.FaceAlignment): + landmarks = face_aligner.get_landmarks_from_image(image) assert landmarks is not None, "No face detected" return landmarks[0] -def get_landmarks(path_image: Path, dtype: torch.dtype = torch.float32, device: torch.device = torch.device("cuda")): +def read_image_and_get_landmarks( + path_image: Path, + dtype: torch.dtype = torch.float32, device: torch.device = torch.device("cuda"), +): face_aligner = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, dtype=dtype, device=device.type) image = imread(path_image) - return _get_face_landmarks(image, face_aligner) + return image, get_face_landmarks(image, face_aligner) diff --git a/src/falign/plot.py b/src/falign/plot.py index a47cd14..7e62adc 100644 --- a/src/falign/plot.py +++ b/src/falign/plot.py @@ -1,17 +1,17 @@ from pathlib import Path -from numpy.typing import NDArray - +import numpy from skimage.io import imsave def imsave_with_landmarks( path: Path, - image: NDArray, landmarks: NDArray, + image: numpy.ndarray, landmarks: numpy.ndarray, size: int = 1, ) -> None: """ Save image with landmarks. + Args: path: The path to save the image. image: The image array. diff --git a/test/test_falign/test_align.py b/test/test_falign/test_align.py new file mode 100644 index 0000000..46d289c --- /dev/null +++ b/test/test_falign/test_align.py @@ -0,0 +1,17 @@ +from pathlib import Path + +from skimage.io import imread, imsave + +from falign.landmarks import read_image_and_get_landmarks +from falign.align import align + + +def test_align(): + dir_gallery = Path(__file__).parent.parent.parent / "gallery" + path_image = dir_gallery / "original.jpg" + + image, landmarks = read_image_and_get_landmarks(path_image) + aligned = align(image, landmarks, height=256, width=256) + assert aligned.shape == (256, 256, 3) + + imsave(dir_gallery / "test_align.jpg", aligned) diff --git a/test/test_falign/test_landmarks.py b/test/test_falign/test_landmarks.py index 38eaaee..beaab0d 100644 --- a/test/test_falign/test_landmarks.py +++ b/test/test_falign/test_landmarks.py @@ -2,19 +2,19 @@ from pathlib import Path from skimage.io import imread -from falign.landmarks import get_landmarks +from falign.landmarks import read_image_and_get_landmarks from falign.plot import imsave_with_landmarks def test_get_landmarks(): dir_gallery = Path(__file__).parent.parent.parent / "gallery" path_image = dir_gallery / "original.jpg" - landmarks = get_landmarks(path_image) + image, landmarks = read_image_and_get_landmarks(path_image) assert landmarks.shape == (68, 2) imsave_with_landmarks( dir_gallery / "test_get_landmarks.jpg", - imread(path_image), landmarks, + image, landmarks, size=5 )