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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: ff83882955f1a1e2a043ff4e71278be9d7217bbe
This commit is contained in:
MediaPipe Team
2021-05-05 14:56:16 -04:00
committed by chuoling
parent ecb5b5f44a
commit a9b643e0f5
210 changed files with 5312 additions and 3838 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/BUILD
+1
View File
@@ -34,6 +34,7 @@ pybind_extension(
deps = [
":builtin_calculators",
"//mediapipe/python/pybind:calculator_graph",
"//mediapipe/python/pybind:image",
"//mediapipe/python/pybind:image_frame",
"//mediapipe/python/pybind:matrix",
"//mediapipe/python/pybind:packet",
mediapipe/python/__init__.py
+2 -1
View File
@@ -1,4 +1,4 @@
# Copyright 2020 The MediaPipe Authors.
# Copyright 2020-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.
@@ -17,6 +17,7 @@
from mediapipe.python._framework_bindings import resource_util
from mediapipe.python._framework_bindings.calculator_graph import CalculatorGraph
from mediapipe.python._framework_bindings.calculator_graph import GraphInputStreamAddMode
from mediapipe.python._framework_bindings.image import Image
from mediapipe.python._framework_bindings.image_frame import ImageFormat
from mediapipe.python._framework_bindings.image_frame import ImageFrame
from mediapipe.python._framework_bindings.matrix import Matrix
mediapipe/python/framework_bindings.cc
+3 -1
View File
@@ -1,4 +1,4 @@
// Copyright 2020 The MediaPipe Authors.
// Copyright 2020-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.
@@ -13,6 +13,7 @@
// limitations under the License.
#include "mediapipe/python/pybind/calculator_graph.h"
#include "mediapipe/python/pybind/image.h"
#include "mediapipe/python/pybind/image_frame.h"
#include "mediapipe/python/pybind/matrix.h"
#include "mediapipe/python/pybind/packet.h"
@@ -27,6 +28,7 @@ namespace python {
PYBIND11_MODULE(_framework_bindings, m) {
ResourceUtilSubmodule(&m);
ImageSubmodule(&m);
ImageFrameSubmodule(&m);
MatrixSubmodule(&m);
TimestampSubmodule(&m);
mediapipe/python/image_test.py
+183
View File
@@ -0,0 +1,183 @@
# 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._framework_bindings.image."""
import gc
import random
import sys
from absl.testing import absltest
import cv2
import mediapipe as mp
import numpy as np
import PIL.Image
# TODO: Add unit tests specifically for memory management.
class ImageTest(absltest.TestCase):
def test_create_image_from_gray_cv_mat(self):
w, h = random.randrange(3, 100), random.randrange(3, 100)
mat = cv2.cvtColor(
np.random.randint(2**8 - 1, size=(h, w, 3), dtype=np.uint8),
cv2.COLOR_RGB2GRAY)
mat[2, 2] = 42
image = mp.Image(image_format=mp.ImageFormat.GRAY8, data=mat)
self.assertTrue(np.array_equal(mat, image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'index dimension mismatch'):
print(image[w, h, 1])
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[w, h])
self.assertEqual(42, image[2, 2])
def test_create_image_from_rgb_cv_mat(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
mat = cv2.cvtColor(
np.random.randint(2**8 - 1, size=(h, w, channels), dtype=np.uint8),
cv2.COLOR_RGB2BGR)
mat[2, 2, 1] = 42
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=mat)
self.assertTrue(np.array_equal(mat, image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[w, h, channels])
self.assertEqual(42, image[2, 2, 1])
def test_create_image_from_rgb48_cv_mat(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
mat = cv2.cvtColor(
np.random.randint(2**16 - 1, size=(h, w, channels), dtype=np.uint16),
cv2.COLOR_RGB2BGR)
mat[2, 2, 1] = 42
image = mp.Image(image_format=mp.ImageFormat.SRGB48, data=mat)
self.assertTrue(np.array_equal(mat, image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[w, h, channels])
self.assertEqual(42, image[2, 2, 1])
def test_create_image_from_gray_pil_image(self):
w, h = random.randrange(3, 100), random.randrange(3, 100)
img = PIL.Image.fromarray(
np.random.randint(2**8 - 1, size=(h, w), dtype=np.uint8), 'L')
image = mp.Image(image_format=mp.ImageFormat.GRAY8, data=np.asarray(img))
self.assertTrue(np.array_equal(np.asarray(img), image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'index dimension mismatch'):
print(image[w, h, 1])
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[w, h])
def test_create_image_from_rgb_pil_image(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
img = PIL.Image.fromarray(
np.random.randint(2**8 - 1, size=(h, w, channels), dtype=np.uint8),
'RGB')
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np.asarray(img))
self.assertTrue(np.array_equal(np.asarray(img), image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[w, h, channels])
def test_create_image_from_rgba64_pil_image(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 4
img = PIL.Image.fromarray(
np.random.randint(2**16 - 1, size=(h, w, channels), dtype=np.uint16),
'RGBA')
image = mp.Image(
image_format=mp.ImageFormat.SRGBA64,
data=np.asarray(img, dtype=np.uint16))
self.assertTrue(np.array_equal(np.asarray(img), image.numpy_view()))
with self.assertRaisesRegex(IndexError, 'out of bounds'):
print(image[1000, 1000, 1000])
def test_image_numby_view(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
mat = cv2.cvtColor(
np.random.randint(2**8 - 1, size=(h, w, channels), dtype=np.uint8),
cv2.COLOR_RGB2BGR)
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=mat)
output_ndarray = image.numpy_view()
self.assertTrue(np.array_equal(mat, image.numpy_view()))
# The output of numpy_view() is a reference to the internal data and it's
# unwritable after creation.
with self.assertRaisesRegex(ValueError,
'assignment destination is read-only'):
output_ndarray[0, 0, 0] = 0
copied_ndarray = np.copy(output_ndarray)
copied_ndarray[0, 0, 0] = 0
def test_cropped_gray8_image(self):
w, h = random.randrange(20, 100), random.randrange(20, 100)
channels, offset = 3, 10
mat = cv2.cvtColor(
np.random.randint(2**8 - 1, size=(h, w, channels), dtype=np.uint8),
cv2.COLOR_RGB2GRAY)
image = mp.Image(
image_format=mp.ImageFormat.GRAY8,
data=np.ascontiguousarray(mat[offset:-offset, offset:-offset]))
self.assertTrue(
np.array_equal(mat[offset:-offset, offset:-offset], image.numpy_view()))
def test_cropped_rgb_image(self):
w, h = random.randrange(20, 100), random.randrange(20, 100)
channels, offset = 3, 10
mat = cv2.cvtColor(
np.random.randint(2**8 - 1, size=(h, w, channels), dtype=np.uint8),
cv2.COLOR_RGB2BGR)
image = mp.Image(
image_format=mp.ImageFormat.SRGB,
data=np.ascontiguousarray(mat[offset:-offset, offset:-offset, :]))
self.assertTrue(
np.array_equal(mat[offset:-offset, offset:-offset, :],
image.numpy_view()))
# For image frames that store contiguous data, the output of numpy_view()
# points to the pixel data of the original image frame object. The life cycle
# of the data array should tie to the image frame object.
def test_image_numpy_view_with_contiguous_data(self):
w, h = 640, 480
mat = np.random.randint(2**8 - 1, size=(h, w, 3), dtype=np.uint8)
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=mat)
self.assertTrue(image.is_contiguous())
initial_ref_count = sys.getrefcount(image)
self.assertTrue(np.array_equal(mat, image.numpy_view()))
# Get 2 data array objects and verify that the image frame's ref count is
# increased by 2.
np_view = image.numpy_view()
self.assertEqual(sys.getrefcount(image), initial_ref_count + 1)
np_view2 = image.numpy_view()
self.assertEqual(sys.getrefcount(image), initial_ref_count + 2)
del np_view
del np_view2
gc.collect()
# After the two data array objects getting destroyed, the current ref count
# should euqal to the initial ref count.
self.assertEqual(sys.getrefcount(image), initial_ref_count)
# For image frames that store non contiguous data, the output of numpy_view()
# stores a copy of the pixel data of the image frame object. The life cycle of
# the data array doesn't tie to the image frame object.
def test_image_numpy_view_with_non_contiguous_data(self):
w, h = 641, 481
mat = np.random.randint(2**8 - 1, size=(h, w, 3), dtype=np.uint8)
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=mat)
self.assertFalse(image.is_contiguous())
initial_ref_count = sys.getrefcount(image)
self.assertTrue(np.array_equal(mat, image.numpy_view()))
np_view = image.numpy_view()
self.assertEqual(sys.getrefcount(image), initial_ref_count)
del np_view
gc.collect()
self.assertEqual(sys.getrefcount(image), initial_ref_count)
if __name__ == '__main__':
absltest.main()
mediapipe/python/packet_creator.py
+103 -1
View File
@@ -21,6 +21,7 @@ import numpy as np
from google.protobuf import message
from mediapipe.python._framework_bindings import _packet_creator
from mediapipe.python._framework_bindings import image
from mediapipe.python._framework_bindings import image_frame
from mediapipe.python._framework_bindings import packet
@@ -115,8 +116,11 @@ def create_image_frame(data: Union[image_frame.ImageFrame, np.ndarray],
raise ValueError(
'The provided image_format doesn\'t match the one from the data arg.')
if copy is not None and not copy:
# Taking a reference will make the created packet be mutable since the
# ImageFrame object can still be manipulated in Python, which voids packet
# immutability.
raise ValueError(
'Creating image frame packet by taking a reference of another image frame object is not supported yet.'
'Creating ImageFrame packet by taking a reference of another ImageFrame object is not supported yet.'
)
# pylint:disable=protected-access
return _packet_creator._create_image_frame_from_image_frame(data)
@@ -144,6 +148,104 @@ def create_image_frame(data: Union[image_frame.ImageFrame, np.ndarray],
# pylint:enable=protected-access
def create_image(data: Union[image.Image, np.ndarray],
*,
image_format: image_frame.ImageFormat = None,
copy: bool = None) -> packet.Packet:
"""Create a MediaPipe Image packet.
A MediaPipe Image packet can be created from an existing MediaPipe
Image object and the data will be realigned and copied into a new
Image object inside of the packet.
A MediaPipe Image packet can also be created from the raw pixel data
represented as a numpy array with one of the uint8, uint16, and float data
types. There are three data ownership modes depending on how the 'copy' arg
is set.
i) Default mode
If copy is not set, mutable data is always copied while the immutable data
is by reference.
ii) Copy mode (safe)
If copy is set to True, the data will be realigned and copied into an
Image object inside of the packet regardless the immutablity of the
original data.
iii) Reference mode (dangerous)
If copy is set to False, the data will be forced to be shared. If the data is
mutable (data.flags.writeable is True), a warning will be raised.
Args:
data: A MediaPipe Image object or the raw pixel data that is represnted as a
numpy ndarray.
image_format: One of the mp.ImageFormat enum types.
copy: Indicate if the packet should copy the data from the numpy nparray.
Returns:
A MediaPipe Image Packet.
Raises:
ValueError:
i) When "data" is a numpy ndarray, "image_format" is not provided or
the "data" array is not c_contiguous in the reference mode.
ii) When "data" is an Image object, the "image_format" arg doesn't
match the image format of the "data" Image object or "copy" is
explicitly set to False.
TypeError: If "image format" doesn't match "data" array's data type.
Examples:
np_array = np.random.randint(255, size=(321, 123, 3), dtype=np.uint8)
# Copy mode by default if the data array is writable.
image_packet = mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB, data=np_array)
# Make the array unwriteable to trigger the reference mode.
np_array.flags.writeable = False
image_packet = mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB, data=np_array)
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=np_array)
image_packet = mp.packet_creator.create_image(image)
"""
if isinstance(data, image.Image):
if image_format is not None and data.image_format != image_format:
raise ValueError(
'The provided image_format doesn\'t match the one from the data arg.')
if copy is not None and not copy:
# Taking a reference will make the created packet be mutable since the
# Image object can still be manipulated in Python, which voids packet
# immutability.
raise ValueError(
'Creating Image packet by taking a reference of another Image object is not supported yet.'
)
# pylint:disable=protected-access
return _packet_creator._create_image_from_image(data)
# pylint:enable=protected-access
else:
if image_format is None:
raise ValueError('Please provide \'image_format\' with \'data\'.')
# If copy arg is not set, copying the data if it's immutable. Otherwise,
# take a reference of the immutable data to avoid data copy.
if copy is None:
copy = True if data.flags.writeable else False
if not copy:
# TODO: Investigate why the first 2 bytes of the data has data
# corruption when "data" is not c_contiguous.
if not data.flags.c_contiguous:
raise ValueError(
'Reference mode is unavailable if \'data\' is not c_contiguous.')
if data.flags.writeable:
warnings.warn(
'\'data\' is still writeable. Taking a reference of the data to create Image packet is dangerous.',
RuntimeWarning, 2)
# pylint:disable=protected-access
return _packet_creator._create_image_from_pixel_data(
image_format, data, copy)
# pylint:enable=protected-access
def create_proto(proto_message: message.Message) -> packet.Packet:
"""Create a MediaPipe protobuf message packet.
mediapipe/python/packet_getter.py
+1
View File
@@ -33,6 +33,7 @@ get_float_list = _packet_getter.get_float_list
get_str_list = _packet_getter.get_str_list
get_packet_list = _packet_getter.get_packet_list
get_str_to_packet_dict = _packet_getter.get_str_to_packet_dict
get_image = _packet_getter.get_image
get_image_frame = _packet_getter.get_image_frame
get_matrix = _packet_getter.get_matrix
mediapipe/python/packet_test.py
+121 -8
View File
@@ -232,34 +232,46 @@ class PacketTest(absltest.TestCase):
self.assertEqual(mp.packet_getter.get_str(output_list['string']), '42')
self.assertEqual(p.timestamp, 100)
def test_uint8_image_frame_packet(self):
def test_uint8_image_packet(self):
uint8_img = np.random.randint(
2**8 - 1,
size=(random.randrange(3, 100), random.randrange(3, 100), 3),
dtype=np.uint8)
p = mp.packet_creator.create_image_frame(
image_frame_packet = mp.packet_creator.create_image_frame(
mp.ImageFrame(image_format=mp.ImageFormat.SRGB, data=uint8_img))
output_image_frame = mp.packet_getter.get_image_frame(p)
output_image_frame = mp.packet_getter.get_image_frame(image_frame_packet)
self.assertTrue(np.array_equal(output_image_frame.numpy_view(), uint8_img))
image_packet = mp.packet_creator.create_image(
mp.Image(image_format=mp.ImageFormat.SRGB, data=uint8_img))
output_image = mp.packet_getter.get_image(image_packet)
self.assertTrue(np.array_equal(output_image.numpy_view(), uint8_img))
def test_uint16_image_frame_packet(self):
def test_uint16_image_packet(self):
uint16_img = np.random.randint(
2**16 - 1,
size=(random.randrange(3, 100), random.randrange(3, 100), 4),
dtype=np.uint16)
p = mp.packet_creator.create_image_frame(
image_frame_packet = mp.packet_creator.create_image_frame(
mp.ImageFrame(image_format=mp.ImageFormat.SRGBA64, data=uint16_img))
output_image_frame = mp.packet_getter.get_image_frame(p)
output_image_frame = mp.packet_getter.get_image_frame(image_frame_packet)
self.assertTrue(np.array_equal(output_image_frame.numpy_view(), uint16_img))
image_packet = mp.packet_creator.create_image(
mp.Image(image_format=mp.ImageFormat.SRGBA64, data=uint16_img))
output_image = mp.packet_getter.get_image(image_packet)
self.assertTrue(np.array_equal(output_image.numpy_view(), uint16_img))
def test_float_image_frame_packet(self):
float_img = np.float32(
np.random.random_sample(
(random.randrange(3, 100), random.randrange(3, 100), 2)))
p = mp.packet_creator.create_image_frame(
image_frame_packet = mp.packet_creator.create_image_frame(
mp.ImageFrame(image_format=mp.ImageFormat.VEC32F2, data=float_img))
output_image_frame = mp.packet_getter.get_image_frame(p)
output_image_frame = mp.packet_getter.get_image_frame(image_frame_packet)
self.assertTrue(np.allclose(output_image_frame.numpy_view(), float_img))
image_packet = mp.packet_creator.create_image(
mp.Image(image_format=mp.ImageFormat.VEC32F2, data=float_img))
output_image = mp.packet_getter.get_image(image_packet)
self.assertTrue(np.array_equal(output_image.numpy_view(), float_img))
def test_image_frame_packet_creation_copy_mode(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
@@ -362,6 +374,107 @@ class PacketTest(absltest.TestCase):
# copy mode.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
def test_image_packet_creation_copy_mode(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
rgb_data = np.random.randint(255, size=(h, w, channels), dtype=np.uint8)
# rgb_data is c_contiguous.
self.assertTrue(rgb_data.flags.c_contiguous)
initial_ref_count = sys.getrefcount(rgb_data)
p = mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB, data=rgb_data)
# copy mode doesn't increase the ref count of the data.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
rgb_data = rgb_data[:, :, ::-1]
# rgb_data is now not c_contiguous. But, copy mode shouldn't be affected.
self.assertFalse(rgb_data.flags.c_contiguous)
initial_ref_count = sys.getrefcount(rgb_data)
p = mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB, data=rgb_data)
# copy mode doesn't increase the ref count of the data.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
output_image = mp.packet_getter.get_image(p)
self.assertEqual(output_image.height, h)
self.assertEqual(output_image.width, w)
self.assertEqual(output_image.channels, channels)
self.assertTrue(np.array_equal(output_image.numpy_view(), rgb_data))
del p
del output_image
gc.collect()
# Destroying the packet also doesn't affect the ref count becuase of the
# copy mode.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
def test_image_packet_creation_reference_mode(self):
w, h, channels = random.randrange(3, 100), random.randrange(3, 100), 3
rgb_data = np.random.randint(255, size=(h, w, channels), dtype=np.uint8)
rgb_data.flags.writeable = False
initial_ref_count = sys.getrefcount(rgb_data)
image_packet = mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB, data=rgb_data)
# Reference mode increase the ref count of the rgb_data by 1.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count + 1)
del image_packet
gc.collect()
# Deleting image_packet should decrese the ref count of rgb_data by 1.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
rgb_data_copy = np.copy(rgb_data)
# rgb_data_copy is a copy of rgb_data and should not increase the ref count.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
text_config = """
node {
calculator: 'PassThroughCalculator'
input_side_packet: "in"
output_side_packet: "out"
}
"""
graph = mp.CalculatorGraph(graph_config=text_config)
graph.start_run(
input_side_packets={
'in':
mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB, data=rgb_data)
})
# reference mode increase the ref count of the rgb_data by 1.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count + 1)
graph.wait_until_done()
output_packet = graph.get_output_side_packet('out')
del rgb_data
del graph
gc.collect()
# The pixel data of the output image frame packet should still be valid
# after the graph and the original rgb_data data are deleted.
self.assertTrue(
np.array_equal(
mp.packet_getter.get_image(output_packet).numpy_view(),
rgb_data_copy))
def test_image_packet_copy_creation_with_cropping(self):
w, h, channels = random.randrange(40, 100), random.randrange(40, 100), 3
channels, offset = 3, 10
rgb_data = np.random.randint(255, size=(h, w, channels), dtype=np.uint8)
initial_ref_count = sys.getrefcount(rgb_data)
p = mp.packet_creator.create_image(
image_format=mp.ImageFormat.SRGB,
data=rgb_data[offset:-offset, offset:-offset, :])
# copy mode doesn't increase the ref count of the data.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
output_image = mp.packet_getter.get_image(p)
self.assertEqual(output_image.height, h - 2 * offset)
self.assertEqual(output_image.width, w - 2 * offset)
self.assertEqual(output_image.channels, channels)
self.assertTrue(
np.array_equal(rgb_data[offset:-offset, offset:-offset, :],
output_image.numpy_view()))
del p
del output_image
gc.collect()
# Destroying the packet also doesn't affect the ref count becuase of the
# copy mode.
self.assertEqual(sys.getrefcount(rgb_data), initial_ref_count)
def test_matrix_packet(self):
np_matrix = np.array([[.1, .2, .3], [.4, .5, .6]])
initial_ref_count = sys.getrefcount(np_matrix)
mediapipe/python/pybind/BUILD
+16 -1
View File
@@ -1,4 +1,4 @@
# Copyright 2020 The MediaPipe Authors.
# Copyright 2020-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.
@@ -36,6 +36,18 @@ pybind_library(
],
)
pybind_library(
name = "image",
srcs = ["image.cc"],
hdrs = ["image.h"],
deps = [
":image_frame_util",
":util",
"//mediapipe/framework:type_map",
"//mediapipe/framework/formats:image",
],
)
pybind_library(
name = "image_frame",
srcs = ["image_frame.cc"],
@@ -51,6 +63,7 @@ pybind_library(
name = "image_frame_util",
hdrs = ["image_frame_util.h"],
deps = [
":util",
"//mediapipe/framework/formats:image_format_cc_proto",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/port:logging",
@@ -88,6 +101,7 @@ pybind_library(
":util",
"//mediapipe/framework:packet",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/port:integral_types",
"@com_google_absl//absl/memory",
@@ -104,6 +118,7 @@ pybind_library(
":util",
"//mediapipe/framework:packet",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/port:integral_types",
],
mediapipe/python/pybind/calculator_graph.cc
+10 -4
View File
@@ -394,13 +394,15 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
calculator_graph.def(
"observe_output_stream",
[](CalculatorGraph* self, const std::string& stream_name,
pybind11::function callback_fn) {
pybind11::function callback_fn, bool observe_timestamp_bounds) {
RaisePyErrorIfNotOk(self->ObserveOutputStream(
stream_name, [callback_fn, stream_name](const Packet& packet) {
stream_name,
[callback_fn, stream_name](const Packet& packet) {
absl::MutexLock lock(&callback_mutex);
callback_fn(stream_name, packet);
return absl::OkStatus();
}));
},
observe_timestamp_bounds));
},
R"doc(Observe the named output stream.
@@ -411,6 +413,8 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
stream_name: The name of the output stream.
callback_fn: The callback function to invoke on every packet emitted by the
output stream.
observe_timestamp_bounds: If true, emits an empty packet at
timestamp_bound -1 when timestamp bound changes.
Raises:
RuntimeError: If the calculator graph isn't initialized or the stream
@@ -422,7 +426,9 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
graph.observe_output_stream('out',
lambda stream_name, packet: out.append(packet))
)doc");
)doc",
py::arg("stream_name"), py::arg("callback_fn"),
py::arg("observe_timestamp_bounds") = false);
calculator_graph.def(
"close",
mediapipe/python/pybind/image.cc
+234
View File
@@ -0,0 +1,234 @@
// 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.
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/python/pybind/image_frame_util.h"
#include "mediapipe/python/pybind/util.h"
#include "pybind11/stl.h"
namespace mediapipe {
namespace python {
namespace py = pybind11;
void ImageSubmodule(pybind11::module* module) {
py::module m = module->def_submodule("image", "MediaPipe image module");
py::options options;
options.disable_function_signatures();
// Image
py::class_<Image> image(
m, "Image",
R"doc(A container for storing an image or a video frame, in one of several formats.
Formats supported by Image are listed in the ImageFormat enum.
Pixels are encoded row-major in an interleaved fashion. Image supports
uint8, uint16, and float as its data types.
Image can be created by copying the data from a numpy ndarray that stores
the pixel data continuously. An Image may realign the input data on its
default alignment boundary during creation. The data in an Image will
become immutable after creation.
Creation examples:
import cv2
cv_mat = cv2.imread(input_file)[:, :, ::-1]
rgb_frame = mp.Image(format=ImageFormat.SRGB, data=cv_mat)
gray_frame = mp.Image(
format=ImageFormat.GRAY, data=cv2.cvtColor(cv_mat, cv2.COLOR_RGB2GRAY))
from PIL import Image
pil_img = Image.new('RGB', (60, 30), color = 'red')
image = mp.Image(
format=mp.ImageFormat.SRGB, data=np.asarray(pil_img))
The pixel data in an Image can be retrieved as a numpy ndarray by calling
`Image.numpy_view()`. The returned numpy ndarray is a reference to the
internal data and itself is unwritable. If the callers want to modify the
numpy ndarray, it's required to obtain a copy of it.
Pixel data retrieval examples:
for channel in range(num_channel):
for col in range(width):
for row in range(height):
print(image[row, col, channel])
output_ndarray = image.numpy_view()
print(output_ndarray[0, 0, 0])
copied_ndarray = np.copy(output_ndarray)
copied_ndarray[0,0,0] = 0
)doc",
py::dynamic_attr());
image
.def(
py::init([](mediapipe::ImageFormat::Format format,
const py::array_t<uint8, py::array::c_style>& data) {
if (format != mediapipe::ImageFormat::GRAY8 &&
format != mediapipe::ImageFormat::SRGB &&
format != mediapipe::ImageFormat::SRGBA) {
throw RaisePyError(PyExc_RuntimeError,
"uint8 image data should be one of the GRAY8, "
"SRGB, and SRGBA MediaPipe image formats.");
}
return Image(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<uint8>(format, data).release())));
}),
R"doc(For uint8 data type, valid ImageFormat are GRAY8, SGRB, and SRGBA.)doc",
py::arg("image_format"), py::arg("data").noconvert())
.def(
py::init([](mediapipe::ImageFormat::Format format,
const py::array_t<uint16, py::array::c_style>& data) {
if (format != mediapipe::ImageFormat::GRAY16 &&
format != mediapipe::ImageFormat::SRGB48 &&
format != mediapipe::ImageFormat::SRGBA64) {
throw RaisePyError(
PyExc_RuntimeError,
"uint16 image data should be one of the GRAY16, "
"SRGB48, and SRGBA64 MediaPipe image formats.");
}
return Image(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<uint16>(format, data).release())));
}),
R"doc(For uint16 data type, valid ImageFormat are GRAY16, SRGB48, and SRGBA64.)doc",
py::arg("image_format"), py::arg("data").noconvert())
.def(
py::init([](mediapipe::ImageFormat::Format format,
const py::array_t<float, py::array::c_style>& data) {
if (format != mediapipe::ImageFormat::VEC32F1 &&
format != mediapipe::ImageFormat::VEC32F2) {
throw RaisePyError(
PyExc_RuntimeError,
"float image data should be either VEC32F1 or VEC32F2 "
"MediaPipe image formats.");
}
return Image(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<float>(format, data).release())));
}),
R"doc(For float data type, valid ImageFormat are VEC32F1 and VEC32F2.)doc",
py::arg("image_format"), py::arg("data").noconvert());
image.def(
"numpy_view",
[](Image& self) {
py::object py_object =
py::cast(self, py::return_value_policy::reference);
// If the image data is contiguous, generates the data pyarray object
// on demand because 1) making a pyarray by referring to the existing
// image pixel data is relatively cheap and 2) caching the pyarray
// object in an attribute of the image is problematic: the image object
// and the data pyarray object refer to each other, which causes gc
// fails to free the pyarray after use.
// For the non-contiguous cases, gets a cached data pyarray object from
// the image pyobject attribute. This optimization is to avoid the
// expensive data realignment and copy operations happening more than
// once.
return self.GetImageFrameSharedPtr()->IsContiguous()
? GenerateDataPyArrayOnDemand(*self.GetImageFrameSharedPtr(),
py_object)
: GetCachedContiguousDataAttr(*self.GetImageFrameSharedPtr(),
py_object);
},
R"doc(Return the image pixel data as an unwritable numpy ndarray.
Realign the pixel data to be stored contiguously and return a reference to the
unwritable numpy ndarray. If the callers want to modify the numpy array data,
it's required to obtain a copy of the ndarray.
Returns:
An unwritable numpy ndarray.
Examples:
output_ndarray = image.numpy_view()
copied_ndarray = np.copy(output_ndarray)
copied_ndarray[0,0,0] = 0
)doc");
image.def(
"__getitem__",
[](Image& self, const std::vector<int>& pos) {
if (pos.size() != 3 && !(pos.size() == 2 && self.channels() == 1)) {
throw RaisePyError(
PyExc_IndexError,
absl::StrCat("Invalid index dimension: ", pos.size()).c_str());
}
py::object py_object =
py::cast(self, py::return_value_policy::reference);
switch (self.GetImageFrameSharedPtr()->ByteDepth()) {
case 1:
return GetValue<uint8>(*self.GetImageFrameSharedPtr(), pos,
py_object);
case 2:
return GetValue<uint16>(*self.GetImageFrameSharedPtr(), pos,
py_object);
case 4:
return GetValue<float>(*self.GetImageFrameSharedPtr(), pos,
py_object);
default:
return py::object();
}
},
R"doc(Use the indexer operators to access pixel data.
Raises:
IndexError: If the index is invalid or out of bounds.
Examples:
for channel in range(num_channel):
for col in range(width):
for row in range(height):
print(image[row, col, channel])
)doc");
image
.def("uses_gpu", &Image::UsesGpu,
R"doc(Return True if data is currently on the GPU.)doc")
.def(
"is_contiguous",
[](Image& self) {
return self.GetImageFrameSharedPtr()->IsContiguous();
},
R"doc(Return True if the pixel data is stored contiguously (without any alignment padding areas).)doc")
.def(
"is_empty",
[](Image& self) { return self.GetImageFrameSharedPtr()->IsEmpty(); },
R"doc(Return True if the pixel data is unallocated.)doc")
.def(
"is_aligned",
[](Image& self, uint32 alignment_boundary) {
return self.GetImageFrameSharedPtr()->IsAligned(alignment_boundary);
},
R"doc(Return True if each row of the data is aligned to alignment boundary, which must be 1 or a power of 2.
Args:
alignment_boundary: An integer.
Returns:
A boolean.
Examples:
image.is_aligned(16)
)doc");
image.def_property_readonly("width", &Image::width)
.def_property_readonly("height", &Image::height)
.def_property_readonly("channels", &Image::channels)
.def_property_readonly("step", &Image::step)
.def_property_readonly("image_format", &Image::image_format);
}
} // namespace python
} // namespace mediapipe
mediapipe/python/pybind/image.h
+28
View File
@@ -0,0 +1,28 @@
// 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.
#ifndef MEDIAPIPE_PYTHON_PYBIND_IMAGE_H_
#define MEDIAPIPE_PYTHON_PYBIND_IMAGE_H_
#include "pybind11/pybind11.h"
namespace mediapipe {
namespace python {
void ImageSubmodule(pybind11::module* module);
} // namespace python
} // namespace mediapipe
#endif // MEDIAPIPE_PYTHON_PYBIND_IMAGE_H_
mediapipe/python/pybind/image_frame.cc
+1 -113
View File
@@ -1,4 +1,4 @@
// Copyright 2020 The MediaPipe Authors.
// Copyright 2020-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.
@@ -18,118 +18,6 @@
namespace mediapipe {
namespace python {
namespace {
template <typename T>
py::array GenerateContiguousDataArrayHelper(const ImageFrame& image_frame,
const py::object& py_object) {
std::vector<int> shape{image_frame.Height(), image_frame.Width()};
if (image_frame.NumberOfChannels() > 1) {
shape.push_back(image_frame.NumberOfChannels());
}
py::array_t<T, py::array::c_style> contiguous_data;
if (image_frame.IsContiguous()) {
contiguous_data = py::array_t<T, py::array::c_style>(
shape, reinterpret_cast<const T*>(image_frame.PixelData()), py_object);
} else {
auto contiguous_data_copy =
absl::make_unique<T[]>(image_frame.Width() * image_frame.Height() *
image_frame.NumberOfChannels());
image_frame.CopyToBuffer(contiguous_data_copy.get(),
image_frame.PixelDataSizeStoredContiguously());
auto capsule = py::capsule(contiguous_data_copy.get(), [](void* data) {
if (data) {
delete[] reinterpret_cast<T*>(data);
}
});
contiguous_data = py::array_t<T, py::array::c_style>(
shape, contiguous_data_copy.release(), capsule);
}
// In both cases, the underlying data is not writable in Python.
py::detail::array_proxy(contiguous_data.ptr())->flags &=
~py::detail::npy_api::NPY_ARRAY_WRITEABLE_;
return contiguous_data;
}
py::array GenerateContiguousDataArray(const ImageFrame& image_frame,
const py::object& py_object) {
switch (image_frame.ChannelSize()) {
case sizeof(uint8):
return GenerateContiguousDataArrayHelper<uint8>(image_frame, py_object)
.cast<py::array>();
case sizeof(uint16):
return GenerateContiguousDataArrayHelper<uint16>(image_frame, py_object)
.cast<py::array>();
case sizeof(float):
return GenerateContiguousDataArrayHelper<float>(image_frame, py_object)
.cast<py::array>();
break;
default:
throw RaisePyError(PyExc_RuntimeError,
"Unsupported image frame channel size. Data is not "
"uint8, uint16, or float?");
}
}
// Generates a contiguous data pyarray object on demand.
// This function only accepts an image frame object that already stores
// contiguous data. The output py::array points to the raw pixel data array of
// the image frame object directly.
py::array GenerateDataPyArrayOnDemand(const ImageFrame& image_frame,
const py::object& py_object) {
if (!image_frame.IsContiguous()) {
throw RaisePyError(PyExc_RuntimeError,
"GenerateDataPyArrayOnDemand must take an ImageFrame "
"object that stores contiguous data.");
}
return GenerateContiguousDataArray(image_frame, py_object);
}
// Gets the cached contiguous data array from the "__contiguous_data" attribute.
// If the attribute doesn't exist, the function calls
// GenerateContiguousDataArray() to generate the contiguous data pyarray object,
// which realigns and copies the data from the original image frame object.
// Then, the data array object is cached in the "__contiguous_data" attribute.
// This function only accepts an image frame object that stores non-contiguous
// data.
py::array GetCachedContiguousDataAttr(const ImageFrame& image_frame,
const py::object& py_object) {
if (image_frame.IsContiguous()) {
throw RaisePyError(PyExc_RuntimeError,
"GetCachedContiguousDataAttr must take an ImageFrame "
"object that stores non-contiguous data.");
}
py::object get_data_attr =
py::getattr(py_object, "__contiguous_data", py::none());
if (image_frame.IsEmpty()) {
throw RaisePyError(PyExc_RuntimeError, "ImageFrame is unallocated.");
}
// If __contiguous_data attr doesn't store data yet, generates the contiguous
// data array object and caches the result.
if (get_data_attr.is_none()) {
py_object.attr("__contiguous_data") =
GenerateContiguousDataArray(image_frame, py_object);
}
return py_object.attr("__contiguous_data").cast<py::array>();
}
template <typename T>
py::object GetValue(const ImageFrame& image_frame, const std::vector<int>& pos,
const py::object& py_object) {
py::array_t<T, py::array::c_style> output_array =
image_frame.IsContiguous()
? GenerateDataPyArrayOnDemand(image_frame, py_object)
: GetCachedContiguousDataAttr(image_frame, py_object);
if (pos.size() == 2) {
return py::cast(static_cast<T>(output_array.at(pos[0], pos[1])));
} else if (pos.size() == 3) {
return py::cast(static_cast<T>(output_array.at(pos[0], pos[1], pos[2])));
}
return py::none();
}
} // namespace
namespace py = pybind11;
mediapipe/python/pybind/image_frame_util.h
+111 -1
View File
@@ -1,4 +1,4 @@
// Copyright 2020 The MediaPipe Authors.
// Copyright 2020-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.
@@ -20,6 +20,7 @@
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/python/pybind/util.h"
#include "pybind11/numpy.h"
#include "pybind11/pybind11.h"
@@ -57,6 +58,115 @@ std::unique_ptr<ImageFrame> CreateImageFrame(
return image_frame;
}
template <typename T>
py::array GenerateContiguousDataArrayHelper(const ImageFrame& image_frame,
const py::object& py_object) {
std::vector<int> shape{image_frame.Height(), image_frame.Width()};
if (image_frame.NumberOfChannels() > 1) {
shape.push_back(image_frame.NumberOfChannels());
}
py::array_t<T, py::array::c_style> contiguous_data;
if (image_frame.IsContiguous()) {
contiguous_data = py::array_t<T, py::array::c_style>(
shape, reinterpret_cast<const T*>(image_frame.PixelData()), py_object);
} else {
auto contiguous_data_copy =
absl::make_unique<T[]>(image_frame.Width() * image_frame.Height() *
image_frame.NumberOfChannels());
image_frame.CopyToBuffer(contiguous_data_copy.get(),
image_frame.PixelDataSizeStoredContiguously());
auto capsule = py::capsule(contiguous_data_copy.get(), [](void* data) {
if (data) {
delete[] reinterpret_cast<T*>(data);
}
});
contiguous_data = py::array_t<T, py::array::c_style>(
shape, contiguous_data_copy.release(), capsule);
}
// In both cases, the underlying data is not writable in Python.
py::detail::array_proxy(contiguous_data.ptr())->flags &=
~py::detail::npy_api::NPY_ARRAY_WRITEABLE_;
return contiguous_data;
}
inline py::array GenerateContiguousDataArray(const ImageFrame& image_frame,
const py::object& py_object) {
switch (image_frame.ChannelSize()) {
case sizeof(uint8):
return GenerateContiguousDataArrayHelper<uint8>(image_frame, py_object)
.cast<py::array>();
case sizeof(uint16):
return GenerateContiguousDataArrayHelper<uint16>(image_frame, py_object)
.cast<py::array>();
case sizeof(float):
return GenerateContiguousDataArrayHelper<float>(image_frame, py_object)
.cast<py::array>();
break;
default:
throw RaisePyError(PyExc_RuntimeError,
"Unsupported image frame channel size. Data is not "
"uint8, uint16, or float?");
}
}
// Generates a contiguous data pyarray object on demand.
// This function only accepts an image frame object that already stores
// contiguous data. The output py::array points to the raw pixel data array of
// the image frame object directly.
inline py::array GenerateDataPyArrayOnDemand(const ImageFrame& image_frame,
const py::object& py_object) {
if (!image_frame.IsContiguous()) {
throw RaisePyError(PyExc_RuntimeError,
"GenerateDataPyArrayOnDemand must take an ImageFrame "
"object that stores contiguous data.");
}
return GenerateContiguousDataArray(image_frame, py_object);
}
// Gets the cached contiguous data array from the "__contiguous_data" attribute.
// If the attribute doesn't exist, the function calls
// GenerateContiguousDataArray() to generate the contiguous data pyarray object,
// which realigns and copies the data from the original image frame object.
// Then, the data array object is cached in the "__contiguous_data" attribute.
// This function only accepts an image frame object that stores non-contiguous
// data.
inline py::array GetCachedContiguousDataAttr(const ImageFrame& image_frame,
const py::object& py_object) {
if (image_frame.IsContiguous()) {
throw RaisePyError(PyExc_RuntimeError,
"GetCachedContiguousDataAttr must take an ImageFrame "
"object that stores non-contiguous data.");
}
py::object get_data_attr =
py::getattr(py_object, "__contiguous_data", py::none());
if (image_frame.IsEmpty()) {
throw RaisePyError(PyExc_RuntimeError, "ImageFrame is unallocated.");
}
// If __contiguous_data attr doesn't store data yet, generates the contiguous
// data array object and caches the result.
if (get_data_attr.is_none()) {
py_object.attr("__contiguous_data") =
GenerateContiguousDataArray(image_frame, py_object);
}
return py_object.attr("__contiguous_data").cast<py::array>();
}
template <typename T>
py::object GetValue(const ImageFrame& image_frame, const std::vector<int>& pos,
const py::object& py_object) {
py::array_t<T, py::array::c_style> output_array =
image_frame.IsContiguous()
? GenerateDataPyArrayOnDemand(image_frame, py_object)
: GetCachedContiguousDataAttr(image_frame, py_object);
if (pos.size() == 2) {
return py::cast(static_cast<T>(output_array.at(pos[0], pos[1])));
} else if (pos.size() == 3) {
return py::cast(static_cast<T>(output_array.at(pos[0], pos[1], pos[2])));
}
return py::none();
}
} // namespace python
} // namespace mediapipe
mediapipe/python/pybind/packet_creator.cc
+40
View File
@@ -16,6 +16,7 @@
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/integral_types.h"
@@ -49,6 +50,28 @@ Packet CreateImageFramePacket(mediapipe::ImageFormat::Format format,
return Packet();
}
Packet CreateImagePacket(mediapipe::ImageFormat::Format format,
const py::array& data, bool copy) {
if (format == mediapipe::ImageFormat::SRGB ||
format == mediapipe::ImageFormat::SRGBA ||
format == mediapipe::ImageFormat::GRAY8) {
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<uint8>(format, data, copy).release())));
} else if (format == mediapipe::ImageFormat::GRAY16 ||
format == mediapipe::ImageFormat::SRGB48 ||
format == mediapipe::ImageFormat::SRGBA64) {
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<uint16>(format, data, copy).release())));
} else if (format == mediapipe::ImageFormat::VEC32F1 ||
format == mediapipe::ImageFormat::VEC32F2) {
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*CreateImageFrame<float>(format, data, copy).release())));
}
throw RaisePyError(PyExc_RuntimeError,
absl::StrCat("Unsupported ImageFormat: ", format).c_str());
return Packet();
}
} // namespace
namespace py = pybind11;
@@ -586,6 +609,10 @@ void InternalPacketCreators(pybind11::module* m) {
py::arg("format"), py::arg("data").noconvert(), py::arg("copy"),
py::return_value_policy::move);
m->def("_create_image_from_pixel_data", &CreateImagePacket, py::arg("format"),
py::arg("data").noconvert(), py::arg("copy"),
py::return_value_policy::move);
m->def(
"_create_image_frame_from_image_frame",
[](ImageFrame& image_frame) {
@@ -598,6 +625,19 @@ void InternalPacketCreators(pybind11::module* m) {
},
py::arg("image_frame").noconvert(), py::return_value_policy::move);
m->def(
"_create_image_from_image",
[](Image& image) {
auto image_frame_copy = absl::make_unique<ImageFrame>();
// Set alignment_boundary to kGlDefaultAlignmentBoundary so that
// both GPU and CPU can process it.
image_frame_copy->CopyFrom(*image.GetImageFrameSharedPtr(),
ImageFrame::kGlDefaultAlignmentBoundary);
return MakePacket<Image>(std::make_shared<ImageFrame>(
std::move(*image_frame_copy.release())));
},
py::arg("image").noconvert(), py::return_value_policy::move);
m->def(
"_create_proto",
[](const std::string& type_name, const py::bytes& serialized_proto) {
mediapipe/python/pybind/packet_getter.cc
+19
View File
@@ -14,6 +14,7 @@
#include "mediapipe/python/pybind/packet_getter.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/integral_types.h"
@@ -322,6 +323,24 @@ void PublicPacketGetters(pybind11::module* m) {
)doc",
py::return_value_policy::reference_internal);
m->def("get_image", &GetContent<Image>,
R"doc(Get the content of a MediaPipe Image Packet as an Image object.
Args:
packet: A MediaPipe Image Packet.
Returns:
A MediaPipe Image object.
Raises:
ValueError: If the Packet doesn't contain Image.
Examples:
packet = packet_creator.create_image(frame)
data = packet_getter.get_image(packet)
)doc",
py::return_value_policy::reference_internal);
m->def(
"get_matrix",
[](const Packet& packet) {
mediapipe/python/solution_base.py
+30 -9
View File
@@ -89,7 +89,8 @@ class _PacketDataType(enum.Enum):
FLOAT = 'float'
FLOAT_LIST = 'float_list'
AUDIO = 'matrix'
IMAGE = 'image_frame'
IMAGE = 'image'
IMAGE_FRAME = 'image_frame'
PROTO = 'proto'
PROTO_LIST = 'proto_list'
@@ -114,7 +115,7 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
'::mediapipe::Matrix':
_PacketDataType.AUDIO,
'::mediapipe::ImageFrame':
_PacketDataType.IMAGE,
_PacketDataType.IMAGE_FRAME,
'::mediapipe::Classification':
_PacketDataType.PROTO,
'::mediapipe::ClassificationList':
@@ -139,6 +140,8 @@ NAME_TO_TYPE: Mapping[str, '_PacketDataType'] = {
_PacketDataType.PROTO,
'::mediapipe::NormalizedLandmarkList':
_PacketDataType.PROTO,
'::mediapipe::Image':
_PacketDataType.IMAGE,
'::std::vector<::mediapipe::Classification>':
_PacketDataType.PROTO_LIST,
'::std::vector<::mediapipe::ClassificationList>':
@@ -242,7 +245,7 @@ class SolutionBase:
self._graph_outputs[stream_name] = output_packet
for stream_name in self._output_stream_type_info.keys():
self._graph.observe_output_stream(stream_name, callback)
self._graph.observe_output_stream(stream_name, callback, True)
input_side_packets = {
name: self._make_packet(self._side_input_type_info[name], data)
@@ -296,12 +299,14 @@ class SolutionBase:
# input.
self._simulated_timestamp += 33333
for stream_name, data in input_dict.items():
if self._input_stream_type_info[stream_name] == _PacketDataType.IMAGE:
input_stream_type = self._input_stream_type_info[stream_name]
if (input_stream_type == _PacketDataType.IMAGE_FRAME or
input_stream_type == _PacketDataType.IMAGE):
if data.shape[2] != RGB_CHANNELS:
raise ValueError('Input image must contain three channel rgb data.')
self._graph.add_packet_to_input_stream(
stream=stream_name,
packet=self._make_packet(_PacketDataType.IMAGE,
packet=self._make_packet(input_stream_type,
data).at(self._simulated_timestamp))
else:
# TODO: Support audio data.
@@ -476,18 +481,34 @@ class SolutionBase:
def _make_packet(self, packet_data_type: _PacketDataType,
data: Any) -> packet.Packet:
if packet_data_type == _PacketDataType.IMAGE:
return packet_creator.create_image_frame(
if (packet_data_type == _PacketDataType.IMAGE_FRAME or
packet_data_type == _PacketDataType.IMAGE):
return getattr(packet_creator, 'create_' + packet_data_type.value)(
data, image_format=image_frame.ImageFormat.SRGB)
else:
return getattr(packet_creator, 'create_' + packet_data_type.value)(data)
def _get_packet_content(self, packet_data_type: _PacketDataType,
output_packet: packet.Packet) -> Any:
"""Gets packet content from a packet by type.
Args:
packet_data_type: The supported packet data type.
output_packet: The packet to get content from.
Returns:
Packet content by packet data type. None to indicate "no output".
"""
if output_packet.is_empty():
return None
if packet_data_type == _PacketDataType.STRING:
return packet_getter.get_str(output_packet)
elif packet_data_type == _PacketDataType.IMAGE:
return packet_getter.get_image_frame(output_packet).numpy_view()
elif (packet_data_type == _PacketDataType.IMAGE_FRAME or
packet_data_type == _PacketDataType.IMAGE):
return getattr(packet_getter, 'get_' +
packet_data_type.value)(output_packet).numpy_view()
else:
return getattr(packet_getter, 'get_' + packet_data_type.value)(
output_packet)
mediapipe/python/solutions/face_detection_test.py
+17 -6
View File
@@ -14,6 +14,8 @@
"""Tests for mediapipe.python.solutions.face_detection."""
import os
import tempfile # pylint: disable=unused-import
from typing import NamedTuple
from absl.testing import absltest
import cv2
@@ -21,16 +23,25 @@ import numpy as np
import numpy.testing as npt
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import drawing_utils as mp_drawing
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
EXPECTED_FACE_KEY_POINTS = [[182, 363], [186, 460], [241, 420], [284, 417],
[199, 295], [198, 502]]
DIFF_THRESHOLD = 5 # pixels
class FaceDetectionTest(absltest.TestCase):
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
for detection in results.detections:
mp_drawing.draw_detection(frame, detection)
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_image_shape(self):
with mp_faces.FaceDetection() as faces:
with self.assertRaisesRegex(
@@ -46,11 +57,11 @@ class FaceDetectionTest(absltest.TestCase):
def test_face(self):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/face.jpg')
image = cv2.flip(cv2.imread(image_path), 1)
image = cv2.imread(image_path)
with mp_faces.FaceDetection(min_detection_confidence=0.5) as faces:
for _ in range(5):
for idx in range(5):
results = faces.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._annotate(image.copy(), results, idx)
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]
mediapipe/python/solutions/face_mesh_test.py
+53 -38
View File
@@ -15,6 +15,8 @@
"""Tests for mediapipe.python.solutions.face_mesh."""
import os
import tempfile # pylint: disable=unused-import
from typing import NamedTuple
from absl.testing import absltest
from absl.testing import parameterized
@@ -23,48 +25,61 @@ import numpy as np
import numpy.testing as npt
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import face_mesh as mp_faces
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 20 # pixels
DIFF_THRESHOLD = 5 # pixels
EYE_INDICES_TO_LANDMARKS = {
33: [176, 350],
7: [177, 353],
163: [178, 357],
144: [179, 362],
145: [179, 369],
153: [179, 376],
154: [178, 382],
155: [177, 386],
133: [177, 388],
246: [175, 352],
161: [174, 355],
160: [172, 360],
159: [170, 367],
158: [171, 374],
157: [172, 381],
173: [175, 386],
263: [176, 475],
249: [177, 471],
390: [177, 467],
373: [178, 462],
374: [179, 454],
380: [179, 448],
381: [178, 441],
382: [177, 437],
362: [177, 435],
466: [175, 473],
388: [173, 469],
387: [171, 464],
386: [170, 457],
385: [171, 450],
384: [172, 443],
398: [175, 438]
33: [178, 345],
7: [179, 348],
163: [178, 352],
144: [179, 357],
145: [179, 365],
153: [179, 371],
154: [178, 378],
155: [177, 381],
133: [177, 383],
246: [175, 347],
161: [174, 350],
160: [172, 355],
159: [170, 362],
158: [171, 368],
157: [172, 375],
173: [175, 380],
263: [176, 467],
249: [177, 464],
390: [177, 460],
373: [178, 455],
374: [179, 448],
380: [179, 441],
381: [178, 435],
382: [177, 432],
362: [177, 430],
466: [175, 465],
388: [173, 462],
387: [171, 457],
386: [170, 450],
385: [171, 444],
384: [172, 437],
398: [175, 432]
}
class FaceMeshTest(parameterized.TestCase):
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)
for face_landmarks in results.multi_face_landmarks:
mp_drawing.draw_landmarks(
image=frame,
landmark_list=face_landmarks,
landmark_drawing_spec=drawing_spec)
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_image_shape(self):
with mp_faces.FaceMesh() as faces:
with self.assertRaisesRegex(
@@ -82,13 +97,13 @@ class FaceMeshTest(parameterized.TestCase):
('video_mode', False, 5))
def test_face(self, static_image_mode: bool, num_frames: int):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/face.jpg')
image = cv2.flip(cv2.imread(image_path), 1)
image = cv2.imread(image_path)
with mp_faces.FaceMesh(
static_image_mode=static_image_mode,
min_detection_confidence=0.5) as faces:
for _ in range(num_frames):
for idx in range(num_frames):
results = faces.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._annotate(image.copy(), results, idx)
multi_face_landmarks = []
for landmarks in results.multi_face_landmarks:
self.assertLen(landmarks.landmark, 468)
@@ -98,9 +113,9 @@ class FaceMeshTest(parameterized.TestCase):
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():
for eye_idx, gt_lds in EYE_INDICES_TO_LANDMARKS.items():
prediction_error = np.abs(
np.asarray(multi_face_landmarks[0][idx]) - np.asarray(gt_lds))
np.asarray(multi_face_landmarks[0][eye_idx]) - np.asarray(gt_lds))
npt.assert_array_less(prediction_error, DIFF_THRESHOLD)
mediapipe/python/solutions/hands.py
+3 -3
View File
@@ -44,7 +44,7 @@ class HandLandmark(enum.IntEnum):
WRIST = 0
THUMB_CMC = 1
THUMB_MCP = 2
THUMB_IP = 3
THUMB_DIP = 3
THUMB_TIP = 4
INDEX_FINGER_MCP = 5
INDEX_FINGER_PIP = 6
@@ -68,8 +68,8 @@ BINARYPB_FILE_PATH = 'mediapipe/modules/hand_landmark/hand_landmark_tracking_cpu
HAND_CONNECTIONS = frozenset([
(HandLandmark.WRIST, HandLandmark.THUMB_CMC),
(HandLandmark.THUMB_CMC, HandLandmark.THUMB_MCP),
(HandLandmark.THUMB_MCP, HandLandmark.THUMB_IP),
(HandLandmark.THUMB_IP, HandLandmark.THUMB_TIP),
(HandLandmark.THUMB_MCP, HandLandmark.THUMB_DIP),
(HandLandmark.THUMB_DIP, HandLandmark.THUMB_TIP),
(HandLandmark.WRIST, HandLandmark.INDEX_FINGER_MCP),
(HandLandmark.INDEX_FINGER_MCP, HandLandmark.INDEX_FINGER_PIP),
(HandLandmark.INDEX_FINGER_PIP, HandLandmark.INDEX_FINGER_DIP),
mediapipe/python/solutions/hands_test.py
+19 -7
View File
@@ -15,6 +15,8 @@
"""Tests for mediapipe.python.solutions.hands."""
import os
import tempfile # pylint: disable=unused-import
from typing import NamedTuple
from absl.testing import absltest
from absl.testing import parameterized
@@ -23,28 +25,38 @@ import numpy as np
import numpy.testing as npt
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import hands as mp_hands
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 20 # pixels
EXPECTED_HAND_COORDINATES_PREDICTION = [[[332, 144], [323, 211], [286, 257],
DIFF_THRESHOLD = 15 # pixels
EXPECTED_HAND_COORDINATES_PREDICTION = [[[345, 144], [323, 211], [286, 257],
[237, 289], [203, 322], [216, 219],
[138, 238], [90, 249], [51, 253],
[204, 177], [115, 184], [60, 187],
[19, 185], [208, 138], [127, 131],
[77, 124], [36, 117], [222, 106],
[159, 92], [124, 79], [93, 68]],
[[43, 570], [56, 504], [94, 459],
[[40, 577], [56, 504], [94, 459],
[146, 429], [182, 397], [167, 496],
[245, 479], [292, 469], [330, 464],
[177, 540], [265, 534], [319, 533],
[360, 536], [172, 581], [252, 587],
[304, 593], [346, 599], [157, 615],
[219, 628], [255, 638], [288, 648]]]
[223, 628], [258, 638], [288, 648]]]
class HandsTest(parameterized.TestCase):
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
for hand_landmarks in results.multi_hand_landmarks:
mp_drawing.draw_landmarks(frame, hand_landmarks,
mp_hands.HAND_CONNECTIONS)
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_image_shape(self):
with mp_hands.Hands() as hands:
with self.assertRaisesRegex(
@@ -63,14 +75,14 @@ class HandsTest(parameterized.TestCase):
('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')
image = cv2.flip(cv2.imread(image_path), 1)
image = cv2.imread(image_path)
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):
for idx in range(num_frames):
results = hands.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._annotate(image.copy(), results, idx)
handedness = [
handedness.classification[0].label
for handedness in results.multi_handedness
mediapipe/python/solutions/holistic.py
+4 -6
View File
@@ -41,7 +41,6 @@ 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'
@@ -60,7 +59,7 @@ class Holistic(SolutionBase):
def __init__(self,
static_image_mode=False,
upper_body_only=False,
model_complexity=1,
smooth_landmarks=True,
min_detection_confidence=0.5,
min_tracking_confidence=0.5):
@@ -70,9 +69,8 @@ class Holistic(SolutionBase):
static_image_mode: Whether to treat the input images as a batch of static
and possibly unrelated images, or a video stream. See details in
https://solutions.mediapipe.dev/holistic#static_image_mode.
upper_body_only: Whether to track the full set of 33 pose landmarks or
only the 25 upper-body pose landmarks. See details in
https://solutions.mediapipe.dev/holistic#upper_body_only.
model_complexity: Complexity of the pose landmark model: 0, 1 or 2. See
details in https://solutions.mediapipe.dev/holistic#model_complexity.
smooth_landmarks: Whether to filter landmarks across different input
images to reduce jitter. See details in
https://solutions.mediapipe.dev/holistic#smooth_landmarks.
@@ -86,7 +84,7 @@ class Holistic(SolutionBase):
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
side_inputs={
'upper_body_only': upper_body_only,
'model_complexity': model_complexity,
'smooth_landmarks': smooth_landmarks and not static_image_mode,
},
calculator_params={
mediapipe/python/solutions/holistic_test.py
+59 -68
View File
@@ -14,6 +14,8 @@
"""Tests for mediapipe.python.solutions.pose."""
import os
import tempfile # pylint: disable=unused-import
from typing import NamedTuple
from absl.testing import absltest
from absl.testing import parameterized
@@ -22,45 +24,38 @@ import numpy as np
import numpy.testing as npt
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import holistic as mp_holistic
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
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]])
EXPECTED_POSE_LANDMARKS = np.array([[782, 243], [791, 232], [796, 233],
[801, 233], [773, 231], [766, 231],
[759, 232], [802, 242], [751, 239],
[791, 258], [766, 258], [830, 301],
[708, 298], [910, 248], [635, 234],
[954, 161], [593, 136], [961, 137],
[583, 110], [952, 132], [592, 106],
[950, 141], [596, 115], [793, 500],
[724, 502], [874, 626], [640, 629],
[965, 756], [542, 760], [962, 779],
[533, 781], [1025, 797], [487, 803]])
EXPECTED_LEFT_HAND_LANDMARKS = np.array([[958, 167], [950, 161], [945, 151],
[945, 141], [947, 134], [945, 136],
[939, 122], [935, 113], [931, 106],
[951, 134], [946, 118], [942, 108],
[938, 100], [957, 135], [954, 120],
[951, 111], [948, 103], [964, 138],
[964, 128], [965, 122], [965, 117]])
EXPECTED_RIGHT_HAND_LANDMARKS = np.array([[590, 135], [602, 125], [609, 114],
[613, 103], [617, 96], [596, 100],
[595, 84], [594, 74], [593, 68],
[588, 100], [586, 84], [585, 73],
[584, 65], [581, 103], [579, 89],
[579, 79], [579, 72], [575, 109],
[571, 99], [570, 93], [569, 87]])
class PoseTest(parameterized.TestCase):
@@ -73,6 +68,22 @@ class PoseTest(parameterized.TestCase):
def _assert_diff_less(self, array1, array2, threshold):
npt.assert_array_less(np.abs(array1 - array2), threshold)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
drawing_spec = mp_drawing.DrawingSpec(thickness=1, circle_radius=1)
mp_drawing.draw_landmarks(
image=frame,
landmark_list=results.face_landmarks,
landmark_drawing_spec=drawing_spec)
mp_drawing.draw_landmarks(frame, results.left_hand_landmarks,
mp_holistic.HAND_CONNECTIONS)
mp_drawing.draw_landmarks(frame, results.right_hand_landmarks,
mp_holistic.HAND_CONNECTIONS)
mp_drawing.draw_landmarks(frame, results.pose_landmarks,
mp_holistic.POSE_CONNECTIONS)
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_image_shape(self):
with mp_holistic.Holistic() as holistic:
with self.assertRaisesRegex(
@@ -86,44 +97,24 @@ class PoseTest(parameterized.TestCase):
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')
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')
@parameterized.named_parameters(('static_lite', True, 0, 3),
('static_full', True, 1, 3),
('static_heavy', True, 2, 3),
('video_lite', False, 0, 3),
('video_full', False, 1, 3),
('video_heavy', False, 2, 3))
def test_on_image(self, static_image_mode, model_complexity, num_frames):
image_path = os.path.join(os.path.dirname(__file__),
'testdata/holistic.jpg')
image = cv2.imread(image_path)
with mp_holistic.Holistic(static_image_mode=static_image_mode) as holistic:
for _ in range(num_frames):
with mp_holistic.Holistic(static_image_mode=static_image_mode,
model_complexity=model_complexity) as holistic:
for idx in range(num_frames):
results = holistic.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._annotate(image.copy(), results, idx)
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks, image.shape),
EXPECTED_FULL_BODY_LANDMARKS,
EXPECTED_POSE_LANDMARKS,
POSE_DIFF_THRESHOLD)
self._assert_diff_less(
self._landmarks_list_to_array(results.left_hand_landmarks,
mediapipe/python/solutions/pose.py
+16 -21
View File
@@ -42,7 +42,7 @@ from mediapipe.python.solution_base import SolutionBase
class PoseLandmark(enum.IntEnum):
"""The 25 (upper-body) pose landmarks."""
"""The 33 pose landmarks."""
NOSE = 0
LEFT_EYE_INNER = 1
LEFT_EYE = 2
@@ -78,7 +78,7 @@ class PoseLandmark(enum.IntEnum):
RIGHT_FOOT_INDEX = 32
BINARYPB_FILE_PATH = 'mediapipe/modules/pose_landmark/pose_landmark_cpu.binarypb'
UPPER_BODY_POSE_CONNECTIONS = frozenset([
POSE_CONNECTIONS = frozenset([
(PoseLandmark.NOSE, PoseLandmark.RIGHT_EYE_INNER),
(PoseLandmark.RIGHT_EYE_INNER, PoseLandmark.RIGHT_EYE),
(PoseLandmark.RIGHT_EYE, PoseLandmark.RIGHT_EYE_OUTER),
@@ -104,21 +104,17 @@ UPPER_BODY_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.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):
@@ -133,7 +129,7 @@ class Pose(SolutionBase):
def __init__(self,
static_image_mode=False,
upper_body_only=False,
model_complexity=1,
smooth_landmarks=True,
min_detection_confidence=0.5,
min_tracking_confidence=0.5):
@@ -143,9 +139,8 @@ class Pose(SolutionBase):
static_image_mode: Whether to treat the input images as a batch of static
and possibly unrelated images, or a video stream. See details in
https://solutions.mediapipe.dev/pose#static_image_mode.
upper_body_only: Whether to track the full set of 33 pose landmarks or
only the 25 upper-body pose landmarks. See details in
https://solutions.mediapipe.dev/pose#upper_body_only.
model_complexity: Complexity of the pose landmark model: 0, 1 or 2. See
details in https://solutions.mediapipe.dev/pose#model_complexity.
smooth_landmarks: Whether to filter landmarks across different input
images to reduce jitter. See details in
https://solutions.mediapipe.dev/pose#smooth_landmarks.
@@ -159,7 +154,7 @@ class Pose(SolutionBase):
super().__init__(
binary_graph_path=BINARYPB_FILE_PATH,
side_inputs={
'upper_body_only': upper_body_only,
'model_complexity': model_complexity,
'smooth_landmarks': smooth_landmarks and not static_image_mode,
},
calculator_params={
mediapipe/python/solutions/pose_test.py
+42 -46
View File
@@ -16,6 +16,7 @@
import json
import os
import tempfile
from typing import NamedTuple
from absl.testing import absltest
from absl.testing import parameterized
@@ -24,30 +25,23 @@ import numpy as np
import numpy.testing as npt
# resources dependency
# undeclared dependency
from mediapipe.python.solutions import drawing_utils as mp_drawing
from mediapipe.python.solutions import pose as mp_pose
TEST_IMAGE_PATH = 'mediapipe/python/solutions/testdata'
DIFF_THRESHOLD = 30 # pixels
EXPECTED_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_POSE_LANDMARKS = np.array([[460, 287], [469, 277], [472, 276],
[475, 276], [464, 277], [463, 277],
[463, 276], [492, 277], [472, 277],
[471, 295], [465, 295], [542, 323],
[448, 318], [619, 319], [372, 313],
[695, 316], [296, 308], [717, 313],
[273, 304], [718, 304], [280, 298],
[709, 307], [289, 303], [521, 470],
[459, 466], [626, 533], [364, 500],
[704, 616], [347, 614], [710, 631],
[357, 633], [737, 625], [306, 639]])
class PoseTest(parameterized.TestCase):
@@ -60,6 +54,13 @@ class PoseTest(parameterized.TestCase):
def _assert_diff_less(self, array1, array2, threshold):
npt.assert_array_less(np.abs(array1 - array2), threshold)
def _annotate(self, frame: np.ndarray, results: NamedTuple, idx: int):
mp_drawing.draw_landmarks(frame, results.pose_landmarks,
mp_pose.POSE_CONNECTIONS)
path = os.path.join(tempfile.gettempdir(), self.id().split('.')[-1] +
'_frame_{}.png'.format(idx))
cv2.imwrite(path, frame)
def test_invalid_image_shape(self):
with mp_pose.Pose() as pose:
with self.assertRaisesRegex(
@@ -73,38 +74,28 @@ class PoseTest(parameterized.TestCase):
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')
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)[:, :2],
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):
@parameterized.named_parameters(('static_lite', True, 0, 3),
('static_full', True, 1, 3),
('static_heavy', True, 2, 3),
('video_lite', False, 0, 3),
('video_full', False, 1, 3),
('video_heavy', False, 2, 3))
def test_on_image(self, static_image_mode, model_complexity, num_frames):
image_path = os.path.join(os.path.dirname(__file__), 'testdata/pose.jpg')
image = cv2.imread(image_path)
with mp_pose.Pose(static_image_mode=static_image_mode) as pose:
for _ in range(num_frames):
with mp_pose.Pose(static_image_mode=static_image_mode,
model_complexity=model_complexity) as pose:
for idx in range(num_frames):
results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
self._annotate(image.copy(), results, idx)
self._assert_diff_less(
self._landmarks_list_to_array(results.pose_landmarks,
image.shape)[:, :2],
EXPECTED_FULL_BODY_LANDMARKS, DIFF_THRESHOLD)
EXPECTED_POSE_LANDMARKS, DIFF_THRESHOLD)
@parameterized.named_parameters(
('full_body', False, 'pose_squats.full_body.npz'),
('upper_body', True, 'pose_squats.upper_body.npz'))
def test_on_video(self, upper_body_only, expected_name):
('full', 1, 'pose_squats.full.npz'))
def test_on_video(self, model_complexity, expected_name):
"""Tests pose models on a video."""
# If set to `True` will dump actual predictions to .npz and JSON files.
dump_predictions = False
@@ -120,8 +111,9 @@ class PoseTest(parameterized.TestCase):
# Predict pose landmarks for each frame.
video_cap = cv2.VideoCapture(video_path)
actual_per_frame = []
with mp_pose.Pose(
static_image_mode=False, upper_body_only=upper_body_only) as pose:
frame_idx = 0
with mp_pose.Pose(static_image_mode=False,
model_complexity=model_complexity) as pose:
while True:
# Get next frame of the video.
success, input_frame = video_cap.read()
@@ -135,6 +127,10 @@ class PoseTest(parameterized.TestCase):
input_frame.shape)
actual_per_frame.append(pose_landmarks)
input_frame = cv2.cvtColor(input_frame, cv2.COLOR_RGB2BGR)
self._annotate(input_frame, result, frame_idx)
frame_idx += 1
actual = np.asarray(actual_per_frame)
if dump_predictions: