diff --git a/mediapipe/framework/stream_handler/BUILD b/mediapipe/framework/stream_handler/BUILD index 8b54ade8..6767a957 100644 --- a/mediapipe/framework/stream_handler/BUILD +++ b/mediapipe/framework/stream_handler/BUILD @@ -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", ], diff --git a/mediapipe/framework/stream_handler/barrier_input_stream_handler.cc b/mediapipe/framework/stream_handler/barrier_input_stream_handler.cc index ece873b1..b483693c 100644 --- a/mediapipe/framework/stream_handler/barrier_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/barrier_input_stream_handler.cc @@ -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 -#include -#include +#include +#include +#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 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 headers_ready_callback, - std::function notification_callback, - std::function schedule_callback, - std::function 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 headers_ready_callback, + std::function notification_callback, + std::function schedule_callback, + std::function 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); diff --git a/mediapipe/framework/stream_handler/barrier_input_stream_handler.h b/mediapipe/framework/stream_handler/barrier_input_stream_handler.h new file mode 100644 index 00000000..55a21d33 --- /dev/null +++ b/mediapipe/framework/stream_handler/barrier_input_stream_handler.h @@ -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 +#include +#include + +#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 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 headers_ready_callback, + std::function notification_callback, + std::function schedule_callback, + std::function 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_ diff --git a/mediapipe/framework/stream_handler/early_close_input_stream_handler.cc b/mediapipe/framework/stream_handler/early_close_input_stream_handler.cc index 983b986c..5c448a34 100644 --- a/mediapipe/framework/stream_handler/early_close_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/early_close_input_stream_handler.cc @@ -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 -#include -#include +#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 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); diff --git a/mediapipe/framework/stream_handler/early_close_input_stream_handler.h b/mediapipe/framework/stream_handler/early_close_input_stream_handler.h new file mode 100644 index 00000000..081954ef --- /dev/null +++ b/mediapipe/framework/stream_handler/early_close_input_stream_handler.h @@ -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 +#include + +#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 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_ diff --git a/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.cc b/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.cc index fd51a738..a2e7be2f 100644 --- a/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.cc @@ -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 +#include #include +#include #include +#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 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 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& packets) override { - InputStreamHandler::AddPackets(id, packets); - absl::MutexLock lock(&erase_mutex_); - if (!pending_) { - EraseSurplusPackets(false); - } +void FixedSizeInputStreamHandler::AddPackets(CollectionItemId id, + const std::list& packets) { + InputStreamHandler::AddPackets(id, packets); + absl::MutexLock lock(&erase_mutex_); + if (!pending_) { + EraseSurplusPackets(false); } +} - void MovePackets(CollectionItemId id, std::list* packets) override { - InputStreamHandler::MovePackets(id, packets); - absl::MutexLock lock(&erase_mutex_); - if (!pending_) { - EraseSurplusPackets(false); - } +void FixedSizeInputStreamHandler::MovePackets(CollectionItemId id, + std::list* 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); diff --git a/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.h b/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.h new file mode 100644 index 00000000..a00bdda5 --- /dev/null +++ b/mediapipe/framework/stream_handler/fixed_size_input_stream_handler.h @@ -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 +#include +#include + +#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 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& packets) override; + + void MovePackets(CollectionItemId id, std::list* 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_ diff --git a/mediapipe/framework/stream_handler/immediate_input_stream_handler.cc b/mediapipe/framework/stream_handler/immediate_input_stream_handler.cc index c34fc96b..2d48c1a7 100644 --- a/mediapipe/framework/stream_handler/immediate_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/immediate_input_stream_handler.cc @@ -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 +#include #include #include +#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 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 headers_ready_callback, - std::function notification_callback, - std::function schedule_callback, - std::function 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 sync_sets_ ABSL_GUARDED_BY(mutex_); - // The input timestamp for each kReadyForProcess input stream. - std::vector ready_timestamps_ ABSL_GUARDED_BY(mutex_); -}; REGISTER_INPUT_STREAM_HANDLER(ImmediateInputStreamHandler); ImmediateInputStreamHandler::ImmediateInputStreamHandler( std::shared_ptr 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) { diff --git a/mediapipe/framework/stream_handler/immediate_input_stream_handler.h b/mediapipe/framework/stream_handler/immediate_input_stream_handler.h new file mode 100644 index 00000000..dd15ad99 --- /dev/null +++ b/mediapipe/framework/stream_handler/immediate_input_stream_handler.h @@ -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 +#include +#include + +#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 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 headers_ready_callback, + std::function notification_callback, + std::function schedule_callback, + std::function 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 sync_sets_ ABSL_GUARDED_BY(mutex_); + // The input timestamp for each kReadyForProcess input stream. + std::vector ready_timestamps_ ABSL_GUARDED_BY(mutex_); +}; + +} // namespace mediapipe + +#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_IMMEDIATE_INPUT_STREAM_HANDLER_H_ diff --git a/mediapipe/framework/stream_handler/mux_input_stream_handler.cc b/mediapipe/framework/stream_handler/mux_input_stream_handler.cc index 209c3b6f..fbf033a4 100644 --- a/mediapipe/framework/stream_handler/mux_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/mux_input_stream_handler.cc @@ -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 + +#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 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(); + 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(); - 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(); - 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(); + 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); diff --git a/mediapipe/framework/stream_handler/mux_input_stream_handler.h b/mediapipe/framework/stream_handler/mux_input_stream_handler.h new file mode 100644 index 00000000..63fdde0e --- /dev/null +++ b/mediapipe/framework/stream_handler/mux_input_stream_handler.h @@ -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 +#include + +#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 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_ diff --git a/mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc b/mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc index 1001d64f..b5b49831 100644 --- a/mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/sync_set_input_stream_handler.cc @@ -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 +#include +#include +#include +#include +#include -// 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 tag_map, - CalculatorContextManager* cc_manager, - const MediaPipeOptions& extendable_options, - bool calculator_run_in_parallel); - - void PrepareForRun(std::function headers_ready_callback, - std::function notification_callback, - std::function schedule_callback, - std::function 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 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 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 headers_ready_callback, std::function notification_callback, std::function schedule_callback, std::function error_callback) { const auto& handler_options = - options_.GetExtension(SyncSetInputStreamHandlerOptions::ext); + options_.GetExtension(mediapipe::SyncSetInputStreamHandlerOptions::ext); { absl::MutexLock lock(&mutex_); sync_sets_.clear(); diff --git a/mediapipe/framework/stream_handler/sync_set_input_stream_handler.h b/mediapipe/framework/stream_handler/sync_set_input_stream_handler.h new file mode 100644 index 00000000..67f1e49a --- /dev/null +++ b/mediapipe/framework/stream_handler/sync_set_input_stream_handler.h @@ -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 +#include +#include +#include + +#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 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 headers_ready_callback, + std::function notification_callback, + std::function schedule_callback, + std::function 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 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_ diff --git a/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.cc b/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.cc index ae075d78..3e68b161 100644 --- a/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.cc +++ b/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.cc @@ -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 +#include +#include #include #include #include +#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 tag_map, - CalculatorContextManager* cc_manager, - const MediaPipeOptions& options, - bool calculator_run_in_parallel); - - void PrepareForRun(std::function headers_ready_callback, - std::function notification_callback, - std::function schedule_callback, - std::function 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 timestamp_offsets_; -}; REGISTER_INPUT_STREAM_HANDLER(TimestampAlignInputStreamHandler); TimestampAlignInputStreamHandler::TimestampAlignInputStreamHandler( std::shared_ptr 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( diff --git a/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h b/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h new file mode 100644 index 00000000..dce8fad9 --- /dev/null +++ b/mediapipe/framework/stream_handler/timestamp_align_input_stream_handler.h @@ -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 +#include +#include + +#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 tag_map, + CalculatorContextManager* cc_manager, + const mediapipe::MediaPipeOptions& options, + bool calculator_run_in_parallel); + + void PrepareForRun(std::function headers_ready_callback, + std::function notification_callback, + std::function schedule_callback, + std::function 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 timestamp_offsets_; +}; + +} // namespace mediapipe + +#endif // MEDIAPIPE_FRAMEWORK_STREAM_HANDLER_TIMESTAMP_ALIGN_INPUT_STREAM_HANDLER_H_