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(
|
cc_library(
|
||||||
name = "presence",
|
name = "presence",
|
||||||
hdrs = ["presence.h"],
|
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