Project import generated by Copybara.
PiperOrigin-RevId: 253489161
This commit is contained in:
MediaPipe Team
@@ -0,0 +1,330 @@
|
||||
# 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.
|
||||
#
|
||||
|
||||
licenses(["notice"]) # Apache 2.0
|
||||
|
||||
package(
|
||||
default_visibility = ["//visibility:private"],
|
||||
features = ["-layering_check"],
|
||||
)
|
||||
|
||||
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_cc_proto_library")
|
||||
|
||||
proto_library(
|
||||
name = "default_input_stream_handler_proto",
|
||||
srcs = ["default_input_stream_handler.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:mediapipe_options_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "fixed_size_input_stream_handler_proto",
|
||||
srcs = ["fixed_size_input_stream_handler.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:mediapipe_options_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "sync_set_input_stream_handler_proto",
|
||||
srcs = ["sync_set_input_stream_handler.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:mediapipe_options_proto"],
|
||||
)
|
||||
|
||||
proto_library(
|
||||
name = "timestamp_align_input_stream_handler_proto",
|
||||
srcs = ["timestamp_align_input_stream_handler.proto"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = ["//mediapipe/framework:mediapipe_options_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "default_input_stream_handler_cc_proto",
|
||||
srcs = ["default_input_stream_handler.proto"],
|
||||
cc_deps = ["//mediapipe/framework:mediapipe_options_cc_proto"],
|
||||
visibility = ["//mediapipe/framework:__subpackages__"],
|
||||
deps = [":default_input_stream_handler_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "fixed_size_input_stream_handler_cc_proto",
|
||||
srcs = ["fixed_size_input_stream_handler.proto"],
|
||||
cc_deps = ["//mediapipe/framework:mediapipe_options_cc_proto"],
|
||||
visibility = ["//mediapipe/framework:__subpackages__"],
|
||||
deps = [":fixed_size_input_stream_handler_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "sync_set_input_stream_handler_cc_proto",
|
||||
srcs = ["sync_set_input_stream_handler.proto"],
|
||||
cc_deps = ["//mediapipe/framework:mediapipe_options_cc_proto"],
|
||||
visibility = ["//mediapipe/framework:__subpackages__"],
|
||||
deps = [":sync_set_input_stream_handler_proto"],
|
||||
)
|
||||
|
||||
mediapipe_cc_proto_library(
|
||||
name = "timestamp_align_input_stream_handler_cc_proto",
|
||||
srcs = ["timestamp_align_input_stream_handler.proto"],
|
||||
cc_deps = ["//mediapipe/framework:mediapipe_options_cc_proto"],
|
||||
visibility = ["//mediapipe/framework:__subpackages__"],
|
||||
deps = [":timestamp_align_input_stream_handler_proto"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "barrier_input_stream_handler",
|
||||
srcs = ["barrier_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "default_input_stream_handler",
|
||||
srcs = ["default_input_stream_handler.cc"],
|
||||
hdrs = ["default_input_stream_handler.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework/stream_handler:default_input_stream_handler_cc_proto",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "early_close_input_stream_handler",
|
||||
srcs = ["early_close_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "fixed_size_input_stream_handler",
|
||||
srcs = ["fixed_size_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
":default_input_stream_handler",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler_cc_proto",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "immediate_input_stream_handler",
|
||||
srcs = ["immediate_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "in_order_output_stream_handler",
|
||||
srcs = ["in_order_output_stream_handler.cc"],
|
||||
hdrs = ["in_order_output_stream_handler.h"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:collection",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework:output_stream_handler",
|
||||
"//mediapipe/framework:output_stream_manager",
|
||||
"//mediapipe/framework:output_stream_shard",
|
||||
"//mediapipe/framework:packet_set",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "mux_input_stream_handler",
|
||||
srcs = ["mux_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "sync_set_input_stream_handler",
|
||||
srcs = ["sync_set_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:collection",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework:packet_set",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/stream_handler:sync_set_input_stream_handler_cc_proto",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "timestamp_align_input_stream_handler",
|
||||
srcs = ["timestamp_align_input_stream_handler.cc"],
|
||||
visibility = ["//visibility:public"],
|
||||
deps = [
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/stream_handler:timestamp_align_input_stream_handler_cc_proto",
|
||||
"//mediapipe/framework/tool:validate_name",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "barrier_input_stream_handler_test",
|
||||
srcs = ["barrier_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":barrier_input_stream_handler",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"//mediapipe/framework/tool:tag_map_helper",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "immediate_input_stream_handler_test",
|
||||
srcs = ["immediate_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":immediate_input_stream_handler",
|
||||
"//mediapipe/framework:calculator_context",
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"//mediapipe/framework/tool:tag_map_helper",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/memory",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "mux_input_stream_handler_test",
|
||||
srcs = ["mux_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":mux_input_stream_handler",
|
||||
"//mediapipe/calculators/core:mux_calculator",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/calculators/core:round_robin_demux_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "set_input_stream_handler_test",
|
||||
srcs = ["set_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":fixed_size_input_stream_handler",
|
||||
":mux_input_stream_handler",
|
||||
"//mediapipe/calculators/core:mux_calculator",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/calculators/core:round_robin_demux_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/port:ret_check",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler_cc_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "default_input_stream_handler_test",
|
||||
srcs = ["default_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":default_input_stream_handler",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "fixed_size_input_stream_handler_test",
|
||||
srcs = ["fixed_size_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":fixed_size_input_stream_handler",
|
||||
"//mediapipe/calculators/core:counting_source_calculator",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/stream_handler:fixed_size_input_stream_handler_cc_proto",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "sync_set_input_stream_handler_test",
|
||||
srcs = ["sync_set_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":sync_set_input_stream_handler",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:test_calculators",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
"//mediapipe/framework/stream_handler:sync_set_input_stream_handler_cc_proto",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/memory",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
)
|
||||
|
||||
cc_test(
|
||||
name = "timestamp_align_input_stream_handler_test",
|
||||
srcs = ["timestamp_align_input_stream_handler_test.cc"],
|
||||
deps = [
|
||||
":timestamp_align_input_stream_handler",
|
||||
"//mediapipe/calculators/core:packet_cloner_calculator",
|
||||
"//mediapipe/calculators/core:pass_through_calculator",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework/port:gtest_main",
|
||||
"//mediapipe/framework/port:parse_text_proto",
|
||||
],
|
||||
)
|
||||
@@ -0,0 +1,95 @@
|
||||
// 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 <algorithm>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Implementation of an input stream handler that considers a node as ready for
|
||||
// Process() if all input streams have a packet available. This implies it must
|
||||
// consider a node as ready for Close() if any input stream is done.
|
||||
class BarrierInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
BarrierInputStreamHandler() = delete;
|
||||
BarrierInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* calculator_context_manager,
|
||||
const MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), calculator_context_manager,
|
||||
options, calculator_run_in_parallel) {}
|
||||
|
||||
void PrepareForRun(
|
||||
std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(::mediapipe::Status)> error_callback) override {
|
||||
InputStreamHandler::PrepareForRun(
|
||||
std::move(headers_ready_callback), std::move(notification_callback),
|
||||
std::move(schedule_callback), std::move(error_callback));
|
||||
for (auto& stream : input_stream_managers_) {
|
||||
stream->DisableTimestamps();
|
||||
}
|
||||
}
|
||||
|
||||
protected:
|
||||
// In BarrierInputStreamHandler, a node is "ready" if:
|
||||
// - any stream is done (need to call Close() in this case), or
|
||||
// - all streams have a packet available.
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
|
||||
DCHECK(min_stream_timestamp);
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
bool all_available = true;
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
if (stream_timestamp == Timestamp::Done()) {
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
all_available = false;
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
CHECK_NE(*min_stream_timestamp, Timestamp::Done());
|
||||
if (all_available) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
auto& stream = input_stream_managers_.Get(id);
|
||||
bool stream_is_done = false;
|
||||
Packet current_packet = stream->PopQueueHead(&stream_is_done);
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(BarrierInputStreamHandler);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,252 @@
|
||||
// Copyright 2018 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 <functional>
|
||||
#include <list>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/base/macros.h"
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/tool/tag_map.h"
|
||||
#include "mediapipe/framework/tool/tag_map_helper.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
class BarrierInputStreamHandlerTest : public ::testing::Test {
|
||||
protected:
|
||||
void SetUp() override {
|
||||
packet_type_.Set<std::string>();
|
||||
headers_ready_callback_ = [this]() {
|
||||
BarrierInputStreamHandlerTest::HeadersReadyNoOp();
|
||||
};
|
||||
notification_callback_ = [this]() {
|
||||
BarrierInputStreamHandlerTest::NotifyNoOp();
|
||||
};
|
||||
schedule_callback_ = std::bind(&BarrierInputStreamHandlerTest::Schedule,
|
||||
this, std::placeholders::_1);
|
||||
error_callback_ = std::bind(&BarrierInputStreamHandlerTest::RecordError,
|
||||
this, std::placeholders::_1);
|
||||
setup_shards_callback_ =
|
||||
std::bind(&BarrierInputStreamHandlerTest::SetupShardsNoOp, this,
|
||||
std::placeholders::_1);
|
||||
queue_full_callback_ =
|
||||
std::bind(&BarrierInputStreamHandlerTest::ReportQueueNoOp, this,
|
||||
std::placeholders::_1, std::placeholders::_2);
|
||||
queue_not_full_callback_ =
|
||||
std::bind(&BarrierInputStreamHandlerTest::ReportQueueNoOp, this,
|
||||
std::placeholders::_1, std::placeholders::_2);
|
||||
|
||||
std::shared_ptr<tool::TagMap> input_tag_map =
|
||||
tool::CreateTagMap({"input_a", "input_b", "input_c"}).ValueOrDie();
|
||||
|
||||
input_stream_managers_.reset(
|
||||
new InputStreamManager[input_tag_map->NumEntries()]);
|
||||
const std::vector<std::string>& names = input_tag_map->Names();
|
||||
for (CollectionItemId id = input_tag_map->BeginId();
|
||||
id < input_tag_map->EndId(); ++id) {
|
||||
const std::string& stream_name = names[id.value()];
|
||||
name_to_id_[stream_name] = id;
|
||||
MEDIAPIPE_ASSERT_OK(input_stream_managers_[id.value()].Initialize(
|
||||
stream_name, &packet_type_, /*back_edge=*/false));
|
||||
}
|
||||
SetupInputStreamHandler(input_tag_map);
|
||||
}
|
||||
|
||||
void SetupInputStreamHandler(
|
||||
const std::shared_ptr<tool::TagMap>& input_tag_map) {
|
||||
calculator_state_ = absl::make_unique<CalculatorState>(
|
||||
"Node", /*node_id=*/0, "Calculator", CalculatorGraphConfig::Node(),
|
||||
nullptr);
|
||||
calculator_context_manager_.Initialize(
|
||||
calculator_state_.get(), input_tag_map,
|
||||
/*output_tag_map=*/tool::CreateTagMap({"output_a"}).ValueOrDie(),
|
||||
/*calculator_run_in_parallel=*/false);
|
||||
|
||||
mediapipe::StatusOr<std::unique_ptr<mediapipe::InputStreamHandler>>
|
||||
status_or_handler = InputStreamHandlerRegistry::CreateByName(
|
||||
"BarrierInputStreamHandler", input_tag_map,
|
||||
&calculator_context_manager_, MediaPipeOptions(),
|
||||
/*calculator_run_in_parallel=*/false);
|
||||
ASSERT_TRUE(status_or_handler.ok());
|
||||
input_stream_handler_ = std::move(status_or_handler.ValueOrDie());
|
||||
MEDIAPIPE_ASSERT_OK(input_stream_handler_->InitializeInputStreamManagers(
|
||||
input_stream_managers_.get()));
|
||||
MEDIAPIPE_ASSERT_OK(
|
||||
calculator_context_manager_.PrepareForRun(setup_shards_callback_));
|
||||
input_stream_handler_->PrepareForRun(headers_ready_callback_,
|
||||
notification_callback_,
|
||||
schedule_callback_, error_callback_);
|
||||
input_stream_handler_->SetQueueSizeCallbacks(queue_full_callback_,
|
||||
queue_not_full_callback_);
|
||||
}
|
||||
|
||||
void HeadersReadyNoOp() {}
|
||||
|
||||
void NotifyNoOp() {}
|
||||
|
||||
void Schedule(CalculatorContext* calculator_context) {
|
||||
CHECK(calculator_context);
|
||||
calculator_context_ = calculator_context;
|
||||
}
|
||||
|
||||
void RecordError(const ::mediapipe::Status& error) {
|
||||
errors_.push_back(error);
|
||||
}
|
||||
|
||||
::mediapipe::Status SetupShardsNoOp(CalculatorContext* calculator_context) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
void ReportQueueNoOp(InputStreamManager* stream, bool* stream_was_full) {}
|
||||
|
||||
PacketType packet_type_;
|
||||
std::function<void()> headers_ready_callback_;
|
||||
std::function<void()> notification_callback_;
|
||||
std::function<void(CalculatorContext*)> schedule_callback_;
|
||||
std::function<void(::mediapipe::Status)> error_callback_;
|
||||
std::function<::mediapipe::Status(CalculatorContext*)> setup_shards_callback_;
|
||||
InputStreamManager::QueueSizeCallback queue_full_callback_;
|
||||
InputStreamManager::QueueSizeCallback queue_not_full_callback_;
|
||||
|
||||
// Vector of errors encountered while using the stream.
|
||||
std::vector<::mediapipe::Status> errors_;
|
||||
|
||||
std::unique_ptr<CalculatorState> calculator_state_;
|
||||
CalculatorContextManager calculator_context_manager_;
|
||||
CalculatorContext* calculator_context_;
|
||||
std::map<std::string, CollectionItemId> name_to_id_;
|
||||
std::unique_ptr<InputStreamHandler> input_stream_handler_;
|
||||
std::unique_ptr<InputStreamManager[]> input_stream_managers_;
|
||||
};
|
||||
|
||||
// This test simulates how CalculatorNode::ProcessNode() uses an
|
||||
// input stream handler and the associated input streams.
|
||||
TEST_F(BarrierInputStreamHandlerTest, SimulateProcessNode) {
|
||||
Timestamp min_stream_timestamp;
|
||||
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
|
||||
std::list<Packet> packets;
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(10)));
|
||||
packets.push_back(Adopt(new std::string("packet 2")).At(Timestamp(30)));
|
||||
packets.push_back(Adopt(new std::string("packet 3")).At(Timestamp(20)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_a"], packets);
|
||||
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(0)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_b"], packets);
|
||||
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(0)));
|
||||
packets.push_back(Adopt(new std::string("packet 2")).At(Timestamp(0)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_c"], packets);
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream =
|
||||
calculator_context_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_FALSE(input_stream.Value().IsEmpty());
|
||||
EXPECT_EQ(input_stream.Value().Get<std::string>(), "packet 1");
|
||||
}
|
||||
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(calculator_context_->InputTimestamp(),
|
||||
&calculator_context_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(calculator_context_);
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream =
|
||||
calculator_context_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_TRUE(input_stream.Value().IsEmpty());
|
||||
}
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 3")).At(Timestamp(0)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_c"], packets);
|
||||
|
||||
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 2")).At(Timestamp(1)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_b"], packets);
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_b"],
|
||||
Timestamp::Done());
|
||||
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream =
|
||||
calculator_context_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_FALSE(input_stream.Value().IsEmpty());
|
||||
EXPECT_EQ(input_stream.Value().Get<std::string>(), "packet 2");
|
||||
}
|
||||
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(calculator_context_->InputTimestamp(),
|
||||
&calculator_context_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(calculator_context_);
|
||||
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream =
|
||||
calculator_context_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_TRUE(input_stream.Value().IsEmpty());
|
||||
}
|
||||
|
||||
EXPECT_FALSE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_a"])
|
||||
->IsEmpty());
|
||||
EXPECT_TRUE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_b"])
|
||||
->IsEmpty());
|
||||
EXPECT_FALSE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_c"])
|
||||
->IsEmpty());
|
||||
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
EXPECT_EQ(Timestamp::Done(), calculator_context_->InputTimestamp());
|
||||
// For Close(), input set should be empty
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream =
|
||||
calculator_context_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_TRUE(input_stream.Value().IsEmpty());
|
||||
}
|
||||
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(calculator_context_->InputTimestamp(),
|
||||
&calculator_context_->Inputs());
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream =
|
||||
calculator_context_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_TRUE(input_stream.Value().IsEmpty());
|
||||
}
|
||||
|
||||
EXPECT_TRUE(errors_.empty());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,81 @@
|
||||
// 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 "mediapipe/framework/stream_handler/default_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include "absl/strings/substitute.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(DefaultInputStreamHandler);
|
||||
|
||||
DefaultInputStreamHandler::DefaultInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), cc_manager, options,
|
||||
calculator_run_in_parallel) {
|
||||
if (options.HasExtension(DefaultInputStreamHandlerOptions::ext)) {
|
||||
SetBatchSize(options.GetExtension(DefaultInputStreamHandlerOptions::ext)
|
||||
.batch_size());
|
||||
}
|
||||
}
|
||||
|
||||
NodeReadiness DefaultInputStreamHandler::GetNodeReadiness(
|
||||
Timestamp* min_stream_timestamp) {
|
||||
DCHECK(min_stream_timestamp);
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
min_bound = std::min(min_bound, stream_timestamp);
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
if (*min_stream_timestamp == Timestamp::Done()) {
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
|
||||
if (min_bound > *min_stream_timestamp) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
CHECK_EQ(min_bound, *min_stream_timestamp);
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
void DefaultInputStreamHandler::FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
auto& stream = input_stream_managers_.Get(id);
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Packet current_packet = stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
CHECK_EQ(num_packets_dropped, 0)
|
||||
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, stream->Name());
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,52 @@
|
||||
// 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_FRAMEWORK_STREAM_HANDLER_DEFAULT_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_DEFAULT_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
// TODO: Move protos in another CL after the C++ code migration.
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/stream_handler/default_input_stream_handler.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Implementation of the "default" input stream handler that is applied on a
|
||||
// given CalculatorGraph when no input stream handler is explicitly specified.
|
||||
class DefaultInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
DefaultInputStreamHandler() = delete;
|
||||
DefaultInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options,
|
||||
bool calculator_run_in_parallel);
|
||||
|
||||
protected:
|
||||
// In DefaultInputStreamHandler, a node is "ready" if:
|
||||
// - all streams are done (need to call Close() in this case), or
|
||||
// - the minimum bound (over all empty streams) is greater than the smallest
|
||||
// timestamp of any stream, which means we have received all the packets
|
||||
// that will be available at the next timestamp.
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_DEFAULT_INPUT_STREAM_HANDLER_H_
|
||||
@@ -0,0 +1,37 @@
|
||||
// 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.
|
||||
//
|
||||
// Forked from
|
||||
// mediapipe/framework/stream_handler/default_input_stream_handler.proto
|
||||
// The forked proto must remain identical to the original proto and should be
|
||||
// ONLY used by mediapipe open source project.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/mediapipe_options.proto";
|
||||
|
||||
// See DefaultInputStreamHandler for documentation.
|
||||
message DefaultInputStreamHandlerOptions {
|
||||
extend MediaPipeOptions {
|
||||
optional DefaultInputStreamHandlerOptions ext = 164312821;
|
||||
}
|
||||
// batch_size determines how many input packets should be collected before
|
||||
// a calculator can process them. Once there are enough packets, Process
|
||||
// method of the Calculator is called sequentially. Currently, batching is not
|
||||
// supported for source nodes but it may be supported in the future.
|
||||
// Therefore, this field should not be specified for source nodes.
|
||||
optional int32 batch_size = 1 [default = 1];
|
||||
}
|
||||
@@ -0,0 +1,283 @@
|
||||
// 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 <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
// This test shows the default behavior of DefaultInputStreamHandler when
|
||||
// batching is disabled.
|
||||
TEST(DefaultInputStreamHandlerTest, NoBatchingWorks) {
|
||||
// A single calculator with two input streams, and two output streams. This
|
||||
// calculator passes all the input packets along.
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input0"
|
||||
input_stream: "input1"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input0"
|
||||
input_stream: "input1"
|
||||
output_stream: "output0"
|
||||
output_stream: "output1"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "DefaultInputStreamHandler"
|
||||
options: {
|
||||
[mediapipe.DefaultInputStreamHandlerOptions.ext]: {
|
||||
batch_size: 1
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
std::vector<Packet> sink_0, sink_1;
|
||||
tool::AddVectorSink("output0", &config, &sink_0);
|
||||
tool::AddVectorSink("output1", &config, &sink_1);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(1)).At(Timestamp(1))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// No packets expected as the second stream is not ready to be processed.
|
||||
EXPECT_EQ(0, sink_0.size());
|
||||
EXPECT_EQ(0, sink_1.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input1", Adopt(new int(2)).At(Timestamp(2))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// First stream can produce output because the timestamp bound of the second
|
||||
// stream is higher.
|
||||
EXPECT_EQ(1, sink_0.size());
|
||||
EXPECT_EQ(0, sink_1.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(2)).At(Timestamp(2))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// Both streams have packets at the same timestamp, therefore both can produce
|
||||
// packets.
|
||||
EXPECT_EQ(2, sink_0.size());
|
||||
EXPECT_EQ(1, sink_1.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
// This test shows the effect of batching on the DefaultInputStreamHandler.
|
||||
TEST(DefaultInputStreamHandlerTest, Batches) {
|
||||
// A single batching calculator with one input stream and one output stream.
|
||||
// This calculator passes all the input packets onto the output streams.
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input0"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input0"
|
||||
output_stream: "output0"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "DefaultInputStreamHandler"
|
||||
options: {
|
||||
[mediapipe.DefaultInputStreamHandlerOptions.ext]: {
|
||||
batch_size: 2
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
std::vector<Packet> sink;
|
||||
tool::AddVectorSink("output0", &config, &sink);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(1)).At(Timestamp(1))));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// There shouldn't be any outputs until a set of two packets is batched.
|
||||
EXPECT_TRUE(sink.empty());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(2)).At(Timestamp(2))));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// There should be two packets, processed during a single invocation.
|
||||
ASSERT_EQ(2, sink.size());
|
||||
EXPECT_THAT(std::vector<int>({sink[0].Get<int>(), sink[1].Get<int>()}),
|
||||
testing::ElementsAre(1, 2));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(3)).At(Timestamp(3))));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// There shouldn't be any outputs until another set of two packets is batched.
|
||||
EXPECT_EQ(2, sink.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(4)).At(Timestamp(4))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// The new batch was complete. There should be two more output packets.
|
||||
ASSERT_EQ(4, sink.size());
|
||||
EXPECT_THAT(std::vector<int>({sink[0].Get<int>(), sink[1].Get<int>(),
|
||||
sink[2].Get<int>(), sink[3].Get<int>()}),
|
||||
testing::ElementsAre(1, 2, 3, 4));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
// This test shows that any packets get flushed (outputted) when the input
|
||||
// streams are closed.
|
||||
TEST(DefaultInputStreamHandlerTest, BatchIsFlushedWhenClosing) {
|
||||
// A single batching calculator with one input stream and one output stream.
|
||||
// This calculator passes all the input packets onto the output streams.
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input0"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input0"
|
||||
output_stream: "output0"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "DefaultInputStreamHandler"
|
||||
options: {
|
||||
[mediapipe.DefaultInputStreamHandlerOptions.ext]: {
|
||||
batch_size: 2
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
std::vector<Packet> sink;
|
||||
tool::AddVectorSink("output0", &config, &sink);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(1)).At(Timestamp(1))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// There shouldn't be any outputs until a set of two packets is batched.
|
||||
EXPECT_TRUE(sink.empty());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(2)).At(Timestamp(2))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// There should be two packets, processed during a single invocation.
|
||||
ASSERT_EQ(2, sink.size());
|
||||
EXPECT_THAT(std::vector<int>({sink[0].Get<int>(), sink[1].Get<int>()}),
|
||||
testing::ElementsAre(1, 2));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(3)).At(Timestamp(3))));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// There shouldn't be any outputs until another set of two packets is batched.
|
||||
EXPECT_EQ(2, sink.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
// When the streams are done, the packets currently being batched should be
|
||||
// flushed out.
|
||||
ASSERT_EQ(3, sink.size());
|
||||
// Batched outputs should be in correct order.
|
||||
EXPECT_THAT(std::vector<int>(
|
||||
{sink[0].Get<int>(), sink[1].Get<int>(), sink[2].Get<int>()}),
|
||||
testing::ElementsAre(1, 2, 3));
|
||||
}
|
||||
|
||||
// This test shows that calculators won't propagate timestamp while they are
|
||||
// batching except for the first timestamp of the batch.
|
||||
TEST(DefaultInputStreamHandlerTest, DoesntPropagateTimestampWhenBatching) {
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input0"
|
||||
input_stream: "input1"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input1"
|
||||
output_stream: "input1_batched"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "DefaultInputStreamHandler"
|
||||
options: {
|
||||
[mediapipe.DefaultInputStreamHandlerOptions.ext]: {
|
||||
batch_size: 3
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input0"
|
||||
input_stream: "input1_batched"
|
||||
output_stream: "output"
|
||||
output_stream: "dummy"
|
||||
})");
|
||||
std::vector<Packet> sink;
|
||||
tool::AddVectorSink("output", &config, &sink);
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(0)).At(Timestamp(0))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
EXPECT_TRUE(sink.empty());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(1)).At(Timestamp(1))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input1", Adopt(new int(1)).At(Timestamp(1))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// Both calculators have packet 1. First node is currently batching and it
|
||||
// propagates the first input timestamp in the batch. Therefore, the
|
||||
// second node should produce output for the packet at 0.
|
||||
EXPECT_EQ(1, sink.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(2)).At(Timestamp(2))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input1", Adopt(new int(2)).At(Timestamp(2))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// Due to batching on the first node, timestamp is not propagated for the
|
||||
// packet at timestamp 2. Therefore, the second node cannot process the packet
|
||||
// at timestamp 1.
|
||||
EXPECT_EQ(1, sink.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input0", Adopt(new int(3)).At(Timestamp(3))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input1", Adopt(new int(3)).At(Timestamp(3))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// Batching is complete on the first node. It produced outputs at timestamp 1,
|
||||
// 2, and 3. The first node can now process the input packets at timestamps 1,
|
||||
// 2, and 3 as well.
|
||||
EXPECT_EQ(4, sink.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
EXPECT_EQ(4, sink.size());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,92 @@
|
||||
// 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 <algorithm>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Implementation of an input stream handler that considers a node as ready for
|
||||
// Close() if any input stream is done.
|
||||
class EarlyCloseInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
EarlyCloseInputStreamHandler() = delete;
|
||||
EarlyCloseInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options,
|
||||
bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), cc_manager, options,
|
||||
calculator_run_in_parallel) {}
|
||||
|
||||
protected:
|
||||
// In EarlyCloseInputStreamHandler, a node is "ready" if:
|
||||
// - any stream is done (need to call Close() in this case), or
|
||||
// - the minimum bound (over all empty streams) is greater than the smallest
|
||||
// timestamp of any stream, which means we have received all the packets
|
||||
// that will be available at the next timestamp.
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
|
||||
DCHECK(min_stream_timestamp);
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
if (stream_timestamp == Timestamp::Done()) {
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
min_bound = std::min(min_bound, stream_timestamp);
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
CHECK_NE(*min_stream_timestamp, Timestamp::Done());
|
||||
|
||||
if (min_bound > *min_stream_timestamp) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
CHECK_EQ(min_bound, *min_stream_timestamp);
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
auto& stream = input_stream_managers_.Get(id);
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Packet current_packet = stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
CHECK_EQ(num_packets_dropped, 0)
|
||||
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, stream->Name());
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(EarlyCloseInputStreamHandler);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,221 @@
|
||||
// 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 <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/stream_handler/default_input_stream_handler.h"
|
||||
// TODO: Move protos in another CL after the C++ code migration.
|
||||
#include "mediapipe/framework/stream_handler/fixed_size_input_stream_handler.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Input stream handler that limits each input queue to a maximum of
|
||||
// target_queue_size packets, discarding older packets as needed. When a
|
||||
// timestamp is dropped from a stream, it is dropped from all others as well.
|
||||
//
|
||||
// For example, a calculator node with one input stream and the following input
|
||||
// stream handler specs:
|
||||
//
|
||||
// node {
|
||||
// calculator: "CalculatorRunningAtOneFps"
|
||||
// input_stream: "packets_streaming_in_at_ten_fps"
|
||||
// input_stream_handler {
|
||||
// input_stream_handler: "FixedSizeInputStreamHandler"
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// will always try to keep the newest packet in the input stream.
|
||||
//
|
||||
// A few details: FixedSizeInputStreamHandler takes action when any stream grows
|
||||
// to trigger_queue_size or larger. It then keeps at most target_queue_size
|
||||
// packets in every InputStreamImpl. Every stream is truncated at the same
|
||||
// timestamp, so that each included timestamp delivers the same packets as
|
||||
// DefaultInputStreamHandler includes.
|
||||
//
|
||||
class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
|
||||
public:
|
||||
FixedSizeInputStreamHandler() = delete;
|
||||
FixedSizeInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options,
|
||||
bool calculator_run_in_parallel)
|
||||
: DefaultInputStreamHandler(std::move(tag_map), cc_manager, options,
|
||||
calculator_run_in_parallel) {
|
||||
const auto& ext =
|
||||
options.GetExtension(FixedSizeInputStreamHandlerOptions::ext);
|
||||
trigger_queue_size_ = ext.trigger_queue_size();
|
||||
target_queue_size_ = ext.target_queue_size();
|
||||
fixed_min_size_ = ext.fixed_min_size();
|
||||
pending_ = false;
|
||||
kept_timestamp_ = Timestamp::Unset();
|
||||
// TODO: Either re-enable SetLatePreparation(true) with
|
||||
// CalculatorContext::InputTimestamp set correctly, or remove the
|
||||
// implementation of SetLatePreparation.
|
||||
}
|
||||
|
||||
private:
|
||||
// Drops packets if all input streams exceed trigger_queue_size.
|
||||
void EraseAllSurplus() EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
|
||||
Timestamp min_timestamp_all_streams = Timestamp::Max();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
// Check whether every InputStreamImpl grew beyond trigger_queue_size.
|
||||
if (stream->QueueSize() < trigger_queue_size_) {
|
||||
return;
|
||||
}
|
||||
Timestamp min_timestamp =
|
||||
stream->GetMinTimestampAmongNLatest(target_queue_size_);
|
||||
|
||||
// Record the min timestamp among the newest target_queue_size_ packets
|
||||
// across all InputStreamImpls.
|
||||
min_timestamp_all_streams =
|
||||
std::min(min_timestamp_all_streams, min_timestamp);
|
||||
}
|
||||
for (auto& stream : input_stream_managers_) {
|
||||
stream->ErasePacketsEarlierThan(min_timestamp_all_streams);
|
||||
}
|
||||
}
|
||||
|
||||
// Returns the latest timestamp allowed before a bound.
|
||||
Timestamp PreviousAllowedInStream(Timestamp bound) {
|
||||
return bound.IsRangeValue() ? bound - 1 : bound;
|
||||
}
|
||||
|
||||
// Returns the lowest timestamp at which a packet may arrive at any stream.
|
||||
Timestamp MinStreamBound() {
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
Timestamp stream_bound = stream->GetMinTimestampAmongNLatest(1);
|
||||
if (stream_bound > Timestamp::Unset()) {
|
||||
stream_bound = stream_bound.NextAllowedInStream();
|
||||
} else {
|
||||
stream_bound = stream->MinTimestampOrBound(nullptr);
|
||||
}
|
||||
min_bound = std::min(min_bound, stream_bound);
|
||||
}
|
||||
return min_bound;
|
||||
}
|
||||
|
||||
// Returns the lowest timestamp of a packet ready to process.
|
||||
Timestamp MinTimestampOrBound() {
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
min_bound = std::min(min_bound, stream->MinTimestampOrBound(nullptr));
|
||||
}
|
||||
return min_bound;
|
||||
}
|
||||
|
||||
// Keeps only the most recent target_queue_size packets in each stream
|
||||
// exceeding trigger_queue_size. Also, discards all packets older than the
|
||||
// first kept timestamp on any stream.
|
||||
void EraseAnySurplus(bool keep_one) EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
|
||||
// Record the most recent first kept timestamp on any stream.
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
int32 queue_size = (stream->QueueSize() >= trigger_queue_size_)
|
||||
? target_queue_size_
|
||||
: trigger_queue_size_ - 1;
|
||||
if (stream->QueueSize() > queue_size) {
|
||||
kept_timestamp_ = std::max(
|
||||
kept_timestamp_, stream->GetMinTimestampAmongNLatest(queue_size + 1)
|
||||
.NextAllowedInStream());
|
||||
}
|
||||
}
|
||||
if (keep_one) {
|
||||
// In order to preserve one viable timestamp, do not truncate past
|
||||
// the timestamp bound of the least current stream.
|
||||
kept_timestamp_ =
|
||||
std::min(kept_timestamp_, PreviousAllowedInStream(MinStreamBound()));
|
||||
}
|
||||
for (auto& stream : input_stream_managers_) {
|
||||
stream->ErasePacketsEarlierThan(kept_timestamp_);
|
||||
}
|
||||
}
|
||||
|
||||
void EraseSurplusPackets(bool keep_one)
|
||||
EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
|
||||
return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one);
|
||||
}
|
||||
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) {
|
||||
DCHECK(min_stream_timestamp);
|
||||
absl::MutexLock lock(&erase_mutex_);
|
||||
// kReadyForProcess is returned only once until FillInputSet completes.
|
||||
// In late_preparation mode, GetNodeReadiness must return kReadyForProcess
|
||||
// exactly once for each input-set produced. Here, GetNodeReadiness
|
||||
// releases just one input-set at a time and then disables input queue
|
||||
// truncation until that promised input-set is consumed.
|
||||
if (pending_) {
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
EraseSurplusPackets(false);
|
||||
NodeReadiness result =
|
||||
DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
|
||||
|
||||
// If a packet has arrived below kept_timestamp_, recalculate.
|
||||
while (*min_stream_timestamp < kept_timestamp_ &&
|
||||
result == NodeReadiness::kReadyForProcess) {
|
||||
EraseSurplusPackets(false);
|
||||
result =
|
||||
DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
|
||||
}
|
||||
pending_ = (result == NodeReadiness::kReadyForProcess);
|
||||
return result;
|
||||
}
|
||||
|
||||
void AddPackets(CollectionItemId id,
|
||||
const std::list<Packet>& packets) override {
|
||||
InputStreamHandler::AddPackets(id, packets);
|
||||
absl::MutexLock lock(&erase_mutex_);
|
||||
if (!pending_) {
|
||||
EraseSurplusPackets(false);
|
||||
}
|
||||
}
|
||||
|
||||
void MovePackets(CollectionItemId id, std::list<Packet>* packets) override {
|
||||
InputStreamHandler::MovePackets(id, packets);
|
||||
absl::MutexLock lock(&erase_mutex_);
|
||||
if (!pending_) {
|
||||
EraseSurplusPackets(false);
|
||||
}
|
||||
}
|
||||
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override {
|
||||
CHECK(input_set);
|
||||
absl::MutexLock lock(&erase_mutex_);
|
||||
if (!pending_) {
|
||||
LOG(ERROR) << "FillInputSet called without GetNodeReadiness.";
|
||||
}
|
||||
// input_timestamp is recalculated here to process the most recent packets.
|
||||
EraseSurplusPackets(true);
|
||||
input_timestamp = MinTimestampOrBound();
|
||||
DefaultInputStreamHandler::FillInputSet(input_timestamp, input_set);
|
||||
pending_ = false;
|
||||
}
|
||||
|
||||
private:
|
||||
int32 trigger_queue_size_;
|
||||
int32 target_queue_size_;
|
||||
bool fixed_min_size_;
|
||||
// Indicates that GetNodeReadiness has returned kReadyForProcess once, and
|
||||
// the corresponding call to FillInputSet has not yet completed.
|
||||
bool pending_ GUARDED_BY(erase_mutex_);
|
||||
// The timestamp used to truncate all input streams.
|
||||
Timestamp kept_timestamp_ GUARDED_BY(erase_mutex_);
|
||||
absl::Mutex erase_mutex_;
|
||||
};
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,38 @@
|
||||
// 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.
|
||||
//
|
||||
// Forked from
|
||||
// mediapipe/framework/stream_handler/fixed_size_input_stream_handler.proto
|
||||
// The forked proto must remain identical to the original proto and should be
|
||||
// ONLY used by mediapipe open source project.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/mediapipe_options.proto";
|
||||
|
||||
// See FixedSizeInputStreamHandler for documentation.
|
||||
message FixedSizeInputStreamHandlerOptions {
|
||||
extend MediaPipeOptions {
|
||||
optional FixedSizeInputStreamHandlerOptions ext = 125744319;
|
||||
}
|
||||
// The queue size at which input queues are truncated.
|
||||
optional int32 trigger_queue_size = 1 [default = 2];
|
||||
// The queue size to which input queues are truncated.
|
||||
optional int32 target_queue_size = 2 [default = 1];
|
||||
// If false, input queues are truncated to at most trigger_queue_size.
|
||||
// If true, input queues are truncated to at least trigger_queue_size.
|
||||
optional bool fixed_min_size = 3 [default = false];
|
||||
}
|
||||
@@ -0,0 +1,363 @@
|
||||
// Copyright 2018 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 <vector>
|
||||
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/integral_types.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/port/threadpool.h"
|
||||
#include "mediapipe/framework/stream_handler/fixed_size_input_stream_handler.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
const int64 kMaxPacketId = 100;
|
||||
const int64 kSlowCalculatorRate = 10;
|
||||
|
||||
// Rate limiter for TestSlowCalculator.
|
||||
ABSL_CONST_INIT absl::Mutex g_source_mutex(absl::kConstInit);
|
||||
int64 g_source_counter GUARDED_BY(g_source_mutex);
|
||||
|
||||
// Rate limiter for TestSourceCalculator.
|
||||
int64 g_slow_counter GUARDED_BY(g_source_mutex);
|
||||
|
||||
// Flag that indicates that the source is done.
|
||||
bool g_source_done GUARDED_BY(g_source_mutex);
|
||||
|
||||
class TestSourceCalculator : public CalculatorBase {
|
||||
public:
|
||||
TestSourceCalculator() : current_packet_id_(0) {}
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Outputs().Index(0).Set<int64>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
absl::MutexLock lock(&g_source_mutex);
|
||||
g_source_counter = 0;
|
||||
g_source_done = false;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
if (current_packet_id_ == kMaxPacketId) {
|
||||
absl::MutexLock lock(&g_source_mutex);
|
||||
g_source_done = true;
|
||||
return tool::StatusStop();
|
||||
}
|
||||
cc->Outputs().Index(0).Add(new int64(0), Timestamp(current_packet_id_));
|
||||
++current_packet_id_;
|
||||
{
|
||||
absl::MutexLock lock(&g_source_mutex);
|
||||
++g_source_counter;
|
||||
g_source_mutex.Await(
|
||||
absl::Condition(this, &TestSourceCalculator::CanProceed));
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
bool CanProceed() const EXCLUSIVE_LOCKS_REQUIRED(g_source_mutex) {
|
||||
return g_source_counter <= kSlowCalculatorRate * g_slow_counter ||
|
||||
g_source_counter <= 1;
|
||||
}
|
||||
int64 current_packet_id_;
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(TestSourceCalculator);
|
||||
|
||||
class TestSlowCalculator : public CalculatorBase {
|
||||
public:
|
||||
TestSlowCalculator() = default;
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
cc->Inputs().Index(0).Set<int64>();
|
||||
cc->Outputs().Index(0).Set<int64>();
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
::mediapipe::Status Open(CalculatorContext* cc) override {
|
||||
absl::MutexLock lock(&g_source_mutex);
|
||||
g_slow_counter = 0;
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
::mediapipe::Status Process(CalculatorContext* cc) override {
|
||||
cc->Outputs().Index(0).Add(new int64(0),
|
||||
cc->Inputs().Index(0).Value().Timestamp());
|
||||
{
|
||||
absl::MutexLock lock(&g_source_mutex);
|
||||
++g_slow_counter;
|
||||
g_source_mutex.Await(
|
||||
absl::Condition(this, &TestSlowCalculator::CanProceed));
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
private:
|
||||
bool CanProceed() const EXCLUSIVE_LOCKS_REQUIRED(g_source_mutex) {
|
||||
return g_source_counter > kSlowCalculatorRate * g_slow_counter ||
|
||||
g_source_done;
|
||||
}
|
||||
};
|
||||
|
||||
REGISTER_CALCULATOR(TestSlowCalculator);
|
||||
|
||||
// Return the values of the timestamps of a vector of Packets.
|
||||
static std::vector<int64> TimestampValues(const std::vector<Packet>& packets) {
|
||||
std::vector<int64> result;
|
||||
for (const Packet& p : packets) {
|
||||
result.push_back(p.Timestamp().Value());
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
void SetFixedMinSize(CalculatorGraphConfig::Node* node, bool fixed_min_size) {
|
||||
node->mutable_input_stream_handler()
|
||||
->mutable_options()
|
||||
->MutableExtension(FixedSizeInputStreamHandlerOptions::ext)
|
||||
->set_fixed_min_size(fixed_min_size);
|
||||
}
|
||||
|
||||
class FixedSizeInputStreamHandlerTest : public ::testing::TestWithParam<bool> {
|
||||
};
|
||||
INSTANTIATE_TEST_SUITE_P(InstantiationFixed, FixedSizeInputStreamHandlerTest,
|
||||
::testing::Values(false, true));
|
||||
|
||||
TEST_P(FixedSizeInputStreamHandlerTest, DropsPackets) {
|
||||
// Sink consumes roughly 10x slower than source produces output which is
|
||||
// simulated with a couple of conditional critical sections. Both calculators
|
||||
// are rate limited in a circular fashion. The source produces 10 packets,
|
||||
// then the TestSlowCalculator consumes one packet, then the source produces
|
||||
// the next 10 packets. One packet is sent initially by the source to start
|
||||
// the processing. The CanProceed conditions for the two calculators are
|
||||
// mutally exclusive to avoid race conditions between them. Queue size is
|
||||
// regulated by FixedSizeInputStreamHandler.
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(node {
|
||||
calculator: "TestSourceCalculator"
|
||||
output_stream: "input_packets"
|
||||
}
|
||||
node {
|
||||
calculator: "TestSlowCalculator"
|
||||
input_stream: "input_packets"
|
||||
output_stream: "output_packets"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "FixedSizeInputStreamHandler"
|
||||
}
|
||||
})");
|
||||
SetFixedMinSize(graph_config.mutable_node(1), GetParam());
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output_packets", &graph_config, &output_packets);
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(graph_config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.Run());
|
||||
|
||||
// The TestSlowCalculator consumes one packet after every tenth packet
|
||||
// is sent. All other packets are dropped by the FixedSizeInputStreamHandler.
|
||||
// The TestSourceCalculator sends 10 more packets after each packet is
|
||||
// consumed. In this way, the TestSlowCalculator consumes and outputs only
|
||||
// every tenth packet.
|
||||
EXPECT_EQ(output_packets.size(), 11);
|
||||
std::vector<int64> expected_ts = {0, 9, 19, 29, 39, 49, 59, 69, 79, 89, 99};
|
||||
EXPECT_THAT(TimestampValues(output_packets),
|
||||
testing::ContainerEq(expected_ts));
|
||||
}
|
||||
|
||||
// A regression test for b/74820963. This test should not time out.
|
||||
TEST_P(FixedSizeInputStreamHandlerTest, DropsPacketsInFullStream) {
|
||||
// CountingSourceCalculator outputs 10 packets at a time. Since max_queue_size
|
||||
// is 10, CountingSourceCalculator will fill up the "input_packets" input
|
||||
// stream of PassThroughCalculator and cause itself to be throttled. When
|
||||
// FixedSizeInputStreamHandler discards the queued packets, it should report
|
||||
// the "queue became non-full" event and unthrottle CountingSourceCalculator.
|
||||
// If FixedSizeInputStreamHandler fails to report the event,
|
||||
// CountingSourceCalculator will stay throttled and the test will time out.
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(max_queue_size: 10
|
||||
node {
|
||||
calculator: "CountingSourceCalculator"
|
||||
input_side_packet: "MAX_COUNT:max_count"
|
||||
input_side_packet: "BATCH_SIZE:batch_size"
|
||||
output_stream: "input_packets"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input_packets"
|
||||
output_stream: "output_packets"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "FixedSizeInputStreamHandler"
|
||||
}
|
||||
})");
|
||||
SetFixedMinSize(graph_config.mutable_node(1), GetParam());
|
||||
std::vector<Packet> output_packets;
|
||||
tool::AddVectorSink("output_packets", &graph_config, &output_packets);
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(
|
||||
graph.Initialize(graph_config, {{"max_count", MakePacket<int>(10)},
|
||||
{"batch_size", MakePacket<int>(10)}}));
|
||||
MEDIAPIPE_ASSERT_OK(graph.Run());
|
||||
}
|
||||
|
||||
// Tests FixedSizeInputStreamHandler with several input streams running
|
||||
// asynchronously in parallel.
|
||||
TEST_P(FixedSizeInputStreamHandlerTest, ParallelWriteAndRead) {
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "in_0"
|
||||
input_stream: "in_1"
|
||||
input_stream: "in_2"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "in_0"
|
||||
input_stream: "in_1"
|
||||
input_stream: "in_2"
|
||||
output_stream: "out_0"
|
||||
output_stream: "out_1"
|
||||
output_stream: "out_2"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "FixedSizeInputStreamHandler"
|
||||
options {
|
||||
[mediapipe.FixedSizeInputStreamHandlerOptions.ext] {
|
||||
trigger_queue_size: 3
|
||||
target_queue_size: 1
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
SetFixedMinSize(graph_config.mutable_node(0), GetParam());
|
||||
std::vector<Packet> output_packets[3];
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
tool::AddVectorSink(absl::StrCat("out_", i), &graph_config,
|
||||
&output_packets[i]);
|
||||
}
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(graph_config, {}));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
{
|
||||
::mediapipe::ThreadPool pool(3);
|
||||
pool.StartWorkers();
|
||||
|
||||
// Start 3 writers.
|
||||
for (int w = 0; w < 3; ++w) {
|
||||
pool.Schedule([&, w]() {
|
||||
std::string stream_name = absl::StrCat("in_", w);
|
||||
for (int i = 0; i < 50; ++i) {
|
||||
Packet p = MakePacket<int>(i).At(Timestamp(i));
|
||||
MEDIAPIPE_EXPECT_OK(graph.AddPacketToInputStream(stream_name, p));
|
||||
absl::SleepFor(absl::Microseconds(100));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
EXPECT_EQ(output_packets[i].size(), output_packets[0].size());
|
||||
for (int j = 0; j < output_packets[i].size(); j++) {
|
||||
EXPECT_EQ(output_packets[i][j].Get<int>(),
|
||||
output_packets[0][j].Get<int>());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Tests dropping of packets that arrive later.
|
||||
// A: 1 2 3[4 5 6]
|
||||
// B: 3 4[5 6 7]
|
||||
// C:
|
||||
// This should cut at 5. If we then send 4 on C, it should be dropped, just as
|
||||
// it would have been if it arrived earlier.
|
||||
TEST_P(FixedSizeInputStreamHandlerTest, LateArrivalDrop) {
|
||||
CalculatorGraphConfig graph_config =
|
||||
ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "in_0"
|
||||
input_stream: "in_1"
|
||||
input_stream: "in_2"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "in_0"
|
||||
input_stream: "in_1"
|
||||
input_stream: "in_2"
|
||||
output_stream: "out_0"
|
||||
output_stream: "out_1"
|
||||
output_stream: "out_2"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "FixedSizeInputStreamHandler"
|
||||
options {
|
||||
[mediapipe.FixedSizeInputStreamHandlerOptions.ext] {
|
||||
trigger_queue_size: 4
|
||||
target_queue_size: 3
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
SetFixedMinSize(graph_config.mutable_node(0), GetParam());
|
||||
std::vector<Packet> output_packets[3];
|
||||
std::string in_streams[3];
|
||||
for (int i = 0; i < 3; ++i) {
|
||||
in_streams[i] = absl::StrCat("in_", i);
|
||||
tool::AddVectorSink(absl::StrCat("out_", i), &graph_config,
|
||||
&output_packets[i]);
|
||||
}
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(graph_config, {}));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
for (int i = 1; i <= 6; i++) {
|
||||
MEDIAPIPE_EXPECT_OK(graph.AddPacketToInputStream(
|
||||
in_streams[0], MakePacket<int>(i).At(Timestamp(i))));
|
||||
}
|
||||
for (int i = 3; i <= 7; i++) {
|
||||
MEDIAPIPE_EXPECT_OK(graph.AddPacketToInputStream(
|
||||
in_streams[1], MakePacket<int>(i).At(Timestamp(i))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
}
|
||||
// At this point everything before ts 5 should be dropped.
|
||||
for (int i = 4; i <= 7; i++) {
|
||||
MEDIAPIPE_EXPECT_OK(graph.AddPacketToInputStream(
|
||||
in_streams[2], MakePacket<int>(i).At(Timestamp(i))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
}
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
|
||||
if (GetParam()) {
|
||||
EXPECT_THAT(TimestampValues(output_packets[0]),
|
||||
testing::ContainerEq(std::vector<int64>{1, 2, 3, 4, 5, 6}));
|
||||
EXPECT_THAT(TimestampValues(output_packets[1]),
|
||||
testing::ContainerEq(std::vector<int64>{3, 4, 5, 6, 7}));
|
||||
EXPECT_THAT(TimestampValues(output_packets[2]),
|
||||
testing::ContainerEq(std::vector<int64>{4, 5, 6, 7}));
|
||||
} else {
|
||||
EXPECT_THAT(TimestampValues(output_packets[0]),
|
||||
testing::ContainerEq(std::vector<int64>{5, 6}));
|
||||
EXPECT_THAT(TimestampValues(output_packets[1]),
|
||||
testing::ContainerEq(std::vector<int64>{5, 6, 7}));
|
||||
EXPECT_THAT(TimestampValues(output_packets[2]),
|
||||
testing::ContainerEq(std::vector<int64>{5, 6, 7}));
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,124 @@
|
||||
// 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 <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// An input stream handler that delivers input packets to the Calculator
|
||||
// immediately, with no dependency between input streams. It also invokes
|
||||
// Calculator::Process when any input stream becomes done.
|
||||
//
|
||||
// NOTE: If packets arrive successively on different input streams with
|
||||
// identical or decreasing timestamps, this input stream handler will
|
||||
// invoke its Calculator with a sequence of InputTimestamps that is
|
||||
// non-increasing. Its Calculator is responsible for accumulating packets
|
||||
// with the required timetamps before processing and delivering output.
|
||||
//
|
||||
class ImmediateInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
ImmediateInputStreamHandler() = delete;
|
||||
ImmediateInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* calculator_context_manager,
|
||||
const MediaPipeOptions& options, bool calculator_run_in_parallel);
|
||||
|
||||
protected:
|
||||
// Returns kReadyForProcess whenever a Packet is available at any of
|
||||
// the input streams, or any input stream becomes done.
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
|
||||
|
||||
// Selects a packet on each stream with an available packet with the
|
||||
// specified timestamp, leaving other input streams unaffected.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
// Record of the last reported timestamp bound for each input stream.
|
||||
mediapipe::internal::Collection<Timestamp> timestamp_bounds_;
|
||||
};
|
||||
REGISTER_INPUT_STREAM_HANDLER(ImmediateInputStreamHandler);
|
||||
|
||||
ImmediateInputStreamHandler::ImmediateInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* calculator_context_manager,
|
||||
const MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(tag_map, calculator_context_manager, options,
|
||||
calculator_run_in_parallel),
|
||||
timestamp_bounds_(std::move(tag_map)) {}
|
||||
|
||||
NodeReadiness ImmediateInputStreamHandler::GetNodeReadiness(
|
||||
Timestamp* input_timestamp) {
|
||||
Timestamp min_stream_timestamp = Timestamp::Done();
|
||||
*input_timestamp = Timestamp::Done();
|
||||
bool stream_became_done = false;
|
||||
|
||||
for (CollectionItemId i = input_stream_managers_.BeginId();
|
||||
i < input_stream_managers_.EndId(); ++i) {
|
||||
const auto& stream = input_stream_managers_.Get(i);
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (!empty) {
|
||||
*input_timestamp = std::min(*input_timestamp, stream_timestamp);
|
||||
}
|
||||
min_stream_timestamp = std::min(min_stream_timestamp, stream_timestamp);
|
||||
if (stream_timestamp != timestamp_bounds_.Get(i)) {
|
||||
if (stream_timestamp == Timestamp::Done()) {
|
||||
stream_became_done = true;
|
||||
}
|
||||
timestamp_bounds_.Get(i) = stream_timestamp;
|
||||
}
|
||||
}
|
||||
|
||||
if (min_stream_timestamp == Timestamp::Done()) {
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
|
||||
if (*input_timestamp < Timestamp::Done()) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
if (stream_became_done) {
|
||||
*input_timestamp = min_stream_timestamp;
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
void ImmediateInputStreamHandler::FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
auto& stream = input_stream_managers_.Get(id);
|
||||
if (stream->QueueHead().Timestamp() == input_timestamp) {
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Packet current_packet = stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
} else {
|
||||
bool empty = false;
|
||||
bool is_done = stream->MinTimestampOrBound(&empty) == Timestamp::Done();
|
||||
AddPacketToShard(&input_set->Get(id), Packet(), is_done);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,343 @@
|
||||
// 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 <functional>
|
||||
#include <list>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/base/macros.h"
|
||||
#include "absl/memory/memory.h"
|
||||
#include "mediapipe/framework/calculator_context.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/tool/tag_map.h"
|
||||
#include "mediapipe/framework/tool/tag_map_helper.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
class ImmediateInputStreamHandlerTest : public ::testing::Test {
|
||||
protected:
|
||||
ImmediateInputStreamHandlerTest() {
|
||||
packet_type_.Set<std::string>();
|
||||
headers_ready_callback_ = [this]() {
|
||||
ImmediateInputStreamHandlerTest::HeadersReadyNoOp();
|
||||
};
|
||||
notification_callback_ = [this]() {
|
||||
ImmediateInputStreamHandlerTest::NotifyNoOp();
|
||||
};
|
||||
schedule_callback_ = std::bind(&ImmediateInputStreamHandlerTest::Schedule,
|
||||
this, std::placeholders::_1);
|
||||
error_callback_ = std::bind(&ImmediateInputStreamHandlerTest::RecordError,
|
||||
this, std::placeholders::_1);
|
||||
setup_shards_callback_ =
|
||||
std::bind(&ImmediateInputStreamHandlerTest::SetupShardsNoOp, this,
|
||||
std::placeholders::_1);
|
||||
queue_full_callback_ =
|
||||
std::bind(&ImmediateInputStreamHandlerTest::ReportQueueNoOp, this,
|
||||
std::placeholders::_1, std::placeholders::_2);
|
||||
queue_not_full_callback_ =
|
||||
std::bind(&ImmediateInputStreamHandlerTest::ReportQueueNoOp, this,
|
||||
std::placeholders::_1, std::placeholders::_2);
|
||||
|
||||
std::shared_ptr<tool::TagMap> input_tag_map =
|
||||
tool::CreateTagMap({"input_a", "input_b", "input_c"}).ValueOrDie();
|
||||
|
||||
input_stream_managers_.reset(
|
||||
new InputStreamManager[input_tag_map->NumEntries()]);
|
||||
const std::vector<std::string>& names = input_tag_map->Names();
|
||||
for (CollectionItemId id = input_tag_map->BeginId();
|
||||
id < input_tag_map->EndId(); ++id) {
|
||||
const std::string& stream_name = names[id.value()];
|
||||
name_to_id_[stream_name] = id;
|
||||
MEDIAPIPE_CHECK_OK(input_stream_managers_[id.value()].Initialize(
|
||||
stream_name, &packet_type_, /*back_edge=*/false));
|
||||
}
|
||||
SetupInputStreamHandler(input_tag_map);
|
||||
}
|
||||
|
||||
void SetupInputStreamHandler(
|
||||
const std::shared_ptr<tool::TagMap>& input_tag_map) {
|
||||
calculator_state_ = absl::make_unique<CalculatorState>(
|
||||
"Node", /*node_id=*/0, "Calculator", CalculatorGraphConfig::Node(),
|
||||
nullptr);
|
||||
cc_manager_.Initialize(
|
||||
calculator_state_.get(), input_tag_map,
|
||||
/*output_tag_map=*/tool::CreateTagMap({"output_a"}).ValueOrDie(),
|
||||
/*calculator_run_in_parallel=*/false);
|
||||
|
||||
mediapipe::StatusOr<std::unique_ptr<mediapipe::InputStreamHandler>>
|
||||
status_or_handler = InputStreamHandlerRegistry::CreateByName(
|
||||
"ImmediateInputStreamHandler", input_tag_map, &cc_manager_,
|
||||
MediaPipeOptions(),
|
||||
/*calculator_run_in_parallel=*/false);
|
||||
ASSERT_TRUE(status_or_handler.ok());
|
||||
input_stream_handler_ = std::move(status_or_handler.ValueOrDie());
|
||||
MEDIAPIPE_ASSERT_OK(input_stream_handler_->InitializeInputStreamManagers(
|
||||
input_stream_managers_.get()));
|
||||
MEDIAPIPE_ASSERT_OK(cc_manager_.PrepareForRun(setup_shards_callback_));
|
||||
input_stream_handler_->PrepareForRun(headers_ready_callback_,
|
||||
notification_callback_,
|
||||
schedule_callback_, error_callback_);
|
||||
input_stream_handler_->SetQueueSizeCallbacks(queue_full_callback_,
|
||||
queue_not_full_callback_);
|
||||
}
|
||||
|
||||
void HeadersReadyNoOp() {}
|
||||
|
||||
void NotifyNoOp() {}
|
||||
|
||||
void Schedule(CalculatorContext* cc) {
|
||||
CHECK(cc);
|
||||
cc_ = cc;
|
||||
}
|
||||
|
||||
void RecordError(const ::mediapipe::Status& error) {
|
||||
errors_.push_back(error);
|
||||
}
|
||||
|
||||
::mediapipe::Status SetupShardsNoOp(CalculatorContext* calculator_context) {
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
void ReportQueueNoOp(InputStreamManager* stream, bool* stream_was_full) {}
|
||||
|
||||
void ExpectPackets(
|
||||
const InputStreamShardSet& input_set,
|
||||
const std::map<std::string, std::string>& expected_values) {
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream = input_set.Get(name_and_id.second);
|
||||
if (::mediapipe::ContainsKey(expected_values, name_and_id.first)) {
|
||||
ASSERT_FALSE(input_stream.Value().IsEmpty());
|
||||
EXPECT_EQ(input_stream.Value().Get<std::string>(),
|
||||
::mediapipe::FindOrDie(expected_values, name_and_id.first));
|
||||
} else {
|
||||
EXPECT_TRUE(input_stream.Value().IsEmpty());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
PacketType packet_type_;
|
||||
std::function<void()> headers_ready_callback_;
|
||||
std::function<void()> notification_callback_;
|
||||
std::function<void(CalculatorContext*)> schedule_callback_;
|
||||
std::function<void(::mediapipe::Status)> error_callback_;
|
||||
std::function<::mediapipe::Status(CalculatorContext*)> setup_shards_callback_;
|
||||
InputStreamManager::QueueSizeCallback queue_full_callback_;
|
||||
InputStreamManager::QueueSizeCallback queue_not_full_callback_;
|
||||
|
||||
// Vector of errors encountered while using the stream.
|
||||
std::vector<::mediapipe::Status> errors_;
|
||||
|
||||
std::unique_ptr<CalculatorState> calculator_state_;
|
||||
CalculatorContextManager cc_manager_;
|
||||
CalculatorContext* cc_;
|
||||
std::map<std::string, CollectionItemId> name_to_id_;
|
||||
std::unique_ptr<InputStreamHandler> input_stream_handler_;
|
||||
std::unique_ptr<InputStreamManager[]> input_stream_managers_;
|
||||
};
|
||||
|
||||
// This test checks that a node is not considered ready for Process() if no
|
||||
// packets have arrived.
|
||||
TEST_F(ImmediateInputStreamHandlerTest, EmptyPacketsNotReady) {
|
||||
Timestamp min_stream_timestamp;
|
||||
ASSERT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
}
|
||||
|
||||
// This test checks that a node is considered ready for Process() if any of the
|
||||
// input streams has a packet available.
|
||||
TEST_F(ImmediateInputStreamHandlerTest, AnyPacketsReady) {
|
||||
Timestamp min_stream_timestamp;
|
||||
std::list<Packet> packets;
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(10)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_a"], packets);
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {{"input_a", "packet 1"}});
|
||||
}
|
||||
|
||||
// This test checks that a node is considered ready for Process() if any of the
|
||||
// input streams has become done.
|
||||
TEST_F(ImmediateInputStreamHandlerTest, StreamDoneReady) {
|
||||
Timestamp min_stream_timestamp;
|
||||
std::list<Packet> packets;
|
||||
|
||||
// One packet arrives, ready for process.
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(10)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_a"], packets);
|
||||
EXPECT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {{"input_a", "packet 1"}});
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
|
||||
// No packets arrive, not ready for process.
|
||||
EXPECT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
|
||||
// Timestamp::Done arrives, ready for process.
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_b"],
|
||||
Timestamp::Done());
|
||||
EXPECT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
|
||||
// No timestamp arrives, not ready for process.
|
||||
EXPECT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
|
||||
// Timestamp 200 arrives, not ready for process.
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_a"],
|
||||
Timestamp(200));
|
||||
EXPECT_FALSE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
|
||||
// Another Timestamp::Done arrives, ready for process.
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_a"],
|
||||
Timestamp::Done());
|
||||
EXPECT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
}
|
||||
|
||||
// This test checks that when any stream is done, the state is ready to close.
|
||||
TEST_F(ImmediateInputStreamHandlerTest, ReadyForClose) {
|
||||
Timestamp min_stream_timestamp;
|
||||
std::list<Packet> packets;
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(1)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_b"], packets);
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_b"],
|
||||
Timestamp::Done());
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {{"input_b", "packet 1"}});
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
|
||||
EXPECT_TRUE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_b"])
|
||||
->IsEmpty());
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_a"],
|
||||
Timestamp::Done());
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_c"],
|
||||
Timestamp::Done());
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
EXPECT_EQ(Timestamp::Done(), cc_->InputTimestamp());
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
EXPECT_TRUE(errors_.empty());
|
||||
}
|
||||
|
||||
// This test simulates how CalculatorNode::ProcessNode() uses an input
|
||||
// stream handler and the associated input streams.
|
||||
TEST_F(ImmediateInputStreamHandlerTest, SimulateProcessNode) {
|
||||
Timestamp min_stream_timestamp;
|
||||
std::list<Packet> packets;
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(10)));
|
||||
packets.push_back(Adopt(new std::string("packet 2")).At(Timestamp(30)));
|
||||
packets.push_back(Adopt(new std::string("packet 3")).At(Timestamp(40)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_a"], packets);
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(10)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_b"], packets);
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 1")).At(Timestamp(10)));
|
||||
packets.push_back(Adopt(new std::string("packet 2")).At(Timestamp(30)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_c"], packets);
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {{"input_a", "packet 1"},
|
||||
{"input_b", "packet 1"},
|
||||
{"input_c", "packet 1"}});
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
|
||||
for (const auto& name_and_id : name_to_id_) {
|
||||
const InputStream& input_stream = cc_->Inputs().Get(name_and_id.second);
|
||||
EXPECT_TRUE(input_stream.Value().IsEmpty());
|
||||
}
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 3")).At(Timestamp(40)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_c"], packets);
|
||||
|
||||
packets.clear();
|
||||
packets.push_back(Adopt(new std::string("packet 2")).At(Timestamp(30)));
|
||||
input_stream_handler_->AddPackets(name_to_id_["input_b"], packets);
|
||||
input_stream_handler_->SetNextTimestampBound(name_to_id_["input_b"],
|
||||
Timestamp::Done());
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(), {{"input_a", "packet 2"},
|
||||
{"input_b", "packet 2"},
|
||||
{"input_c", "packet 2"}});
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
|
||||
EXPECT_FALSE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_a"])
|
||||
->IsEmpty());
|
||||
EXPECT_TRUE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_b"])
|
||||
->IsEmpty());
|
||||
EXPECT_FALSE(
|
||||
input_stream_handler_->GetInputStreamManager(name_to_id_["input_c"])
|
||||
->IsEmpty());
|
||||
ASSERT_TRUE(input_stream_handler_->ScheduleInvocations(
|
||||
/*max_allowance=*/1, &min_stream_timestamp));
|
||||
ExpectPackets(cc_->Inputs(),
|
||||
{{"input_a", "packet 3"}, {"input_c", "packet 3"}});
|
||||
// FinalizeInputSet() is a no-op.
|
||||
input_stream_handler_->FinalizeInputSet(cc_->InputTimestamp(),
|
||||
&cc_->Inputs());
|
||||
input_stream_handler_->ClearCurrentInputs(cc_);
|
||||
|
||||
ExpectPackets(cc_->Inputs(), {});
|
||||
|
||||
EXPECT_TRUE(errors_.empty());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,133 @@
|
||||
// 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 "mediapipe/framework/stream_handler/in_order_output_stream_handler.h"
|
||||
|
||||
#include "mediapipe/framework/collection.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/output_stream_shard.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
REGISTER_OUTPUT_STREAM_HANDLER(InOrderOutputStreamHandler);
|
||||
|
||||
void InOrderOutputStreamHandler::PropagationLoop() {
|
||||
CHECK_EQ(propagation_state_, kIdle);
|
||||
Timestamp context_timestamp;
|
||||
CalculatorContext* calculator_context;
|
||||
if (!calculator_context_manager_->HasActiveContexts()) {
|
||||
propagation_state_ = kPropagatingBound;
|
||||
} else {
|
||||
calculator_context = calculator_context_manager_->GetFrontCalculatorContext(
|
||||
&context_timestamp);
|
||||
if (!completed_input_timestamps_.empty()) {
|
||||
Timestamp completed_timestamp = *completed_input_timestamps_.begin();
|
||||
if (context_timestamp != completed_timestamp) {
|
||||
CHECK_LT(context_timestamp, completed_timestamp);
|
||||
return;
|
||||
}
|
||||
propagation_state_ = kPropagatingPackets;
|
||||
}
|
||||
}
|
||||
|
||||
while (propagation_state_ != kIdle) {
|
||||
if (propagation_state_ == kPropagatingPackets) {
|
||||
PropagatePackets(&calculator_context, &context_timestamp);
|
||||
} else {
|
||||
CHECK_EQ(kPropagatingBound, propagation_state_);
|
||||
PropagationBound(&calculator_context, &context_timestamp);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void InOrderOutputStreamHandler::PropagatePackets(
|
||||
CalculatorContext** calculator_context, Timestamp* context_timestamp) {
|
||||
timestamp_mutex_.Unlock();
|
||||
// Propagates packets without holding timestamp_mutex_.
|
||||
PropagateOutputPackets(*context_timestamp, &(*calculator_context)->Outputs());
|
||||
calculator_context_manager_->RecycleCalculatorContext();
|
||||
timestamp_mutex_.Lock();
|
||||
completed_input_timestamps_.erase(completed_input_timestamps_.begin());
|
||||
// The first check is for performance reasons (it's cheaper).
|
||||
// Note that completed_input_timestamps_ is a subset of the input
|
||||
// timestamps of the active contexts. Therefore, the second check
|
||||
// covers the first check.
|
||||
if (completed_input_timestamps_.empty() &&
|
||||
!calculator_context_manager_->HasActiveContexts()) {
|
||||
// If task_timestamp_bound_ is not greater than context_timestamp + 1,
|
||||
// timestamp propagation isn't necessary since the bound of the
|
||||
// downstream input streams has been updated to a larger value
|
||||
// already. Timestamp propagation will be skipped, and the
|
||||
// propagation process is completed.
|
||||
if (task_timestamp_bound_ <= context_timestamp->NextAllowedInStream()) {
|
||||
propagation_state_ = kIdle;
|
||||
return;
|
||||
}
|
||||
propagation_state_ = kPropagatingBound;
|
||||
return;
|
||||
}
|
||||
*calculator_context =
|
||||
calculator_context_manager_->GetFrontCalculatorContext(context_timestamp);
|
||||
if (!completed_input_timestamps_.empty() &&
|
||||
*context_timestamp == *completed_input_timestamps_.begin()) {
|
||||
// Continues propagating output packets if the smallest completed
|
||||
// input timestamp is equal to the input timestamp of the earliest
|
||||
// active calculator context.
|
||||
return;
|
||||
}
|
||||
propagation_state_ = kIdle;
|
||||
}
|
||||
|
||||
void InOrderOutputStreamHandler::PropagationBound(
|
||||
CalculatorContext** calculator_context, Timestamp* context_timestamp) {
|
||||
Timestamp bound_to_propagate = task_timestamp_bound_;
|
||||
timestamp_mutex_.Unlock();
|
||||
// Timestamp bound propagation without holding timestamp_mutex_.
|
||||
TryPropagateTimestampBound(bound_to_propagate);
|
||||
timestamp_mutex_.Lock();
|
||||
if (propagation_state_ == kPropagatingBound) {
|
||||
// There is no invocation completed and no newly arrived timestamp
|
||||
// bound during the timestamp bound propagation. So the propagation
|
||||
// process is completed.
|
||||
propagation_state_ = kIdle;
|
||||
return;
|
||||
}
|
||||
// Some recent changes require the propagation thread to recheck if any
|
||||
// new packets can be propagated.
|
||||
CHECK_EQ(propagation_state_, kPropagationPending);
|
||||
// task_timestamp_bound_ was updated while the propagation thread was
|
||||
// doing timestamp propagation. This thread will redo timestamp
|
||||
// propagation for the new task_timestamp_bound_.
|
||||
if (!calculator_context_manager_->HasActiveContexts()) {
|
||||
CHECK_LT(bound_to_propagate, task_timestamp_bound_);
|
||||
propagation_state_ = kPropagatingBound;
|
||||
return;
|
||||
}
|
||||
*calculator_context =
|
||||
calculator_context_manager_->GetFrontCalculatorContext(context_timestamp);
|
||||
if (completed_input_timestamps_.empty() ||
|
||||
*context_timestamp != *completed_input_timestamps_.begin()) {
|
||||
// If there is no newly completed invocation or the newly arrived packets
|
||||
// are not ready for propagation, the propagation process is completed.
|
||||
propagation_state_ = kIdle;
|
||||
return;
|
||||
} else {
|
||||
// Found new packets to be propagated, and will redo packets
|
||||
// propagation.
|
||||
propagation_state_ = kPropagatingPackets;
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,55 @@
|
||||
// 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_FRAMEWORK_STREAM_HANDLER_IN_ORDER_OUTPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IN_ORDER_OUTPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
// TODO: Move protos in another CL after the C++ code migration.
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/output_stream_handler.h"
|
||||
#include "mediapipe/framework/packet_set.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
#include "mediapipe/framework/tool/tag_map.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// InOrderOutputStreamHandler supports both sequential and parallel processing
|
||||
// of input packets, and will deliver the output packets in increasing timestamp
|
||||
// order.
|
||||
class InOrderOutputStreamHandler : public OutputStreamHandler {
|
||||
public:
|
||||
InOrderOutputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* calculator_context_manager,
|
||||
const MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: OutputStreamHandler(std::move(tag_map), calculator_context_manager,
|
||||
options, calculator_run_in_parallel) {}
|
||||
|
||||
private:
|
||||
void PropagationLoop() EXCLUSIVE_LOCKS_REQUIRED(timestamp_mutex_) final;
|
||||
|
||||
void PropagatePackets(CalculatorContext** calculator_context,
|
||||
Timestamp* context_timestamp)
|
||||
EXCLUSIVE_LOCKS_REQUIRED(timestamp_mutex_);
|
||||
|
||||
void PropagationBound(CalculatorContext** calculator_context,
|
||||
Timestamp* context_timestamp)
|
||||
EXCLUSIVE_LOCKS_REQUIRED(timestamp_mutex_);
|
||||
};
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IN_ORDER_OUTPUT_STREAM_HANDLER_H_
|
||||
@@ -0,0 +1,134 @@
|
||||
// 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 "absl/strings/substitute.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/port/logging.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// Implementation of the input stream handler for the MuxCalculator.
|
||||
//
|
||||
// One of the input streams is the control stream; all the other input streams
|
||||
// are data streams. To make MuxInputStreamHandler work properly, the tag of the
|
||||
// input streams must obey the following rules:
|
||||
// Let N be the number of input streams. Data streams must use tag "INPUT" with
|
||||
// index 0, ..., N - 2; the control stream must use tag "SELECT".
|
||||
//
|
||||
// The control stream carries packets of type 'int'. The 'int' value in a
|
||||
// control stream packet must be a valid index in the range 0, ..., N - 2 and
|
||||
// select the data stream at that index. The selected data stream must have a
|
||||
// packet with the same timestamp as the control stream packet.
|
||||
//
|
||||
// When the control stream is done, GetNodeReadiness() returns
|
||||
// NodeReadiness::kReadyForClose.
|
||||
//
|
||||
// TODO: pass the input stream tags to the MuxInputStreamHandler
|
||||
// constructor so that it can refer to input streams by tag. See b/30125118.
|
||||
class MuxInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
MuxInputStreamHandler() = delete;
|
||||
MuxInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options,
|
||||
bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), cc_manager, options,
|
||||
calculator_run_in_parallel) {}
|
||||
|
||||
protected:
|
||||
// In MuxInputStreamHandler, a node is "ready" if:
|
||||
// - the control stream is done (need to call Close() in this case), or
|
||||
// - we have received the packets on the control stream and the selected data
|
||||
// stream at the next timestamp.
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
|
||||
DCHECK(min_stream_timestamp);
|
||||
absl::MutexLock lock(&input_streams_mutex_);
|
||||
|
||||
const auto& control_stream =
|
||||
input_stream_managers_.Get(input_stream_managers_.EndId() - 1);
|
||||
bool empty;
|
||||
*min_stream_timestamp = control_stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
if (*min_stream_timestamp == Timestamp::Done()) {
|
||||
// Calculator is done if the control input stream is done.
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
// Calculator is not ready to run if the control input stream is empty.
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
Packet control_packet = control_stream->QueueHead();
|
||||
CHECK(!control_packet.IsEmpty());
|
||||
int control_value = control_packet.Get<int>();
|
||||
CHECK_LE(0, control_value);
|
||||
CHECK_LT(control_value, input_stream_managers_.NumEntries() - 1);
|
||||
|
||||
const auto& data_stream = input_stream_managers_.Get(
|
||||
input_stream_managers_.BeginId() + control_value);
|
||||
Timestamp stream_timestamp = data_stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
CHECK_LE(stream_timestamp, *min_stream_timestamp);
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
CHECK_EQ(stream_timestamp, *min_stream_timestamp);
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
absl::MutexLock lock(&input_streams_mutex_);
|
||||
|
||||
const CollectionItemId control_stream_id =
|
||||
input_stream_managers_.EndId() - 1;
|
||||
auto& control_stream = input_stream_managers_.Get(control_stream_id);
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Packet control_packet = control_stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
CHECK_EQ(num_packets_dropped, 0)
|
||||
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, control_stream->Name());
|
||||
CHECK(!control_packet.IsEmpty());
|
||||
int control_value = control_packet.Get<int>();
|
||||
AddPacketToShard(&input_set->Get(control_stream_id),
|
||||
std::move(control_packet), stream_is_done);
|
||||
|
||||
const CollectionItemId data_stream_id =
|
||||
input_stream_managers_.BeginId() + control_value;
|
||||
CHECK_LE(input_stream_managers_.BeginId(), data_stream_id);
|
||||
CHECK_LT(data_stream_id, control_stream_id);
|
||||
auto& data_stream = input_stream_managers_.Get(data_stream_id);
|
||||
stream_is_done = false;
|
||||
Packet data_packet = data_stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
CHECK_EQ(num_packets_dropped, 0)
|
||||
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, data_stream->Name());
|
||||
AddPacketToShard(&input_set->Get(data_stream_id), std::move(data_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
|
||||
private:
|
||||
// Must be acquired when manipulating the control and data streams to ensure
|
||||
// we have a consistent view of the two streams.
|
||||
absl::Mutex input_streams_mutex_;
|
||||
};
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(MuxInputStreamHandler);
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,91 @@
|
||||
// Copyright 2018 The MediaPipe Authors.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
// A regression test for b/31620439. MuxInputStreamHandler's accesses to the
|
||||
// control and data streams should be atomic so that it has a consistent view
|
||||
// of the two streams. None of the CHECKs in the GetNodeReadiness() method of
|
||||
// MuxInputStreamHandler should fail when running this test.
|
||||
TEST(MuxInputStreamHandlerTest, AtomicAccessToControlAndDataStreams) {
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input"
|
||||
node {
|
||||
calculator: "RoundRobinDemuxCalculator"
|
||||
input_stream: "input"
|
||||
output_stream: "OUTPUT:0:input0"
|
||||
output_stream: "OUTPUT:1:input1"
|
||||
output_stream: "OUTPUT:2:input2"
|
||||
output_stream: "OUTPUT:3:input3"
|
||||
output_stream: "OUTPUT:4:input4"
|
||||
output_stream: "SELECT:select"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input0"
|
||||
output_stream: "output0"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input1"
|
||||
output_stream: "output1"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input2"
|
||||
output_stream: "output2"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input3"
|
||||
output_stream: "output3"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input4"
|
||||
output_stream: "output4"
|
||||
}
|
||||
node {
|
||||
calculator: "MuxCalculator"
|
||||
input_stream: "INPUT:0:output0"
|
||||
input_stream: "INPUT:1:output1"
|
||||
input_stream: "INPUT:2:output2"
|
||||
input_stream: "INPUT:3:output3"
|
||||
input_stream: "INPUT:4:output4"
|
||||
input_stream: "SELECT:select"
|
||||
output_stream: "OUTPUT:output"
|
||||
input_stream_handler { input_stream_handler: "MuxInputStreamHandler" }
|
||||
})");
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
for (int i = 0; i < 2000; ++i) {
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(new int(i)).At(Timestamp(i))));
|
||||
}
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,253 @@
|
||||
// Copyright 2018 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 <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/ret_check.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/port/threadpool.h"
|
||||
#include "mediapipe/framework/stream_handler/fixed_size_input_stream_handler.pb.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
// Copied from mux_input_stream_handler_test.cc, and removed ISH from the graph
|
||||
//
|
||||
// A regression test for b/31620439. MuxInputStreamHandler's accesses to the
|
||||
// control and data streams should be atomic so that it has a consistent view
|
||||
// of the two streams. None of the CHECKs in the GetNodeReadiness() method of
|
||||
// MuxInputStreamHandler should fail when running this test.
|
||||
TEST(MuxInputStreamHandlerTest, AtomicAccessToControlAndDataStreams) {
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input"
|
||||
node {
|
||||
calculator: "RoundRobinDemuxCalculator"
|
||||
input_stream: "input"
|
||||
output_stream: "OUTPUT:0:input0"
|
||||
output_stream: "OUTPUT:1:input1"
|
||||
output_stream: "OUTPUT:2:input2"
|
||||
output_stream: "OUTPUT:3:input3"
|
||||
output_stream: "OUTPUT:4:input4"
|
||||
output_stream: "SELECT:select"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input0"
|
||||
output_stream: "output0"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input1"
|
||||
output_stream: "output1"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input2"
|
||||
output_stream: "output2"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input3"
|
||||
output_stream: "output3"
|
||||
}
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "input4"
|
||||
output_stream: "output4"
|
||||
}
|
||||
node {
|
||||
calculator: "MuxCalculator"
|
||||
input_stream: "INPUT:0:output0"
|
||||
input_stream: "INPUT:1:output1"
|
||||
input_stream: "INPUT:2:output2"
|
||||
input_stream: "INPUT:3:output3"
|
||||
input_stream: "INPUT:4:output4"
|
||||
input_stream: "SELECT:select"
|
||||
output_stream: "OUTPUT:output"
|
||||
# MuxInputStreamHandler set in GetContract().
|
||||
})");
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
for (int i = 0; i < 2000; ++i) {
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input", Adopt(new int(i)).At(Timestamp(i))));
|
||||
}
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
|
||||
// Copied from pass_through_calculator.cc, and modified to specify
|
||||
// InputStreamHandler.
|
||||
//
|
||||
// A Calculator that simply passes its input Packets and header through,
|
||||
// unchanged. The inputs may be specified by tag or index. The outputs
|
||||
// must match the inputs exactly. Any number of input side packets may
|
||||
// also be specified. If output side packets are specified, they must
|
||||
// match the input side packets exactly and the Calculator passes its
|
||||
// input side packets through, unchanged. Otherwise, the input side
|
||||
// packets will be ignored (allowing PassThroughCalculator to be used to
|
||||
// test internal behavior). Any options may be specified and will be
|
||||
// ignored.
|
||||
class FixedPassThroughCalculator : public CalculatorBase {
|
||||
public:
|
||||
static ::mediapipe::Status GetContract(CalculatorContract* cc) {
|
||||
if (!cc->Inputs().TagMap()->SameAs(*cc->Outputs().TagMap())) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Input and output streams to PassThroughCalculator must use "
|
||||
"matching tags and indexes.");
|
||||
}
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
cc->Inputs().Get(id).SetAny();
|
||||
cc->Outputs().Get(id).SetSameAs(&cc->Inputs().Get(id));
|
||||
}
|
||||
for (CollectionItemId id = cc->InputSidePackets().BeginId();
|
||||
id < cc->InputSidePackets().EndId(); ++id) {
|
||||
cc->InputSidePackets().Get(id).SetAny();
|
||||
}
|
||||
if (cc->OutputSidePackets().NumEntries() != 0) {
|
||||
if (!cc->InputSidePackets().TagMap()->SameAs(
|
||||
*cc->OutputSidePackets().TagMap())) {
|
||||
return ::mediapipe::InvalidArgumentError(
|
||||
"Input and output side packets to PassThroughCalculator must use "
|
||||
"matching tags and indexes.");
|
||||
}
|
||||
for (CollectionItemId id = cc->InputSidePackets().BeginId();
|
||||
id < cc->InputSidePackets().EndId(); ++id) {
|
||||
cc->OutputSidePackets().Get(id).SetSameAs(
|
||||
&cc->InputSidePackets().Get(id));
|
||||
}
|
||||
}
|
||||
|
||||
// Assign this calculator's InputStreamHandler and options.
|
||||
cc->SetInputStreamHandler("FixedSizeInputStreamHandler");
|
||||
MediaPipeOptions options;
|
||||
options.MutableExtension(FixedSizeInputStreamHandlerOptions::ext)
|
||||
->set_fixed_min_size(2);
|
||||
options.MutableExtension(FixedSizeInputStreamHandlerOptions::ext)
|
||||
->set_trigger_queue_size(2);
|
||||
options.MutableExtension(FixedSizeInputStreamHandlerOptions::ext)
|
||||
->set_target_queue_size(2);
|
||||
cc->SetInputStreamHandlerOptions(options);
|
||||
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Open(CalculatorContext* cc) final {
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
if (!cc->Inputs().Get(id).Header().IsEmpty()) {
|
||||
cc->Outputs().Get(id).SetHeader(cc->Inputs().Get(id).Header());
|
||||
}
|
||||
}
|
||||
if (cc->OutputSidePackets().NumEntries() != 0) {
|
||||
for (CollectionItemId id = cc->InputSidePackets().BeginId();
|
||||
id < cc->InputSidePackets().EndId(); ++id) {
|
||||
cc->OutputSidePackets().Get(id).Set(cc->InputSidePackets().Get(id));
|
||||
}
|
||||
}
|
||||
cc->SetOffset(TimestampDiff(0));
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
::mediapipe::Status Process(CalculatorContext* cc) final {
|
||||
cc->GetCounter("PassThrough")->Increment();
|
||||
if (cc->Inputs().NumEntries() == 0) {
|
||||
return tool::StatusStop();
|
||||
}
|
||||
for (CollectionItemId id = cc->Inputs().BeginId();
|
||||
id < cc->Inputs().EndId(); ++id) {
|
||||
if (!cc->Inputs().Get(id).IsEmpty()) {
|
||||
VLOG(3) << "Passing " << cc->Inputs().Get(id).Name() << " to "
|
||||
<< cc->Outputs().Get(id).Name() << " at "
|
||||
<< cc->InputTimestamp().DebugString();
|
||||
cc->Outputs().Get(id).AddPacket(cc->Inputs().Get(id).Value());
|
||||
}
|
||||
}
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
};
|
||||
REGISTER_CALCULATOR(FixedPassThroughCalculator);
|
||||
|
||||
// Copied from fixed_size_input_stream_handler_test.cc, and modified graph.
|
||||
//
|
||||
// Tests FixedSizeInputStreamHandler with several input streams running
|
||||
// asynchronously in parallel.
|
||||
TEST(FixedSizeInputStreamHandlerTest, ParallelWriteAndRead) {
|
||||
#define NUM_STREAMS 4
|
||||
CalculatorGraphConfig graph_config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "in_0"
|
||||
input_stream: "in_1"
|
||||
input_stream: "in_2"
|
||||
input_stream: "in_3"
|
||||
node {
|
||||
calculator: "FixedPassThroughCalculator"
|
||||
input_stream: "in_0"
|
||||
input_stream: "in_1"
|
||||
input_stream: "in_2"
|
||||
input_stream: "in_3"
|
||||
output_stream: "out_0"
|
||||
output_stream: "out_1"
|
||||
output_stream: "out_2"
|
||||
output_stream: "out_3"
|
||||
# FixedSizeInputStreamHandler set in GetContract()
|
||||
})");
|
||||
std::vector<Packet> output_packets[NUM_STREAMS];
|
||||
for (int i = 0; i < NUM_STREAMS; ++i) {
|
||||
tool::AddVectorSink(absl::StrCat("out_", i), &graph_config,
|
||||
&output_packets[i]);
|
||||
}
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(graph_config, {}));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
{
|
||||
::mediapipe::ThreadPool pool(NUM_STREAMS);
|
||||
pool.StartWorkers();
|
||||
|
||||
// Start writers.
|
||||
for (int w = 0; w < NUM_STREAMS; ++w) {
|
||||
pool.Schedule([&, w]() {
|
||||
std::string stream_name = absl::StrCat("in_", w);
|
||||
for (int i = 0; i < 50; ++i) {
|
||||
Packet p = MakePacket<int>(i).At(Timestamp(i));
|
||||
MEDIAPIPE_EXPECT_OK(graph.AddPacketToInputStream(stream_name, p));
|
||||
absl::SleepFor(absl::Microseconds(100));
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
for (int i = 0; i < NUM_STREAMS; ++i) {
|
||||
EXPECT_EQ(output_packets[i].size(), output_packets[0].size());
|
||||
for (int j = 0; j < output_packets[i].size(); j++) {
|
||||
EXPECT_EQ(output_packets[i][j].Get<int>(),
|
||||
output_packets[0][j].Get<int>());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,209 @@
|
||||
// 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 <algorithm>
|
||||
|
||||
// TODO: Move protos in another CL after the C++ code migration.
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/packet_set.h"
|
||||
#include "mediapipe/framework/stream_handler/sync_set_input_stream_handler.pb.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
#include "mediapipe/framework/tool/tag_map.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// An input stream handler which separates the inputs into sets which
|
||||
// are each independently synchronized. For example, if 5 inputs are
|
||||
// present, then the first three can be grouped (and will be synchronized
|
||||
// as if they were in a calculator with only those three streams) and the
|
||||
// remaining 2 streams can be independently grouped. The calculator will
|
||||
// always be called with all the available packets from a single sync set
|
||||
// (never more than one). The input timestamps seen by the calculator
|
||||
// will be ordered sequentially for each sync set but may jump around
|
||||
// between sync sets.
|
||||
class SyncSetInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
SyncSetInputStreamHandler() = delete;
|
||||
SyncSetInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& extendable_options,
|
||||
bool calculator_run_in_parallel);
|
||||
|
||||
void PrepareForRun(
|
||||
std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(::mediapipe::Status)> error_callback) override;
|
||||
|
||||
protected:
|
||||
// In SyncSetInputStreamHandler, a node is "ready" if any
|
||||
// of its sync sets are ready in the traditional sense (See
|
||||
// DefaultInputStreamHandler).
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
private:
|
||||
absl::Mutex mutex_;
|
||||
// The ids of each set of inputs.
|
||||
std::vector<std::vector<CollectionItemId>> sync_sets_ GUARDED_BY(mutex_);
|
||||
// The index of the ready sync set. A value of -1 indicates that no
|
||||
// sync sets are ready.
|
||||
int ready_sync_set_index_ GUARDED_BY(mutex_) = -1;
|
||||
// The timestamp at which the sync set is ready. If no sync set is
|
||||
// ready then this variable should be Timestamp::Done() .
|
||||
Timestamp ready_timestamp_ GUARDED_BY(mutex_);
|
||||
};
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(SyncSetInputStreamHandler);
|
||||
|
||||
SyncSetInputStreamHandler::SyncSetInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& extendable_options, bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), cc_manager, extendable_options,
|
||||
calculator_run_in_parallel) {}
|
||||
|
||||
void SyncSetInputStreamHandler::PrepareForRun(
|
||||
std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(::mediapipe::Status)> error_callback) {
|
||||
const auto& handler_options =
|
||||
options_.GetExtension(SyncSetInputStreamHandlerOptions::ext);
|
||||
{
|
||||
absl::MutexLock lock(&mutex_);
|
||||
sync_sets_.clear();
|
||||
std::set<CollectionItemId> used_ids;
|
||||
for (const auto& sync_set : handler_options.sync_set()) {
|
||||
sync_sets_.emplace_back();
|
||||
CHECK_LT(0, sync_set.tag_index_size());
|
||||
for (const auto& tag_index : sync_set.tag_index()) {
|
||||
std::string tag;
|
||||
int index;
|
||||
MEDIAPIPE_CHECK_OK(tool::ParseTagIndex(tag_index, &tag, &index));
|
||||
CollectionItemId id = input_stream_managers_.GetId(tag, index);
|
||||
CHECK(id.IsValid()) << "stream \"" << tag_index << "\" is not found.";
|
||||
CHECK(!::mediapipe::ContainsKey(used_ids, id))
|
||||
<< "stream \"" << tag_index << "\" is in more than one sync set.";
|
||||
used_ids.insert(id);
|
||||
sync_sets_.back().push_back(id);
|
||||
}
|
||||
}
|
||||
std::vector<CollectionItemId> remaining_ids;
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
if (!::mediapipe::ContainsKey(used_ids, id)) {
|
||||
remaining_ids.push_back(id);
|
||||
}
|
||||
}
|
||||
if (!remaining_ids.empty()) {
|
||||
sync_sets_.push_back(std::move(remaining_ids));
|
||||
}
|
||||
ready_sync_set_index_ = -1;
|
||||
ready_timestamp_ = Timestamp::Done();
|
||||
}
|
||||
|
||||
InputStreamHandler::PrepareForRun(
|
||||
std::move(headers_ready_callback), std::move(notification_callback),
|
||||
std::move(schedule_callback), std::move(error_callback));
|
||||
}
|
||||
|
||||
NodeReadiness SyncSetInputStreamHandler::GetNodeReadiness(
|
||||
Timestamp* min_stream_timestamp) {
|
||||
DCHECK(min_stream_timestamp);
|
||||
absl::MutexLock lock(&mutex_);
|
||||
if (ready_sync_set_index_ >= 0) {
|
||||
*min_stream_timestamp = ready_timestamp_;
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
for (int sync_set_index = 0; sync_set_index < sync_sets_.size();
|
||||
++sync_set_index) {
|
||||
const std::vector<CollectionItemId>& sync_set = sync_sets_[sync_set_index];
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (CollectionItemId id : sync_set) {
|
||||
const auto& stream = input_stream_managers_.Get(id);
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
min_bound = std::min(min_bound, stream_timestamp);
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
if (*min_stream_timestamp == Timestamp::Done()) {
|
||||
// This sync set is done, remove it. Note that this invalidates
|
||||
// sync set indexes higher than sync_set_index. However, we are
|
||||
// guaranteed that we were not ready before entering the outer
|
||||
// loop, so even if we are ready now, ready_sync_set_index_ must
|
||||
// be less than the current value of sync_set_index.
|
||||
sync_sets_.erase(sync_sets_.begin() + sync_set_index);
|
||||
--sync_set_index;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (min_bound > *min_stream_timestamp) {
|
||||
if (*min_stream_timestamp < ready_timestamp_) {
|
||||
// Store the timestamp and corresponding sync set index for the
|
||||
// sync set with the earliest arrival timestamp.
|
||||
ready_timestamp_ = *min_stream_timestamp;
|
||||
ready_sync_set_index_ = sync_set_index;
|
||||
}
|
||||
} else {
|
||||
CHECK_EQ(min_bound, *min_stream_timestamp);
|
||||
}
|
||||
}
|
||||
if (ready_sync_set_index_ >= 0) {
|
||||
*min_stream_timestamp = ready_timestamp_;
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
if (sync_sets_.empty()) {
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
// TODO The value of *min_stream_timestamp is undefined in this case.
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
void SyncSetInputStreamHandler::FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) {
|
||||
// Assume that all current packets are already cleared.
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
absl::MutexLock lock(&mutex_);
|
||||
CHECK_LE(0, ready_sync_set_index_);
|
||||
CHECK_EQ(input_timestamp, ready_timestamp_);
|
||||
// Set the input streams for the ready sync set.
|
||||
for (CollectionItemId id : sync_sets_[ready_sync_set_index_]) {
|
||||
const auto& stream = input_stream_managers_.Get(id);
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Packet current_packet = stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
CHECK_EQ(num_packets_dropped, 0)
|
||||
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, stream->Name());
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
ready_sync_set_index_ = -1;
|
||||
ready_timestamp_ = Timestamp::Done();
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,50 @@
|
||||
// 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.
|
||||
//
|
||||
// Forked from
|
||||
// mediapipe/framework/stream_handler/sync_set_input_stream_handler.proto
|
||||
// The forked proto must remain identical to the original proto and should be
|
||||
// ONLY used by mediapipe open source project.
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/mediapipe_options.proto";
|
||||
|
||||
message SyncSetInputStreamHandlerOptions {
|
||||
extend MediaPipeOptions {
|
||||
optional SyncSetInputStreamHandlerOptions ext = 155816273;
|
||||
}
|
||||
message SyncSet {
|
||||
// A description of the streams which will be synchronized together.
|
||||
// This description uses the Calculator visible specification of
|
||||
// a stream. The format is a tag, then an index with both being
|
||||
// optional. If the tag is missing it is assumed to be "" and if
|
||||
// the index is missing then it is assumed to be 0. If the index
|
||||
// is provided then a colon (':') must be used.
|
||||
// Examples:
|
||||
// "TAG" -> tag "TAG", index 0
|
||||
// "" -> tag "", index 0
|
||||
// ":0" -> tag "", index 0
|
||||
// ":3" -> tag "", index 3
|
||||
// "VIDEO:0" -> tag "VIDEO", index 0
|
||||
// "VIDEO:2" -> tag "VIDEO", index 2
|
||||
repeated string tag_index = 1;
|
||||
}
|
||||
// Each synchronization set describes a collection of inputs which
|
||||
// must be provided together to the calculator. Any streams which
|
||||
// are not in any sync_set will be grouped into a (default) sync set.
|
||||
repeated SyncSet sync_set = 1;
|
||||
}
|
||||
@@ -0,0 +1,327 @@
|
||||
// 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 <memory>
|
||||
#include <random>
|
||||
#include <tuple>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
// TODO: Move protos in another CL after the C++ code migration.
|
||||
#include "absl/base/macros.h"
|
||||
#include "absl/memory/memory.h"
|
||||
#include "absl/strings/str_cat.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
#include "mediapipe/framework/stream_handler/sync_set_input_stream_handler.pb.h"
|
||||
|
||||
using RandomEngine = std::mt19937_64;
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
// The type LambdaCalculator takes.
|
||||
typedef std::function<::mediapipe::Status(const InputStreamShardSet&,
|
||||
OutputStreamShardSet*)>
|
||||
ProcessFunction;
|
||||
|
||||
// Helper function to create a tuple (inside an initializer list).
|
||||
std::tuple<std::string, Timestamp, std::vector<std::string>> CommandTuple(
|
||||
std::string stream, Timestamp timestamp,
|
||||
std::vector<std::string> expected) {
|
||||
return std::make_tuple(stream, timestamp, expected);
|
||||
}
|
||||
|
||||
// Function to take the inputs and produce a diagnostic output std::string
|
||||
// and output a packet with a diagnostic output std::string which includes
|
||||
// the input timestamp and the ids of each input which is present.
|
||||
::mediapipe::Status InputsToDebugString(const InputStreamShardSet& inputs,
|
||||
OutputStreamShardSet* outputs) {
|
||||
std::string output;
|
||||
Timestamp output_timestamp;
|
||||
for (CollectionItemId id = inputs.BeginId(); id < inputs.EndId(); ++id) {
|
||||
if (!inputs.Get(id).IsEmpty()) {
|
||||
if (output.empty()) {
|
||||
output_timestamp = inputs.Get(id).Value().Timestamp();
|
||||
if (output_timestamp.IsSpecialValue()) {
|
||||
output = output_timestamp.DebugString();
|
||||
} else {
|
||||
output =
|
||||
absl::StrCat("Timestamp(", output_timestamp.DebugString(), ")");
|
||||
}
|
||||
}
|
||||
absl::StrAppend(&output, ",", id.value());
|
||||
}
|
||||
}
|
||||
Packet output_packet;
|
||||
ABSL_CONST_INIT static absl::Mutex mu(absl::kConstInit);
|
||||
static Timestamp static_timestamp = Timestamp(0);
|
||||
{
|
||||
absl::MutexLock lock(&mu);
|
||||
output_packet = MakePacket<std::string>(output).At(static_timestamp);
|
||||
++static_timestamp;
|
||||
}
|
||||
// TODO Output at output_timestamp once unordered output stream
|
||||
// handlers are allowed.
|
||||
outputs->Index(0).AddPacket(output_packet);
|
||||
return ::mediapipe::OkStatus();
|
||||
}
|
||||
|
||||
TEST(SyncSetInputStreamHandlerTest, OrdinaryOperation) {
|
||||
CalculatorGraphConfig config = ParseTextProtoOrDie<CalculatorGraphConfig>(
|
||||
R"(
|
||||
input_stream: "a"
|
||||
input_stream: "b"
|
||||
input_stream: "c"
|
||||
input_stream: "d"
|
||||
input_stream: "e"
|
||||
input_stream: "f"
|
||||
input_stream: "g"
|
||||
input_stream: "h"
|
||||
node {
|
||||
calculator: "LambdaCalculator"
|
||||
input_stream: "a"
|
||||
input_stream: "b"
|
||||
input_stream: "c"
|
||||
input_stream: "d"
|
||||
input_stream: "e"
|
||||
input_stream: "f"
|
||||
input_stream: "g"
|
||||
input_stream: "h"
|
||||
output_stream: "output"
|
||||
input_side_packet: "lambda"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "SyncSetInputStreamHandler"
|
||||
options {
|
||||
[mediapipe.SyncSetInputStreamHandlerOptions.ext] {
|
||||
sync_set {
|
||||
# TODO Update this to use a mix of indexes
|
||||
# and tags once the framework supports it.
|
||||
tag_index: ":0"
|
||||
tag_index: ":2"
|
||||
tag_index: ":4"
|
||||
}
|
||||
sync_set { tag_index: ":1" tag_index: ":3" }
|
||||
sync_set { tag_index: ":5" }
|
||||
sync_set { tag_index: ":6" }
|
||||
sync_set { tag_index: ":7" }
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
// The sync sets by stream name and CollectionItemId.
|
||||
// {a, c, e}, {b, d}, {f}, {g}, {h}
|
||||
// {0, 2, 4}, {1, 3}, {5}, {6}, {7}
|
||||
|
||||
// The tuple is an "command" which consists of the stream name to
|
||||
// add a packet to, the input timestamp of the packet, and a list of
|
||||
// output summaries expected. Keep the list of commands separate for
|
||||
// each sync set, so that we can combine them in different ways later
|
||||
// (better testing their independence).
|
||||
std::vector<
|
||||
std::vector<std::tuple<std::string, Timestamp, std::vector<std::string>>>>
|
||||
command_sets;
|
||||
command_sets.push_back({
|
||||
CommandTuple("a", Timestamp(0), {}),
|
||||
CommandTuple("c", Timestamp(0), {}),
|
||||
CommandTuple("a", Timestamp(10), {}),
|
||||
CommandTuple("e", Timestamp(0), {"Timestamp(0),0,2,4"}),
|
||||
CommandTuple("c", Timestamp(10), {}),
|
||||
CommandTuple("e", Timestamp(10), {"Timestamp(10),0,2,4"}),
|
||||
CommandTuple("e", Timestamp(20), {}),
|
||||
CommandTuple("a", Timestamp(20), {}),
|
||||
CommandTuple("c", Timestamp(20), {"Timestamp(20),0,2,4"}),
|
||||
CommandTuple("c", Timestamp(30), {}),
|
||||
CommandTuple("a", Timestamp(30), {}),
|
||||
CommandTuple("a", Timestamp(40), {}),
|
||||
CommandTuple("a", Timestamp(50), {}),
|
||||
CommandTuple("c", Timestamp(40), {}),
|
||||
CommandTuple("a", Timestamp::Done(), {}),
|
||||
CommandTuple("e", Timestamp(40),
|
||||
{"Timestamp(30),0,2", "Timestamp(40),0,2,4"}),
|
||||
CommandTuple("c", Timestamp(50), {}),
|
||||
CommandTuple("e", Timestamp(50), {"Timestamp(50),0,2,4"}),
|
||||
CommandTuple("c", Timestamp(60), {}),
|
||||
CommandTuple("c", Timestamp(70), {}),
|
||||
CommandTuple("c", Timestamp::Done(), {}),
|
||||
CommandTuple("e", Timestamp::Done(),
|
||||
{"Timestamp(60),2", "Timestamp(70),2"}),
|
||||
});
|
||||
|
||||
command_sets.push_back({
|
||||
CommandTuple("b", Timestamp(-300), {}), //
|
||||
CommandTuple("b", Timestamp(-200), {}), //
|
||||
CommandTuple("b", Timestamp(-100), {}), //
|
||||
CommandTuple("d", Timestamp(-200),
|
||||
{"Timestamp(-300),1", "Timestamp(-200),1,3"}),
|
||||
CommandTuple("d", Timestamp(-20), {"Timestamp(-100),1"}), //
|
||||
CommandTuple("d", Timestamp(-10), {}), //
|
||||
CommandTuple("b", Timestamp(0), {"Timestamp(-20),3", "Timestamp(-10),3"}),
|
||||
CommandTuple("d", Timestamp(0), {"Timestamp(0),1,3"}), //
|
||||
CommandTuple("d", Timestamp(10), {}), //
|
||||
CommandTuple("b", Timestamp(10), {"Timestamp(10),1,3"}), //
|
||||
CommandTuple("b", Timestamp(20), {}), //
|
||||
CommandTuple("d", Timestamp(200), {"Timestamp(20),1"}), //
|
||||
CommandTuple("b", Timestamp(100), {"Timestamp(100),1"}), //
|
||||
CommandTuple("b", Timestamp(200), {"Timestamp(200),1,3"}), //
|
||||
CommandTuple("b", Timestamp(250), {}), //
|
||||
CommandTuple("b", Timestamp(300), {}), //
|
||||
CommandTuple("d", Timestamp::Done(),
|
||||
{"Timestamp(250),1", "Timestamp(300),1"}),
|
||||
CommandTuple("b", Timestamp::Done(), {}),
|
||||
});
|
||||
|
||||
std::vector<std::tuple<std::string, Timestamp, std::vector<std::string>>>
|
||||
temp_commands;
|
||||
for (Timestamp t = Timestamp(-350); t < Timestamp(350); t += 35) {
|
||||
temp_commands.push_back(CommandTuple(
|
||||
"f", t, {absl::StrCat("Timestamp(", t.DebugString(), "),5")}));
|
||||
}
|
||||
temp_commands.push_back(CommandTuple("f", Timestamp::Done(), {}));
|
||||
command_sets.push_back(temp_commands);
|
||||
|
||||
command_sets.push_back(
|
||||
{CommandTuple("g", Timestamp::PreStream(),
|
||||
{absl::StrCat(Timestamp::PreStream().DebugString(), ",6")}),
|
||||
CommandTuple("g", Timestamp::Done(), {})});
|
||||
|
||||
command_sets.push_back(
|
||||
{CommandTuple(
|
||||
"h", Timestamp::PostStream(),
|
||||
{absl::StrCat(Timestamp::PostStream().DebugString(), ",7")}),
|
||||
CommandTuple("h", Timestamp::Done(), {})});
|
||||
|
||||
int num_commands = 0;
|
||||
std::vector<int> cummulative_num_commands;
|
||||
for (int i = 0; i < command_sets.size(); ++i) {
|
||||
num_commands += command_sets[i].size();
|
||||
cummulative_num_commands.push_back(num_commands);
|
||||
}
|
||||
|
||||
RandomEngine rng(testing::UnitTest::GetInstance()->random_seed());
|
||||
for (int iter = 0; iter < 1000; ++iter) {
|
||||
LOG(INFO) << "Starting command shuffling iteration " << iter;
|
||||
|
||||
// Merge the commands for each sync set together into a serial list.
|
||||
// This is done by randomly choosing which list to grab from next.
|
||||
std::vector<std::tuple<std::string, Timestamp, std::vector<std::string>>>
|
||||
shuffled_commands;
|
||||
std::vector<int> current_positions(command_sets.size(), 0);
|
||||
while (shuffled_commands.size() < num_commands) {
|
||||
// Weight the index chosen by how many commands are in each set.
|
||||
int rand_num = rng() % num_commands;
|
||||
int command_set_index;
|
||||
for (command_set_index = 0;
|
||||
rand_num >= cummulative_num_commands[command_set_index];
|
||||
++command_set_index) {
|
||||
// Find the index corresponding to this weighted random number.
|
||||
}
|
||||
// Add the command to the list if they haven't already all been used.
|
||||
if (current_positions[command_set_index] <
|
||||
command_sets[command_set_index].size()) {
|
||||
shuffled_commands.push_back(
|
||||
command_sets[command_set_index]
|
||||
[current_positions[command_set_index]]);
|
||||
++current_positions[command_set_index];
|
||||
VLOG(2) << "ShuffledCommand (" << std::get<0>(shuffled_commands.back())
|
||||
<< ", Timestamp(" << std::get<1>(shuffled_commands.back())
|
||||
<< "))";
|
||||
}
|
||||
}
|
||||
|
||||
CalculatorGraph graph;
|
||||
|
||||
// Remove one* of the sync sets from the configuration, forcing it's
|
||||
// streams into the default sync set, which is otherwise empty.
|
||||
// * Actually, also have a possibility of not removing any.
|
||||
CalculatorGraphConfig modified_config = config;
|
||||
auto* repeated_field =
|
||||
modified_config.mutable_node(0)
|
||||
->mutable_input_stream_handler()
|
||||
->mutable_options()
|
||||
->MutableExtension(SyncSetInputStreamHandlerOptions::ext)
|
||||
->mutable_sync_set();
|
||||
int index_to_remove = rng() % (repeated_field->size() + 1);
|
||||
if (index_to_remove != repeated_field->size()) {
|
||||
repeated_field->SwapElements(index_to_remove, repeated_field->size() - 1);
|
||||
repeated_field->RemoveLast();
|
||||
}
|
||||
std::shuffle(repeated_field->begin(), repeated_field->end(), rng);
|
||||
|
||||
VLOG(2) << "Modified configuration: " << modified_config.DebugString();
|
||||
|
||||
// Setup and run the graph.
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(
|
||||
modified_config,
|
||||
{{"lambda", MakePacket<ProcessFunction>(InputsToDebugString)}}));
|
||||
std::deque<Packet> outputs;
|
||||
MEDIAPIPE_ASSERT_OK(
|
||||
graph.ObserveOutputStream("output", [&outputs](const Packet& packet) {
|
||||
outputs.push_back(packet);
|
||||
return ::mediapipe::OkStatus();
|
||||
}));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
for (int command_index = 0; command_index < shuffled_commands.size();
|
||||
/* command_index is incremented by the inner loop. */) {
|
||||
int initial_command_index = command_index;
|
||||
int command_batch_size = rng() % 10;
|
||||
std::vector<std::string> expected_strings;
|
||||
// Push in a batch of commands.
|
||||
for (; command_index < shuffled_commands.size() &&
|
||||
command_index < initial_command_index + command_batch_size;
|
||||
++command_index) {
|
||||
const auto& tup = shuffled_commands[command_index];
|
||||
const std::string& stream_name = std::get<0>(tup);
|
||||
Timestamp timestamp = std::get<1>(tup);
|
||||
expected_strings.insert(expected_strings.end(),
|
||||
std::get<2>(tup).begin(),
|
||||
std::get<2>(tup).end());
|
||||
|
||||
VLOG(1) << "Adding (" << stream_name << ", Timestamp: " << timestamp
|
||||
<< ")";
|
||||
if (timestamp == Timestamp::Done()) {
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseInputStream(stream_name));
|
||||
} else {
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
stream_name, MakePacket<int>(0).At(timestamp)));
|
||||
}
|
||||
}
|
||||
// Ensure that we produce all packets which we can.
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
|
||||
// Check the output strings (ignoring order, since calculator may
|
||||
// have run in parallel).
|
||||
// TODO Actually enable parallel process calls.
|
||||
std::vector<std::string> actual_strings;
|
||||
for (const Packet& output : outputs) {
|
||||
actual_strings.push_back(output.Get<std::string>());
|
||||
VLOG(1) << "Expecting \"" << actual_strings.back() << "\"";
|
||||
}
|
||||
if (actual_strings.empty()) {
|
||||
VLOG(1) << "Expecting nothing.";
|
||||
}
|
||||
outputs.clear();
|
||||
EXPECT_THAT(actual_strings,
|
||||
testing::UnorderedElementsAreArray(expected_strings));
|
||||
}
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,233 @@
|
||||
// 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 <algorithm>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.pb.h"
|
||||
#include "mediapipe/framework/timestamp.h"
|
||||
#include "mediapipe/framework/tool/validate_name.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
// The input streams must have the same time unit but may have different time
|
||||
// origins (also called epochs). The timestamp_base_tag_index option
|
||||
// designates an input stream as the timestamp base.
|
||||
//
|
||||
// TimestampAlignInputStreamHandler operates in two phases:
|
||||
//
|
||||
// 1. Pre-initialization: In this phase, the input stream handler passes
|
||||
// through input packets in the timestamp base input stream, but buffers the
|
||||
// input packets in all other input streams. This phase ends when the input
|
||||
// stream handler has an input packet in every input stream. It uses the
|
||||
// the timestamps of these input packets to calculate the timestamp offset of
|
||||
// each input stream with respect to the timestamp base input stream. The
|
||||
// timestamp offsets are saved for use in the next phase.
|
||||
//
|
||||
// 2. Post-initialization: In this phase, the input stream handler behaves
|
||||
// like the DefaultInputStreamHandler, except that timestamp offsets are
|
||||
// applied to the packet timestamps.
|
||||
class TimestampAlignInputStreamHandler : public InputStreamHandler {
|
||||
public:
|
||||
TimestampAlignInputStreamHandler() = delete;
|
||||
TimestampAlignInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
|
||||
CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options,
|
||||
bool calculator_run_in_parallel);
|
||||
|
||||
void PrepareForRun(
|
||||
std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(::mediapipe::Status)> error_callback) override;
|
||||
|
||||
protected:
|
||||
// In TimestampAlignInputStreamHandler, a node is "ready" if:
|
||||
// - before the timestamp offsets are initialized: we have received a packet
|
||||
// in the timestamp base input stream, or
|
||||
// - after the timestamp offsets are initialized: the minimum bound (over
|
||||
// all empty streams) is greater than the smallest timestamp of any
|
||||
// stream, which means we have received all the packets that will be
|
||||
// available at the next timestamp, or
|
||||
// - all streams are done (need to call Close() in this case).
|
||||
// Note that all packet timestamps and timestamp bounds are aligned with the
|
||||
// timestamp base.
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
private:
|
||||
CollectionItemId timestamp_base_stream_id_;
|
||||
|
||||
absl::Mutex mutex_;
|
||||
bool offsets_initialized_ GUARDED_BY(mutex_) = false;
|
||||
std::vector<TimestampDiff> timestamp_offsets_;
|
||||
};
|
||||
REGISTER_INPUT_STREAM_HANDLER(TimestampAlignInputStreamHandler);
|
||||
|
||||
TimestampAlignInputStreamHandler::TimestampAlignInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
|
||||
const MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), cc_manager, options,
|
||||
calculator_run_in_parallel),
|
||||
timestamp_offsets_(input_stream_managers_.NumEntries()) {
|
||||
const auto& handler_options =
|
||||
options.GetExtension(TimestampAlignInputStreamHandlerOptions::ext);
|
||||
std::string tag;
|
||||
int index;
|
||||
MEDIAPIPE_CHECK_OK(tool::ParseTagIndex(
|
||||
handler_options.timestamp_base_tag_index(), &tag, &index));
|
||||
timestamp_base_stream_id_ = input_stream_managers_.GetId(tag, index);
|
||||
CHECK(timestamp_base_stream_id_.IsValid())
|
||||
<< "stream \"" << handler_options.timestamp_base_tag_index()
|
||||
<< "\" is not found.";
|
||||
timestamp_offsets_[timestamp_base_stream_id_.value()] = 0;
|
||||
}
|
||||
|
||||
void TimestampAlignInputStreamHandler::PrepareForRun(
|
||||
std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(::mediapipe::Status)> error_callback) {
|
||||
{
|
||||
absl::MutexLock lock(&mutex_);
|
||||
offsets_initialized_ = (input_stream_managers_.NumEntries() == 1);
|
||||
}
|
||||
|
||||
InputStreamHandler::PrepareForRun(
|
||||
std::move(headers_ready_callback), std::move(notification_callback),
|
||||
std::move(schedule_callback), std::move(error_callback));
|
||||
}
|
||||
|
||||
NodeReadiness TimestampAlignInputStreamHandler::GetNodeReadiness(
|
||||
Timestamp* min_stream_timestamp) {
|
||||
DCHECK(min_stream_timestamp);
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
|
||||
{
|
||||
absl::MutexLock lock(&mutex_);
|
||||
if (!offsets_initialized_) {
|
||||
bool timestamp_base_empty;
|
||||
*min_stream_timestamp =
|
||||
input_stream_managers_.Get(timestamp_base_stream_id_)
|
||||
->MinTimestampOrBound(×tamp_base_empty);
|
||||
if (timestamp_base_empty) {
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
int unknown_non_base_stream_count = 0;
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
if (id == timestamp_base_stream_id_) {
|
||||
continue;
|
||||
}
|
||||
const auto& stream = input_stream_managers_.Get(id);
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
++unknown_non_base_stream_count;
|
||||
} else {
|
||||
timestamp_offsets_[id.value()] =
|
||||
*min_stream_timestamp - stream_timestamp;
|
||||
}
|
||||
}
|
||||
if (unknown_non_base_stream_count == 0) {
|
||||
offsets_initialized_ = true;
|
||||
}
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
}
|
||||
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
const auto& stream = input_stream_managers_.Get(id);
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
if (stream_timestamp.IsRangeValue()) {
|
||||
stream_timestamp += timestamp_offsets_[id.value()];
|
||||
}
|
||||
if (empty) {
|
||||
min_bound = std::min(min_bound, stream_timestamp);
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
if (*min_stream_timestamp == Timestamp::Done()) {
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
|
||||
if (min_bound > *min_stream_timestamp) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
CHECK_EQ(min_bound, *min_stream_timestamp);
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
void TimestampAlignInputStreamHandler::FillInputSet(
|
||||
Timestamp input_timestamp, InputStreamShardSet* input_set) {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
{
|
||||
absl::MutexLock lock(&mutex_);
|
||||
if (!offsets_initialized_) {
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
const auto& stream = input_stream_managers_.Get(id);
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Packet current_packet;
|
||||
if (id == timestamp_base_stream_id_) {
|
||||
current_packet = stream->PopPacketAtTimestamp(
|
||||
input_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
CHECK_EQ(num_packets_dropped, 0) << absl::Substitute(
|
||||
"Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, stream->Name());
|
||||
}
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < input_stream_managers_.EndId(); ++id) {
|
||||
auto& stream = input_stream_managers_.Get(id);
|
||||
int num_packets_dropped = 0;
|
||||
bool stream_is_done = false;
|
||||
Timestamp stream_timestamp =
|
||||
input_timestamp - timestamp_offsets_[id.value()];
|
||||
Packet current_packet = stream->PopPacketAtTimestamp(
|
||||
stream_timestamp, &num_packets_dropped, &stream_is_done);
|
||||
if (!current_packet.IsEmpty()) {
|
||||
CHECK_EQ(current_packet.Timestamp(), stream_timestamp);
|
||||
current_packet = current_packet.At(input_timestamp);
|
||||
}
|
||||
CHECK_EQ(num_packets_dropped, 0)
|
||||
<< absl::Substitute("Dropped $0 packet(s) on input stream \"$1\".",
|
||||
num_packets_dropped, stream->Name());
|
||||
AddPacketToShard(&input_set->Get(id), std::move(current_packet),
|
||||
stream_is_done);
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace mediapipe
|
||||
@@ -0,0 +1,16 @@
|
||||
syntax = "proto2";
|
||||
|
||||
package mediapipe;
|
||||
|
||||
import "mediapipe/framework/mediapipe_options.proto";
|
||||
|
||||
message TimestampAlignInputStreamHandlerOptions {
|
||||
extend MediaPipeOptions {
|
||||
optional TimestampAlignInputStreamHandlerOptions ext = 190104979;
|
||||
}
|
||||
|
||||
// The TAG:index of the input stream used as the timestamp base.
|
||||
// TimestampAlignInputStreamHandler aligns the packet timestamps of all
|
||||
// other input streams with the packet timestamps of this input stream.
|
||||
optional string timestamp_base_tag_index = 1;
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
#include <vector>
|
||||
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/port/gmock.h"
|
||||
#include "mediapipe/framework/port/gtest.h"
|
||||
#include "mediapipe/framework/port/parse_text_proto.h"
|
||||
#include "mediapipe/framework/port/status_matchers.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
namespace {
|
||||
|
||||
TEST(TimestampAlignInputStreamHandlerTest, Initialization) {
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input_video"
|
||||
input_stream: "input_camera"
|
||||
node {
|
||||
calculator: "PassThroughCalculator"
|
||||
input_stream: "VIDEO:input_video"
|
||||
input_stream: "CAMERA:input_camera"
|
||||
output_stream: "VIDEO:output_video"
|
||||
output_stream: "CAMERA:output_camera"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "TimestampAlignInputStreamHandler"
|
||||
options: {
|
||||
[mediapipe.TimestampAlignInputStreamHandlerOptions.ext]: {
|
||||
timestamp_base_tag_index: "CAMERA"
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
std::vector<Packet> sink_video, sink_camera;
|
||||
tool::AddVectorSink("output_video", &config, &sink_video);
|
||||
tool::AddVectorSink("output_camera", &config, &sink_camera);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_camera", Adopt(new int(1)).At(Timestamp(101))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// The timestamp base stream's packet is output immediately.
|
||||
EXPECT_EQ(0, sink_video.size());
|
||||
ASSERT_EQ(1, sink_camera.size());
|
||||
EXPECT_EQ(1, sink_camera[0].Get<int>());
|
||||
EXPECT_EQ(Timestamp(101), sink_camera[0].Timestamp());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_camera", Adopt(new int(2)).At(Timestamp(102))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// The timestamp base stream's packet is output immediately.
|
||||
EXPECT_EQ(0, sink_video.size());
|
||||
ASSERT_EQ(2, sink_camera.size());
|
||||
EXPECT_EQ(2, sink_camera[1].Get<int>());
|
||||
EXPECT_EQ(Timestamp(102), sink_camera[1].Timestamp());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_video", Adopt(new int(1)).At(Timestamp(1))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// No packet is output. The packet added to input_video is buffered in the
|
||||
// input stream.
|
||||
EXPECT_EQ(0, sink_video.size());
|
||||
EXPECT_EQ(2, sink_camera.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_camera", Adopt(new int(3)).At(Timestamp(103))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// Both input streams have a packet. The following equivalence of timestamps
|
||||
// is established:
|
||||
// input_video input_camera
|
||||
// 1 103
|
||||
//
|
||||
// The input stream handler is now initialized. From now on, it operates
|
||||
// like the default input stream handler except that timestamp offsets are
|
||||
// applied.
|
||||
ASSERT_EQ(1, sink_video.size());
|
||||
ASSERT_EQ(3, sink_camera.size());
|
||||
EXPECT_EQ(1, sink_video[0].Get<int>());
|
||||
EXPECT_EQ(Timestamp(103), sink_video[0].Timestamp());
|
||||
EXPECT_EQ(3, sink_camera[2].Get<int>());
|
||||
EXPECT_EQ(Timestamp(103), sink_camera[2].Timestamp());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_camera", Adopt(new int(4)).At(Timestamp(104))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// The timestamp base stream does not receive special treatment now.
|
||||
EXPECT_EQ(1, sink_video.size());
|
||||
EXPECT_EQ(3, sink_camera.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_video", Adopt(new int(4)).At(Timestamp(4))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
EXPECT_EQ(1, sink_video.size());
|
||||
ASSERT_EQ(4, sink_camera.size());
|
||||
EXPECT_EQ(4, sink_camera[3].Get<int>());
|
||||
EXPECT_EQ(Timestamp(104), sink_camera[3].Timestamp());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
ASSERT_EQ(2, sink_video.size());
|
||||
EXPECT_EQ(4, sink_camera.size());
|
||||
EXPECT_EQ(4, sink_video[1].Get<int>());
|
||||
EXPECT_EQ(Timestamp(106), sink_video[1].Timestamp());
|
||||
}
|
||||
|
||||
TEST(TimestampAlignInputStreamHandlerTest, TickRate) {
|
||||
CalculatorGraphConfig config =
|
||||
::mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"(
|
||||
input_stream: "input_video"
|
||||
input_stream: "input_camera"
|
||||
node {
|
||||
calculator: "PacketClonerCalculator"
|
||||
input_stream: "input_camera"
|
||||
input_stream: "input_video"
|
||||
input_stream: "input_video"
|
||||
output_stream: "output_camera"
|
||||
output_stream: "output_video"
|
||||
input_stream_handler {
|
||||
input_stream_handler: "TimestampAlignInputStreamHandler"
|
||||
options: {
|
||||
[mediapipe.TimestampAlignInputStreamHandlerOptions.ext]: {
|
||||
timestamp_base_tag_index: ":0" # input_camera
|
||||
}
|
||||
}
|
||||
}
|
||||
})");
|
||||
std::vector<Packet> sink_video, sink_camera;
|
||||
tool::AddVectorSink("output_video", &config, &sink_video);
|
||||
tool::AddVectorSink("output_camera", &config, &sink_camera);
|
||||
|
||||
CalculatorGraph graph;
|
||||
MEDIAPIPE_ASSERT_OK(graph.Initialize(config));
|
||||
MEDIAPIPE_ASSERT_OK(graph.StartRun({}));
|
||||
|
||||
// Video timestamps start from 0 seconds. Video frame rate is 2 fps.
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_video", Adopt(new int(0)).At(Timestamp(0))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// No packets expected as the timestamp base stream has not seen any packet.
|
||||
EXPECT_EQ(0, sink_video.size());
|
||||
EXPECT_EQ(0, sink_camera.size());
|
||||
|
||||
// Camera timestamps start from 100 seconds. Camera frame rate is 1 fps.
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_camera", Adopt(new int(0)).At(Timestamp(100000000))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
// Both input streams have a packet. The following equivalence of timestamps
|
||||
// is established:
|
||||
// input_video input_camera
|
||||
// 0 100000000
|
||||
ASSERT_EQ(1, sink_video.size());
|
||||
ASSERT_EQ(1, sink_camera.size());
|
||||
EXPECT_EQ(0, sink_video[0].Get<int>());
|
||||
EXPECT_EQ(Timestamp(100000000), sink_video[0].Timestamp());
|
||||
EXPECT_EQ(0, sink_camera[0].Get<int>());
|
||||
EXPECT_EQ(Timestamp(100000000), sink_camera[0].Timestamp());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_video", Adopt(new int(1)).At(Timestamp(500000))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
EXPECT_EQ(1, sink_video.size());
|
||||
EXPECT_EQ(1, sink_camera.size());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_video", Adopt(new int(2)).At(Timestamp(1000000))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.AddPacketToInputStream(
|
||||
"input_camera", Adopt(new int(1)).At(Timestamp(101000000))));
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilIdle());
|
||||
ASSERT_EQ(3, sink_video.size());
|
||||
ASSERT_EQ(3, sink_camera.size());
|
||||
EXPECT_EQ(1, sink_video[1].Get<int>());
|
||||
EXPECT_EQ(Timestamp(100500000), sink_video[1].Timestamp());
|
||||
EXPECT_EQ(0, sink_camera[1].Get<int>());
|
||||
EXPECT_EQ(Timestamp(100500000), sink_camera[1].Timestamp());
|
||||
EXPECT_EQ(2, sink_video[2].Get<int>());
|
||||
EXPECT_EQ(Timestamp(101000000), sink_video[2].Timestamp());
|
||||
EXPECT_EQ(1, sink_camera[2].Get<int>());
|
||||
EXPECT_EQ(Timestamp(101000000), sink_camera[2].Timestamp());
|
||||
|
||||
MEDIAPIPE_ASSERT_OK(graph.CloseAllInputStreams());
|
||||
MEDIAPIPE_ASSERT_OK(graph.WaitUntilDone());
|
||||
ASSERT_EQ(3, sink_video.size());
|
||||
ASSERT_EQ(3, sink_camera.size());
|
||||
}
|
||||
|
||||
} // namespace
|
||||
} // namespace mediapipe
|
||||
Reference in New Issue
Block a user