51 lines
2.1 KiB
Python
51 lines
2.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
from typing import Any
|
|
import mediapipe as mp
|
|
from basicsr.utils.download_util import load_file_from_url
|
|
|
|
class FaceMask:
|
|
def __init__(self) -> None:
|
|
BaseOptions = mp.tasks.BaseOptions
|
|
FaceLandmarker = mp.tasks.vision.FaceLandmarker
|
|
FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
|
|
VisionRunningMode = mp.tasks.vision.RunningMode
|
|
|
|
face_landmarks_detector_path = load_file_from_url(url="https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/latest/face_landmarker.task",
|
|
model_dir="weights",
|
|
file_name="face_landmarker.task")
|
|
options = FaceLandmarkerOptions(
|
|
base_options=BaseOptions(model_asset_path=face_landmarks_detector_path),
|
|
running_mode=VisionRunningMode.IMAGE)
|
|
self.face_landmarks_detector = FaceLandmarker.create_from_options(options)
|
|
|
|
def __call__(self,image,*args: Any, **kwds: Any) -> Any:
|
|
"""
|
|
Calculate face mask from image. This is done by
|
|
|
|
Args:
|
|
image: numpy array of an image
|
|
Returns:
|
|
A uint8 numpy array with the same height and width of the input image, containing a binary mask of the face in the image
|
|
"""
|
|
# initialize mask
|
|
mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
|
|
|
|
# detect face landmarks
|
|
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
|
|
detection = self.face_landmarks_detector.detect(mp_image)
|
|
|
|
if len(detection.face_landmarks) == 0:
|
|
# no face detected - set mask to all of the image
|
|
mask[:] = 1
|
|
return mask
|
|
|
|
# extract landmarks coordinates
|
|
face_coords = np.array([[lm.x * image.shape[1], lm.y * image.shape[0]] for lm in detection.face_landmarks[0]])
|
|
|
|
# calculate convex hull from face coordinates
|
|
convex_hull = cv2.convexHull(face_coords.astype(np.float32))
|
|
|
|
# apply convex hull to mask
|
|
return cv2.fillPoly(mask, pts=[convex_hull.squeeze().astype(np.int32)], color=1)
|