From 7ab3d70aa4239039814d32bd5585be67ccda0656 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Wed, 4 Oct 2023 21:56:55 -0700 Subject: [PATCH] Add stream API presence utils. PiperOrigin-RevId: 570901832 --- mediapipe/framework/api2/stream/BUILD | 22 +++++++++++++ mediapipe/framework/api2/stream/presence.h | 19 +++++++++++ .../framework/api2/stream/presence_test.cc | 33 +++++++++++++++++++ 3 files changed, 74 insertions(+) create mode 100644 mediapipe/framework/api2/stream/presence.h create mode 100644 mediapipe/framework/api2/stream/presence_test.cc diff --git a/mediapipe/framework/api2/stream/BUILD b/mediapipe/framework/api2/stream/BUILD index bca4d2b2e..86e9052dd 100644 --- a/mediapipe/framework/api2/stream/BUILD +++ b/mediapipe/framework/api2/stream/BUILD @@ -237,6 +237,28 @@ cc_test( ], ) +cc_library( + name = "presence", + hdrs = ["presence.h"], + deps = [ + "//mediapipe/calculators/core:packet_presence_calculator", + "//mediapipe/framework/api2:builder", + ], +) + +cc_test( + name = "presence_test", + srcs = ["presence_test.cc"], + deps = [ + ":presence", + "//mediapipe/framework:calculator_cc_proto", + "//mediapipe/framework/api2:builder", + "//mediapipe/framework/port:gtest", + "//mediapipe/framework/port:gtest_main", + "//mediapipe/framework/port:parse_text_proto", + ], +) + cc_library( name = "rect_transformation", srcs = ["rect_transformation.cc"], diff --git a/mediapipe/framework/api2/stream/presence.h b/mediapipe/framework/api2/stream/presence.h new file mode 100644 index 000000000..9b6d7fac0 --- /dev/null +++ b/mediapipe/framework/api2/stream/presence.h @@ -0,0 +1,19 @@ +#ifndef MEDIAPIPE_FRAMEWORK_API2_STREAM_PRESENCE_H_ +#define MEDIAPIPE_FRAMEWORK_API2_STREAM_PRESENCE_H_ + +#include "mediapipe/framework/api2/builder.h" + +namespace mediapipe::api2::builder { + +// Updates @graph to emit a stream containing `bool` packets, where each packet +// indicates whether @stream has a packet with corresponding timestamp or not. +template +Stream IsPresent(Stream stream, Graph& graph) { + auto& presence_node = graph.AddNode("PacketPresenceCalculator"); + stream.ConnectTo(presence_node.In("PACKET")); + return presence_node.Out("PRESENCE").Cast(); +} + +} // namespace mediapipe::api2::builder + +#endif // MEDIAPIPE_FRAMEWORK_API2_STREAM_PRESENCE_H_ diff --git a/mediapipe/framework/api2/stream/presence_test.cc b/mediapipe/framework/api2/stream/presence_test.cc new file mode 100644 index 000000000..b7d9c9cf7 --- /dev/null +++ b/mediapipe/framework/api2/stream/presence_test.cc @@ -0,0 +1,33 @@ +#include "mediapipe/framework/api2/stream/presence.h" + +#include "mediapipe/framework/api2/builder.h" +#include "mediapipe/framework/calculator.pb.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(Presence, VerifyConfig) { + Graph graph; + + Stream stream = graph.In("STREAM").Cast(); + stream.SetName("stream_to_check"); + Stream is_present_stream = IsPresent(stream, graph); + is_present_stream.SetName("is_present_stream"); + + EXPECT_THAT( + graph.GetConfig(), + EqualsProto(mediapipe::ParseTextProtoOrDie(R"pb( + node { + calculator: "PacketPresenceCalculator" + input_stream: "PACKET:stream_to_check" + output_stream: "PRESENCE:is_present_stream" + } + input_stream: "STREAM:stream_to_check" + )pb"))); +} + +} // namespace +} // namespace mediapipe::api2::builder