Internal change
PiperOrigin-RevId: 477538515
This commit is contained in:
MediaPipe Team
committed by
Sebastian Schmidt
Sebastian Schmidt
parent
6cdc6443b6
commit
f8af41b1eb
@@ -821,6 +821,19 @@ cc_library(
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
mediapipe_cc_test(
|
||||
name = "switch_demux_calculator_test",
|
||||
srcs = ["switch_demux_calculator_test.cc"],
|
||||
deps = [
|
||||
":container_util",
|
||||
":switch_demux_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "switch_mux_calculator",
|
||||
srcs = ["switch_mux_calculator.cc"],
|
||||
|
||||
@@ -129,12 +129,12 @@ absl::Status SwitchDemuxCalculator::Open(CalculatorContext* cc) {
|
||||
// Relay side packets to all channels.
|
||||
// Note: This is necessary because Calculator::Open only proceeds when every
|
||||
// anticipated side-packet arrives.
|
||||
int channel_count = tool::ChannelCount(cc->OutputSidePackets().TagMap());
|
||||
int side_channel_count = tool::ChannelCount(cc->OutputSidePackets().TagMap());
|
||||
for (const std::string& tag : ChannelTags(cc->OutputSidePackets().TagMap())) {
|
||||
int num_entries = cc->InputSidePackets().NumEntries(tag);
|
||||
for (int index = 0; index < num_entries; ++index) {
|
||||
Packet input = cc->InputSidePackets().Get(tag, index);
|
||||
for (int channel = 0; channel < channel_count; ++channel) {
|
||||
for (int channel = 0; channel < side_channel_count; ++channel) {
|
||||
std::string output_tag = tool::ChannelTag(tag, channel);
|
||||
auto output_id = cc->OutputSidePackets().GetId(output_tag, index);
|
||||
if (output_id.IsValid()) {
|
||||
@@ -143,6 +143,23 @@ absl::Status SwitchDemuxCalculator::Open(CalculatorContext* cc) {
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Relay headers to all channels.
|
||||
int output_channel_count = tool::ChannelCount(cc->Outputs().TagMap());
|
||||
for (const std::string& tag : ChannelTags(cc->Outputs().TagMap())) {
|
||||
int num_entries = cc->Inputs().NumEntries(tag);
|
||||
for (int index = 0; index < num_entries; ++index) {
|
||||
auto& input = cc->Inputs().Get(tag, index);
|
||||
if (input.Header().IsEmpty()) continue;
|
||||
for (int channel = 0; channel < output_channel_count; ++channel) {
|
||||
std::string output_tag = tool::ChannelTag(tag, channel);
|
||||
auto output_id = cc->Outputs().GetId(output_tag, index);
|
||||
if (output_id.IsValid()) {
|
||||
cc->Outputs().Get(output_tag, index).SetHeader(input.Header());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return absl::OkStatus();
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
// 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_framework.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/logging.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 SwitchDemuxCalculator.
|
||||
class SwitchDemuxCalculatorTest : public ::testing::Test {
|
||||
protected:
|
||||
SwitchDemuxCalculatorTest() {}
|
||||
~SwitchDemuxCalculatorTest() override {}
|
||||
void SetUp() override {}
|
||||
void TearDown() override {}
|
||||
|
||||
// Defines a SwitchDemuxCalculator CalculatorGraphConfig::Node.
|
||||
CalculatorGraphConfig::Node BuildNodeConfig() {
|
||||
CalculatorGraphConfig::Node result;
|
||||
*result.mutable_calculator() = "SwitchDemuxCalculator";
|
||||
*result.add_input_stream() = "SELECT:select";
|
||||
for (int c = 0; c < 2; ++c) {
|
||||
*result.add_output_stream() =
|
||||
absl::StrCat(tool::ChannelTag("FRAME", c), ":frame_", c);
|
||||
*result.add_output_stream() =
|
||||
absl::StrCat(tool::ChannelTag("MASK", c), ":mask_", c);
|
||||
}
|
||||
*result.add_input_stream() = "FRAME:frame";
|
||||
*result.add_input_stream() = "MASK:mask";
|
||||
return result;
|
||||
}
|
||||
};
|
||||
|
||||
// Shows the SwitchMuxCalculator is available.
|
||||
TEST_F(SwitchDemuxCalculatorTest, IsRegistered) {
|
||||
EXPECT_TRUE(CalculatorBaseRegistry::IsRegistered("SwitchDemuxCalculator"));
|
||||
}
|
||||
|
||||
TEST_F(SwitchDemuxCalculatorTest, BasicDataFlow) {
|
||||
CalculatorGraphConfig::Node node_config = BuildNodeConfig();
|
||||
CalculatorGraph graph = BuildCalculatorGraph(node_config);
|
||||
std::vector<Packet> output_frames0;
|
||||
EXPECT_TRUE(graph
|
||||
.ObserveOutputStream("frame_0",
|
||||
[&](const Packet& p) {
|
||||
output_frames0.push_back(p);
|
||||
return absl::OkStatus();
|
||||
})
|
||||
.ok());
|
||||
std::vector<Packet> output_frames1;
|
||||
EXPECT_TRUE(graph
|
||||
.ObserveOutputStream("frame_1",
|
||||
[&](const Packet& p) {
|
||||
output_frames1.push_back(p);
|
||||
return absl::OkStatus();
|
||||
})
|
||||
.ok());
|
||||
EXPECT_TRUE(
|
||||
graph.StartRun({}, {{"frame", MakePacket<std::string>("frame_header")}})
|
||||
.ok());
|
||||
|
||||
// Finalize input for the "mask" input stream.
|
||||
EXPECT_TRUE(graph.CloseInputStream("mask").ok());
|
||||
|
||||
// Channel 0 is selected just before corresponding packets arrive.
|
||||
EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(0, 1)).ok());
|
||||
EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(0, 10)).ok());
|
||||
EXPECT_TRUE(graph.AddPacketToInputStream("frame", pack("p0_t10", 10)).ok());
|
||||
EXPECT_TRUE(graph.WaitUntilIdle().ok());
|
||||
EXPECT_EQ(output_frames0.size(), 1);
|
||||
EXPECT_EQ(output_frames1.size(), 0);
|
||||
EXPECT_EQ(output_frames0[0].Get<std::string>(), "p0_t10");
|
||||
|
||||
// Channel 1 is selected just before corresponding packets arrive.
|
||||
EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(1, 11)).ok());
|
||||
EXPECT_TRUE(graph.AddPacketToInputStream("select", pack(1, 20)).ok());
|
||||
EXPECT_TRUE(graph.AddPacketToInputStream("frame", pack("p1_t20", 20)).ok());
|
||||
EXPECT_TRUE(graph.WaitUntilIdle().ok());
|
||||
EXPECT_EQ(output_frames0.size(), 1);
|
||||
EXPECT_EQ(output_frames1.size(), 1);
|
||||
EXPECT_EQ(output_frames1[0].Get<std::string>(), "p1_t20");
|
||||
|
||||
EXPECT_EQ(
|
||||
graph.FindOutputStreamManager("frame_0")->Header().Get<std::string>(),
|
||||
"frame_header");
|
||||
EXPECT_EQ(
|
||||
graph.FindOutputStreamManager("frame_1")->Header().Get<std::string>(),
|
||||
"frame_header");
|
||||
|
||||
EXPECT_TRUE(graph.CloseAllPacketSources().ok());
|
||||
EXPECT_TRUE(graph.WaitUntilDone().ok());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
Reference in New Issue
Block a user