cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Expose stream handlers in headers to allow dynamic registration for superusers

Browse Source 
PiperOrigin-RevId: 556988288
This commit is contained in:
MediaPipe Team 2023-08-14 19:57:30 -07:00 committed by Copybara-Service
parent b6f5414b3d
commit 0da296536b
15 changed files with 1019 additions and 617 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

59
mediapipe/framework/stream_handler/BUILD
View File

@ -53,8 +53,16 @@ mediapipe_proto_library(
cc_library( cc_library(
name = "barrier_input_stream_handler", name = "barrier_input_stream_handler",
srcs = ["barrier_input_stream_handler.cc"], srcs = ["barrier_input_stream_handler.cc"],
hdrs = ["barrier_input_stream_handler.h"],
deps = [ deps = [
"//mediapipe/framework:calculator_context_manager",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_handler", "//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, alwayslink = 1,
) )
@ -74,8 +82,15 @@ cc_library(
cc_library( cc_library(
name = "early_close_input_stream_handler", name = "early_close_input_stream_handler",
srcs = ["early_close_input_stream_handler.cc"], srcs = ["early_close_input_stream_handler.cc"],
hdrs = ["early_close_input_stream_handler.h"],
deps = [ deps = [
"//mediapipe/framework:calculator_context_manager",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_handler", "//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", "@com_google_absl//absl/strings",
], ],
alwayslink = 1, alwayslink = 1,
@ -84,10 +99,21 @@ cc_library(
cc_library( cc_library(
name = "fixed_size_input_stream_handler", name = "fixed_size_input_stream_handler",
srcs = ["fixed_size_input_stream_handler.cc"], srcs = ["fixed_size_input_stream_handler.cc"],
hdrs = ["fixed_size_input_stream_handler.h"],
deps = [ deps = [
":default_input_stream_handler", ":default_input_stream_handler",
":fixed_size_input_stream_handler_cc_proto", ":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: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, alwayslink = 1,
) )
@ -95,8 +121,18 @@ cc_library(
cc_library( cc_library(
name = "immediate_input_stream_handler", name = "immediate_input_stream_handler",
srcs = ["immediate_input_stream_handler.cc"], srcs = ["immediate_input_stream_handler.cc"],
hdrs = ["immediate_input_stream_handler.h"],
deps = [ deps = [
"//mediapipe/framework:calculator_context_manager",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_handler", "//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, alwayslink = 1,
) )
@ -122,9 +158,13 @@ cc_library(
cc_library( cc_library(
name = "mux_input_stream_handler", name = "mux_input_stream_handler",
srcs = ["mux_input_stream_handler.cc"], srcs = ["mux_input_stream_handler.cc"],
hdrs = ["mux_input_stream_handler.h"],
deps = [ deps = [
"//mediapipe/framework:calculator_context_manager",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_handler", "//mediapipe/framework:input_stream_handler",
"//mediapipe/framework/port:logging", "@com_google_absl//absl/log:check",
"@com_google_absl//absl/strings", "@com_google_absl//absl/strings",
"@com_google_absl//absl/synchronization", "@com_google_absl//absl/synchronization",
], ],
@ -134,16 +174,22 @@ cc_library(
cc_library( cc_library(
name = "sync_set_input_stream_handler", name = "sync_set_input_stream_handler",
srcs = ["sync_set_input_stream_handler.cc"], srcs = ["sync_set_input_stream_handler.cc"],
hdrs = ["sync_set_input_stream_handler.h"],
deps = [ deps = [
":sync_set_input_stream_handler_cc_proto", ":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:collection_item_id",
"//mediapipe/framework:input_stream_handler", "//mediapipe/framework:input_stream_handler",
"//mediapipe/framework:mediapipe_options_cc_proto", "//mediapipe/framework:mediapipe_options_cc_proto",
"//mediapipe/framework:packet_set", "//mediapipe/framework:packet_set",
"//mediapipe/framework:timestamp", "//mediapipe/framework:timestamp",
"//mediapipe/framework/port:map_util",
"//mediapipe/framework/port:status",
"//mediapipe/framework/tool:tag_map", "//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", "@com_google_absl//absl/synchronization",
], ],
alwayslink = 1, alwayslink = 1,
@ -152,12 +198,19 @@ cc_library(
cc_library( cc_library(
name = "timestamp_align_input_stream_handler", name = "timestamp_align_input_stream_handler",
srcs = ["timestamp_align_input_stream_handler.cc"], srcs = ["timestamp_align_input_stream_handler.cc"],
hdrs = ["timestamp_align_input_stream_handler.h"],
deps = [ deps = [
":timestamp_align_input_stream_handler_cc_proto", ":timestamp_align_input_stream_handler_cc_proto",
"//mediapipe/framework:calculator_context_manager",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id", "//mediapipe/framework:collection_item_id",
"//mediapipe/framework:input_stream_handler", "//mediapipe/framework:input_stream_handler",
"//mediapipe/framework:mediapipe_options_cc_proto",
"//mediapipe/framework:timestamp", "//mediapipe/framework:timestamp",
"//mediapipe/framework/tool:validate_name", "//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/strings",
"@com_google_absl//absl/synchronization", "@com_google_absl//absl/synchronization",
], ],

116
mediapipe/framework/stream_handler/barrier_input_stream_handler.cc
View File

@ -11,84 +11,70 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include "mediapipe/framework/stream_handler/barrier_input_stream_handler.h"
#include <algorithm> #include <functional>
#include <memory> #include <utility>
#include <vector>
#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" #include "mediapipe/framework/input_stream_handler.h"
namespace mediapipe { namespace mediapipe {
// Implementation of an input stream handler that considers a node as ready for void BarrierInputStreamHandler::PrepareForRun(
// Process() if all input streams have a packet available. This implies it must std::function<void()> headers_ready_callback,
// consider a node as ready for Close() if any input stream is done. std::function<void()> notification_callback,
class BarrierInputStreamHandler : public InputStreamHandler { std::function<void(CalculatorContext*)> schedule_callback,
public: std::function<void(absl::Status)> error_callback) {
BarrierInputStreamHandler() = delete; InputStreamHandler::PrepareForRun(
BarrierInputStreamHandler( std::move(headers_ready_callback), std::move(notification_callback),
std::shared_ptr<tool::TagMap> tag_map, std::move(schedule_callback), std::move(error_callback));
CalculatorContextManager* calculator_context_manager, for (auto& stream : input_stream_managers_) {
const MediaPipeOptions& options, bool calculator_run_in_parallel) stream->DisableTimestamps();
: 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();
}
} }
}
protected: NodeReadiness BarrierInputStreamHandler::GetNodeReadiness(
// In BarrierInputStreamHandler, a node is "ready" if: Timestamp* min_stream_timestamp) {
// - any stream is done (need to call Close() in this case), or DCHECK(min_stream_timestamp);
// - all streams have a packet available. *min_stream_timestamp = Timestamp::Done();
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override { bool all_available = true;
DCHECK(min_stream_timestamp); for (const auto& stream : input_stream_managers_) {
*min_stream_timestamp = Timestamp::Done(); bool empty;
bool all_available = true; Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
for (const auto& stream : input_stream_managers_) { if (empty) {
bool empty; if (stream_timestamp == Timestamp::Done()) {
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty); *min_stream_timestamp = Timestamp::Done();
if (empty) { return NodeReadiness::kReadyForClose;
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); 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. CHECK_NE(*min_stream_timestamp, Timestamp::Done());
void FillInputSet(Timestamp input_timestamp, if (all_available) {
InputStreamShardSet* input_set) override { return NodeReadiness::kReadyForProcess;
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);
}
} }
}; 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); REGISTER_INPUT_STREAM_HANDLER(BarrierInputStreamHandler);

64
mediapipe/framework/stream_handler/barrier_input_stream_handler.h Normal file
View File

@ -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_

115
mediapipe/framework/stream_handler/early_close_input_stream_handler.cc
View File

@ -1,4 +1,4 @@
// Copyright 2019 The MediaPipe Authors. // Copyright 2023 The MediaPipe Authors.
// //
// Licensed under the Apache License, Version 2.0 (the "License"); // Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with 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. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include "mediapipe/framework/stream_handler/early_close_input_stream_handler.h"
#include <algorithm> #include <algorithm>
#include <memory>
#include <vector>
#include "absl/log/check.h"
#include "absl/strings/substitute.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" #include "mediapipe/framework/input_stream_handler.h"
namespace mediapipe { namespace mediapipe {
// Implementation of an input stream handler that considers a node as ready for // In EarlyCloseInputStreamHandler, a node is "ready" if:
// Close() if any input stream is done. // - any stream is done (need to call Close() in this case), or
class EarlyCloseInputStreamHandler : public InputStreamHandler { // - the minimum bound (over all empty streams) is greater than the smallest
public: // timestamp of any stream, which means we have received all the packets
EarlyCloseInputStreamHandler() = delete; // that will be available at the next timestamp.
EarlyCloseInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map, NodeReadiness EarlyCloseInputStreamHandler::GetNodeReadiness(
CalculatorContextManager* cc_manager, Timestamp* min_stream_timestamp) {
const MediaPipeOptions& options, DCHECK(min_stream_timestamp);
bool calculator_run_in_parallel) *min_stream_timestamp = Timestamp::Done();
: InputStreamHandler(std::move(tag_map), cc_manager, options, Timestamp min_bound = Timestamp::Done();
calculator_run_in_parallel) {} for (const auto& stream : input_stream_managers_) {
bool empty;
protected: Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty);
// In EarlyCloseInputStreamHandler, a node is "ready" if: if (empty) {
// - any stream is done (need to call Close() in this case), or if (stream_timestamp == Timestamp::Done()) {
// - the minimum bound (over all empty streams) is greater than the smallest *min_stream_timestamp = Timestamp::Done();
// timestamp of any stream, which means we have received all the packets return NodeReadiness::kReadyForClose;
// 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); 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. CHECK_NE(*min_stream_timestamp, Timestamp::Done());
void FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) override { if (min_bound > *min_stream_timestamp) {
CHECK(input_timestamp.IsAllowedInStream()); return NodeReadiness::kReadyForProcess;
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_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); REGISTER_INPUT_STREAM_HANDLER(EarlyCloseInputStreamHandler);

56
mediapipe/framework/stream_handler/early_close_input_stream_handler.h Normal file
View File

@ -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_

342
mediapipe/framework/stream_handler/fixed_size_input_stream_handler.cc
View File

@ -1,4 +1,4 @@
// Copyright 2019 The MediaPipe Authors. // Copyright 2023 The MediaPipe Authors.
// //
// Licensed under the Apache License, Version 2.0 (the "License"); // Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with 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. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include "mediapipe/framework/stream_handler/fixed_size_input_stream_handler.h"
#include <algorithm>
#include <list>
#include <memory> #include <memory>
#include <utility>
#include <vector> #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" #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/stream_handler/fixed_size_input_stream_handler.pb.h"
#include "mediapipe/framework/tool/tag_map.h"
namespace mediapipe { namespace mediapipe {
// Input stream handler that limits each input queue to a maximum of FixedSizeInputStreamHandler::FixedSizeInputStreamHandler(
// target_queue_size packets, discarding older packets as needed. When a std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager,
// timestamp is dropped from a stream, it is dropped from all others as well. const mediapipe::MediaPipeOptions& options, bool calculator_run_in_parallel)
// : DefaultInputStreamHandler(std::move(tag_map), cc_manager, options,
// For example, a calculator node with one input stream and the following input calculator_run_in_parallel) {
// stream handler specs: const auto& ext =
// options.GetExtension(mediapipe::FixedSizeInputStreamHandlerOptions::ext);
// node { trigger_queue_size_ = ext.trigger_queue_size();
// calculator: "CalculatorRunningAtOneFps" target_queue_size_ = ext.target_queue_size();
// input_stream: "packets_streaming_in_at_ten_fps" fixed_min_size_ = ext.fixed_min_size();
// input_stream_handler { pending_ = false;
// input_stream_handler: "FixedSizeInputStreamHandler" kept_timestamp_ = Timestamp::Unset();
// } // TODO: Either re-enable SetLatePreparation(true) with
// } // CalculatorContext::InputTimestamp set correctly, or remove the
// // implementation of SetLatePreparation.
// 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: void FixedSizeInputStreamHandler::EraseAllSurplus() {
// Drops packets if all input streams exceed trigger_queue_size. Timestamp min_timestamp_all_streams = Timestamp::Max();
void EraseAllSurplus() ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) { for (const auto& stream : input_stream_managers_) {
Timestamp min_timestamp_all_streams = Timestamp::Max(); // Check whether every InputStreamImpl grew beyond trigger_queue_size.
for (const auto& stream : input_stream_managers_) { if (stream->QueueSize() < trigger_queue_size_) {
// Check whether every InputStreamImpl grew beyond trigger_queue_size. return;
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_) { Timestamp min_timestamp =
stream->ErasePacketsEarlierThan(min_timestamp_all_streams); 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());
} }
} }
if (keep_one) {
// Returns the latest timestamp allowed before a bound. // In order to preserve one viable timestamp, do not truncate past
Timestamp PreviousAllowedInStream(Timestamp bound) { // the timestamp bound of the least current stream.
return bound.IsRangeValue() ? bound - 1 : bound; kept_timestamp_ =
std::min(kept_timestamp_, PreviousAllowedInStream(MinStreamBound()));
} }
for (auto& stream : input_stream_managers_) {
// Returns the lowest timestamp at which a packet may arrive at any stream. stream->ErasePacketsEarlierThan(kept_timestamp_);
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. void FixedSizeInputStreamHandler::EraseSurplusPackets(bool keep_one) {
Timestamp MinTimestampToProcess() { return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one);
Timestamp min_bound = Timestamp::Done(); }
for (const auto& stream : input_stream_managers_) {
bool empty; NodeReadiness FixedSizeInputStreamHandler::GetNodeReadiness(
Timestamp stream_timestamp = stream->MinTimestampOrBound(&empty); Timestamp* min_stream_timestamp) {
// If we're using the stream's *bound*, we only want to process up to the DCHECK(min_stream_timestamp);
// packet *before* the bound, because a packet may still arrive at that absl::MutexLock lock(&erase_mutex_);
// time. // kReadyForProcess is returned only once until FillInputSet completes.
if (empty) { // In late_preparation mode, GetNodeReadiness must return kReadyForProcess
stream_timestamp = PreviousAllowedInStream(stream_timestamp); // exactly once for each input-set produced. Here, GetNodeReadiness
} // releases just one input-set at a time and then disables input queue
min_bound = std::min(min_bound, stream_timestamp); // truncation until that promised input-set is consumed.
} if (pending_) {
return min_bound; return NodeReadiness::kNotReady;
} }
EraseSurplusPackets(false);
NodeReadiness result =
DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
// Keeps only the most recent target_queue_size packets in each stream // If a packet has arrived below kept_timestamp_, recalculate.
// exceeding trigger_queue_size. Also, discards all packets older than the while (*min_stream_timestamp < kept_timestamp_ &&
// first kept timestamp on any stream. result == NodeReadiness::kReadyForProcess) {
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;
}
EraseSurplusPackets(false); EraseSurplusPackets(false);
NodeReadiness result = result = DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
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;
} }
pending_ = (result == NodeReadiness::kReadyForProcess);
return result;
}
void AddPackets(CollectionItemId id, void FixedSizeInputStreamHandler::AddPackets(CollectionItemId id,
const std::list<Packet>& packets) override { const std::list<Packet>& packets) {
InputStreamHandler::AddPackets(id, packets); InputStreamHandler::AddPackets(id, packets);
absl::MutexLock lock(&erase_mutex_); absl::MutexLock lock(&erase_mutex_);
if (!pending_) { if (!pending_) {
EraseSurplusPackets(false); EraseSurplusPackets(false);
}
} }
}
void MovePackets(CollectionItemId id, std::list<Packet>* packets) override { void FixedSizeInputStreamHandler::MovePackets(CollectionItemId id,
InputStreamHandler::MovePackets(id, packets); std::list<Packet>* packets) {
absl::MutexLock lock(&erase_mutex_); InputStreamHandler::MovePackets(id, packets);
if (!pending_) { absl::MutexLock lock(&erase_mutex_);
EraseSurplusPackets(false); if (!pending_) {
} EraseSurplusPackets(false);
} }
}
void FillInputSet(Timestamp input_timestamp, void FixedSizeInputStreamHandler::FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) override { InputStreamShardSet* input_set) {
CHECK(input_set); CHECK(input_set);
absl::MutexLock lock(&erase_mutex_); absl::MutexLock lock(&erase_mutex_);
if (!pending_) { if (!pending_) {
LOG(ERROR) << "FillInputSet called without GetNodeReadiness."; 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;
} }
// input_timestamp is recalculated here to process the most recent packets.
private: EraseSurplusPackets(true);
int32_t trigger_queue_size_; input_timestamp = MinTimestampToProcess();
int32_t target_queue_size_; DefaultInputStreamHandler::FillInputSet(input_timestamp, input_set);
bool fixed_min_size_; pending_ = false;
// 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_;
};
REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler); REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler);

108
mediapipe/framework/stream_handler/fixed_size_input_stream_handler.h Normal file
View File

@ -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_

56
mediapipe/framework/stream_handler/immediate_input_stream_handler.cc
View File

@ -11,65 +11,33 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include "mediapipe/framework/stream_handler/immediate_input_stream_handler.h"
#include <algorithm>
#include <functional>
#include <memory> #include <memory>
#include <vector> #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/input_stream_handler.h"
#include "mediapipe/framework/mediapipe_options.pb.h"
#include "mediapipe/framework/tool/tag_map.h"
namespace mediapipe { namespace mediapipe {
using SyncSet = InputStreamHandler::SyncSet; 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); REGISTER_INPUT_STREAM_HANDLER(ImmediateInputStreamHandler);
ImmediateInputStreamHandler::ImmediateInputStreamHandler( ImmediateInputStreamHandler::ImmediateInputStreamHandler(
std::shared_ptr<tool::TagMap> tag_map, std::shared_ptr<tool::TagMap> tag_map,
CalculatorContextManager* calculator_context_manager, 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, : InputStreamHandler(tag_map, calculator_context_manager, options,
calculator_run_in_parallel) { calculator_run_in_parallel) {
for (auto id = tag_map->BeginId(); id < tag_map->EndId(); ++id) { for (auto id = tag_map->BeginId(); id < tag_map->EndId(); ++id) {

77
mediapipe/framework/stream_handler/immediate_input_stream_handler.h Normal file
View File

@ -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_

230
mediapipe/framework/stream_handler/mux_input_stream_handler.cc
View File

@ -11,156 +11,124 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // 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/strings/substitute.h"
#include "absl/synchronization/mutex.h" #include "absl/synchronization/mutex.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h" #include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/input_stream_handler.h" #include "mediapipe/framework/input_stream_handler.h"
#include "mediapipe/framework/port/logging.h"
namespace mediapipe { namespace mediapipe {
// Implementation of the input stream handler for the MuxCalculator. CollectionItemId MuxInputStreamHandler::GetControlStreamId() const {
// return input_stream_managers_.EndId() - 1;
// 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 void MuxInputStreamHandler::RemoveOutdatedDataPackets(Timestamp timestamp) {
// input streams must obey the following rules: const CollectionItemId control_stream_id = GetControlStreamId();
// Let N be the number of input streams. Data streams must use tag "INPUT" with for (CollectionItemId id = input_stream_managers_.BeginId();
// index 0, ..., N - 2; the control stream must use tag "SELECT". id < control_stream_id; ++id) {
// input_stream_managers_.Get(id)->ErasePacketsEarlierThan(timestamp);
// 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;
} }
void RemoveOutdatedDataPackets(Timestamp timestamp) { }
const CollectionItemId control_stream_id = GetControlStreamId();
for (CollectionItemId id = input_stream_managers_.BeginId(); // In MuxInputStreamHandler, a node is "ready" if:
id < control_stream_id; ++id) { // - the control stream is done (need to call Close() in this case), or
input_stream_managers_.Get(id)->ErasePacketsEarlierThan(timestamp); // - 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: Packet control_packet = control_stream->QueueHead();
// In MuxInputStreamHandler, a node is "ready" if: CHECK(!control_packet.IsEmpty());
// - the control stream is done (need to call Close() in this case), or int control_value = control_packet.Get<int>();
// - we have received the packets on the control stream and the selected data CHECK_LE(0, control_value);
// stream at the next timestamp. CHECK_LT(control_value, input_stream_managers_.NumEntries() - 1);
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override { const auto& data_stream = input_stream_managers_.Get(
DCHECK(min_stream_timestamp); input_stream_managers_.BeginId() + control_value);
absl::MutexLock lock(&input_streams_mutex_);
const auto& control_stream = Timestamp stream_timestamp = data_stream->MinTimestampOrBound(&empty);
input_stream_managers_.Get(GetControlStreamId()); if (empty) {
bool empty; if (stream_timestamp <= *min_stream_timestamp) {
*min_stream_timestamp = control_stream->MinTimestampOrBound(&empty); // "data_stream" didn't receive a packet corresponding to the current
// "control_stream" packet yet.
// 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.
return NodeReadiness::kNotReady; return NodeReadiness::kNotReady;
} }
// "data_stream" timestamp bound update detected.
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);
return NodeReadiness::kReadyForProcess; return NodeReadiness::kReadyForProcess;
} }
if (stream_timestamp > *min_stream_timestamp) {
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess. // The earliest packet "data_stream" holds corresponds to a control packet
void FillInputSet(Timestamp input_timestamp, // yet to arrive, which means there won't be a "data_stream" packet
InputStreamShardSet* input_set) override { // corresponding to the current "control_stream" packet, which should be
CHECK(input_timestamp.IsAllowedInStream()); // indicated as timestamp boun update.
CHECK(input_set); return NodeReadiness::kReadyForProcess;
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());
} }
CHECK_EQ(stream_timestamp, *min_stream_timestamp);
return NodeReadiness::kReadyForProcess;
}
private: // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
// Must be acquired when manipulating the control and data streams to ensure void MuxInputStreamHandler::FillInputSet(Timestamp input_timestamp,
// we have a consistent view of the two streams. InputStreamShardSet* input_set) {
absl::Mutex input_streams_mutex_; 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); REGISTER_INPUT_STREAM_HANDLER(MuxInputStreamHandler);

80
mediapipe/framework/stream_handler/mux_input_stream_handler.h Normal file
View File

@ -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_

77
mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc
View File

@ -11,91 +11,36 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // 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/log/check.h"
#include "absl/strings/substitute.h"
#include "absl/synchronization/mutex.h" #include "absl/synchronization/mutex.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h" #include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/input_stream_handler.h" #include "mediapipe/framework/input_stream_handler.h"
#include "mediapipe/framework/mediapipe_options.pb.h"
#include "mediapipe/framework/packet_set.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/stream_handler/sync_set_input_stream_handler.pb.h"
#include "mediapipe/framework/timestamp.h" #include "mediapipe/framework/timestamp.h"
#include "mediapipe/framework/tool/tag_map.h"
namespace mediapipe { 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); 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( void SyncSetInputStreamHandler::PrepareForRun(
std::function<void()> headers_ready_callback, std::function<void()> headers_ready_callback,
std::function<void()> notification_callback, std::function<void()> notification_callback,
std::function<void(CalculatorContext*)> schedule_callback, std::function<void(CalculatorContext*)> schedule_callback,
std::function<void(absl::Status)> error_callback) { std::function<void(absl::Status)> error_callback) {
const auto& handler_options = const auto& handler_options =
options_.GetExtension(SyncSetInputStreamHandlerOptions::ext); options_.GetExtension(mediapipe::SyncSetInputStreamHandlerOptions::ext);
{ {
absl::MutexLock lock(&mutex_); absl::MutexLock lock(&mutex_);
sync_sets_.clear(); sync_sets_.clear();

97
mediapipe/framework/stream_handler/sync_set_input_stream_handler.h Normal file
View File

@ -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_

68
mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.cc
View File

@ -12,85 +12,39 @@
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include "mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h"
#include <algorithm> #include <algorithm>
#include <functional>
#include <memory>
#include <string> #include <string>
#include <utility> #include <utility>
#include <vector> #include <vector>
#include "absl/log/check.h"
#include "absl/strings/substitute.h" #include "absl/strings/substitute.h"
#include "absl/synchronization/mutex.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/collection_item_id.h"
#include "mediapipe/framework/input_stream_handler.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/stream_handler/timestamp_align_input_stream_handler.pb.h"
#include "mediapipe/framework/timestamp.h" #include "mediapipe/framework/timestamp.h"
#include "mediapipe/framework/tool/validate_name.h" #include "mediapipe/framework/tool/validate_name.h"
namespace mediapipe { 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); REGISTER_INPUT_STREAM_HANDLER(TimestampAlignInputStreamHandler);
TimestampAlignInputStreamHandler::TimestampAlignInputStreamHandler( TimestampAlignInputStreamHandler::TimestampAlignInputStreamHandler(
std::shared_ptr<tool::TagMap> tag_map, CalculatorContextManager* cc_manager, 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, : InputStreamHandler(std::move(tag_map), cc_manager, options,
calculator_run_in_parallel), calculator_run_in_parallel),
timestamp_offsets_(input_stream_managers_.NumEntries()) { timestamp_offsets_(input_stream_managers_.NumEntries()) {
const auto& handler_options = const auto& handler_options = options.GetExtension(
options.GetExtension(TimestampAlignInputStreamHandlerOptions::ext); mediapipe::TimestampAlignInputStreamHandlerOptions::ext);
std::string tag; std::string tag;
int index; int index;
MEDIAPIPE_CHECK_OK(tool::ParseTagIndex( MEDIAPIPE_CHECK_OK(tool::ParseTagIndex(

91
mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h Normal file
View File

@ -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_