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

Project import generated by Copybara.

Browse Source 
GitOrigin-RevId: 796203faee20d7aae2876aac8ca5a1827dee4fe3
This commit is contained in:
MediaPipe Team
2019-09-30 11:26:36 -07:00
committed by jqtang
parent 412ab42d1f
commit a2a63e3876
122 changed files with 7330 additions and 2016 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mediapipe/framework/BUILD
+35
View File
@@ -1134,6 +1134,7 @@ cc_library(
"//mediapipe/framework/port:status",
"//mediapipe/framework/port:statusor",
"//mediapipe/framework/tool:calculator_graph_template_cc_proto",
"//mediapipe/framework/tool:options_util",
"//mediapipe/framework/tool:template_expander",
"@com_google_absl//absl/base:core_headers",
"@com_google_absl//absl/memory",
@@ -1499,13 +1500,47 @@ cc_test(
deps = [
":calculator_context",
":calculator_framework",
":test_calculators",
":thread_pool_executor",
":timestamp",
":type_map",
"//mediapipe/calculators/core:counting_source_calculator",
"//mediapipe/calculators/core:mux_calculator",
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:status",
"//mediapipe/framework/stream_handler:barrier_input_stream_handler",
"//mediapipe/framework/stream_handler:early_close_input_stream_handler",
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler",
"//mediapipe/framework/stream_handler:immediate_input_stream_handler",
"//mediapipe/framework/stream_handler:mux_input_stream_handler",
"//mediapipe/framework/tool:sink",
],
)
cc_test(
name = "calculator_graph_side_packet_test",
size = "small",
srcs = [
"calculator_graph_side_packet_test.cc",
],
visibility = ["//visibility:public"],
deps = [
":calculator_framework",
":test_calculators",
"//mediapipe/calculators/core:counting_source_calculator",
"//mediapipe/calculators/core:mux_calculator",
"//mediapipe/calculators/core:pass_through_calculator",
"//mediapipe/framework:calculator_cc_proto",
"//mediapipe/framework/port:gtest_main",
"//mediapipe/framework/port:logging",
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:sink",
"@com_google_absl//absl/time",
],
)
mediapipe/framework/calculator_graph_bounds_test.cc
+776 -1
View File
@@ -21,11 +21,258 @@
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_matchers.h"
#include "mediapipe/framework/thread_pool_executor.h"
#include "mediapipe/framework/timestamp.h"
namespace mediapipe {
namespace {
typedef std::function<::mediapipe::Status(CalculatorContext* cc)>
CalculatorContextFunction;
// A simple Semaphore for synchronizing test threads.
class AtomicSemaphore {
public:
AtomicSemaphore(int64_t supply) : supply_(supply) {}
void Acquire(int64_t amount) {
while (supply_.fetch_sub(amount) - amount < 0) {
Release(amount);
}
}
void Release(int64_t amount) { supply_ += amount; }
private:
std::atomic<int64_t> supply_;
};
// A mediapipe::Executor that signals the start and finish of each task.
class CountingExecutor : public Executor {
public:
CountingExecutor(int num_threads, std::function<void()> start_callback,
std::function<void()> finish_callback)
: thread_pool_(num_threads),
start_callback_(std::move(start_callback)),
finish_callback_(std::move(finish_callback)) {
thread_pool_.StartWorkers();
}
void Schedule(std::function<void()> task) override {
start_callback_();
thread_pool_.Schedule([this, task] {
task();
finish_callback_();
});
}
private:
ThreadPool thread_pool_;
std::function<void()> start_callback_;
std::function<void()> finish_callback_;
};
// A Calculator that adds the integer values in the packets in all the input
// streams and outputs the sum to the output stream.
class IntAdderCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
cc->Inputs().Index(i).Set<int>();
}
cc->Outputs().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->SetOffset(TimestampDiff(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
int sum = 0;
for (int i = 0; i < cc->Inputs().NumEntries(); ++i) {
sum += cc->Inputs().Index(i).Get<int>();
}
cc->Outputs().Index(0).Add(new int(sum), cc->InputTimestamp());
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(IntAdderCalculator);
template <typename InputType>
class TypedSinkCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<InputType>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) override {
return ::mediapipe::OkStatus();
}
};
typedef TypedSinkCalculator<std::string> StringSinkCalculator;
typedef TypedSinkCalculator<int> IntSinkCalculator;
REGISTER_CALCULATOR(StringSinkCalculator);
REGISTER_CALCULATOR(IntSinkCalculator);
// A Calculator that passes an input packet through if it contains an even
// integer.
class EvenIntFilterCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Outputs().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
int value = cc->Inputs().Index(0).Get<int>();
if (value % 2 == 0) {
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
} else {
cc->Outputs().Index(0).SetNextTimestampBound(
cc->InputTimestamp().NextAllowedInStream());
}
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(EvenIntFilterCalculator);
// A Calculator that passes packets through or not, depending on a second
// input. The first input stream's packets are only propagated if the second
// input stream carries the value true.
class ValveCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Inputs().Index(1).Set<bool>();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->Outputs().Index(0).SetHeader(cc->Inputs().Index(0).Header());
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
if (cc->Inputs().Index(1).Get<bool>()) {
cc->GetCounter("PassThrough")->Increment();
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
} else {
cc->GetCounter("Block")->Increment();
// The next timestamp bound is the minimum timestamp that the next packet
// can have, so, if we want to inform the downstream that no packet at
// InputTimestamp() is coming, we need to set it to the next value.
// We could also just call SetOffset(TimestampDiff(0)) in Open, and then
// we would not have to call this manually.
cc->Outputs().Index(0).SetNextTimestampBound(
cc->InputTimestamp().NextAllowedInStream());
}
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(ValveCalculator);
// A Calculator that simply passes its input Packets and header through,
// but shifts the timestamp.
class TimeShiftCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
cc->InputSidePackets().Index(0).Set<TimestampDiff>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
// Input: arbitrary Packets.
// Output: copy of the input.
cc->Outputs().Index(0).SetHeader(cc->Inputs().Index(0).Header());
shift_ = cc->InputSidePackets().Index(0).Get<TimestampDiff>();
cc->SetOffset(shift_);
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
cc->GetCounter("PassThrough")->Increment();
cc->Outputs().Index(0).AddPacket(
cc->Inputs().Index(0).Value().At(cc->InputTimestamp() + shift_));
return ::mediapipe::OkStatus();
}
private:
TimestampDiff shift_;
};
REGISTER_CALCULATOR(TimeShiftCalculator);
// A source calculator that alternates between outputting an integer (0, 1, 2,
// ..., 100) and setting the next timestamp bound. The timestamps of the output
// packets and next timestamp bounds are 0, 10, 20, 30, ...
//
// T=0 Output 0
// T=10 Set timestamp bound
// T=20 Output 1
// T=30 Set timestamp bound
// ...
// T=2000 Output 100
class OutputAndBoundSourceCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Outputs().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) override {
counter_ = 0;
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) override {
Timestamp timestamp(counter_);
if (counter_ % 20 == 0) {
cc->Outputs().Index(0).AddPacket(
MakePacket<int>(counter_ / 20).At(timestamp));
} else {
cc->Outputs().Index(0).SetNextTimestampBound(timestamp);
}
if (counter_ == 2000) {
return tool::StatusStop();
}
counter_ += 10;
return ::mediapipe::OkStatus();
}
private:
int counter_;
};
REGISTER_CALCULATOR(OutputAndBoundSourceCalculator);
// A calculator that outputs an initial packet of value 0 at time 0 in the
// Open() method, and then delays each input packet by 20 time units in the
// Process() method. The input stream and output stream have the integer type.
class Delay20Calculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Outputs().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->SetOffset(TimestampDiff(20));
cc->Outputs().Index(0).AddPacket(MakePacket<int>(0).At(Timestamp(0)));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
const Packet& packet = cc->Inputs().Index(0).Value();
Timestamp timestamp = packet.Timestamp() + 20;
cc->Outputs().Index(0).AddPacket(packet.At(timestamp));
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(Delay20Calculator);
class CustomBoundCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
@@ -45,8 +292,280 @@ class CustomBoundCalculator : public CalculatorBase {
};
REGISTER_CALCULATOR(CustomBoundCalculator);
// Test that SetNextTimestampBound propagates.
TEST(CalculatorGraph, SetNextTimestampBoundPropagation) {
CalculatorGraph graph;
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'in'
input_stream: 'gate'
node {
calculator: 'ValveCalculator'
input_stream: 'in'
input_stream: 'gate'
output_stream: 'gated'
}
node {
calculator: 'PassThroughCalculator'
input_stream: 'gated'
output_stream: 'passed'
}
node {
calculator: 'TimeShiftCalculator'
input_stream: 'passed'
output_stream: 'shifted'
input_side_packet: 'shift'
}
node {
calculator: 'MergeCalculator'
input_stream: 'in'
input_stream: 'shifted'
output_stream: 'merged'
}
node {
name: 'merged_output'
calculator: 'PassThroughCalculator'
input_stream: 'merged'
output_stream: 'out'
}
)");
Timestamp timestamp = Timestamp(0);
auto send_inputs = [&graph, &timestamp](int input, bool pass) {
++timestamp;
MP_EXPECT_OK(graph.AddPacketToInputStream(
"in", MakePacket<int>(input).At(timestamp)));
MP_EXPECT_OK(graph.AddPacketToInputStream(
"gate", MakePacket<bool>(pass).At(timestamp)));
};
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({{"shift", MakePacket<TimestampDiff>(0)}}));
auto pass_counter =
graph.GetCounterFactory()->GetCounter("ValveCalculator-PassThrough");
auto block_counter =
graph.GetCounterFactory()->GetCounter("ValveCalculator-Block");
auto merged_counter =
graph.GetCounterFactory()->GetCounter("merged_output-PassThrough");
send_inputs(1, true);
send_inputs(2, true);
send_inputs(3, false);
send_inputs(4, false);
MP_ASSERT_OK(graph.WaitUntilIdle());
// Verify that MergeCalculator was able to run even when the gated branch
// was blocked.
EXPECT_EQ(2, pass_counter->Get());
EXPECT_EQ(2, block_counter->Get());
EXPECT_EQ(4, merged_counter->Get());
send_inputs(5, true);
send_inputs(6, false);
MP_ASSERT_OK(graph.WaitUntilIdle());
EXPECT_EQ(3, pass_counter->Get());
EXPECT_EQ(3, block_counter->Get());
EXPECT_EQ(6, merged_counter->Get());
MP_ASSERT_OK(graph.CloseAllInputStreams());
MP_ASSERT_OK(graph.WaitUntilDone());
// Now test with time shift
MP_ASSERT_OK(graph.StartRun({{"shift", MakePacket<TimestampDiff>(-1)}}));
send_inputs(7, true);
MP_ASSERT_OK(graph.WaitUntilIdle());
// The merger should have run only once now, at timestamp 6, with inputs
// <null, 7>. If we do not respect the offset and unblock the merger for
// timestamp 7 too, then it will have run twice, with 6: <null,7> and
// 7: <7, null>.
EXPECT_EQ(4, pass_counter->Get());
EXPECT_EQ(3, block_counter->Get());
EXPECT_EQ(7, merged_counter->Get());
MP_ASSERT_OK(graph.CloseAllInputStreams());
MP_ASSERT_OK(graph.WaitUntilDone());
EXPECT_EQ(4, pass_counter->Get());
EXPECT_EQ(3, block_counter->Get());
EXPECT_EQ(8, merged_counter->Get());
}
// Both input streams of the calculator node have the same next timestamp
// bound. One input stream has a packet at that timestamp. The other input
// stream is empty. We should not run the Process() method of the node in this
// case.
TEST(CalculatorGraph, NotAllInputPacketsAtNextTimestampBoundAvailable) {
//
// in0_unfiltered in1_to_be_filtered
// | |
// | V
// | +-----------------------+
// | |EvenIntFilterCalculator|
// | +-----------------------+
// | |
// \ /
// \ / in1_filtered
// \ /
// | |
// V V
// +------------------+
// |IntAdderCalculator|
// +------------------+
// |
// V
// out
//
CalculatorGraph graph;
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'in0_unfiltered'
input_stream: 'in1_to_be_filtered'
node {
calculator: 'EvenIntFilterCalculator'
input_stream: 'in1_to_be_filtered'
output_stream: 'in1_filtered'
}
node {
calculator: 'IntAdderCalculator'
input_stream: 'in0_unfiltered'
input_stream: 'in1_filtered'
output_stream: 'out'
}
)");
std::vector<Packet> packet_dump;
tool::AddVectorSink("out", &config, &packet_dump);
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
Timestamp timestamp = Timestamp(0);
// We send an integer with timestamp 1 to the in0_unfiltered input stream of
// the IntAdderCalculator. We then send an even integer with timestamp 1 to
// the EvenIntFilterCalculator. This packet will go through and
// the IntAdderCalculator will run. The next timestamp bounds of both the
// input streams of the IntAdderCalculator will become 2.
++timestamp; // Timestamp 1.
MP_EXPECT_OK(graph.AddPacketToInputStream("in0_unfiltered",
MakePacket<int>(1).At(timestamp)));
MP_EXPECT_OK(graph.AddPacketToInputStream("in1_to_be_filtered",
MakePacket<int>(2).At(timestamp)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(1, packet_dump.size());
EXPECT_EQ(3, packet_dump[0].Get<int>());
// We send an odd integer with timestamp 2 to the EvenIntFilterCalculator.
// This packet will be filtered out and the next timestamp bound of the
// in1_filtered input stream of the IntAdderCalculator will become 3.
++timestamp; // Timestamp 2.
MP_EXPECT_OK(graph.AddPacketToInputStream("in1_to_be_filtered",
MakePacket<int>(3).At(timestamp)));
// We send an integer with timestamp 3 to the in0_unfiltered input stream of
// the IntAdderCalculator. MediaPipe should propagate the next timestamp bound
// across the IntAdderCalculator but should not run its Process() method.
++timestamp; // Timestamp 3.
MP_EXPECT_OK(graph.AddPacketToInputStream("in0_unfiltered",
MakePacket<int>(3).At(timestamp)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(1, packet_dump.size());
// We send an even integer with timestamp 3 to the IntAdderCalculator. This
// packet will go through and the IntAdderCalculator will run.
MP_EXPECT_OK(graph.AddPacketToInputStream("in1_to_be_filtered",
MakePacket<int>(4).At(timestamp)));
MP_ASSERT_OK(graph.WaitUntilIdle());
ASSERT_EQ(2, packet_dump.size());
EXPECT_EQ(7, packet_dump[1].Get<int>());
MP_ASSERT_OK(graph.CloseAllInputStreams());
MP_ASSERT_OK(graph.WaitUntilDone());
EXPECT_EQ(2, packet_dump.size());
}
TEST(CalculatorGraph, PropagateBoundLoop) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: 'OutputAndBoundSourceCalculator'
output_stream: 'integers'
}
node {
calculator: 'IntAdderCalculator'
input_stream: 'integers'
input_stream: 'old_sum'
input_stream_info: {
tag_index: ':1' # 'old_sum'
back_edge: true
}
output_stream: 'sum'
input_stream_handler {
input_stream_handler: 'EarlyCloseInputStreamHandler'
}
}
node {
calculator: 'Delay20Calculator'
input_stream: 'sum'
output_stream: 'old_sum'
}
)");
std::vector<Packet> packet_dump;
tool::AddVectorSink("sum", &config, &packet_dump);
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.Run());
ASSERT_EQ(101, packet_dump.size());
int sum = 0;
for (int i = 0; i < 101; ++i) {
sum += i;
EXPECT_EQ(sum, packet_dump[i].Get<int>());
EXPECT_EQ(Timestamp(i * 20), packet_dump[i].Timestamp());
}
}
TEST(CalculatorGraph, CheckBatchProcessingBoundPropagation) {
// The timestamp bound sent by OutputAndBoundSourceCalculator shouldn't be
// directly propagated to the output stream when PassThroughCalculator has
// anything in its default calculator context for batch processing. Otherwise,
// the sink calculator's input stream should report packet timestamp
// mismatches.
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: 'OutputAndBoundSourceCalculator'
output_stream: 'integers'
}
node {
calculator: 'PassThroughCalculator'
input_stream: 'integers'
output_stream: 'output'
input_stream_handler {
input_stream_handler: "DefaultInputStreamHandler"
options: {
[mediapipe.DefaultInputStreamHandlerOptions.ext]: {
batch_size: 10
}
}
}
}
node { calculator: 'IntSinkCalculator' input_stream: 'output' }
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.Run());
}
// Shows that ImmediateInputStreamHandler allows bounds propagation.
TEST(CalculatorGraphBounds, ImmediateHandlerBounds) {
TEST(CalculatorGraphBoundsTest, ImmediateHandlerBounds) {
// CustomBoundCalculator produces only timestamp bounds.
// The first PassThroughCalculator propagates bounds using SetOffset(0).
// The second PassthroughCalculator delivers an output packet whenever the
@@ -101,5 +620,261 @@ TEST(CalculatorGraphBounds, ImmediateHandlerBounds) {
EXPECT_EQ(output_packets.size(), 4);
}
// A Calculator that only sets timestamp bound by SetOffset().
class OffsetBoundCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Outputs().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->SetOffset(0);
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(OffsetBoundCalculator);
// A Calculator that produces a packet for each call to Process.
class BoundToPacketCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).Set<int>();
cc->Outputs().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
cc->Outputs().Index(0).AddPacket(Adopt(new int(33)));
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(BoundToPacketCalculator);
// Verifies that SetOffset still propagates when Process is called and
// produces no output packets.
TEST(CalculatorGraphBoundsTest, OffsetBoundPropagation) {
// OffsetBoundCalculator produces only timestamp bounds.
// The PassthroughCalculator delivers an output packet whenever the
// OffsetBoundCalculator delivers a timestamp bound.
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'input'
node {
calculator: 'OffsetBoundCalculator'
input_stream: 'input'
output_stream: 'bounds'
}
node {
calculator: 'PassThroughCalculator'
input_stream: 'bounds'
input_stream: 'input'
output_stream: 'bounds_output'
output_stream: 'output'
}
)");
CalculatorGraph graph;
std::vector<Packet> output_packets;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.ObserveOutputStream("output", [&](const Packet& p) {
output_packets.push_back(p);
return ::mediapipe::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.WaitUntilIdle());
// Add four packets into the graph.
constexpr int kNumInputs = 4;
for (int i = 0; i < kNumInputs; ++i) {
Packet p = MakePacket<int>(33).At(Timestamp(i));
MP_ASSERT_OK(graph.AddPacketToInputStream("input", p));
}
// Four packets arrive at the output only if timestamp bounds are propagated.
MP_ASSERT_OK(graph.WaitUntilIdle());
EXPECT_EQ(output_packets.size(), kNumInputs);
// Shutdown the graph.
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows that bounds changes alone do not invoke Process.
// Note: Bounds changes alone will invoke Process eventually
// when SetOffset is cleared, see: go/mediapipe-realtime-graph.
TEST(CalculatorGraphBoundsTest, BoundWithoutInputPackets) {
// OffsetBoundCalculator produces only timestamp bounds.
// The BoundToPacketCalculator delivers an output packet whenever the
// OffsetBoundCalculator delivers a timestamp bound.
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'input'
node {
calculator: 'OffsetBoundCalculator'
input_stream: 'input'
output_stream: 'bounds'
}
node {
calculator: 'BoundToPacketCalculator'
input_stream: 'bounds'
output_stream: 'output'
}
)");
CalculatorGraph graph;
std::vector<Packet> output_packets;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.ObserveOutputStream("output", [&](const Packet& p) {
output_packets.push_back(p);
return ::mediapipe::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.WaitUntilIdle());
// Add four packets into the graph.
constexpr int kNumInputs = 4;
for (int i = 0; i < kNumInputs; ++i) {
Packet p = MakePacket<int>(33).At(Timestamp(i));
MP_ASSERT_OK(graph.AddPacketToInputStream("input", p));
}
// No packets arrive, because updated timestamp bounds do not invoke
// BoundToPacketCalculator::Process.
MP_ASSERT_OK(graph.WaitUntilIdle());
EXPECT_EQ(output_packets.size(), 0);
// Shutdown the graph.
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
// Shows that when fixed-size-input-stream-hanlder drops packets,
// no timetamp bounds are announced.
TEST(CalculatorGraphBoundsTest, FixedSizeHandlerBounds) {
// LambdaCalculator with FixedSizeInputStreamHandler will drop packets
// while it is busy. Timetamps for the dropped packets are only relevant
// when SetOffset is active on the LambdaCalculator.
// The PassthroughCalculator delivers an output packet whenever the
// LambdaCalculator delivers a timestamp bound.
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: 'input'
input_side_packet: 'open_function'
input_side_packet: 'process_function'
node {
calculator: 'LambdaCalculator'
input_stream: 'input'
output_stream: 'thinned'
input_side_packet: 'OPEN:open_fn'
input_side_packet: 'PROCESS:process_fn'
input_stream_handler {
input_stream_handler: "FixedSizeInputStreamHandler"
}
}
node {
calculator: 'PassThroughCalculator'
input_stream: 'thinned'
input_stream: 'input'
output_stream: 'thinned_output'
output_stream: 'output'
}
)");
CalculatorGraph graph;
// The task_semaphore counts the number of running tasks.
constexpr int kTaskSupply = 10;
AtomicSemaphore task_semaphore(/*supply=*/kTaskSupply);
// This executor invokes a callback at the start and finish of each task.
auto executor = std::make_shared<CountingExecutor>(
4, /*start_callback=*/[&]() { task_semaphore.Acquire(1); },
/*finish_callback=*/[&]() { task_semaphore.Release(1); });
MP_ASSERT_OK(graph.SetExecutor(/*name=*/"", executor));
// Monitor output from the graph.
MP_ASSERT_OK(graph.Initialize(config));
std::vector<Packet> outputs;
MP_ASSERT_OK(graph.ObserveOutputStream("output", [&](const Packet& p) {
outputs.push_back(p);
return ::mediapipe::OkStatus();
}));
std::vector<Packet> thinned_outputs;
MP_ASSERT_OK(
graph.ObserveOutputStream("thinned_output", [&](const Packet& p) {
thinned_outputs.push_back(p);
return ::mediapipe::OkStatus();
}));
// The enter_semaphore is used to wait for LambdaCalculator::Process.
// The exit_semaphore blocks and unblocks LambdaCalculator::Process.
AtomicSemaphore enter_semaphore(0);
AtomicSemaphore exit_semaphore(0);
CalculatorContextFunction open_fn = [&](CalculatorContext* cc) {
cc->SetOffset(0);
return ::mediapipe::OkStatus();
};
CalculatorContextFunction process_fn = [&](CalculatorContext* cc) {
enter_semaphore.Release(1);
exit_semaphore.Acquire(1);
cc->Outputs().Index(0).AddPacket(cc->Inputs().Index(0).Value());
return ::mediapipe::OkStatus();
};
MP_ASSERT_OK(graph.StartRun({
{"open_fn", Adopt(new auto(open_fn))},
{"process_fn", Adopt(new auto(process_fn))},
}));
MP_ASSERT_OK(graph.WaitUntilIdle());
// Add four packets into the graph.
constexpr int kNumInputs = 4;
for (int i = 0; i < kNumInputs; ++i) {
Packet p = MakePacket<int>(33).At(Timestamp(i));
MP_ASSERT_OK(graph.AddPacketToInputStream("input", p));
}
// Wait until only the LambdaCalculator is running,
// by wating until the task_semaphore has only one token occupied.
// At this point 2 packets were dropped by the FixedSizeInputStreamHandler.
task_semaphore.Acquire(kTaskSupply - 1);
task_semaphore.Release(kTaskSupply - 1);
// No timestamp bounds and no packets are emitted yet.
EXPECT_EQ(outputs.size(), 0);
EXPECT_EQ(thinned_outputs.size(), 0);
// Allow the first LambdaCalculator::Process call to complete.
// Wait for the second LambdaCalculator::Process call to begin.
// Wait until only the LambdaCalculator is running.
enter_semaphore.Acquire(1);
exit_semaphore.Release(1);
enter_semaphore.Acquire(1);
task_semaphore.Acquire(kTaskSupply - 1);
task_semaphore.Release(kTaskSupply - 1);
// Only one timestamp bound and one packet are emitted.
EXPECT_EQ(outputs.size(), 1);
EXPECT_EQ(thinned_outputs.size(), 1);
// Allow the second LambdaCalculator::Process call to complete.
exit_semaphore.Release(1);
MP_ASSERT_OK(graph.WaitUntilIdle());
// Packets 1 and 2 were dropped by the FixedSizeInputStreamHandler.
EXPECT_EQ(thinned_outputs.size(), 2);
EXPECT_EQ(thinned_outputs[0].Timestamp(), Timestamp(0));
EXPECT_EQ(thinned_outputs[1].Timestamp(), Timestamp(kNumInputs - 1));
EXPECT_EQ(outputs.size(), kNumInputs);
MP_ASSERT_OK(graph.CloseAllPacketSources());
MP_ASSERT_OK(graph.WaitUntilDone());
}
} // namespace
} // namespace mediapipe
mediapipe/framework/calculator_graph_side_packet_test.cc
+747
View File
@@ -0,0 +1,747 @@
// 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.
#include <map>
#include <memory>
#include <string>
#include <vector>
#include "absl/time/clock.h"
#include "absl/time/time.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/port/canonical_errors.h"
#include "mediapipe/framework/port/gmock.h"
#include "mediapipe/framework/port/gtest.h"
#include "mediapipe/framework/port/logging.h"
#include "mediapipe/framework/port/parse_text_proto.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_matchers.h"
namespace mediapipe {
namespace {
// Takes an input stream packet and passes it (with timestamp removed) as an
// output side packet.
class OutputSidePacketInProcessCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->OutputSidePackets().Index(0).SetSameAs(&cc->Inputs().Index(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
cc->OutputSidePackets().Index(0).Set(
cc->Inputs().Index(0).Value().At(Timestamp::Unset()));
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(OutputSidePacketInProcessCalculator);
// Takes an input stream packet and counts the number of the packets it
// receives. Outputs the total number of packets as a side packet in Close.
class CountAndOutputSummarySidePacketInCloseCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->OutputSidePackets().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
++count_;
return ::mediapipe::OkStatus();
}
::mediapipe::Status Close(CalculatorContext* cc) final {
cc->OutputSidePackets().Index(0).Set(
MakePacket<int>(count_).At(Timestamp::Unset()));
return ::mediapipe::OkStatus();
}
int count_ = 0;
};
REGISTER_CALCULATOR(CountAndOutputSummarySidePacketInCloseCalculator);
// Takes an input stream packet and passes it (with timestamp intact) as an
// output side packet. This triggers an error in the graph.
class OutputSidePacketWithTimestampCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->OutputSidePackets().Index(0).SetSameAs(&cc->Inputs().Index(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
cc->OutputSidePackets().Index(0).Set(cc->Inputs().Index(0).Value());
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(OutputSidePacketWithTimestampCalculator);
// Generates an output side packet containing the integer 1.
class IntegerOutputSidePacketCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->OutputSidePackets().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->OutputSidePackets().Index(0).Set(MakePacket<int>(1));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
LOG(FATAL) << "Not reached.";
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(IntegerOutputSidePacketCalculator);
// Generates an output side packet containing the sum of the two integer input
// side packets.
class SidePacketAdderCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->InputSidePackets().Index(0).Set<int>();
cc->InputSidePackets().Index(1).Set<int>();
cc->OutputSidePackets().Index(0).Set<int>();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->OutputSidePackets().Index(0).Set(
MakePacket<int>(cc->InputSidePackets().Index(1).Get<int>() +
cc->InputSidePackets().Index(0).Get<int>()));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
LOG(FATAL) << "Not reached.";
return ::mediapipe::OkStatus();
}
};
REGISTER_CALCULATOR(SidePacketAdderCalculator);
// Produces an output packet with the PostStream timestamp containing the
// input side packet.
class SidePacketToStreamPacketCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->InputSidePackets().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->InputSidePackets().Index(0));
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
cc->Outputs().Index(0).AddPacket(
cc->InputSidePackets().Index(0).At(Timestamp::PostStream()));
cc->Outputs().Index(0).Close();
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
return ::mediapipe::tool::StatusStop();
}
};
REGISTER_CALCULATOR(SidePacketToStreamPacketCalculator);
// Packet generator for an arbitrary unit64 packet.
class Uint64PacketGenerator : public PacketGenerator {
public:
static ::mediapipe::Status FillExpectations(
const PacketGeneratorOptions& extendable_options,
PacketTypeSet* input_side_packets, PacketTypeSet* output_side_packets) {
output_side_packets->Index(0).Set<uint64>();
return ::mediapipe::OkStatus();
}
static ::mediapipe::Status Generate(
const PacketGeneratorOptions& extendable_options,
const PacketSet& input_side_packets, PacketSet* output_side_packets) {
output_side_packets->Index(0) = Adopt(new uint64(15LL << 32 | 5));
return ::mediapipe::OkStatus();
}
};
REGISTER_PACKET_GENERATOR(Uint64PacketGenerator);
TEST(CalculatorGraph, OutputSidePacketInProcess) {
const int64 offset = 100;
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "offset"
node {
calculator: "OutputSidePacketInProcessCalculator"
input_stream: "offset"
output_side_packet: "offset"
}
node {
calculator: "SidePacketToStreamPacketCalculator"
output_stream: "output"
input_side_packet: "offset"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
std::vector<Packet> output_packets;
MP_ASSERT_OK(graph.ObserveOutputStream(
"output", [&output_packets](const Packet& packet) {
output_packets.push_back(packet);
return ::mediapipe::OkStatus();
}));
// Run the graph twice.
for (int run = 0; run < 2; ++run) {
output_packets.clear();
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"offset", MakePacket<TimestampDiff>(offset).At(Timestamp(0))));
MP_ASSERT_OK(graph.CloseInputStream("offset"));
MP_ASSERT_OK(graph.WaitUntilDone());
ASSERT_EQ(1, output_packets.size());
EXPECT_EQ(offset, output_packets[0].Get<TimestampDiff>().Value());
}
}
// A PacketGenerator that simply passes its input Packets through
// unchanged. The inputs may be specified by tag or index. The outputs
// must match the inputs exactly. Any options may be specified and will
// also be ignored.
class PassThroughGenerator : public PacketGenerator {
public:
static ::mediapipe::Status FillExpectations(
const PacketGeneratorOptions& extendable_options, PacketTypeSet* inputs,
PacketTypeSet* outputs) {
if (!inputs->TagMap()->SameAs(*outputs->TagMap())) {
return ::mediapipe::InvalidArgumentError(
"Input and outputs to PassThroughGenerator must use the same tags "
"and indexes.");
}
for (CollectionItemId id = inputs->BeginId(); id < inputs->EndId(); ++id) {
inputs->Get(id).SetAny();
outputs->Get(id).SetSameAs(&inputs->Get(id));
}
return ::mediapipe::OkStatus();
}
static ::mediapipe::Status Generate(
const PacketGeneratorOptions& extendable_options,
const PacketSet& input_side_packets, PacketSet* output_side_packets) {
for (CollectionItemId id = input_side_packets.BeginId();
id < input_side_packets.EndId(); ++id) {
output_side_packets->Get(id) = input_side_packets.Get(id);
}
return ::mediapipe::OkStatus();
}
};
REGISTER_PACKET_GENERATOR(PassThroughGenerator);
TEST(CalculatorGraph, SharePacketGeneratorGraph) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: 'CountingSourceCalculator'
output_stream: 'count1'
input_side_packet: 'MAX_COUNT:max_count1'
}
node {
calculator: 'CountingSourceCalculator'
output_stream: 'count2'
input_side_packet: 'MAX_COUNT:max_count2'
}
node {
calculator: 'CountingSourceCalculator'
output_stream: 'count3'
input_side_packet: 'MAX_COUNT:max_count3'
}
node {
calculator: 'CountingSourceCalculator'
output_stream: 'count4'
input_side_packet: 'MAX_COUNT:max_count4'
}
node {
calculator: 'PassThroughCalculator'
input_side_packet: 'MAX_COUNT:max_count5'
output_side_packet: 'MAX_COUNT:max_count6'
}
node {
calculator: 'CountingSourceCalculator'
output_stream: 'count5'
input_side_packet: 'MAX_COUNT:max_count6'
}
packet_generator {
packet_generator: 'PassThroughGenerator'
input_side_packet: 'max_count1'
output_side_packet: 'max_count2'
}
packet_generator {
packet_generator: 'PassThroughGenerator'
input_side_packet: 'max_count4'
output_side_packet: 'max_count5'
}
)");
// At this point config is a standard config which specifies both
// calculators and packet_factories/packet_genators. The following
// code is an example of reusing side packets across a number of
// CalculatorGraphs. It is particularly informative to note how each
// side packet is created.
//
// max_count1 is set for all graphs by a PacketFactory in the config.
// The side packet is created by generator_graph.InitializeGraph().
//
// max_count2 is set for all graphs by a PacketGenerator in the config.
// The side packet is created by generator_graph.InitializeGraph()
// because max_count1 is available at that time.
//
// max_count3 is set for all graphs by directly being specified as an
// argument to generator_graph.InitializeGraph().
//
// max_count4 is set per graph because it is directly specified as an
// argument to generator_graph.ProcessGraph().
//
// max_count5 is set per graph by a PacketGenerator which is run when
// generator_graph.ProcessGraph() is run (because max_count4 isn't
// available until then).
// Before anything else, split the graph config into two parts, one
// with the PacketFactory and PacketGenerator config and the other
// with the Calculator config.
CalculatorGraphConfig calculator_config = config;
calculator_config.clear_packet_factory();
calculator_config.clear_packet_generator();
CalculatorGraphConfig generator_config = config;
generator_config.clear_node();
// Next, create a ValidatedGraphConfig for both configs.
ValidatedGraphConfig validated_calculator_config;
MP_ASSERT_OK(validated_calculator_config.Initialize(calculator_config));
ValidatedGraphConfig validated_generator_config;
MP_ASSERT_OK(validated_generator_config.Initialize(generator_config));
// Create a PacketGeneratorGraph. Side packets max_count1, max_count2,
// and max_count3 are created upon initialization.
// Note that validated_generator_config must outlive generator_graph.
PacketGeneratorGraph generator_graph;
MP_ASSERT_OK(
generator_graph.Initialize(&validated_generator_config, nullptr,
{{"max_count1", MakePacket<int>(10)},
{"max_count3", MakePacket<int>(20)}}));
ASSERT_THAT(generator_graph.BasePackets(),
testing::ElementsAre(testing::Key("max_count1"),
testing::Key("max_count2"),
testing::Key("max_count3")));
// Create a bunch of graphs.
std::vector<std::unique_ptr<CalculatorGraph>> graphs;
for (int i = 0; i < 100; ++i) {
graphs.emplace_back(absl::make_unique<CalculatorGraph>());
// Do not pass extra side packets here.
// Note that validated_calculator_config must outlive the graph.
MP_ASSERT_OK(graphs.back()->Initialize(calculator_config, {}));
}
// Run a bunch of graphs, reusing side packets max_count1, max_count2,
// and max_count3. The side packet max_count4 is added per run,
// and triggers the execution of a packet generator which generates
// max_count5.
for (int i = 0; i < 100; ++i) {
std::map<std::string, Packet> all_side_packets;
// Creates max_count4 and max_count5.
MP_ASSERT_OK(generator_graph.RunGraphSetup(
{{"max_count4", MakePacket<int>(30 + i)}}, &all_side_packets));
ASSERT_THAT(all_side_packets,
testing::ElementsAre(
testing::Key("max_count1"), testing::Key("max_count2"),
testing::Key("max_count3"), testing::Key("max_count4"),
testing::Key("max_count5")));
// Pass all the side packets prepared by generator_graph here.
MP_ASSERT_OK(graphs[i]->Run(all_side_packets));
// TODO Verify the actual output.
}
// Destroy all the graphs.
graphs.clear();
}
TEST(CalculatorGraph, OutputSidePacketAlreadySet) {
const int64 offset = 100;
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "offset"
node {
calculator: "OutputSidePacketInProcessCalculator"
input_stream: "offset"
output_side_packet: "offset"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
// Send two input packets to cause OutputSidePacketInProcessCalculator to
// set the output side packet twice.
MP_ASSERT_OK(graph.AddPacketToInputStream(
"offset", MakePacket<TimestampDiff>(offset).At(Timestamp(0))));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"offset", MakePacket<TimestampDiff>(offset).At(Timestamp(1))));
MP_ASSERT_OK(graph.CloseInputStream("offset"));
::mediapipe::Status status = graph.WaitUntilDone();
EXPECT_EQ(status.code(), ::mediapipe::StatusCode::kAlreadyExists);
EXPECT_THAT(status.message(), testing::HasSubstr("was already set."));
}
TEST(CalculatorGraph, OutputSidePacketWithTimestamp) {
const int64 offset = 100;
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "offset"
node {
calculator: "OutputSidePacketWithTimestampCalculator"
input_stream: "offset"
output_side_packet: "offset"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
// The OutputSidePacketWithTimestampCalculator neglects to clear the
// timestamp in the input packet when it copies the input packet to the
// output side packet. The timestamp value should appear in the error
// message.
MP_ASSERT_OK(graph.AddPacketToInputStream(
"offset", MakePacket<TimestampDiff>(offset).At(Timestamp(237))));
MP_ASSERT_OK(graph.CloseInputStream("offset"));
::mediapipe::Status status = graph.WaitUntilDone();
EXPECT_EQ(status.code(), ::mediapipe::StatusCode::kInvalidArgument);
EXPECT_THAT(status.message(), testing::HasSubstr("has a timestamp 237."));
}
TEST(CalculatorGraph, OutputSidePacketConsumedBySourceNode) {
const int max_count = 10;
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "max_count"
node {
calculator: "OutputSidePacketInProcessCalculator"
input_stream: "max_count"
output_side_packet: "max_count"
}
node {
calculator: "CountingSourceCalculator"
output_stream: "count"
input_side_packet: "MAX_COUNT:max_count"
}
node {
calculator: "PassThroughCalculator"
input_stream: "count"
output_stream: "output"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
std::vector<Packet> output_packets;
MP_ASSERT_OK(graph.ObserveOutputStream(
"output", [&output_packets](const Packet& packet) {
output_packets.push_back(packet);
return ::mediapipe::OkStatus();
}));
MP_ASSERT_OK(graph.StartRun({}));
// Wait until the graph is idle so that
// Scheduler::TryToScheduleNextSourceLayer() gets called.
// Scheduler::TryToScheduleNextSourceLayer() should not activate source
// nodes that haven't been opened. We can't call graph.WaitUntilIdle()
// because the graph has a source node.
absl::SleepFor(absl::Milliseconds(10));
MP_ASSERT_OK(graph.AddPacketToInputStream(
"max_count", MakePacket<int>(max_count).At(Timestamp(0))));
MP_ASSERT_OK(graph.CloseInputStream("max_count"));
MP_ASSERT_OK(graph.WaitUntilDone());
ASSERT_EQ(max_count, output_packets.size());
for (int i = 0; i < output_packets.size(); ++i) {
EXPECT_EQ(i, output_packets[i].Get<int>());
EXPECT_EQ(Timestamp(i), output_packets[i].Timestamp());
}
}
// Returns the first packet of the input stream.
class FirstPacketFilterCalculator : public CalculatorBase {
public:
FirstPacketFilterCalculator() {}
~FirstPacketFilterCalculator() override {}
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
cc->Inputs().Index(0).SetAny();
cc->Outputs().Index(0).SetSameAs(&cc->Inputs().Index(0));
return ::mediapipe::OkStatus();
}
::mediapipe::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 ::mediapipe::OkStatus();
}
private:
bool seen_first_packet_ = false;
};
REGISTER_CALCULATOR(FirstPacketFilterCalculator);
TEST(CalculatorGraph, SourceLayerInversion) {
// There are three CountingSourceCalculators, indexed 0, 1, and 2. Each of
// them outputs 10 packets.
//
// CountingSourceCalculator 0 should output 0, 1, 2, 3, ..., 9.
// CountingSourceCalculator 1 should output 100, 101, 102, 103, ..., 109.
// CountingSourceCalculator 2 should output 0, 100, 200, 300, ..., 900.
// However, there is a source layer inversion.
// CountingSourceCalculator 0 is in source layer 0.
// CountingSourceCalculator 1 is in source layer 1.
// CountingSourceCalculator 2 is in source layer 0, but consumes an output
// side packet generated by a downstream calculator of
// CountingSourceCalculator 1.
//
// This graph will deadlock when CountingSourceCalculator 0 runs to
// completion and CountingSourceCalculator 1 cannot be activated because
// CountingSourceCalculator 2 cannot be opened.
const int max_count = 10;
const int initial_value1 = 100;
// Set num_threads to 1 to force sequential execution for deterministic
// outputs.
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
num_threads: 1
node {
calculator: "CountingSourceCalculator"
output_stream: "count0"
input_side_packet: "MAX_COUNT:max_count"
source_layer: 0
}
node {
calculator: "CountingSourceCalculator"
output_stream: "count1"
input_side_packet: "MAX_COUNT:max_count"
input_side_packet: "INITIAL_VALUE:initial_value1"
source_layer: 1
}
node {
calculator: "FirstPacketFilterCalculator"
input_stream: "count1"
output_stream: "first_count1"
}
node {
calculator: "OutputSidePacketInProcessCalculator"
input_stream: "first_count1"
output_side_packet: "increment2"
}
node {
calculator: "CountingSourceCalculator"
output_stream: "count2"
input_side_packet: "MAX_COUNT:max_count"
input_side_packet: "INCREMENT:increment2"
source_layer: 0
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(
config, {{"max_count", MakePacket<int>(max_count)},
{"initial_value1", MakePacket<int>(initial_value1)}}));
::mediapipe::Status status = graph.Run();
EXPECT_EQ(status.code(), ::mediapipe::StatusCode::kUnknown);
EXPECT_THAT(status.message(), testing::HasSubstr("deadlock"));
}
// Tests a graph of packet-generator-like calculators, which have no input
// streams and no output streams.
TEST(CalculatorGraph, PacketGeneratorLikeCalculators) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
node {
calculator: "IntegerOutputSidePacketCalculator"
output_side_packet: "one"
}
node {
calculator: "IntegerOutputSidePacketCalculator"
output_side_packet: "another_one"
}
node {
calculator: "SidePacketAdderCalculator"
input_side_packet: "one"
input_side_packet: "another_one"
output_side_packet: "two"
}
node {
calculator: "IntegerOutputSidePacketCalculator"
output_side_packet: "yet_another_one"
}
node {
calculator: "SidePacketAdderCalculator"
input_side_packet: "two"
input_side_packet: "yet_another_one"
output_side_packet: "three"
}
node {
calculator: "SidePacketToStreamPacketCalculator"
input_side_packet: "three"
output_stream: "output"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
std::vector<Packet> output_packets;
MP_ASSERT_OK(graph.ObserveOutputStream(
"output", [&output_packets](const Packet& packet) {
output_packets.push_back(packet);
return ::mediapipe::OkStatus();
}));
MP_ASSERT_OK(graph.Run());
ASSERT_EQ(1, output_packets.size());
EXPECT_EQ(3, output_packets[0].Get<int>());
EXPECT_EQ(Timestamp::PostStream(), output_packets[0].Timestamp());
}
TEST(CalculatorGraph, OutputSummarySidePacketInClose) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_packets"
node {
calculator: "CountAndOutputSummarySidePacketInCloseCalculator"
input_stream: "input_packets"
output_side_packet: "num_of_packets"
}
node {
calculator: "SidePacketToStreamPacketCalculator"
input_side_packet: "num_of_packets"
output_stream: "output"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
std::vector<Packet> output_packets;
MP_ASSERT_OK(graph.ObserveOutputStream(
"output", [&output_packets](const Packet& packet) {
output_packets.push_back(packet);
return ::mediapipe::OkStatus();
}));
// Run the graph twice.
int max_count = 100;
for (int run = 0; run < 1; ++run) {
output_packets.clear();
MP_ASSERT_OK(graph.StartRun({}));
for (int i = 0; i < max_count; ++i) {
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_packets", MakePacket<int>(i).At(Timestamp(i))));
}
MP_ASSERT_OK(graph.CloseInputStream("input_packets"));
MP_ASSERT_OK(graph.WaitUntilDone());
ASSERT_EQ(1, output_packets.size());
EXPECT_EQ(max_count, output_packets[0].Get<int>());
EXPECT_EQ(Timestamp::PostStream(), output_packets[0].Timestamp());
}
}
TEST(CalculatorGraph, GetOutputSidePacket) {
CalculatorGraphConfig config =
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
input_stream: "input_packets"
node {
calculator: "CountAndOutputSummarySidePacketInCloseCalculator"
input_stream: "input_packets"
output_side_packet: "num_of_packets"
}
packet_generator {
packet_generator: "Uint64PacketGenerator"
output_side_packet: "output_uint64"
}
packet_generator {
packet_generator: "IntSplitterPacketGenerator"
input_side_packet: "input_uint64"
output_side_packet: "output_uint32_pair"
}
)");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
// Check a packet generated by the PacketGenerator, which is available after
// graph initialization, can be fetched before graph starts.
::mediapipe::StatusOr<Packet> status_or_packet =
graph.GetOutputSidePacket("output_uint64");
MP_ASSERT_OK(status_or_packet);
EXPECT_EQ(Timestamp::Unset(), status_or_packet.ValueOrDie().Timestamp());
// IntSplitterPacketGenerator is missing its input side packet and we
// won't be able to get its output side packet now.
status_or_packet = graph.GetOutputSidePacket("output_uint32_pair");
EXPECT_EQ(::mediapipe::StatusCode::kUnavailable,
status_or_packet.status().code());
// Run the graph twice.
int max_count = 100;
std::map<std::string, Packet> extra_side_packets;
extra_side_packets.insert({"input_uint64", MakePacket<uint64>(1123)});
for (int run = 0; run < 1; ++run) {
MP_ASSERT_OK(graph.StartRun(extra_side_packets));
status_or_packet = graph.GetOutputSidePacket("output_uint32_pair");
MP_ASSERT_OK(status_or_packet);
EXPECT_EQ(Timestamp::Unset(), status_or_packet.ValueOrDie().Timestamp());
for (int i = 0; i < max_count; ++i) {
MP_ASSERT_OK(graph.AddPacketToInputStream(
"input_packets", MakePacket<int>(i).At(Timestamp(i))));
}
MP_ASSERT_OK(graph.CloseInputStream("input_packets"));
// Should return NOT_FOUND for invalid side packets.
status_or_packet = graph.GetOutputSidePacket("unknown");
EXPECT_FALSE(status_or_packet.ok());
EXPECT_EQ(::mediapipe::StatusCode::kNotFound,
status_or_packet.status().code());
// Should return UNAVAILABLE before graph is done for valid non-base
// packets.
status_or_packet = graph.GetOutputSidePacket("num_of_packets");
EXPECT_FALSE(status_or_packet.ok());
EXPECT_EQ(::mediapipe::StatusCode::kUnavailable,
status_or_packet.status().code());
// Should stil return a base even before graph is done.
status_or_packet = graph.GetOutputSidePacket("output_uint64");
MP_ASSERT_OK(status_or_packet);
EXPECT_EQ(Timestamp::Unset(), status_or_packet.ValueOrDie().Timestamp());
MP_ASSERT_OK(graph.WaitUntilDone());
// Check packets are available after graph is done.
status_or_packet = graph.GetOutputSidePacket("num_of_packets");
MP_ASSERT_OK(status_or_packet);
EXPECT_EQ(max_count, status_or_packet.ValueOrDie().Get<int>());
EXPECT_EQ(Timestamp::Unset(), status_or_packet.ValueOrDie().Timestamp());
// Should still return a base packet after graph is done.
status_or_packet = graph.GetOutputSidePacket("output_uint64");
MP_ASSERT_OK(status_or_packet);
EXPECT_EQ(Timestamp::Unset(), status_or_packet.ValueOrDie().Timestamp());
// Should still return a non-base packet after graph is done.
status_or_packet = graph.GetOutputSidePacket("output_uint32_pair");
MP_ASSERT_OK(status_or_packet);
EXPECT_EQ(Timestamp::Unset(), status_or_packet.ValueOrDie().Timestamp());
}
}
} // namespace
} // namespace mediapipe
mediapipe/framework/calculator_graph_test.cc
+1 -1046
View File
File diff suppressed because it is too large Load Diff
mediapipe/framework/formats/matrix_data.proto
+3 -3
View File
@@ -34,7 +34,7 @@ message MatrixData {
ROW_MAJOR = 1;
}
// Order in which the data are stored. Implicitly defaults to COLUMN_MAJOR,
// which matches the default for mediapipe::Matrix and Eigen::Matrix*.
optional Layout layout = 4;
// Order in which the data are stored. Defaults to COLUMN_MAJOR, which matches
// the default for mediapipe::Matrix and Eigen::Matrix*.
optional Layout layout = 4 [default = COLUMN_MAJOR];
}
mediapipe/framework/graph_validation_test.cc
+7 -5
View File
@@ -154,11 +154,13 @@ TEST(ValidatedGraphConfigTest, InitializeTemplateFromProtos) {
}
)");
auto options = ParseTextProtoOrDie<Subgraph::SubgraphOptions>(R"(
[mediapipe.TemplateSubgraphOptions.ext]: {
dict: {
arg: {
key: "in_name"
value: { str: "stream_9" }
options: {
[mediapipe.TemplateSubgraphOptions.ext]: {
dict: {
arg: {
key: "in_name"
value: { str: "stream_9" }
}
}
}
})");
mediapipe/framework/subgraph.cc
+2 -2
View File
@@ -44,8 +44,8 @@ TemplateSubgraph::~TemplateSubgraph() {}
::mediapipe::StatusOr<CalculatorGraphConfig> TemplateSubgraph::GetConfig(
const Subgraph::SubgraphOptions& options) {
const TemplateDict& arguments =
options.GetExtension(TemplateSubgraphOptions::ext).dict();
TemplateDict arguments =
Subgraph::GetOptions<mediapipe::TemplateSubgraphOptions>(options).dict();
tool::TemplateExpander expander;
CalculatorGraphConfig config;
MP_RETURN_IF_ERROR(expander.ExpandTemplates(arguments, templ_, &config));
mediapipe/framework/subgraph.h
+8 -1
View File
@@ -24,6 +24,7 @@
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/statusor.h"
#include "mediapipe/framework/tool/calculator_graph_template.pb.h"
#include "mediapipe/framework/tool/options_util.h"
namespace mediapipe {
@@ -32,7 +33,7 @@ namespace mediapipe {
// the graph is running.
class Subgraph {
public:
using SubgraphOptions = CalculatorOptions;
using SubgraphOptions = CalculatorGraphConfig::Node;
Subgraph();
virtual ~Subgraph();
// Returns the config to use for one instantiation of the subgraph. The
@@ -42,6 +43,12 @@ class Subgraph {
// TODO: make this static?
virtual ::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) = 0;
// Returns options of a specific type.
template <typename T>
static T GetOptions(Subgraph::SubgraphOptions supgraph_options) {
return tool::OptionsMap().Initialize(supgraph_options).Get<T>();
}
};
using SubgraphRegistry = GlobalFactoryRegistry<std::unique_ptr<Subgraph>>;
mediapipe/framework/test_calculators.cc
+37 -5
View File
@@ -548,8 +548,12 @@ typedef std::function<::mediapipe::Status(const InputStreamShardSet&,
OutputStreamShardSet*)>
ProcessFunction;
// A callback function for Calculator::Open, Process, or Close.
typedef std::function<::mediapipe::Status(CalculatorContext* cc)>
CalculatorContextFunction;
// A Calculator that runs a testing callback function in Process,
// which is specified as an input side packet.
// Open, or Close, which is specified as an input side packet.
class LambdaCalculator : public CalculatorBase {
public:
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
@@ -561,21 +565,49 @@ class LambdaCalculator : public CalculatorBase {
id < cc->Outputs().EndId(); ++id) {
cc->Outputs().Get(id).SetAny();
}
cc->InputSidePackets().Index(0).Set<ProcessFunction>();
if (cc->InputSidePackets().HasTag("") > 0) {
cc->InputSidePackets().Tag("").Set<ProcessFunction>();
}
for (std::string tag : {"OPEN", "PROCESS", "CLOSE"}) {
if (cc->InputSidePackets().HasTag(tag)) {
cc->InputSidePackets().Tag(tag).Set<CalculatorContextFunction>();
}
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status Open(CalculatorContext* cc) final {
callback_ = cc->InputSidePackets().Index(0).Get<ProcessFunction>();
if (cc->InputSidePackets().HasTag("OPEN")) {
return GetContextFn(cc, "OPEN")(cc);
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status Process(CalculatorContext* cc) final {
return callback_(cc->Inputs(), &(cc->Outputs()));
if (cc->InputSidePackets().HasTag("PROCESS")) {
return GetContextFn(cc, "PROCESS")(cc);
}
if (cc->InputSidePackets().HasTag("") > 0) {
return GetProcessFn(cc, "")(cc->Inputs(), &cc->Outputs());
}
return ::mediapipe::OkStatus();
}
::mediapipe::Status Close(CalculatorContext* cc) final {
if (cc->InputSidePackets().HasTag("CLOSE")) {
return GetContextFn(cc, "CLOSE")(cc);
}
return ::mediapipe::OkStatus();
}
private:
ProcessFunction callback_;
ProcessFunction GetProcessFn(CalculatorContext* cc, std::string tag) {
return cc->InputSidePackets().Tag(tag).Get<ProcessFunction>();
}
CalculatorContextFunction GetContextFn(CalculatorContext* cc,
std::string tag) {
return cc->InputSidePackets().Tag(tag).Get<CalculatorContextFunction>();
}
};
REGISTER_CALCULATOR(LambdaCalculator);
mediapipe/framework/testdata/BUILD
Vendored
+8
View File
@@ -55,6 +55,14 @@ proto_library(
deps = ["@com_google_protobuf//:any_proto"],
)
mediapipe_cc_proto_library(
name = "zoo_mutator_cc_proto",
srcs = ["zoo_mutator.proto"],
cc_deps = ["@com_google_protobuf//:cc_wkt_protos"],
visibility = ["//mediapipe:__subpackages__"],
deps = [":zoo_mutator_proto"],
)
proto_library(
name = "zoo_mutation_calculator_proto",
srcs = ["zoo_mutation_calculator.proto"],
mediapipe/framework/tool/subgraph_expansion.cc
+3 -3
View File
@@ -237,9 +237,9 @@ static ::mediapipe::Status PrefixNames(int subgraph_index,
for (auto it = subgraph_nodes_start; it != nodes->end(); ++it) {
const auto& node = *it;
MP_RETURN_IF_ERROR(ValidateSubgraphFields(node));
ASSIGN_OR_RETURN(auto subgraph, graph_registry->CreateByName(
config->package(), node.calculator(),
&node.options()));
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(ConnectSubgraphStreams(node, &subgraph));
subgraphs.push_back(subgraph);
mediapipe/framework/tool/subgraph_expansion_test.cc
+2 -2
View File
@@ -128,8 +128,8 @@ class NodeChainSubgraph : public Subgraph {
public:
::mediapipe::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) override {
const mediapipe::NodeChainSubgraphOptions& opts =
options.GetExtension(mediapipe::NodeChainSubgraphOptions::ext);
auto opts =
Subgraph::GetOptions<mediapipe::NodeChainSubgraphOptions>(options);
const ProtoString& node_type = opts.node_type();
int chain_length = opts.chain_length();
RET_CHECK(!node_type.empty());