talksik/mediapipe
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 6e4aff1cc351be3ae4537b677f36d139ee50ce09
This commit is contained in:
MediaPipe Team
2021-03-25 22:09:18 -04:00
committed by chuoling
parent a92cff7a60
commit 7c331ad58b
175 changed files with 4804 additions and 1325 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/BUILD
+68 -21
View File
@@ -23,7 +23,6 @@ package(default_visibility = ["//visibility:private"])
package_group(
name = "mediapipe_internal",
packages = [
"//java/com/google/mediapipe/framework/...",
"//mediapipe/...",
],
)
@@ -78,21 +77,19 @@ mediapipe_proto_library(
mediapipe_proto_library(
name = "mediapipe_options_proto",
srcs = ["mediapipe_options.proto"],
visibility = ["//mediapipe/framework:__subpackages__"],
visibility = [":mediapipe_internal"],
)
mediapipe_proto_library(
name = "packet_factory_proto",
srcs = ["packet_factory.proto"],
visibility = ["//mediapipe/framework:__subpackages__"],
visibility = [":mediapipe_internal"],
)
mediapipe_proto_library(
name = "packet_generator_proto",
srcs = ["packet_generator.proto"],
visibility = [
"//mediapipe:__subpackages__",
],
visibility = [":mediapipe_internal"],
)
mediapipe_proto_library(
@@ -105,7 +102,7 @@ mediapipe_proto_library(
mediapipe_proto_library(
name = "status_handler_proto",
srcs = ["status_handler.proto"],
visibility = ["//mediapipe/framework:__subpackages__"],
visibility = [":mediapipe_internal"],
deps = ["//mediapipe/framework:mediapipe_options_proto"],
)
@@ -274,14 +271,17 @@ cc_library(
],
deps = [
":calculator_base",
":calculator_node",
":counter_factory",
":delegating_executor",
":mediapipe_profiling",
":executor",
":graph_output_stream",
":graph_service",
":graph_service_manager",
":input_stream_manager",
":input_stream_shard",
":graph_service",
":output_side_packet_impl",
":output_stream",
":output_stream_manager",
":output_stream_poller",
@@ -303,29 +303,27 @@ cc_library(
"//mediapipe/framework:packet_generator_cc_proto",
"//mediapipe/framework:status_handler_cc_proto",
"//mediapipe/framework:thread_pool_executor_cc_proto",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"//mediapipe/gpu:graph_support",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/container:fixed_array",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
":calculator_node",
":output_side_packet_impl",
"//mediapipe/framework/profiler:graph_profiler",
"//mediapipe/framework/tool:fill_packet_set",
"//mediapipe/framework/tool:status_util",
"//mediapipe/framework/tool:tag_map",
"//mediapipe/framework/tool:validate",
"//mediapipe/framework/tool:validate_name",
"//mediapipe/framework/port:core_proto",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:source_location",
"//mediapipe/framework/port:status",
"//mediapipe/framework/profiler:graph_profiler",
"//mediapipe/framework/tool:fill_packet_set",
"//mediapipe/framework/tool:status_util",
"//mediapipe/framework/tool:tag_map",
"//mediapipe/framework/tool:validate",
"//mediapipe/framework/tool:validate_name",
"//mediapipe/gpu:graph_support",
"//mediapipe/util:cpu_util",
] + select({
"//conditions:default": [
@@ -336,6 +334,28 @@ cc_library(
}),
)
cc_library(
name = "graph_service_manager",
srcs = ["graph_service_manager.cc"],
hdrs = ["graph_service_manager.h"],
visibility = [":mediapipe_internal"],
deps = [
":graph_service",
"//mediapipe/framework:packet",
"@com_google_absl//absl/status",
],
)
cc_test(
name = "graph_service_manager_test",
srcs = ["graph_service_manager_test.cc"],
deps = [
":graph_service_manager",
"//mediapipe/framework:packet",
"//mediapipe/framework/port:gtest_main",
],
)
cc_library(
name = "calculator_node",
srcs = ["calculator_node.cc"],
@@ -425,6 +445,7 @@ cc_library(
":counter",
":counter_factory",
":graph_service",
":graph_service_manager",
":input_stream",
":output_stream",
":packet",
@@ -977,6 +998,8 @@ cc_library(
hdrs = ["subgraph.h"],
visibility = ["//visibility:public"],
deps = [
":graph_service",
":graph_service_manager",
":port",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:mediapipe_options_cc_proto",
@@ -989,6 +1012,8 @@ cc_library(
"//mediapipe/framework/tool:template_expander",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@com_google_absl//absl/types:optional",
],
)
@@ -1008,7 +1033,7 @@ cc_library(
"//mediapipe/framework/port:status",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/strings",
"@eigen_archive//:eigen",
"@eigen_archive//:eigen3",
],
alwayslink = 1,
)
@@ -1102,6 +1127,7 @@ cc_library(
deps = [
":calculator_base",
":calculator_contract",
":graph_service_manager",
":legacy_calculator_support",
":packet",
":packet_generator",
@@ -1136,6 +1162,24 @@ cc_library(
],
)
cc_test(
name = "validated_graph_config_test",
srcs = ["validated_graph_config_test.cc"],
deps = [
":calculator_framework",
":graph_service",
":graph_service_manager",
":validated_graph_config",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework/api2:node",
"//mediapipe/framework/api2:port",
"//mediapipe/framework/deps:message_matchers",
"//mediapipe/framework/port:gtest_main",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
],
)
cc_library(
name = "graph_validation",
hdrs = ["graph_validation.h"],
@@ -1591,13 +1635,16 @@ cc_test(
srcs = ["subgraph_test.cc"],
deps = [
":calculator_framework",
":graph_service_manager",
":subgraph",
":test_calculators",
"//mediapipe/calculators/core:constant_side_packet_calculator",
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/tool:sink",
"//mediapipe/framework/tool/testdata:dub_quad_test_subgraph",
"@com_google_absl//absl/strings:str_format",
],
)
mediapipe/framework/calculator_context.cc
+2 -2
View File
@@ -41,9 +41,9 @@ Counter* CalculatorContext::GetCounter(const std::string& name) {
return calculator_state_->GetCounter(name);
}
CounterSet* CalculatorContext::GetCounterSet() {
CounterFactory* CalculatorContext::GetCounterFactory() {
CHECK(calculator_state_);
return calculator_state_->GetCounterSet();
return calculator_state_->GetCounterFactory();
}
const PacketSet& CalculatorContext::InputSidePackets() const {
mediapipe/framework/calculator_context.h
+2 -19
View File
@@ -76,7 +76,7 @@ class CalculatorContext {
// Returns the counter set, which can be used to create new counters.
// No prefix is added to counters created in this way.
CounterSet* GetCounterSet();
CounterFactory* GetCounterFactory();
// Returns the current input timestamp, or Timestamp::Unset if there are
// no input packets.
@@ -113,26 +113,9 @@ class CalculatorContext {
return calculator_state_->GetSharedProfilingContext().get();
}
template <typename T>
class ServiceBinding {
public:
bool IsAvailable() {
return calculator_state_->IsServiceAvailable(service_);
}
T& GetObject() { return calculator_state_->GetServiceObject(service_); }
ServiceBinding(CalculatorState* calculator_state,
const GraphService<T>& service)
: calculator_state_(calculator_state), service_(service) {}
private:
CalculatorState* calculator_state_;
const GraphService<T>& service_;
};
template <typename T>
ServiceBinding<T> Service(const GraphService<T>& service) {
return ServiceBinding<T>(calculator_state_, service);
return ServiceBinding<T>(calculator_state_->GetServiceObject(service));
}
private:
mediapipe/framework/calculator_graph.cc
+21 -35
View File
@@ -36,6 +36,7 @@
#include "mediapipe/framework/calculator_base.h"
#include "mediapipe/framework/counter_factory.h"
#include "mediapipe/framework/delegating_executor.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/input_stream_manager.h"
#include "mediapipe/framework/mediapipe_profiling.h"
#include "mediapipe/framework/packet_generator.h"
@@ -392,7 +393,8 @@ absl::Status CalculatorGraph::Initialize(
const CalculatorGraphConfig& input_config,
const std::map<std::string, Packet>& side_packets) {
auto validated_graph = absl::make_unique<ValidatedGraphConfig>();
MP_RETURN_IF_ERROR(validated_graph->Initialize(input_config));
MP_RETURN_IF_ERROR(validated_graph->Initialize(
input_config, /*graph_registry=*/nullptr, &service_manager_));
return Initialize(std::move(validated_graph), side_packets);
}
@@ -402,8 +404,8 @@ absl::Status CalculatorGraph::Initialize(
const std::map<std::string, Packet>& side_packets,
const std::string& graph_type, const Subgraph::SubgraphOptions* options) {
auto validated_graph = absl::make_unique<ValidatedGraphConfig>();
MP_RETURN_IF_ERROR(validated_graph->Initialize(input_configs, input_templates,
graph_type, options));
MP_RETURN_IF_ERROR(validated_graph->Initialize(
input_configs, input_templates, graph_type, options, &service_manager_));
return Initialize(std::move(validated_graph), side_packets);
}
@@ -509,19 +511,15 @@ absl::Status CalculatorGraph::StartRun(
#if !MEDIAPIPE_DISABLE_GPU
absl::Status CalculatorGraph::SetGpuResources(
std::shared_ptr<::mediapipe::GpuResources> resources) {
RET_CHECK(!ContainsKey(service_packets_, kGpuService.key))
auto gpu_service = service_manager_.GetServiceObject(kGpuService);
RET_CHECK_EQ(gpu_service, nullptr)
<< "The GPU resources have already been configured.";
service_packets_[kGpuService.key] =
MakePacket<std::shared_ptr<::mediapipe::GpuResources>>(
std::move(resources));
return absl::OkStatus();
return service_manager_.SetServiceObject(kGpuService, std::move(resources));
}
std::shared_ptr<::mediapipe::GpuResources> CalculatorGraph::GetGpuResources()
const {
auto service_iter = service_packets_.find(kGpuService.key);
if (service_iter == service_packets_.end()) return nullptr;
return service_iter->second.Get<std::shared_ptr<::mediapipe::GpuResources>>();
return service_manager_.GetServiceObject(kGpuService);
}
absl::StatusOr<std::map<std::string, Packet>> CalculatorGraph::PrepareGpu(
@@ -536,8 +534,7 @@ absl::StatusOr<std::map<std::string, Packet>> CalculatorGraph::PrepareGpu(
}
}
if (uses_gpu) {
auto service_iter = service_packets_.find(kGpuService.key);
bool has_service = service_iter != service_packets_.end();
auto gpu_resources = service_manager_.GetServiceObject(kGpuService);
auto legacy_sp_iter = side_packets.find(kGpuSharedSidePacketName);
// Workaround for b/116875321: CalculatorRunner provides an empty packet,
@@ -545,15 +542,12 @@ absl::StatusOr<std::map<std::string, Packet>> CalculatorGraph::PrepareGpu(
bool has_legacy_sp = legacy_sp_iter != side_packets.end() &&
!legacy_sp_iter->second.IsEmpty();
std::shared_ptr<::mediapipe::GpuResources> gpu_resources;
if (has_service) {
if (gpu_resources) {
if (has_legacy_sp) {
LOG(WARNING)
<< "::mediapipe::GpuSharedData provided as a side packet while the "
<< "graph already had one; ignoring side packet";
}
gpu_resources = service_iter->second
.Get<std::shared_ptr<::mediapipe::GpuResources>>();
update_sp = true;
} else {
if (has_legacy_sp) {
@@ -564,8 +558,8 @@ absl::StatusOr<std::map<std::string, Packet>> CalculatorGraph::PrepareGpu(
ASSIGN_OR_RETURN(gpu_resources, ::mediapipe::GpuResources::Create());
update_sp = true;
}
service_packets_[kGpuService.key] =
MakePacket<std::shared_ptr<::mediapipe::GpuResources>>(gpu_resources);
MP_RETURN_IF_ERROR(
service_manager_.SetServiceObject(kGpuService, gpu_resources));
}
// Create or replace the legacy side packet if needed.
@@ -682,8 +676,10 @@ absl::Status CalculatorGraph::PrepareForRun(
std::placeholders::_1, std::placeholders::_2);
node.SetQueueSizeCallbacks(queue_size_callback, queue_size_callback);
scheduler_.AssignNodeToSchedulerQueue(&node);
// TODO: update calculator node to use GraphServiceManager
// instead of service packets?
const absl::Status result = node.PrepareForRun(
current_run_side_packets_, service_packets_,
current_run_side_packets_, service_manager_.ServicePackets(),
std::bind(&internal::Scheduler::ScheduleNodeForOpen, &scheduler_,
&node),
std::bind(&internal::Scheduler::AddNodeToSourcesQueue, &scheduler_,
@@ -811,6 +807,11 @@ absl::Status CalculatorGraph::AddPacketToInputStreamInternal(
CHECK_GE(node_id, validated_graph_->CalculatorInfos().size());
{
absl::MutexLock lock(&full_input_streams_mutex_);
if (full_input_streams_.empty()) {
return mediapipe::FailedPreconditionErrorBuilder(MEDIAPIPE_LOC)
<< "CalculatorGraph::AddPacketToInputStream() is called before "
"StartRun()";
}
if (graph_input_stream_add_mode_ ==
GraphInputStreamAddMode::ADD_IF_NOT_FULL) {
if (has_error_) {
@@ -1170,21 +1171,6 @@ void CalculatorGraph::Pause() { scheduler_.Pause(); }
void CalculatorGraph::Resume() { scheduler_.Resume(); }
absl::Status CalculatorGraph::SetServicePacket(const GraphServiceBase& service,
Packet p) {
// TODO: check that the graph has not been started!
service_packets_[service.key] = std::move(p);
return absl::OkStatus();
}
Packet CalculatorGraph::GetServicePacket(const GraphServiceBase& service) {
auto it = service_packets_.find(service.key);
if (it == service_packets_.end()) {
return {};
}
return it->second;
}
absl::Status CalculatorGraph::SetExecutorInternal(
const std::string& name, std::shared_ptr<Executor> executor) {
if (!executors_.emplace(name, executor).second) {
mediapipe/framework/calculator_graph.h
+10 -8
View File
@@ -38,6 +38,7 @@
#include "mediapipe/framework/executor.h"
#include "mediapipe/framework/graph_output_stream.h"
#include "mediapipe/framework/graph_service.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/mediapipe_profiling.h"
#include "mediapipe/framework/output_side_packet_impl.h"
#include "mediapipe/framework/output_stream.h"
@@ -377,19 +378,20 @@ class CalculatorGraph {
template <typename T>
absl::Status SetServiceObject(const GraphService<T>& service,
std::shared_ptr<T> object) {
return SetServicePacket(service,
MakePacket<std::shared_ptr<T>>(std::move(object)));
// TODO: check that the graph has not been started!
return service_manager_.SetServiceObject(service, object);
}
template <typename T>
std::shared_ptr<T> GetServiceObject(const GraphService<T>& service) {
Packet p = GetServicePacket(service);
if (p.IsEmpty()) return nullptr;
return p.Get<std::shared_ptr<T>>();
return service_manager_.GetServiceObject(service);
}
// Only the Java API should call this directly.
absl::Status SetServicePacket(const GraphServiceBase& service, Packet p);
absl::Status SetServicePacket(const GraphServiceBase& service, Packet p) {
// TODO: check that the graph has not been started!
return service_manager_.SetServicePacket(service, p);
}
private:
// GraphRunState is used as a parameter in the function CallStatusHandlers.
@@ -523,7 +525,6 @@ class CalculatorGraph {
// status before taking any action.
void UpdateThrottledNodes(InputStreamManager* stream, bool* stream_was_full);
Packet GetServicePacket(const GraphServiceBase& service);
#if !MEDIAPIPE_DISABLE_GPU
// Owns the legacy GpuSharedData if we need to create one for backwards
// compatibility.
@@ -598,7 +599,8 @@ class CalculatorGraph {
// The processed input side packet map for this run.
std::map<std::string, Packet> current_run_side_packets_;
std::map<std::string, Packet> service_packets_;
// Object to manage graph services.
GraphServiceManager service_manager_;
// Vector of errors encountered while running graph. Always use RecordError()
// to add an error to this vector.
mediapipe/framework/calculator_graph_bounds_test.cc
+32
View File
@@ -1361,6 +1361,38 @@ TEST(CalculatorGraphBoundsTest, ProcessTimestampBounds_Passthrough) {
MP_ASSERT_OK(graph.WaitUntilDone());
}
TEST(CalculatorGraphBoundsTest, PostStreamPacketToSetProcessTimestampBound) {
std::string config_str = R"(
input_stream: "input_0"
node {
calculator: "ProcessBoundToPacketCalculator"
input_stream: "input_0"
output_stream: "output_0"
}
)";
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(config_str);
CalculatorGraph graph;
std::vector<Packet> output_0_packets;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.ObserveOutputStream("output_0", [&](const Packet& p) {
output_0_packets.push_back(p);
return absl::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.WaitUntilIdle());
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_0", MakePacket<int>(0).At(Timestamp::PostStream())));
MP_ASSERT_OK(graph.WaitUntilIdle());
EXPECT_EQ(output_0_packets.size(), 1);
EXPECT_EQ(output_0_packets[0].Timestamp(), Timestamp::PostStream());
// Shutdown the graph.
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// A Calculator that sends a timestamp bound for every other input.
class OccasionalBoundCalculator : public CalculatorBase {
public:
mediapipe/framework/calculator_graph_test.cc
-251
View File
@@ -4356,256 +4356,5 @@ TEST(CalculatorGraph, GraphInputStreamWithTag) {
ASSERT_EQ(5, packet_dump.size());
}
// Returns the first packet of the input stream.
class FirstPacketFilterCalculator : public CalculatorBase {
public:
FirstPacketFilterCalculator() {}
~FirstPacketFilterCalculator() override {}
static absl::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
return absl::OkStatus();
}
absl::Status Process(CalculatorContext* cc) override {
if (!seen_first_packet_) {
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
cc->Outputs().Index(0).Close();
seen_first_packet_ = true;
}
return absl::OkStatus();
}
private:
bool seen_first_packet_ = false;
};
REGISTER_CALCULATOR(FirstPacketFilterCalculator);
constexpr int kDefaultMaxCount = 1000;
TEST(CalculatorGraph, TestPollPacket) {
CalculatorGraphConfig config;
CalculatorGraphConfig::Node* node = config.add_node();
node->set_calculator("CountingSourceCalculator");
node->add_output_stream("output");
node->add_input_side_packet("MAX_COUNT:max_count");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
auto status_or_poller = graph.AddOutputStreamPoller("output");
ASSERT_TRUE(status_or_poller.ok());
OutputStreamPoller poller = std::move(status_or_poller.value());
MP_ASSERT_OK(
graph.StartRun({{"max_count", MakePacket<int>(kDefaultMaxCount)}}));
Packet packet;
int num_packets = 0;
while (poller.Next(&packet)) {
EXPECT_EQ(num_packets, packet.Get<int>());
++num_packets;
}
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
EXPECT_FALSE(poller.Next(&packet));
EXPECT_EQ(kDefaultMaxCount, num_packets);
}
TEST(CalculatorGraph, TestOutputStreamPollerDesiredQueueSize) {
CalculatorGraphConfig config;
CalculatorGraphConfig::Node* node = config.add_node();
node->set_calculator("CountingSourceCalculator");
node->add_output_stream("output");
node->add_input_side_packet("MAX_COUNT:max_count");
for (int queue_size = 1; queue_size < 10; ++queue_size) {
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
auto status_or_poller = graph.AddOutputStreamPoller("output");
ASSERT_TRUE(status_or_poller.ok());
OutputStreamPoller poller = std::move(status_or_poller.value());
poller.SetMaxQueueSize(queue_size);
MP_ASSERT_OK(
graph.StartRun({{"max_count", MakePacket<int>(kDefaultMaxCount)}}));
Packet packet;
int num_packets = 0;
while (poller.Next(&packet)) {
EXPECT_EQ(num_packets, packet.Get<int>());
++num_packets;
}
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
EXPECT_FALSE(poller.Next(&packet));
EXPECT_EQ(kDefaultMaxCount, num_packets);
}
}
TEST(CalculatorGraph, TestPollPacketsFromMultipleStreams) {
CalculatorGraphConfig config;
CalculatorGraphConfig::Node* node1 = config.add_node();
node1->set_calculator("CountingSourceCalculator");
node1->add_output_stream("stream1");
node1->add_input_side_packet("MAX_COUNT:max_count");
CalculatorGraphConfig::Node* node2 = config.add_node();
node2->set_calculator("PassThroughCalculator");
node2->add_input_stream("stream1");
node2->add_output_stream("stream2");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
auto status_or_poller1 = graph.AddOutputStreamPoller("stream1");
ASSERT_TRUE(status_or_poller1.ok());
OutputStreamPoller poller1 = std::move(status_or_poller1.value());
auto status_or_poller2 = graph.AddOutputStreamPoller("stream2");
ASSERT_TRUE(status_or_poller2.ok());
OutputStreamPoller poller2 = std::move(status_or_poller2.value());
MP_ASSERT_OK(
graph.StartRun({{"max_count", MakePacket<int>(kDefaultMaxCount)}}));
Packet packet1;
Packet packet2;
int num_packets1 = 0;
int num_packets2 = 0;
int running_pollers = 2;
while (running_pollers > 0) {
if (poller1.Next(&packet1)) {
EXPECT_EQ(num_packets1++, packet1.Get<int>());
} else {
--running_pollers;
}
if (poller2.Next(&packet2)) {
EXPECT_EQ(num_packets2++, packet2.Get<int>());
} else {
--running_pollers;
}
}
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
EXPECT_FALSE(poller1.Next(&packet1));
EXPECT_FALSE(poller2.Next(&packet2));
EXPECT_EQ(kDefaultMaxCount, num_packets1);
EXPECT_EQ(kDefaultMaxCount, num_packets2);
}
// Ensure that when a custom input stream handler is used to handle packets from
// input streams, an error message is outputted with the appropriate link to
// resolve the issue when the calculator doesn't handle inputs in monotonically
// increasing order of timestamps.
TEST(CalculatorGraph, SimpleMuxCalculatorWithCustomInputStreamHandler) {
CalculatorGraph graph;
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'input0'
input_stream: 'input1'
node {
calculator: 'SimpleMuxCalculator'
input_stream: 'input0'
input_stream: 'input1'
input_stream_handler {
input_stream_handler: "ImmediateInputStreamHandler"
}
output_stream: 'output'
}
)");
std::vector<Packet> packet_dump;
tool::AddVectorSink("output", &config, &packet_dump);
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
// Send packets to input stream "input0" at timestamps 0 and 1 consecutively.
Timestamp input0_timestamp = Timestamp(0);
MP_EXPECT_OK(graph.AddPacketToInputStream(
"input0", MakePacket<int>(1).At(input0_timestamp)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(1, packet_dump.size());
EXPECT_EQ(1, packet_dump[0].Get<int>());
++input0_timestamp;
MP_EXPECT_OK(graph.AddPacketToInputStream(
"input0", MakePacket<int>(3).At(input0_timestamp)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(2, packet_dump.size());
EXPECT_EQ(3, packet_dump[1].Get<int>());
// Send a packet to input stream "input1" at timestamp 0 after sending two
// packets at timestamps 0 and 1 to input stream "input0". This will result
// in a mismatch in timestamps as the SimpleMuxCalculator doesn't handle
// inputs from all streams in monotonically increasing order of timestamps.
Timestamp input1_timestamp = Timestamp(0);
MP_EXPECT_OK(graph.AddPacketToInputStream(
"input1", MakePacket<int>(2).At(input1_timestamp)));
absl::Status run_status = graph.WaitUntilIdle();
EXPECT_THAT(
run_status.ToString(),
testing::AllOf(
// The core problem.
testing::HasSubstr("timestamp mismatch on a calculator"),
testing::HasSubstr(
"timestamps that are not strictly monotonically increasing"),
// Link to the possible solution.
testing::HasSubstr("ImmediateInputStreamHandler class comment")));
}
void DoTestMultipleGraphRuns(absl::string_view input_stream_handler,
bool select_packet) {
std::string graph_proto = absl::StrFormat(R"(
input_stream: 'input'
input_stream: 'select'
node {
calculator: 'PassThroughCalculator'
input_stream: 'input'
input_stream: 'select'
input_stream_handler {
input_stream_handler: "%s"
}
output_stream: 'output'
output_stream: 'select_out'
}
)",
input_stream_handler.data());
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(graph_proto);
std::vector<Packet> packet_dump;
tool::AddVectorSink("output", &config, &packet_dump);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
struct Run {
Timestamp timestamp;
int value;
};
std::vector<Run> runs = {{.timestamp = Timestamp(2000), .value = 2},
{.timestamp = Timestamp(1000), .value = 1}};
for (const Run& run : runs) {
MP_ASSERT_OK(graph.StartRun({}));
if (select_packet) {
MP_EXPECT_OK(graph.AddPacketToInputStream(
"select", MakePacket<int>(0).At(run.timestamp)));
}
MP_EXPECT_OK(graph.AddPacketToInputStream(
"input", MakePacket<int>(run.value).At(run.timestamp)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(1, packet_dump.size());
EXPECT_EQ(run.value, packet_dump[0].Get<int>());
EXPECT_EQ(run.timestamp, packet_dump[0].Timestamp());
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
packet_dump.clear();
}
}
TEST(CalculatorGraph, MultipleRunsWithDifferentInputStreamHandlers) {
DoTestMultipleGraphRuns("BarrierInputStreamHandler", true);
DoTestMultipleGraphRuns("DefaultInputStreamHandler", true);
DoTestMultipleGraphRuns("EarlyCloseInputStreamHandler", true);
DoTestMultipleGraphRuns("FixedSizeInputStreamHandler", true);
DoTestMultipleGraphRuns("ImmediateInputStreamHandler", false);
DoTestMultipleGraphRuns("MuxInputStreamHandler", true);
DoTestMultipleGraphRuns("SyncSetInputStreamHandler", true);
DoTestMultipleGraphRuns("TimestampAlignInputStreamHandler", true);
}
} // namespace
} // namespace mediapipe
mediapipe/framework/calculator_node.cc
+4 -4
View File
@@ -408,13 +408,13 @@ absl::Status CalculatorNode::PrepareForRun(
validated_graph_->CalculatorInfos()[node_id_].Contract();
for (const auto& svc_req : contract.ServiceRequests()) {
const auto& req = svc_req.second;
std::string key{req.Service().key};
auto it = service_packets.find(key);
auto it = service_packets.find(req.Service().key);
if (it == service_packets.end()) {
RET_CHECK(req.IsOptional())
<< "required service '" << key << "' was not provided";
<< "required service '" << req.Service().key << "' was not provided";
} else {
calculator_state_->SetServicePacket(key, it->second);
MP_RETURN_IF_ERROR(
calculator_state_->SetServicePacket(req.Service(), it->second));
}
}
mediapipe/framework/calculator_state.cc
+2 -6
View File
@@ -61,13 +61,9 @@ Counter* CalculatorState::GetCounter(const std::string& name) {
return counter_factory_->GetCounter(absl::StrCat(NodeName(), "-", name));
}
CounterSet* CalculatorState::GetCounterSet() {
CounterFactory* CalculatorState::GetCounterFactory() {
CHECK(counter_factory_);
return counter_factory_->GetCounterSet();
}
void CalculatorState::SetServicePacket(const std::string& key, Packet packet) {
service_packets_[key] = std::move(packet);
return counter_factory_;
}
} // namespace mediapipe
mediapipe/framework/calculator_state.h
+8 -10
View File
@@ -27,6 +27,7 @@
#include "mediapipe/framework/counter.h"
#include "mediapipe/framework/counter_factory.h"
#include "mediapipe/framework/graph_service.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/packet_set.h"
#include "mediapipe/framework/port.h"
@@ -81,7 +82,7 @@ class CalculatorState {
// Returns a counter set, which can be passed to other classes, to generate
// counters. NOTE: This differs from GetCounter, in that the counters
// created by this counter set do not have the NodeName prefix.
CounterSet* GetCounterSet();
CounterFactory* GetCounterFactory();
std::shared_ptr<ProfilingContext> GetSharedProfilingContext() const {
return profiling_context_;
@@ -99,17 +100,14 @@ class CalculatorState {
counter_factory_ = counter_factory;
}
void SetServicePacket(const std::string& key, Packet packet);
bool IsServiceAvailable(const GraphServiceBase& service) {
return ContainsKey(service_packets_, service.key);
absl::Status SetServicePacket(const GraphServiceBase& service,
Packet packet) {
return graph_service_manager_.SetServicePacket(service, packet);
}
template <typename T>
T& GetServiceObject(const GraphService<T>& service) {
auto it = service_packets_.find(service.key);
CHECK(it != service_packets_.end());
return *it->second.template Get<std::shared_ptr<T>>();
std::shared_ptr<T> GetServiceObject(const GraphService<T>& service) {
return graph_service_manager_.GetServiceObject(service);
}
private:
@@ -129,7 +127,7 @@ class CalculatorState {
// The graph tracing and profiling interface.
std::shared_ptr<ProfilingContext> profiling_context_;
std::map<std::string, Packet> service_packets_;
GraphServiceManager graph_service_manager_;
////////////////////////////////////////
// Variables which ARE cleared by ResetBetweenRuns().
mediapipe/framework/deps/ret_check.h
+1 -1
View File
@@ -37,7 +37,7 @@ inline StatusBuilder RetCheckImpl(const absl::Status& status,
const char* condition,
mediapipe::source_location location) {
if (ABSL_PREDICT_TRUE(status.ok()))
return mediapipe::StatusBuilder(OkStatus(), location);
return mediapipe::StatusBuilder(absl::OkStatus(), location);
return RetCheckFailSlowPath(location, condition, status);
}
mediapipe/framework/deps/status.cc
+1 -1
View File
@@ -18,7 +18,7 @@
namespace mediapipe {
std::ostream& operator<<(std::ostream& os, const Status& x) {
std::ostream& operator<<(std::ostream& os, const absl::Status& x) {
os << x.ToString();
return os;
}
mediapipe/framework/deps/status_macros.h
+2 -2
View File
@@ -194,10 +194,10 @@ namespace status_macro_internal {
// that declares a variable.
class StatusAdaptorForMacros {
public:
StatusAdaptorForMacros(const Status& status, const char* file, int line)
StatusAdaptorForMacros(const absl::Status& status, const char* file, int line)
: builder_(status, file, line) {}
StatusAdaptorForMacros(Status&& status, const char* file, int line)
StatusAdaptorForMacros(absl::Status&& status, const char* file, int line)
: builder_(std::move(status), file, line) {}
StatusAdaptorForMacros(const StatusBuilder& builder, const char* /* file */,
mediapipe/framework/encode_binary_proto.bzl
+14 -8
View File
@@ -79,13 +79,10 @@ def _get_proto_provider(dep):
def _encode_binary_proto_impl(ctx):
"""Implementation of the encode_binary_proto rule."""
all_protos = depset()
for dep in ctx.attr.deps:
provider = _get_proto_provider(dep)
all_protos = depset(
direct = [],
transitive = [all_protos, provider.transitive_sources],
)
all_protos = depset(
direct = [],
transitive = [_get_proto_provider(dep).transitive_sources for dep in ctx.attr.deps],
)
textpb = ctx.file.input
binarypb = ctx.outputs.output or ctx.actions.declare_file(
@@ -120,7 +117,7 @@ def _encode_binary_proto_impl(ctx):
data_runfiles = ctx.runfiles(transitive_files = output_depset),
)]
encode_binary_proto = rule(
_encode_binary_proto = rule(
implementation = _encode_binary_proto_impl,
attrs = {
"_proto_compiler": attr.label(
@@ -142,6 +139,15 @@ encode_binary_proto = rule(
},
)
def encode_binary_proto(name, input, message_type, deps, **kwargs):
_encode_binary_proto(
name = name,
input = input,
message_type = message_type,
deps = deps,
**kwargs
)
def _generate_proto_descriptor_set_impl(ctx):
"""Implementation of the generate_proto_descriptor_set rule."""
all_protos = depset(transitive = [
mediapipe/framework/formats/BUILD
+3 -1
View File
@@ -114,7 +114,7 @@ cc_library(
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"@eigen_archive//:eigen",
"@eigen_archive//:eigen3",
],
)
@@ -260,9 +260,11 @@ mediapipe_register_type(
include_headers = ["mediapipe/framework/formats/landmark.pb.h"],
types = [
"::mediapipe::Landmark",
"::mediapipe::LandmarkList",
"::mediapipe::NormalizedLandmark",
"::mediapipe::NormalizedLandmarkList",
"::std::vector<::mediapipe::Landmark>",
"::std::vector<::mediapipe::LandmarkList>",
"::std::vector<::mediapipe::NormalizedLandmark>",
"::std::vector<::mediapipe::NormalizedLandmarkList>",
],
mediapipe/framework/formats/classification.proto
+2
View File
@@ -31,6 +31,8 @@ message Classification {
optional float score = 2;
// Label or name of the class.
optional string label = 3;
// Optional human-readable string for display purposes.
optional string display_name = 4;
}
// Group of Classification protos.
mediapipe/framework/formats/image.h
+7 -6
View File
@@ -78,6 +78,12 @@ class Image {
pixel_mutex_ = std::make_shared<absl::Mutex>();
}
// CPU getters.
const ImageFrameSharedPtr& GetImageFrameSharedPtr() const {
if (use_gpu_ == true) ConvertToCpu();
return image_frame_;
}
// Creates an Image representing the same image content as the input GPU
// buffer in platform-specific representations.
#if !MEDIAPIPE_DISABLE_GPU
@@ -95,13 +101,8 @@ class Image {
gpu_buffer_ = gpu_buffer;
pixel_mutex_ = std::make_shared<absl::Mutex>();
}
#endif // !MEDIAPIPE_DISABLE_GPU
const ImageFrameSharedPtr& GetImageFrameSharedPtr() const {
if (use_gpu_ == true) ConvertToCpu();
return image_frame_;
}
#if !MEDIAPIPE_DISABLE_GPU
// GPU getters.
#if MEDIAPIPE_GPU_BUFFER_USE_CV_PIXEL_BUFFER
CVPixelBufferRef GetCVPixelBufferRef() const {
if (use_gpu_ == false) ConvertToGpu();
mediapipe/framework/formats/landmark.proto
+2 -2
View File
@@ -47,8 +47,8 @@ message LandmarkList {
repeated Landmark landmark = 1;
}
// A normalized version of above Landmark proto. All coordiates should be within
// [0, 1].
// A normalized version of above Landmark proto. All coordinates should be
// within [0, 1].
message NormalizedLandmark {
optional float x = 1;
optional float y = 2;
mediapipe/framework/formats/motion/BUILD
+1 -1
View File
@@ -67,11 +67,11 @@ cc_test(
deps = [
":optical_flow_field",
"//mediapipe/framework/deps:file_path",
"//mediapipe/framework/port:commandlineflags",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:logging",
"@com_google_absl//absl/flags:flag",
"@org_tensorflow//tensorflow/core:framework",
],
)
mediapipe/framework/formats/motion/optical_flow_field_test.cc
+1 -1
View File
@@ -18,8 +18,8 @@
#include <memory>
#include <string>
#include "absl/flags/flag.h"
#include "mediapipe/framework/deps/file_path.h"
#include "mediapipe/framework/port/commandlineflags.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/integral_types.h"
mediapipe/framework/graph_service.h
+13
View File
@@ -41,6 +41,19 @@ struct GraphService : public GraphServiceBase {
constexpr GraphService(const char* key) : GraphServiceBase(key) {}
};
template <typename T>
class ServiceBinding {
public:
bool IsAvailable() { return service_ != nullptr; }
T& GetObject() { return *service_; }
ServiceBinding() {}
explicit ServiceBinding(std::shared_ptr<T> service) : service_(service) {}
private:
std::shared_ptr<T> service_;
};
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_GRAPH_SERVICE_H_
mediapipe/framework/graph_service_manager.cc
+21
View File
@@ -0,0 +1,21 @@
#include "mediapipe/framework/graph_service_manager.h"
namespace mediapipe {
absl::Status GraphServiceManager::SetServicePacket(
const GraphServiceBase& service, Packet p) {
// TODO: check service is already set?
service_packets_[service.key] = std::move(p);
return absl::OkStatus();
}
Packet GraphServiceManager::GetServicePacket(
const GraphServiceBase& service) const {
auto it = service_packets_.find(service.key);
if (it == service_packets_.end()) {
return {};
}
return it->second;
}
} // namespace mediapipe
mediapipe/framework/graph_service_manager.h
+42
View File
@@ -0,0 +1,42 @@
#ifndef MEDIAPIPE_FRAMEWORK_GRAPH_SERVICE_MANAGER_H_
#define MEDIAPIPE_FRAMEWORK_GRAPH_SERVICE_MANAGER_H_
#include <memory>
#include "absl/status/status.h"
#include "mediapipe/framework/graph_service.h"
#include "mediapipe/framework/packet.h"
namespace mediapipe {
class GraphServiceManager {
public:
template <typename T>
absl::Status SetServiceObject(const GraphService<T>& service,
std::shared_ptr<T> object) {
return SetServicePacket(service,
MakePacket<std::shared_ptr<T>>(std::move(object)));
}
absl::Status SetServicePacket(const GraphServiceBase& service, Packet p);
template <typename T>
std::shared_ptr<T> GetServiceObject(const GraphService<T>& service) const {
Packet p = GetServicePacket(service);
if (p.IsEmpty()) return nullptr;
return p.Get<std::shared_ptr<T>>();
}
const std::map<std::string, Packet>& ServicePackets() {
return service_packets_;
}
private:
Packet GetServicePacket(const GraphServiceBase& service) const;
std::map<std::string, Packet> service_packets_;
};
} // namespace mediapipe
#endif // MEDIAPIPE_FRAMEWORK_GRAPH_SERVICE_MANAGER_H_
mediapipe/framework/graph_service_manager_test.cc
+53
View File
@@ -0,0 +1,53 @@
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
TEST(GraphServiceManager, SetGetServiceObject) {
GraphServiceManager service_manager;
constexpr GraphService<int> kIntService("mediapipe::IntService");
EXPECT_EQ(service_manager.GetServiceObject(kIntService), nullptr);
MP_EXPECT_OK(service_manager.SetServiceObject(kIntService,
std::make_shared<int>(100)));
ASSERT_NE(service_manager.GetServiceObject(kIntService), nullptr);
EXPECT_EQ(*service_manager.GetServiceObject(kIntService), 100);
}
TEST(GraphServiceManager, SetServicePacket) {
GraphServiceManager service_manager;
constexpr GraphService<int> kIntService("mediapipe::IntService");
MP_EXPECT_OK(service_manager.SetServicePacket(
kIntService,
mediapipe::MakePacket<std::shared_ptr<int>>(std::make_shared<int>(100))));
ASSERT_NE(service_manager.GetServiceObject(kIntService), nullptr);
EXPECT_EQ(*service_manager.GetServiceObject(kIntService), 100);
}
TEST(GraphServiceManager, ServicePackets) {
GraphServiceManager service_manager;
EXPECT_TRUE(service_manager.ServicePackets().empty());
constexpr GraphService<int> kIntService("mediapipe::IntService");
MP_EXPECT_OK(service_manager.SetServiceObject(kIntService,
std::make_shared<int>(100)));
EXPECT_EQ(service_manager.ServicePackets().size(), 1);
ASSERT_NE(service_manager.ServicePackets().find(kIntService.key),
service_manager.ServicePackets().end());
EXPECT_EQ(*service_manager.ServicePackets()
.at(kIntService.key)
.Get<std::shared_ptr<int>>(),
100);
}
} // namespace mediapipe
mediapipe/framework/input_stream_handler.cc
+8 -3
View File
@@ -365,9 +365,14 @@ NodeReadiness SyncSet::GetReadiness(Timestamp* min_stream_timestamp) {
}
} else {
// Any unprocessed input_ts can be processed.
// Note that (min_bound - 1) is the highest fully settled timestamp.
Timestamp input_timestamp =
std::min(min_packet, min_bound.PreviousAllowedInStream());
// The settled timestamp is the highest timestamp at which no future packets
// can arrive. Timestamp::PostStream is treated specially because it is
// omitted by Timestamp::PreviousAllowedInStream.
Timestamp settled =
(min_packet == Timestamp::PostStream() && min_bound > min_packet)
? min_packet
: min_bound.PreviousAllowedInStream();
Timestamp input_timestamp = std::min(min_packet, settled);
if (input_timestamp >
std::max(last_processed_ts_, Timestamp::Unstarted())) {
*min_stream_timestamp = input_timestamp;
mediapipe/framework/port/BUILD
-12
View File
@@ -89,18 +89,6 @@ cc_library(
],
)
cc_library(
name = "commandlineflags",
hdrs = [
"commandlineflags.h",
],
visibility = ["//visibility:public"],
deps = [
"//third_party:glog",
"@com_google_absl//absl/flags:flag",
],
)
cc_library(
name = "core_proto",
hdrs = [
mediapipe/framework/port/commandlineflags.h
-30
View File
@@ -1,30 +0,0 @@
// Copyright 2019 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_PORT_COMMANDLINEFLAGS_H_
#define MEDIAPIPE_PORT_COMMANDLINEFLAGS_H_
#include "gflags/gflags.h"
namespace absl {
template <typename T>
T GetFlag(const T& f) {
return f;
}
template <typename T, typename U>
void SetFlag(T* f, const U& u) {
*f = u;
}
} // namespace absl
#endif // MEDIAPIPE_PORT_COMMANDLINEFLAGS_H_
mediapipe/framework/profiler/BUILD
+1 -1
View File
@@ -235,7 +235,6 @@ cc_test(
"//mediapipe/framework/deps:clock",
"//mediapipe/framework/deps:message_matchers",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:commandlineflags",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:integral_types",
@@ -247,6 +246,7 @@ cc_test(
"//mediapipe/framework/tool:simulation_clock",
"//mediapipe/framework/tool:simulation_clock_executor",
"//mediapipe/framework/tool:status_util",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/time",
],
)
mediapipe/framework/profiler/graph_tracer_test.cc
+1 -1
View File
@@ -21,6 +21,7 @@
#include <utility>
#include <vector>
#include "absl/flags/flag.h"
#include "absl/time/time.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
@@ -28,7 +29,6 @@
#include "mediapipe/framework/deps/clock.h"
#include "mediapipe/framework/deps/message_matchers.h"
#include "mediapipe/framework/port/advanced_proto_inc.h"
#include "mediapipe/framework/port/commandlineflags.h"
#include "mediapipe/framework/port/file_helpers.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
mediapipe/framework/profiler/reporter/BUILD
+1 -1
View File
@@ -31,7 +31,6 @@ cc_library(
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_profile_cc_proto",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:commandlineflags",
"//mediapipe/framework/port:file_helpers",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:re2",
@@ -39,6 +38,7 @@ cc_library(
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:statusor",
"@com_google_absl//absl/container:btree",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
mediapipe/framework/subgraph.cc
+6 -6
View File
@@ -93,17 +93,17 @@ bool GraphRegistry::IsRegistered(const std::string& ns,
absl::StatusOr<CalculatorGraphConfig> GraphRegistry::CreateByName(
const std::string& ns, const std::string& type_name,
const Subgraph::SubgraphOptions* options) const {
Subgraph::SubgraphOptions graph_options;
if (options) {
graph_options = *options;
}
SubgraphContext* context) const {
absl::StatusOr<std::unique_ptr<Subgraph>> maker =
local_factories_.IsRegistered(ns, type_name)
? local_factories_.Invoke(ns, type_name)
: global_factories_->Invoke(ns, type_name);
MP_RETURN_IF_ERROR(maker.status());
return maker.value()->GetConfig(graph_options);
if (context != nullptr) {
return maker.value()->GetConfig(context);
}
SubgraphContext default_context;
return maker.value()->GetConfig(&default_context);
}
} // namespace mediapipe
mediapipe/framework/subgraph.h
+64 -3
View File
@@ -19,8 +19,12 @@
#include "absl/base/macros.h"
#include "absl/memory/memory.h"
#include "absl/status/status.h"
#include "absl/types/optional.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/deps/registration.h"
#include "mediapipe/framework/graph_service.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/framework/tool/calculator_graph_template.pb.h"
@@ -28,6 +32,51 @@
namespace mediapipe {
class SubgraphContext {
public:
SubgraphContext() : SubgraphContext(nullptr, nullptr) {}
// @node and/or @service_manager can be nullptr.
SubgraphContext(const CalculatorGraphConfig::Node* node,
const GraphServiceManager* service_manager)
: default_node_(node ? absl::nullopt
: absl::optional<CalculatorGraphConfig::Node>(
CalculatorGraphConfig::Node())),
original_node_(node ? *node : default_node_.value()),
default_service_manager_(
service_manager
? absl::nullopt
: absl::optional<GraphServiceManager>(GraphServiceManager())),
service_manager_(service_manager ? *service_manager
: default_service_manager_.value()),
options_map_(std::move(tool::OptionsMap().Initialize(original_node_))) {
}
template <typename T>
const T& Options() {
return options_map_.Get<T>();
}
const CalculatorGraphConfig::Node& OriginalNode() { return original_node_; }
template <typename T>
ServiceBinding<T> Service(const GraphService<T>& service) const {
return ServiceBinding<T>(service_manager_.GetServiceObject(service));
}
private:
// Populated if node is not provided during construction.
const absl::optional<CalculatorGraphConfig::Node> default_node_;
const CalculatorGraphConfig::Node& original_node_;
// Populated if service manager is not provided during construction.
const absl::optional<GraphServiceManager> default_service_manager_;
const GraphServiceManager& service_manager_;
tool::OptionsMap options_map_;
};
// Instances of this class are responsible for providing a subgraph config.
// They are only used during graph construction. They do not stay alive once
// the graph is running.
@@ -36,13 +85,25 @@ class Subgraph {
using SubgraphOptions = CalculatorGraphConfig::Node;
Subgraph();
virtual ~Subgraph();
// Returns the config to use for one instantiation of the subgraph. The
// nodes and generators in this config will replace the subgraph node in
// the parent graph.
// Subclasses may use the options argument to parameterize the config.
// Subclasses may use `SubgraphContext*` param to parameterize the config.
// TODO: make this static?
virtual absl::StatusOr<CalculatorGraphConfig> GetConfig(SubgraphContext* sc) {
if (sc == nullptr) {
return GetConfig(SubgraphOptions{});
}
return GetConfig(sc->OriginalNode());
}
// Kept for backward compatibility - please override `GetConfig` taking
// `SubgraphContext*` param.
virtual absl::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) = 0;
const SubgraphOptions& options) {
return absl::UnimplementedError("Not implemented.");
}
// Returns options of a specific type.
template <typename T>
@@ -120,7 +181,7 @@ class GraphRegistry {
// Returns the specified graph config.
absl::StatusOr<CalculatorGraphConfig> CreateByName(
const std::string& ns, const std::string& type_name,
const Subgraph::SubgraphOptions* options = nullptr) const;
SubgraphContext* context = nullptr) const;
static GraphRegistry global_graph_registry;
mediapipe/framework/subgraph_test.cc
+58
View File
@@ -14,11 +14,16 @@
#include "mediapipe/framework/subgraph.h"
#include <string>
#include "absl/strings/str_format.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
// Because of portability issues, we include this directly.
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status_matchers.h" // NOLINT(build/deprecated)
namespace mediapipe {
@@ -75,5 +80,58 @@ TEST_F(SubgraphTest, LinkedSubgraph) {
TestGraphEnclosing("DubQuadTestSubgraph");
}
const mediapipe::GraphService<std::string> kStringTestService{
"mediapipe::StringTestService"};
class EmitSideServiceStringTestSubgraph : public Subgraph {
public:
absl::StatusOr<CalculatorGraphConfig> GetConfig(
mediapipe::SubgraphContext* sc) override {
auto string_service = sc->Service(kStringTestService);
RET_CHECK(string_service.IsAvailable()) << "Service not available";
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
absl::StrFormat(R"(
output_side_packet: "string"
node {
calculator: "ConstantSidePacketCalculator"
output_side_packet: "PACKET:string"
options: {
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
packet { string_value: "%s" }
}
}
}
)",
string_service.GetObject()));
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(EmitSideServiceStringTestSubgraph);
TEST(SubgraphServicesTest, EmitStringFromTestService) {
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
output_side_packet: "str"
node {
calculator: "EmitSideServiceStringTestSubgraph"
output_side_packet: "str"
}
)");
Packet side_string;
tool::AddSidePacketSink("str", &config, &side_string);
CalculatorGraph graph;
// It's important that service object is set before Initialize()
MP_ASSERT_OK(graph.SetServiceObject(
kStringTestService, std::make_shared<std::string>("Expected STRING")));
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.WaitUntilIdle());
MP_ASSERT_OK(graph.WaitUntilDone());
EXPECT_EQ(side_string.Get<std::string>(), "Expected STRING");
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/BUILD
+28 -1
View File
@@ -19,6 +19,7 @@ load(
"data_as_c_string",
"mediapipe_binary_graph",
)
load("//mediapipe/framework:mediapipe_cc_test.bzl", "mediapipe_cc_test")
licenses(["notice"])
@@ -35,9 +36,10 @@ cc_library(
deps = [
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework/port:advanced_proto",
"//mediapipe/framework/port:commandlineflags",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/flags:parse",
],
)
@@ -150,6 +152,28 @@ cc_library(
],
)
mediapipe_cc_test(
name = "options_util_test",
size = "small",
srcs = ["options_util_test.cc"],
data = [":node_chain_subgraph.proto"],
requires_full_emulation = False,
deps = [
":options_util",
"//mediapipe/calculators/core:flow_limiter_calculator",
"//mediapipe/calculators/core:flow_limiter_calculator_cc_proto",
"//mediapipe/framework:basic_types_registration",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:calculator_runner",
"//mediapipe/framework:validated_graph_config",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
"//mediapipe/framework/testdata:night_light_calculator_cc_proto",
"//mediapipe/util:header_util",
],
)
cc_library(
name = "packet_util",
hdrs = ["packet_util.h"],
@@ -227,6 +251,7 @@ cc_library(
":name_util",
":tag_map",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:graph_service_manager",
"//mediapipe/framework:packet_generator",
"//mediapipe/framework:packet_generator_cc_proto",
"//mediapipe/framework:port",
@@ -560,7 +585,9 @@ cc_test(
":subgraph_expansion",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:graph_service_manager",
"//mediapipe/framework:mediapipe_options_cc_proto",
"//mediapipe/framework:packet",
"//mediapipe/framework:packet_generator_cc_proto",
"//mediapipe/framework:packet_set",
"//mediapipe/framework:packet_type",
mediapipe/framework/tool/subgraph_expansion.cc
+7 -4
View File
@@ -25,6 +25,7 @@
#include "absl/memory/memory.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/packet_generator.pb.h"
#include "mediapipe/framework/port.h"
#include "mediapipe/framework/port/core_proto_inc.h"
@@ -273,7 +274,8 @@ absl::Status ConnectSubgraphStreams(
}
absl::Status ExpandSubgraphs(CalculatorGraphConfig* config,
const GraphRegistry* graph_registry) {
const GraphRegistry* graph_registry,
const GraphServiceManager* service_manager) {
graph_registry =
graph_registry ? graph_registry : &GraphRegistry::global_graph_registry;
RET_CHECK(config);
@@ -292,9 +294,10 @@ absl::Status ExpandSubgraphs(CalculatorGraphConfig* config,
int node_id = it - nodes->begin();
std::string node_name = CanonicalNodeName(*config, node_id);
MP_RETURN_IF_ERROR(ValidateSubgraphFields(node));
ASSIGN_OR_RETURN(auto subgraph,
graph_registry->CreateByName(config->package(),
node.calculator(), &node));
SubgraphContext subgraph_context(&node, service_manager);
ASSIGN_OR_RETURN(auto subgraph, graph_registry->CreateByName(
config->package(), node.calculator(),
&subgraph_context));
MP_RETURN_IF_ERROR(PrefixNames(node_name, &subgraph));
MP_RETURN_IF_ERROR(ConnectSubgraphStreams(node, &subgraph));
subgraphs.push_back(subgraph);
mediapipe/framework/tool/subgraph_expansion.h
+5 -2
View File
@@ -19,6 +19,7 @@
#include "absl/strings/string_view.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/port/proto_ns.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/subgraph.h"
@@ -68,8 +69,10 @@ absl::Status ConnectSubgraphStreams(
// Replaces subgraph nodes in the given config with the contents of the
// corresponding subgraphs. Nested subgraphs are retrieved from the
// graph registry and expanded recursively.
absl::Status ExpandSubgraphs(CalculatorGraphConfig* config,
const GraphRegistry* graph_registry = nullptr);
absl::Status ExpandSubgraphs(
CalculatorGraphConfig* config,
const GraphRegistry* graph_registry = nullptr,
const GraphServiceManager* service_manager = nullptr);
// Creates a graph wrapping the provided node and exposing all of its
// connections
mediapipe/framework/tool/subgraph_expansion_test.cc
+38
View File
@@ -19,6 +19,8 @@
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/message_matchers.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/packet_set.h"
#include "mediapipe/framework/packet_type.h"
#include "mediapipe/framework/port/gmock.h"
@@ -526,5 +528,41 @@ TEST(SubgraphExpansionTest, ExecutorFieldOfNodeInSubgraphPreserved) {
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
const mediapipe::GraphService<std::string> kStringTestService{
"mediapipe::StringTestService"};
class GraphServicesClientTestSubgraph : public Subgraph {
public:
absl::StatusOr<CalculatorGraphConfig> GetConfig(
SubgraphContext* sc) override {
auto string_service = sc->Service(kStringTestService);
RET_CHECK(string_service.IsAvailable()) << "Service not available";
CalculatorGraphConfig config;
config.add_node()->set_calculator(string_service.GetObject());
return config;
}
};
REGISTER_MEDIAPIPE_GRAPH(GraphServicesClientTestSubgraph);
TEST(SubgraphExpansionTest, GraphServicesUsage) {
CalculatorGraphConfig supergraph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node { calculator: "GraphServicesClientTestSubgraph" }
)");
CalculatorGraphConfig expected_graph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
name: "graphservicesclienttestsubgraph__ExpectedNode"
calculator: "ExpectedNode"
}
)");
GraphServiceManager service_manager;
MP_ASSERT_OK(service_manager.SetServiceObject(
kStringTestService, std::make_shared<std::string>("ExpectedNode")));
MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph, /*graph_registry=*/nullptr,
&service_manager));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/text_to_binary_graph.cc
+8 -8
View File
@@ -19,19 +19,19 @@
#include <fstream>
#include <string>
#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/port/advanced_proto_inc.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/commandlineflags.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
DEFINE_string(proto_source, "",
"The template source file containing CalculatorGraphConfig "
"protobuf text with inline template params.");
DEFINE_string(
proto_output, "",
"An output template file in binary CalculatorGraphTemplate form.");
ABSL_FLAG(std::string, proto_source, "",
"The template source file containing CalculatorGraphConfig "
"protobuf text with inline template params.");
ABSL_FLAG(std::string, proto_output, "",
"An output template file in binary CalculatorGraphTemplate form.");
#define EXIT_IF_ERROR(status) \
if (!status.ok()) { \
@@ -92,7 +92,7 @@ absl::Status WriteFile(const std::string& proto_output, bool write_text,
int main(int argc, char** argv) {
google::InitGoogleLogging(argv[0]);
gflags::ParseCommandLineFlags(&argc, &argv, true);
absl::ParseCommandLine(argc, argv);
// Validate command line options.
absl::Status status;
mediapipe/framework/validated_graph_config.cc
+17 -10
View File
@@ -21,6 +21,7 @@
#include "absl/strings/substitute.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_base.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/legacy_calculator_support.h"
#include "mediapipe/framework/packet_generator.h"
#include "mediapipe/framework/packet_generator.pb.h"
@@ -142,10 +143,11 @@ absl::Status AddPredefinedExecutorConfigs(CalculatorGraphConfig* graph_config) {
absl::Status PerformBasicTransforms(
const CalculatorGraphConfig& input_graph_config,
const GraphRegistry* graph_registry,
const GraphServiceManager* service_manager,
CalculatorGraphConfig* output_graph_config) {
*output_graph_config = input_graph_config;
MP_RETURN_IF_ERROR(
tool::ExpandSubgraphs(output_graph_config, graph_registry));
MP_RETURN_IF_ERROR(tool::ExpandSubgraphs(output_graph_config, graph_registry,
service_manager));
MP_RETURN_IF_ERROR(AddPredefinedExecutorConfigs(output_graph_config));
@@ -344,7 +346,8 @@ absl::Status NodeTypeInfo::Initialize(
absl::Status ValidatedGraphConfig::Initialize(
const CalculatorGraphConfig& input_config,
const GraphRegistry* graph_registry) {
const GraphRegistry* graph_registry,
const GraphServiceManager* service_manager) {
RET_CHECK(!initialized_)
<< "ValidatedGraphConfig can be initialized only once.";
@@ -353,8 +356,8 @@ absl::Status ValidatedGraphConfig::Initialize(
<< input_config.DebugString();
#endif
MP_RETURN_IF_ERROR(
PerformBasicTransforms(input_config, graph_registry, &config_));
MP_RETURN_IF_ERROR(PerformBasicTransforms(input_config, graph_registry,
service_manager, &config_));
// Initialize the basic node information.
MP_RETURN_IF_ERROR(InitializeGeneratorInfo());
@@ -429,18 +432,22 @@ absl::Status ValidatedGraphConfig::Initialize(
absl::Status ValidatedGraphConfig::Initialize(
const std::string& graph_type, const Subgraph::SubgraphOptions* options,
const GraphRegistry* graph_registry) {
const GraphRegistry* graph_registry,
const GraphServiceManager* service_manager) {
graph_registry =
graph_registry ? graph_registry : &GraphRegistry::global_graph_registry;
auto status_or_config = graph_registry->CreateByName("", graph_type, options);
SubgraphContext subgraph_context(options, service_manager);
auto status_or_config =
graph_registry->CreateByName("", graph_type, &subgraph_context);
MP_RETURN_IF_ERROR(status_or_config.status());
return Initialize(status_or_config.value(), graph_registry);
return Initialize(status_or_config.value(), graph_registry, service_manager);
}
absl::Status ValidatedGraphConfig::Initialize(
const std::vector<CalculatorGraphConfig>& input_configs,
const std::vector<CalculatorGraphTemplate>& input_templates,
const std::string& graph_type, const Subgraph::SubgraphOptions* options) {
const std::string& graph_type, const Subgraph::SubgraphOptions* arguments,
const GraphServiceManager* service_manager) {
GraphRegistry graph_registry;
for (auto& config : input_configs) {
graph_registry.Register(config.type(), config);
@@ -448,7 +455,7 @@ absl::Status ValidatedGraphConfig::Initialize(
for (auto& templ : input_templates) {
graph_registry.Register(templ.config().type(), templ);
}
return Initialize(graph_type, options, &graph_registry);
return Initialize(graph_type, arguments, &graph_registry, service_manager);
}
absl::Status ValidatedGraphConfig::InitializeCalculatorInfo() {
mediapipe/framework/validated_graph_config.h
+7 -3
View File
@@ -21,6 +21,7 @@
#include "absl/container/flat_hash_set.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_contract.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/packet_generator.pb.h"
#include "mediapipe/framework/packet_type.h"
#include "mediapipe/framework/port/map_util.h"
@@ -195,7 +196,8 @@ class ValidatedGraphConfig {
// before any other functions. Subgraphs are specified through the
// global graph registry or an optional local graph registry.
absl::Status Initialize(const CalculatorGraphConfig& input_config,
const GraphRegistry* graph_registry = nullptr);
const GraphRegistry* graph_registry = nullptr,
const GraphServiceManager* service_manager = nullptr);
// Initializes the ValidatedGraphConfig from registered graph and subgraph
// configs. Subgraphs are retrieved from the specified graph registry or from
@@ -203,7 +205,8 @@ class ValidatedGraphConfig {
// specifying its type in |graph_type|.
absl::Status Initialize(const std::string& graph_type,
const Subgraph::SubgraphOptions* options = nullptr,
const GraphRegistry* graph_registry = nullptr);
const GraphRegistry* graph_registry = nullptr,
const GraphServiceManager* service_manager = nullptr);
// Initializes the ValidatedGraphConfig from the specified graph and subgraph
// configs. Template graph and subgraph configs can be specified through
@@ -215,7 +218,8 @@ class ValidatedGraphConfig {
const std::vector<CalculatorGraphConfig>& input_configs,
const std::vector<CalculatorGraphTemplate>& input_templates,
const std::string& graph_type = "",
const Subgraph::SubgraphOptions* arguments = nullptr);
const Subgraph::SubgraphOptions* arguments = nullptr,
const GraphServiceManager* service_manager = nullptr);
// Returns true if the ValidatedGraphConfig has been initialized.
bool Initialized() const { return initialized_; }
mediapipe/framework/validated_graph_config_test.cc
+165
View File
@@ -0,0 +1,165 @@
#include "mediapipe/framework/validated_graph_config.h"
#include <string_view>
#include "absl/status/status.h"
#include "absl/strings/ascii.h"
#include "absl/strings/str_cat.h"
#include "mediapipe/framework/api2/node.h"
#include "mediapipe/framework/api2/port.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/deps/message_matchers.h"
#include "mediapipe/framework/graph_service.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
class NoOp : public mediapipe::api2::Node {
public:
static constexpr mediapipe::api2::Input<int>::Optional kInputNotNeeded{"NN"};
static constexpr mediapipe::api2::Output<int>::Optional kOutputNotNeeded{
"NN"};
MEDIAPIPE_NODE_CONTRACT(kInputNotNeeded, kOutputNotNeeded);
absl::Status Process(CalculatorContext* cc) override {
return absl::OkStatus();
}
};
using CalculatorA = NoOp;
MEDIAPIPE_REGISTER_NODE(CalculatorA);
using CalculatorB = NoOp;
MEDIAPIPE_REGISTER_NODE(CalculatorB);
using CalculatorC = NoOp;
MEDIAPIPE_REGISTER_NODE(CalculatorC);
CalculatorGraphConfig ExpectedConfig(const std::string& node_name) {
CalculatorGraphConfig config;
config.add_node()->set_calculator(node_name);
config.add_executor();
return config;
}
CalculatorGraphConfig ExpectedConfigExpandedFromGraph(
const std::string& graph_name, const std::string& node_name) {
CalculatorGraphConfig config;
auto* node = config.add_node();
node->set_calculator(node_name);
node->set_name(
absl::StrCat(absl::AsciiStrToLower(graph_name), "__", node_name));
config.add_executor();
return config;
}
class AlwaysCalculatorALegacySubgraph : public Subgraph {
absl::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) override {
return ExpectedConfig("CalculatorA");
}
};
REGISTER_MEDIAPIPE_GRAPH(AlwaysCalculatorALegacySubgraph);
TEST(ValidatedGraphConfigTest, InitializeByTypeLegacySubgraphHardcoded) {
ValidatedGraphConfig config;
MP_EXPECT_OK(config.Initialize("AlwaysCalculatorALegacySubgraph",
/*options=*/nullptr,
/*graph_registry=*/nullptr,
/*service_manager=*/nullptr));
ASSERT_TRUE(config.Initialized());
EXPECT_THAT(config.Config(), EqualsProto(ExpectedConfig("CalculatorA")));
}
TEST(ValidatedGraphConfigTest, InitializeLegacySubgraphHardcoded) {
CalculatorGraphConfig graph;
graph.add_node()->set_calculator("AlwaysCalculatorALegacySubgraph");
ValidatedGraphConfig config;
MP_EXPECT_OK(config.Initialize(graph,
/*graph_registry=*/nullptr,
/*service_manager=*/nullptr));
ASSERT_TRUE(config.Initialized());
EXPECT_THAT(config.Config(),
EqualsProto(ExpectedConfigExpandedFromGraph(
"AlwaysCalculatorALegacySubgraph", "CalculatorA")));
}
class AlwaysCalculatorASubgraph : public Subgraph {
absl::StatusOr<CalculatorGraphConfig> GetConfig(
SubgraphContext* sc) override {
return ExpectedConfig("CalculatorA");
}
};
REGISTER_MEDIAPIPE_GRAPH(AlwaysCalculatorASubgraph);
TEST(ValidatedGraphConfigTest, InitializeByTypeSubgraphHardcoded) {
ValidatedGraphConfig config;
MP_EXPECT_OK(config.Initialize("AlwaysCalculatorASubgraph",
/*options=*/nullptr,
/*graph_registry=*/nullptr,
/*service_manager=*/nullptr));
ASSERT_TRUE(config.Initialized());
EXPECT_THAT(config.Config(), EqualsProto(ExpectedConfig("CalculatorA")));
}
TEST(ValidatedGraphConfigTest, InitializeSubgraphHardcoded) {
CalculatorGraphConfig graph;
graph.add_node()->set_calculator("AlwaysCalculatorASubgraph");
ValidatedGraphConfig config;
MP_EXPECT_OK(config.Initialize(graph,
/*graph_registry=*/nullptr,
/*service_manager=*/nullptr));
ASSERT_TRUE(config.Initialized());
EXPECT_THAT(config.Config(),
EqualsProto(ExpectedConfigExpandedFromGraph(
"AlwaysCalculatorASubgraph", "CalculatorA")));
}
const mediapipe::GraphService<std::string> kStringTestService{
"mediapipe::StringTestService"};
class TestServiceSubgraph : public Subgraph {
absl::StatusOr<CalculatorGraphConfig> GetConfig(
SubgraphContext* sc) override {
return ExpectedConfig(sc->Service(kStringTestService).GetObject());
}
};
REGISTER_MEDIAPIPE_GRAPH(TestServiceSubgraph);
TEST(ValidatedGraphConfigTest, InitializeByTypeSubgraphWithServiceCalculatorB) {
for (const std::string& calculator_name :
{"CalculatorA", "CalculatorB", "CalculatorC"}) {
ValidatedGraphConfig config;
GraphServiceManager service_manager;
MP_ASSERT_OK(service_manager.SetServiceObject(
kStringTestService, std::make_shared<std::string>(calculator_name)));
MP_EXPECT_OK(config.Initialize("TestServiceSubgraph",
/*options=*/nullptr,
/*graph_registry=*/nullptr,
/*service_manager=*/&service_manager));
ASSERT_TRUE(config.Initialized());
EXPECT_THAT(config.Config(), EqualsProto(ExpectedConfig(calculator_name)));
}
}
TEST(ValidatedGraphConfigTest, InitializeSubgraphWithServiceCalculatorB) {
for (const std::string& calculator_name :
{"CalculatorA", "CalculatorB", "CalculatorC"}) {
CalculatorGraphConfig graph;
graph.add_node()->set_calculator("TestServiceSubgraph");
ValidatedGraphConfig config;
GraphServiceManager service_manager;
MP_ASSERT_OK(service_manager.SetServiceObject(
kStringTestService, std::make_shared<std::string>(calculator_name)));
MP_EXPECT_OK(config.Initialize(graph,
/*graph_registry=*/nullptr,
/*service_manager=*/&service_manager));
ASSERT_TRUE(config.Initialized());
EXPECT_THAT(config.Config(), EqualsProto(ExpectedConfigExpandedFromGraph(
"TestServiceSubgraph", calculator_name)));
}
}
} // namespace mediapipe