Improve loop calculator documentation and add additional specializations

The documentation is confusing since it was unclear where e.g. BeginLoopWithIterableCalculator comes from. Also, input_to_loop_body wasn't connected to anything. loop_internal_ts is more confusing than helpful. This CL cleans up the docs. It also adds specializations for CPU image and GPU texture buffers to be used for the recompose effect.

PiperOrigin-RevId: 532083714
This commit is contained in:
MediaPipe Team 2023-05-15 05:53:16 -07:00 committed by Copybara-Service
parent 12ba644393
commit fc9538533c
5 changed files with 62 additions and 51 deletions

View File

@ -192,17 +192,18 @@ cc_library(
"//mediapipe/framework:calculator_context",
"//mediapipe/framework:calculator_contract",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework:packet",
"//mediapipe/framework/formats:detection_cc_proto",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/gpu:gpu_buffer",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
],
alwayslink = 1,
)
@ -215,18 +216,20 @@ cc_library(
"//mediapipe/framework:calculator_context",
"//mediapipe/framework:calculator_contract",
"//mediapipe/framework:calculator_framework",
"//mediapipe/framework:collection_item_id",
"//mediapipe/framework/formats:classification_cc_proto",
"//mediapipe/framework/formats:detection_cc_proto",
"//mediapipe/framework/formats:image",
"//mediapipe/framework/formats:image_frame",
"//mediapipe/framework/formats:landmark_cc_proto",
"//mediapipe/framework/formats:matrix",
"//mediapipe/framework/formats:rect_cc_proto",
"//mediapipe/framework/formats:tensor",
"//mediapipe/framework/port:integral_types",
"//mediapipe/framework/port:ret_check",
"//mediapipe/framework/port:status",
"//mediapipe/gpu:gpu_buffer",
"//mediapipe/util:render_data_cc_proto",
"@com_google_absl//absl/memory",
"@com_google_absl//absl/status",
"@org_tensorflow//tensorflow/lite:framework",
],
alwayslink = 1,

View File

@ -17,10 +17,12 @@
#include <vector>
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/gpu/gpu_buffer.h"
namespace mediapipe {
@ -60,4 +62,14 @@ REGISTER_CALCULATOR(BeginLoopUint64tCalculator);
typedef BeginLoopCalculator<std::vector<Tensor>> BeginLoopTensorCalculator;
REGISTER_CALCULATOR(BeginLoopTensorCalculator);
// A calculator to process std::vector<mediapipe::ImageFrame>.
typedef BeginLoopCalculator<std::vector<ImageFrame>>
BeginLoopImageFrameCalculator;
REGISTER_CALCULATOR(BeginLoopImageFrameCalculator);
// A calculator to process std::vector<mediapipe::GpuBuffer>.
typedef BeginLoopCalculator<std::vector<GpuBuffer>>
BeginLoopGpuBufferCalculator;
REGISTER_CALCULATOR(BeginLoopGpuBufferCalculator);
} // namespace mediapipe

View File

