cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Internal change

Browse Source 
PiperOrigin-RevId: 478093259
This commit is contained in:
MediaPipe Team 2022-09-30 15:25:57 -07:00 committed by Copybara-Service
parent 46a5117c6d
commit a3dc91fafe
1 changed files with 32 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
mediapipe/tasks/cc/components/utils/gate.h
View File

@ -26,15 +26,12 @@ namespace tasks {
namespace components { namespace components {
namespace utils { namespace utils {
using ::mediapipe::api2::builder::SideSource;
using ::mediapipe::api2::builder::Source;
// Utility class that simplifies allowing (gating) multiple streams. // Utility class that simplifies allowing (gating) multiple streams.
class AllowGate { class AllowGate {
public: public:
AllowGate(Source<bool> allow, mediapipe::api2::builder::Graph& graph) AllowGate(api2::builder::Source<bool> allow, api2::builder::Graph& graph)
: node_(AddSourceGate(allow, graph)) {} : node_(AddSourceGate(allow, graph)) {}
AllowGate(SideSource<bool> allow, mediapipe::api2::builder::Graph& graph) AllowGate(api2::builder::SideSource<bool> allow, api2::builder::Graph& graph)
: node_(AddSideSourceGate(allow, graph)) {} : node_(AddSideSourceGate(allow, graph)) {}
// Move-only // Move-only
@ -42,39 +39,40 @@ class AllowGate {
AllowGate& operator=(AllowGate&& allow_gate) = default; AllowGate& operator=(AllowGate&& allow_gate) = default;
template <typename T> template <typename T>
Source<T> Allow(Source<T> source) { api2::builder::Source<T> Allow(api2::builder::Source<T> source) {
source >> node_.In(index_); source >> node_.In(index_);
return node_.Out(index_++).Cast<T>(); return node_.Out(index_++).Cast<T>();
} }
private: private:
template <typename T> template <typename T>
static mediapipe::api2::builder::GenericNode& AddSourceGate( static api2::builder::GenericNode& AddSourceGate(
T allow, mediapipe::api2::builder::Graph& graph) { T allow, api2::builder::Graph& graph) {
auto& gate_node = graph.AddNode("GateCalculator"); auto& gate_node = graph.AddNode("GateCalculator");
allow >> gate_node.In("ALLOW"); allow >> gate_node.In("ALLOW");
return gate_node; return gate_node;
} }
template <typename T> template <typename T>
static mediapipe::api2::builder::GenericNode& AddSideSourceGate( static api2::builder::GenericNode& AddSideSourceGate(
T allow, mediapipe::api2::builder::Graph& graph) { T allow, api2::builder::Graph& graph) {
auto& gate_node = graph.AddNode("GateCalculator"); auto& gate_node = graph.AddNode("GateCalculator");
allow >> gate_node.SideIn("ALLOW"); allow >> gate_node.SideIn("ALLOW");
return gate_node; return gate_node;
} }
mediapipe::api2::builder::GenericNode& node_; api2::builder::GenericNode& node_;
int index_ = 0; int index_ = 0;
}; };
// Utility class that simplifies disallowing (gating) multiple streams. // Utility class that simplifies disallowing (gating) multiple streams.
class DisallowGate { class DisallowGate {
public: public:
DisallowGate(Source<bool> disallow, mediapipe::api2::builder::Graph& graph) DisallowGate(api2::builder::Source<bool> disallow,
api2::builder::Graph& graph)
: node_(AddSourceGate(disallow, graph)) {} : node_(AddSourceGate(disallow, graph)) {}
DisallowGate(SideSource<bool> disallow, DisallowGate(api2::builder::SideSource<bool> disallow,
mediapipe::api2::builder::Graph& graph) api2::builder::Graph& graph)
: node_(AddSideSourceGate(disallow, graph)) {} : node_(AddSideSourceGate(disallow, graph)) {}
// Move-only // Move-only
@ -82,15 +80,15 @@ class DisallowGate {
DisallowGate& operator=(DisallowGate&& disallow_gate) = default; DisallowGate& operator=(DisallowGate&& disallow_gate) = default;
template <typename T> template <typename T>
Source<T> Disallow(Source<T> source) { api2::builder::Source<T> Disallow(api2::builder::Source<T> source) {
source >> node_.In(index_); source >> node_.In(index_);
return node_.Out(index_++).Cast<T>(); return node_.Out(index_++).Cast<T>();
} }
private: private:
template <typename T> template <typename T>
static mediapipe::api2::builder::GenericNode& AddSourceGate( static api2::builder::GenericNode& AddSourceGate(
T disallow, mediapipe::api2::builder::Graph& graph) { T disallow, api2::builder::Graph& graph) {
auto& gate_node = graph.AddNode("GateCalculator"); auto& gate_node = graph.AddNode("GateCalculator");
auto& gate_node_opts = auto& gate_node_opts =
gate_node.GetOptions<mediapipe::GateCalculatorOptions>(); gate_node.GetOptions<mediapipe::GateCalculatorOptions>();
@ -104,8 +102,8 @@ class DisallowGate {
} }
template <typename T> template <typename T>
static mediapipe::api2::builder::GenericNode& AddSideSourceGate( static api2::builder::GenericNode& AddSideSourceGate(
T disallow, mediapipe::api2::builder::Graph& graph) { T disallow, api2::builder::Graph& graph) {
auto& gate_node = graph.AddNode("GateCalculator"); auto& gate_node = graph.AddNode("GateCalculator");
auto& gate_node_opts = auto& gate_node_opts =
gate_node.GetOptions<mediapipe::GateCalculatorOptions>(); gate_node.GetOptions<mediapipe::GateCalculatorOptions>();
@ -114,39 +112,43 @@ class DisallowGate {
return gate_node; return gate_node;
} }
mediapipe::api2::builder::GenericNode& node_; api2::builder::GenericNode& node_;
int index_ = 0; int index_ = 0;
}; };
// Updates graph to drop @value stream packet if corresponding @condition stream // Updates graph to drop @value stream packet if corresponding @condition stream
// packet holds true. // packet holds true.
template <class T> template <class T>
Source<T> DisallowIf(Source<T> value, Source<bool> condition, api2::builder::Source<T> DisallowIf(api2::builder::Source<T> value,
mediapipe::api2::builder::Graph& graph) { api2::builder::Source<bool> condition,
api2::builder::Graph& graph) {
return DisallowGate(condition, graph).Disallow(value); return DisallowGate(condition, graph).Disallow(value);
} }
// Updates graph to drop @value stream packet if corresponding @condition stream // Updates graph to drop @value stream packet if corresponding @condition stream
// packet holds true. // packet holds true.
template <class T> template <class T>
Source<T> DisallowIf(Source<T> value, SideSource<bool> condition, api2::builder::Source<T> DisallowIf(api2::builder::Source<T> value,
mediapipe::api2::builder::Graph& graph) { api2::builder::SideSource<bool> condition,
api2::builder::Graph& graph) {
return DisallowGate(condition, graph).Disallow(value); return DisallowGate(condition, graph).Disallow(value);
} }
// Updates graph to pass through @value stream packet if corresponding // Updates graph to pass through @value stream packet if corresponding
// @condition stream packet holds true. // @allow stream packet holds true.
template <class T> template <class T>
Source<T> AllowIf(Source<T> value, Source<bool> allow, api2::builder::Source<T> AllowIf(api2::builder::Source<T> value,
mediapipe::api2::builder::Graph& graph) { api2::builder::Source<bool> allow,
api2::builder::Graph& graph) {
return AllowGate(allow, graph).Allow(value); return AllowGate(allow, graph).Allow(value);
} }
// Updates graph to pass through @value stream packet if corresponding // Updates graph to pass through @value stream packet if corresponding
// @condition side stream packet holds true. // @allow side stream packet holds true.
template <class T> template <class T>
Source<T> AllowIf(Source<T> value, SideSource<bool> allow, api2::builder::Source<T> AllowIf(api2::builder::Source<T> value,
mediapipe::api2::builder::Graph& graph) { api2::builder::SideSource<bool> allow,
api2::builder::Graph& graph) {
return AllowGate(allow, graph).Allow(value); return AllowGate(allow, graph).Allow(value);
} }