Add stream API merge utils.
PiperOrigin-RevId: 571124981
This commit is contained in:
parent
a1e542fc16
commit
b3f9587bc2
|
@ -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"],
|
||||
|
|
20
mediapipe/framework/api2/stream/merge.h
Normal file
20
mediapipe/framework/api2/stream/merge.h
Normal file
|
@ -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
mediapipe/framework/api2/stream/merge_test.cc
Normal file
35
mediapipe/framework/api2/stream/merge_test.cc
Normal file
|
@ -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
|
Loading…
Reference in New Issue
Block a user