@ -15,47 +15,57 @@
#ifndef MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
#define MEDIAPIPE_CALCULATORS_CORE_BEGIN_LOOP_CALCULATOR_H_
#include "absl/status/status.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_contract.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/packet.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
#include "mediapipe/framework/port/status_macros.h"
namespace mediapipe {
// Calculator for implementing loops on iterable collections inside a MediaPipe
// graph.
// graph. Assume InputIterT is an iterable for type InputT, and OutputIterT is
// an iterable for type OutputT, e.g. vector<InputT> and vector<OutputT>.
// First, instantiate specializations in the loop calculators' implementations
// if missing:
// BeginLoopInputTCalculator = BeginLoopCalculator<InputIterT>
// EndLoopOutputTCalculator = EndLoopCalculator<OutputIterT>
// Then, the following graph transforms an item of type InputIterT to an
// OutputIterT by applying InputToOutputConverter to every element:
//
// It is designed to be used like:
//
// node {
// calculator: "BeginLoopWithIterableCalculator"
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
// node { # Type @timestamp
// calculator: "BeginLoopInputTCalculator"
// input_stream: "ITERABLE:input_iterable" # InputIterT @iterable_ts
// input_stream: "CLONE:extra_input" # ExtraT @extra_ts
// output_stream: "ITEM:input_iterator" # InputT @loop_internal_ts
// output_stream: "CLONE:cloned_extra_input" # ExtraT @loop_internal_ts
// output_stream: "BATCH_END:iterable_ts" # Timestamp @loop_internal_ts
// }
//
// node {
// calculator: "ElementToBlaConverterSubgraph"
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
// calculator: "InputToOutputConverter"
// input_stream: "INPUT:input_iterator" # InputT @loop_internal_ts
// input_stream: "EXTRA:cloned_extra_input" # ExtraT @loop_internal_ts
// output_stream: "OUTPUT:output_iterator" # OutputT @loop_internal_ts
// }
//
// node {
// calculator: "EndLoopWithOutputCalculator"
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
// output_stream: "ITERABLE:aggregated_result" # IterableU @ext_ts
// calculator: "EndLoopOutputTCalculator"
// input_stream: "ITEM:output_iterator" # OutputT @loop_internal_ts
// input_stream: "BATCH_END:iterable_ts" # Timestamp @loop_internal_ts
// output_stream: "ITERABLE:output_iterable" # OutputIterT @iterable_ts
// }
//
// The resulting 'output_iterable' has the same timestamp as 'input_iterable'.
// The output packets of this calculator are part of the loop body and have
// loop-internal timestamps that are unrelated to the input iterator timestamp.
//
// Input streams tagged with "CLONE" are cloned to the corresponding output
// streams at loop timestamps. This ensures that a MediaPipe graph or sub-graph
// can run multiple times, once per element in the "ITERABLE" for each pakcet
// clone of the packets in the "CLONE" input streams.
// streams at loop-internal timestamps. This ensures that a MediaPipe graph or
// sub-graph can run multiple times, once per element in the "ITERABLE" for each
// packet clone of the packets in the "CLONE" input streams. Think of CLONEd
// inputs as loop-wide constants.
template <typename IterableT>
class BeginLoopCalculator : public CalculatorBase {
using ItemT = typename IterableT::value_type;

View File

@ -19,10 +19,12 @@
#include "mediapipe/framework/formats/classification.pb.h"
#include "mediapipe/framework/formats/detection.pb.h"
#include "mediapipe/framework/formats/image.h"
#include "mediapipe/framework/formats/image_frame.h"
#include "mediapipe/framework/formats/landmark.pb.h"
#include "mediapipe/framework/formats/matrix.h"
#include "mediapipe/framework/formats/rect.pb.h"
#include "mediapipe/framework/formats/tensor.h"
#include "mediapipe/gpu/gpu_buffer.h"
#include "mediapipe/util/render_data.pb.h"
#include "tensorflow/lite/interpreter.h"
@ -68,6 +70,12 @@ REGISTER_CALCULATOR(EndLoopMatrixCalculator);
typedef EndLoopCalculator<std::vector<Tensor>> EndLoopTensorCalculator;
REGISTER_CALCULATOR(EndLoopTensorCalculator);
typedef EndLoopCalculator<std::vector<ImageFrame>> EndLoopImageFrameCalculator;
REGISTER_CALCULATOR(EndLoopImageFrameCalculator);
typedef EndLoopCalculator<std::vector<GpuBuffer>> EndLoopGpuBufferCalculator;
REGISTER_CALCULATOR(EndLoopGpuBufferCalculator);
typedef EndLoopCalculator<std::vector<::mediapipe::Image>>
EndLoopImageCalculator;
REGISTER_CALCULATOR(EndLoopImageCalculator);

View File

@ -17,13 +17,11 @@
#include <type_traits>
#include "absl/status/status.h"
#include "mediapipe/framework/calculator_context.h"
#include "mediapipe/framework/calculator_contract.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/collection_item_id.h"
#include "mediapipe/framework/port/integral_types.h"
#include "mediapipe/framework/port/ret_check.h"
#include "mediapipe/framework/port/status.h"
namespace mediapipe {
@ -33,27 +31,7 @@ namespace mediapipe {
// from the "BATCH_END" tagged input stream, it emits the aggregated results
// at the original timestamp contained in the "BATCH_END" input stream.
//
// It is designed to be used like:
//
// node {
// calculator: "BeginLoopWithIterableCalculator"
// input_stream: "ITERABLE:input_iterable" # IterableT @ext_ts
// output_stream: "ITEM:input_element" # ItemT @loop_internal_ts
// output_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
// }
//
// node {
// calculator: "ElementToBlaConverterSubgraph"
// input_stream: "ITEM:input_to_loop_body" # ItemT @loop_internal_ts
// output_stream: "BLA:output_of_loop_body" # ItemU @loop_internal_ts
// }
//
// node {
// calculator: "EndLoopWithOutputCalculator"
// input_stream: "ITEM:output_of_loop_body" # ItemU @loop_internal_ts
// input_stream: "BATCH_END:ext_ts" # Timestamp @loop_internal_ts
// output_stream: "ITERABLE:aggregated_result" # IterableU @ext_ts
// }
// See BeginLoopCalculator for a usage example.
template <typename IterableT>
class EndLoopCalculator : public CalculatorBase {
using ItemT = typename IterableT::value_type;
@ -79,7 +57,7 @@ class EndLoopCalculator : public CalculatorBase {
}
// Try to consume the item and move it into the collection. If the items
// are not consumable, then try to copy them instead. If the items are
// not copiable, then an error will be returned.
// not copyable, then an error will be returned.
auto item_ptr_or = cc->Inputs().Tag("ITEM").Value().Consume<ItemT>();
if (item_ptr_or.ok()) {
input_stream_collection_->push_back(std::move(*item_ptr_or.value()));