Expose stream handlers in headers to allow dynamic registration for superusers

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

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",
], ],

View File

@ -11,33 +11,24 @@
// 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
// consider a node as ready for Close() if any input stream is done.
class BarrierInputStreamHandler : public InputStreamHandler {
public:
BarrierInputStreamHandler() = delete;
BarrierInputStreamHandler(
std::shared_ptr<tool::TagMap> tag_map,
CalculatorContextManager* calculator_context_manager,
const MediaPipeOptions& options, bool calculator_run_in_parallel)
: InputStreamHandler(std::move(tag_map), calculator_context_manager,
options, calculator_run_in_parallel) {}
void PrepareForRun(
std::function<void()> headers_ready_callback, std::function<void()> 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) override { std::function<void(absl::Status)> error_callback) {
InputStreamHandler::PrepareForRun( InputStreamHandler::PrepareForRun(
std::move(headers_ready_callback), std::move(notification_callback), std::move(headers_ready_callback), std::move(notification_callback),
std::move(schedule_callback), std::move(error_callback)); std::move(schedule_callback), std::move(error_callback));
@ -46,11 +37,8 @@ class BarrierInputStreamHandler : public InputStreamHandler {
} }
} }
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
// - all streams have a packet available.
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override {
DCHECK(min_stream_timestamp); DCHECK(min_stream_timestamp);
*min_stream_timestamp = Timestamp::Done(); *min_stream_timestamp = Timestamp::Done();
bool all_available = true; bool all_available = true;
@ -74,9 +62,8 @@ class BarrierInputStreamHandler : public InputStreamHandler {
return NodeReadiness::kNotReady; return NodeReadiness::kNotReady;
} }
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess. void BarrierInputStreamHandler::FillInputSet(Timestamp input_timestamp,
void FillInputSet(Timestamp input_timestamp, InputStreamShardSet* input_set) {
InputStreamShardSet* input_set) override {
CHECK(input_timestamp.IsAllowedInStream()); CHECK(input_timestamp.IsAllowedInStream());
CHECK(input_set); CHECK(input_set);
for (CollectionItemId id = input_stream_managers_.BeginId(); for (CollectionItemId id = input_stream_managers_.BeginId();
@ -88,7 +75,6 @@ class BarrierInputStreamHandler : public InputStreamHandler {
stream_is_done); stream_is_done);
} }
} }
};
REGISTER_INPUT_STREAM_HANDLER(BarrierInputStreamHandler); REGISTER_INPUT_STREAM_HANDLER(BarrierInputStreamHandler);

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_

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,35 +11,25 @@
// 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
// Close() if any input stream is done.
class EarlyCloseInputStreamHandler : public InputStreamHandler {
public:
EarlyCloseInputStreamHandler() = delete;
EarlyCloseInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
CalculatorContextManager* cc_manager,
const MediaPipeOptions& options,
bool calculator_run_in_parallel)
: InputStreamHandler(std::move(tag_map), cc_manager, options,
calculator_run_in_parallel) {}
protected:
// In EarlyCloseInputStreamHandler, a node is "ready" if: // In EarlyCloseInputStreamHandler, a node is "ready" if:
// - any stream is done (need to call Close() in this case), or // - any stream is done (need to call Close() in this case), or
// - the minimum bound (over all empty streams) is greater than the smallest // - the minimum bound (over all empty streams) is greater than the smallest
// timestamp of any stream, which means we have received all the packets // timestamp of any stream, which means we have received all the packets
// that will be available at the next timestamp. // that will be available at the next timestamp.
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override { NodeReadiness EarlyCloseInputStreamHandler::GetNodeReadiness(
Timestamp* min_stream_timestamp) {
DCHECK(min_stream_timestamp); DCHECK(min_stream_timestamp);
*min_stream_timestamp = Timestamp::Done(); *min_stream_timestamp = Timestamp::Done();
Timestamp min_bound = Timestamp::Done(); Timestamp min_bound = Timestamp::Done();
@ -67,8 +57,8 @@ class EarlyCloseInputStreamHandler : public InputStreamHandler {
} }
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess. // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
void FillInputSet(Timestamp input_timestamp, void EarlyCloseInputStreamHandler::FillInputSet(
InputStreamShardSet* input_set) override { Timestamp input_timestamp, InputStreamShardSet* input_set) {
CHECK(input_timestamp.IsAllowedInStream()); CHECK(input_timestamp.IsAllowedInStream());
CHECK(input_set); CHECK(input_set);
for (CollectionItemId id = input_stream_managers_.BeginId(); for (CollectionItemId id = input_stream_managers_.BeginId();
@ -85,7 +75,6 @@ class EarlyCloseInputStreamHandler : public InputStreamHandler {
stream_is_done); stream_is_done);
} }
} }
};
REGISTER_INPUT_STREAM_HANDLER(EarlyCloseInputStreamHandler); REGISTER_INPUT_STREAM_HANDLER(EarlyCloseInputStreamHandler);

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_

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,50 +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/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)
//
// For example, a calculator node with one input stream and the following input
// stream handler specs:
//
// node {
// calculator: "CalculatorRunningAtOneFps"
// input_stream: "packets_streaming_in_at_ten_fps"
// input_stream_handler {
// input_stream_handler: "FixedSizeInputStreamHandler"
// }
// }
//
// will always try to keep the newest packet in the input stream.
//
// A few details: FixedSizeInputStreamHandler takes action when any stream grows
// to trigger_queue_size or larger. It then keeps at most target_queue_size
// packets in every InputStreamImpl. Every stream is truncated at the same
// timestamp, so that each included timestamp delivers the same packets as
// DefaultInputStreamHandler includes.
//
class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
public:
FixedSizeInputStreamHandler() = delete;
FixedSizeInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
CalculatorContextManager* cc_manager,
const MediaPipeOptions& options,
bool calculator_run_in_parallel)
: DefaultInputStreamHandler(std::move(tag_map), cc_manager, options, : DefaultInputStreamHandler(std::move(tag_map), cc_manager, options,
calculator_run_in_parallel) { calculator_run_in_parallel) {
const auto& ext = const auto& ext =
options.GetExtension(FixedSizeInputStreamHandlerOptions::ext); options.GetExtension(mediapipe::FixedSizeInputStreamHandlerOptions::ext);
trigger_queue_size_ = ext.trigger_queue_size(); trigger_queue_size_ = ext.trigger_queue_size();
target_queue_size_ = ext.target_queue_size(); target_queue_size_ = ext.target_queue_size();
fixed_min_size_ = ext.fixed_min_size(); fixed_min_size_ = ext.fixed_min_size();
@ -65,9 +51,7 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
// implementation of SetLatePreparation. // implementation of SetLatePreparation.
} }
private: void FixedSizeInputStreamHandler::EraseAllSurplus() {
// 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(); Timestamp min_timestamp_all_streams = Timestamp::Max();
for (const auto& stream : input_stream_managers_) { for (const auto& stream : input_stream_managers_) {
// Check whether every InputStreamImpl grew beyond trigger_queue_size. // Check whether every InputStreamImpl grew beyond trigger_queue_size.
@ -87,13 +71,12 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
} }
} }
// Returns the latest timestamp allowed before a bound. Timestamp FixedSizeInputStreamHandler::PreviousAllowedInStream(
Timestamp PreviousAllowedInStream(Timestamp bound) { Timestamp bound) {
return bound.IsRangeValue() ? bound - 1 : bound; return bound.IsRangeValue() ? bound - 1 : bound;
} }
// Returns the lowest timestamp at which a packet may arrive at any stream. Timestamp FixedSizeInputStreamHandler::MinStreamBound() {
Timestamp MinStreamBound() {
Timestamp min_bound = Timestamp::Done(); Timestamp min_bound = Timestamp::Done();
for (const auto& stream : input_stream_managers_) { for (const auto& stream : input_stream_managers_) {
Timestamp stream_bound = stream->GetMinTimestampAmongNLatest(1); Timestamp stream_bound = stream->GetMinTimestampAmongNLatest(1);
@ -107,8 +90,7 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
return min_bound; return min_bound;
} }
// Returns the lowest timestamp of a packet ready to process. Timestamp FixedSizeInputStreamHandler::MinTimestampToProcess() {
Timestamp MinTimestampToProcess() {
Timestamp min_bound = Timestamp::Done(); Timestamp min_bound = Timestamp::Done();
for (const auto& stream : input_stream_managers_) { for (const auto& stream : input_stream_managers_) {
bool empty; bool empty;
@ -124,11 +106,7 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
return min_bound; return min_bound;
} }
// Keeps only the most recent target_queue_size packets in each stream void FixedSizeInputStreamHandler::EraseAnySurplus(bool keep_one) {
// 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. // Record the most recent first kept timestamp on any stream.
for (const auto& stream : input_stream_managers_) { for (const auto& stream : input_stream_managers_) {
int32_t queue_size = (stream->QueueSize() >= trigger_queue_size_) int32_t queue_size = (stream->QueueSize() >= trigger_queue_size_)
@ -151,12 +129,12 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
} }
} }
void EraseSurplusPackets(bool keep_one) void FixedSizeInputStreamHandler::EraseSurplusPackets(bool keep_one) {
ABSL_EXCLUSIVE_LOCKS_REQUIRED(erase_mutex_) {
return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one); return (fixed_min_size_) ? EraseAllSurplus() : EraseAnySurplus(keep_one);
} }
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override { NodeReadiness FixedSizeInputStreamHandler::GetNodeReadiness(
Timestamp* min_stream_timestamp) {
DCHECK(min_stream_timestamp); DCHECK(min_stream_timestamp);
absl::MutexLock lock(&erase_mutex_); absl::MutexLock lock(&erase_mutex_);
// kReadyForProcess is returned only once until FillInputSet completes. // kReadyForProcess is returned only once until FillInputSet completes.
@ -175,15 +153,14 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
while (*min_stream_timestamp < kept_timestamp_ && while (*min_stream_timestamp < kept_timestamp_ &&
result == NodeReadiness::kReadyForProcess) { result == NodeReadiness::kReadyForProcess) {
EraseSurplusPackets(false); EraseSurplusPackets(false);
result = result = DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
DefaultInputStreamHandler::GetNodeReadiness(min_stream_timestamp);
} }
pending_ = (result == NodeReadiness::kReadyForProcess); pending_ = (result == NodeReadiness::kReadyForProcess);
return result; 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_) {
@ -191,7 +168,8 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
} }
} }
void MovePackets(CollectionItemId id, std::list<Packet>* packets) override { void FixedSizeInputStreamHandler::MovePackets(CollectionItemId id,
std::list<Packet>* packets) {
InputStreamHandler::MovePackets(id, packets); InputStreamHandler::MovePackets(id, packets);
absl::MutexLock lock(&erase_mutex_); absl::MutexLock lock(&erase_mutex_);
if (!pending_) { if (!pending_) {
@ -199,8 +177,8 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
} }
} }
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_) {
@ -213,18 +191,6 @@ class FixedSizeInputStreamHandler : public DefaultInputStreamHandler {
pending_ = false; pending_ = false;
} }
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_;
};
REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler); REGISTER_INPUT_STREAM_HANDLER(FixedSizeInputStreamHandler);
} // namespace mediapipe } // namespace mediapipe

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_

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) {

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_

View File

@ -11,48 +11,23 @@
// 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 {
//
// One of the input streams is the control stream; all the other input streams
// are data streams. To make MuxInputStreamHandler work properly, the tag of the
// input streams must obey the following rules:
// Let N be the number of input streams. Data streams must use tag "INPUT" with
// index 0, ..., N - 2; the control stream must use tag "SELECT".
//
// The control stream carries packets of type 'int'. The 'int' value in a
// control stream packet must be a valid index in the range 0, ..., N - 2 and
// select the data stream at that index. The selected data stream must have a
// packet with the same timestamp as the control stream packet.
//
// When the control stream is done, GetNodeReadiness() returns
// NodeReadiness::kReadyForClose.
//
// TODO: pass the input stream tags to the MuxInputStreamHandler
// constructor so that it can refer to input streams by tag. See b/30125118.
class MuxInputStreamHandler : public InputStreamHandler {
public:
MuxInputStreamHandler() = delete;
MuxInputStreamHandler(std::shared_ptr<tool::TagMap> tag_map,
CalculatorContextManager* cc_manager,
const MediaPipeOptions& options,
bool calculator_run_in_parallel)
: InputStreamHandler(std::move(tag_map), cc_manager, options,
calculator_run_in_parallel) {}
private:
CollectionItemId GetControlStreamId() const {
return input_stream_managers_.EndId() - 1; return input_stream_managers_.EndId() - 1;
} }
void RemoveOutdatedDataPackets(Timestamp timestamp) { void MuxInputStreamHandler::RemoveOutdatedDataPackets(Timestamp timestamp) {
const CollectionItemId control_stream_id = GetControlStreamId(); const CollectionItemId control_stream_id = GetControlStreamId();
for (CollectionItemId id = input_stream_managers_.BeginId(); for (CollectionItemId id = input_stream_managers_.BeginId();
id < control_stream_id; ++id) { id < control_stream_id; ++id) {
@ -60,23 +35,22 @@ class MuxInputStreamHandler : public InputStreamHandler {
} }
} }
protected:
// In MuxInputStreamHandler, a node is "ready" if: // In MuxInputStreamHandler, a node is "ready" if:
// - the control stream is done (need to call Close() in this case), or // - 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 // - we have received the packets on the control stream and the selected data
// stream at the next timestamp. // stream at the next timestamp.
NodeReadiness GetNodeReadiness(Timestamp* min_stream_timestamp) override { NodeReadiness MuxInputStreamHandler::GetNodeReadiness(
Timestamp* min_stream_timestamp) {
DCHECK(min_stream_timestamp); DCHECK(min_stream_timestamp);
absl::MutexLock lock(&input_streams_mutex_); absl::MutexLock lock(&input_streams_mutex_);
const auto& control_stream = const auto& control_stream = input_stream_managers_.Get(GetControlStreamId());
input_stream_managers_.Get(GetControlStreamId());
bool empty; bool empty;
*min_stream_timestamp = control_stream->MinTimestampOrBound(&empty); *min_stream_timestamp = control_stream->MinTimestampOrBound(&empty);
// Data streams may contain some outdated packets which failed to be popped // Data streams may contain some outdated packets which failed to be popped
// out during "FillInputSet". (This handler doesn't sync input streams, // out during "FillInputSet". (This handler doesn't sync input streams,
// hence "FillInputSet" can be triggerred before every input stream is // hence "FillInputSet" can be triggered before every input stream is
// filled with packets corresponding to the same timestamp.) // filled with packets corresponding to the same timestamp.)
RemoveOutdatedDataPackets(*min_stream_timestamp); RemoveOutdatedDataPackets(*min_stream_timestamp);
if (empty) { if (empty) {
@ -118,8 +92,8 @@ class MuxInputStreamHandler : public InputStreamHandler {
} }
// Only invoked when associated GetNodeReadiness() returned kReadyForProcess. // Only invoked when associated GetNodeReadiness() returned kReadyForProcess.
void FillInputSet(Timestamp input_timestamp, void MuxInputStreamHandler::FillInputSet(Timestamp input_timestamp,
InputStreamShardSet* input_set) override { InputStreamShardSet* input_set) {
CHECK(input_timestamp.IsAllowedInStream()); CHECK(input_timestamp.IsAllowedInStream());
CHECK(input_set); CHECK(input_set);
absl::MutexLock lock(&input_streams_mutex_); absl::MutexLock lock(&input_streams_mutex_);
@ -156,12 +130,6 @@ class MuxInputStreamHandler : public InputStreamHandler {
RemoveOutdatedDataPackets(input_timestamp.NextAllowedInStream()); RemoveOutdatedDataPackets(input_timestamp.NextAllowedInStream());
} }
private:
// Must be acquired when manipulating the control and data streams to ensure
// we have a consistent view of the two streams.
absl::Mutex input_streams_mutex_;
};
REGISTER_INPUT_STREAM_HANDLER(MuxInputStreamHandler); REGISTER_INPUT_STREAM_HANDLER(MuxInputStreamHandler);
} // namespace mediapipe } // namespace mediapipe

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_

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();

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_

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(

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_