mediapipe/mediapipe2/python/solutions/face_mesh.py
2021-06-10 23:01:19 +00:00

239 lines
6.4 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 FaceMesh."""
from typing import NamedTuple
import numpy as np
from mediapipe.calculators.core import constant_side_packet_calculator_pb2
# pylint: disable=unused-import
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 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
BINARYPB_FILE_PATH = 'mediapipe/modules/face_landmark/face_landmark_front_cpu.binarypb'
FACE_CONNECTIONS = frozenset([
# Lips.
(61, 146),
(146, 91),
(91, 181),
(181, 84),
(84, 17),
(17, 314),
(314, 405),
(405, 321),
(321, 375),
(375, 291),
(61, 185),
(185, 40),
(40, 39),
(39, 37),
(37, 0),
(0, 267),
(267, 269),
(269, 270),
(270, 409),
(409, 291),
(78, 95),
(95, 88),
(88, 178),
(178, 87),
(87, 14),
(14, 317),
(317, 402),
(402, 318),
(318, 324),
(324, 308),
(78, 191),
(191, 80),
(80, 81),
(81, 82),
(82, 13),
(13, 312),
(312, 311),
(311, 310),
(310, 415),
(415, 308),
# Left eye.
(263, 249),
(249, 390),
(390, 373),
(373, 374),
(374, 380),
(380, 381),
(381, 382),
(382, 362),
(263, 466),
(466, 388),
(388, 387),
(387, 386),
(386, 385),
(385, 384),
(384, 398),
(398, 362),
# Left eyebrow.
(276, 283),
(283, 282),
(282, 295),
(295, 285),
(300, 293),
(293, 334),
(334, 296),
(296, 336),
# Right eye.
(33, 7),
(7, 163),
(163, 144),
(144, 145),
(145, 153),
(153, 154),
(154, 155),
(155, 133),
(33, 246),
(246, 161),
(161, 160),
(160, 159),
(159, 158),
(158, 157),
(157, 173),
(173, 133),
# Right eyebrow.
(46, 53),
(53, 52),
(52, 65),
(65, 55),
(70, 63),
(63, 105),
(105, 66),
(66, 107),
# Face oval.
(10, 338),
(338, 297),
(297, 332),
(332, 284),
(284, 251),
(251, 389),
(389, 356),
(356, 454),
(454, 323),
(323, 361),
(361, 288),
(288, 397),
(397, 365),
(365, 379),
(379, 378),
(378, 400),
(400, 377),
(377, 152),
(152, 148),
(148, 176),
(176, 149),
(149, 150),
(150, 136),
(136, 172),
(172, 58),
(58, 132),
(132, 93),
(93, 234),
(234, 127),
(127, 162),
(162, 21),
(21, 54),
(54, 103),
(103, 67),
(67, 109),
(109, 10)
])
class FaceMesh(SolutionBase):
"""MediaPipe FaceMesh.
MediaPipe FaceMesh 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,
min_detection_confidence=0.5,
min_tracking_confidence=0.5):
"""Initializes a MediaPipe FaceMesh 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.
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,
},
calculator_params={
'ConstantSidePacketCalculator.packet': [
constant_side_packet_calculator_pb2
.ConstantSidePacketCalculatorOptions.ConstantSidePacket(
bool_value=not static_image_mode)
],
'facedetectionfrontcpu__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})