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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 27c70b5fe62ab71189d358ca122ee4b19c817a8f
This commit is contained in:
MediaPipe Team
2021-07-27 19:36:32 -04:00
committed by chuoling
parent 374f5e2e7e
commit 50c92c6623
158 changed files with 4704 additions and 621 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/image_frame_test.py
+1 -1
View File
@@ -95,7 +95,7 @@ class ImageFrameTest(absltest.TestCase):
'RGBA')
image_frame = mp.ImageFrame(
image_format=mp.ImageFormat.SRGBA64,
data=np.asarray(img, dtype=np.uint16))
data=np.asarray(img).astype(np.uint16))
self.assertTrue(np.array_equal(np.asarray(img), image_frame.numpy_view()))
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image_frame[1000, 1000, 1000])
mediapipe/python/image_test.py
+1 -1
View File
@@ -93,7 +93,7 @@ class ImageTest(absltest.TestCase):
'RGBA')
image = mp.Image(
image_format=mp.ImageFormat.SRGBA64,
data=np.asarray(img, dtype=np.uint16))
data=np.asarray(img).astype(np.uint16))
self.assertTrue(np.array_equal(np.asarray(img), image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[1000, 1000, 1000])
mediapipe/python/solution_base.py
+27 -12
View File
@@ -29,6 +29,7 @@ from typing import Any, Iterable, List, Mapping, NamedTuple, Optional, Union
import numpy as np
from google.protobuf import descriptor
from google.protobuf import message
# resources dependency
# pylint: disable=unused-import
# pylint: enable=unused-import
@@ -120,6 +121,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO,
'::mediapipe::ClassificationList':
_PacketDataType.PROTO,
'::mediapipe::ClassificationListCollection':
_PacketDataType.PROTO,
'::mediapipe::Detection':
_PacketDataType.PROTO,
'::mediapipe::DetectionList':
@@ -128,6 +131,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO,
'::mediapipe::LandmarkList':
_PacketDataType.PROTO,
'::mediapipe::LandmarkListCollection':
_PacketDataType.PROTO,
'::mediapipe::NormalizedLandmark':
_PacketDataType.PROTO,
'::mediapipe::FrameAnnotation':
@@ -140,6 +145,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO,
'::mediapipe::NormalizedLandmarkList':
_PacketDataType.PROTO,
'::mediapipe::NormalizedLandmarkListCollection':
_PacketDataType.PROTO,
'::mediapipe::Image':
_PacketDataType.IMAGE,
'::std::vector<::mediapipe::Classification>':
@@ -257,17 +264,19 @@ class SolutionBase:
# types from "_input_stream_type_info" and then auto generate the process
# method signature by "inspect.Signature" in __init__.
def process(
self, input_data: Union[np.ndarray, Mapping[str,
np.ndarray]]) -> NamedTuple:
self, input_data: Union[np.ndarray, Mapping[str, Union[np.ndarray,
message.Message]]]
) -> NamedTuple:
"""Processes a set of RGB image data and output SolutionOutputs.
Args:
input_data: Either a single numpy ndarray object representing the solo
image input of a graph or a mapping from the stream name to the image
data that represents every input streams of a graph.
image input of a graph or a mapping from the stream name to the image or
proto data that represents every input streams of a graph.
Raises:
NotImplementedError: If input_data contains non image data.
NotImplementedError: If input_data contains audio data or a list of proto
objects.
RuntimeError: If the underlying graph occurs any error.
ValueError: If the input image data is not three channel RGB.
@@ -300,8 +309,15 @@ class SolutionBase:
self._simulated_timestamp += 33333
for stream_name, data in input_dict.items():
input_stream_type = self._input_stream_type_info[stream_name]
if (input_stream_type == _PacketDataType.IMAGE_FRAME or
input_stream_type == _PacketDataType.IMAGE):
if (input_stream_type == _PacketDataType.PROTO_LIST or
input_stream_type == _PacketDataType.AUDIO):
# TODO: Support audio data.
raise NotImplementedError(
f'SolutionBase can only process non-audio and non-proto-list data. '
f'{self._input_stream_type_info[stream_name].name} '
f'type is not supported yet.')
elif (input_stream_type == _PacketDataType.IMAGE_FRAME or
input_stream_type == _PacketDataType.IMAGE):
if data.shape[2] != RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
self._graph.add_packet_to_input_stream(
@@ -309,11 +325,10 @@ class SolutionBase:
packet=self._make_packet(input_stream_type,
data).at(self._simulated_timestamp))
else:
# TODO: Support audio data.
raise NotImplementedError(
f'SolutionBase can only process image data. '
f'{self._input_stream_type_info[stream_name].name} '
f'type is not supported yet.')
self._graph.add_packet_to_input_stream(
stream=stream_name,
packet=self._make_packet(input_stream_type,
data).at(self._simulated_timestamp))
self._graph.wait_until_idle()
# Create a NamedTuple object where the field names are mapping to the graph
mediapipe/python/solution_base_test.py
+33 -2
View File
@@ -93,7 +93,37 @@ class SolutionBaseTest(parameterized.TestCase):
with self.assertRaisesRegex(error_type, error_message):
solution_base.SolutionBase(graph_config=config_proto)
def test_invalid_input_data_type(self):
def test_valid_input_data_type_proto(self):
text_config = """
input_stream: 'input_detections'
output_stream: 'output_detections'
node {
calculator: 'DetectionUniqueIdCalculator'
input_stream: 'DETECTION_LIST:input_detections'
output_stream: 'DETECTION_LIST:output_detections'
}
"""
config_proto = text_format.Parse(text_config,
calculator_pb2.CalculatorGraphConfig())
with solution_base.SolutionBase(graph_config=config_proto) as solution:
input_detections = detection_pb2.DetectionList()
detection_1 = input_detections.detection.add()
text_format.Parse('score: 0.5', detection_1)
detection_2 = input_detections.detection.add()
text_format.Parse('score: 0.8', detection_2)
results = solution.process({'input_detections': input_detections})
self.assertTrue(hasattr(results, 'output_detections'))
self.assertLen(results.output_detections.detection, 2)
expected_detection_1 = detection_pb2.Detection()
text_format.Parse('score: 0.5, detection_id: 1', expected_detection_1)
expected_detection_2 = detection_pb2.Detection()
text_format.Parse('score: 0.8, detection_id: 2', expected_detection_2)
self.assertEqual(results.output_detections.detection[0],
expected_detection_1)
self.assertEqual(results.output_detections.detection[1],
expected_detection_2)
def test_invalid_input_data_type_proto_vector(self):
text_config = """
input_stream: 'input_detections'
output_stream: 'output_detections'
@@ -110,7 +140,8 @@ class SolutionBaseTest(parameterized.TestCase):
text_format.Parse('score: 0.5', detection)
with self.assertRaisesRegex(
NotImplementedError,
'SolutionBase can only process image data. PROTO_LIST type is not supported.'
'SolutionBase can only process non-audio and non-proto-list data. '
+ 'PROTO_LIST type is not supported.'
):
solution.process({'input_detections': detection})
mediapipe/python/solutions/__init__.py
+1
View File
@@ -14,6 +14,7 @@
"""MediaPipe Solutions Python API."""
import mediapipe.python.solutions.drawing_styles
import mediapipe.python.solutions.drawing_utils
import mediapipe.python.solutions.face_detection
import mediapipe.python.solutions.face_mesh
mediapipe/python/solutions/drawing_styles.py
+146
View File
@@ -0,0 +1,146 @@
# 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 requi_RED 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 solution drawing styles."""
from typing import Mapping, Tuple
from mediapipe.python.solutions.drawing_utils import DrawingSpec
from mediapipe.python.solutions.hands import HandLandmark
_RADIUS = 5
_RED = (54, 67, 244)
_GREEN = (118, 230, 0)
_BLUE = (192, 101, 21)
_YELLOW = (0, 204, 255)
_GRAY = (174, 164, 144)
_PURPLE = (128, 64, 128)
_PEACH = (180, 229, 255)
# Hands
_THICKNESS_WRIST_MCP = 3
_THICKNESS_FINGER = 2
_THICKNESS_DOT = -1
# Hand landmarks
_PALM_LANMARKS = (HandLandmark.WRIST, HandLandmark.THUMB_CMC,
HandLandmark.INDEX_FINGER_MCP, HandLandmark.MIDDLE_FINGER_MCP,
HandLandmark.RING_FINGER_MCP, HandLandmark.PINKY_MCP)
_THUMP_LANDMARKS = (HandLandmark.THUMB_MCP, HandLandmark.THUMB_IP,
HandLandmark.THUMB_TIP)
_INDEX_FINGER_LANDMARKS = (HandLandmark.INDEX_FINGER_PIP,
HandLandmark.INDEX_FINGER_DIP,
HandLandmark.INDEX_FINGER_TIP)
_MIDDLE_FINGER_LANDMARKS = (HandLandmark.MIDDLE_FINGER_PIP,
HandLandmark.MIDDLE_FINGER_DIP,
HandLandmark.MIDDLE_FINGER_TIP)
_RING_FINGER_LANDMARKS = (HandLandmark.RING_FINGER_PIP,
HandLandmark.RING_FINGER_DIP,
HandLandmark.RING_FINGER_TIP)
_PINKY_FINGER_LANDMARKS = (HandLandmark.PINKY_PIP, HandLandmark.PINKY_DIP,
HandLandmark.PINKY_TIP)
_HAND_LANDMARK_STYLE = {
_PALM_LANMARKS:
DrawingSpec(
color=_RED, thickness=_THICKNESS_DOT, circle_radius=_RADIUS),
_THUMP_LANDMARKS:
DrawingSpec(
color=_PEACH, thickness=_THICKNESS_DOT, circle_radius=_RADIUS),
_INDEX_FINGER_LANDMARKS:
DrawingSpec(
color=_PURPLE, thickness=_THICKNESS_DOT, circle_radius=_RADIUS),
_MIDDLE_FINGER_LANDMARKS:
DrawingSpec(
color=_YELLOW, thickness=_THICKNESS_DOT, circle_radius=_RADIUS),
_RING_FINGER_LANDMARKS:
DrawingSpec(
color=_GREEN, thickness=_THICKNESS_DOT, circle_radius=_RADIUS),
_PINKY_FINGER_LANDMARKS:
DrawingSpec(
color=_BLUE, thickness=_THICKNESS_DOT, circle_radius=_RADIUS),
}
# Hand connections
_PALM_CONNECTIONS = ((HandLandmark.WRIST, HandLandmark.THUMB_CMC),
(HandLandmark.WRIST, HandLandmark.INDEX_FINGER_MCP),
(HandLandmark.MIDDLE_FINGER_MCP,
HandLandmark.RING_FINGER_MCP),
(HandLandmark.RING_FINGER_MCP, HandLandmark.PINKY_MCP),
(HandLandmark.INDEX_FINGER_MCP,
HandLandmark.MIDDLE_FINGER_MCP), (HandLandmark.WRIST,
HandLandmark.PINKY_MCP))
_THUMB_CONNECTIONS = ((HandLandmark.THUMB_CMC, HandLandmark.THUMB_MCP),
(HandLandmark.THUMB_MCP, HandLandmark.THUMB_IP),
(HandLandmark.THUMB_IP, HandLandmark.THUMB_TIP))
_INDEX_FINGER_CONNECTIONS = ((HandLandmark.INDEX_FINGER_MCP,
HandLandmark.INDEX_FINGER_PIP),
(HandLandmark.INDEX_FINGER_PIP,
HandLandmark.INDEX_FINGER_DIP),
(HandLandmark.INDEX_FINGER_DIP,
HandLandmark.INDEX_FINGER_TIP))
_MIDDLE_FINGER_CONNECTIONS = ((HandLandmark.MIDDLE_FINGER_MCP,
HandLandmark.MIDDLE_FINGER_PIP),
(HandLandmark.MIDDLE_FINGER_PIP,
HandLandmark.MIDDLE_FINGER_DIP),
(HandLandmark.MIDDLE_FINGER_DIP,
HandLandmark.MIDDLE_FINGER_TIP))
_RING_FINGER_CONNECTIONS = ((HandLandmark.RING_FINGER_MCP,
HandLandmark.RING_FINGER_PIP),
(HandLandmark.RING_FINGER_PIP,
HandLandmark.RING_FINGER_DIP),
(HandLandmark.RING_FINGER_DIP,
HandLandmark.RING_FINGER_TIP))
_PINKY_FINGER_CONNECTIONS = ((HandLandmark.PINKY_MCP, HandLandmark.PINKY_PIP),
(HandLandmark.PINKY_PIP, HandLandmark.PINKY_DIP),
(HandLandmark.PINKY_DIP, HandLandmark.PINKY_TIP))
_HAND_CONNECTION_STYLE = {
_PALM_CONNECTIONS:
DrawingSpec(color=_GRAY, thickness=_THICKNESS_WRIST_MCP),
_THUMB_CONNECTIONS:
DrawingSpec(color=_PEACH, thickness=_THICKNESS_FINGER),
_INDEX_FINGER_CONNECTIONS:
DrawingSpec(color=_PURPLE, thickness=_THICKNESS_FINGER),
_MIDDLE_FINGER_CONNECTIONS:
DrawingSpec(color=_YELLOW, thickness=_THICKNESS_FINGER),
_RING_FINGER_CONNECTIONS:
DrawingSpec(color=_GREEN, thickness=_THICKNESS_FINGER),
_PINKY_FINGER_CONNECTIONS:
DrawingSpec(color=_BLUE, thickness=_THICKNESS_FINGER)
}
def get_default_hand_landmark_style() -> Mapping[int, DrawingSpec]:
"""Returns the default hand landmark drawing style.
Returns:
A mapping from each hand landmark to the default drawing spec.
"""
hand_landmark_style = {}
for k, v in _HAND_LANDMARK_STYLE.items():
for landmark in k:
hand_landmark_style[landmark] = v
return hand_landmark_style
def get_default_hand_connection_style(
) -> Mapping[Tuple[int, int], DrawingSpec]:
"""Returns the default hand connection drawing style.
Returns:
A mapping from each hand connection to the default drawing spec.
"""
hand_connection_style = {}
for k, v in _HAND_CONNECTION_STYLE.items():
for connection in k:
hand_connection_style[connection] = v
return hand_connection_style
mediapipe/python/solutions/drawing_utils.py
+21 -11
View File
@@ -15,7 +15,7 @@
"""MediaPipe solution drawing utils."""
import math
from typing import List, Optional, Tuple, Union
from typing import List, Mapping, Optional, Tuple, Union
import cv2
import dataclasses
@@ -119,8 +119,12 @@ def draw_landmarks(
image: np.ndarray,
landmark_list: landmark_pb2.NormalizedLandmarkList,
connections: Optional[List[Tuple[int, int]]] = None,
landmark_drawing_spec: DrawingSpec = DrawingSpec(color=RED_COLOR),
connection_drawing_spec: DrawingSpec = DrawingSpec()):
landmark_drawing_spec: Union[DrawingSpec,
Mapping[int, DrawingSpec]] = DrawingSpec(
color=RED_COLOR),
connection_drawing_spec: Union[DrawingSpec,
Mapping[Tuple[int, int],
DrawingSpec]] = DrawingSpec()):
"""Draws the landmarks and the connections on the image.
Args:
@@ -129,9 +133,11 @@ def draw_landmarks(
the image.
connections: A list of landmark index tuples that specifies how landmarks to
be connected in the drawing.
landmark_drawing_spec: A DrawingSpec object that specifies the landmarks'
drawing settings such as color, line thickness, and circle radius.
connection_drawing_spec: A DrawingSpec object that specifies the
landmark_drawing_spec: Either a DrawingSpec object or a mapping from
hand landmarks to the DrawingSpecs that specifies the landmarks' drawing
settings such as color, line thickness, and circle radius.
connection_drawing_spec: Either a DrawingSpec object or a mapping from
hand connections to the DrawingSpecs that specifies the
connections' drawing settings such as color and line thickness.
Raises:
@@ -165,14 +171,18 @@ def draw_landmarks(
raise ValueError(f'Landmark index is out of range. Invalid connection '
f'from landmark #{start_idx} to landmark #{end_idx}.')
if start_idx in idx_to_coordinates and end_idx in idx_to_coordinates:
drawing_spec = connection_drawing_spec[connection] if isinstance(
connection_drawing_spec, Mapping) else connection_drawing_spec
cv2.line(image, idx_to_coordinates[start_idx],
idx_to_coordinates[end_idx], connection_drawing_spec.color,
connection_drawing_spec.thickness)
idx_to_coordinates[end_idx], drawing_spec.color,
drawing_spec.thickness)
# Draws landmark points after finishing the connection lines, which is
# aesthetically better.
for landmark_px in idx_to_coordinates.values():
cv2.circle(image, landmark_px, landmark_drawing_spec.circle_radius,
landmark_drawing_spec.color, landmark_drawing_spec.thickness)
for idx, landmark_px in idx_to_coordinates.items():
drawing_spec = landmark_drawing_spec[idx] if isinstance(
landmark_drawing_spec, Mapping) else landmark_drawing_spec
cv2.circle(image, landmark_px, drawing_spec.circle_radius,
drawing_spec.color, drawing_spec.thickness)
def draw_axis(
mediapipe/python/solutions/hands_test.py
+5 -2
View File
@@ -26,6 +26,7 @@ import numpy.testing as npt
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import drawing_styles
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import hands as mp_hands
@@ -51,8 +52,10 @@ class HandsTest(parameterized.TestCase):
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
for hand_landmarks in results.multi_hand_landmarks:
mp_drawing.draw_landmarks(frame, hand_landmarks,
mp_hands.HAND_CONNECTIONS)
mp_drawing.draw_landmarks(
frame, hand_landmarks, mp_hands.HAND_CONNECTIONS,
drawing_styles.get_default_hand_landmark_style(),
drawing_styles.get_default_hand_connection_style())
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
mediapipe/python/solutions/objectron_test.py
+1 -1
View File
@@ -26,7 +26,7 @@ import numpy.testing as npt
from mediapipe.python.solutions import objectron as mp_objectron
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 30 # pixels
DIFF_THRESHOLD = 35 # pixels
EXPECTED_BOX_COORDINATES_PREDICTION = [[[236, 413], [408, 474], [135, 457],
[383, 505], [80, 478], [408, 345],
[130, 347], [384, 355], [72, 353]],