Project import generated by Copybara.

GitOrigin-RevId: d073f8e21be2fcc0e503cb97c6695078b6b75310
This commit is contained in:
MediaPipe Team
2021-02-27 03:30:05 -05:00
committed by chuoling
parent 39309bedba
commit 350fbb2100
755 changed files with 16391 additions and 11075 deletions
+3
View File
@@ -47,6 +47,8 @@ pybind_extension(
"//mediapipe/framework/formats:classification_registration",
"//mediapipe/framework/formats:detection_registration",
"//mediapipe/framework/formats:landmark_registration",
"//mediapipe/framework/formats:rect_registration",
"//mediapipe/modules/objectron/calculators:annotation_registration",
],
)
@@ -64,6 +66,7 @@ cc_library(
"//mediapipe/modules/face_landmark:face_landmark_front_cpu",
"//mediapipe/modules/hand_landmark:hand_landmark_tracking_cpu",
"//mediapipe/modules/holistic_landmark:holistic_landmark_cpu",
"//mediapipe/modules/objectron:objectron_cpu",
"//mediapipe/modules/palm_detection:palm_detection_cpu",
"//mediapipe/modules/pose_detection:pose_detection_cpu",
"//mediapipe/modules/pose_landmark:pose_landmark_by_roi_cpu",
+3 -3
View File
@@ -376,7 +376,7 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
calculator_graph.def(
"get_combined_error_message",
[](CalculatorGraph* self) {
mediapipe::Status error_status;
absl::Status error_status;
if (self->GetCombinedErrors(&error_status) && !error_status.ok()) {
return error_status.ToString();
}
@@ -400,7 +400,7 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
// Acquire a mutex so that only one callback_fn can run at once.
absl::MutexLock lock(&callback_mutex);
callback_fn(stream_name, packet);
return mediapipe::OkStatus();
return absl::OkStatus();
}));
},
R"doc(Observe the named output stream.
@@ -438,7 +438,7 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
[](CalculatorGraph* self, const std::string& packet_name) {
auto status_or_packet = self->GetOutputSidePacket(packet_name);
RaisePyErrorIfNotOk(status_or_packet.status());
return status_or_packet.ValueOrDie();
return status_or_packet.value();
},
R"doc(Get output side packet by name after the graph is done.
+2 -2
View File
@@ -602,7 +602,7 @@ void InternalPacketCreators(pybind11::module* m) {
"_create_proto",
[](const std::string& type_name, const py::bytes& serialized_proto) {
using packet_internal::HolderBase;
mediapipe::StatusOr<std::unique_ptr<HolderBase>> maybe_holder =
absl::StatusOr<std::unique_ptr<HolderBase>> maybe_holder =
packet_internal::MessageHolderRegistry::CreateByName(type_name);
if (!maybe_holder.ok()) {
throw RaisePyError(
@@ -612,7 +612,7 @@ void InternalPacketCreators(pybind11::module* m) {
}
// Creates a Packet with the concrete C++ payload type.
std::unique_ptr<HolderBase> message_holder =
std::move(maybe_holder).ValueOrDie();
std::move(maybe_holder).value();
auto* copy = const_cast<proto_ns::MessageLite*>(
message_holder->GetProtoMessageLite());
copy->ParseFromString(std::string(serialized_proto));
+5 -5
View File
@@ -358,7 +358,7 @@ void InternalPacketGetters(pybind11::module* m) {
[](Packet& packet) {
auto proto_vector = packet.GetVectorOfProtoMessageLitePtrs();
RaisePyErrorIfNotOk(proto_vector.status());
return proto_vector.ValueOrDie().size();
return proto_vector.value().size();
},
py::return_value_policy::move);
@@ -367,10 +367,10 @@ void InternalPacketGetters(pybind11::module* m) {
[](Packet& packet) {
auto proto_vector = packet.GetVectorOfProtoMessageLitePtrs();
RaisePyErrorIfNotOk(proto_vector.status());
if (proto_vector.ValueOrDie().empty()) {
if (proto_vector.value().empty()) {
return std::string();
}
return proto_vector.ValueOrDie()[0]->GetTypeName();
return proto_vector.value()[0]->GetTypeName();
},
py::return_value_policy::move);
@@ -391,10 +391,10 @@ void InternalPacketGetters(pybind11::module* m) {
[](Packet& packet) {
auto proto_vector = packet.GetVectorOfProtoMessageLitePtrs();
RaisePyErrorIfNotOk(proto_vector.status());
int size = proto_vector.ValueOrDie().size();
int size = proto_vector.value().size();
std::vector<py::bytes> results;
results.reserve(size);
for (const proto_ns::MessageLite* ptr : proto_vector.ValueOrDie()) {
for (const proto_ns::MessageLite* ptr : proto_vector.value()) {
results.push_back(py::bytes(ptr->SerializeAsString()));
}
return results;
+1 -1
View File
@@ -45,7 +45,7 @@ inline PyObject* StatusCodeToPyError(const ::absl::StatusCode& code) {
}
}
inline void RaisePyErrorIfNotOk(const mediapipe::Status& status) {
inline void RaisePyErrorIfNotOk(const absl::Status& status) {
if (!status.ok()) {
throw RaisePyError(StatusCodeToPyError(status.code()),
status.message().data());
@@ -98,7 +98,7 @@ void ValidatedGraphConfigSubmodule(pybind11::module* module) {
[](ValidatedGraphConfig& self, const std::string& stream_name) {
auto status_or_type_name = self.RegisteredStreamTypeName(stream_name);
RaisePyErrorIfNotOk(status_or_type_name.status());
return status_or_type_name.ValueOrDie();
return status_or_type_name.value();
},
R"doc(Return the registered type name of the specified stream if it can be determined.
@@ -122,7 +122,7 @@ void ValidatedGraphConfigSubmodule(pybind11::module* module) {
auto status_or_type_name =
self.RegisteredSidePacketTypeName(side_packet_name);
RaisePyErrorIfNotOk(status_or_type_name.status());
return status_or_type_name.ValueOrDie();
return status_or_type_name.value();
},
R"doc(Return the registered type name of the specified side packet if it can be determined.
+23 -9
View File
@@ -45,6 +45,8 @@ from mediapipe.calculators.util import thresholding_calculator_pb2
from mediapipe.framework.formats import classification_pb2
from mediapipe.framework.formats import landmark_pb2
from mediapipe.framework.formats import rect_pb2
from mediapipe.modules.objectron.calculators import annotation_data_pb2
from mediapipe.modules.objectron.calculators import lift_2d_frame_annotation_to_3d_calculator_pb2
# pylint: enable=unused-import
from mediapipe.python._framework_bindings import calculator_graph
from mediapipe.python._framework_bindings import image_frame
@@ -71,6 +73,9 @@ CALCULATOR_TO_OPTIONS = {
'TensorsToDetectionsCalculator':
tensors_to_detections_calculator_pb2
.TensorsToDetectionsCalculatorOptions,
'Lift2DFrameAnnotationTo3DCalculator':
lift_2d_frame_annotation_to_3d_calculator_pb2
.Lift2DFrameAnnotationTo3DCalculatorOptions,
}
@@ -120,6 +125,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO,
'::mediapipe::NormalizedLandmark':
_PacketDataType.PROTO,
'::mediapipe::FrameAnnotation':
_PacketDataType.PROTO,
'::mediapipe::Trigger':
_PacketDataType.PROTO,
'::mediapipe::Rect':
@@ -157,15 +164,14 @@ class SolutionBase:
shutdown.
Example usage:
hand_tracker = solution_base.SolutionBase(
binary_graph_path='mediapipe/modules/hand_landmark/hand_landmark_tracking_cpu.binarypb',
side_inputs={'num_hands': 2})
# Read an image and convert the BGR image to RGB.
input_image = cv2.cvtColor(cv2.imread('/tmp/hand.png'), COLOR_BGR2RGB)
results = hand_tracker.process(input_image)
print(results.palm_detections)
print(results.multi_hand_landmarks)
hand_tracker.close()
with solution_base.SolutionBase(
binary_graph_path='mediapipe/modules/hand_landmark/hand_landmark_tracking_cpu.binarypb',
side_inputs={'num_hands': 2}) as hand_tracker:
# Read an image and convert the BGR image to RGB.
input_image = cv2.cvtColor(cv2.imread('/tmp/hand.png'), COLOR_BGR2RGB)
results = hand_tracker.process(input_image)
print(results.palm_detections)
print(results.multi_hand_landmarks)
"""
def __init__(
@@ -479,3 +485,11 @@ class SolutionBase:
else:
return getattr(packet_getter, 'get_' + packet_data_type.value)(
output_packet)
def __enter__(self):
"""A "with" statement support."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Closes all the input sources and the graph."""
self.close()
+16 -17
View File
@@ -105,14 +105,14 @@ class SolutionBaseTest(parameterized.TestCase):
"""
config_proto = text_format.Parse(text_config,
calculator_pb2.CalculatorGraphConfig())
solution = solution_base.SolutionBase(graph_config=config_proto)
detection = detection_pb2.Detection()
text_format.Parse('score: 0.5', detection)
with self.assertRaisesRegex(
NotImplementedError,
'SolutionBase can only process image data. PROTO_LIST type is not supported.'
):
solution.process({'input_detections': detection})
with solution_base.SolutionBase(graph_config=config_proto) as solution:
detection = detection_pb2.Detection()
text_format.Parse('score: 0.5', detection)
with self.assertRaisesRegex(
NotImplementedError,
'SolutionBase can only process image data. PROTO_LIST type is not supported.'
):
solution.process({'input_detections': detection})
def test_invalid_input_image_data(self):
text_config = """
@@ -131,10 +131,10 @@ class SolutionBaseTest(parameterized.TestCase):
"""
config_proto = text_format.Parse(text_config,
calculator_pb2.CalculatorGraphConfig())
solution = solution_base.SolutionBase(graph_config=config_proto)
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
solution.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
with solution_base.SolutionBase(graph_config=config_proto) as solution:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
solution.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
@parameterized.named_parameters(('graph_without_side_packets', """
input_stream: 'image_in'
@@ -272,15 +272,14 @@ class SolutionBaseTest(parameterized.TestCase):
side_inputs=None,
calculator_params=None):
input_image = np.arange(27, dtype=np.uint8).reshape(3, 3, 3)
solution = solution_base.SolutionBase(
with solution_base.SolutionBase(
graph_config=config_proto,
side_inputs=side_inputs,
calculator_params=calculator_params)
outputs = solution.process(input_image)
calculator_params=calculator_params) as solution:
outputs = solution.process(input_image)
outputs2 = solution.process({'image_in': input_image})
self.assertTrue(np.array_equal(input_image, outputs.image_out))
outputs2 = solution.process({'image_in': input_image})
self.assertTrue(np.array_equal(input_image, outputs2.image_out))
solution.close()
if __name__ == '__main__':
+2
View File
@@ -15,7 +15,9 @@
"""MediaPipe Solutions Python API."""
import mediapipe.python.solutions.drawing_utils
import mediapipe.python.solutions.face_detection
import mediapipe.python.solutions.face_mesh
import mediapipe.python.solutions.hands
import mediapipe.python.solutions.holistic
import mediapipe.python.solutions.objectron
import mediapipe.python.solutions.pose
+113
View File
@@ -21,6 +21,8 @@ import cv2
import dataclasses
import numpy as np
from mediapipe.framework.formats import detection_pb2
from mediapipe.framework.formats import location_data_pb2
from mediapipe.framework.formats import landmark_pb2
PRESENCE_THRESHOLD = 0.5
@@ -58,6 +60,57 @@ def _normalized_to_pixel_coordinates(
return x_px, y_px
def draw_detection(
image: np.ndarray,
detection: detection_pb2.Detection,
keypoint_drawing_spec: DrawingSpec = DrawingSpec(color=RED_COLOR),
bbox_drawing_spec: DrawingSpec = DrawingSpec()):
"""Draws the detction bounding box and keypoints on the image.
Args:
image: A three channel RGB image represented as numpy ndarray.
detection: A detection proto message to be annotated on the image.
keypoint_drawing_spec: A DrawingSpec object that specifies the keypoints'
drawing settings such as color, line thickness, and circle radius.
bbox_drawing_spec: A DrawingSpec object that specifies the bounding box's
drawing settings such as color and line thickness.
Raises:
ValueError: If one of the followings:
a) If the input image is not three channel RGB.
b) If the location data is not relative data.
"""
if not detection.location_data:
return
if image.shape[2] != RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
image_rows, image_cols, _ = image.shape
location = detection.location_data
if location.format != location_data_pb2.LocationData.RELATIVE_BOUNDING_BOX:
raise ValueError(
'LocationData must be relative for this drawing funtion to work.')
# Draws keypoints.
for keypoint in location.relative_keypoints:
keypoint_px = _normalized_to_pixel_coordinates(keypoint.x, keypoint.y,
image_cols, image_rows)
cv2.circle(image, keypoint_px, keypoint_drawing_spec.circle_radius,
keypoint_drawing_spec.color, keypoint_drawing_spec.thickness)
# Draws bounding box if exists.
if not location.HasField('relative_bounding_box'):
return
relative_bounding_box = location.relative_bounding_box
rect_start_point = _normalized_to_pixel_coordinates(
relative_bounding_box.xmin, relative_bounding_box.ymin, image_cols,
image_rows)
rect_end_point = _normalized_to_pixel_coordinates(
relative_bounding_box.xmin + relative_bounding_box.width,
relative_bounding_box.ymin + +relative_bounding_box.height, image_cols,
image_rows)
cv2.rectangle(image, rect_start_point, rect_end_point,
bbox_drawing_spec.color, bbox_drawing_spec.thickness)
def draw_landmarks(
image: np.ndarray,
landmark_list: landmark_pb2.NormalizedLandmarkList,
@@ -116,3 +169,63 @@ def draw_landmarks(
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)
def draw_axis(
image: np.ndarray,
rotation: np.ndarray,
translation: np.ndarray,
focal_length: Tuple[float, float] = (1.0, 1.0),
principal_point: Tuple[float, float] = (0.0, 0.0),
axis_length: float = 0.1,
x_axis_drawing_spec: DrawingSpec = DrawingSpec(color=(0, 0, 255)),
y_axis_drawing_spec: DrawingSpec = DrawingSpec(color=(0, 128, 0)),
z_axis_drawing_spec: DrawingSpec = DrawingSpec(color=(255, 0, 0))):
"""Draws the 3D axis on the image.
Args:
image: A three channel RGB image represented as numpy ndarray.
rotation: Rotation matrix from object to camera coordinate frame.
translation: Translation vector from object to camera coordinate frame.
focal_length: camera focal length along x and y directions.
principal_point: camera principal point in x and y.
axis_length: length of the axis in the drawing.
x_axis_drawing_spec: A DrawingSpec object that specifies the x axis
drawing settings such as color, line thickness.
y_axis_drawing_spec: A DrawingSpec object that specifies the y axis
drawing settings such as color, line thickness.
z_axis_drawing_spec: A DrawingSpec object that specifies the z axis
drawing settings such as color, line thickness.
Raises:
ValueError: If one of the followings:
a) If the input image is not three channel RGB.
"""
if image.shape[2] != RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
image_rows, image_cols, _ = image.shape
# Create axis points in camera coordinate frame.
axis_world = np.float32([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]])
axis_cam = np.matmul(rotation, axis_length*axis_world.T).T + translation
x = axis_cam[..., 0]
y = axis_cam[..., 1]
z = axis_cam[..., 2]
# Project 3D points to NDC space.
fx, fy = focal_length
px, py = principal_point
x_ndc = -fx * x / z + px
y_ndc = -fy * y / z + py
# Convert from NDC space to image space.
x_im = np.int32((1 + x_ndc) * 0.5 * image_cols)
y_im = np.int32((1 - y_ndc) * 0.5 * image_rows)
# Draw xyz axis on the image.
origin = (x_im[0], y_im[0])
x_axis = (x_im[1], y_im[1])
y_axis = (x_im[2], y_im[2])
z_axis = (x_im[3], y_im[3])
image = cv2.arrowedLine(image, origin, x_axis, x_axis_drawing_spec.color,
x_axis_drawing_spec.thickness)
image = cv2.arrowedLine(image, origin, y_axis, y_axis_drawing_spec.color,
y_axis_drawing_spec.thickness)
image = cv2.arrowedLine(image, origin, z_axis, z_axis_drawing_spec.color,
z_axis_drawing_spec.thickness)
@@ -21,11 +21,13 @@ import numpy as np
from google.protobuf import text_format
from mediapipe.framework.formats import detection_pb2
from mediapipe.framework.formats import landmark_pb2
from mediapipe.python.solutions import drawing_utils
DEFAULT_BBOX_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_CONNECTION_DRAWING_SPEC = drawing_utils.DrawingSpec()
DEFAULT_LANDMARK_DRAWING_SPEC = drawing_utils.DrawingSpec(color=(0, 0, 255))
DEFAULT_CIRCLE_DRAWING_SPEC = drawing_utils.DrawingSpec(color=(0, 0, 255))
class DrawingUtilTest(parameterized.TestCase):
@@ -35,6 +37,9 @@ class DrawingUtilTest(parameterized.TestCase):
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
drawing_utils.draw_landmarks(image, landmark_pb2.NormalizedLandmarkList())
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
drawing_utils.draw_detection(image, detection_pb2.Detection())
def test_invalid_connection(self):
landmark_list = text_format.Parse(
@@ -44,6 +49,46 @@ class DrawingUtilTest(parameterized.TestCase):
with self.assertRaisesRegex(ValueError, 'Landmark index is out of range.'):
drawing_utils.draw_landmarks(image, landmark_list, [(0, 2)])
def test_unqualified_detection(self):
detection = text_format.Parse('location_data {format: GLOBAL}',
detection_pb2.Detection())
image = np.arange(27, dtype=np.uint8).reshape(3, 3, 3)
with self.assertRaisesRegex(ValueError, 'LocationData must be relative'):
drawing_utils.draw_detection(image, detection)
def test_draw_keypoints_only(self):
detection = text_format.Parse(
'location_data {'
' format: RELATIVE_BOUNDING_BOX'
' relative_keypoints {x: 0 y: 1}'
' relative_keypoints {x: 1 y: 0}}', detection_pb2.Detection())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
cv2.circle(expected_result, (0, 99),
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, (99, 0),
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_detection(image, detection)
np.testing.assert_array_equal(image, expected_result)
def test_draw_bboxs_only(self):
detection = text_format.Parse(
'location_data {'
' format: RELATIVE_BOUNDING_BOX'
' relative_bounding_box {xmin: 0 ymin: 0 width: 1 height: 1}}',
detection_pb2.Detection())
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
cv2.rectangle(expected_result, (0, 0), (99, 99),
DEFAULT_BBOX_DRAWING_SPEC.color,
DEFAULT_BBOX_DRAWING_SPEC.thickness)
drawing_utils.draw_detection(image, detection)
np.testing.assert_array_equal(image, expected_result)
@parameterized.named_parameters(
('landmark_list_has_only_one_element', 'landmark {x: 0.1 y: 0.1}'),
('second_landmark_is_invisible',
@@ -54,9 +99,9 @@ class DrawingUtilTest(parameterized.TestCase):
image = np.zeros((100, 100, 3), np.uint8)
expected_result = np.copy(image)
cv2.circle(expected_result, (10, 10),
DEFAULT_LANDMARK_DRAWING_SPEC.circle_radius,
DEFAULT_LANDMARK_DRAWING_SPEC.color,
DEFAULT_LANDMARK_DRAWING_SPEC.thickness)
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_landmarks(image, landmark_list)
np.testing.assert_array_equal(image, expected_result)
@@ -77,13 +122,13 @@ class DrawingUtilTest(parameterized.TestCase):
DEFAULT_CONNECTION_DRAWING_SPEC.color,
DEFAULT_CONNECTION_DRAWING_SPEC.thickness)
cv2.circle(expected_result, start_point,
DEFAULT_LANDMARK_DRAWING_SPEC.circle_radius,
DEFAULT_LANDMARK_DRAWING_SPEC.color,
DEFAULT_LANDMARK_DRAWING_SPEC.thickness)
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, end_point,
DEFAULT_LANDMARK_DRAWING_SPEC.circle_radius,
DEFAULT_LANDMARK_DRAWING_SPEC.color,
DEFAULT_LANDMARK_DRAWING_SPEC.thickness)
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_landmarks(
image=image, landmark_list=landmark_list, connections=[(0, 1)])
np.testing.assert_array_equal(image, expected_result)
@@ -100,13 +145,13 @@ class DrawingUtilTest(parameterized.TestCase):
DEFAULT_CONNECTION_DRAWING_SPEC.color,
DEFAULT_CONNECTION_DRAWING_SPEC.thickness)
cv2.circle(expected_result, start_point,
DEFAULT_LANDMARK_DRAWING_SPEC.circle_radius,
DEFAULT_LANDMARK_DRAWING_SPEC.color,
DEFAULT_LANDMARK_DRAWING_SPEC.thickness)
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
cv2.circle(expected_result, end_point,
DEFAULT_LANDMARK_DRAWING_SPEC.circle_radius,
DEFAULT_LANDMARK_DRAWING_SPEC.color,
DEFAULT_LANDMARK_DRAWING_SPEC.thickness)
DEFAULT_CIRCLE_DRAWING_SPEC.circle_radius,
DEFAULT_CIRCLE_DRAWING_SPEC.color,
DEFAULT_CIRCLE_DRAWING_SPEC.thickness)
drawing_utils.draw_landmarks(
image=image, landmark_list=landmark_list, connections=[(0, 1)])
np.testing.assert_array_equal(image, expected_result)
@@ -0,0 +1,103 @@
# 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 Face Detection."""
import enum
from typing import NamedTuple, Union
import numpy as np
from mediapipe.framework.formats import detection_pb2
from mediapipe.framework.formats import location_data_pb2
# pylint: disable=unused-import
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_detections_calculator_pb2
from mediapipe.calculators.tflite import ssd_anchors_calculator_pb2
from mediapipe.calculators.util import non_max_suppression_calculator_pb2
# pylint: enable=unused-import
from mediapipe.python.solution_base import SolutionBase
BINARYPB_FILE_PATH = 'mediapipe/modules/face_detection/face_detection_front_cpu.binarypb'
def get_key_point(
detection: detection_pb2.Detection, key_point_enum: 'FaceKeyPoint'
) -> Union[None, location_data_pb2.LocationData.RelativeKeypoint]:
"""A convenience method to return a face key point by the FaceKeyPoint type.
Args:
detection: A detection proto message that contains face key points.
key_point_enum: A FaceKeyPoint type.
Returns:
A RelativeKeypoint proto message.
"""
if not detection or not detection.location_data:
return None
return detection.location_data.relative_keypoints[key_point_enum]
class FaceKeyPoint(enum.IntEnum):
"""The enum type of the six face detection key points."""
RIGHT_EYE = 0
LEFT_EYE = 1
NOSE_TIP = 2
MOUTH_CENTER = 3
RIGHT_EAR_TRAGION = 4
LEFT_EAR_TRAGION = 5
class FaceDetection(SolutionBase):
"""MediaPipe Face Detection.
MediaPipe Face Detection processes an RGB image and returns a list of the
detected face location data.
Please refer to
https://solutions.mediapipe.dev/face_detection#python-solution-api
for usage examples.
"""
def __init__(self, min_detection_confidence=0.5):
"""Initializes a MediaPipe Face Detection object.
Args:
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_detection#min_detection_confidence.
"""
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
calculator_params={
'facedetectionfrontcommon__TensorsToDetectionsCalculator.min_score_thresh':
min_detection_confidence,
},
outputs=['detections'])
def process(self, image: np.ndarray) -> NamedTuple:
"""Processes an RGB image and returns a list of the detected face location data.
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 "detections" field that contains a list of the
detected face location data.
"""
return super().process(input_data={'image': image})
@@ -0,0 +1,67 @@
# 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.face_detection."""
import os
from absl.testing import absltest
import cv2
import numpy as np
import numpy.testing as npt
# resources dependency
from mediapipe.python.solutions import face_detection as mp_faces
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
EXPECTED_FACE_KEY_POINTS = [[182, 368], [186, 467], [236, 416], [284, 415],
[203, 310], [212, 521]]
DIFF_THRESHOLD = 10 # pixels
class FaceDetectionTest(absltest.TestCase):
def test_invalid_image_shape(self):
with mp_faces.FaceDetection() as faces:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
faces.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
with mp_faces.FaceDetection(min_detection_confidence=0.5) as faces:
results = faces.process(image)
self.assertIsNone(results.detections)
def test_face(self):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/face.jpg')
image = cv2.flip(cv2.imread(image_path), 1)
with mp_faces.FaceDetection(min_detection_confidence=0.5) as faces:
for _ in range(5):
results = faces.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
location_data = results.detections[0].location_data
x = [keypoint.x for keypoint in location_data.relative_keypoints]
y = [keypoint.y for keypoint in location_data.relative_keypoints]
face_keypoints = np.transpose(np.stack((y, x))) * image.shape[0:2]
prediction_error = np.abs(
np.asarray(face_keypoints) - np.asarray(EXPECTED_FACE_KEY_POINTS))
self.assertLen(results.detections, 1)
self.assertLen(location_data.relative_keypoints, 6)
npt.assert_array_less(prediction_error, DIFF_THRESHOLD)
if __name__ == '__main__':
absltest.main()
+1 -1
View File
@@ -227,7 +227,7 @@ class FaceMesh(SolutionBase):
image: An RGB image represented as a numpy ndarray.
Raises:
RuntimeError: If the underlying graph occurs any error.
RuntimeError: If the underlying graph throws any error.
ValueError: If the input image is not three channel RGB.
Returns:
+28 -32
View File
@@ -26,7 +26,7 @@ import numpy.testing as npt
from mediapipe.python.solutions import face_mesh as mp_faces
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLOD = 20
DIFF_THRESHOLD = 20 # pixels
EYE_INDICES_TO_LANDMARKS = {
33: [176, 350],
7: [177, 353],
@@ -66,46 +66,42 @@ EYE_INDICES_TO_LANDMARKS = {
class FaceMeshTest(parameterized.TestCase):
def test_invalid_image_shape(self):
faces = mp_faces.FaceMesh()
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
faces.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
with mp_faces.FaceMesh() as faces:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
faces.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
faces = mp_faces.FaceMesh()
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = faces.process(image)
self.assertIsNone(results.multi_face_landmarks)
faces.close()
with mp_faces.FaceMesh() as faces:
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = faces.process(image)
self.assertIsNone(results.multi_face_landmarks)
@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')
faces = mp_faces.FaceMesh(
static_image_mode=static_image_mode, min_detection_confidence=0.5)
image = cv2.flip(cv2.imread(image_path), 1)
def process_one_frame():
results = faces.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
multi_face_landmarks = []
for landmarks in results.multi_face_landmarks:
self.assertLen(landmarks.landmark, 468)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
face_landmarks = np.transpose(np.stack((y, x))) * image.shape[0:2]
multi_face_landmarks.append(face_landmarks)
self.assertLen(multi_face_landmarks, 1)
# Verify the eye landmarks are correct as sanity check.
for idx, gt_lds in EYE_INDICES_TO_LANDMARKS.items():
prediction_error = np.abs(
np.asarray(multi_face_landmarks[0][idx]) - np.asarray(gt_lds))
npt.assert_array_less(prediction_error, DIFF_THRESHOLOD)
for _ in range(num_frames):
process_one_frame()
faces.close()
with mp_faces.FaceMesh(
static_image_mode=static_image_mode,
min_detection_confidence=0.5) as faces:
for _ in range(num_frames):
results = faces.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
multi_face_landmarks = []
for landmarks in results.multi_face_landmarks:
self.assertLen(landmarks.landmark, 468)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
face_landmarks = np.transpose(np.stack((y, x))) * image.shape[0:2]
multi_face_landmarks.append(face_landmarks)
self.assertLen(multi_face_landmarks, 1)
# Verify the eye landmarks are correct as sanity check.
for idx, gt_lds in EYE_INDICES_TO_LANDMARKS.items():
prediction_error = np.abs(
np.asarray(multi_face_landmarks[0][idx]) - np.asarray(gt_lds))
npt.assert_array_less(prediction_error, DIFF_THRESHOLD)
if __name__ == '__main__':
+1 -1
View File
@@ -151,7 +151,7 @@ class Hands(SolutionBase):
image: An RGB image represented as a numpy ndarray.
Raises:
RuntimeError: If the underlying graph occurs any error.
RuntimeError: If the underlying graph throws any error.
ValueError: If the input image is not three channel RGB.
Returns:
+34 -39
View File
@@ -26,7 +26,7 @@ import numpy.testing as npt
from mediapipe.python.solutions import hands as mp_hands
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLOD = 20
DIFF_THRESHOLD = 20 # pixels
EXPECTED_HAND_COORDINATES_PREDICTION = [[[332, 144], [323, 211], [286, 257],
[237, 289], [203, 322], [216, 219],
[138, 238], [90, 249], [51, 253],
@@ -46,53 +46,48 @@ EXPECTED_HAND_COORDINATES_PREDICTION = [[[332, 144], [323, 211], [286, 257],
class HandsTest(parameterized.TestCase):
def test_invalid_image_shape(self):
hands = mp_hands.Hands()
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
hands.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
with mp_hands.Hands() as hands:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
hands.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
hands = mp_hands.Hands()
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = hands.process(image)
self.assertIsNone(results.multi_hand_landmarks)
self.assertIsNone(results.multi_handedness)
hands.close()
with mp_hands.Hands() as hands:
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = hands.process(image)
self.assertIsNone(results.multi_hand_landmarks)
self.assertIsNone(results.multi_handedness)
@parameterized.named_parameters(('static_image_mode', True, 1),
('video_mode', False, 5))
def test_multi_hands(self, static_image_mode, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/hands.jpg')
hands = mp_hands.Hands(
static_image_mode=static_image_mode,
max_num_hands=2,
min_detection_confidence=0.5)
image = cv2.flip(cv2.imread(image_path), 1)
def process_one_frame():
results = hands.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
handedness = [
handedness.classification[0].label
for handedness in results.multi_handedness
]
self.assertLen(handedness, 2)
multi_hand_coordinates = []
for landmarks in results.multi_hand_landmarks:
self.assertLen(landmarks.landmark, 21)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
hand_coordinates = np.transpose(np.stack((y, x))) * image.shape[0:2]
multi_hand_coordinates.append(hand_coordinates)
self.assertLen(multi_hand_coordinates, 2)
prediction_error = np.abs(
np.asarray(multi_hand_coordinates) -
np.asarray(EXPECTED_HAND_COORDINATES_PREDICTION))
npt.assert_array_less(prediction_error, DIFF_THRESHOLOD)
for _ in range(num_frames):
process_one_frame()
hands.close()
with mp_hands.Hands(
static_image_mode=static_image_mode,
max_num_hands=2,
min_detection_confidence=0.5) as hands:
for _ in range(num_frames):
results = hands.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
handedness = [
handedness.classification[0].label
for handedness in results.multi_handedness
]
multi_hand_coordinates = []
for landmarks in results.multi_hand_landmarks:
self.assertLen(landmarks.landmark, 21)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
hand_coordinates = np.transpose(np.stack((y, x))) * image.shape[0:2]
multi_hand_coordinates.append(hand_coordinates)
self.assertLen(handedness, 2)
self.assertLen(multi_hand_coordinates, 2)
prediction_error = np.abs(
np.asarray(multi_hand_coordinates) -
np.asarray(EXPECTED_HAND_COORDINATES_PREDICTION))
npt.assert_array_less(prediction_error, DIFF_THRESHOLD)
if __name__ == '__main__':
+2 -1
View File
@@ -41,6 +41,7 @@ from mediapipe.python.solutions.hands import HAND_CONNECTIONS
from mediapipe.python.solutions.hands import HandLandmark
from mediapipe.python.solutions.pose import POSE_CONNECTIONS
from mediapipe.python.solutions.pose import PoseLandmark
from mediapipe.python.solutions.pose import UPPER_BODY_POSE_CONNECTIONS
# pylint: enable=unused-import
BINARYPB_FILE_PATH = 'mediapipe/modules/holistic_landmark/holistic_landmark_cpu.binarypb'
@@ -111,7 +112,7 @@ class Holistic(SolutionBase):
image: An RGB image represented as a numpy ndarray.
Raises:
RuntimeError: If the underlying graph occurs any error.
RuntimeError: If the underlying graph throws any error.
ValueError: If the input image is not three channel RGB.
Returns:
+92 -92
View File
@@ -13,7 +13,6 @@
# limitations under the License.
"""Tests for mediapipe.python.solutions.pose."""
import math
import os
from absl.testing import absltest
@@ -26,117 +25,118 @@ import numpy.testing as npt
from mediapipe.python.solutions import holistic as mp_holistic
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
POSE_DIFF_THRESHOLOD = 30 # pixels
HAND_DIFF_THRESHOLOD = 10 # pixels
EXPECTED_POSE_COORDINATES_PREDICTION = [[593, 645], [593, 626], [599, 621],
[605, 617], [575, 637], [569, 640],
[563, 643], [621, 616], [565, 652],
[617, 652], [595, 667], [714, 662],
[567, 749], [792, 559], [497, 844],
[844, 435], [407, 906], [866, 403],
[381, 921], [859, 392], [366, 922],
[850, 405], [381, 918], [707, 948],
[631, 940], [582, 1122], [599, 1097],
[495, 1277], [641, 1239], [485, 1300],
[658, 1257], [453, 1332], [626, 1308]]
EXPECTED_LEFT_HAND_COORDINATES_PREDICTION = [[843, 404], [862, 395], [876, 383],
[887, 369], [896, 359], [854, 367],
[868, 347], [879, 346], [885, 349],
[843, 362], [859, 341], [871, 340],
[878, 344], [837, 361], [849, 341],
[859, 338], [867, 339], [834, 361],
[841, 346], [848, 342], [854, 341]]
EXPECTED_RIGHT_HAND_COORDINATES_PREDICTION = [[391, 934], [371,
930], [354, 930],
[340, 934], [328,
939], [350, 938],
[339, 946], [347,
951], [355, 952],
[356, 946], [346,
955], [358, 956],
[366, 953], [361,
952], [354, 959],
[364, 958], [372,
954], [366, 957],
[359, 963], [364, 962],
[368, 960]]
POSE_DIFF_THRESHOLD = 30 # pixels
HAND_DIFF_THRESHOLD = 30 # pixels
EXPECTED_UPPER_BODY_LANDMARKS = np.array([[457, 289], [465, 278], [467, 278],
[470, 277], [461, 279], [461, 279],
[461, 279], [485, 277], [474, 278],
[468, 296], [463, 297], [542, 324],
[449, 327], [614, 321], [376, 318],
[680, 322], [312, 310], [697, 320],
[293, 305], [699, 314], [289, 302],
[693, 316], [296, 305], [515, 451],
[467, 453]])
EXPECTED_FULL_BODY_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]])
EXPECTED_LEFT_HAND_LANDMARKS = np.array([[698, 314], [712, 314], [721, 314],
[727, 314], [732, 313], [728, 309],
[738, 309], [745, 308], [751, 307],
[724, 310], [735, 309], [742, 309],
[747, 307], [719, 312], [727, 313],
[729, 312], [731, 311], [713, 315],
[717, 315], [719, 314], [719, 313]])
EXPECTED_RIGHT_HAND_LANDMARKS = np.array([[293, 307], [284, 306], [277, 304],
[271, 303], [266, 303], [271, 302],
[261, 302], [254, 301], [247, 299],
[272, 303], [261, 303], [253, 301],
[245, 299], [275, 304], [266, 303],
[258, 302], [252, 300], [279, 305],
[273, 305], [268, 304], [263, 303]])
class PoseTest(parameterized.TestCase):
def _verify_output_landmarks(self, landmark_list, image_shape, num_landmarks,
expected_results, diff_thresholds):
self.assertLen(landmark_list.landmark, num_landmarks)
image_rows, image_cols, _ = image_shape
pose_coordinates = [(math.floor(landmark.x * image_cols),
math.floor(landmark.y * image_rows))
for landmark in landmark_list.landmark]
prediction_error = np.abs(
np.asarray(pose_coordinates) -
np.asarray(expected_results[:num_landmarks]))
npt.assert_array_less(prediction_error, diff_thresholds)
def _landmarks_list_to_array(self, landmark_list, image_shape):
rows, cols, _ = image_shape
return np.asarray([(lmk.x * cols, lmk.y * rows)
for lmk in landmark_list.landmark])
def _assert_diff_less(self, array1, array2, threshold):
npt.assert_array_less(np.abs(array1 - array2), threshold)
def test_invalid_image_shape(self):
holistic = mp_holistic.Holistic()
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
holistic.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
with mp_holistic.Holistic() as holistic:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
holistic.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
holistic = mp_holistic.Holistic()
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = holistic.process(image)
self.assertIsNone(results.pose_landmarks)
holistic.close()
with mp_holistic.Holistic() as holistic:
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = holistic.process(image)
self.assertIsNone(results.pose_landmarks)
@parameterized.named_parameters(('static_image_mode', True, 3),
('video_mode', False, 3))
def test_upper_body_model(self, static_image_mode, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/pose.jpg')
holistic = mp_holistic.Holistic(
static_image_mode=static_image_mode, upper_body_only=True)
image = cv2.imread(image_path)
for _ in range(num_frames):
results = holistic.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._verify_output_landmarks(results.pose_landmarks, image.shape, 25,
EXPECTED_POSE_COORDINATES_PREDICTION,
POSE_DIFF_THRESHOLOD)
self._verify_output_landmarks(results.left_hand_landmarks, image.shape,
21,
EXPECTED_LEFT_HAND_COORDINATES_PREDICTION,
HAND_DIFF_THRESHOLOD)
self._verify_output_landmarks(results.right_hand_landmarks, image.shape,
21,
EXPECTED_RIGHT_HAND_COORDINATES_PREDICTION,
HAND_DIFF_THRESHOLOD)
# TODO: Verify the correctness of the face landmarks.
self.assertLen(results.face_landmarks.landmark, 468)
holistic.close()
with mp_holistic.Holistic(
static_image_mode=static_image_mode, upper_body_only=True) as holistic:
image = cv2.imread(image_path)
for _ in range(num_frames):
results = holistic.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_UPPER_BODY_LANDMARKS,
POSE_DIFF_THRESHOLD)
self._assert_diff_less(
self._landmarks_list_to_array(results.left_hand_landmarks,
image.shape),
EXPECTED_LEFT_HAND_LANDMARKS,
HAND_DIFF_THRESHOLD)
self._assert_diff_less(
self._landmarks_list_to_array(results.right_hand_landmarks,
image.shape),
EXPECTED_RIGHT_HAND_LANDMARKS,
HAND_DIFF_THRESHOLD)
# TODO: Verify the correctness of the face landmarks.
self.assertLen(results.face_landmarks.landmark, 468)
@parameterized.named_parameters(('static_image_mode', True, 3),
('video_mode', False, 3))
def test_full_body_model(self, static_image_mode, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/pose.jpg')
holistic = mp_holistic.Holistic(static_image_mode=static_image_mode)
image = cv2.imread(image_path)
for _ in range(num_frames):
results = holistic.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._verify_output_landmarks(results.pose_landmarks, image.shape, 33,
EXPECTED_POSE_COORDINATES_PREDICTION,
POSE_DIFF_THRESHOLOD)
self._verify_output_landmarks(results.left_hand_landmarks, image.shape,
21,
EXPECTED_LEFT_HAND_COORDINATES_PREDICTION,
HAND_DIFF_THRESHOLOD)
self._verify_output_landmarks(results.right_hand_landmarks, image.shape,
21,
EXPECTED_RIGHT_HAND_COORDINATES_PREDICTION,
HAND_DIFF_THRESHOLOD)
# TODO: Verify the correctness of the face landmarks.
self.assertLen(results.face_landmarks.landmark, 468)
holistic.close()
with mp_holistic.Holistic(static_image_mode=static_image_mode) as holistic:
for _ in range(num_frames):
results = holistic.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_FULL_BODY_LANDMARKS,
POSE_DIFF_THRESHOLD)
self._assert_diff_less(
self._landmarks_list_to_array(results.left_hand_landmarks,
image.shape),
EXPECTED_LEFT_HAND_LANDMARKS,
HAND_DIFF_THRESHOLD)
self._assert_diff_less(
self._landmarks_list_to_array(results.right_hand_landmarks,
image.shape),
EXPECTED_RIGHT_HAND_LANDMARKS,
HAND_DIFF_THRESHOLD)
# TODO: Verify the correctness of the face landmarks.
self.assertLen(results.face_landmarks.landmark, 468)
if __name__ == '__main__':
+278
View File
@@ -0,0 +1,278 @@
# 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 Objectron."""
import enum
from typing import List, Tuple, NamedTuple, Optional
import attr
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_detections_calculator_pb2
from mediapipe.calculators.tensor import tensors_to_floats_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 collection_has_min_size_calculator_pb2
from mediapipe.calculators.util import detection_label_id_to_text_calculator_pb2
from mediapipe.calculators.util import detections_to_rects_calculator_pb2
from mediapipe.calculators.util import landmark_projection_calculator_pb2
from mediapipe.calculators.util import local_file_contents_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
from mediapipe.framework.formats import landmark_pb2
from mediapipe.modules.objectron.calculators import annotation_data_pb2
from mediapipe.modules.objectron.calculators import frame_annotation_to_rect_calculator_pb2
from mediapipe.modules.objectron.calculators import lift_2d_frame_annotation_to_3d_calculator_pb2
# pylint: enable=unused-import
from mediapipe.python.solution_base import SolutionBase
class BoxLandmark(enum.IntEnum):
"""The 9 3D box landmarks."""
#
# 3 + + + + + + + + 7
# +\ +\ UP
# + \ + \
# + \ + \ |
# + 4 + + + + + + + + 8 | y
# + + + + |
# + + + + |
# + + (0) + + .------- x
# + + + + \
# 1 + + + + + + + + 5 + \
# \ + \ + \ z
# \ + \ + \
# \+ \+
# 2 + + + + + + + + 6
CENTER = 0
BACK_BOTTOM_LEFT = 1
FRONT_BOTTOM_LEFT = 2
BACK_TOP_LEFT = 3
FRONT_TOP_LEFT = 4
BACK_BOTTOM_RIGHT = 5
FRONT_BOTTOM_RIGHT = 6
BACK_TOP_RIGHT = 7
FRONT_TOP_RIGHT = 8
BINARYPB_FILE_PATH = 'mediapipe/modules/objectron/objectron_cpu.binarypb'
BOX_CONNECTIONS = frozenset([
(BoxLandmark.BACK_BOTTOM_LEFT, BoxLandmark.FRONT_BOTTOM_LEFT),
(BoxLandmark.BACK_BOTTOM_LEFT, BoxLandmark.BACK_TOP_LEFT),
(BoxLandmark.BACK_BOTTOM_LEFT, BoxLandmark.BACK_BOTTOM_RIGHT),
(BoxLandmark.FRONT_BOTTOM_LEFT, BoxLandmark.FRONT_TOP_LEFT),
(BoxLandmark.FRONT_BOTTOM_LEFT, BoxLandmark.FRONT_BOTTOM_RIGHT),
(BoxLandmark.BACK_TOP_LEFT, BoxLandmark.FRONT_TOP_LEFT),
(BoxLandmark.BACK_TOP_LEFT, BoxLandmark.BACK_TOP_RIGHT),
(BoxLandmark.FRONT_TOP_LEFT, BoxLandmark.FRONT_TOP_RIGHT),
(BoxLandmark.BACK_BOTTOM_RIGHT, BoxLandmark.FRONT_BOTTOM_RIGHT),
(BoxLandmark.BACK_BOTTOM_RIGHT, BoxLandmark.BACK_TOP_RIGHT),
(BoxLandmark.FRONT_BOTTOM_RIGHT, BoxLandmark.FRONT_TOP_RIGHT),
(BoxLandmark.BACK_TOP_RIGHT, BoxLandmark.FRONT_TOP_RIGHT),
])
@attr.s(auto_attribs=True)
class ObjectronModel(object):
model_path: str
label_name: str
@attr.s(auto_attribs=True, frozen=True)
class ShoeModel(ObjectronModel):
model_path: str = ('mediapipe/modules/objectron/'
'object_detection_3d_sneakers.tflite')
label_name: str = 'Footwear'
@attr.s(auto_attribs=True, frozen=True)
class ChairModel(ObjectronModel):
model_path: str = ('mediapipe/modules/objectron/'
'object_detection_3d_chair.tflite')
label_name: str = 'Chair'
@attr.s(auto_attribs=True, frozen=True)
class CameraModel(ObjectronModel):
model_path: str = ('mediapipe/modules/objectron/'
'object_detection_3d_camera.tflite')
label_name: str = 'Camera'
@attr.s(auto_attribs=True, frozen=True)
class CupModel(ObjectronModel):
model_path: str = ('mediapipe/modules/objectron/'
'object_detection_3d_cup.tflite')
label_name: str = 'Coffee cup, Mug'
_MODEL_DICT = {
'Shoe': ShoeModel(),
'Chair': ChairModel(),
'Cup': CupModel(),
'Camera': CameraModel()
}
def GetModelByName(name: str) -> ObjectronModel:
if name not in _MODEL_DICT:
raise ValueError(f'{name} is not a valid model name for Objectron.')
return _MODEL_DICT[name]
@attr.s(auto_attribs=True)
class ObjectronOutputs(object):
landmarks_2d: landmark_pb2.NormalizedLandmarkList
landmarks_3d: landmark_pb2.LandmarkList
rotation: np.ndarray
translation: np.ndarray
scale: np.ndarray
class Objectron(SolutionBase):
"""MediaPipe Objectron.
MediaPipe Objectron processes an RGB image and returns the 3D box landmarks
and 2D rectangular bounding box of each detected object.
"""
def __init__(self,
static_image_mode: bool = False,
max_num_objects: int = 5,
min_detection_confidence: float = 0.5,
min_tracking_confidence: float = 0.99,
model_name: str = 'Shoe',
focal_length: Tuple[float, float] = (1.0, 1.0),
principal_point: Tuple[float, float] = (0.0, 0.0),
image_size: Optional[Tuple[int, int]] = None,
):
"""Initializes a MediaPipe Objectron class.
Args:
static_image_mode: Whether to treat the input images as a batch of static
and possibly unrelated images, or a video stream.
max_num_objects: Maximum number of objects to detect.
min_detection_confidence: Minimum confidence value ([0.0, 1.0]) for object
detection to be considered successful.
min_tracking_confidence: Minimum confidence value ([0.0, 1.0]) for the
box landmarks to be considered tracked successfully.
model_name: Name of model to use for predicting box landmarks, currently
support {'Shoe', 'Chair', 'Cup', 'Camera'}.
focal_length: Camera focal length `(fx, fy)`, by default is defined in NDC
space. To use focal length (fx_pixel, fy_pixel) in pixel space, users
should provide image_size = (image_width, image_height) to enable
conversions inside the API.
principal_point: Camera principal point (px, py), by default is defined in
NDC space. To use principal point (px_pixel, py_pixel) in pixel space,
users should provide image_size = (image_width, image_height) to enable
conversions inside the API.
image_size (Optional): size (image_width, image_height) of the input image
, ONLY needed when use focal_length and principal_point in pixel space.
"""
# Get Camera parameters.
fx, fy = focal_length
px, py = principal_point
if image_size is not None:
half_width = image_size[0] / 2.0
half_height = image_size[1] / 2.0
fx = fx / half_width
fy = fy / half_height
px = - (px - half_width) / half_width
py = - (py - half_height) / half_height
# Create and init model.
model = GetModelByName(model_name)
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
side_inputs={
'box_landmark_model_path': model.model_path,
'allowed_labels': model.label_name,
'max_num_objects': max_num_objects,
},
calculator_params={
'ConstantSidePacketCalculator.packet': [
constant_side_packet_calculator_pb2
.ConstantSidePacketCalculatorOptions.ConstantSidePacket(
bool_value=not static_image_mode)
],
('objectdetectionoidv4subgraph'
'__TensorsToDetectionsCalculator.min_score_thresh'):
min_detection_confidence,
('boxlandmarksubgraph__ThresholdingCalculator'
'.threshold'):
min_tracking_confidence,
('Lift2DFrameAnnotationTo3DCalculator'
'.normalized_focal_x'): fx,
('Lift2DFrameAnnotationTo3DCalculator'
'.normalized_focal_y'): fy,
('Lift2DFrameAnnotationTo3DCalculator'
'.normalized_principal_point_x'): px,
('Lift2DFrameAnnotationTo3DCalculator'
'.normalized_principal_point_y'): py,
},
outputs=['detected_objects'])
def process(self, image: np.ndarray) -> NamedTuple:
"""Processes an RGB image and returns the box landmarks and rectangular bounding box of each detected object.
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 "detected_objects" field that contains a list
of detected 3D bounding boxes. Each detected box is represented as an
"ObjectronOutputs" instance.
"""
results = super().process(input_data={'image': image})
if results.detected_objects:
results.detected_objects = self._convert_format(results.detected_objects)
else:
results.detected_objects = None
return results
def _convert_format(
self,
inputs: annotation_data_pb2.FrameAnnotation) -> List[ObjectronOutputs]:
new_outputs = list()
for annotation in inputs.annotations:
# Get 3d object pose.
rotation = np.reshape(np.array(annotation.rotation), (3, 3))
translation = np.array(annotation.translation)
scale = np.array(annotation.scale)
# Get 2d/3d landmakrs.
landmarks_2d = landmark_pb2.NormalizedLandmarkList()
landmarks_3d = landmark_pb2.LandmarkList()
for keypoint in annotation.keypoints:
point_2d = keypoint.point_2d
landmarks_2d.landmark.add(x=point_2d.x, y=point_2d.y)
point_3d = keypoint.point_3d
landmarks_3d.landmark.add(x=point_3d.x, y=point_3d.y, z=point_3d.z)
# Add to objectron outputs.
new_outputs.append(ObjectronOutputs(landmarks_2d, landmarks_3d,
rotation, translation, scale=scale))
return new_outputs
@@ -0,0 +1,81 @@
# 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.
"""Tests for mediapipe.python.solutions.objectron."""
import os
from absl.testing import absltest
from absl.testing import parameterized
import cv2
import numpy as np
import numpy.testing as npt
# resources dependency
from mediapipe.python.solutions import objectron as mp_objectron
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 30 # pixels
EXPECTED_BOX_COORDINATES_PREDICTION = [[[236, 413], [408, 474], [135, 457],
[383, 505], [80, 478], [408, 345],
[130, 347], [384, 355], [72, 353]],
[[241, 206], [411, 279], [131, 280],
[392, 249], [78, 252], [412, 155],
[140, 178], [396, 105], [89, 137]]]
class ObjectronTest(parameterized.TestCase):
def test_invalid_image_shape(self):
with mp_objectron.Objectron() as objectron:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
objectron.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
with mp_objectron.Objectron() as objectron:
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = objectron.process(image)
self.assertIsNone(results.detected_objects)
@parameterized.named_parameters(('static_image_mode', True, 1),
('video_mode', False, 5))
def test_multi_objects(self, static_image_mode, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/shoes.jpg')
image = cv2.imread(image_path)
with mp_objectron.Objectron(
static_image_mode=static_image_mode,
max_num_objects=2,
min_detection_confidence=0.5) as objectron:
for _ in range(num_frames):
results = objectron.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
multi_box_coordinates = []
for detected_object in results.detected_objects:
landmarks = detected_object.landmarks_2d
self.assertLen(landmarks.landmark, 9)
x = [landmark.x for landmark in landmarks.landmark]
y = [landmark.y for landmark in landmarks.landmark]
box_coordinates = np.transpose(np.stack((y, x))) * image.shape[0:2]
multi_box_coordinates.append(box_coordinates)
self.assertLen(multi_box_coordinates, 2)
prediction_error = np.abs(
np.asarray(multi_box_coordinates) -
np.asarray(EXPECTED_BOX_COORDINATES_PREDICTION))
npt.assert_array_less(prediction_error, DIFF_THRESHOLD)
if __name__ == '__main__':
absltest.main()
+16 -13
View File
@@ -77,7 +77,7 @@ class PoseLandmark(enum.IntEnum):
RIGHT_FOOT_INDEX = 32
BINARYPB_FILE_PATH = 'mediapipe/modules/pose_landmark/pose_landmark_cpu.binarypb'
POSE_CONNECTIONS = frozenset([
UPPER_BODY_POSE_CONNECTIONS = frozenset([
(PoseLandmark.NOSE, PoseLandmark.RIGHT_EYE_INNER),
(PoseLandmark.RIGHT_EYE_INNER, PoseLandmark.RIGHT_EYE),
(PoseLandmark.RIGHT_EYE, PoseLandmark.RIGHT_EYE_OUTER),
@@ -103,18 +103,21 @@ POSE_CONNECTIONS = frozenset([
(PoseLandmark.RIGHT_SHOULDER, PoseLandmark.RIGHT_HIP),
(PoseLandmark.LEFT_SHOULDER, PoseLandmark.LEFT_HIP),
(PoseLandmark.RIGHT_HIP, PoseLandmark.LEFT_HIP),
(PoseLandmark.RIGHT_HIP, PoseLandmark.LEFT_HIP),
(PoseLandmark.RIGHT_HIP, PoseLandmark.RIGHT_KNEE),
(PoseLandmark.LEFT_HIP, PoseLandmark.LEFT_KNEE),
(PoseLandmark.RIGHT_KNEE, PoseLandmark.RIGHT_ANKLE),
(PoseLandmark.LEFT_KNEE, PoseLandmark.LEFT_ANKLE),
(PoseLandmark.RIGHT_ANKLE, PoseLandmark.RIGHT_HEEL),
(PoseLandmark.LEFT_ANKLE, PoseLandmark.LEFT_HEEL),
(PoseLandmark.RIGHT_HEEL, PoseLandmark.RIGHT_FOOT_INDEX),
(PoseLandmark.LEFT_HEEL, PoseLandmark.LEFT_FOOT_INDEX),
(PoseLandmark.RIGHT_ANKLE, PoseLandmark.RIGHT_FOOT_INDEX),
(PoseLandmark.LEFT_ANKLE, PoseLandmark.LEFT_FOOT_INDEX),
])
POSE_CONNECTIONS = frozenset.union(
UPPER_BODY_POSE_CONNECTIONS,
frozenset([
(PoseLandmark.RIGHT_HIP, PoseLandmark.RIGHT_KNEE),
(PoseLandmark.LEFT_HIP, PoseLandmark.LEFT_KNEE),
(PoseLandmark.RIGHT_KNEE, PoseLandmark.RIGHT_ANKLE),
(PoseLandmark.LEFT_KNEE, PoseLandmark.LEFT_ANKLE),
(PoseLandmark.RIGHT_ANKLE, PoseLandmark.RIGHT_HEEL),
(PoseLandmark.LEFT_ANKLE, PoseLandmark.LEFT_HEEL),
(PoseLandmark.RIGHT_HEEL, PoseLandmark.RIGHT_FOOT_INDEX),
(PoseLandmark.LEFT_HEEL, PoseLandmark.LEFT_FOOT_INDEX),
(PoseLandmark.RIGHT_ANKLE, PoseLandmark.RIGHT_FOOT_INDEX),
(PoseLandmark.LEFT_ANKLE, PoseLandmark.LEFT_FOOT_INDEX),
]))
class Pose(SolutionBase):
@@ -178,7 +181,7 @@ class Pose(SolutionBase):
image: An RGB image represented as a numpy ndarray.
Raises:
RuntimeError: If the underlying graph occurs any error.
RuntimeError: If the underlying graph throws any error.
ValueError: If the input image is not three channel RGB.
Returns:
+53 -46
View File
@@ -13,7 +13,6 @@
# limitations under the License.
"""Tests for mediapipe.python.solutions.pose."""
import math
import os
from absl.testing import absltest
@@ -26,71 +25,79 @@ import numpy.testing as npt
from mediapipe.python.solutions import pose as mp_pose
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLOD = 30 # pixels
EXPECTED_POSE_COORDINATES_PREDICTION = [[593, 645], [593, 626], [599, 621],
[605, 617], [575, 637], [569, 640],
[563, 643], [621, 616], [565, 652],
[617, 652], [595, 667], [714, 662],
[567, 749], [792, 559], [497, 844],
[844, 435], [407, 906], [866, 403],
[381, 921], [859, 392], [366, 922],
[850, 405], [381, 918], [707, 948],
[631, 940], [582, 1122], [599, 1097],
[495, 1277], [641, 1239], [485, 1300],
[658, 1257], [453, 1332], [626, 1308]]
DIFF_THRESHOLD = 30 # pixels
EXPECTED_UPPER_BODY_LANDMARKS = np.array([[457, 289], [465, 278], [467, 278],
[470, 277], [461, 279], [461, 279],
[461, 279], [485, 277], [474, 278],
[468, 296], [463, 297], [542, 324],
[449, 327], [614, 321], [376, 318],
[680, 322], [312, 310], [697, 320],
[293, 305], [699, 314], [289, 302],
[693, 316], [296, 305], [515, 451],
[467, 453]])
EXPECTED_FULL_BODY_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]])
class PoseTest(parameterized.TestCase):
def _verify_output_landmarks(self, landmark_list, image_shape, num_landmarks):
self.assertLen(landmark_list.landmark, num_landmarks)
image_rows, image_cols, _ = image_shape
pose_coordinates = [(math.floor(landmark.x * image_cols),
math.floor(landmark.y * image_rows))
for landmark in landmark_list.landmark]
prediction_error = np.abs(
np.asarray(pose_coordinates) -
np.asarray(EXPECTED_POSE_COORDINATES_PREDICTION[:num_landmarks]))
npt.assert_array_less(prediction_error, DIFF_THRESHOLOD)
def _landmarks_list_to_array(self, landmark_list, image_shape):
rows, cols, _ = image_shape
return np.asarray([(lmk.x * cols, lmk.y * rows)
for lmk in landmark_list.landmark])
def _assert_diff_less(self, array1, array2, threshold):
npt.assert_array_less(np.abs(array1 - array2), threshold)
def test_invalid_image_shape(self):
pose = mp_pose.Pose()
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
pose.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
with mp_pose.Pose() as pose:
with self.assertRaisesRegex(
ValueError, 'Input image must contain three channel rgb data.'):
pose.process(np.arange(36, dtype=np.uint8).reshape(3, 3, 4))
def test_blank_image(self):
pose = mp_pose.Pose()
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = pose.process(image)
self.assertIsNone(results.pose_landmarks)
pose.close()
with mp_pose.Pose() as pose:
image = np.zeros([100, 100, 3], dtype=np.uint8)
image.fill(255)
results = pose.process(image)
self.assertIsNone(results.pose_landmarks)
@parameterized.named_parameters(('static_image_mode', True, 3),
('video_mode', False, 3))
def test_upper_body_model(self, static_image_mode, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/pose.jpg')
pose = mp_pose.Pose(static_image_mode=static_image_mode,
upper_body_only=True)
image = cv2.imread(image_path)
for _ in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._verify_output_landmarks(results.pose_landmarks, image.shape, 25)
pose.close()
with mp_pose.Pose(
static_image_mode=static_image_mode, upper_body_only=True) as pose:
image = cv2.imread(image_path)
for _ in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_UPPER_BODY_LANDMARKS,
DIFF_THRESHOLD)
@parameterized.named_parameters(('static_image_mode', True, 3),
('video_mode', False, 3))
def test_full_body_model(self, static_image_mode, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/pose.jpg')
pose = mp_pose.Pose(static_image_mode=static_image_mode)
image = cv2.imread(image_path)
for _ in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._verify_output_landmarks(results.pose_landmarks, image.shape, 33)
pose.close()
with mp_pose.Pose(static_image_mode=static_image_mode) as pose:
for _ in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_FULL_BODY_LANDMARKS,
DIFF_THRESHOLD)
if __name__ == '__main__':