Add functions for checking the existence of options in subgraphs and calculators.

PiperOrigin-RevId: 513689742
This commit is contained in:
MediaPipe Team
2023-03-02 18:01:07 -08:00
committed by Copybara-Service
parent 91d53cd181
commit 3837c92fd5
11 changed files with 224 additions and 0 deletions
+14
View File
@@ -181,6 +181,20 @@ cc_library(
],
)
cc_test(
name = "options_map_test",
srcs = ["options_map_test.cc"],
deps = [
":options_map",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
"//mediapipe/framework/testdata:night_light_calculator_cc_proto",
"//mediapipe/framework/testdata:night_light_calculator_options_lib",
],
)
mediapipe_proto_library(
name = "field_data_proto",
srcs = ["field_data.proto"],
+33
View File
@@ -28,6 +28,18 @@ struct IsExtension {
static constexpr bool value = (sizeof(test<T>(0)) == sizeof(char));
};
template <class T,
typename std::enable_if<IsExtension<T>::value, int>::type = 0>
bool HasExtension(const CalculatorOptions& options) {
return options.HasExtension(T::ext);
}
template <class T,
typename std::enable_if<!IsExtension<T>::value, int>::type = 0>
bool HasExtension(const CalculatorOptions& options) {
return false;
}
template <class T,
typename std::enable_if<IsExtension<T>::value, int>::type = 0>
T* GetExtension(CalculatorOptions& options) {
@@ -124,6 +136,27 @@ class OptionsMap {
return *result;
}
template <class T>
bool Has() const {
if (options_.Has<T>()) {
return true;
}
if (node_config_->has_options()) {
return HasExtension<T>(node_config_->options());
}
#if defined(MEDIAPIPE_PROTO_LITE) && defined(MEDIAPIPE_PROTO_THIRD_PARTY)
// protobuf::Any is unavailable with third_party/protobuf:protobuf-lite.
#else
for (const mediapipe::protobuf::Any& options :
node_config_->node_options()) {
if (options.Is<T>()) {
return true;
}
}
#endif
return false;
}
CalculatorGraphConfig::Node* node_config_;
TypeMap options_;
};
@@ -0,0 +1,88 @@
// 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/framework/tool/options_map.h"
#include <unistd.h>
#include <memory>
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_macros.h"
#include "mediapipe/framework/testdata/night_light_calculator.pb.h"
namespace mediapipe {
namespace tool {
namespace {
TEST(OptionsMapTest, QueryNotFound) {
CalculatorGraphConfig::Node node =
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "NightLightCalculator"
input_side_packet: "input_value"
output_stream: "values"
)pb");
OptionsMap options;
options.Initialize(node);
EXPECT_FALSE(options.Has<mediapipe::NightLightCalculatorOptions>());
}
TEST(OptionsMapTest, QueryFound) {
CalculatorGraphConfig::Node node =
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "NightLightCalculator"
input_side_packet: "input_value"
output_stream: "values"
options {
[mediapipe.NightLightCalculatorOptions.ext] {
base_timestamp: 123
output_header: PASS_HEADER
jitter: 0.123
}
}
)pb");
OptionsMap options;
options.Initialize(node);
EXPECT_TRUE(options.Has<mediapipe::NightLightCalculatorOptions>());
EXPECT_EQ(
options.Get<mediapipe::NightLightCalculatorOptions>().base_timestamp()[0],
123);
}
TEST(MutableOptionsMapTest, InsertAndQueryFound) {
CalculatorGraphConfig::Node node =
ParseTextProtoOrDie<CalculatorGraphConfig::Node>(R"pb(
calculator: "NightLightCalculator"
input_side_packet: "input_value"
output_stream: "values"
)pb");
MutableOptionsMap options;
options.Initialize(node);
EXPECT_FALSE(options.Has<mediapipe::NightLightCalculatorOptions>());
mediapipe::NightLightCalculatorOptions night_light_options;
night_light_options.add_base_timestamp(123);
options.Set(night_light_options);
EXPECT_TRUE(options.Has<mediapipe::NightLightCalculatorOptions>());
EXPECT_EQ(
options.Get<mediapipe::NightLightCalculatorOptions>().base_timestamp()[0],
123);
}
} // namespace
} // namespace tool
} // namespace mediapipe