graph utils functions.
PiperOrigin-RevId: 510513793
This commit is contained in:
parent
1a60a0e2d6
commit
37a825c98d
|
@ -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",
|
||||
],
|
||||
)
|
||||
|
|
64
mediapipe/util/graph_builder_utils.cc
Normal file
64
mediapipe/util/graph_builder_utils.cc
Normal 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
|
38
mediapipe/util/graph_builder_utils.h
Normal file
38
mediapipe/util/graph_builder_utils.h
Normal 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_
|
49
mediapipe/util/graph_builder_utils_test.cc
Normal file
49
mediapipe/util/graph_builder_utils_test.cc
Normal 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
|
Loading…
Reference in New Issue
Block a user