Expose stream handlers in headers to allow dynamic registration for superusers
PiperOrigin-RevId: 556988288
This commit is contained in:
parent
b6f5414b3d
commit
0da296536b
|
@ -53,8 +53,16 @@ mediapipe_proto_library(
|
|||
cc_library(
|
||||
name = "barrier_input_stream_handler",
|
||||
srcs = ["barrier_input_stream_handler.cc"],
|
||||
hdrs = ["barrier_input_stream_handler.h"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/status",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
@ -74,8 +82,15 @@ cc_library(
|
|||
cc_library(
|
||||
name = "early_close_input_stream_handler",
|
||||
srcs = ["early_close_input_stream_handler.cc"],
|
||||
hdrs = ["early_close_input_stream_handler.h"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/strings",
|
||||
],
|
||||
alwayslink = 1,
|
||||
|
@ -84,10 +99,21 @@ cc_library(
|
|||
cc_library(
|
||||
name = "fixed_size_input_stream_handler",
|
||||
srcs = ["fixed_size_input_stream_handler.cc"],
|
||||
hdrs = ["fixed_size_input_stream_handler.h"],
|
||||
deps = [
|
||||
":default_input_stream_handler",
|
||||
":fixed_size_input_stream_handler_cc_proto",
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework:packet",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/log",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
@ -95,8 +121,18 @@ cc_library(
|
|||
cc_library(
|
||||
name = "immediate_input_stream_handler",
|
||||
srcs = ["immediate_input_stream_handler.cc"],
|
||||
hdrs = ["immediate_input_stream_handler.h"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
alwayslink = 1,
|
||||
)
|
||||
|
@ -122,9 +158,13 @@ cc_library(
|
|||
cc_library(
|
||||
name = "mux_input_stream_handler",
|
||||
srcs = ["mux_input_stream_handler.cc"],
|
||||
hdrs = ["mux_input_stream_handler.h"],
|
||||
deps = [
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework/port:logging",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
|
@ -134,16 +174,22 @@ cc_library(
|
|||
cc_library(
|
||||
name = "sync_set_input_stream_handler",
|
||||
srcs = ["sync_set_input_stream_handler.cc"],
|
||||
hdrs = ["sync_set_input_stream_handler.h"],
|
||||
deps = [
|
||||
":sync_set_input_stream_handler_cc_proto",
|
||||
"//mediapipe/framework:collection",
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//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/port:map_util",
|
||||
"//mediapipe/framework/port:status",
|
||||
"//mediapipe/framework/tool:tag_map",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
alwayslink = 1,
|
||||
|
@ -152,12 +198,19 @@ cc_library(
|
|||
cc_library(
|
||||
name = "timestamp_align_input_stream_handler",
|
||||
srcs = ["timestamp_align_input_stream_handler.cc"],
|
||||
hdrs = ["timestamp_align_input_stream_handler.h"],
|
||||
deps = [
|
||||
":timestamp_align_input_stream_handler_cc_proto",
|
||||
"//mediapipe/framework:calculator_context_manager",
|
||||
"//mediapipe/framework:calculator_framework",
|
||||
"//mediapipe/framework:collection_item_id",
|
||||
"//mediapipe/framework:input_stream_handler",
|
||||
"//mediapipe/framework:mediapipe_options_cc_proto",
|
||||
"//mediapipe/framework:timestamp",
|
||||
"//mediapipe/framework/tool:validate_name",
|
||||
"@com_google_absl//absl/base:core_headers",
|
||||
"@com_google_absl//absl/log:check",
|
||||
"@com_google_absl//absl/status",
|
||||
"@com_google_absl//absl/strings",
|
||||
"@com_google_absl//absl/synchronization",
|
||||
],
|
||||
|
|
|
@ -11,84 +11,70 @@
|
|||
// 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/barrier_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
#include <functional>
|
||||
#include <utility>
|
||||
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#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(absl::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();
|
||||
}
|
||||
void BarrierInputStreamHandler::PrepareForRun(
|
||||
std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(absl::Status)> error_callback) {
|
||||
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;
|
||||
NodeReadiness BarrierInputStreamHandler::GetNodeReadiness(
|
||||
Timestamp* min_stream_timestamp) {
|
||||
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;
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
all_available = false;
|
||||
}
|
||||
|
||||
CHECK_NE(*min_stream_timestamp, Timestamp::Done());
|
||||
if (all_available) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
return NodeReadiness::kNotReady;
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
// 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);
|
||||
}
|
||||
CHECK_NE(*min_stream_timestamp, Timestamp::Done());
|
||||
if (all_available) {
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
};
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
|
||||
void BarrierInputStreamHandler::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);
|
||||
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);
|
||||
|
||||
|
|
|
@ -0,0 +1,64 @@
|
|||
|
||||
// Copyright 2023 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_BARRIER_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_BARRIER_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "absl/status/status.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/tool/tag_map.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 mediapipe::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(absl::Status)> error_callback) override;
|
||||
|
||||
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;
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_BARRIER_INPUT_STREAM_HANDLER_H_
|
|
@ -1,4 +1,4 @@
|
|||
// Copyright 2019 The MediaPipe Authors.
|
||||
// Copyright 2023 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.
|
||||
|
@ -11,81 +11,70 @@
|
|||
// 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/early_close_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.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);
|
||||
// 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 EarlyCloseInputStreamHandler::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) {
|
||||
if (stream_timestamp == Timestamp::Done()) {
|
||||
*min_stream_timestamp = Timestamp::Done();
|
||||
return NodeReadiness::kReadyForClose;
|
||||
}
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
min_bound = std::min(min_bound, 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;
|
||||
*min_stream_timestamp = std::min(*min_stream_timestamp, stream_timestamp);
|
||||
}
|
||||
|
||||
// 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);
|
||||
}
|
||||
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 EarlyCloseInputStreamHandler::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);
|
||||
}
|
||||
}
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(EarlyCloseInputStreamHandler);
|
||||
|
||||
|
|
|
@ -0,0 +1,56 @@
|
|||
// Copyright 2023 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_EARLY_CLOSE_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_EARLY_CLOSE_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/tool/tag_map.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 mediapipe::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;
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_EARLY_CLOSE_INPUT_STREAM_HANDLER_H_
|
|
@ -1,4 +1,4 @@
|
|||
// Copyright 2019 The MediaPipe Authors.
|
||||
// Copyright 2023 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.
|
||||
|
@ -11,219 +11,185 @@
|
|||
// 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/fixed_size_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <list>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/log/log.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/packet.h"
|
||||
#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"
|
||||
#include "mediapipe/framework/tool/tag_map.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.
|
||||
}
|
||||
FixedSizeInputStreamHandler::FixedSizeInputStreamHandler(
|
||||
std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
|
||||
const mediapipe::MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: DefaultInputStreamHandler(std::move(tag_map), cc_manager, options,
|
||||
calculator_run_in_parallel) {
|
||||
const auto& ext =
|
||||
options.GetExtension(mediapipe::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() ABSL_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);
|
||||
void FixedSizeInputStreamHandler::EraseAllSurplus() {
|
||||
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;
|
||||
}
|
||||
for (auto& stream : input_stream_managers_) {
|
||||
stream->ErasePacketsEarlierThan(min_timestamp_all_streams);
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
Timestamp FixedSizeInputStreamHandler::PreviousAllowedInStream(
|
||||
Timestamp bound) {
|
||||
return bound.IsRangeValue() ? bound - 1 : bound;
|
||||
}
|
||||
|
||||
Timestamp FixedSizeInputStreamHandler::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;
|
||||
}
|
||||
|
||||
Timestamp FixedSizeInputStreamHandler::MinTimestampToProcess() {
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
// If we're using the stream's *bound*, we only want to process up to the
|
||||
// packet *before* the bound, because a packet may still arrive at that
|
||||
// time.
|
||||
if (empty) {
|
||||
stream_timestamp = PreviousAllowedInStream(stream_timestamp);
|
||||
}
|
||||
min_bound = std::min(min_bound, stream_timestamp);
|
||||
}
|
||||
return min_bound;
|
||||
}
|
||||
|
||||
void FixedSizeInputStreamHandler::EraseAnySurplus(bool keep_one) {
|
||||
// Record the most recent first kept timestamp on any stream.
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
int32_t 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());
|
||||
}
|
||||
}
|
||||
|
||||
// Returns the latest timestamp allowed before a bound.
|
||||
Timestamp PreviousAllowedInStream(Timestamp bound) {
|
||||
return bound.IsRangeValue() ? bound - 1 : bound;
|
||||
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()));
|
||||
}
|
||||
|
||||
// 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;
|
||||
for (auto& stream : input_stream_managers_) {
|
||||
stream->ErasePacketsEarlierThan(kept_timestamp_);
|
||||
}
|
||||
}
|
||||
|
||||
// Returns the lowest timestamp of a packet ready to process.
|
||||
Timestamp MinTimestampToProcess() {
|
||||
Timestamp min_bound = Timestamp::Done();
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
bool empty;
|
||||
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
|
||||
// If we're using the stream's *bound*, we only want to process up to the
|
||||
// packet *before* the bound, because a packet may still arrive at that
|
||||
// time.
|
||||
if (empty) {
|
||||
stream_timestamp = PreviousAllowedInStream(stream_timestamp);
|
||||
}
|
||||
min_bound = std::min(min_bound, stream_timestamp);
|
||||
}
|
||||
return min_bound;
|
||||
void FixedSizeInputStreamHandler::EraseSurplusPackets(bool keep_one) {
|
||||
return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one);
|
||||
}
|
||||
|
||||
NodeReadiness FixedSizeInputStreamHandler::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);
|
||||
|
||||
// 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)
|
||||
ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
|
||||
// Record the most recent first kept timestamp on any stream.
|
||||
for (const auto& stream : input_stream_managers_) {
|
||||
int32_t 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)
|
||||
ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
|
||||
return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one);
|
||||
}
|
||||
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
|
||||
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;
|
||||
}
|
||||
// If a packet has arrived below kept_timestamp_, recalculate.
|
||||
while (*min_stream_timestamp < kept_timestamp_ &&
|
||||
result == NodeReadiness::kReadyForProcess) {
|
||||
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;
|
||||
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 FixedSizeInputStreamHandler::AddPackets(CollectionItemId id,
|
||||
const std::list<Packet>& packets) {
|
||||
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 FixedSizeInputStreamHandler::MovePackets(CollectionItemId id,
|
||||
std::list<Packet>* packets) {
|
||||
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 = MinTimestampToProcess();
|
||||
DefaultInputStreamHandler::FillInputSet(input_timestamp, input_set);
|
||||
pending_ = false;
|
||||
void FixedSizeInputStreamHandler::FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) {
|
||||
CHECK(input_set);
|
||||
absl::MutexLock lock(&erase_mutex_);
|
||||
if (!pending_) {
|
||||
LOG(ERROR) << "FillInputSet called without GetNodeReadiness.";
|
||||
}
|
||||
|
||||
private:
|
||||
int32_t trigger_queue_size_;
|
||||
int32_t 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_ ABSL_GUARDED_BY(erase_mutex_);
|
||||
// The timestamp used to truncate all input streams.
|
||||
Timestamp kept_timestamp_ ABSL_GUARDED_BY(erase_mutex_);
|
||||
absl::Mutex erase_mutex_;
|
||||
};
|
||||
// input_timestamp is recalculated here to process the most recent packets.
|
||||
EraseSurplusPackets(true);
|
||||
input_timestamp = MinTimestampToProcess();
|
||||
DefaultInputStreamHandler::FillInputSet(input_timestamp, input_set);
|
||||
pending_ = false;
|
||||
}
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler);
|
||||
|
||||
|
|
|
@ -0,0 +1,108 @@
|
|||
// Copyright 2023 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_FIXED_SIZE_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_FIXED_SIZE_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <cstdint>
|
||||
#include <list>
|
||||
#include <memory>
|
||||
|
||||
#include "absl/base/thread_annotations.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/stream_handler/default_input_stream_handler.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);
|
||||
|
||||
private:
|
||||
// Drops packets if all input streams exceed trigger_queue_size.
|
||||
void EraseAllSurplus() ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_);
|
||||
|
||||
// Returns the latest timestamp allowed before a bound.
|
||||
Timestamp PreviousAllowedInStream(Timestamp bound);
|
||||
|
||||
// Returns the lowest timestamp at which a packet may arrive at any stream.
|
||||
Timestamp MinStreamBound();
|
||||
|
||||
// Returns the lowest timestamp of a packet ready to process.
|
||||
Timestamp MinTimestampToProcess();
|
||||
|
||||
// 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)
|
||||
ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_);
|
||||
|
||||
void EraseSurplusPackets(bool keep_one)
|
||||
ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_);
|
||||
|
||||
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override;
|
||||
|
||||
void AddPackets(CollectionItemId id,
|
||||
const std::list<Packet>& packets) override;
|
||||
|
||||
void MovePackets(CollectionItemId id, std::list<Packet>* packets) override;
|
||||
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
private:
|
||||
int32_t trigger_queue_size_;
|
||||
int32_t 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_ ABSL_GUARDED_BY(erase_mutex_);
|
||||
// The timestamp used to truncate all input streams.
|
||||
Timestamp kept_timestamp_ ABSL_GUARDED_BY(erase_mutex_);
|
||||
absl::Mutex erase_mutex_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_FIXED_SIZE_INPUT_STREAM_HANDLER_H_
|
|
@ -11,65 +11,33 @@
|
|||
// 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/immediate_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/tool/tag_map.h"
|
||||
|
||||
namespace mediapipe {
|
||||
|
||||
using SyncSet = InputStreamHandler::SyncSet;
|
||||
|
||||
// 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:
|
||||
// Reinitializes this InputStreamHandler before each CalculatorGraph run.
|
||||
void PrepareForRun(std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(absl::Status)> error_callback) override;
|
||||
|
||||
// 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;
|
||||
|
||||
// Returns the number of sync-sets maintained by this input-handler.
|
||||
int SyncSetCount() override;
|
||||
|
||||
absl::Mutex mutex_;
|
||||
// The packet-set builder for each input stream.
|
||||
std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
|
||||
// The input timestamp for each kReadyForProcess input stream.
|
||||
std::vector<Timestamp> ready_timestamps_ ABSL_GUARDED_BY(mutex_);
|
||||
};
|
||||
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)
|
||||
const mediapipe::MediaPipeOptions& options, bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(tag_map, calculator_context_manager, options,
|
||||
calculator_run_in_parallel) {
|
||||
for (auto id = tag_map->BeginId(); id < tag_map->EndId(); ++id) {
|
||||
|
|
|
@ -0,0 +1,77 @@
|
|||
// Copyright 2023 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_IMMEDIATE_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IMMEDIATE_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/base/thread_annotations.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/tool/tag_map.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 timestamps 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:
|
||||
// Reinitializes this InputStreamHandler before each CalculatorGraph run.
|
||||
void PrepareForRun(std::function<void()> headers_ready_callback,
|
||||
std::function<void()> notification_callback,
|
||||
std::function<void(CalculatorContext*)> schedule_callback,
|
||||
std::function<void(absl::Status)> error_callback) override;
|
||||
|
||||
// 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;
|
||||
|
||||
// Returns the number of sync-sets maintained by this input-handler.
|
||||
int SyncSetCount() override;
|
||||
|
||||
absl::Mutex mutex_;
|
||||
// The packet-set builder for each input stream.
|
||||
std::vector<SyncSet> sync_sets_ ABSL_GUARDED_BY(mutex_);
|
||||
// The input timestamp for each kReadyForProcess input stream.
|
||||
std::vector<Timestamp> ready_timestamps_ ABSL_GUARDED_BY(mutex_);
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IMMEDIATE_INPUT_STREAM_HANDLER_H_
|
|
@ -11,156 +11,124 @@
|
|||
// 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/mux_input_stream_handler.h"
|
||||
|
||||
#include <utility>
|
||||
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.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) {}
|
||||
|
||||
private:
|
||||
CollectionItemId GetControlStreamId() const {
|
||||
return input_stream_managers_.EndId() - 1;
|
||||
CollectionItemId MuxInputStreamHandler::GetControlStreamId() const {
|
||||
return input_stream_managers_.EndId() - 1;
|
||||
}
|
||||
void MuxInputStreamHandler::RemoveOutdatedDataPackets(Timestamp timestamp) {
|
||||
const CollectionItemId control_stream_id = GetControlStreamId();
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < control_stream_id; ++id) {
|
||||
input_stream_managers_.Get(id)->ErasePacketsEarlierThan(timestamp);
|
||||
}
|
||||
void RemoveOutdatedDataPackets(Timestamp timestamp) {
|
||||
const CollectionItemId control_stream_id = GetControlStreamId();
|
||||
for (CollectionItemId id = input_stream_managers_.BeginId();
|
||||
id < control_stream_id; ++id) {
|
||||
input_stream_managers_.Get(id)->ErasePacketsEarlierThan(timestamp);
|
||||
}
|
||||
|
||||
// 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 MuxInputStreamHandler::GetNodeReadiness(
|
||||
Timestamp* min_stream_timestamp) {
|
||||
DCHECK(min_stream_timestamp);
|
||||
absl::MutexLock lock(&input_streams_mutex_);
|
||||
|
||||
const auto& control_stream = input_stream_managers_.Get(GetControlStreamId());
|
||||
bool empty;
|
||||
*min_stream_timestamp = control_stream->MinTimestampOrBound(&empty);
|
||||
|
||||
// Data streams may contain some outdated packets which failed to be popped
|
||||
// out during "FillInputSet". (This handler doesn't sync input streams,
|
||||
// hence "FillInputSet" can be triggered before every input stream is
|
||||
// filled with packets corresponding to the same timestamp.)
|
||||
RemoveOutdatedDataPackets(*min_stream_timestamp);
|
||||
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;
|
||||
}
|
||||
|
||||
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_);
|
||||
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);
|
||||
|
||||
const auto& control_stream =
|
||||
input_stream_managers_.Get(GetControlStreamId());
|
||||
bool empty;
|
||||
*min_stream_timestamp = control_stream->MinTimestampOrBound(&empty);
|
||||
|
||||
// Data streams may contain some outdated packets which failed to be popped
|
||||
// out during "FillInputSet". (This handler doesn't sync input streams,
|
||||
// hence "FillInputSet" can be triggerred before every input stream is
|
||||
// filled with packets corresponding to the same timestamp.)
|
||||
RemoveOutdatedDataPackets(*min_stream_timestamp);
|
||||
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.
|
||||
Timestamp stream_timestamp = data_stream->MinTimestampOrBound(&empty);
|
||||
if (empty) {
|
||||
if (stream_timestamp <= *min_stream_timestamp) {
|
||||
// "data_stream" didn't receive a packet corresponding to the current
|
||||
// "control_stream" packet yet.
|
||||
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) {
|
||||
if (stream_timestamp <= *min_stream_timestamp) {
|
||||
// "data_stream" didn't receive a packet corresponding to the current
|
||||
// "control_stream" packet yet.
|
||||
return NodeReadiness::kNotReady;
|
||||
}
|
||||
// "data_stream" timestamp bound update detected.
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
if (stream_timestamp > *min_stream_timestamp) {
|
||||
// The earliest packet "data_stream" holds corresponds to a control packet
|
||||
// yet to arrive, which means there won't be a "data_stream" packet
|
||||
// corresponding to the current "control_stream" packet, which should be
|
||||
// indicated as timestamp boun update.
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
CHECK_EQ(stream_timestamp, *min_stream_timestamp);
|
||||
// "data_stream" timestamp bound update detected.
|
||||
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 = GetControlStreamId();
|
||||
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);
|
||||
|
||||
// Discard old packets on data streams.
|
||||
RemoveOutdatedDataPackets(input_timestamp.NextAllowedInStream());
|
||||
if (stream_timestamp > *min_stream_timestamp) {
|
||||
// The earliest packet "data_stream" holds corresponds to a control packet
|
||||
// yet to arrive, which means there won't be a "data_stream" packet
|
||||
// corresponding to the current "control_stream" packet, which should be
|
||||
// indicated as timestamp boun update.
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
CHECK_EQ(stream_timestamp, *min_stream_timestamp);
|
||||
return NodeReadiness::kReadyForProcess;
|
||||
}
|
||||
|
||||
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_;
|
||||
};
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void MuxInputStreamHandler::FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) {
|
||||
CHECK(input_timestamp.IsAllowedInStream());
|
||||
CHECK(input_set);
|
||||
absl::MutexLock lock(&input_streams_mutex_);
|
||||
|
||||
const CollectionItemId control_stream_id = GetControlStreamId();
|
||||
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);
|
||||
|
||||
// Discard old packets on data streams.
|
||||
RemoveOutdatedDataPackets(input_timestamp.NextAllowedInStream());
|
||||
}
|
||||
|
||||
REGISTER_INPUT_STREAM_HANDLER(MuxInputStreamHandler);
|
||||
|
||||
|
|
|
@ -0,0 +1,80 @@
|
|||
// Copyright 2023 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_MUX_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_MUX_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.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) {}
|
||||
|
||||
private:
|
||||
CollectionItemId GetControlStreamId() const;
|
||||
void RemoveOutdatedDataPackets(Timestamp timestamp);
|
||||
|
||||
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;
|
||||
|
||||
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
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_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_MUX_INPUT_STREAM_HANDLER_H_
|
|
@ -11,91 +11,36 @@
|
|||
// 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/sync_set_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <set>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
// TODO: Move protos in another CL after the C++ code migration.
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.h"
|
||||
#include "mediapipe/framework/packet_set.h"
|
||||
#include "mediapipe/framework/port/map_util.h"
|
||||
#include "mediapipe/framework/port/status.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(absl::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.
|
||||
// Populates packets for the ready sync-set, and populates timestamp bounds
|
||||
// for all sync-sets.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
// Populates timestamp bounds for streams outside the ready sync-set.
|
||||
void FillInputBounds(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set)
|
||||
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
|
||||
|
||||
// Returns the number of sync-sets maintained by this input-handler.
|
||||
int SyncSetCount() override;
|
||||
|
||||
private:
|
||||
absl::Mutex mutex_;
|
||||
// The ids of each set of inputs.
|
||||
std::vector<SyncSet> sync_sets_ ABSL_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_ ABSL_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_ ABSL_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(absl::Status)> error_callback) {
|
||||
const auto& handler_options =
|
||||
options_.GetExtension(SyncSetInputStreamHandlerOptions::ext);
|
||||
options_.GetExtension(mediapipe::SyncSetInputStreamHandlerOptions::ext);
|
||||
{
|
||||
absl::MutexLock lock(&mutex_);
|
||||
sync_sets_.clear();
|
||||
|
|
|
@ -0,0 +1,97 @@
|
|||
// Copyright 2023 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_SYNC_SET_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_SYNC_SET_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/base/thread_annotations.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.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 mediapipe::MediaPipeOptions& extendable_options,
|
||||
bool calculator_run_in_parallel)
|
||||
: InputStreamHandler(std::move(tag_map), cc_manager, extendable_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(absl::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.
|
||||
// Populates packets for the ready sync-set, and populates timestamp bounds
|
||||
// for all sync-sets.
|
||||
void FillInputSet(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set) override;
|
||||
|
||||
// Populates timestamp bounds for streams outside the ready sync-set.
|
||||
void FillInputBounds(Timestamp input_timestamp,
|
||||
InputStreamShardSet* input_set)
|
||||
ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
|
||||
|
||||
// Returns the number of sync-sets maintained by this input-handler.
|
||||
int SyncSetCount() override;
|
||||
|
||||
private:
|
||||
absl::Mutex mutex_;
|
||||
// The ids of each set of inputs.
|
||||
std::vector<SyncSet> sync_sets_ ABSL_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_ ABSL_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_ ABSL_GUARDED_BY(mutex_);
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_SYNC_SET_INPUT_STREAM_HANDLER_H_
|
|
@ -12,85 +12,39 @@
|
|||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/log/check.h"
|
||||
#include "absl/strings/substitute.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.h"
|
||||
#include "mediapipe/framework/collection_item_id.h"
|
||||
#include "mediapipe/framework/input_stream_handler.h"
|
||||
#include "mediapipe/framework/mediapipe_options.pb.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(absl::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_ ABSL_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)
|
||||
const mediapipe::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);
|
||||
const auto& handler_options = options.GetExtension(
|
||||
mediapipe::TimestampAlignInputStreamHandlerOptions::ext);
|
||||
std::string tag;
|
||||
int index;
|
||||
MEDIAPIPE_CHECK_OK(tool::ParseTagIndex(
|
||||
|
|
|
@ -0,0 +1,91 @@
|
|||
// Copyright 2023 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_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_
|
||||
#define MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_
|
||||
|
||||
#include <functional>
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "absl/base/thread_annotations.h"
|
||||
#include "absl/status/status.h"
|
||||
#include "absl/synchronization/mutex.h"
|
||||
#include "mediapipe/framework/calculator_context_manager.h"
|
||||
#include "mediapipe/framework/calculator_framework.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"
|
||||
|
||||
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 mediapipe::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(absl::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_ ABSL_GUARDED_BY(mutex_) = false;
|
||||
std::vector<TimestampDiff> timestamp_offsets_;
|
||||
};
|
||||
|
||||
} // namespace mediapipe
|
||||
|
||||
#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_
|
Loading…
Reference in New Issue
Block a user