Spaces:
Build error
Build error
File size: 3,499 Bytes
752c636 f5562ad |
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 |
import cv2
import itertools
import numpy as np
import mediapipe as mp
import matplotlib.pyplot as plt
import PIL
import os
import shutil
class Aligner(mp.solutions.face_mesh.FaceMesh):
"""Inherits from MediaPipe Face Mesh."""
def __init__(
self,
static_image_mode: bool = True,
max_num_faces: int = 1,
refine_landmarks: bool = False,
min_detection_confidence: float = 0.5,
min_tracking_confidence: float = 0.5
):
"""Initializes a Image Aligner object.
Unlike MediaPipe Face Mesh we set `static_image_mode` to `True` as we only
require image manipulation.
Args:
static_image_mode: Whether to treat the input images as a batch of static
and possibly unrelated images, or a video stream.
max_num_faces: Maximum number of faces to detect.
refine_landmarks: Whether to further refine the landmark coordinates
around the eyes and lips, and output additional landmarks around the
irises. Default to False.
min_detection_confidence: Minimum confidence value ([0.0, 1.0]) for face
detection to be considered successful.
min_tracking_confidence: Minimum confidence value ([0.0, 1.0]) for the
face landmarks to be considered tracked successfully.
"""
super().__init__(
static_image_mode=static_image_mode,
max_num_faces=max_num_faces,
refine_landmarks=refine_landmarks,
min_detection_confidence=min_detection_confidence,
min_tracking_confidence=min_tracking_confidence
)
self._left_eye_idx = list(
set(itertools.chain(*mp.solutions.face_mesh.FACEMESH_LEFT_EYE))
)[7]
self._right_eye_idx = list(
set(itertools.chain(*mp.solutions.face_mesh.FACEMESH_RIGHT_EYE))
)[4]
def _aligner(self, /, img: np.ndarray) -> np.ndarray:
"""Private helper function to align the given image parallel to the x-axis.
This function creates a line between the left and right eye points and tries
to align that line parallel to the x-axis, thus aligning the complete image.
Args:
img: Image to align parallel to the x-axis.
"""
fm = self.process(img)
if fm is None:
return None
points = []
h, w, _ = img.shape
face_landmarks = fm.multi_face_landmarks[0]
le_x_coord = int(
np.clip(face_landmarks.landmark[self._left_eye_idx].x * w, 0, w)
)
le_y_coord = int(
np.clip(face_landmarks.landmark[self._left_eye_idx].y * h, 0, h)
)
p0 = np.array((le_x_coord, le_y_coord), dtype=np.float64)
re_x_coord = int(
np.clip(face_landmarks.landmark[self._right_eye_idx].x * w, 0, w)
)
re_y_coord = int(
np.clip(face_landmarks.landmark[self._right_eye_idx].y * h, 0, h)
)
p1 = np.array((re_x_coord, re_y_coord), dtype=np.float64)
h = abs(p0[1] - p1[1])
w = abs(p0[0] - p1[0])
# Get the angle between the x-axis and the line joining the eye points.
theta = np.arctan(h / w)
angle = (theta * 180) / np.pi
if p0[0] < p1[0]:
direction = 1 if p0[1] < p1[1] else -1
else:
direction = 1 if p1[1] < p0[1] else -1
angle *= direction
img = PIL.Image.fromarray(img)
return np.array(img.rotate(angle))
def align(self, /, imgs: tuple[np.ndarray]) -> list[np.ndarray]:
"""Aligns the given set of images parallel to the x-axis on the image plane.
Args:
imgs: Images to align parallel to the x-axis on the image place.
"""
return [self._aligner(img) for img in imgs]
|