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