talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 08c2016a4df5aef571b464a4d4491f38c6b2af10
This commit is contained in:
MediaPipe Team
2021-06-03 17:04:35 -04:00
committed by chuoling
parent ae05ad04b3
commit 8b57bf879b
118 changed files with 3999 additions and 391 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/BUILD
+1
View File
@@ -72,5 +72,6 @@ cc_library(
"//mediapipe/modules/pose_detection:pose_detection_cpu",
"//mediapipe/modules/pose_landmark:pose_landmark_by_roi_cpu",
"//mediapipe/modules/pose_landmark:pose_landmark_cpu",
"//mediapipe/modules/selfie_segmentation:selfie_segmentation_cpu",
],
)
mediapipe/python/solutions/__init__.py
+1
View File
@@ -21,3 +21,4 @@ import mediapipe.python.solutions.hands
import mediapipe.python.solutions.holistic
import mediapipe.python.solutions.objectron
import mediapipe.python.solutions.pose
import mediapipe.python.solutions.selfie_segmentation
mediapipe/python/solutions/drawing_utils.py
+2 -2
View File
@@ -15,7 +15,7 @@
"""MediaPipe solution drawing utils."""
import math
from typing import List, Tuple, Union
from typing import List, Optional, Tuple, Union
import cv2
import dataclasses
@@ -116,7 +116,7 @@ def draw_detection(
def draw_landmarks(
image: np.ndarray,
landmark_list: landmark_pb2.NormalizedLandmarkList,
connections: List[Tuple[int, int]] = None,
connections: Optional[List[Tuple[int, int]]] = None,
landmark_drawing_spec: DrawingSpec = DrawingSpec(color=RED_COLOR),
connection_drawing_spec: DrawingSpec = DrawingSpec()):
"""Draws the landmarks and the connections on the image.
mediapipe/python/solutions/face_detection_test.py
+2 -1
View File
@@ -56,7 +56,8 @@ class FaceDetectionTest(absltest.TestCase):
self.assertIsNone(results.detections)
def test_face(self):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/face.jpg')
image_path = os.path.join(os.path.dirname(__file__),
'testdata/portrait.jpg')
image = cv2.imread(image_path)
with mp_faces.FaceDetection(min_detection_confidence=0.5) as faces:
for idx in range(5):
mediapipe/python/solutions/face_mesh_test.py
+2 -1
View File
@@ -96,7 +96,8 @@ class FaceMeshTest(parameterized.TestCase):
@parameterized.named_parameters(('static_image_mode', True, 1),
('video_mode', False, 5))
def test_face(self, static_image_mode: bool, num_frames: int):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/face.jpg')
image_path = os.path.join(os.path.dirname(__file__),
'testdata/portrait.jpg')
image = cv2.imread(image_path)
with mp_faces.FaceMesh(
static_image_mode=static_image_mode,
mediapipe/python/solutions/pose_test.py
+12 -12
View File
@@ -30,18 +30,18 @@ from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import pose as mp_pose
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 30 # pixels
EXPECTED_POSE_LANDMARKS = np.array([[460, 287], [469, 277], [472, 276],
[475, 276], [464, 277], [463, 277],
[463, 276], [492, 277], [472, 277],
[471, 295], [465, 295], [542, 323],
[448, 318], [619, 319], [372, 313],
[695, 316], [296, 308], [717, 313],
[273, 304], [718, 304], [280, 298],
[709, 307], [289, 303], [521, 470],
[459, 466], [626, 533], [364, 500],
[704, 616], [347, 614], [710, 631],
[357, 633], [737, 625], [306, 639]])
DIFF_THRESHOLD = 15 # pixels
EXPECTED_POSE_LANDMARKS = np.array([[460, 283], [467, 273], [471, 273],
[474, 273], [465, 273], [465, 273],
[466, 273], [491, 277], [480, 277],
[470, 294], [465, 294], [545, 319],
[453, 329], [622, 323], [375, 316],
[696, 316], [299, 307], [719, 316],
[278, 306], [721, 311], [274, 304],
[713, 313], [283, 306], [520, 476],
[467, 471], [612, 550], [358, 490],
[701, 613], [349, 611], [709, 624],
[363, 630], [730, 633], [303, 628]])
class PoseTest(parameterized.TestCase):
mediapipe/python/solutions/selfie_segmentation.py
+76
View File
@@ -0,0 +1,76 @@
# Copyright 2021 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 Selfie Segmentation."""
from typing import NamedTuple
import numpy as np
# The following imports are needed because python pb2 silently discards
# unknown protobuf fields.
# pylint: disable=unused-import
from mediapipe.calculators.core import constant_side_packet_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_segmentation_calculator_pb2
from mediapipe.calculators.util import local_file_contents_calculator_pb2
from mediapipe.framework.tool import switch_container_pb2
# pylint: enable=unused-import
from mediapipe.python.solution_base import SolutionBase
BINARYPB_FILE_PATH = 'mediapipe/modules/selfie_segmentation/selfie_segmentation_cpu.binarypb'
class SelfieSegmentation(SolutionBase):
"""MediaPipe Selfie Segmentation.
MediaPipe Selfie Segmentation processes an RGB image and returns a
segmentation mask.
Please refer to
https://solutions.mediapipe.dev/selfie_segmentation#python-solution-api for
usage examples.
"""
def __init__(self, model_selection=0):
"""Initializes a MediaPipe Selfie Segmentation object.
Args:
model_selection: 0 or 1. 0 to select a general-purpose model, and 1 to
select a model more optimized for landscape images. See details in
https://solutions.mediapipe.dev/selfie_segmentation#model_selection.
"""
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
side_inputs={
'model_selection': model_selection,
},
outputs=['segmentation_mask'])
def process(self, image: np.ndarray) -> NamedTuple:
"""Processes an RGB image and returns a segmentation mask.
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 "segmentation_mask" field that contains a float
type 2d np array representing the mask.
"""
return super().process(input_data={'image': image})
mediapipe/python/solutions/selfie_segmentation_test.py
+68
View File
@@ -0,0 +1,68 @@
# Copyright 2021 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.
"""Tests for mediapipe.python.solutions.selfie_segmentation."""
import os
from absl.testing import absltest
from absl.testing import parameterized
import cv2
import numpy as np
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import selfie_segmentation as mp_selfie_segmentation
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
class SelfieSegmentationTest(parameterized.TestCase):
def _draw(self, frame: np.ndarray, mask: np.ndarray):
frame = np.minimum(frame, np.stack((mask,) * 3, axis=-1))
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] + '.png')
cv2.imwrite(path, frame)
def test_invalid_image_shape(self):
with mp_selfie_segmentation.SelfieSegmentation() as selfie_segmentation:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
selfie_segmentation.process(
np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
with mp_selfie_segmentation.SelfieSegmentation() as selfie_segmentation:
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = selfie_segmentation.process(image)
normalized_segmentation_mask = (results.segmentation_mask *
255).astype(int)
self.assertLess(np.amax(normalized_segmentation_mask), 1)
@parameterized.named_parameters(('general', 0), ('landscape', 1))
def test_segmentation(self, model_selection):
image_path = os.path.join(os.path.dirname(__file__),
'testdata/portrait.jpg')
image = cv2.imread(image_path)
with mp_selfie_segmentation.SelfieSegmentation(
model_selection=model_selection) as selfie_segmentation:
results = selfie_segmentation.process(
cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
normalized_segmentation_mask = (results.segmentation_mask *
255).astype(int)
self._draw(image.copy(), normalized_segmentation_mask)
if __name__ == '__main__':
absltest.main()