graph utils functions.

PiperOrigin-RevId: 510513793
This commit is contained in:
MediaPipe Team 2023-02-17 13:50:38 -08:00 committed by Copybara-Service
parent 1a60a0e2d6
commit 37a825c98d
4 changed files with 174 additions and 0 deletions

View File

@ -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",
],
)

View File

@ -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 <string>
#include <utility>
#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

View File

@ -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 <string>
#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_

View File

@ -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