Add gate utility functions.
PiperOrigin-RevId: 478026407
This commit is contained in:
parent
af2ad1abbe
commit
46a5117c6d
|
@ -42,3 +42,16 @@ cc_test(
|
|||
"//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