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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: bbbbcb4f5174dea33525729ede47c770069157cd
This commit is contained in:
MediaPipe Team
2021-10-18 17:00:29 -04:00
committed by chuoling
parent 33d683c671
commit 1faeaae7e5
75 changed files with 1944 additions and 560 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/tool/BUILD
+1
View File
@@ -724,6 +724,7 @@ cc_test(
srcs = ["subgraph_expansion_test.cc"],
deps = [
":node_chain_subgraph_cc_proto",
":node_chain_subgraph_options_lib",
":subgraph_expansion",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework:calculator_framework",
mediapipe/framework/tool/proto_util_lite.cc
+14 -12
View File
@@ -23,6 +23,8 @@
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/type_map.h"
#define RET_CHECK_NO_LOG(cond) RET_CHECK(cond).SetNoLogging()
namespace mediapipe {
namespace tool {
@@ -47,13 +49,13 @@ absl::Status ReadFieldValue(uint32 tag, CodedInputStream* in,
WireFormatLite::WireType wire_type = WireFormatLite::GetTagWireType(tag);
if (IsLengthDelimited(wire_type)) {
uint32 length;
RET_CHECK(in->ReadVarint32(&length));
RET_CHECK(in->ReadString(result, length));
RET_CHECK_NO_LOG(in->ReadVarint32(&length));
RET_CHECK_NO_LOG(in->ReadString(result, length));
} else {
std::string field_data;
StringOutputStream sos(&field_data);
CodedOutputStream cos(&sos);
RET_CHECK(WireFormatLite::SkipField(in, tag, &cos));
RET_CHECK_NO_LOG(WireFormatLite::SkipField(in, tag, &cos));
// Skip the tag written by SkipField.
int tag_size = CodedOutputStream::VarintSize32(tag);
cos.Trim();
@@ -67,13 +69,13 @@ absl::Status ReadPackedValues(WireFormatLite::WireType wire_type,
CodedInputStream* in,
std::vector<std::string>* field_values) {
uint32 data_size;
RET_CHECK(in->ReadVarint32(&data_size));
RET_CHECK_NO_LOG(in->ReadVarint32(&data_size));
// fake_tag encodes the wire-type for calls to WireFormatLite::SkipField.
uint32 fake_tag = WireFormatLite::MakeTag(1, wire_type);
while (data_size > 0) {
std::string number;
MP_RETURN_IF_ERROR(ReadFieldValue(fake_tag, in, &number));
RET_CHECK_LE(number.size(), data_size);
RET_CHECK_NO_LOG(number.size() <= data_size);
field_values->push_back(number);
data_size -= number.size();
}
@@ -98,7 +100,7 @@ absl::Status GetFieldValues(uint32 field_id, WireFormatLite::WireType wire_type,
field_values->push_back(value);
}
} else {
RET_CHECK(WireFormatLite::SkipField(in, tag, out));
RET_CHECK_NO_LOG(WireFormatLite::SkipField(in, tag, out));
}
}
return absl::OkStatus();
@@ -157,12 +159,12 @@ absl::Status ProtoUtilLite::ReplaceFieldRange(
MP_RETURN_IF_ERROR(access.SetMessage(*message));
std::vector<std::string>& v = *access.mutable_field_values();
if (!proto_path.empty()) {
RET_CHECK(index >= 0 && index < v.size());
RET_CHECK_NO_LOG(index >= 0 && index < v.size());
MP_RETURN_IF_ERROR(ReplaceFieldRange(&v[index], proto_path, length,
field_type, field_values));
} else {
RET_CHECK(index >= 0 && index <= v.size());
RET_CHECK(index + length >= 0 && index + length <= v.size());
RET_CHECK_NO_LOG(index >= 0 && index <= v.size());
RET_CHECK_NO_LOG(index + length >= 0 && index + length <= v.size());
v.erase(v.begin() + index, v.begin() + index + length);
v.insert(v.begin() + index, field_values.begin(), field_values.end());
}
@@ -184,12 +186,12 @@ absl::Status ProtoUtilLite::GetFieldRange(
MP_RETURN_IF_ERROR(access.SetMessage(message));
std::vector<std::string>& v = *access.mutable_field_values();
if (!proto_path.empty()) {
RET_CHECK(index >= 0 && index < v.size());
RET_CHECK_NO_LOG(index >= 0 && index < v.size());
MP_RETURN_IF_ERROR(
GetFieldRange(v[index], proto_path, length, field_type, field_values));
} else {
RET_CHECK(index >= 0 && index <= v.size());
RET_CHECK(index + length >= 0 && index + length <= v.size());
RET_CHECK_NO_LOG(index >= 0 && index <= v.size());
RET_CHECK_NO_LOG(index + length >= 0 && index + length <= v.size());
field_values->insert(field_values->begin(), v.begin() + index,
v.begin() + index + length);
}
mediapipe/framework/tool/subgraph_expansion.cc
+3 -1
View File
@@ -274,12 +274,14 @@ absl::Status ConnectSubgraphStreams(
absl::Status ExpandSubgraphs(CalculatorGraphConfig* config,
const GraphRegistry* graph_registry,
const Subgraph::SubgraphOptions* graph_options,
const GraphServiceManager* service_manager) {
graph_registry =
graph_registry ? graph_registry : &GraphRegistry::global_graph_registry;
RET_CHECK(config);
MP_RETURN_IF_ERROR(mediapipe::tool::DefineGraphOptions(
CalculatorGraphConfig::Node(), config));
graph_options ? *graph_options : CalculatorGraphConfig::Node(), config));
auto* nodes = config->mutable_node();
while (1) {
auto subgraph_nodes_start = std::stable_partition(
mediapipe/framework/tool/subgraph_expansion.h
+1
View File
@@ -72,6 +72,7 @@ absl::Status ConnectSubgraphStreams(
absl::Status ExpandSubgraphs(
CalculatorGraphConfig* config,
const GraphRegistry* graph_registry = nullptr,
const Subgraph::SubgraphOptions* graph_options = nullptr,
const GraphServiceManager* service_manager = nullptr);
// Creates a graph wrapping the provided node and exposing all of its
mediapipe/framework/tool/subgraph_expansion_test.cc
+102
View File
@@ -560,9 +560,111 @@ TEST(SubgraphExpansionTest, GraphServicesUsage) {
MP_ASSERT_OK(service_manager.SetServiceObject(
kStringTestService, std::make_shared<std::string>("ExpectedNode")));
MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph, /*graph_registry=*/nullptr,
/*graph_options=*/nullptr,
&service_manager));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
// Shows SubgraphOptions consumed by GraphRegistry::CreateByName.
TEST(SubgraphExpansionTest, SubgraphOptionsUsage) {
EXPECT_TRUE(SubgraphRegistry::IsRegistered("NodeChainSubgraph"));
GraphRegistry graph_registry;
// CalculatorGraph::Initialize passes the SubgraphOptions into:
// (1) GraphRegistry::CreateByName("NodeChainSubgraph", options)
// (2) tool::ExpandSubgraphs(&config, options)
auto graph_options =
mediapipe::ParseTextProtoOrDie<Subgraph::SubgraphOptions>(R"pb(
options {
[mediapipe.NodeChainSubgraphOptions.ext] {
node_type: "DoubleIntCalculator"
chain_length: 3
}
})pb");
SubgraphContext context(&graph_options, /*service_manager=*/nullptr);
// "NodeChainSubgraph" consumes graph_options only in CreateByName.
auto subgraph_status =
graph_registry.CreateByName("", "NodeChainSubgraph", &context);
MP_ASSERT_OK(subgraph_status);
auto subgraph = std::move(subgraph_status).value();
MP_ASSERT_OK(
tool::ExpandSubgraphs(&subgraph, &graph_registry, &graph_options));
CalculatorGraphConfig expected_graph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
node {
calculator: "DoubleIntCalculator"
input_stream: "stream_0"
output_stream: "stream_1"
}
node {
calculator: "DoubleIntCalculator"
input_stream: "stream_1"
output_stream: "stream_2"
}
node {
calculator: "DoubleIntCalculator"
input_stream: "stream_2"
output_stream: "stream_3"
}
input_stream: "INPUT:stream_0"
output_stream: "OUTPUT:stream_3"
)pb");
EXPECT_THAT(subgraph, mediapipe::EqualsProto(expected_graph));
}
// Shows SubgraphOptions consumed by tool::ExpandSubgraphs.
TEST(SubgraphExpansionTest, SimpleSubgraphOptionsUsage) {
EXPECT_TRUE(SubgraphRegistry::IsRegistered("NodeChainSubgraph"));
GraphRegistry graph_registry;
auto moon_options =
mediapipe::ParseTextProtoOrDie<Subgraph::SubgraphOptions>(R"pb(
options {
[mediapipe.NodeChainSubgraphOptions.ext] {
node_type: "DoubleIntCalculator"
chain_length: 3
}
})pb");
auto moon_subgraph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
type: "MoonSubgraph"
graph_options: {
[type.googleapis.com/mediapipe.NodeChainSubgraphOptions] {}
}
node: {
calculator: "MoonCalculator"
node_options: {
[type.googleapis.com/mediapipe.NodeChainSubgraphOptions] {}
}
option_value: "chain_length:options/chain_length"
}
)pb");
// The moon_options are copied into the graph_options of moon_subgraph.
MP_ASSERT_OK(
tool::ExpandSubgraphs(&moon_subgraph, &graph_registry, &moon_options));
// The field chain_length is copied from moon_options into MoonCalculator.
CalculatorGraphConfig expected_graph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
node {
calculator: "MoonCalculator"
node_options {
[type.googleapis.com/mediapipe.NodeChainSubgraphOptions] {
chain_length: 3
}
}
option_value: "chain_length:options/chain_length"
}
type: "MoonSubgraph"
graph_options {
[type.googleapis.com/mediapipe.NodeChainSubgraphOptions] {}
}
)pb");
EXPECT_THAT(moon_subgraph, mediapipe::EqualsProto(expected_graph));
}
} // namespace
} // namespace mediapipe
mediapipe/framework/tool/switch_container_test.cc
+6 -4
View File
@@ -150,12 +150,13 @@ void RunTestContainer(CalculatorGraphConfig supergraph,
const int packet_count = 10;
// Send int value packets at {10K, 20K, 30K, ..., 100K}.
for (uint64 t = 1; t <= packet_count; ++t) {
MP_EXPECT_OK(graph.AddPacketToInputStream(
"foo", MakePacket<int>(t).At(Timestamp(t * 10000))));
if (send_bounds) {
MP_EXPECT_OK(graph.AddPacketToInputStream(
"enable", MakePacket<bool>(true).At(Timestamp(t * 10000))));
MP_ASSERT_OK(graph.WaitUntilIdle());
}
MP_EXPECT_OK(graph.AddPacketToInputStream(
"foo", MakePacket<int>(t).At(Timestamp(t * 10000))));
MP_ASSERT_OK(graph.WaitUntilIdle());
// The inputs are sent to the input stream "foo", they should pass through.
EXPECT_EQ(out_foo.size(), t);
@@ -175,12 +176,13 @@ void RunTestContainer(CalculatorGraphConfig supergraph,
// Send int value packets at {110K, 120K, ..., 200K}.
for (uint64 t = 11; t <= packet_count * 2; ++t) {
MP_EXPECT_OK(graph.AddPacketToInputStream(
"foo", MakePacket<int>(t).At(Timestamp(t * 10000))));
if (send_bounds) {
MP_EXPECT_OK(graph.AddPacketToInputStream(
"enable", MakePacket<bool>(false).At(Timestamp(t * 10000))));
MP_ASSERT_OK(graph.WaitUntilIdle());
}
MP_EXPECT_OK(graph.AddPacketToInputStream(
"foo", MakePacket<int>(t).At(Timestamp(t * 10000))));
MP_ASSERT_OK(graph.WaitUntilIdle());
// The inputs are sent to the input stream "foo", they should pass through.
EXPECT_EQ(out_foo.size(), t);