Project import generated by Copybara.
GitOrigin-RevId: d073f8e21be2fcc0e503cb97c6695078b6b75310
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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__':
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
@@ -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__':
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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__':
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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__':
|
||||
|
||||
@@ -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()
|
||||
@@ -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:
|
||||
|
||||
@@ -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__':
|
||||
|
||||
Reference in New Issue
Block a user