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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 1e13be30e2c6838d4a2ff768a39c414bc80534bb
This commit is contained in:
MediaPipe Team
2022-09-06 21:46:17 +00:00
committed by Sebastian Schmidt
parent 63e679d99c
commit 4dc4b19ddb
639 changed files with 71327 additions and 2078 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/tool/BUILD
+27 -6
View File
@@ -53,6 +53,7 @@ bzl_library(
"build_defs.bzl",
],
visibility = [
"//mediapipe/app/pursuit/wasm:__subpackages__",
"//mediapipe/app/xeno/catalog:__subpackages__",
"//mediapipe/framework:__subpackages__",
],
@@ -768,7 +769,6 @@ cc_library(
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:gtest",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
@@ -829,24 +829,40 @@ cc_library(
],
deps = [
":container_util",
":options_util",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_shard",
"//mediapipe/framework:output_stream_shard",
"//mediapipe/framework/deps:mathutil",
"//mediapipe/framework/formats:video_stream_header",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
"//mediapipe/framework/tool:switch_container_cc_proto",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
mediapipe_cc_test(
name = "switch_mux_calculator_test",
srcs = ["switch_mux_calculator_test.cc"],
deps = [
":container_util",
":switch_mux_calculator",
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:subgraph",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
"@com_google_absl//absl/strings",
],
)
mediapipe_proto_library(
name = "switch_container_proto",
srcs = ["switch_container.proto"],
@@ -867,13 +883,17 @@ cc_library(
":subgraph_expansion",
":switch_demux_calculator",
":switch_mux_calculator",
"//mediapipe/calculators/core:packet_sequencer_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:mediapipe_options_cc_proto",
"//mediapipe/framework:calculator_options_cc_proto",
"//mediapipe/framework:subgraph",
"//mediapipe/framework/deps:registration",
"//mediapipe/framework/port:core_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:switch_container_cc_proto",
"@com_google_absl//absl/strings",
],
alwayslink = 1,
)
@@ -890,6 +910,7 @@ cc_test(
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:stream_handler_cc_proto",
"//mediapipe/framework:subgraph",
"//mediapipe/framework:test_calculators",
"//mediapipe/framework/port:gtest_main",
mediapipe/framework/tool/mediapipe_files.bzl
+23
View File
@@ -0,0 +1,23 @@
"""Build rule to depend on files downloaded from GCS."""
# buildifier: disable=unnamed-macro
def mediapipe_files(srcs):
"""Links file from GCS with the current directory.
Args:
srcs: the names of the mediapipe_file target, which is also the name of
the MediaPipe file in external_files.bzl. For example, if `name` is Foo,
`mediapipe_file` will create a link to the downloaded file
"@com_google_mediapipe_Foo_tfile" to the current directory as
"Foo.tflite".
"""
for src in srcs:
archive_name = "com_google_mediapipe_%s" % src.replace("/", "_").replace(".", "_")
native.genrule(
name = "%s_ln" % archive_name,
srcs = ["@%s//file" % archive_name],
outs = [src],
output_to_bindir = 1,
cmd = "ln $< $@",
)
mediapipe/framework/tool/options_field_util.cc
+3 -3
View File
@@ -380,7 +380,7 @@ absl::StatusOr<std::vector<FieldData>> GetFieldValues(
int index = tail.empty() ? head.index : std::max(0, head.index);
if ((int)results.size() <= index) {
return absl::OutOfRangeError(absl::StrCat(
"Missing feild value: ", head.field ? head.field->name() : "#",
"Missing field value: ", head.field ? head.field->name() : "#",
" at index: ", index));
}
if (!tail.empty()) {
@@ -402,7 +402,7 @@ absl::StatusOr<FieldData> GetField(const FieldData& message_data,
if (results.empty()) {
FieldPathEntry tail = field_path.back();
return absl::OutOfRangeError(absl::StrCat(
"Missing feild value: ", tail.field ? tail.field->name() : "##",
"Missing field value: ", tail.field ? tail.field->name() : "##",
" at index: ", tail.index));
}
return results[0];
@@ -414,7 +414,7 @@ absl::Status SetFieldValues(FieldData& message_data,
const std::vector<FieldData>& values) {
if (field_path.empty()) {
if (values.empty()) {
return absl::InvalidArgumentError("Missing feild value.");
return absl::InvalidArgumentError("Missing field value.");
}
message_data = values[0];
return absl::OkStatus();
mediapipe/framework/tool/options_util_test.cc
+2 -2
View File
@@ -345,7 +345,7 @@ absl::Status Equals(std::vector<FieldData> b1, std::vector<FieldData> b2) {
return absl::OkStatus();
}
// Unit-tests for graph options feild accessors from options_field_util.
// Unit-tests for graph options field accessors from options_field_util.
class OptionsFieldUtilTest : public ::testing::Test {
protected:
void SetUp() override {}
@@ -581,7 +581,7 @@ TEST_F(OptionsFieldUtilTest, MergeFieldValuesInt) {
MergeFieldValues(node_data, path, {AsFieldData(55), AsFieldData(66)});
EXPECT_EQ(status.code(), absl::StatusCode::kOutOfRange);
EXPECT_THAT(status.message(),
HasSubstr("Missing feild value: num_lights at index: 1"));
HasSubstr("Missing field value: num_lights at index: 1"));
}
// Tests MergeFieldValues applied to a protobuf field.
mediapipe/framework/tool/switch_container.cc
+64 -9
View File
@@ -12,13 +12,18 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <fstream>
#include <iostream>
#include <sstream>
#include <functional>
#include <map>
#include <memory>
#include <set>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/mediapipe_options.pb.h"
#include "mediapipe/framework/calculator_options.pb.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
@@ -98,6 +103,18 @@ CalculatorGraphConfig::Node* BuildMuxNode(
return result;
}
// Returns a PacketSequencerCalculator node.
CalculatorGraphConfig::Node* BuildTimestampNode(CalculatorGraphConfig* config,
bool synchronize_io) {
CalculatorGraphConfig::Node* result = config->add_node();
*result->mutable_calculator() = "PacketSequencerCalculator";
if (synchronize_io) {
*result->mutable_input_stream_handler()->mutable_input_stream_handler() =
"DefaultInputStreamHandler";
}
return result;
}
// Copies options from one node to another.
void CopyOptions(const CalculatorGraphConfig::Node& source,
CalculatorGraphConfig::Node* dest) {
@@ -214,6 +231,14 @@ absl::Status ValidateContract(
return absl::OkStatus();
}
// Returns true if a set of streams references a certain tag name.
bool HasTag(const proto_ns::RepeatedPtrField<std::string>& streams,
std::string tag) {
std::map<TagIndex, std::string> tags;
ParseTags(streams, &tags);
return tags.count({tag, 0}) > 0;
}
absl::StatusOr<CalculatorGraphConfig> SwitchContainer::GetConfig(
const Subgraph::SubgraphOptions& options) {
CalculatorGraphConfig config;
@@ -232,32 +257,62 @@ absl::StatusOr<CalculatorGraphConfig> SwitchContainer::GetConfig(
ParseTags(container_streams.input_side_packet(), &side_input_tags);
ParseTags(container_streams.output_side_packet(), &side_output_tags);
CalculatorGraphConfig::Node* select_node = nullptr;
CalculatorGraphConfig::Node* enable_node = nullptr;
std::string select_stream = "SELECT:gate_select";
std::string enable_stream = "ENABLE:gate_enable";
// Add a PacketSequencerCalculator node for "SELECT" or "ENABLE" streams.
bool synchronize_io =
Subgraph::GetOptions<mediapipe::SwitchContainerOptions>(options)
.synchronize_io();
if (HasTag(container_node.input_stream(), "SELECT")) {
select_node = BuildTimestampNode(&config, synchronize_io);
select_node->add_input_stream("INPUT:gate_select");
select_node->add_output_stream("OUTPUT:gate_select_timed");
select_stream = "SELECT:gate_select_timed";
}
if (HasTag(container_node.input_stream(), "ENABLE")) {
enable_node = BuildTimestampNode(&config, synchronize_io);
enable_node->add_input_stream("INPUT:gate_enable");
enable_node->add_output_stream("OUTPUT:gate_enable_timed");
enable_stream = "ENABLE:gate_enable_timed";
}
// Add a graph node for the demux, mux.
auto demux = BuildDemuxNode(input_tags, container_node, &config);
CopyOptions(container_node, demux);
ClearContainerOptions(demux);
demux->add_input_stream("SELECT:gate_select");
demux->add_input_stream("ENABLE:gate_enable");
demux->add_input_stream(select_stream);
demux->add_input_stream(enable_stream);
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_stream(select_stream);
mux->add_input_stream(enable_stream);
mux->add_input_side_packet("SELECT:gate_select");
mux->add_input_side_packet("ENABLE:gate_enable");
// Add input streams for graph and demux.
// Add input streams for graph and demux and the timestamper.
config.add_input_stream("SELECT:gate_select");
config.add_input_stream("ENABLE:gate_enable");
config.add_input_side_packet("SELECT:gate_select");
config.add_input_side_packet("ENABLE:gate_enable");
int tick_index = 0;
for (const auto& p : input_tags) {
std::string stream = CatStream(p.first, p.second);
config.add_input_stream(stream);
demux->add_input_stream(stream);
TagIndex tick_tag{"TICK", tick_index++};
if (select_node) {
select_node->add_input_stream(CatStream(tick_tag, p.second));
}
if (enable_node) {
enable_node->add_input_stream(CatStream(tick_tag, p.second));
}
}
// Add output streams for graph and mux.
mediapipe/framework/tool/switch_container_test.cc
+22 -5
View File
@@ -12,7 +12,13 @@
// See the License for the specific language governing permissions and
// limitations under the License.
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "absl/strings/str_replace.h"
#include "absl/strings/string_view.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/gmock.h"
@@ -161,7 +167,9 @@ void RunTestContainer(CalculatorGraphConfig supergraph,
// i.e. the one containing the PassThroughCalculator should output the
// input values without changing them.
EXPECT_EQ(out_bar.size(), t);
EXPECT_EQ(out_bar.back().Get<int>(), t);
if (!out_bar.empty()) {
EXPECT_EQ(out_bar.back().Get<int>(), t);
}
}
if (!send_bounds) {
@@ -187,7 +195,9 @@ void RunTestContainer(CalculatorGraphConfig supergraph,
// i.e. the one containing the TripleIntCalculator should output the values
// after tripling them.
EXPECT_EQ(out_bar.size(), t);
EXPECT_EQ(out_bar.back().Get<int>(), t * 3);
if (!out_bar.empty()) {
EXPECT_EQ(out_bar.back().Get<int>(), t * 3);
}
}
MP_ASSERT_OK(graph.CloseAllInputStreams());
@@ -236,10 +246,17 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
CalculatorGraphConfig supergraph = SubnodeContainerExample();
CalculatorGraphConfig expected_graph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
node {
name: "switchcontainer__PacketSequencerCalculator"
calculator: "PacketSequencerCalculator"
input_stream: "INPUT:enable"
input_stream: "TICK:foo"
output_stream: "OUTPUT:switchcontainer__gate_enable_timed"
}
node {
name: "switchcontainer__SwitchDemuxCalculator"
calculator: "SwitchDemuxCalculator"
input_stream: "ENABLE:enable"
input_stream: "ENABLE:switchcontainer__gate_enable_timed"
input_stream: "foo"
output_stream: "C0__:switchcontainer__c0__foo"
output_stream: "C1__:switchcontainer__c1__foo"
@@ -262,7 +279,7 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
node {
name: "switchcontainer__SwitchMuxCalculator"
calculator: "SwitchMuxCalculator"
input_stream: "ENABLE:enable"
input_stream: "ENABLE:switchcontainer__gate_enable_timed"
input_stream: "C0__:switchcontainer__c0__bar"
input_stream: "C1__:switchcontainer__c1__bar"
output_stream: "bar"
@@ -281,7 +298,7 @@ TEST(SwitchContainerTest, ApplyToSubnodes) {
input_stream: "enable"
input_side_packet: "timezone"
)pb");
expected_graph = OrderNodes(expected_graph, {4, 0, 3, 1, 2});
expected_graph = OrderNodes(expected_graph, {5, 0, 1, 4, 2, 3});
MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
mediapipe/framework/tool/switch_mux_calculator.cc
+116 -82
View File
@@ -13,11 +13,15 @@
// limitations under the License.
#include <algorithm>
#include <memory>
#include <functional>
#include <iterator>
#include <map>
#include <queue>
#include <set>
#include <string>
#include <type_traits>
#include <utility>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/input_stream_shard.h"
@@ -66,20 +70,26 @@ class SwitchMuxCalculator : public CalculatorBase {
absl::Status Open(CalculatorContext* cc) override;
absl::Status Process(CalculatorContext* cc) override;
private:
// Stores any new input channel history.
void RecordChannel(CalculatorContext* cc);
// Temporarily enqueues every new packet or timestamp bounds.
void RecordPackets(CalculatorContext* cc);
// Immediately sends any packets or timestamp bounds for settled timestamps.
void SendActivePackets(CalculatorContext* cc);
private:
int channel_index_;
std::set<std::string> channel_tags_;
mediapipe::SwitchContainerOptions options_;
// This is used to keep around packets that we've received but not
// relayed yet (because we may not know which channel we should yet be using
// when synchronized_io flag is set).
std::map<Timestamp, std::map<CollectionItemId, Packet>> packet_history_;
// relayed yet (because we may not know which channel we should yet be using).
std::map<CollectionItemId, std::queue<Packet>> packet_queue_;
// Historical channel index values for timestamps where we don't have all
// packets available yet (when synchronized_io flag is set).
// packets available yet.
std::map<Timestamp, int> channel_history_;
// Number of output steams that we already processed for the current output
// timestamp.
int current_processed_stream_count_ = 0;
};
REGISTER_CALCULATOR(SwitchMuxCalculator);
@@ -133,10 +143,28 @@ absl::Status SwitchMuxCalculator::GetContract(CalculatorContract* cc) {
return absl::OkStatus();
}
// Returns the last delivered timestamp for an input stream.
Timestamp SettledTimestamp(const InputStreamShard& input) {
return input.Value().Timestamp();
}
// Returns the last delivered timestamp for channel selection.
Timestamp ChannelSettledTimestamp(CalculatorContext* cc) {
Timestamp result = Timestamp::Done();
if (cc->Inputs().HasTag("ENABLE")) {
result = SettledTimestamp(cc->Inputs().Tag("ENABLE"));
} else if (cc->Inputs().HasTag("SELECT")) {
result = SettledTimestamp(cc->Inputs().Tag("SELECT"));
}
return result;
}
absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) {
// Initialize channel_index_ and channel_history_.
options_ = cc->Options<mediapipe::SwitchContainerOptions>();
channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
channel_tags_ = ChannelTags(cc->Inputs().TagMap());
channel_history_[Timestamp::Unset()] = channel_index_;
// Relay side packets only from channel_index_.
for (const std::string& tag : ChannelTags(cc->InputSidePackets().TagMap())) {
@@ -150,85 +178,91 @@ absl::Status SwitchMuxCalculator::Open(CalculatorContext* cc) {
return absl::OkStatus();
}
absl::Status SwitchMuxCalculator::Process(CalculatorContext* cc) {
// Update the input channel index if specified.
channel_index_ = tool::GetChannelIndex(*cc, channel_index_);
void SwitchMuxCalculator::RecordChannel(CalculatorContext* cc) {
Timestamp channel_settled = ChannelSettledTimestamp(cc);
int new_channel_index = tool::GetChannelIndex(*cc, channel_index_);
if (options_.synchronize_io()) {
// Start with adding input signals into channel_history_ and packet_history_
if (cc->Inputs().HasTag("ENABLE") &&
!cc->Inputs().Tag("ENABLE").IsEmpty()) {
channel_history_[cc->Inputs().Tag("ENABLE").Value().Timestamp()] =
channel_index_;
// Enque any new input channel and its activation timestamp.
if (channel_settled == cc->InputTimestamp() &&
new_channel_index != channel_index_) {
channel_index_ = new_channel_index;
channel_history_[channel_settled] = channel_index_;
}
}
void SwitchMuxCalculator::RecordPackets(CalculatorContext* cc) {
auto select_id = cc->Inputs().GetId("SELECT", 0);
auto enable_id = cc->Inputs().GetId("ENABLE", 0);
for (auto id = cc->Inputs().BeginId(); id < cc->Inputs().EndId(); ++id) {
if (id == select_id || id == enable_id) continue;
Packet packet = cc->Inputs().Get(id).Value();
// Enque any new packet or timestamp bound.
if (packet.Timestamp() == cc->InputTimestamp()) {
packet_queue_[id].push(packet);
}
if (cc->Inputs().HasTag("SELECT") &&
!cc->Inputs().Tag("SELECT").IsEmpty()) {
channel_history_[cc->Inputs().Tag("SELECT").Value().Timestamp()] =
channel_index_;
}
for (auto input_id = cc->Inputs().BeginId();
input_id < cc->Inputs().EndId(); ++input_id) {
auto& entry = cc->Inputs().Get(input_id);
if (entry.IsEmpty()) {
continue;
}
packet_history_[entry.Value().Timestamp()][input_id] = entry.Value();
}
// Now check if we have enough information to produce any outputs.
while (!channel_history_.empty()) {
// Look at the oldest unprocessed timestamp.
auto it = channel_history_.begin();
auto& packets = packet_history_[it->first];
int total_streams = 0;
// Loop over all outputs to see if we have anything new that we can relay.
for (const std::string& tag : channel_tags_) {
for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
++total_streams;
auto input_id =
cc->Inputs().GetId(tool::ChannelTag(tag, it->second), index);
auto packet_it = packets.find(input_id);
if (packet_it != packets.end()) {
cc->Outputs().Get(tag, index).AddPacket(packet_it->second);
++current_processed_stream_count_;
} else if (it->first <
cc->Inputs().Get(input_id).Value().Timestamp()) {
// Getting here means that input stream that corresponds to this
// output at the timestamp we're trying to process right now has
// already advanced beyond this timestamp. This means that we will
// shouldn't expect a packet for this timestamp anymore, and we can
// safely advance timestamp on the output.
cc->Outputs()
.Get(tag, index)
.SetNextTimestampBound(it->first.NextAllowedInStream());
++current_processed_stream_count_;
}
}
}
if (current_processed_stream_count_ == total_streams) {
// There's nothing else to wait for at the current timestamp, do the
// cleanup and move on to the next one.
packet_history_.erase(it->first);
channel_history_.erase(it);
current_processed_stream_count_ = 0;
} else {
// We're still missing some packets for the current timestamp. Clean up
// those that we just relayed and let the rest wait until the next
// Process() call.
packets.clear();
break;
}
}
} else {
// Relay packets and timestamps only from channel_index_.
}
}
void SwitchMuxCalculator::SendActivePackets(CalculatorContext* cc) {
Timestamp expired_history;
// Iterate through the recent active input channels.
for (auto it = channel_history_.begin(); it != channel_history_.end(); ++it) {
int channel = it->second;
Timestamp channel_start = it->first;
Timestamp channel_end =
(std::next(it) == channel_history_.end())
? ChannelSettledTimestamp(cc).NextAllowedInStream()
: std::next(it)->first;
Timestamp stream_settled = Timestamp::Done();
for (const std::string& tag : channel_tags_) {
for (int index = 0; index < cc->Outputs().NumEntries(tag); ++index) {
auto& output = cc->Outputs().Get(tag, index);
std::string input_tag = tool::ChannelTag(tag, channel_index_);
auto& input = cc->Inputs().Get(input_tag, index);
tool::Relay(input, &output);
std::string input_tag = tool::ChannelTag(tag, channel);
for (int index = 0; index < cc->Inputs().NumEntries(input_tag); ++index) {
CollectionItemId input_id = cc->Inputs().GetId(input_tag, index);
OutputStreamShard& output = cc->Outputs().Get(tag, index);
std::queue<Packet>& q = packet_queue_[input_id];
// Send any packets or bounds from a recent active input channel.
while (!q.empty() && q.front().Timestamp() < channel_end) {
if (q.front().Timestamp() >= channel_start) {
output.AddPacket(q.front());
}
q.pop();
}
stream_settled = std::min(stream_settled,
SettledTimestamp(cc->Inputs().Get(input_id)));
}
}
// A history entry is expired only if all streams have advanced past it.
if (stream_settled.NextAllowedInStream() < channel_end ||
std::next(it) == channel_history_.end()) {
break;
}
expired_history = channel_start;
// Discard any packets or bounds from recent inactive input channels.
for (auto id = cc->Inputs().BeginId(); id < cc->Inputs().EndId(); ++id) {
std::queue<Packet>& q = packet_queue_[id];
while (!q.empty() && q.front().Timestamp() < channel_end) {
q.pop();
}
}
}
// Discard any expired channel history entries.
if (expired_history != Timestamp::Unset()) {
channel_history_.erase(channel_history_.begin(),
std::next(channel_history_.find(expired_history)));
}
}
absl::Status SwitchMuxCalculator::Process(CalculatorContext* cc) {
// Normally packets will arrive on the active channel and will be passed
// through immediately. In the less common case in which the active input
// channel is not known for an input packet timestamp, the input packet is
// queued until the active channel becomes known.
RecordChannel(cc);
RecordPackets(cc);
SendActivePackets(cc);
return absl::OkStatus();
}
mediapipe/framework/tool/switch_mux_calculator_test.cc
+329
View File
@@ -0,0 +1,329 @@
// Copyright 2022 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 <string>
#include <vector>
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.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/ret_check.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/tool/container_util.h"
namespace mediapipe {
namespace {
// Returns a CalculatorGraph to run a single calculator.
CalculatorGraph BuildCalculatorGraph(CalculatorGraphConfig::Node node_config) {
CalculatorGraphConfig config;
*config.add_node() = node_config;
*config.mutable_input_stream() = node_config.input_stream();
*config.mutable_output_stream() = node_config.output_stream();
*config.mutable_input_side_packet() = node_config.input_side_packet();
*config.mutable_output_side_packet() = node_config.output_side_packet();
return CalculatorGraph(config);
}
// Creates a string packet.
Packet pack(std::string data, int timestamp) {
return MakePacket<std::string>(data).At(Timestamp(timestamp));
}
// Creates an int packet.
Packet pack(int data, int timestamp) {
return MakePacket<int>(data).At(Timestamp(timestamp));
}
// Tests showing packet channel synchronization through SwitchMuxCalculator.
class SwitchMuxCalculatorTest : public ::testing::Test {
protected:
SwitchMuxCalculatorTest() {}
~SwitchMuxCalculatorTest() override {}
void SetUp() override {}
void TearDown() override {}
// Defines a SwitchMuxCalculator CalculatorGraphConfig::Node.
CalculatorGraphConfig::Node BuildNodeConfig() {
CalculatorGraphConfig::Node result;
*result.mutable_calculator() = "SwitchMuxCalculator";
*result.add_input_stream() = "SELECT:select";
for (int c = 0; c < 3; ++c) {
*result.add_input_stream() =
absl::StrCat(tool::ChannelTag("FRAME", c), ":frame_", c);
*result.add_input_stream() =
absl::StrCat(tool::ChannelTag("MASK", c), ":mask_", c);
}
*result.add_output_stream() = "FRAME:frame";
*result.add_output_stream() = "MASK:mask";
return result;
}
};
// Shows the SwitchMuxCalculator is available.
TEST_F(SwitchMuxCalculatorTest, IsRegistered) {
EXPECT_TRUE(CalculatorBaseRegistry::IsRegistered("SwitchMuxCalculator"));
}
// Shows that channels are queued until packets arrive.
TEST_F(SwitchMuxCalculatorTest, ChannelEarly) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// All channels are specified before any frame packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
// The packet at timestamp 10 is passed from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t10", 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p0_t10");
// The packet at timestamp 20 is passed from channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 2);
EXPECT_EQ(output_frames[1].Get<std::string>(), "p1_t20");
// The packet at timestamp 30 is passed from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 3);
EXPECT_EQ(output_frames[2].Get<std::string>(), "p0_t30");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows that packets are queued until channels are specified.
TEST_F(SwitchMuxCalculatorTest, ChannelsLate) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// All frame packets arrive before any channels are specified.
// All packets are queued awaiting channel choices.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t10", 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 0);
// The packet at timestamp 10 is released from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p0_t10");
// The packet at timestamp 20 is released from channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 2);
EXPECT_EQ(output_frames[1].Get<std::string>(), "p1_t20");
// The packet at timestamp 30 is released from channel 0.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 3);
EXPECT_EQ(output_frames[2].Get<std::string>(), "p0_t30");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows that no packets are queued when channels are specified just in time.
TEST_F(SwitchMuxCalculatorTest, ChannelsOnTime) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// Channel 0 is selected just before corresponding packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t10", 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p0_t10");
// Channel 1 is selected just before corresponding packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 2);
EXPECT_EQ(output_frames[1].Get<std::string>(), "p1_t20");
// Channel 0 is selected just before corresponding packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 30)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 3);
EXPECT_EQ(output_frames[2].Get<std::string>(), "p0_t30");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows what happens if the last packet from a channel never arrives.
TEST_F(SwitchMuxCalculatorTest, ChannelNeverCompletes) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// Channel 0 is selected, but it's closing packet never arrives.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 0);
// Channel 1 is selected, but we still wait for channel 0 to finish.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 0);
// Finally channel 0 advances, and channel 1 can be delivered.
// Note that "p0_t15" is discarded because its channel is deselected.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t15", 15)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
EXPECT_EQ(output_frames[0].Get<std::string>(), "p1_t20");
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows how output is sequenced when one channel is delayed substantially.
// Queues up "SELECT" packets for channel 0, 1, and 2.
// Queues up "frame" packets for channel 0 and 2.
// The output packets from channel 1, 2, and 0 wait for channel 1.
TEST_F(SwitchMuxCalculatorTest, OneChannelIsSlow) {
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
CalculatorGraph graph = BuildCalculatorGraph(node_config);
std::vector<Packet> output_frames;
MP_ASSERT_OK(graph.ObserveOutputStream("frame", [&](const Packet& p) {
output_frames.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Finalize input for the "mask" input stream.
MP_ASSERT_OK(graph.CloseInputStream("mask_0"));
MP_ASSERT_OK(graph.CloseInputStream("mask_1"));
MP_ASSERT_OK(graph.CloseInputStream("mask_2"));
// Channel 0 is selected, and some packets arrive.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 1)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 10)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t20", 10)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 1 is selected, but its packets are delayed.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 11)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t20", 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_2", pack("p2_t20", 20)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 2 is selected, packets arrive, but wait for channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(2, 21)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(2, 30)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_2", pack("p2_t30", 30)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 0 is selected again, packets arrive, but wait for channel 1.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 31)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(0, 40)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_0", pack("p0_t40", 40)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Channel 1 is selected again, but its packets are still delayed.
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 41)));
MP_ASSERT_OK(graph.AddPacketToInputStream("select", pack(1, 50)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 1);
// Finally, the delayed packets from channel 1 arrive.
// And all packets for all five "SELECT"" inetervals are delivered.
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t20", 20)));
MP_ASSERT_OK(graph.AddPacketToInputStream("frame_1", pack("p1_t50", 50)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(output_frames.size(), 5);
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/template_expander.cc
+1 -1
View File
@@ -594,7 +594,7 @@ class TemplateExpanderImpl {
return AsArgument(static_cast<double>(b));
}
// Convert between a proto feild value and a template argument.
// Convert between a proto field value and a template argument.
absl::Status AsFieldValues(const std::vector<TemplateArgument>& args,
FieldType field_type,
std::vector<FieldValue>* result) {
mediapipe/framework/tool/test_util.cc
-1
View File
@@ -34,7 +34,6 @@
#include "mediapipe/framework/formats/image_format.pb.h"
#include "mediapipe/framework/port/advanced_proto_inc.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/proto_ns.h"
#include "mediapipe/framework/port/ret_check.h"