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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 4cee4a2c2317fb190680c17e31ebbb03bb73b71c
This commit is contained in:
MediaPipe Team
2020-09-17 11:09:17 -04:00
committed by chuoling
parent 1db91b550a
commit a908d668c7
142 changed files with 40878 additions and 574 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/python/BUILD
+1
View File
@@ -49,5 +49,6 @@ pybind_extension(
"//mediapipe/python/pybind:packet_getter",
"//mediapipe/python/pybind:resource_util",
"//mediapipe/python/pybind:timestamp",
"//mediapipe/python/pybind:validated_graph_config",
],
)
mediapipe/python/__init__.py
+1
View File
@@ -22,5 +22,6 @@ from mediapipe.python._framework_bindings.image_frame import ImageFrame
from mediapipe.python._framework_bindings.matrix import Matrix
from mediapipe.python._framework_bindings.packet import Packet
from mediapipe.python._framework_bindings.timestamp import Timestamp
from mediapipe.python._framework_bindings.validated_graph_config import ValidatedGraphConfig
import mediapipe.python.packet_creator
import mediapipe.python.packet_getter
mediapipe/python/calculator_graph_test.py
+37
View File
@@ -127,6 +127,43 @@ class GraphTest(absltest.TestCase):
self.assertEqual(mp.packet_getter.get_str(out[0]), 'hello world')
self.assertEqual(mp.packet_getter.get_str(out[1]), 'hello world')
def testGraphValidationAndInitialization(self):
text_config = """
max_queue_size: 1
input_stream: 'in'
output_stream: 'out'
node {
calculator: 'PassThroughCalculator'
input_stream: 'in'
output_stream: 'out'
}
"""
hello_world_packet = mp.packet_creator.create_string('hello world')
out = []
validated_graph_config = mp.ValidatedGraphConfig()
self.assertFalse(validated_graph_config.initialized())
validated_graph_config.initialize(graph_config=text_config)
self.assertTrue(validated_graph_config.initialized())
graph = mp.CalculatorGraph(validated_graph_config=validated_graph_config)
graph.observe_output_stream('out', lambda _, packet: out.append(packet))
graph.start_run()
graph.add_packet_to_input_stream(
stream='in', packet=hello_world_packet.at(0))
graph.add_packet_to_input_stream(
stream='in', packet=hello_world_packet, timestamp=1)
graph.close()
self.assertEqual(graph.graph_input_stream_add_mode,
mp.GraphInputStreamAddMode.WAIT_TILL_NOT_FULL)
self.assertEqual(graph.max_queue_size, 1)
self.assertFalse(graph.has_error())
self.assertLen(out, 2)
self.assertEqual(out[0].timestamp, 0)
self.assertEqual(out[1].timestamp, 1)
self.assertEqual(mp.packet_getter.get_str(out[0]), 'hello world')
self.assertEqual(mp.packet_getter.get_str(out[1]), 'hello world')
def testInsertPacketsWithSameTimestamp(self):
text_config = """
max_queue_size: 1
mediapipe/python/framework_bindings.cc
+2
View File
@@ -20,6 +20,7 @@
#include "mediapipe/python/pybind/packet_getter.h"
#include "mediapipe/python/pybind/resource_util.h"
#include "mediapipe/python/pybind/timestamp.h"
#include "mediapipe/python/pybind/validated_graph_config.h"
namespace mediapipe {
namespace python {
@@ -33,6 +34,7 @@ PYBIND11_MODULE(_framework_bindings, m) {
PacketCreatorSubmodule(&m);
PacketGetterSubmodule(&m);
CalculatorGraphSubmodule(&m);
ValidatedGraphConfigSubmodule(&m);
}
} // namespace python
mediapipe/python/image_frame_test.py
+41
View File
@@ -14,7 +14,9 @@
"""Tests for mediapipe.python._framework_bindings.image_frame."""
import gc
import random
import sys
from absl.testing import absltest
import cv2
import mediapipe as mp
@@ -140,6 +142,45 @@ class ImageFrameTest(absltest.TestCase):
np.array_equal(mat[offset:-offset, offset:-offset, :],
image_frame.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 testImageFrameNumpyViewWithContiguousData(self):
w, h = 640, 480
mat = np.random.randint(2**8 - 1, size=(h, w, 3), dtype=np.uint8)
image_frame = mp.ImageFrame(image_format=mp.ImageFormat.SRGB, data=mat)
self.assertTrue(image_frame.is_contiguous())
initial_ref_count = sys.getrefcount(image_frame)
self.assertTrue(np.array_equal(mat, image_frame.numpy_view()))
# Get 2 data array objects and verify that the image frame's ref count is
# increased by 2.
np_view = image_frame.numpy_view()
self.assertEqual(sys.getrefcount(image_frame), initial_ref_count + 1)
np_view2 = image_frame.numpy_view()
self.assertEqual(sys.getrefcount(image_frame), 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_frame), 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 testImageFrameNumpyViewWithNonContiguousData(self):
w, h = 641, 481
mat = np.random.randint(2**8 - 1, size=(h, w, 3), dtype=np.uint8)
image_frame = mp.ImageFrame(image_format=mp.ImageFormat.SRGB, data=mat)
self.assertFalse(image_frame.is_contiguous())
initial_ref_count = sys.getrefcount(image_frame)
self.assertTrue(np.array_equal(mat, image_frame.numpy_view()))
np_view = image_frame.numpy_view()
self.assertEqual(sys.getrefcount(image_frame), initial_ref_count)
del np_view
gc.collect()
self.assertEqual(sys.getrefcount(image_frame), initial_ref_count)
if __name__ == '__main__':
absltest.main()
mediapipe/python/pybind/BUILD
+13 -1
View File
@@ -27,7 +27,6 @@ pybind_library(
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_graph",
"//mediapipe/framework:packet",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:map_util",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
@@ -135,7 +134,20 @@ pybind_library(
name = "util",
hdrs = ["util.h"],
deps = [
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:timestamp",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:status",
],
)
pybind_library(
name = "validated_graph_config",
srcs = ["validated_graph_config.cc"],
hdrs = ["validated_graph_config.h"],
deps = [
":util",
"//mediapipe/framework:validated_graph_config",
"//mediapipe/framework/port:parse_text_proto",
],
)
mediapipe/python/pybind/calculator_graph.cc
+15 -26
View File
@@ -19,7 +19,6 @@
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_graph.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/map_util.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"
@@ -70,39 +69,25 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
// TODO: Support graph initialization with graph templates and
// subgraph.
calculator_graph.def(
py::init([](py::args args, py::kwargs kwargs) {
if (!args.empty()) {
throw RaisePyError(PyExc_RuntimeError,
"Invalid position input arguments.");
}
py::init([](py::kwargs kwargs) {
bool init_with_binary_graph = false;
bool init_with_graph_proto = false;
bool init_with_validated_graph_config = false;
CalculatorGraphConfig graph_config_proto;
for (const auto& kw : kwargs) {
const std::string& key = kw.first.cast<std::string>();
if (key == "binary_graph_path") {
init_with_binary_graph = true;
std::string file_name(kw.second.cast<py::object>().str());
auto status = file::Exists(file_name);
if (!status.ok()) {
throw RaisePyError(PyExc_FileNotFoundError,
status.message().data());
}
std::string graph_config_string;
RaisePyErrorIfNotOk(
file::GetContents(file_name, &graph_config_string));
if (!graph_config_proto.ParseFromArray(
graph_config_string.c_str(),
graph_config_string.length())) {
throw RaisePyError(
PyExc_RuntimeError,
absl::StrCat("Failed to parse the binary graph: ", file_name)
.c_str());
}
graph_config_proto = ReadCalculatorGraphConfigFromFile(file_name);
} else if (key == "graph_config") {
init_with_graph_proto = true;
graph_config_proto =
ParseProto<CalculatorGraphConfig>(kw.second.cast<py::object>());
} else if (key == "validated_graph_config") {
init_with_validated_graph_config = true;
graph_config_proto =
py::cast<ValidatedGraphConfig*>(kw.second)->Config();
} else {
throw RaisePyError(
PyExc_RuntimeError,
@@ -110,12 +95,15 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
}
}
if (!(init_with_binary_graph ^ init_with_graph_proto)) {
if ((init_with_binary_graph ? 1 : 0) + (init_with_graph_proto ? 1 : 0) +
(init_with_validated_graph_config ? 1 : 0) !=
1) {
throw RaisePyError(
PyExc_ValueError,
"Please provide \'binary_graph\' to initialize the graph with"
" binary graph or provide \'graph_config\' to initialize the "
" with graph config proto.");
" with graph config proto or provide \'validated_graph_config\' "
" to initialize the with ValidatedGraphConfig object.");
}
auto calculator_graph = absl::make_unique<CalculatorGraph>();
RaisePyErrorIfNotOk(calculator_graph->Initialize(graph_config_proto));
@@ -127,6 +115,7 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
binary_graph_path: The path to a binary mediapipe graph file (.binarypb).
graph_config: A single CalculatorGraphConfig proto message or its text proto
format.
validated_graph_config: A ValidatedGraphConfig object.
Raises:
FileNotFoundError: If the binary graph file can't be found.
@@ -136,11 +125,11 @@ void CalculatorGraphSubmodule(pybind11::module* module) {
// TODO: Return a Python CalculatorGraphConfig instead.
calculator_graph.def_property_readonly(
"config",
"text_config",
[](const CalculatorGraph& self) { return self.Config().DebugString(); });
calculator_graph.def_property_readonly(
"serialized_config", [](const CalculatorGraph& self) {
"binary_config", [](const CalculatorGraph& self) {
return py::bytes(self.Config().SerializeAsString());
});
mediapipe/python/pybind/image_frame.cc
+68 -27
View File
@@ -21,19 +21,16 @@ namespace python {
namespace {
template <typename T>
py::array GenerateContiguousDataArray(const ImageFrame& image_frame,
const py::object& py_object) {
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()) {
// TODO: Create contiguous_data without copying ata.
// It's possible to achieve this with the help of py::capsule.
// Reference: https://github.com/pybind/pybind11/issues/1042,
contiguous_data = py::array_t<T, py::array::c_style>(
shape, reinterpret_cast<const T*>(image_frame.PixelData()));
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() *
@@ -55,35 +52,65 @@ py::array GenerateContiguousDataArray(const ImageFrame& image_frame,
return contiguous_data;
}
py::array GetContiguousDataAttr(const ImageFrame& image_frame,
const py::object& py_object) {
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 already stores data, return the cached results.
if (!get_data_attr.is_none()) {
return get_data_attr.cast<py::array>();
}
py::array GenerateContiguousDataArray(const ImageFrame& image_frame,
const py::object& py_object) {
switch (image_frame.ChannelSize()) {
case sizeof(uint8):
py_object.attr("__contiguous_data") =
GenerateContiguousDataArray<uint8>(image_frame, py_object);
break;
return GenerateContiguousDataArrayHelper<uint8>(image_frame, py_object)
.cast<py::array>();
case sizeof(uint16):
py_object.attr("__contiguous_data") =
GenerateContiguousDataArray<uint16>(image_frame, py_object);
break;
return GenerateContiguousDataArrayHelper<uint16>(image_frame, py_object)
.cast<py::array>();
case sizeof(float):
py_object.attr("__contiguous_data") =
GenerateContiguousDataArray<float>(image_frame, py_object);
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>();
}
@@ -91,7 +118,9 @@ 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 =
GetContiguousDataAttr(image_frame, py_object);
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) {
@@ -243,7 +272,19 @@ void ImageFrameSubmodule(pybind11::module* module) {
[](ImageFrame& self) {
py::object py_object =
py::cast(self, py::return_value_policy::reference);
return GetContiguousDataAttr(self, py_object);
// If the image frame data is contiguous, generates the data pyarray
// object on demand because 1) making a pyarray by referring to the
// existing image frame pixel data is relatively cheap and 2) caching
// the pyarray object in an attribute of the image frame is problematic:
// the image frame 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 frame pyobject attribute. This optimization is to avoid the
// expensive data realignment and copy operations happening more than
// once.
return self.IsContiguous()
? GenerateDataPyArrayOnDemand(self, py_object)
: GetCachedContiguousDataAttr(self, py_object);
},
R"doc(Return the image frame pixel data as an unwritable numpy ndarray.
mediapipe/python/pybind/util.h
+21
View File
@@ -15,6 +15,8 @@
#ifndef MEDIAPIPE_PYTHON_PYBIND_UTIL_H_
#define MEDIAPIPE_PYTHON_PYBIND_UTIL_H_
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/timestamp.h"
#include "pybind11/pybind11.h"
@@ -86,6 +88,25 @@ inline std::string TimestampValueString(const Timestamp& timestamp) {
}
}
// Reads a CalculatorGraphConfig from a file. If failed, raises a PyError.
inline ::mediapipe::CalculatorGraphConfig ReadCalculatorGraphConfigFromFile(
const std::string& file_name) {
::mediapipe::CalculatorGraphConfig graph_config_proto;
auto status = file::Exists(file_name);
if (!status.ok()) {
throw RaisePyError(PyExc_FileNotFoundError, status.message().data());
}
std::string graph_config_string;
RaisePyErrorIfNotOk(file::GetContents(file_name, &graph_config_string));
if (!graph_config_proto.ParseFromArray(graph_config_string.c_str(),
graph_config_string.length())) {
throw RaisePyError(
PyExc_RuntimeError,
absl::StrCat("Failed to parse the binary graph: ", file_name).c_str());
}
return graph_config_proto;
}
} // namespace python
} // namespace mediapipe
mediapipe/python/pybind/validated_graph_config.cc
+153
View File
@@ -0,0 +1,153 @@
// 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.
#include "mediapipe/python/pybind/validated_graph_config.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/validated_graph_config.h"
#include "mediapipe/python/pybind/util.h"
namespace mediapipe {
namespace python {
namespace py = pybind11;
void ValidatedGraphConfigSubmodule(pybind11::module* module) {
py::module m = module->def_submodule(
"validated_graph_config", "MediaPipe validated graph config module.");
// Validated Graph Config
py::class_<ValidatedGraphConfig> validated_graph_config(
m, "ValidatedGraphConfig",
R"doc(A class to validate and canonicalize a CalculatorGraphConfig.)doc");
validated_graph_config.def(py::init())
.def(
"initialize",
[](ValidatedGraphConfig* self, py::kwargs kwargs) {
bool init_with_binary_graph = false;
bool init_with_graph_proto = false;
CalculatorGraphConfig graph_config_proto;
for (const auto& kw : kwargs) {
const std::string& key = kw.first.cast<std::string>();
if (key == "binary_graph_path") {
init_with_binary_graph = true;
std::string file_name(kw.second.cast<py::object>().str());
graph_config_proto =
ReadCalculatorGraphConfigFromFile(file_name);
} else if (key == "graph_config") {
init_with_graph_proto = true;
if (!ParseTextProto<CalculatorGraphConfig>(
kw.second.cast<py::object>().str(),
&graph_config_proto)) {
throw RaisePyError(
PyExc_RuntimeError,
absl::StrCat(
"Failed to parse: ",
std::string(kw.second.cast<py::object>().str()))
.c_str());
}
} else {
throw RaisePyError(
PyExc_RuntimeError,
absl::StrCat("Unknown kwargs input argument: ", key)
.c_str());
}
}
RaisePyErrorIfNotOk(self->Initialize(graph_config_proto));
},
R"doc(Initialize ValidatedGraphConfig with a CalculatorGraphConfig.
Args:
binary_graph_path: The path to a binary mediapipe graph file (.binarypb).
graph_config: A single CalculatorGraphConfig proto message or its text proto
format.
Raises:
FileNotFoundError: If the binary graph file can't be found.
ValueError: If the input arguments prvoided are more than needed or the
graph validation process contains error.
Examples:
validated_graph_config = mp.ValidatedGraphConfig()
validated_graph_config.initialize(graph_config=text_config)
)doc");
validated_graph_config.def(
"registered_stream_type_name",
[](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();
},
R"doc(Return the registered type name of the specified stream if it can be determined.
Args:
stream_name: The input/output stream name.
Returns:
The registered packet type name of the input/output stream.
Raises:
ValueError: If the input/output stream cannot be found.
Examples:
validated_graph_config.registered_stream_type_name('stream_name')
)doc");
validated_graph_config.def(
"registered_side_packet_type_name",
[](ValidatedGraphConfig& self, const std::string& side_packet_name) {
auto status_or_type_name =
self.RegisteredSidePacketTypeName(side_packet_name);
RaisePyErrorIfNotOk(status_or_type_name.status());
return status_or_type_name.ValueOrDie();
},
R"doc(Return the registered type name of the specified side packet if it can be determined.
Args:
side_packet_name: The input/output side packet name.
Returns:
The registered packet type name of the input/output side packet.
Raises:
ValueError: If the input/output side packet cannot be found.
Examples:
validated_graph_config.registered_side_packet_type_name('side_packet')
)doc");
// TODO: Return a Python CalculatorGraphConfig instead.
validated_graph_config.def_property_readonly(
"text_config", [](const ValidatedGraphConfig& self) {
return self.Config().DebugString();
});
validated_graph_config.def_property_readonly(
"binary_config", [](const ValidatedGraphConfig& self) {
return py::bytes(self.Config().SerializeAsString());
});
validated_graph_config.def(
"initialized",
[](const ValidatedGraphConfig& self) { return self.Initialized(); },
R"doc(Indicate if ValidatedGraphConfig is initialized.)doc");
}
} // namespace python
} // namespace mediapipe
mediapipe/python/pybind/validated_graph_config.h
+28
View File
@@ -0,0 +1,28 @@
// 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.
#ifndef MEDIAPIPE_PYTHON_PYBIND_VALIDATED_GRAPH_CONFIG_H_
#define MEDIAPIPE_PYTHON_PYBIND_VALIDATED_GRAPH_CONFIG_H_
#include "pybind11/pybind11.h"
namespace mediapipe {
namespace python {
void ValidatedGraphConfigSubmodule(pybind11::module* module);
} // namespace python
} // namespace mediapipe
#endif // MEDIAPIPE_PYTHON_PYBIND_VALIDATED_GRAPH_CONFIG_H_