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

Project import generated by Copybara.

Browse Source 
PiperOrigin-RevId: 263889205
This commit is contained in:
MediaPipe Team
2019-08-16 18:56:48 -07:00
committed by jqtang
parent dc40414468
commit 294687295d
443 changed files with 33160 additions and 2011 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/tool/BUILD
+122 -3
View File
@@ -17,18 +17,57 @@ licenses(["notice"]) # Apache 2.0
package(default_visibility = ["//visibility:private"])
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
load(
"//mediapipe/framework/port:build_config.bzl",
"mediapipe_cc_proto_library",
)
load(
"//mediapipe/framework/tool:mediapipe_graph.bzl",
"data_as_c_string",
"mediapipe_binary_graph",
)
exports_files([
"simple_subgraph_template.cc",
])
cc_library(
name = "text_to_binary_graph",
srcs = ["text_to_binary_graph.cc"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:commandlineflags",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
],
)
proto_library(
name = "calculator_graph_template_proto",
srcs = ["calculator_graph_template.proto"],
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:calculator_options_proto",
"//mediapipe/framework:calculator_proto",
"//mediapipe/framework/deps:proto_descriptor_proto",
],
)
java_proto_library(
name = "calculator_graph_template_java_proto",
visibility = ["//visibility:public"],
deps = [":calculator_graph_template_proto"],
)
java_lite_proto_library(
name = "calculator_graph_template_java_proto_lite",
strict_deps = 0,
visibility = ["//visibility:public"],
deps = [":calculator_graph_template_proto"],
)
proto_library(
name = "source_proto",
srcs = ["source.proto"],
@@ -43,7 +82,10 @@ mediapipe_cc_proto_library(
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework/deps:proto_descriptor_cc_proto",
],
visibility = ["//mediapipe/framework:__subpackages__"],
visibility = [
"//mediapipe/framework:__subpackages__",
"//mediapipe/java/com/google/mediapipe/framework:__subpackages__",
],
deps = [":calculator_graph_template_proto"],
)
@@ -55,6 +97,15 @@ mediapipe_cc_proto_library(
deps = [":source_proto"],
)
cc_binary(
name = "encode_as_c_string",
srcs = ["encode_as_c_string.cc"],
visibility = ["//visibility:public"],
deps = [
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "fill_packet_set",
srcs = ["fill_packet_set.cc"],
@@ -464,8 +515,8 @@ cc_test(
deps = [
":simulation_clock",
":simulation_clock_executor",
"//mediapipe/calculators/core:flow_limiter_calculator",
"//mediapipe/calculators/core:immediate_mux_calculator",
"//mediapipe/calculators/core:real_time_flow_limiter_calculator",
"//mediapipe/calculators/core:round_robin_demux_calculator",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:executor",
@@ -483,3 +534,71 @@ cc_test(
"@com_google_absl//absl/strings",
],
)
mediapipe_binary_graph(
name = "test_binarypb",
graph = "//mediapipe/framework/tool/testdata:test_graph",
output_name = "test.binarypb",
visibility = ["//visibility:private"],
)
data_as_c_string(
name = "test_binarypb_inc",
testonly = 1,
srcs = [":test_binarypb"],
outs = ["test_binarypb.inc"],
)
proto_library(
name = "node_chain_subgraph_proto",
srcs = ["node_chain_subgraph.proto"],
deps = [
"//mediapipe/framework:calculator_proto",
],
)
mediapipe_cc_proto_library(
name = "node_chain_subgraph_cc_proto",
srcs = ["node_chain_subgraph.proto"],
cc_deps = ["//mediapipe/framework:calculator_cc_proto"],
visibility = ["//mediapipe:__subpackages__"],
deps = [":node_chain_subgraph_proto"],
)
cc_test(
name = "data_as_c_string_test",
srcs = [
"data_as_c_string_test.cc",
":test_binarypb_inc",
],
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/port:gtest_main",
],
)
cc_test(
name = "subgraph_expansion_test",
size = "small",
srcs = ["subgraph_expansion_test.cc"],
deps = [
":node_chain_subgraph_cc_proto",
":subgraph_expansion",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:mediapipe_options_cc_proto",
"//mediapipe/framework:packet_generator_cc_proto",
"//mediapipe/framework:packet_set",
"//mediapipe/framework:packet_type",
"//mediapipe/framework:status_handler",
"//mediapipe/framework:subgraph",
"//mediapipe/framework/deps:message_matchers",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/tool/testdata:dub_quad_test_subgraph",
"//mediapipe/framework/tool/testdata:nested_test_subgraph",
"@com_google_absl//absl/strings",
],
)
mediapipe/framework/tool/calculator_graph_template.proto
+3
View File
@@ -5,6 +5,9 @@ package mediapipe;
import "mediapipe/framework/calculator.proto";
import "mediapipe/framework/deps/proto_descriptor.proto";
option java_package = "com.google.mediapipe.proto";
option java_outer_classname = "GraphTemplateProto";
// A template rule or a template rule argument expression.
message TemplateExpression {
// A template parameter name or a literal value.
mediapipe/framework/tool/data_as_c_string_test.cc
+33
View File
@@ -0,0 +1,33 @@
// Copyright 2019 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/calculator_framework.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
namespace mediapipe {
namespace {
static const char my_graph[] =
#include "mediapipe/framework/tool/test_binarypb.inc"
; // NOLINT(whitespace/semicolon)
TEST(DataAsCString, CanDecodeCalculatorGraphConfig) {
CalculatorGraphConfig config;
bool success = config.ParseFromArray(my_graph, sizeof(my_graph) - 1);
EXPECT_TRUE(success);
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/encode_as_c_string.cc
+62
View File
@@ -0,0 +1,62 @@
// Copyright 2019 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.
//
// This program takes one input file and encodes its contents as a C++
// std::string, which can be included in a C++ source file. It is similar to
// filewrapper (and borrows some of its code), but simpler.
#include <fstream>
#include <iostream>
#include <memory>
#include "absl/strings/escaping.h"
int main(int argc, char** argv) {
if (argc != 2) {
std::cerr << "usage: encode_as_c_string input_file\n";
return 1;
}
const std::string input_name = argv[1];
std::ifstream input_file(input_name,
std::ios_base::in | std::ios_base::binary);
if (!input_file.is_open()) {
std::cerr << "cannot open '" << input_name << "'\n";
return 2;
}
constexpr int kBufSize = 4096;
std::unique_ptr<char[]> buf(new char[kBufSize]);
std::cout << "\"";
int line_len = 1;
while (1) {
input_file.read(buf.get(), kBufSize);
int count = input_file.gcount();
if (count == 0) break;
for (int i = 0; i < count; ++i) {
std::string out = absl::CEscape(absl::string_view(&buf[i], 1));
if (line_len + out.size() > 79) {
std::cout << "\"\n\"";
line_len = 1;
}
std::cout << out;
line_len += out.size();
}
}
input_file.close();
if (!input_file.eof()) {
std::cerr << "error reading '" << input_name << "'\n";
return 2;
}
std::cout << "\"\n";
return 0;
}
mediapipe/framework/tool/mediapipe_graph.bzl
+111 -50
View File
@@ -16,26 +16,11 @@ Example:
"""
load("//mediapipe/framework:encode_binary_proto.bzl", "encode_binary_proto", "generate_proto_descriptor_set")
load("//mediapipe/framework:transitive_protos.bzl", "transitive_proto_cc_libs", "transitive_proto_descriptor_sets", "transitive_protos")
load("//mediapipe/framework:transitive_protos.bzl", "transitive_protos")
load("//mediapipe/framework/deps:expand_template.bzl", "expand_template")
def mediapipe_binary_graph(name, graph = None, output_name = None, deps = [], testonly = None, **kwargs):
"""Converts a graph from text format to binary format.
Args:
name: the name of the encode_binary_proto rule generated by this macro.
graph: the BUILD label of a text-format MediaPipe graph.
output_name: the name of the file to which the binary serialization is
written.
deps: the BUILD labels of dependencies that provide any additional message
types used by the graph. The basic messages defined in calculator.proto
are always available, but any custom types (e.g. specific calculator
options) should be provided here. It is sufficient to provide targets
that depend on the required protos indirectly: this macro examines the
entire dependency tree, and does not build any dependencies except for
the protos it finds.
testonly: pass 1 if the graph is to be used only for tests.
**kwargs: any other arguments valid for encode_binary_proto.
"""
def mediapipe_binary_graph(name, graph = None, output_name = None, deps = [], testonly = False, **kwargs):
"""Converts a graph from text format to binary format."""
if not graph:
fail("No input graph file specified.")
@@ -44,48 +29,124 @@ def mediapipe_binary_graph(name, graph = None, output_name = None, deps = [], te
fail("Must specify the output_name.")
transitive_protos(
name = name + "_gather_protos",
name = name + "_gather_cc_protos",
deps = deps,
testonly = testonly,
)
# This collects descriptor sets for tools that need them.
transitive_proto_descriptor_sets(
name = name + "_gather_proto_descriptor_sets",
deps = deps,
testonly = testonly,
)
# This collects the generated .a libraries for tools that need them.
transitive_proto_cc_libs(
name = name + "_gather_proto_libs",
deps = deps,
testonly = testonly,
)
# This generates a single descriptor set with a single invocation of the proto compiler.
# May be faster than using the descriptor sets from proto_library.
# We always pass at least the calculator proto since the proto compiler would fail
# if it were passed no protos.
generate_proto_descriptor_set(
name = name + "_proto_descriptor_set",
# Compile a simple proto parser binary using the deps.
native.cc_binary(
name = name + "_text_to_binary_graph",
visibility = ["//visibility:private"],
deps = [
name + "_gather_protos",
"//mediapipe/framework:calculator_proto",
"//mediapipe/framework/tool:text_to_binary_graph",
name + "_gather_cc_protos",
],
tags = ["manual"],
testonly = testonly,
)
return encode_binary_proto(
# Invoke the proto parser binary.
native.genrule(
name = name,
deps = [
name + "_gather_protos",
"//mediapipe/framework:calculator_proto",
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler_proto",
srcs = [graph],
outs = [output_name],
cmd = (
"$(location " + name + "_text_to_binary_graph" + ") " +
("--proto_source=$(location %s) " % graph) +
("--proto_output=\"$@\" ")
),
tools = [name + "_text_to_binary_graph"],
testonly = testonly,
)
def data_as_c_string(
name,
srcs,
outs = None,
testonly = None):
"""Encodes the data from a file as a C string literal.
This produces a text file containing the quoted C string literal. It can be
included directly in a C++ source file.
Args:
name: The name of the rule.
srcs: A list containing a single item, the file to encode.
outs: A list containing a single item, the name of the output text file.
Defaults to the rule name.
testonly: pass 1 if the graph is to be used only for tests.
"""
if len(srcs) != 1:
fail("srcs must be a single-element list")
if outs == None:
outs = [name]
native.genrule(
name = name,
srcs = srcs,
outs = outs,
cmd = "$(location //mediapipe/framework/tool:encode_as_c_string) \"$<\" > \"$@\"",
tools = ["//mediapipe/framework/tool:encode_as_c_string"],
testonly = testonly,
)
def mediapipe_simple_subgraph(
name,
register_as,
graph,
deps = [],
visibility = None,
testonly = None,
**kwargs):
"""Defines a registered subgraph for inclusion in other graphs.
Args:
name: name of the subgraph target to define.
register_as: name used to invoke this graph in supergraphs. Should be in
CamelCase.
graph: the BUILD label of a text-format MediaPipe graph.
deps: any calculators or subgraphs used by this graph.
visibility: The list of packages the subgraph should be visible to.
testonly: pass 1 if the graph is to be used only for tests.
**kwargs: Remaining keyword args, forwarded to cc_library.
"""
graph_base_name = graph.replace(":", "/").split("/")[-1].rsplit(".", 1)[0]
mediapipe_binary_graph(
name = name + "_graph",
graph = graph,
output_name = graph_base_name + ".binarypb",
deps = deps,
testonly = testonly,
)
data_as_c_string(
name = name + "_inc",
srcs = [graph_base_name + ".binarypb"],
outs = [graph_base_name + ".inc"],
)
# cc_library for a linked mediapipe graph.
expand_template(
name = name + "_linked_cc",
template = "//mediapipe/framework/tool:simple_subgraph_template.cc",
out = name + "_linked.cc",
substitutions = {
"{{SUBGRAPH_CLASS_NAME}}": register_as,
"{{SUBGRAPH_INC_FILE_PATH}}": native.package_name() + "/" + graph_base_name + ".inc",
},
testonly = testonly,
)
native.cc_library(
name = name,
srcs = [
name + "_linked.cc",
graph_base_name + ".inc",
],
message_type = "mediapipe.CalculatorGraphConfig",
input = graph,
output = output_name,
deps = [
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:subgraph",
] + deps,
alwayslink = 1,
visibility = visibility,
testonly = testonly,
**kwargs
)
mediapipe/framework/tool/node_chain_subgraph.proto
+19
View File
@@ -0,0 +1,19 @@
syntax = "proto2";
package mediapipe;
import "mediapipe/framework/calculator.proto";
// Options for NodeChainSubgraph.
message NodeChainSubgraphOptions {
extend CalculatorOptions {
optional NodeChainSubgraphOptions ext = 167210579;
}
// The type of the node. The node must have exactly one input stream and
// exactly one output stream.
optional string node_type = 1;
// How many copies of the node should be chained in series.
optional int32 chain_length = 2;
}
mediapipe/framework/tool/simple_subgraph_template.cc
+43
View File
@@ -0,0 +1,43 @@
// Copyright 2019 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.
//
// This template is used by the mediapipe_simple_subgraph macro in
// //mediapipe/framework/tool/mediapipe_graph.bzl
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/subgraph.h"
namespace mediapipe {
static const char binary_graph[] =
#include "{{SUBGRAPH_INC_FILE_PATH}}"
; // NOLINT(whitespace/semicolon)
class {{SUBGRAPH_CLASS_NAME}} : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& /*options*/) {
CalculatorGraphConfig config;
// Note: this is a binary protobuf serialization, and may include NUL
// bytes. The trailing NUL added to the std::string literal should be excluded.
if (config.ParseFromArray(binary_graph, sizeof(binary_graph) - 1)) {
return config;
} else {
return ::mediapipe::InternalError("Could not parse subgraph.");
}
}
};
REGISTER_MEDIAPIPE_GRAPH({{SUBGRAPH_CLASS_NAME}});
} // namespace mediapipe
mediapipe/framework/tool/simulation_clock_test.cc
+1 -1
View File
@@ -46,7 +46,7 @@ class SimulationClockTest : public ::testing::Test {
graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_packets_0"
node {
calculator: 'RealTimeFlowLimiterCalculator'
calculator: 'FlowLimiterCalculator'
input_stream_handler {
input_stream_handler: 'ImmediateInputStreamHandler'
}
mediapipe/framework/tool/subgraph_expansion_test.cc
+529
View File
@@ -0,0 +1,529 @@
// Copyright 2019 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/tool/subgraph_expansion.h"
#include <memory>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/message_matchers.h"
#include "mediapipe/framework/packet_set.h"
#include "mediapipe/framework/packet_type.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/proto_ns.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/status_handler.h"
#include "mediapipe/framework/subgraph.h"
#include "mediapipe/framework/tool/node_chain_subgraph.pb.h"
namespace mediapipe {
namespace {
class SimpleTestCalculator : public CalculatorBase {
public:
::mediapipe::Status Process(CalculatorContext* cc) override {
return ::mediapipe::OkStatus();
}
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
for (PacketType& type : cc->Inputs()) {
type.Set<int>();
}
for (PacketType& type : cc->Outputs()) {
type.Set<int>();
}
for (PacketType& type : cc->InputSidePackets()) {
type.Set<int>();
}
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(SimpleTestCalculator);
typedef SimpleTestCalculator SomeSourceCalculator;
typedef SimpleTestCalculator SomeSinkCalculator;
typedef SimpleTestCalculator SomeRegularCalculator;
typedef SimpleTestCalculator SomeAggregator;
REGISTER_CALCULATOR(SomeSourceCalculator);
REGISTER_CALCULATOR(SomeSinkCalculator);
REGISTER_CALCULATOR(SomeRegularCalculator);
REGISTER_CALCULATOR(SomeAggregator);
class TestSubgraph : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& /*options*/) override {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "DATA:input_1"
node {
name: "regular_node"
calculator: "SomeRegularCalculator"
input_stream: "input_1"
output_stream: "stream_a"
input_side_packet: "side"
}
node {
name: "simple_sink"
calculator: "SomeSinkCalculator"
input_stream: "stream_a"
}
packet_generator {
packet_generator: "SomePacketGenerator"
output_side_packet: "side"
}
)");
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(TestSubgraph);
class PacketFactoryTestSubgraph : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& /*options*/) override {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "DATA:input_1"
node {
name: "regular_node"
calculator: "SomeRegularCalculator"
input_stream: "input_1"
output_stream: "stream_a"
input_side_packet: "side"
}
node {
name: "simple_sink"
calculator: "SomeSinkCalculator"
input_stream: "stream_a"
}
packet_factory {
packet_factory: "SomePacketFactory"
output_side_packet: "side"
}
)");
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(PacketFactoryTestSubgraph);
// This subgraph chains copies of the specified node in series. The node type
// and the number of copies of the node are specified in subgraph options.
class NodeChainSubgraph : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) override {
const mediapipe::NodeChainSubgraphOptions& opts =
options.GetExtension(mediapipe::NodeChainSubgraphOptions::ext);
const ProtoString& node_type = opts.node_type();
int chain_length = opts.chain_length();
RET_CHECK(!node_type.empty());
RET_CHECK_GT(chain_length, 0);
CalculatorGraphConfig config;
config.add_input_stream("INPUT:stream_0");
config.add_output_stream(absl::StrCat("OUTPUT:stream_", chain_length));
for (int i = 0; i < chain_length; ++i) {
CalculatorGraphConfig::Node* node = config.add_node();
node->set_calculator(node_type);
node->add_input_stream(absl::StrCat("stream_", i));
node->add_output_stream(absl::StrCat("stream_", i + 1));
}
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(NodeChainSubgraph);
// A subgraph used in the ExecutorFieldOfNodeInSubgraphPreserved test. The
// subgraph contains a node with the executor field "custom_thread_pool".
class NodeWithExecutorSubgraph : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) override {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "INPUT:foo"
output_stream: "OUTPUT:bar"
node {
calculator: "PassThroughCalculator"
input_stream: "foo"
output_stream: "bar"
executor: "custom_thread_pool"
}
)");
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(NodeWithExecutorSubgraph);
// A subgraph used in the ExecutorFieldOfNodeInSubgraphPreserved test. The
// subgraph contains a NodeWithExecutorSubgraph.
class EnclosingSubgraph : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) override {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "IN:in"
output_stream: "OUT:out"
node {
calculator: "NodeWithExecutorSubgraph"
input_stream: "INPUT:in"
output_stream: "OUTPUT:out"
}
)");
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(EnclosingSubgraph);
TEST(SubgraphExpansionTest, TransformStreamNames) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "SomeSinkCalculator"
input_stream: "input_1"
input_stream: "VIDEO:input_2"
input_stream: "AUDIO:0:input_3"
input_stream: "AUDIO:1:input_4"
}
)");
CalculatorGraphConfig expected_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "SomeSinkCalculator"
input_stream: "input_1_foo"
input_stream: "VIDEO:input_2_foo"
input_stream: "AUDIO:0:input_3_foo"
input_stream: "AUDIO:1:input_4_foo"
}
)");
auto add_foo = [](absl::string_view s) { return absl::StrCat(s, "_foo"); };
MEDIAPIPE_EXPECT_OK(tool::TransformStreamNames(
(*config.mutable_node())[0].mutable_input_stream(), add_foo));
EXPECT_THAT(config, mediapipe::EqualsProto(expected_config));
}
TEST(SubgraphExpansionTest, TransformNames) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_1"
node {
calculator: "SomeRegularCalculator"
name: "bob"
input_stream: "input_1"
input_stream: "VIDEO:input_2"
input_stream: "AUDIO:0:input_3"
input_stream: "AUDIO:1:input_4"
output_stream: "output_1"
}
node {
calculator: "SomeRegularCalculator"
input_stream: "output_1"
output_stream: "output_2"
}
)");
CalculatorGraphConfig expected_config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "__sg0_input_1"
node {
calculator: "SomeRegularCalculator"
name: "__sg0_bob"
input_stream: "__sg0_input_1"
input_stream: "VIDEO:__sg0_input_2"
input_stream: "AUDIO:0:__sg0_input_3"
input_stream: "AUDIO:1:__sg0_input_4"
output_stream: "__sg0_output_1"
}
node {
calculator: "SomeRegularCalculator"
input_stream: "__sg0_output_1"
output_stream: "__sg0_output_2"
}
)");
auto add_prefix = [](absl::string_view s) {
return absl::StrCat("__sg0_", s);
};
MEDIAPIPE_EXPECT_OK(tool::TransformNames(&config, add_prefix));
EXPECT_THAT(config, mediapipe::EqualsProto(expected_config));
}
TEST(SubgraphExpansionTest, FindCorrespondingStreams) {
CalculatorGraphConfig config1 =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_1"
input_stream: "VIDEO:input_2"
input_stream: "AUDIO:0:input_3"
input_stream: "AUDIO:1:input_4"
)");
CalculatorGraphConfig config2 =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "SomeSubgraph"
input_stream: "foo"
input_stream: "VIDEO:bar"
input_stream: "AUDIO:0:baz"
input_stream: "AUDIO:1:qux"
}
)");
std::map<std::string, std::string> stream_map;
MEDIAPIPE_EXPECT_OK(tool::FindCorrespondingStreams(
&stream_map, config1.input_stream(), config2.node()[0].input_stream()));
EXPECT_THAT(stream_map,
testing::UnorderedElementsAre(testing::Pair("input_1", "foo"),
testing::Pair("input_2", "bar"),
testing::Pair("input_3", "baz"),
testing::Pair("input_4", "qux")));
}
TEST(SubgraphExpansionTest, FindCorrespondingStreamsNonexistentTag) {
// The VIDEO tag does not exist in the subgraph.
CalculatorGraphConfig config1 =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_1"
input_stream: "AUDIO:0:input_3"
input_stream: "AUDIO:1:input_4"
)");
CalculatorGraphConfig config2 =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "SomeSubgraph"
input_stream: "foo"
input_stream: "VIDEO:bar"
input_stream: "AUDIO:0:baz"
input_stream: "AUDIO:1:qux"
}
)");
std::map<std::string, std::string> stream_map;
auto status = tool::FindCorrespondingStreams(
&stream_map, config1.input_stream(), config2.node()[0].input_stream());
EXPECT_THAT(status.message(),
testing::AllOf(
// Problematic tag.
testing::HasSubstr("VIDEO"),
// Error.
testing::HasSubstr("does not exist")));
}
TEST(SubgraphExpansionTest, FindCorrespondingStreamsTooFewIndexes) {
// The AUDIO tag has too few indexes in the subgraph (1 vs. 2).
CalculatorGraphConfig config1 =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_1"
input_stream: "VIDEO:input_2"
input_stream: "AUDIO:0:input_3"
)");
CalculatorGraphConfig config2 =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "SomeSubgraph"
input_stream: "foo"
input_stream: "VIDEO:bar"
input_stream: "AUDIO:0:baz"
input_stream: "AUDIO:1:qux"
}
)");
std::map<std::string, std::string> stream_map;
auto status = tool::FindCorrespondingStreams(
&stream_map, config1.input_stream(), config2.node()[0].input_stream());
EXPECT_THAT(status.message(),
testing::AllOf(
// Problematic tag.
testing::HasSubstr("AUDIO"),
// Error.
testing::HasSubstr(" 2 "), testing::HasSubstr(" 1 ")));
}
TEST(SubgraphExpansionTest, ConnectSubgraphStreams) {
CalculatorGraphConfig subgraph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "A:input_1"
input_stream: "B:input_2"
output_stream: "O:output_2"
input_side_packet: "SI:side_input"
output_side_packet: "SO:side_output"
node {
calculator: "SomeRegularCalculator"
input_stream: "input_1"
input_stream: "VIDEO:input_2"
input_side_packet: "side_input"
output_stream: "output_1"
}
node {
calculator: "SomeRegularCalculator"
input_stream: "input_1"
input_stream: "output_1"
output_stream: "output_2"
}
packet_generator {
packet_generator: "SomeGenerator"
input_side_packet: "side_input"
output_side_packet: "side_output"
}
)");
CalculatorGraphConfig supergraph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "SomeSubgraph"
input_stream: "A:foo"
input_stream: "B:bar"
output_stream: "O:foobar"
input_side_packet: "SI:flip"
output_side_packet: "SO:flop"
}
)");
// Note: graph input streams, output streams, and side packets on the
// subgraph are not changed because they are going to be discarded anyway.
CalculatorGraphConfig expected_subgraph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "A:input_1"
input_stream: "B:input_2"
output_stream: "O:output_2"
input_side_packet: "SI:side_input"
output_side_packet: "SO:side_output"
node {
calculator: "SomeRegularCalculator"
input_stream: "foo"
input_stream: "VIDEO:bar"
input_side_packet: "flip"
output_stream: "output_1"
}
node {
calculator: "SomeRegularCalculator"
input_stream: "foo"
input_stream: "output_1"
output_stream: "foobar"
}
packet_generator {
packet_generator: "SomeGenerator"
input_side_packet: "flip"
output_side_packet: "flop"
}
)");
MEDIAPIPE_EXPECT_OK(
tool::ConnectSubgraphStreams(supergraph.node()[0], &subgraph));
EXPECT_THAT(subgraph, mediapipe::EqualsProto(expected_subgraph));
}
TEST(SubgraphExpansionTest, ExpandSubgraphs) {
CalculatorGraphConfig supergraph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
name: "simple_source"
calculator: "SomeSourceCalculator"
output_stream: "foo"
}
node { calculator: "TestSubgraph" input_stream: "DATA:foo" }
)");
CalculatorGraphConfig expected_graph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
name: "simple_source"
calculator: "SomeSourceCalculator"
output_stream: "foo"
}
node {
name: "__sg0_regular_node"
calculator: "SomeRegularCalculator"
input_stream: "foo"
output_stream: "__sg0_stream_a"
input_side_packet: "__sg0_side"
}
node {
name: "__sg0_simple_sink"
calculator: "SomeSinkCalculator"
input_stream: "__sg0_stream_a"
}
packet_generator {
packet_generator: "SomePacketGenerator"
output_side_packet: "__sg0_side"
}
)");
MEDIAPIPE_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
TEST(SubgraphExpansionTest, ValidateSubgraphFields) {
CalculatorGraphConfig supergraph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
name: "simple_source"
calculator: "SomeSourceCalculator"
output_stream: "foo"
}
node {
name: "foo_subgraph"
calculator: "TestSubgraph"
input_stream: "DATA:foo"
buffer_size_hint: -1 # This field is only applicable to calculators.
}
)");
::mediapipe::Status s1 = tool::ValidateSubgraphFields(supergraph.node(1));
EXPECT_EQ(s1.code(), ::mediapipe::StatusCode::kInvalidArgument);
EXPECT_THAT(s1.message(), testing::HasSubstr("foo_subgraph"));
::mediapipe::Status s2 = tool::ExpandSubgraphs(&supergraph);
EXPECT_EQ(s2.code(), ::mediapipe::StatusCode::kInvalidArgument);
EXPECT_THAT(s2.message(), testing::HasSubstr("foo_subgraph"));
}
// A test that captures the use case of CL 191001940. The "executor" field of
// a node inside a subgraph should be preserved, not mapped or mangled. This
// test will help us detect breakage of this use case when we implement
// subgraph executor support in the future.
TEST(SubgraphExpansionTest, ExecutorFieldOfNodeInSubgraphPreserved) {
CalculatorGraphConfig supergraph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input"
executor {
name: "custom_thread_pool"
type: "ThreadPoolExecutor"
options {
[mediapipe.ThreadPoolExecutorOptions.ext] { num_threads: 4 }
}
}
node {
calculator: "EnclosingSubgraph"
input_stream: "IN:input"
output_stream: "OUT:output"
}
)");
CalculatorGraphConfig expected_graph =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input"
executor {
name: "custom_thread_pool"
type: "ThreadPoolExecutor"
options {
[mediapipe.ThreadPoolExecutorOptions.ext] { num_threads: 4 }
}
}
node {
calculator: "PassThroughCalculator"
input_stream: "input"
output_stream: "output"
executor: "custom_thread_pool"
}
)");
MEDIAPIPE_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/testdata/BUILD
Vendored
+57
View File
@@ -0,0 +1,57 @@
# Copyright 2019 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.
#
licenses(["notice"]) # Apache 2.0
package(default_visibility = ["//mediapipe:__subpackages__"])
load(
"//mediapipe/framework/tool:mediapipe_graph.bzl",
"mediapipe_simple_subgraph",
)
filegroup(
name = "test_graph",
srcs = ["test.pbtxt"],
)
exports_files([
"test.pbtxt",
"dub_quad_test_subgraph.pbtxt",
"nested_test_subgraph.pbtxt",
])
mediapipe_simple_subgraph(
name = "dub_quad_test_subgraph",
testonly = 1,
graph = "dub_quad_test_subgraph.pbtxt",
register_as = "DubQuadTestSubgraph",
visibility = ["//visibility:public"],
deps = [
"//mediapipe/framework:test_calculators",
],
)
mediapipe_simple_subgraph(
name = "nested_test_subgraph",
testonly = 1,
graph = "nested_test_subgraph.pbtxt",
register_as = "NestedTestSubgraph",
visibility = ["//visibility:public"],
deps = [
":dub_quad_test_subgraph",
"//mediapipe/framework:test_calculators",
],
)
mediapipe/framework/tool/testdata/dub_quad_test_subgraph.pbtxt
Vendored
+13
View File
@@ -0,0 +1,13 @@
input_stream: "INTS:ints"
output_stream: "DUBS:doubled"
output_stream: "QUADS:quadrupled"
node {
calculator: "DoubleIntCalculator"
input_stream: "ints"
output_stream: "doubled"
}
node {
calculator: "DoubleIntCalculator"
input_stream: "doubled"
output_stream: "quadrupled"
}
mediapipe/framework/tool/testdata/nested_test_subgraph.pbtxt
Vendored
+20
View File
@@ -0,0 +1,20 @@
input_stream: "INTS:ints"
output_stream: "DUBS:doubled"
output_stream: "QUADS:quadrupled"
output_stream: "OCTS:octupled"
node {
calculator: "DubQuadTestSubgraph"
input_stream: "INTS:ints"
output_stream: "DUBS:doubled"
output_stream: "QUADS:quadrupled"
}
node {
calculator: "DoubleIntCalculator"
input_stream: "quadrupled"
output_stream: "octupled"
# The following is to ensure we handle NULs correctly.
input_stream_info {
tag_index: ":0" # 'quadrupled'
back_edge: false # The false boolean value is encoded as a zero byte.
}
}
mediapipe/framework/tool/testdata/test.pbtxt
Vendored
+5
View File
@@ -0,0 +1,5 @@
node {
calculator: "PassThroughCalculator"
input_stream: "in"
output_stream: "out"
}
mediapipe/framework/tool/text_to_binary_graph.cc
+111
View File
@@ -0,0 +1,111 @@
// Copyright 2019 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.
//
// A command line utility to parse a text proto and output a binary proto.
#include <stdlib.h>
#include <fstream>
#include <string>
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/port/advanced_proto_inc.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/commandlineflags.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
DEFINE_string(proto_source, "",
"The template source file containing CalculatorGraphConfig "
"protobuf text with inline template params.");
DEFINE_string(
proto_output, "",
"An output template file in binary CalculatorGraphTemplate form.");
#define EXIT_IF_ERROR(status) \
if (!status.ok()) { \
LOG(ERROR) << status; \
return EXIT_FAILURE; \
}
namespace mediapipe {
mediapipe::Status ReadProto(proto_ns::io::ZeroCopyInputStream* in,
bool read_text, const std::string& source,
proto_ns::Message* result) {
if (read_text) {
RET_CHECK(proto_ns::TextFormat::Parse(in, result))
<< "could not parse text proto: " << source;
} else {
RET_CHECK(result->ParseFromZeroCopyStream(in))
<< "could not parse binary proto: " << source;
}
return mediapipe::OkStatus();
}
mediapipe::Status WriteProto(const proto_ns::Message& message, bool write_text,
const std::string& dest,
proto_ns::io::ZeroCopyOutputStream* out) {
if (write_text) {
RET_CHECK(proto_ns::TextFormat::Print(message, out))
<< "could not write text proto to: " << dest;
} else {
RET_CHECK(message.SerializeToZeroCopyStream(out))
<< "could not write binary proto to: " << dest;
}
return mediapipe::OkStatus();
}
// Read a proto from a text or a binary file.
mediapipe::Status ReadFile(const std::string& proto_source, bool read_text,
proto_ns::Message* result) {
std::ifstream ifs(proto_source);
proto_ns::io::IstreamInputStream in(&ifs);
RETURN_IF_ERROR(ReadProto(&in, read_text, proto_source, result));
return mediapipe::OkStatus();
}
// Write a proto to a text or a binary file.
mediapipe::Status WriteFile(const std::string& proto_output, bool write_text,
const proto_ns::Message& message) {
std::ofstream ofs(proto_output, std::ofstream::out | std::ofstream::trunc);
proto_ns::io::OstreamOutputStream out(&ofs);
RETURN_IF_ERROR(WriteProto(message, write_text, proto_output, &out));
return mediapipe::OkStatus();
}
} // namespace mediapipe
int main(int argc, char** argv) {
google::InitGoogleLogging(argv[0]);
gflags::ParseCommandLineFlags(&argc, &argv, true);
// Validate command line options.
mediapipe::Status status;
if (FLAGS_proto_source.empty()) {
status.Update(
::mediapipe::InvalidArgumentError("--proto_source must be specified"));
}
if (FLAGS_proto_output.empty()) {
status.Update(
::mediapipe::InvalidArgumentError("--proto_output must be specified"));
}
if (!status.ok()) {
return EXIT_FAILURE;
}
mediapipe::CalculatorGraphConfig config;
EXIT_IF_ERROR(mediapipe::ReadFile(FLAGS_proto_source, true, &config));
EXIT_IF_ERROR(mediapipe::WriteFile(FLAGS_proto_output, false, config));
return EXIT_SUCCESS;
}