Implement align

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

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)

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.

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)

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
     )