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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 2146b10f0a498f665f246e16033b686c7947b92d
This commit is contained in:
MediaPipe Team
2021-05-10 16:42:02 -04:00
committed by chuoling
parent a9b643e0f5
commit 017c1dc7ea
52 changed files with 708 additions and 298 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/tool/BUILD
+2
View File
@@ -681,6 +681,7 @@ cc_library(
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:switch_container_cc_proto",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
@@ -705,6 +706,7 @@ cc_library(
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:switch_container_cc_proto",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
mediapipe/framework/tool/switch_container.cc
+39
View File
@@ -82,6 +82,8 @@ CalculatorGraphConfig::Node* BuildDemuxNode(
CalculatorGraphConfig* config) {
CalculatorGraphConfig::Node* result = config->add_node();
*result->mutable_calculator() = "SwitchDemuxCalculator";
*result->mutable_input_stream_handler()->mutable_input_stream_handler() =
"ImmediateInputStreamHandler";
return result;
}
@@ -91,9 +93,42 @@ CalculatorGraphConfig::Node* BuildMuxNode(
CalculatorGraphConfig* config) {
CalculatorGraphConfig::Node* result = config->add_node();
*result->mutable_calculator() = "SwitchMuxCalculator";
*result->mutable_input_stream_handler()->mutable_input_stream_handler() =
"ImmediateInputStreamHandler";
return result;
}
// Copies options from one node to another.
void CopyOptions(const CalculatorGraphConfig::Node& source,
CalculatorGraphConfig::Node* dest) {
if (source.has_options()) {
*dest->mutable_options() = source.options();
}
*dest->mutable_node_options() = source.node_options();
}
// Clears options that are consumed by the container and not forwarded.
void ClearContainerOptions(SwitchContainerOptions* result) {
result->clear_contained_node();
}
// Clears options that are consumed by the container and not forwarded.
void ClearContainerOptions(CalculatorGraphConfig::Node* dest) {
if (dest->has_options() &&
dest->mutable_options()->HasExtension(SwitchContainerOptions::ext)) {
ClearContainerOptions(
dest->mutable_options()->MutableExtension(SwitchContainerOptions::ext));
}
for (google::protobuf::Any& a : *dest->mutable_node_options()) {
if (a.Is<SwitchContainerOptions>()) {
SwitchContainerOptions extension;
a.UnpackTo(&extension);
ClearContainerOptions(&extension);
a.PackFrom(extension);
}
}
}
// Returns an unused name similar to a specified name.
std::string UniqueName(std::string name, std::set<std::string>* names) {
CHECK(names != nullptr);
@@ -199,12 +234,16 @@ absl::StatusOr<CalculatorGraphConfig> SwitchContainer::GetConfig(
// Add a graph node for the demux, mux.
auto demux = BuildDemuxNode(input_tags, &config);
CopyOptions(container_node, demux);
ClearContainerOptions(demux);
demux->add_input_stream("SELECT:gate_select");
demux->add_input_stream("ENABLE:gate_enable");
demux->add_input_side_packet("SELECT:gate_select");
demux->add_input_side_packet("ENABLE:gate_enable");
auto mux = BuildMuxNode(output_tags, &config);
CopyOptions(container_node, mux);
ClearContainerOptions(mux);
mux->add_input_stream("SELECT:gate_select");
mux->add_input_stream("ENABLE:gate_enable");
mux->add_input_side_packet("SELECT:gate_select");
mediapipe/framework/tool/switch_container_test.cc
+93
View File
@@ -225,6 +225,12 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
input_stream: "foo"
output_stream: "C0__:switchcontainer__c0__foo"
output_stream: "C1__:switchcontainer__c1__foo"
options {
[mediapipe.SwitchContainerOptions.ext] {}
}
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
}
node {
name: "switchcontainer__TripleIntCalculator"
@@ -245,6 +251,12 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
input_stream: "C0__:switchcontainer__c0__bar"
input_stream: "C1__:switchcontainer__c1__bar"
output_stream: "bar"
options {
[mediapipe.SwitchContainerOptions.ext] {}
}
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
}
node {
calculator: "PassThroughCalculator"
@@ -270,6 +282,75 @@ TEST(SwitchContainerTest, RunsWithSubnodes) {
RunTestContainer(supergraph);
}
// Shows the SwitchContainer does not allow input_stream_handler overwrite.
TEST(SwitchContainerTest, ValidateInputStreamHandler) {
EXPECT_TRUE(SubgraphRegistry::IsRegistered("SwitchContainer"));
CalculatorGraph graph;
CalculatorGraphConfig supergraph = SideSubnodeContainerExample();
*supergraph.mutable_input_stream_handler()->mutable_input_stream_handler() =
"DefaultInputStreamHandler";
MP_ASSERT_OK(graph.Initialize(supergraph, {}));
CalculatorGraphConfig expected_graph = mediapipe::ParseTextProtoOrDie<
CalculatorGraphConfig>(R"pb(
node {
name: "switchcontainer__SwitchDemuxCalculator"
calculator: "SwitchDemuxCalculator"
input_side_packet: "ENABLE:enable"
input_side_packet: "foo"
output_side_packet: "C0__:switchcontainer__c0__foo"
output_side_packet: "C1__:switchcontainer__c1__foo"
options {
[mediapipe.SwitchContainerOptions.ext] {}
}
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
}
node {
name: "switchcontainer__TripleIntCalculator"
calculator: "TripleIntCalculator"
input_side_packet: "switchcontainer__c0__foo"
output_side_packet: "switchcontainer__c0__bar"
input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
}
node {
name: "switchcontainer__PassThroughCalculator"
calculator: "PassThroughCalculator"
input_side_packet: "switchcontainer__c1__foo"
output_side_packet: "switchcontainer__c1__bar"
input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
}
node {
name: "switchcontainer__SwitchMuxCalculator"
calculator: "SwitchMuxCalculator"
input_side_packet: "ENABLE:enable"
input_side_packet: "C0__:switchcontainer__c0__bar"
input_side_packet: "C1__:switchcontainer__c1__bar"
output_side_packet: "bar"
options {
[mediapipe.SwitchContainerOptions.ext] {}
}
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
}
node {
calculator: "PassThroughCalculator"
input_side_packet: "foo"
input_side_packet: "bar"
output_side_packet: "output_foo"
output_side_packet: "output_bar"
input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
}
input_stream_handler { input_stream_handler: "DefaultInputStreamHandler" }
executor {}
input_side_packet: "foo"
input_side_packet: "enable"
output_side_packet: "output_bar"
)pb");
EXPECT_THAT(graph.Config(), mediapipe::EqualsProto(expected_graph));
}
// Shows the SwitchContainer container applied to a pair of simple subnodes.
TEST(SwitchContainerTest, ApplyToSideSubnodes) {
EXPECT_TRUE(SubgraphRegistry::IsRegistered("SwitchContainer"));
@@ -286,6 +367,12 @@ TEST(SwitchContainerTest, ApplyToSideSubnodes) {
input_side_packet: "foo"
output_side_packet: "C0__:switchcontainer__c0__foo"
output_side_packet: "C1__:switchcontainer__c1__foo"
options {
[mediapipe.SwitchContainerOptions.ext] {}
}
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
}
node {
name: "switchcontainer__TripleIntCalculator"
@@ -306,6 +393,12 @@ TEST(SwitchContainerTest, ApplyToSideSubnodes) {
input_side_packet: "C0__:switchcontainer__c0__bar"
input_side_packet: "C1__:switchcontainer__c1__bar"
output_side_packet: "bar"
options {
[mediapipe.SwitchContainerOptions.ext] {}
}
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
}
node {
calculator: "PassThroughCalculator"
mediapipe/framework/tool/switch_demux_calculator.cc
+4 -10
View File
@@ -70,17 +70,11 @@ REGISTER_CALCULATOR(SwitchDemuxCalculator);
absl::Status SwitchDemuxCalculator::GetContract(CalculatorContract* cc) {
// Allow any one of kSelectTag, kEnableTag.
if (cc->Inputs().HasTag(kSelectTag)) {
cc->Inputs().Tag(kSelectTag).Set<int>();
} else if (cc->Inputs().HasTag(kEnableTag)) {
cc->Inputs().Tag(kEnableTag).Set<bool>();
}
cc->Inputs().Tag(kSelectTag).Set<int>().Optional();
cc->Inputs().Tag(kEnableTag).Set<bool>().Optional();
// Allow any one of kSelectTag, kEnableTag.
if (cc->InputSidePackets().HasTag(kSelectTag)) {
cc->InputSidePackets().Tag(kSelectTag).Set<int>();
} else if (cc->InputSidePackets().HasTag(kEnableTag)) {
cc->InputSidePackets().Tag(kEnableTag).Set<bool>();
}
cc->InputSidePackets().Tag(kSelectTag).Set<int>().Optional();
cc->InputSidePackets().Tag(kEnableTag).Set<bool>().Optional();
// Set the types for all output channels to corresponding input types.
std::set<std::string> channel_tags = ChannelTags(cc->Outputs().TagMap());
mediapipe/framework/tool/switch_mux_calculator.cc
+4 -10
View File
@@ -73,17 +73,11 @@ REGISTER_CALCULATOR(SwitchMuxCalculator);
absl::Status SwitchMuxCalculator::GetContract(CalculatorContract* cc) {
// Allow any one of kSelectTag, kEnableTag.
if (cc->Inputs().HasTag(kSelectTag)) {
cc->Inputs().Tag(kSelectTag).Set<int>();
} else if (cc->Inputs().HasTag(kEnableTag)) {
cc->Inputs().Tag(kEnableTag).Set<bool>();
}
cc->Inputs().Tag(kSelectTag).Set<int>().Optional();
cc->Inputs().Tag(kEnableTag).Set<bool>().Optional();
// Allow any one of kSelectTag, kEnableTag.
if (cc->InputSidePackets().HasTag(kSelectTag)) {
cc->InputSidePackets().Tag(kSelectTag).Set<int>();
} else if (cc->InputSidePackets().HasTag(kEnableTag)) {
cc->InputSidePackets().Tag(kEnableTag).Set<bool>();
}
cc->InputSidePackets().Tag(kSelectTag).Set<int>().Optional();
cc->InputSidePackets().Tag(kEnableTag).Set<bool>().Optional();
// Set the types for all input channels to corresponding output types.
std::set<std::string> channel_tags = ChannelTags(cc->Inputs().TagMap());