From 37a825c98df29442ef49b70c3ee375dcb628e9df Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Fri, 17 Feb 2023 13:50:38 -0800 Subject: [PATCH] graph utils functions. PiperOrigin-RevId: 510513793 --- mediapipe/util/BUILD | 23 ++++++++ mediapipe/util/graph_builder_utils.cc | 64 ++++++++++++++++++++++ mediapipe/util/graph_builder_utils.h | 38 +++++++++++++ mediapipe/util/graph_builder_utils_test.cc | 49 +++++++++++++++++ 4 files changed, 174 insertions(+) create mode 100644 mediapipe/util/graph_builder_utils.cc create mode 100644 mediapipe/util/graph_builder_utils.h create mode 100644 mediapipe/util/graph_builder_utils_test.cc diff --git a/mediapipe/util/BUILD b/mediapipe/util/BUILD index 55556955..59d03361 100644 --- a/mediapipe/util/BUILD +++ b/mediapipe/util/BUILD @@ -387,3 +387,26 @@ cc_library( "//mediapipe/framework/port:opencv_imgproc", ], ) + +cc_library( + name = "graph_builder_utils", + srcs = ["graph_builder_utils.cc"], + hdrs = ["graph_builder_utils.h"], + visibility = ["//visibility:public"], + deps = [ + "//mediapipe/framework:calculator_cc_proto", + "//mediapipe/framework:calculator_framework", + "@com_google_absl//absl/strings", + ], +) + +cc_test( + name = "graph_builder_utils_test", + srcs = ["graph_builder_utils_test.cc"], + deps = [ + ":graph_builder_utils", + "//mediapipe/framework:calculator_cc_proto", + "//mediapipe/framework/port:gtest", + "//mediapipe/framework/port:gtest_main", + ], +) diff --git a/mediapipe/util/graph_builder_utils.cc b/mediapipe/util/graph_builder_utils.cc new file mode 100644 index 00000000..d54573cb --- /dev/null +++ b/mediapipe/util/graph_builder_utils.cc @@ -0,0 +1,64 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mediapipe/util/graph_builder_utils.h" + +#include +#include + +#include "absl/strings/match.h" +#include "absl/strings/str_cat.h" +#include "absl/strings/string_view.h" +#include "mediapipe/framework/calculator.pb.h" +#include "mediapipe/framework/calculator_framework.h" + +namespace mediapipe { +namespace { + +bool StartsWithTag(absl::string_view name, absl::string_view tag) { + constexpr absl::string_view kDelimiter(":"); + return absl::StartsWith(name, absl::StrCat(tag, kDelimiter)); +} + +} // namespace + +bool HasInput(const CalculatorGraphConfig::Node& node, absl::string_view tag) { + for (int i = 0; i < node.input_stream_size(); ++i) { + if (StartsWithTag(node.input_stream(i), tag)) { + return true; + } + } + return false; +} + +bool HasSideInput(const CalculatorGraphConfig::Node& node, + absl::string_view tag) { + for (int i = 0; i < node.input_side_packet_size(); ++i) { + if (StartsWithTag(node.input_side_packet(i), tag)) { + return true; + } + } + return false; +} + +bool HasOutput(const CalculatorGraphConfig::Node& node, absl::string_view tag) { + for (int i = 0; i < node.output_stream_size(); ++i) { + if (StartsWithTag(node.output_stream(i), tag)) { + return true; + } + } + return false; +} + +} // namespace mediapipe diff --git a/mediapipe/util/graph_builder_utils.h b/mediapipe/util/graph_builder_utils.h new file mode 100644 index 00000000..250e79a3 --- /dev/null +++ b/mediapipe/util/graph_builder_utils.h @@ -0,0 +1,38 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef MEDIAPIPE_UTIL_GRAPH_BUILDER_UTILS_H_ +#define MEDIAPIPE_UTIL_GRAPH_BUILDER_UTILS_H_ + +#include + +#include "absl/strings/string_view.h" +#include "mediapipe/framework/calculator.pb.h" +#include "mediapipe/framework/calculator_framework.h" + +namespace mediapipe { + +// Checks if @node has input with the specified @tag. +bool HasInput(const CalculatorGraphConfig::Node& node, absl::string_view tag); + +// Checks if @node has input side-packet with the specified @tag. +bool HasSideInput(const CalculatorGraphConfig::Node& node, + absl::string_view tag); + +// Checks if @node has output with the specified @tag. +bool HasOutput(const CalculatorGraphConfig::Node& node, absl::string_view tag); + +} // namespace mediapipe + +#endif // MEDIAPIPE_UTIL_GRAPH_BUILDER_UTILS_H_ diff --git a/mediapipe/util/graph_builder_utils_test.cc b/mediapipe/util/graph_builder_utils_test.cc new file mode 100644 index 00000000..f8e4490d --- /dev/null +++ b/mediapipe/util/graph_builder_utils_test.cc @@ -0,0 +1,49 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "mediapipe/util/graph_builder_utils.h" + +#include "mediapipe/framework/calculator.pb.h" +#include "mediapipe/framework/port/gmock.h" +#include "mediapipe/framework/port/gtest.h" + +namespace mediapipe { +namespace { + +TEST(GraphUtils, HasInput) { + CalculatorGraphConfig::Node node; + node.add_input_stream("SOME_TAG:some_name"); + + EXPECT_TRUE(HasInput(node, "SOME_TAG")); + EXPECT_FALSE(HasInput(node, "SOME")); +} + +TEST(GraphUtils, HasSideInput) { + CalculatorGraphConfig::Node node; + node.add_input_side_packet("SOME_TAG:some_name"); + + EXPECT_TRUE(HasSideInput(node, "SOME_TAG")); + EXPECT_FALSE(HasSideInput(node, "SOME")); +} + +TEST(GraphUtils, HasOutput) { + CalculatorGraphConfig::Node node; + node.add_output_stream("SOME_TAG:some_name"); + + EXPECT_TRUE(HasOutput(node, "SOME_TAG")); + EXPECT_FALSE(HasOutput(node, "SOME")); +} + +} // namespace +} // namespace mediapipe