From fc9538533c16751ca5e8df754441aa54bb79a9e6 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Mon, 15 May 2023 05:53:16 -0700 Subject: [PATCH] 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 --- mediapipe/calculators/core/BUILD | 11 ++-- .../calculators/core/begin_loop_calculator.cc | 12 +++++ .../calculators/core/begin_loop_calculator.h | 54 +++++++++++-------- .../calculators/core/end_loop_calculator.cc | 8 +++ .../calculators/core/end_loop_calculator.h | 28 ++-------- 5 files changed, 62 insertions(+), 51 deletions(-) diff --git a/mediapipe/calculators/core/BUILD b/mediapipe/calculators/core/BUILD index d5ba6c74f..3294c0383 100644 --- a/mediapipe/calculators/core/BUILD +++ b/mediapipe/calculators/core/BUILD @@ -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, diff --git a/mediapipe/calculators/core/begin_loop_calculator.cc b/mediapipe/calculators/core/begin_loop_calculator.cc index 441c66937..ac74bb382 100644 --- a/mediapipe/calculators/core/begin_loop_calculator.cc +++ b/mediapipe/calculators/core/begin_loop_calculator.cc @@ -17,10 +17,12 @@ #include #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> BeginLoopTensorCalculator; REGISTER_CALCULATOR(BeginLoopTensorCalculator); +// A calculator to process std::vector. +typedef BeginLoopCalculator> + BeginLoopImageFrameCalculator; +REGISTER_CALCULATOR(BeginLoopImageFrameCalculator); + +// A calculator to process std::vector. +typedef BeginLoopCalculator> + BeginLoopGpuBufferCalculator; +REGISTER_CALCULATOR(BeginLoopGpuBufferCalculator); + } // namespace mediapipe diff --git a/mediapipe/calculators/core/begin_loop_calculator.h b/mediapipe/calculators/core/begin_loop_calculator.h index 81fff39da..c0b3022d4 100644 --- a/mediapipe/calculators/core/begin_loop_calculator.h +++ b/mediapipe/calculators/core/begin_loop_calculator.h @@ -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 and vector. +// First, instantiate specializations in the loop calculators' implementations +// if missing: +// BeginLoopInputTCalculator = BeginLoopCalculator +// EndLoopOutputTCalculator = EndLoopCalculator +// 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 class BeginLoopCalculator : public CalculatorBase { using ItemT = typename IterableT::value_type; diff --git a/mediapipe/calculators/core/end_loop_calculator.cc b/mediapipe/calculators/core/end_loop_calculator.cc index b3b889ecd..dea2b6cad 100644 --- a/mediapipe/calculators/core/end_loop_calculator.cc +++ b/mediapipe/calculators/core/end_loop_calculator.cc @@ -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> EndLoopTensorCalculator; REGISTER_CALCULATOR(EndLoopTensorCalculator); +typedef EndLoopCalculator> EndLoopImageFrameCalculator; +REGISTER_CALCULATOR(EndLoopImageFrameCalculator); + +typedef EndLoopCalculator> EndLoopGpuBufferCalculator; +REGISTER_CALCULATOR(EndLoopGpuBufferCalculator); + typedef EndLoopCalculator> EndLoopImageCalculator; REGISTER_CALCULATOR(EndLoopImageCalculator); diff --git a/mediapipe/calculators/core/end_loop_calculator.h b/mediapipe/calculators/core/end_loop_calculator.h index 2598194e6..1e258f046 100644 --- a/mediapipe/calculators/core/end_loop_calculator.h +++ b/mediapipe/calculators/core/end_loop_calculator.h @@ -17,13 +17,11 @@ #include +#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 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(); if (item_ptr_or.ok()) { input_stream_collection_->push_back(std::move(*item_ptr_or.value()));