mediapipe/mediapipe/python/solutions/face_mesh.py
MediaPipe Team 33d683c671 Project import generated by Copybara.
GitOrigin-RevId: 373e3ac1e5839befd95bf7d73ceff3c5f1171969
2021-10-06 14:27:49 -07:00

125 lines
5.7 KiB
Python

# Copyright 2020 The MediaPipe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""MediaPipe Face Mesh."""
from typing import NamedTuple
import numpy as np
# pylint: disable=unused-import
from mediapipe.calculators.core import constant_side_packet_calculator_pb2
from mediapipe.calculators.core import gate_calculator_pb2
from mediapipe.calculators.core import split_vector_calculator_pb2
from mediapipe.calculators.tensor import image_to_tensor_calculator_pb2
from mediapipe.calculators.tensor import inference_calculator_pb2
from mediapipe.calculators.tensor import tensors_to_classification_calculator_pb2
from mediapipe.calculators.tensor import tensors_to_detections_calculator_pb2
from mediapipe.calculators.tensor import tensors_to_landmarks_calculator_pb2
from mediapipe.calculators.tflite import ssd_anchors_calculator_pb2
from mediapipe.calculators.util import association_calculator_pb2
from mediapipe.calculators.util import detections_to_rects_calculator_pb2
from mediapipe.calculators.util import landmarks_refinement_calculator_pb2
from mediapipe.calculators.util import logic_calculator_pb2
from mediapipe.calculators.util import non_max_suppression_calculator_pb2
from mediapipe.calculators.util import rect_transformation_calculator_pb2
from mediapipe.calculators.util import thresholding_calculator_pb2
# pylint: enable=unused-import
from mediapipe.python.solution_base import SolutionBase
# pylint: disable=unused-import
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_CONTOURS
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_FACE_OVAL
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_IRISES
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LEFT_EYE
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LEFT_EYEBROW
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LEFT_IRIS
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_LIPS
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_EYE
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_EYEBROW
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_IRIS
from mediapipe.python.solutions.face_mesh_connections import FACEMESH_TESSELATION
# pylint: enable=unused-import
FACEMESH_NUM_LANDMARKS = 468
FACEMESH_NUM_LANDMARKS_WITH_IRISES = 478
_BINARYPB_FILE_PATH = 'mediapipe/modules/face_landmark/face_landmark_front_cpu.binarypb'
class FaceMesh(SolutionBase):
"""MediaPipe Face Mesh.
MediaPipe Face Mesh processes an RGB image and returns the face landmarks on
each detected face.
Please refer to https://solutions.mediapipe.dev/face_mesh#python-solution-api
for usage examples.
"""
def __init__(self,
static_image_mode=False,
max_num_faces=1,
refine_landmarks=False,
min_detection_confidence=0.5,
min_tracking_confidence=0.5):
"""Initializes a MediaPipe Face Mesh object.
Args:
static_image_mode: Whether to treat the input images as a batch of static
and possibly unrelated images, or a video stream. See details in
https://solutions.mediapipe.dev/face_mesh#static_image_mode.
max_num_faces: Maximum number of faces to detect. See details in
https://solutions.mediapipe.dev/face_mesh#max_num_faces.
refine_landmarks: Whether to further refine the landmark coordinates
around the eyes and lips, and output additional landmarks around the
irises. Default to False. See details in
https://solutions.mediapipe.dev/face_mesh#refine_landmarks.
min_detection_confidence: Minimum confidence value ([0.0, 1.0]) for face
detection to be considered successful. See details in
https://solutions.mediapipe.dev/face_mesh#min_detection_confidence.
min_tracking_confidence: Minimum confidence value ([0.0, 1.0]) for the
face landmarks to be considered tracked successfully. See details in
https://solutions.mediapipe.dev/face_mesh#min_tracking_confidence.
"""
super().__init__(
binary_graph_path=_BINARYPB_FILE_PATH,
side_inputs={
'num_faces': max_num_faces,
'with_attention': refine_landmarks,
'use_prev_landmarks': not static_image_mode,
},
calculator_params={
'facedetectionshortrangecpu__facedetectionshortrangecommon__TensorsToDetectionsCalculator.min_score_thresh':
min_detection_confidence,
'facelandmarkcpu__ThresholdingCalculator.threshold':
min_tracking_confidence,
},
outputs=['multi_face_landmarks'])
def process(self, image: np.ndarray) -> NamedTuple:
"""Processes an RGB image and returns the face landmarks on each detected face.
Args:
image: An RGB image represented as a numpy ndarray.
Raises:
RuntimeError: If the underlying graph throws any error.
ValueError: If the input image is not three channel RGB.
Returns:
A NamedTuple object with a "multi_face_landmarks" field that contains the
face landmarks on each detected face.
"""
return super().process(input_data={'image': image})