import logging import cv2 import numpy as np logger = logging.getLogger(__name__) # pylint: disable=invalid-name class Mask(): """ Parent class for masks the output mask will be .mask channels: 1, 3 or 4: 1 - Returns a single channel mask 3 - Returns a 3 channel mask 4 - Returns the original image with the mask in the alpha channel """ def __init__(self, landmarks, face, channels=4): logger.info("Initializing %s: (face_shape: %s, channels: %s, landmarks: %s)", self.__class__.__name__, face.shape, channels, landmarks) self.landmarks = landmarks self.face = face self.channels = channels mask = self.build_mask() self.mask = self.merge_mask(mask) logger.info("Initialized %s", self.__class__.__name__) def build_mask(self): """ Override to build the mask """ raise NotImplementedError def merge_mask(self, mask): """ Return the mask in requested shape """ logger.info("mask_shape: %s", mask.shape) assert self.channels in (1, 3, 4), "Channels should be 1, 3 or 4" assert mask.shape[2] == 1 and mask.ndim == 3, "Input mask be 3 dimensions with 1 channel" if self.channels == 3: retval = np.tile(mask, 3) elif self.channels == 4: retval = np.concatenate((self.face, mask), -1) else: retval = mask logger.info("Final mask shape: %s", retval.shape) return retval class dfl_full(Mask): # pylint: disable=invalid-name """ DFL facial mask """ def build_mask(self): mask = np.zeros(self.face.shape[0:2] + (1, ), dtype=np.float32) nose_ridge = (self.landmarks[27:31], self.landmarks[33:34]) jaw = (self.landmarks[0:17], self.landmarks[48:68], self.landmarks[0:1], self.landmarks[8:9], self.landmarks[16:17]) eyes = (self.landmarks[17:27], self.landmarks[0:1], self.landmarks[27:28], self.landmarks[16:17], self.landmarks[33:34]) parts = [jaw, nose_ridge, eyes] for item in parts: merged = np.concatenate(item) cv2.fillConvexPoly(mask, cv2.convexHull(merged), 255.) # pylint: disable=no-member return mask class components(Mask): # pylint: disable=invalid-name """ Component model mask """ def build_mask(self): mask = np.zeros(self.face.shape[0:2] + (1, ), dtype=np.float32) r_jaw = (self.landmarks[0:9], self.landmarks[17:18]) l_jaw = (self.landmarks[8:17], self.landmarks[26:27]) r_cheek = (self.landmarks[17:20], self.landmarks[8:9]) l_cheek = (self.landmarks[24:27], self.landmarks[8:9]) nose_ridge = (self.landmarks[19:25], self.landmarks[8:9],) r_eye = (self.landmarks[17:22], self.landmarks[27:28], self.landmarks[31:36], self.landmarks[8:9]) l_eye = (self.landmarks[22:27], self.landmarks[27:28], self.landmarks[31:36], self.landmarks[8:9]) nose = (self.landmarks[27:31], self.landmarks[31:36]) parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose] for item in parts: merged = np.concatenate(item) cv2.fillConvexPoly(mask, cv2.convexHull(merged), 255.) # pylint: disable=no-member return mask class extended(Mask): # pylint: disable=invalid-name """ Extended mask Based on components mask. Attempts to extend the eyebrow points up the forehead """ def build_mask(self): mask = np.zeros(self.face.shape[0:2] + (1, ), dtype=np.float32) landmarks = self.landmarks.copy() # mid points between the side of face and eye point ml_pnt = (landmarks[36] + landmarks[0]) // 2 mr_pnt = (landmarks[16] + landmarks[45]) // 2 # mid points between the mid points and eye ql_pnt = (landmarks[36] + ml_pnt) // 2 qr_pnt = (landmarks[45] + mr_pnt) // 2 # Top of the eye arrays bot_l = np.array((ql_pnt, landmarks[36], landmarks[37], landmarks[38], landmarks[39])) bot_r = np.array((landmarks[42], landmarks[43], landmarks[44], landmarks[45], qr_pnt)) # Eyebrow arrays top_l = landmarks[17:22] top_r = landmarks[22:27] # Adjust eyebrow arrays landmarks[17:22] = top_l + ((top_l - bot_l) // 2) landmarks[22:27] = top_r + ((top_r - bot_r) // 2) r_jaw = (landmarks[0:9], landmarks[17:18]) l_jaw = (landmarks[8:17], landmarks[26:27]) r_cheek = (landmarks[17:20], landmarks[8:9]) l_cheek = (landmarks[24:27], landmarks[8:9]) nose_ridge = (landmarks[19:25], landmarks[8:9],) r_eye = (landmarks[17:22], landmarks[27:28], landmarks[31:36], landmarks[8:9]) l_eye = (landmarks[22:27], landmarks[27:28], landmarks[31:36], landmarks[8:9]) nose = (landmarks[27:31], landmarks[31:36]) parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose] for item in parts: merged = np.concatenate(item) cv2.fillConvexPoly(mask, cv2.convexHull(merged), 255.) # pylint: disable=no-member return mask class facehull(Mask): # pylint: disable=invalid-name """ Basic face hull mask """ def build_mask(self): mask = np.zeros(self.face.shape[0:2] + (1, ), dtype=np.float32) hull = cv2.convexHull( # pylint: disable=no-member np.array(self.landmarks).reshape((-1, 2))) cv2.fillConvexPoly(mask, hull, 255.0, lineType=cv2.LINE_AA) # pylint: disable=no-member return mask