Skip to content

Commit

Permalink
Add stream API merge utils.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 571124981
  • Loading branch information
MediaPipe Team authored and copybara-github committed Oct 5, 2023
1 parent a1e542f commit b3f9587
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 0 deletions.
22 changes: 22 additions & 0 deletions mediapipe/framework/api2/stream/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,28 @@ cc_test(
],
)

cc_library(
name = "merge",
hdrs = ["merge.h"],
deps = [
"//mediapipe/calculators/core:merge_calculator",
"//mediapipe/framework/api2:builder",
],
)

cc_test(
name = "merge_test",
srcs = ["merge_test.cc"],
deps = [
":merge",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/api2:builder",
"//mediapipe/framework/port:gtest",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
],
)

cc_library(
name = "presence",
hdrs = ["presence.h"],
Expand Down
20 changes: 20 additions & 0 deletions mediapipe/framework/api2/stream/merge.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
#ifndef MEDIAPIPE_FRAMEWORK_API2_STREAM_MERGE_H_
#define MEDIAPIPE_FRAMEWORK_API2_STREAM_MERGE_H_

#include "mediapipe/framework/api2/builder.h"

namespace mediapipe::api2::builder {

// Updates @graph to choose @a stream if it's available (not empty stream at
// specific timestamp) or @b stream otherwise.
template <typename T>
Stream<T> Merge(Stream<T> a, Stream<T> b, Graph& graph) {
auto& merge_node = graph.AddNode("MergeCalculator");
a.ConnectTo(merge_node.In("")[0]);
b.ConnectTo(merge_node.In("")[1]);
return merge_node.Out("").template Cast<T>();
}

} // namespace mediapipe::api2::builder

#endif // MEDIAPIPE_FRAMEWORK_API2_STREAM_MERGE_H_
35 changes: 35 additions & 0 deletions mediapipe/framework/api2/stream/merge_test.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
#include "mediapipe/framework/api2/stream/merge.h"

#include "mediapipe/framework/api2/builder.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"

namespace mediapipe::api2::builder {
namespace {

TEST(Merge, VerifyConfig) {
mediapipe::api2::builder::Graph graph;

Stream<int> input_a = graph.In("INPUT_A").Cast<int>();
Stream<int> input_b = graph.In("INPUT_B").Cast<int>();
Stream<int> input = Merge(input_a, input_b, graph);
input.SetName("input");

EXPECT_THAT(
graph.GetConfig(),
EqualsProto(mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
node {
calculator: "MergeCalculator"
input_stream: "__stream_0"
input_stream: "__stream_1"
output_stream: "input"
}
input_stream: "INPUT_A:__stream_0"
input_stream: "INPUT_B:__stream_1"
)pb")));
}

} // namespace
} // namespace mediapipe::api2::builder

0 comments on commit b3f9587

Please sign in to comment.