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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 852dfb05d450167899c0dd5ef7c45622a12e865b
This commit is contained in:
MediaPipe Team
2020-02-10 14:13:25 -08:00
committed by Hadon Nash
parent d144e564d8
commit de4fbc10e6
100 changed files with 1664 additions and 628 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/tool/subgraph_expansion.cc
+20 -17
View File
@@ -35,6 +35,7 @@
#include "mediapipe/framework/port/status_macros.h"
#include "mediapipe/framework/status_handler.pb.h"
#include "mediapipe/framework/subgraph.h"
#include "mediapipe/framework/tool/name_util.h"
#include "mediapipe/framework/tool/tag_map.h"
namespace mediapipe {
@@ -95,6 +96,13 @@ namespace tool {
config->mutable_output_side_packet()}) {
MP_RETURN_IF_ERROR(TransformStreamNames(streams, transform));
}
std::vector<std::string> node_names(config->node_size());
for (int node_id = 0; node_id < config->node_size(); ++node_id) {
node_names[node_id] = CanonicalNodeName(*config, node_id);
}
for (int node_id = 0; node_id < config->node_size(); ++node_id) {
config->mutable_node(node_id)->set_name(transform(node_names[node_id]));
}
for (auto& node : *config->mutable_node()) {
for (auto* streams :
{node.mutable_input_stream(), node.mutable_output_stream(),
@@ -102,9 +110,6 @@ namespace tool {
node.mutable_output_side_packet()}) {
MP_RETURN_IF_ERROR(TransformStreamNames(streams, transform));
}
if (!node.name().empty()) {
node.set_name(transform(node.name()));
}
}
for (auto& generator : *config->mutable_packet_generator()) {
for (auto* streams : {generator.mutable_input_side_packet(),
@@ -120,21 +125,18 @@ namespace tool {
}
// Adds a prefix to the name of each stream, side packet and node in the
// config. Each call to this method should use a different subgraph_index
// to produce a different numerical prefix. For example:
// 1, { foo, bar } --PrefixNames-> { __sg_1_foo, __sg_1_bar }
// 2, { foo, bar } --PrefixNames-> { __sg_2_foo, __sg_2_bar }
// config. Each call to this method should use a different prefix. For example:
// 1, { foo, bar } --PrefixNames-> { qsg__foo, qsg__bar }
// 2, { foo, bar } --PrefixNames-> { rsg__foo, rsg__bar }
// This means that two copies of the same subgraph will not interfere with
// each other.
static ::mediapipe::Status PrefixNames(int subgraph_index,
static ::mediapipe::Status PrefixNames(std::string prefix,
CalculatorGraphConfig* config) {
// TODO: prefix with subgraph name instead (see cl/157677233
// discussion).
// TODO: since we expand nested subgraphs outside-in, we should
// append the prefix to the existing prefix, if any. This is unimportant
// with the meaningless prefix we use now, but it should be considered
// when prefixing with names.
std::string prefix = absl::StrCat("__sg", subgraph_index, "_");
std::transform(prefix.begin(), prefix.end(), prefix.begin(), ::tolower);
std::replace(prefix.begin(), prefix.end(), '.', '_');
std::replace(prefix.begin(), prefix.end(), ' ', '_');
std::replace(prefix.begin(), prefix.end(), ':', '_');
absl::StrAppend(&prefix, "__");
auto add_prefix = [&prefix](absl::string_view s) {
return absl::StrCat(prefix, s);
};
@@ -271,7 +273,6 @@ static ::mediapipe::Status PrefixNames(int subgraph_index,
graph_registry ? graph_registry : &GraphRegistry::global_graph_registry;
RET_CHECK(config);
auto* nodes = config->mutable_node();
int subgraph_counter = 0;
while (1) {
auto subgraph_nodes_start = std::stable_partition(
nodes->begin(), nodes->end(),
@@ -283,11 +284,13 @@ static ::mediapipe::Status PrefixNames(int subgraph_index,
std::vector<CalculatorGraphConfig> subgraphs;
for (auto it = subgraph_nodes_start; it != nodes->end(); ++it) {
const auto& node = *it;
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));
MP_RETURN_IF_ERROR(PrefixNames(subgraph_counter++, &subgraph));
MP_RETURN_IF_ERROR(PrefixNames(node_name, &subgraph));
MP_RETURN_IF_ERROR(ConnectSubgraphStreams(node, &subgraph));
subgraphs.push_back(subgraph);
}