Add gate utility functions.
PiperOrigin-RevId: 478026407
This commit is contained in:
MediaPipe Team
Sebastian Schmidt
parent
af2ad1abbe
commit
46a5117c6d
|
|
@ -42,3 +42,16 @@ cc_test(
|
||||||
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
|
"//mediapipe/tasks/cc/components/containers/proto:embeddings_cc_proto",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
cc_library(
|
||||||
|
name = "gate",
|
||||||
|
hdrs = ["gate.h"],
|
||||||
|
deps = [
|
||||||
|
"//mediapipe/calculators/core:gate_calculator",
|
||||||
|
"//mediapipe/calculators/core:gate_calculator_cc_proto",
|
||||||
|
"//mediapipe/framework/api2:builder",
|
||||||
|
"//mediapipe/framework/api2:port",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
# TODO: Enable this test
|
||||||
|
|
|
||||||
158
mediapipe/tasks/cc/components/utils/gate.h
Normal file
158
mediapipe/tasks/cc/components/utils/gate.h
Normal file
|
|
@ -0,0 +1,158 @@
|
||||||
|
/* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
|
||||||
|
|
||||||
|
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_TASKS_CC_COMPONENTS_UTILS_GATE_H_
|
||||||
|
#define MEDIAPIPE_TASKS_CC_COMPONENTS_UTILS_GATE_H_
|
||||||
|
|
||||||
|
#include <utility>
|
||||||
|
|
||||||
|
#include "mediapipe/calculators/core/gate_calculator.pb.h"
|
||||||
|
#include "mediapipe/framework/api2/builder.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
namespace tasks {
|
||||||
|
namespace components {
|
||||||
|
namespace utils {
|
||||||
|
|
||||||
|
using ::mediapipe::api2::builder::SideSource;
|
||||||
|
using ::mediapipe::api2::builder::Source;
|
||||||
|
|
||||||
|
// Utility class that simplifies allowing (gating) multiple streams.
|
||||||
|
class AllowGate {
|
||||||
|
public:
|
||||||
|
AllowGate(Source<bool> allow, mediapipe::api2::builder::Graph& graph)
|
||||||
|
: node_(AddSourceGate(allow, graph)) {}
|
||||||
|
AllowGate(SideSource<bool> allow, mediapipe::api2::builder::Graph& graph)
|
||||||
|
: node_(AddSideSourceGate(allow, graph)) {}
|
||||||
|
|
||||||
|
// Move-only
|
||||||
|
AllowGate(AllowGate&& allow_gate) = default;
|
||||||
|
AllowGate& operator=(AllowGate&& allow_gate) = default;
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
Source<T> Allow(Source<T> source) {
|
||||||
|
source >> node_.In(index_);
|
||||||
|
return node_.Out(index_++).Cast<T>();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
template <typename T>
|
||||||
|
static mediapipe::api2::builder::GenericNode& AddSourceGate(
|
||||||
|
T allow, mediapipe::api2::builder::Graph& graph) {
|
||||||
|
auto& gate_node = graph.AddNode("GateCalculator");
|
||||||
|
allow >> gate_node.In("ALLOW");
|
||||||
|
return gate_node;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
static mediapipe::api2::builder::GenericNode& AddSideSourceGate(
|
||||||
|
T allow, mediapipe::api2::builder::Graph& graph) {
|
||||||
|
auto& gate_node = graph.AddNode("GateCalculator");
|
||||||
|
allow >> gate_node.SideIn("ALLOW");
|
||||||
|
return gate_node;
|
||||||
|
}
|
||||||
|
|
||||||
|
mediapipe::api2::builder::GenericNode& node_;
|
||||||
|
int index_ = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
// Utility class that simplifies disallowing (gating) multiple streams.
|
||||||
|
class DisallowGate {
|
||||||
|
public:
|
||||||
|
DisallowGate(Source<bool> disallow, mediapipe::api2::builder::Graph& graph)
|
||||||
|
: node_(AddSourceGate(disallow, graph)) {}
|
||||||
|
DisallowGate(SideSource<bool> disallow,
|
||||||
|
mediapipe::api2::builder::Graph& graph)
|
||||||
|
: node_(AddSideSourceGate(disallow, graph)) {}
|
||||||
|
|
||||||
|
// Move-only
|
||||||
|
DisallowGate(DisallowGate&& disallow_gate) = default;
|
||||||
|
DisallowGate& operator=(DisallowGate&& disallow_gate) = default;
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
Source<T> Disallow(Source<T> source) {
|
||||||
|
source >> node_.In(index_);
|
||||||
|
return node_.Out(index_++).Cast<T>();
|
||||||
|
}
|
||||||
|
|
||||||
|
private:
|
||||||
|
template <typename T>
|
||||||
|
static mediapipe::api2::builder::GenericNode& AddSourceGate(
|
||||||
|
T disallow, mediapipe::api2::builder::Graph& graph) {
|
||||||
|
auto& gate_node = graph.AddNode("GateCalculator");
|
||||||
|
auto& gate_node_opts =
|
||||||
|
gate_node.GetOptions<mediapipe::GateCalculatorOptions>();
|
||||||
|
// Supposedly, the most popular configuration for MediaPipe Tasks team
|
||||||
|
// graphs. Hence, intentionally hard coded to catch and verify any other use
|
||||||
|
// case (should help to workout a common approach and have a recommended way
|
||||||
|
// of blocking streams).
|
||||||
|
gate_node_opts.set_empty_packets_as_allow(true);
|
||||||
|
disallow >> gate_node.In("DISALLOW");
|
||||||
|
return gate_node;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <typename T>
|
||||||
|
static mediapipe::api2::builder::GenericNode& AddSideSourceGate(
|
||||||
|
T disallow, mediapipe::api2::builder::Graph& graph) {
|
||||||
|
auto& gate_node = graph.AddNode("GateCalculator");
|
||||||
|
auto& gate_node_opts =
|
||||||
|
gate_node.GetOptions<mediapipe::GateCalculatorOptions>();
|
||||||
|
gate_node_opts.set_empty_packets_as_allow(true);
|
||||||
|
disallow >> gate_node.SideIn("DISALLOW");
|
||||||
|
return gate_node;
|
||||||
|
}
|
||||||
|
|
||||||
|
mediapipe::api2::builder::GenericNode& node_;
|
||||||
|
int index_ = 0;
|
||||||
|
};
|
||||||
|
|
||||||
|
// Updates graph to drop @value stream packet if corresponding @condition stream
|
||||||
|
// packet holds true.
|
||||||
|
template <class T>
|
||||||
|
Source<T> DisallowIf(Source<T> value, Source<bool> condition,
|
||||||
|
mediapipe::api2::builder::Graph& graph) {
|
||||||
|
return DisallowGate(condition, graph).Disallow(value);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Updates graph to drop @value stream packet if corresponding @condition stream
|
||||||
|
// packet holds true.
|
||||||
|
template <class T>
|
||||||
|
Source<T> DisallowIf(Source<T> value, SideSource<bool> condition,
|
||||||
|
mediapipe::api2::builder::Graph& graph) {
|
||||||
|
return DisallowGate(condition, graph).Disallow(value);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Updates graph to pass through @value stream packet if corresponding
|
||||||
|
// @condition stream packet holds true.
|
||||||
|
template <class T>
|
||||||
|
Source<T> AllowIf(Source<T> value, Source<bool> allow,
|
||||||
|
mediapipe::api2::builder::Graph& graph) {
|
||||||
|
return AllowGate(allow, graph).Allow(value);
|
||||||
|
}
|
||||||
|
|
||||||
|
// Updates graph to pass through @value stream packet if corresponding
|
||||||
|
// @condition side stream packet holds true.
|
||||||
|
template <class T>
|
||||||
|
Source<T> AllowIf(Source<T> value, SideSource<bool> allow,
|
||||||
|
mediapipe::api2::builder::Graph& graph) {
|
||||||
|
return AllowGate(allow, graph).Allow(value);
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace utils
|
||||||
|
} // namespace components
|
||||||
|
} // namespace tasks
|
||||||
|
} // namespace mediapipe
|
||||||
|
|
||||||
|
#endif // MEDIAPIPE_TASKS_CC_COMPONENTS_UTILS_GATE_H_
|
||||||
229
mediapipe/tasks/cc/components/utils/gate_test.cc
Normal file
229
mediapipe/tasks/cc/components/utils/gate_test.cc
Normal file
|
|
@ -0,0 +1,229 @@
|
||||||
|
/* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
|
||||||
|
|
||||||
|
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/tasks/cc/components/utils/gate.h"
|
||||||
|
|
||||||
|
#include "mediapipe/framework/api2/builder.h"
|
||||||
|
#include "mediapipe/framework/calculator_framework.h"
|
||||||
|
#include "mediapipe/framework/calculator_graph.h"
|
||||||
|
#include "mediapipe/framework/port/gmock.h"
|
||||||
|
#include "mediapipe/framework/port/gtest.h"
|
||||||
|
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||||
|
|
||||||
|
namespace mediapipe {
|
||||||
|
namespace tasks {
|
||||||
|
namespace components {
|
||||||
|
namespace utils {
|
||||||
|
namespace {
|
||||||
|
|
||||||
|
using ::mediapipe::api2::builder::SideSource;
|
||||||
|
using ::mediapipe::api2::builder::Source;
|
||||||
|
|
||||||
|
TEST(DisallowGate, VerifyConfig) {
|
||||||
|
mediapipe::api2::builder::Graph graph;
|
||||||
|
|
||||||
|
Source<bool> condition = graph.In("CONDITION").Cast<bool>();
|
||||||
|
Source<int> value1 = graph.In("VALUE_1").Cast<int>();
|
||||||
|
Source<int> value2 = graph.In("VALUE_2").Cast<int>();
|
||||||
|
Source<int> value3 = graph.In("VALUE_3").Cast<int>();
|
||||||
|
|
||||||
|
DisallowGate gate(condition, graph);
|
||||||
|
gate.Disallow(value1).SetName("gated_stream1");
|
||||||
|
gate.Disallow(value2).SetName("gated_stream2");
|
||||||
|
gate.Disallow(value3).SetName("gated_stream3");
|
||||||
|
|
||||||
|
EXPECT_THAT(graph.GetConfig(),
|
||||||
|
testing::EqualsProto(
|
||||||
|
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||||
|
node {
|
||||||
|
calculator: "GateCalculator"
|
||||||
|
input_stream: "__stream_1"
|
||||||
|
input_stream: "__stream_2"
|
||||||
|
input_stream: "__stream_3"
|
||||||
|
input_stream: "DISALLOW:__stream_0"
|
||||||
|
output_stream: "gated_stream1"
|
||||||
|
output_stream: "gated_stream2"
|
||||||
|
output_stream: "gated_stream3"
|
||||||
|
options {
|
||||||
|
[mediapipe.GateCalculatorOptions.ext] {
|
||||||
|
empty_packets_as_allow: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
input_stream: "CONDITION:__stream_0"
|
||||||
|
input_stream: "VALUE_1:__stream_1"
|
||||||
|
input_stream: "VALUE_2:__stream_2"
|
||||||
|
input_stream: "VALUE_3:__stream_3"
|
||||||
|
)pb")));
|
||||||
|
|
||||||
|
CalculatorGraph calcualtor_graph;
|
||||||
|
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(DisallowIf, VerifyConfig) {
|
||||||
|
mediapipe::api2::builder::Graph graph;
|
||||||
|
|
||||||
|
Source<int> value = graph.In("VALUE").Cast<int>();
|
||||||
|
Source<bool> condition = graph.In("CONDITION").Cast<bool>();
|
||||||
|
|
||||||
|
auto gated_stream = DisallowIf(value, condition, graph);
|
||||||
|
gated_stream.SetName("gated_stream");
|
||||||
|
|
||||||
|
EXPECT_THAT(graph.GetConfig(),
|
||||||
|
testing::EqualsProto(
|
||||||
|
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||||
|
node {
|
||||||
|
calculator: "GateCalculator"
|
||||||
|
input_stream: "__stream_1"
|
||||||
|
input_stream: "DISALLOW:__stream_0"
|
||||||
|
output_stream: "gated_stream"
|
||||||
|
options {
|
||||||
|
[mediapipe.GateCalculatorOptions.ext] {
|
||||||
|
empty_packets_as_allow: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
input_stream: "CONDITION:__stream_0"
|
||||||
|
input_stream: "VALUE:__stream_1"
|
||||||
|
)pb")));
|
||||||
|
|
||||||
|
CalculatorGraph calcualtor_graph;
|
||||||
|
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(DisallowIf, VerifyConfigWithSideCondition) {
|
||||||
|
mediapipe::api2::builder::Graph graph;
|
||||||
|
|
||||||
|
Source<int> value = graph.In("VALUE").Cast<int>();
|
||||||
|
SideSource<bool> condition = graph.SideIn("CONDITION").Cast<bool>();
|
||||||
|
|
||||||
|
auto gated_stream = DisallowIf(value, condition, graph);
|
||||||
|
gated_stream.SetName("gated_stream");
|
||||||
|
|
||||||
|
EXPECT_THAT(graph.GetConfig(),
|
||||||
|
testing::EqualsProto(
|
||||||
|
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||||
|
node {
|
||||||
|
calculator: "GateCalculator"
|
||||||
|
input_stream: "__stream_0"
|
||||||
|
output_stream: "gated_stream"
|
||||||
|
input_side_packet: "DISALLOW:__side_packet_1"
|
||||||
|
options {
|
||||||
|
[mediapipe.GateCalculatorOptions.ext] {
|
||||||
|
empty_packets_as_allow: true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
input_stream: "VALUE:__stream_0"
|
||||||
|
input_side_packet: "CONDITION:__side_packet_1"
|
||||||
|
)pb")));
|
||||||
|
|
||||||
|
CalculatorGraph calcualtor_graph;
|
||||||
|
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(AllowGate, VerifyConfig) {
|
||||||
|
mediapipe::api2::builder::Graph graph;
|
||||||
|
|
||||||
|
Source<bool> condition = graph.In("CONDITION").Cast<bool>();
|
||||||
|
Source<int> value1 = graph.In("VALUE_1").Cast<int>();
|
||||||
|
Source<int> value2 = graph.In("VALUE_2").Cast<int>();
|
||||||
|
Source<int> value3 = graph.In("VALUE_3").Cast<int>();
|
||||||
|
|
||||||
|
AllowGate gate(condition, graph);
|
||||||
|
gate.Allow(value1).SetName("gated_stream1");
|
||||||
|
gate.Allow(value2).SetName("gated_stream2");
|
||||||
|
gate.Allow(value3).SetName("gated_stream3");
|
||||||
|
|
||||||
|
EXPECT_THAT(graph.GetConfig(),
|
||||||
|
testing::EqualsProto(
|
||||||
|
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||||
|
node {
|
||||||
|
calculator: "GateCalculator"
|
||||||
|
input_stream: "__stream_1"
|
||||||
|
input_stream: "__stream_2"
|
||||||
|
input_stream: "__stream_3"
|
||||||
|
input_stream: "ALLOW:__stream_0"
|
||||||
|
output_stream: "gated_stream1"
|
||||||
|
output_stream: "gated_stream2"
|
||||||
|
output_stream: "gated_stream3"
|
||||||
|
}
|
||||||
|
input_stream: "CONDITION:__stream_0"
|
||||||
|
input_stream: "VALUE_1:__stream_1"
|
||||||
|
input_stream: "VALUE_2:__stream_2"
|
||||||
|
input_stream: "VALUE_3:__stream_3"
|
||||||
|
)pb")));
|
||||||
|
|
||||||
|
CalculatorGraph calcualtor_graph;
|
||||||
|
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(AllowIf, VerifyConfig) {
|
||||||
|
mediapipe::api2::builder::Graph graph;
|
||||||
|
|
||||||
|
Source<int> value = graph.In("VALUE").Cast<int>();
|
||||||
|
Source<bool> condition = graph.In("CONDITION").Cast<bool>();
|
||||||
|
|
||||||
|
auto gated_stream = AllowIf(value, condition, graph);
|
||||||
|
gated_stream.SetName("gated_stream");
|
||||||
|
|
||||||
|
EXPECT_THAT(graph.GetConfig(),
|
||||||
|
testing::EqualsProto(
|
||||||
|
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||||
|
node {
|
||||||
|
calculator: "GateCalculator"
|
||||||
|
input_stream: "__stream_1"
|
||||||
|
input_stream: "ALLOW:__stream_0"
|
||||||
|
output_stream: "gated_stream"
|
||||||
|
}
|
||||||
|
input_stream: "CONDITION:__stream_0"
|
||||||
|
input_stream: "VALUE:__stream_1"
|
||||||
|
)pb")));
|
||||||
|
|
||||||
|
CalculatorGraph calcualtor_graph;
|
||||||
|
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||||
|
}
|
||||||
|
|
||||||
|
TEST(AllowIf, VerifyConfigWithSideConition) {
|
||||||
|
mediapipe::api2::builder::Graph graph;
|
||||||
|
|
||||||
|
Source<int> value = graph.In("VALUE").Cast<int>();
|
||||||
|
SideSource<bool> condition = graph.SideIn("CONDITION").Cast<bool>();
|
||||||
|
|
||||||
|
auto gated_stream = AllowIf(value, condition, graph);
|
||||||
|
gated_stream.SetName("gated_stream");
|
||||||
|
|
||||||
|
EXPECT_THAT(graph.GetConfig(),
|
||||||
|
testing::EqualsProto(
|
||||||
|
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
|
||||||
|
node {
|
||||||
|
calculator: "GateCalculator"
|
||||||
|
input_stream: "__stream_0"
|
||||||
|
output_stream: "gated_stream"
|
||||||
|
input_side_packet: "ALLOW:__side_packet_1"
|
||||||
|
}
|
||||||
|
input_stream: "VALUE:__stream_0"
|
||||||
|
input_side_packet: "CONDITION:__side_packet_1"
|
||||||
|
)pb")));
|
||||||
|
|
||||||
|
CalculatorGraph calcualtor_graph;
|
||||||
|
MP_EXPECT_OK(calcualtor_graph.Initialize(graph.GetConfig()));
|
||||||
|
}
|
||||||
|
|
||||||
|
} // namespace
|
||||||
|
} // namespace utils
|
||||||
|
} // namespace components
|
||||||
|
} // namespace tasks
|
||||||
|
} // namespace mediapipe
|
||||||
Loading…
Reference in New Issue
Block a user