Internal change

PiperOrigin-RevId: 477538515
2022-09-28 21:32:36 +00:00
parent 6cdc6443b6
commit f8af41b1eb
236 changed files with 12865 additions and 1923 deletions
@@ -55,6 +55,14 @@ mediapipe_proto_library(
 cc_library(
    name = "audio_to_tensor_calculator",
    srcs = ["audio_to_tensor_calculator.cc"],
+    copts = select({
+        # b/215212850
+        "//mediapipe:apple": [
+            "-x objective-c++",
+            "-fobjc-arc",
+        ],
+        "//conditions:default": [],
+    }),
    visibility = [
        "//mediapipe/framework:mediapipe_internal",
    ],
@@ -67,13 +75,16 @@ cc_library(
        "//mediapipe/framework/formats:matrix",
        "//mediapipe/framework/formats:tensor",
        "//mediapipe/framework/formats:time_series_header_cc_proto",
+        "//mediapipe/framework/port:ret_check",
        "//mediapipe/util:time_series_util",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_audio_tools//audio/dsp:resampler_q",
+        "@com_google_audio_tools//audio/dsp:window_functions",
        "@org_tensorflow//tensorflow/lite/c:common",
+        "@pffft",
    ],
    alwayslink = 1,
 )
@@ -83,6 +94,7 @@ cc_test(
    srcs = ["audio_to_tensor_calculator_test.cc"],
    deps = [
        ":audio_to_tensor_calculator",
+        ":audio_to_tensor_calculator_cc_proto",
        "//mediapipe/framework:calculator_cc_proto",
        "//mediapipe/framework:calculator_framework",
        "//mediapipe/framework:timestamp",
@@ -97,6 +109,58 @@ cc_test(
    ],
 )

+mediapipe_proto_library(
+    name = "feedback_tensors_calculator_proto",
+    srcs = ["feedback_tensors_calculator.proto"],
+    visibility = [
+        "//mediapipe/framework:mediapipe_internal",
+    ],
+    deps = [
+        "//mediapipe/framework:calculator_options_proto",
+        "//mediapipe/framework:calculator_proto",
+    ],
+)
+
+cc_library(
+    name = "feedback_tensors_calculator",
+    srcs = ["feedback_tensors_calculator.cc"],
+    copts = select({
+        # b/215212850
+        "//mediapipe:apple": [
+            "-x objective-c++",
+            "-fobjc-arc",
+        ],
+        "//conditions:default": [],
+    }),
+    visibility = [
+        "//mediapipe/framework:mediapipe_internal",
+    ],
+    deps = [
+        ":feedback_tensors_calculator_cc_proto",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework/api2:node",
+        "//mediapipe/framework/formats:tensor",
+        "@com_google_absl//absl/status",
+    ],
+    alwayslink = 1,
+)
+
+cc_test(
+    name = "feedback_tensors_calculator_test",
+    srcs = ["feedback_tensors_calculator_test.cc"],
+    deps = [
+        ":feedback_tensors_calculator",
+        ":feedback_tensors_calculator_cc_proto",
+        "//mediapipe/framework:calculator_cc_proto",
+        "//mediapipe/framework:calculator_framework",
+        "//mediapipe/framework:timestamp",
+        "//mediapipe/framework/formats:tensor",
+        "//mediapipe/framework/port:gtest_main",
+        "//mediapipe/framework/port:parse_text_proto",
+        "@org_tensorflow//tensorflow/lite/c:common",
+    ],
+)
+
 mediapipe_proto_library(
    name = "inference_calculator_proto",
    srcs = ["inference_calculator.proto"],
@@ -346,6 +410,10 @@ cc_library(
    }),
 )

+# This target provides the InferenceCalculator and a default set of implementations tailored for the
+# current build platforms. More implementations can be added as separate dependencies to a client;
+# for clients that want a narrower set of implementations than the default should see the comment on
+# inference_calculator_interface.
 cc_library(
    name = "inference_calculator",
    visibility = ["//visibility:public"],
@@ -12,9 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

-#include <math.h>
-
 #include <algorithm>
+#include <cmath>
 #include <cstring>
 #include <memory>
 #include <string>
@@ -26,6 +25,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "audio/dsp/resampler_q.h"
+#include "audio/dsp/window_functions.h"
 #include "mediapipe/calculators/tensor/audio_to_tensor_calculator.pb.h"
 #include "mediapipe/framework/api2/node.h"
 #include "mediapipe/framework/api2/packet.h"
@@ -34,19 +34,60 @@
 #include "mediapipe/framework/formats/matrix.h"
 #include "mediapipe/framework/formats/tensor.h"
 #include "mediapipe/framework/formats/time_series_header.pb.h"
+#include "mediapipe/framework/port/ret_check.h"
 #include "mediapipe/util/time_series_util.h"
+#include "pffft.h"

 namespace mediapipe {
 namespace api2 {
+namespace {
+
+using Options = ::mediapipe::AudioToTensorCalculatorOptions;
+using FlushMode = Options::FlushMode;
+
+std::vector<float> HannWindow(int window_size, bool sqrt_hann) {
+  std::vector<float> hann_window(window_size);
+  audio_dsp::HannWindow().GetPeriodicSamples(window_size, &hann_window);
+  if (sqrt_hann) {
+    absl::c_transform(hann_window, hann_window.begin(),
+                      [](double x) { return std::sqrt(x); });
+  }
+  return hann_window;
+}
+
+// PFFFT only supports transforms for inputs of length N of the form
+// N = (2^a)*(3^b)*(5^c) where b >=0 and c >= 0 and a >= 5 for the real FFT.
+bool IsValidFftSize(int size) {
+  if (size <= 0) {
+    return false;
+  }
+  constexpr int kFactors[] = {2, 3, 5};
+  int factorization[] = {0, 0, 0};
+  int n = static_cast<int>(size);
+  for (int i = 0; i < 3; ++i) {
+    while (n % kFactors[i] == 0) {
+      n = n / kFactors[i];
+      ++factorization[i];
+    }
+  }
+  return factorization[0] >= 5 && n == 1;
+}
+
+}  // namespace

 // Converts audio buffers into tensors, possibly with resampling, buffering
 // and framing, according to specified inputs and options. All input audio
 // buffers will be first resampled from the input sample rate to the target
 // sample rate if they are not equal. The resampled audio data (with the
 // buffered samples from the previous runs in the streaming mode) will be broken
-// into fixed-sized, possibly overlapping frames. Finally, all frames will be
-// converted to and outputted as MediaPipe Tensors. The last output tensor will
-// be zero-padding if the remaining samples are insufficient.
+// into fixed-sized, possibly overlapping frames. If the calculator is not asked
+// to perform fft (the fft_size is not set in the calculator options), all
+// frames will be converted to and outputted as MediaPipe Tensors. The last
+// output tensor will be zero-padding if the remaining samples are insufficient.
+// Otherwise, when the fft_size is set and valid, the calculator will perform
+// fft on the fixed-sized audio frames, the complex DFT results will be
+// converted to and outputted as 2D MediaPipe float Tensors where the first
+// rows are the DFT real parts and the second rows are the DFT imagery parts.
 //
 // This calculator assumes that the input timestamps refer to the first
 // sample in each Matrix. The output timestamps follow this same convention.
@@ -86,11 +127,15 @@ namespace api2 {
 // Outputs:
 //   TENSORS - std::vector<Tensor>
 //     Vector containing a single Tensor that represents a fix-sized audio
-//     frame.
+//     frame or the complex DFT results.
 //   TIMESTAMPS - std::vector<Timestamp> @Optional
 //     Vector containing the output timestamps emitted by the current Process()
 //     invocation. In the non-streaming mode, the vector contains all of the
 //     output timestamps for an input audio buffer.
+//   DC_AND_NYQUIST - std::pair<float, float> @Optional.
+//     A pair of dc component and nyquest component. Only can be connected when
+//     the calculator performs fft (the fft_size is set in the calculator
+//     options).
 //
 // Example:
 // node {
@@ -116,12 +161,14 @@ class AudioToTensorCalculator : public Node {
  // such as sample rate.
  static constexpr Input<double>::Optional kAudioSampleRateIn{"SAMPLE_RATE"};
  static constexpr Output<std::vector<Tensor>> kTensorsOut{"TENSORS"};
+  static constexpr Output<std::pair<float, float>>::Optional kDcAndNyquistOut{
+      "DC_AND_NYQUIST"};
  // A vector of the output timestamps emitted by the current Process()
  // invocation. The packet timestamp is the last emitted timestamp.
  static constexpr Output<std::vector<Timestamp>>::Optional kTimestampsOut{
      "TIMESTAMPS"};
  MEDIAPIPE_NODE_CONTRACT(kAudioIn, kAudioSampleRateIn, kTensorsOut,
-                          kTimestampsOut);
+                          kDcAndNyquistOut, kTimestampsOut);

  static absl::Status UpdateContract(CalculatorContract* cc);
  absl::Status Open(CalculatorContext* cc);
@@ -138,6 +185,9 @@ class AudioToTensorCalculator : public Node {
  int frame_step_;
  bool stream_mode_;
  bool check_inconsistent_timestamps_;
+  int padding_samples_before_;
+  int padding_samples_after_;
+  FlushMode flush_mode_;
  Timestamp initial_timestamp_ = Timestamp::Unstarted();
  int64 cumulative_input_samples_ = 0;
  Timestamp next_output_timestamp_ = Timestamp::Unstarted();
@@ -151,22 +201,33 @@ class AudioToTensorCalculator : public Node {
  Matrix sample_buffer_;
  int processed_buffer_cols_ = 0;

+  // The internal state of the FFT library.
+  PFFFT_Setup* fft_state_ = nullptr;
+  int fft_size_ = 0;
+  std::vector<float> fft_window_;
+  std::vector<float, Eigen::aligned_allocator<float>> fft_input_buffer_;
+  // pffft requires memory to work with to avoid using the stack.
+  std::vector<float, Eigen::aligned_allocator<float>> fft_workplace_;
+  std::vector<float, Eigen::aligned_allocator<float>> fft_output_;
+
  absl::Status ProcessStreamingData(CalculatorContext* cc, const Matrix& input);
  absl::Status ProcessNonStreamingData(CalculatorContext* cc,
                                       const Matrix& input);

  absl::Status SetupStreamingResampler(double input_sample_rate_);
  void AppendToSampleBuffer(Matrix buffer_to_append);
+  void AppendZerosToSampleBuffer(int num_samples);

  absl::StatusOr<std::vector<Tensor>> ConvertToTensor(
-      const Matrix& frame_to_convert);
-  absl::Status OutputTensors(const Matrix& buffer, bool should_flush,
+      const Matrix& block, std::vector<int> tensor_dims);
+  absl::Status OutputTensor(const Matrix& block, Timestamp timestamp,
+                            CalculatorContext* cc);
+  absl::Status ProcessBuffer(const Matrix& buffer, bool should_flush,
                             CalculatorContext* cc);
 };

 absl::Status AudioToTensorCalculator::UpdateContract(CalculatorContract* cc) {
-  const auto& options =
-      cc->Options<mediapipe::AudioToTensorCalculatorOptions>();
+  const auto& options = cc->Options<Options>();
  if (!options.has_num_channels() || !options.has_num_samples() ||
      !options.has_target_sample_rate()) {
    return absl::InvalidArgumentError(
@@ -174,13 +235,21 @@ absl::Status AudioToTensorCalculator::UpdateContract(CalculatorContract* cc) {
        "`num_channels`, `num_samples`, and `target_sample_rate`.");
  }
  if (options.stream_mode()) {
-    // Explicitly disables tiemstamp offset to disallow the timestamp bound
+    // Explicitly disables timestamp offset to disallow the timestamp bound
    // from the input streams to be propagated to the output streams.
    // In the streaming mode, the output timestamp bound is based on
    // next_output_timestamp_, which can be smaller than the current input
    // timestamps.
    cc->SetTimestampOffset(TimestampDiff::Unset());
  }
+  if (options.padding_samples_before() < 0 ||
+      options.padding_samples_after() < 0) {
+    return absl::InvalidArgumentError("Negative zero padding unsupported");
+  }
+  if (options.flush_mode() != Options::ENTIRE_TAIL_AT_TIMESTAMP_MAX &&
+      options.flush_mode() != Options::PROCEED_AS_USUAL) {
+    return absl::InvalidArgumentError("Unsupported flush mode");
+  }
  return absl::OkStatus();
 }

@@ -202,6 +271,9 @@ absl::Status AudioToTensorCalculator::Open(CalculatorContext* cc) {
    check_inconsistent_timestamps_ = options.check_inconsistent_timestamps();
    sample_buffer_.resize(num_channels_, Eigen::NoChange);
  }
+  padding_samples_before_ = options.padding_samples_before();
+  padding_samples_after_ = options.padding_samples_after();
+  flush_mode_ = options.flush_mode();

  RET_CHECK(kAudioSampleRateIn(cc).IsConnected() ^
            !kAudioIn(cc).Header().IsEmpty())
@@ -217,6 +289,25 @@ absl::Status AudioToTensorCalculator::Open(CalculatorContext* cc) {
      source_sample_rate_ = input_header.sample_rate();
    }
  }
+  AppendZerosToSampleBuffer(padding_samples_before_);
+  if (options.has_fft_size()) {
+    RET_CHECK(IsValidFftSize(options.fft_size()))
+        << "FFT size must be of the form fft_size = (2^a)*(3^b)*(5^c) where b "
+           ">=0 and c >= 0 and a >= 5, the requested fft size is "
+        << options.fft_size();
+    RET_CHECK_EQ(1, num_channels_)
+        << "Currently only support applying FFT on mono channel.";
+    fft_size_ = options.fft_size();
+    fft_state_ = pffft_new_setup(fft_size_, PFFFT_REAL);
+    fft_window_ = HannWindow(fft_size_, /* sqrt_hann = */ false);
+    fft_input_buffer_.resize(fft_size_);
+    fft_workplace_.resize(fft_size_);
+    fft_output_.resize(fft_size_);
+  } else {
+    RET_CHECK(!kDcAndNyquistOut(cc).IsConnected())
+        << "The DC_AND_NYQUIST output stream can only be connected when the "
+           "calculator outputs fft tensors";
+  }
  return absl::OkStatus();
 }

@@ -262,7 +353,12 @@ absl::Status AudioToTensorCalculator::Close(CalculatorContext* cc) {
    resampler_->Flush(&resampled_buffer);
    AppendToSampleBuffer(std::move(resampled_buffer));
  }
-  return OutputTensors(sample_buffer_, /*should_flush=*/true, cc);
+  AppendZerosToSampleBuffer(padding_samples_after_);
+  MP_RETURN_IF_ERROR(ProcessBuffer(sample_buffer_, /*should_flush=*/true, cc));
+  if (fft_state_) {
+    pffft_destroy_setup(fft_state_);
+  }
+  return absl::OkStatus();
 }

 absl::Status AudioToTensorCalculator::ProcessStreamingData(
@@ -303,7 +399,7 @@ absl::Status AudioToTensorCalculator::ProcessStreamingData(
    }
  }

-  MP_RETURN_IF_ERROR(OutputTensors(sample_buffer_, /*should_flush=*/false, cc));
+  MP_RETURN_IF_ERROR(ProcessBuffer(sample_buffer_, /*should_flush=*/false, cc));
  // Removes the processed samples from the global sample buffer.
  sample_buffer_ = Matrix(sample_buffer_.rightCols(sample_buffer_.cols() -
                                                   processed_buffer_cols_ - 1));
@@ -323,9 +419,9 @@ absl::Status AudioToTensorCalculator::ProcessNonStreamingData(
        input_frame);
    Eigen::Map<const Matrix> matrix_mapping(resampled.data(), num_channels_,
                                            resampled.size() / num_channels_);
-    return OutputTensors(matrix_mapping, /*should_flush=*/true, cc);
+    return ProcessBuffer(matrix_mapping, /*should_flush=*/true, cc);
  }
-  return OutputTensors(input_frame, /*should_flush=*/true, cc);
+  return ProcessBuffer(input_frame, /*should_flush=*/true, cc);
 }

 absl::Status AudioToTensorCalculator::SetupStreamingResampler(
@@ -344,6 +440,16 @@ absl::Status AudioToTensorCalculator::SetupStreamingResampler(
  return absl::OkStatus();
 }

+void AudioToTensorCalculator::AppendZerosToSampleBuffer(int num_samples) {
+  CHECK_GE(num_samples, 0);  // Ensured by `UpdateContract`.
+  if (num_samples == 0) {
+    return;
+  }
+  sample_buffer_.conservativeResize(Eigen::NoChange,
+                                    sample_buffer_.cols() + num_samples);
+  sample_buffer_.rightCols(num_samples).setZero();
+}
+
 void AudioToTensorCalculator::AppendToSampleBuffer(Matrix buffer_to_append) {
  sample_buffer_.conservativeResize(
      Eigen::NoChange, sample_buffer_.cols() + buffer_to_append.cols());
@@ -351,49 +457,89 @@ void AudioToTensorCalculator::AppendToSampleBuffer(Matrix buffer_to_append) {
 }

 absl::StatusOr<std::vector<Tensor>> AudioToTensorCalculator::ConvertToTensor(
-    const Matrix& frame_to_convert) {
-  Tensor tensor(Tensor::ElementType::kFloat32,
-                Tensor::Shape({num_channels_, num_samples_}));
+    const Matrix& block, std::vector<int> tensor_dims) {
+  Tensor tensor(Tensor::ElementType::kFloat32, Tensor::Shape(tensor_dims));
  auto buffer_view = tensor.GetCpuWriteView();
-  if (frame_to_convert.size() < num_channels_ * num_samples_) {
+  int total_size = 1;
+  for (int dim : tensor_dims) {
+    total_size *= dim;
+  }
+  if (block.size() < total_size) {
    std::memset(buffer_view.buffer<float>(), 0, tensor.bytes());
  }
-  std::memcpy(buffer_view.buffer<float>(), frame_to_convert.data(),
-              frame_to_convert.size() * sizeof(float));
+  std::memcpy(buffer_view.buffer<float>(), block.data(),
+              block.size() * sizeof(float));
  std::vector<Tensor> tensor_vector;
  tensor_vector.push_back(std::move(tensor));
  return tensor_vector;
 }

-absl::Status AudioToTensorCalculator::OutputTensors(const Matrix& buffer,
+absl::Status AudioToTensorCalculator::OutputTensor(const Matrix& block,
+                                                   Timestamp timestamp,
+                                                   CalculatorContext* cc) {
+  std::vector<Tensor> output_tensor;
+  if (fft_state_) {
+    Eigen::VectorXf time_series_data =
+        Eigen::VectorXf::Map(block.data(), block.size());
+    //  Window on input audio prior to FFT.
+    std::transform(time_series_data.begin(), time_series_data.end(),
+                   fft_window_.begin(), fft_input_buffer_.begin(),
+                   std::multiplies<float>());
+    pffft_transform_ordered(fft_state_, fft_input_buffer_.data(),
+                            fft_output_.data(), fft_workplace_.data(),
+                            PFFFT_FORWARD);
+    if (kDcAndNyquistOut(cc).IsConnected()) {
+      kDcAndNyquistOut(cc).Send(std::make_pair(fft_output_[0], fft_output_[1]),
+                                timestamp);
+    }
+    Matrix fft_output_matrix =
+        Eigen::Map<const Matrix>(fft_output_.data() + 2, 1, fft_size_ - 2);
+    fft_output_matrix.conservativeResize(Eigen::NoChange, fft_size_);
+    // The last two elements are the DFT Nyquist values.
+    fft_output_matrix(fft_size_ - 2) = fft_output_[1];  // Nyquist real part
+    fft_output_matrix(fft_size_ - 1) = 0.0f;            // Nyquist imagery part
+    ASSIGN_OR_RETURN(output_tensor,
+                     ConvertToTensor(fft_output_matrix, {2, fft_size_ / 2}));
+  } else {
+    ASSIGN_OR_RETURN(output_tensor,
+                     ConvertToTensor(block, {num_channels_, num_samples_}));
+  }
+  kTensorsOut(cc).Send(std::move(output_tensor), timestamp);
+  return absl::OkStatus();
+}
+
+absl::Status AudioToTensorCalculator::ProcessBuffer(const Matrix& buffer,
                                                    bool should_flush,
                                                    CalculatorContext* cc) {
+  const bool should_flush_at_timestamp_max =
+      stream_mode_ && should_flush &&
+      flush_mode_ == Options::ENTIRE_TAIL_AT_TIMESTAMP_MAX;
  int next_frame_first_col = 0;
  std::vector<Timestamp> timestamps;
-  while ((!stream_mode_ || !should_flush) &&
-         next_frame_first_col + num_samples_ <= buffer.cols()) {
-    ASSIGN_OR_RETURN(auto output_tensor, ConvertToTensor(buffer.block(
-                                             0, next_frame_first_col,
-                                             num_channels_, num_samples_)));
-    kTensorsOut(cc).Send(std::move(output_tensor), next_output_timestamp_);
-    timestamps.push_back(next_output_timestamp_);
-    next_output_timestamp_ += round(frame_step_ / target_sample_rate_ *
-                                    Timestamp::kTimestampUnitsPerSecond);
-    next_frame_first_col += frame_step_;
+  if (!should_flush_at_timestamp_max) {
+    while (next_frame_first_col + num_samples_ <= buffer.cols()) {
+      MP_RETURN_IF_ERROR(OutputTensor(
+          buffer.block(0, next_frame_first_col, num_channels_, num_samples_),
+          next_output_timestamp_, cc));
+      timestamps.push_back(next_output_timestamp_);
+      next_output_timestamp_ += round(frame_step_ / target_sample_rate_ *
+                                      Timestamp::kTimestampUnitsPerSecond);
+      next_frame_first_col += frame_step_;
+    }
  }
  if (should_flush && next_frame_first_col < buffer.cols()) {
-    ASSIGN_OR_RETURN(auto output_tensor,
-                     ConvertToTensor(buffer.block(
-                         0, next_frame_first_col, num_channels_,
-                         std::min(num_samples_,
-                                  (int)buffer.cols() - next_frame_first_col))));
    // In the streaming mode, the flush happens in Close() and a packet at
    // Timestamp::Max() will be emitted. In the non-streaming mode, each
    // Process() invocation will process the entire buffer completely.
-    Timestamp timestamp =
-        stream_mode_ ? Timestamp::Max() : next_output_timestamp_;
+    Timestamp timestamp = should_flush_at_timestamp_max
+                              ? Timestamp::Max()
+                              : next_output_timestamp_;
+    MP_RETURN_IF_ERROR(OutputTensor(
+        buffer.block(
+            0, next_frame_first_col, num_channels_,
+            std::min(num_samples_, (int)buffer.cols() - next_frame_first_col)),
+        timestamp, cc));
    timestamps.push_back(timestamp);
-    kTensorsOut(cc).Send(std::move(output_tensor), timestamp);
  }
  if (kTimestampsOut(cc).IsConnected()) {
    Timestamp timestamp = timestamps.back();
@@ -44,4 +44,28 @@ message AudioToTensorCalculatorOptions {
  // Set to false to disable checks for jitter in timestamp values. Useful with
  // live audio input.
  optional bool check_inconsistent_timestamps = 6 [default = true];
+
+  // Size of the fft in number of bins. If set, the calculator outputs fft
+  // tensors.
+  optional int64 fft_size = 7;
+
+  // The amount of padding samples to add before the audio after resampling.
+  // Note that the timestamps shift. Currently, only zero padding is supported.
+  optional int64 padding_samples_before = 8;
+
+  // The amount of padding samples to add after the audio after resampling.
+  // Currently, only zero padding is supported.
+  optional int64 padding_samples_after = 9;
+
+  // Determines the "flushing" behavior in stream mode.
+  enum FlushMode {
+    // Unspecified (causes an error). Won't be used because of the default.
+    NONE = 0;
+    // Emit a packet with the entire remainder at `Timestamp::Max`.
+    ENTIRE_TAIL_AT_TIMESTAMP_MAX = 1;
+    // Continue emitting framed packets with relevant timestamps.
+    PROCEED_AS_USUAL = 2;
+  }
+
+  optional FlushMode flush_mode = 10 [default = ENTIRE_TAIL_AT_TIMESTAMP_MAX];
 }
@@ -12,13 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.

-#include <cmath>
 #include <memory>
 #include <string>
 #include <vector>

 #include "absl/strings/substitute.h"
 #include "audio/dsp/resampler_q.h"
+#include "mediapipe/calculators/tensor/audio_to_tensor_calculator.pb.h"
 #include "mediapipe/framework/api2/packet.h"
 #include "mediapipe/framework/calculator.pb.h"
 #include "mediapipe/framework/calculator_framework.h"
@@ -32,6 +32,14 @@
 namespace mediapipe {
 namespace {

+using ::testing::Not;
+using Options = ::mediapipe::AudioToTensorCalculatorOptions;
+using FlushMode = Options::FlushMode;
+
+int DivideRoundedUp(int dividend, int divisor) {
+  return (dividend + divisor - 1) / divisor;
+}
+
 std::unique_ptr<Matrix> CreateTestMatrix(int num_channels, int num_samples,
                                         int timestamp) {
  auto matrix = std::make_unique<Matrix>(num_channels, num_samples);
@@ -292,16 +300,17 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
    num_iterations_ = num_iterations;
  }

-  int GetExpectedNumOfSamples() {
-    Matrix* expected_matrix =
-        resampled_buffer_ ? resampled_buffer_.get() : sample_buffer_.get();
-    return expected_matrix->cols();
-  }
+  int GetExpectedNumOfSamples() { return output_sample_buffer_->cols(); }

  void Run(int num_samples, int num_overlapping_samples,
-           double resampling_factor) {
+           double resampling_factor, int padding_before = 0,
+           int padding_after = 0, bool expect_init_error = false) {
    double input_sample_rate = 10000;
    double target_sample_rate = input_sample_rate * resampling_factor;
+    FlushMode flush_mode = (padding_before != 0 || padding_after != 0)
+                               ? Options::PROCEED_AS_USUAL
+                               : Options::ENTIRE_TAIL_AT_TIMESTAMP_MAX;
+
    auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(
        absl::Substitute(R"(
        input_stream: "audio"
@@ -319,16 +328,25 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
              num_overlapping_samples: $1
              target_sample_rate: $2
              stream_mode:true
+              padding_samples_before: $3
+              padding_samples_after: $4
+              flush_mode: $5
            }
          }
        }
        )",
                         /*$0=*/num_samples, /*$1=*/num_overlapping_samples,
-                         /*$2=*/target_sample_rate));
+                         /*$2=*/target_sample_rate, /*$3=*/padding_before,
+                         /*$4=*/padding_after, /*$5=*/flush_mode));
    tool::AddVectorSink("tensors", &graph_config, &tensors_packets_);

    // Run the graph.
-    MP_ASSERT_OK(graph_.Initialize(graph_config));
+    const absl::Status init_status = graph_.Initialize(graph_config);
+    if (expect_init_error) {
+      EXPECT_THAT(init_status, Not(IsOk()));
+      return;
+    }
+    MP_ASSERT_OK(init_status);
    MP_ASSERT_OK(graph_.StartRun({}));
    for (int i = 0; i < num_iterations_; ++i) {
      Timestamp input_timestamp(Timestamp::kTimestampUnitsPerSecond * i);
@@ -345,8 +363,18 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
    }
    MP_ASSERT_OK(graph_.CloseAllInputStreams());
    MP_ASSERT_OK(graph_.WaitUntilIdle());
-    if (resampling_factor != 1) {
-      resampled_buffer_ = ResampleBuffer(*sample_buffer_, resampling_factor);
+    if (resampling_factor == 1) {
+      output_sample_buffer_ = std::make_unique<Matrix>(*sample_buffer_);
+    } else {
+      output_sample_buffer_ =
+          ResampleBuffer(*sample_buffer_, resampling_factor);
+    }
+    if (padding_before != 0 || padding_after != 0) {
+      Matrix padded = Matrix::Zero(
+          2, padding_before + output_sample_buffer_->cols() + padding_after);
+      padded.block(0, padding_before, 2, output_sample_buffer_->cols()) =
+          *output_sample_buffer_;
+      output_sample_buffer_->swap(padded);
    }
  }

@@ -372,15 +400,13 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
    auto buffer = output_tensor.GetCpuReadView().buffer<float>();
    int num_values = output_tensor.shape().num_elements();
    std::vector<float> output_floats(buffer, buffer + num_values);
-    Matrix* expected_matrix =
-        resampled_buffer_ ? resampled_buffer_.get() : sample_buffer_.get();
    for (int i = 0; i < num_values; ++i) {
-      if (i + sample_offset >= expected_matrix->size()) {
+      if (i + sample_offset >= output_sample_buffer_->size()) {
        EXPECT_FLOAT_EQ(output_floats[i], 0);
      } else {
        EXPECT_NEAR(output_floats[i],
-                    expected_matrix->coeff((i + sample_offset) % 2,
-                                           (i + sample_offset) / 2),
+                    output_sample_buffer_->coeff((i + sample_offset) % 2,
+                                                 (i + sample_offset) / 2),
                    0.001)
            << "i=" << i << ", sample_offset=" << sample_offset
            << ", packet index=" << index;
@@ -391,7 +417,8 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {

  // Fully close graph at end, otherwise calculator+tensors are destroyed
  // after calling WaitUntilDone().
-  void CloseGraph() { MP_EXPECT_OK(graph_.WaitUntilDone()); }
+  absl::Status TryCloseGraph() { return graph_.WaitUntilDone(); }
+  void CloseGraph() { MP_EXPECT_OK(TryCloseGraph()); }

 private:
  int input_buffer_num_samples_ = 10;
@@ -399,7 +426,7 @@ class AudioToTensorCalculatorStreamingModeTest : public ::testing::Test {
  CalculatorGraph graph_;
  std::vector<Packet> tensors_packets_;
  std::unique_ptr<Matrix> sample_buffer_;
-  std::unique_ptr<Matrix> resampled_buffer_;
+  std::unique_ptr<Matrix> output_sample_buffer_;
 };

 TEST_F(AudioToTensorCalculatorStreamingModeTest,
@@ -408,7 +435,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest,
      /*resampling_factor=*/1.0f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/10,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 5),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 5),
      /*timestamp_interval=*/500,
      /*output_last_at_close=*/false);
  CloseGraph();
@@ -419,7 +446,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, OutputRemainingInCloseMethod) {
      /*resampling_factor=*/1.0f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/12,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 6),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 6),
      /*timestamp_interval=*/600,
      /*output_last_at_close=*/true);
  CloseGraph();
@@ -431,7 +458,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, OutputOverlappingFp32Tensors) {
      /*resampling_factor=*/1.0f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/16,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 8),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 8),
      /*timestamp_interval=*/800,
      /*output_last_at_close=*/true);
  CloseGraph();
@@ -443,7 +470,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, Downsampling) {
      /*resampling_factor=*/0.5f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/512,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 256),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 256),
      /*timestamp_interval=*/51200,
      /*output_last_at_close=*/true);
  CloseGraph();
@@ -455,7 +482,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, DownsamplingWithOverlapping) {
      /*resampling_factor=*/0.5f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/384,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 192),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 192),
      /*timestamp_interval=*/38400,
      /*output_last_at_close=*/true);
  CloseGraph();
@@ -467,7 +494,7 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, Upsampling) {
      /*resampling_factor=*/2.0f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/512,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 256),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 256),
      /*timestamp_interval=*/12800,
      /*output_last_at_close=*/true);
  CloseGraph();
@@ -479,12 +506,33 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest, UpsamplingWithOverlapping) {
      /*resampling_factor=*/2.0f);
  CheckTensorsOutputPackets(
      /*sample_offset=*/384,
-      /*num_packets=*/std::ceil((float)GetExpectedNumOfSamples() / 192),
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 192),
      /*timestamp_interval=*/9600,
      /*output_last_at_close=*/true);
  CloseGraph();
 }

+TEST_F(AudioToTensorCalculatorStreamingModeTest,
+       UpsamplingWithOverlappingAndPadding) {
+  SetInputBufferNumSamplesPerChannel(1024);
+  Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
+      /*resampling_factor=*/2.0f, /*padding_before=*/13, /*padding_after=*/999);
+  CheckTensorsOutputPackets(
+      /*sample_offset=*/384,
+      /*num_packets=*/DivideRoundedUp(GetExpectedNumOfSamples(), 192),
+      /*timestamp_interval=*/9600,
+      /*output_last_at_close=*/false);
+  CloseGraph();
+}
+
+TEST_F(AudioToTensorCalculatorStreamingModeTest, NegativePaddingUnsupported) {
+  SetInputBufferNumSamplesPerChannel(1024);
+  Run(/*num_samples=*/256, /*num_overlapping_samples=*/64,
+      /*resampling_factor=*/2.0f, /*padding_before=*/13, /*padding_after=*/-3,
+      /*expect_init_error=*/true);
+  EXPECT_THAT(TryCloseGraph(), Not(IsOk()));
+}
+
 TEST_F(AudioToTensorCalculatorStreamingModeTest,
       OnlyOutputInCloseIfNoSufficientSamples) {
  SetNumIterations(1);
@@ -498,5 +546,122 @@ TEST_F(AudioToTensorCalculatorStreamingModeTest,
  CloseGraph();
 }

+class AudioToTensorCalculatorFftTest : public ::testing::Test {
+ protected:
+  // Creates an audio matrix containing a single sample of 1.0 at a specified
+  // offset.
+  std::unique_ptr<Matrix> CreateImpulseSignalData(int64 num_samples,
+                                                  int impulse_offset_idx) {
+    Matrix impulse = Matrix::Zero(1, num_samples);
+    impulse(0, impulse_offset_idx) = 1.0;
+    return std::make_unique<Matrix>(std::move(impulse));
+  }
+
+  void ConfigGraph(int num_channels, int num_samples,
+                   int num_overlapping_samples, double sample_rate,
+                   int fft_size) {
+    graph_config_ = ParseTextProtoOrDie<CalculatorGraphConfig>(
+        absl::Substitute(R"(
+        input_stream: "audio"
+        input_stream: "sample_rate"
+        output_stream: "tensors"
+        output_stream: "dc_and_nyquist"
+        node {
+          calculator: "AudioToTensorCalculator"
+          input_stream: "AUDIO:audio"
+          input_stream: "SAMPLE_RATE:sample_rate"
+          output_stream: "TENSORS:tensors"
+          output_stream: "DC_AND_NYQUIST:dc_and_nyquist"
+          options {
+            [mediapipe.AudioToTensorCalculatorOptions.ext] {
+              num_channels: $0
+              num_samples: $1
+              num_overlapping_samples: $2
+              target_sample_rate: $3
+              fft_size: $4
+            }
+          }
+        }
+        )",
+                         /*$0=*/num_channels,
+                         /*$1=*/num_samples,
+                         /*$2=*/num_overlapping_samples,
+                         /*$3=*/sample_rate, /*$4=*/fft_size));
+    std::vector<Packet> tensors_packets;
+    tool::AddVectorSink("tensors", &graph_config_, &tensors_packets_);
+    std::vector<Packet> dc_and_nyquist_packets;
+    tool::AddVectorSink("dc_and_nyquist", &graph_config_,
+                        &dc_and_nyquist_packets_);
+  }
+
+  void RunGraph(std::unique_ptr<Matrix> input_data, double sample_rate) {
+    MP_ASSERT_OK(graph_.Initialize(graph_config_));
+    MP_ASSERT_OK(graph_.StartRun({}));
+    MP_ASSERT_OK(graph_.AddPacketToInputStream(
+        "sample_rate", MakePacket<double>(sample_rate).At(Timestamp(0))));
+    MP_ASSERT_OK(graph_.AddPacketToInputStream(
+        "audio", MakePacket<Matrix>(*input_data).At(Timestamp(0))));
+    MP_ASSERT_OK(graph_.CloseAllInputStreams());
+    MP_ASSERT_OK(graph_.WaitUntilIdle());
+    ASSERT_EQ(tensors_packets_.size(), dc_and_nyquist_packets_.size());
+  }
+
+  // Fully close graph at end, otherwise calculator+tensors are destroyed
+  // after calling WaitUntilDone().
+  void CloseGraph() { MP_EXPECT_OK(graph_.WaitUntilDone()); }
+
+  std::vector<Packet> tensors_packets_;
+  std::vector<Packet> dc_and_nyquist_packets_;
+  CalculatorGraphConfig graph_config_;
+  CalculatorGraph graph_;
+};
+
+TEST_F(AudioToTensorCalculatorFftTest, TestInvalidFftSize) {
+  ConfigGraph(1, 320, 160, 16000, 103);
+  MP_ASSERT_OK(graph_.Initialize(graph_config_));
+  MP_ASSERT_OK(graph_.StartRun({}));
+  auto status = graph_.WaitUntilIdle();
+  EXPECT_EQ(status.code(), absl::StatusCode::kInternal);
+  EXPECT_THAT(status.message(),
+              ::testing::HasSubstr("FFT size must be of the form"));
+}
+
+TEST_F(AudioToTensorCalculatorFftTest, TestInvalidNumChannels) {
+  ConfigGraph(3, 320, 160, 16000, 256);
+  MP_ASSERT_OK(graph_.Initialize(graph_config_));
+  MP_ASSERT_OK(graph_.StartRun({}));
+  auto status = graph_.WaitUntilIdle();
+  EXPECT_EQ(status.code(), absl::StatusCode::kInternal);
+  EXPECT_THAT(
+      status.message(),
+      ::testing::HasSubstr("only support applying FFT on mono channel"));
+}
+
+TEST_F(AudioToTensorCalculatorFftTest, TestImpulseSignal) {
+  constexpr double sample_rate = 16000;
+  ConfigGraph(1, 320, 160, sample_rate, 320);
+  RunGraph(CreateImpulseSignalData(320, 160), sample_rate);
+  for (int i = 0; i < tensors_packets_.size(); ++i) {
+    const auto& tensors = tensors_packets_[i].Get<std::vector<Tensor>>();
+    ASSERT_EQ(1, tensors.size());
+    const Tensor& output_tensor =
+        tensors_packets_[0].Get<std::vector<Tensor>>()[0];
+    auto* buffer = output_tensor.GetCpuReadView().buffer<float>();
+    int num_values = output_tensor.shape().num_elements();
+    const std::vector<float> output_floats(buffer, buffer + num_values);
+    // Impulse signal should have (approximately) const power across all
+    // frequency bins.
+    const auto& pair =
+        dc_and_nyquist_packets_[i].Get<std::pair<float, float>>();
+    EXPECT_FLOAT_EQ(pair.first, 1.0f);
+    EXPECT_FLOAT_EQ(pair.second, 1.0f);
+    for (int j = 0; j < num_values / 2; ++j) {
+      std::complex<float> cf(output_floats[j * 2], output_floats[j * 2 + 1]);
+      EXPECT_FLOAT_EQ(std::norm(cf), 1.0f);
+    }
+  }
+  CloseGraph();
+}
+
 }  // namespace
 }  // namespace mediapipe
@@ -0,0 +1,165 @@
+// Copyright 2022 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.
+
+#include <algorithm>
+#include <memory>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "mediapipe/calculators/tensor/feedback_tensors_calculator.pb.h"
+#include "mediapipe/framework/api2/node.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/tensor.h"
+
+namespace mediapipe {
+namespace api2 {
+
+namespace {
+constexpr char kInputTensorsTag[] = "INPUT_TENSORS";
+constexpr char kFeedbackTensorsTag[] = "FEEDBACK_TENSORS";
+constexpr char kOutputTensorsTag[] = "TENSORS";
+
+using Tensors = std::vector<Tensor>;
+}  // namespace
+
+// FeedbackTensorsCalculator groups the input and the feedback (typically
+// recurrent neural network cell state output tensors from the previous run)
+// tensor vectors as the input tensor vector for the next recurrent model cell
+// inference. On the first step, the feedback tensor is filled with zeros to
+// jumpstart the loop.
+class FeedbackTensorsCalculator : public Node {
+ public:
+  static constexpr Input<Tensors> kFeedbackTensorsIn{kFeedbackTensorsTag};
+  static constexpr Input<Tensors> kInputTensorsIn{kInputTensorsTag};
+  static constexpr Output<Tensors> kTensorsOut{kOutputTensorsTag};
+
+  MEDIAPIPE_NODE_CONTRACT(kFeedbackTensorsIn, kInputTensorsIn, kTensorsOut);
+
+  static absl::Status GetContract(CalculatorContract* cc) {
+    cc->SetProcessTimestampBounds(true);
+    return absl::OkStatus();
+  }
+
+  absl::Status Open(CalculatorContext* cc) override {
+    const auto& options =
+        cc->Options<mediapipe::FeedbackTensorsCalculatorOptions>();
+
+    const auto& shape_dims = options.feedback_tensor_shape().dims();
+    feedback_tensor_shape_.dims.assign(shape_dims.begin(), shape_dims.end());
+    feedback_tensor_size_ = feedback_tensor_shape_.num_elements();
+
+    num_feedback_tensors_ = options.num_feedback_tensors();
+
+    feedback_tensors_location_ = options.location();
+
+    return absl::OkStatus();
+  }
+
+  absl::Status Process(CalculatorContext* cc) override {
+    if (feedback_tensors_location_ ==
+        mediapipe::FeedbackTensorsCalculatorOptions::NONE) {
+      kTensorsOut(cc).Send(kInputTensorsIn(cc).packet().As<Tensors>());
+      return absl::OkStatus();
+    }
+
+    std::vector<Tensor> outputs;
+    switch (feedback_tensors_location_) {
+      case mediapipe::FeedbackTensorsCalculatorOptions::PREPENDED:
+        MP_RETURN_IF_ERROR(AddFeedbackTensors(cc, outputs));
+        MP_RETURN_IF_ERROR(AddInputTensors(cc, outputs));
+        break;
+      case mediapipe::FeedbackTensorsCalculatorOptions::APPENDED:
+        MP_RETURN_IF_ERROR(AddInputTensors(cc, outputs));
+        MP_RETURN_IF_ERROR(AddFeedbackTensors(cc, outputs));
+        break;
+      default:
+        return absl::InvalidArgumentError(
+            "Unsupported feedback tensors location");
+    }
+    kTensorsOut(cc).Send(std::move(outputs));
+    return absl::OkStatus();
+  }
+
+ private:
+  absl::Status AddInputTensors(CalculatorContext* cc,
+                               std::vector<Tensor>& outputs) {
+    absl::StatusOr<std::unique_ptr<std::vector<Tensor>>> input_tensors =
+        cc->Inputs()
+            .Tag(kInputTensorsTag)
+            .Value()
+            .Consume<std::vector<Tensor>>();
+    if (!input_tensors.ok()) {
+      return absl::InternalError("The input tensors packet is not consumable");
+    }
+    RET_CHECK(*input_tensors);
+    std::vector<Tensor>& inputs = **input_tensors;
+    outputs.insert(outputs.end(), std::make_move_iterator(inputs.begin()),
+                   std::make_move_iterator(inputs.end()));
+    return absl::OkStatus();
+  }
+
+  absl::Status AddFeedbackTensors(CalculatorContext* cc,
+                                  std::vector<Tensor>& outputs) {
+    if (first_run_) {
+      for (int index = 0; index < num_feedback_tensors_; ++index) {
+        Tensor initial_feedback_tensor(Tensor::ElementType::kFloat32,
+                                       feedback_tensor_shape_);
+        float* data = initial_feedback_tensor.GetCpuWriteView().buffer<float>();
+        std::fill_n(data, feedback_tensor_size_, 0.0f);
+        outputs.push_back(std::move(initial_feedback_tensor));
+      }
+      first_run_ = false;
+      return absl::OkStatus();
+    }
+
+    if (num_feedback_tensors_ != kFeedbackTensorsIn(cc)->size()) {
+      return absl::InvalidArgumentError(
+          "The number of tensors fed back differs from the configuration");
+    }
+    absl::StatusOr<std::unique_ptr<std::vector<Tensor>>> feedback_tensors =
+        cc->Inputs()
+            .Tag(kFeedbackTensorsTag)
+            .Value()
+            .Consume<std::vector<Tensor>>();
+    if (!feedback_tensors.ok()) {
+      return absl::InternalError(
+          "The feedback tensors packet is not consumable");
+    }
+    RET_CHECK(*feedback_tensors);
+    std::vector<Tensor>& feedbacks = **feedback_tensors;
+    for (const auto& feedback : feedbacks) {
+      if (feedback.shape().dims != feedback_tensor_shape_.dims) {
+        return absl::InvalidArgumentError(
+            "The shape of a tensor fed back differs from the configuration");
+      }
+    }
+    outputs.insert(outputs.end(), std::make_move_iterator(feedbacks.begin()),
+                   std::make_move_iterator(feedbacks.end()));
+
+    return absl::OkStatus();
+  }
+
+  Tensor::Shape feedback_tensor_shape_;
+  int num_feedback_tensors_ = 0;
+  mediapipe::FeedbackTensorsCalculatorOptions::FeedbackTensorsLocation
+      feedback_tensors_location_;
+
+  int feedback_tensor_size_ = 0;
+  bool first_run_ = true;
+};
+
+MEDIAPIPE_REGISTER_NODE(FeedbackTensorsCalculator);
+
+}  // namespace api2
+}  // namespace mediapipe
@@ -0,0 +1,47 @@
+// Copyright 2022 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.
+
+syntax = "proto2";
+
+package mediapipe;
+
+import "mediapipe/framework/calculator.proto";
+
+message FeedbackTensorsCalculatorOptions {
+  extend mediapipe.CalculatorOptions {
+    optional FeedbackTensorsCalculatorOptions ext = 474496252;
+  }
+
+  // Represents the dimensions of a tensor starting from the outermost size.
+  message TensorShape {
+    repeated int32 dims = 1 [packed = true];
+  }
+
+  // The shape of the feedback tensors to add.
+  optional TensorShape feedback_tensor_shape = 1;
+  // The number of the feedback tensors to add.
+  optional int32 num_feedback_tensors = 2 [default = 1];
+
+  enum FeedbackTensorsLocation {
+    // The feedback tensors will not be added.
+    NONE = 0;
+    // The feedback tensors will be added before the input tensors.
+    PREPENDED = 1;
+    // The feedback tensors will be added after the input tensors.
+    APPENDED = 2;
+  }
+
+  // Determines the location of the feedback tensor(s) in the output vector.
+  optional FeedbackTensorsLocation location = 3 [default = APPENDED];
+}
@@ -0,0 +1,389 @@
+// Copyright 2022 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.
+
+#include <functional>
+#include <initializer_list>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "mediapipe/calculators/tensor/feedback_tensors_calculator.pb.h"
+#include "mediapipe/framework/calculator.pb.h"
+#include "mediapipe/framework/calculator_framework.h"
+#include "mediapipe/framework/formats/tensor.h"
+#include "mediapipe/framework/port/gmock.h"
+#include "mediapipe/framework/port/gtest.h"
+#include "mediapipe/framework/port/parse_text_proto.h"
+#include "mediapipe/framework/port/status_matchers.h"
+#include "mediapipe/framework/timestamp.h"
+
+namespace mediapipe {
+namespace {
+
+using ::mediapipe::CalculatorGraphConfig;
+using ::testing::ElementsAreArray;
+using ::testing::Not;
+using Tensors = std::vector<Tensor>;
+
+template <typename T>
+struct TensorElementType {
+  static constexpr Tensor::ElementType value = Tensor::ElementType::kNone;
+};
+
+template <>
+struct TensorElementType<float> {
+  static constexpr Tensor::ElementType value = Tensor::ElementType::kFloat32;
+};
+
+template <>
+struct TensorElementType<std::int8_t> {
+  static constexpr Tensor::ElementType value = Tensor::ElementType::kInt8;
+};
+
+template <>
+struct TensorElementType<std::uint8_t> {
+  static constexpr Tensor::ElementType value = Tensor::ElementType::kUInt8;
+};
+
+template <>
+struct TensorElementType<std::int32_t> {
+  static constexpr Tensor::ElementType value = Tensor::ElementType::kInt32;
+};
+
+template <typename T>
+Tensor MakeTensor(std::initializer_list<int> shape,
+                  std::initializer_list<T> values) {
+  Tensor tensor(TensorElementType<T>::value, shape);
+  CHECK_EQ(values.size(), tensor.shape().num_elements())
+      << "The size of `values` is incompatible with `shape`";
+  absl::c_copy(values, tensor.GetCpuWriteView().buffer<T>());
+  return tensor;
+}
+
+template <typename T>
+void ValidateTensor(const Tensor& tensor,
+                    const std::vector<int>& expected_shape,
+                    const std::vector<T>& expected_values) {
+  ASSERT_EQ(tensor.element_type(), TensorElementType<T>::value);
+  EXPECT_EQ(tensor.shape().dims, expected_shape);
+  EXPECT_EQ(tensor.shape().num_elements(), expected_values.size());
+
+  auto* tensor_buffer = tensor.GetCpuReadView().buffer<T>();
+  const std::vector<T> tensor_values(
+      tensor_buffer, tensor_buffer + tensor.shape().num_elements());
+  EXPECT_THAT(tensor_values, ElementsAreArray(expected_values));
+}
+
+TEST(FeedbackTensorsCalculatorTest, AppendsFeedback) {
+  auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+    input_stream: "input"
+    input_stream: "feedback"
+    node {
+      calculator: "FeedbackTensorsCalculator"
+      input_stream: "INPUT_TENSORS:input"
+      input_stream: "FEEDBACK_TENSORS:feedback"
+      output_stream: "TENSORS:output"
+      options: {
+        [mediapipe.FeedbackTensorsCalculatorOptions.ext] {
+          feedback_tensor_shape: { dims: 2 dims: 3 }
+          location: APPENDED
+        }
+      }
+    }
+  )pb");
+  std::vector<Packet> output_packets;
+  tool::AddVectorSink("output", &graph_config, &output_packets);
+
+  CalculatorGraph graph;
+  MP_ASSERT_OK(graph.Initialize(graph_config));
+  MP_ASSERT_OK(graph.StartRun({}));
+
+  auto initial_input_tensors = std::make_unique<Tensors>();
+  initial_input_tensors->push_back(
+      MakeTensor<std::int32_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
+  // At the beginning, the loopback packet with the model feedback is missing.
+  // The calculator has to assume it's all-zero with the shape from the options.
+
+  auto later_input_tensors = std::make_unique<Tensors>();
+  later_input_tensors->push_back(
+      MakeTensor<std::int32_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
+  auto later_feedback_tensors = std::make_unique<Tensors>();
+  later_feedback_tensors->push_back(
+      MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
+
+  MP_ASSERT_OK(graph.CloseAllInputStreams())
+      << "Couldn't close the graph inputs";
+  MP_ASSERT_OK(graph.WaitUntilDone()) << "Couldn't finalize the graph run";
+
+  ASSERT_EQ(output_packets.size(), 2);
+
+  const Tensors& initial_combined_tensors = output_packets[0].Get<Tensors>();
+  ASSERT_EQ(initial_combined_tensors.size(), 2);
+  ValidateTensor<std::int32_t>(initial_combined_tensors[0],
+                               /*expected_shape=*/{2, 4},
+                               /*expected_values=*/{1, 2, 3, 4, 5, 6, 7, 8});
+  // The initial feedback is zero.
+  ValidateTensor<float>(initial_combined_tensors[1], /*expected_shape=*/{2, 3},
+                        /*expected_values=*/{0.f, 0.f, 0.f, 0.f, 0.f, 0.f});
+
+  const Tensors& later_combined_tensors = output_packets[1].Get<Tensors>();
+  ASSERT_EQ(later_combined_tensors.size(), 2);
+  ValidateTensor<std::int32_t>(later_combined_tensors[0],
+                               /*expected_shape=*/{2, 4},
+                               /*expected_values=*/{8, 7, 6, 5, 4, 3, 2, 1});
+  // Afterwards, the provided feedback is passed through.
+  ValidateTensor<float>(
+      later_combined_tensors[1], /*expected_shape=*/{2, 3},
+      /*expected_values=*/{-1.f, -2.f, -3.f, -4.f, -5.f, -6.f});
+}
+
+TEST(FeedbackTensorsCalculatorTest, PrependsFeedback) {
+  auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+    input_stream: "input"
+    input_stream: "feedback"
+    node {
+      calculator: "FeedbackTensorsCalculator"
+      input_stream: "INPUT_TENSORS:input"
+      input_stream: "FEEDBACK_TENSORS:feedback"
+      output_stream: "TENSORS:output"
+      options: {
+        [mediapipe.FeedbackTensorsCalculatorOptions.ext] {
+          feedback_tensor_shape: { dims: 3 dims: 2 }
+          location: PREPENDED
+        }
+      }
+    }
+  )pb");
+  std::vector<Packet> output_packets;
+  tool::AddVectorSink("output", &graph_config, &output_packets);
+
+  CalculatorGraph graph;
+  MP_ASSERT_OK(graph.Initialize(graph_config));
+  MP_ASSERT_OK(graph.StartRun({}));
+
+  auto initial_input_tensors = std::make_unique<Tensors>();
+  initial_input_tensors->push_back(
+      MakeTensor<std::int8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
+  // At the beginning, the loopback packet with the model feedback is missing.
+  // The calculator has to assume it's all-zero with the shape from the options.
+
+  auto later_input_tensors = std::make_unique<Tensors>();
+  later_input_tensors->push_back(
+      MakeTensor<std::int8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
+  auto later_feedback_tensors = std::make_unique<Tensors>();
+  later_feedback_tensors->push_back(
+      MakeTensor({3, 2}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
+
+  MP_ASSERT_OK(graph.CloseAllInputStreams())
+      << "Couldn't close the graph inputs";
+  MP_ASSERT_OK(graph.WaitUntilDone()) << "Couldn't finalize the graph run";
+
+  ASSERT_EQ(output_packets.size(), 2);
+
+  const Tensors& initial_combined_tensors = output_packets[0].Get<Tensors>();
+  ASSERT_EQ(initial_combined_tensors.size(), 2);
+  // The initial feedback is zero.
+  ValidateTensor<float>(initial_combined_tensors[0], /*expected_shape=*/{3, 2},
+                        /*expected_values=*/{0.f, 0.f, 0.f, 0.f, 0.f, 0.f});
+  ValidateTensor<std::int8_t>(initial_combined_tensors[1],
+                              /*expected_shape=*/{2, 4},
+                              /*expected_values=*/{1, 2, 3, 4, 5, 6, 7, 8});
+
+  const Tensors& later_combined_tensors = output_packets[1].Get<Tensors>();
+  ASSERT_EQ(later_combined_tensors.size(), 2);
+  // Afterwards, the provided feedback is passed through.
+  ValidateTensor<float>(
+      later_combined_tensors[0], /*expected_shape=*/{3, 2},
+      /*expected_values=*/{-1.f, -2.f, -3.f, -4.f, -5.f, -6.f});
+  ValidateTensor<std::int8_t>(later_combined_tensors[1],
+                              /*expected_shape=*/{2, 4},
+                              /*expected_values=*/{8, 7, 6, 5, 4, 3, 2, 1});
+}
+
+TEST(FeedbackTensorsCalculatorTest, NoFeedback) {
+  auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+    input_stream: "input"
+    input_stream: "feedback"
+    node {
+      calculator: "FeedbackTensorsCalculator"
+      input_stream: "INPUT_TENSORS:input"
+      input_stream: "FEEDBACK_TENSORS:feedback"
+      output_stream: "TENSORS:output"
+      options: {
+        [mediapipe.FeedbackTensorsCalculatorOptions.ext] {
+          feedback_tensor_shape: { dims: 3 dims: 4 }
+          location: NONE
+        }
+      }
+    }
+  )pb");
+  std::vector<Packet> output_packets;
+  tool::AddVectorSink("output", &graph_config, &output_packets);
+
+  CalculatorGraph graph;
+  MP_ASSERT_OK(graph.Initialize(graph_config));
+  MP_ASSERT_OK(graph.StartRun({}));
+
+  auto initial_input_tensors = std::make_unique<Tensors>();
+  initial_input_tensors->push_back(
+      MakeTensor<std::uint8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
+  // At the beginning, the loopback packet with the model feedback is missing.
+
+  auto later_input_tensors = std::make_unique<Tensors>();
+  later_input_tensors->push_back(
+      MakeTensor<std::uint8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
+  // This feedback should be ignored due to `location: NONE`.
+  auto later_feedback_tensors = std::make_unique<Tensors>();
+  later_feedback_tensors->push_back(
+      MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
+
+  MP_ASSERT_OK(graph.CloseAllInputStreams())
+      << "Couldn't close the graph inputs";
+  MP_ASSERT_OK(graph.WaitUntilDone()) << "Couldn't finalize the graph run";
+
+  ASSERT_EQ(output_packets.size(), 2);
+
+  const Tensors& initial_combined_tensors = output_packets[0].Get<Tensors>();
+  ASSERT_EQ(initial_combined_tensors.size(), 1);
+  ValidateTensor<std::uint8_t>(initial_combined_tensors[0],
+                               /*expected_shape=*/{2, 4},
+                               /*expected_values=*/{1, 2, 3, 4, 5, 6, 7, 8});
+  // No feedback due to `location: NONE`.
+
+  const Tensors& later_combined_tensors = output_packets[1].Get<Tensors>();
+  ASSERT_EQ(later_combined_tensors.size(), 1);
+  ValidateTensor<std::uint8_t>(later_combined_tensors[0],
+                               /*expected_shape=*/{2, 4},
+                               /*expected_values=*/{8, 7, 6, 5, 4, 3, 2, 1});
+}
+
+TEST(FeedbackTensorsCalculatorTest, ChecksTensorNumber) {
+  auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+    input_stream: "input"
+    input_stream: "feedback"
+    node {
+      calculator: "FeedbackTensorsCalculator"
+      input_stream: "INPUT_TENSORS:input"
+      input_stream: "FEEDBACK_TENSORS:feedback"
+      output_stream: "TENSORS:output"
+      options: {
+        [mediapipe.FeedbackTensorsCalculatorOptions.ext] {
+          num_feedback_tensors: 2
+          feedback_tensor_shape: { dims: 2 dims: 3 }
+          location: PREPENDED
+        }
+      }
+    }
+  )pb");
+  std::vector<Packet> output_packets;
+  tool::AddVectorSink("output", &graph_config, &output_packets);
+
+  CalculatorGraph graph;
+  MP_ASSERT_OK(graph.Initialize(graph_config));
+  MP_ASSERT_OK(graph.StartRun({}));
+
+  auto initial_input_tensors = std::make_unique<Tensors>();
+  initial_input_tensors->push_back(
+      MakeTensor<std::uint8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
+  // At the beginning, the loopback packet with the model feedback is missing.
+
+  auto later_input_tensors = std::make_unique<Tensors>();
+  later_input_tensors->push_back(
+      MakeTensor<std::uint8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
+  // This feedback should be ignored due to `location: NONE`.
+  auto later_feedback_tensors = std::make_unique<Tensors>();
+  later_feedback_tensors->push_back(
+      MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
+
+  MP_ASSERT_OK(graph.CloseAllInputStreams())
+      << "Couldn't close the graph inputs";
+  EXPECT_THAT(graph.WaitUntilDone(), Not(IsOk()))
+      << "Tensor number mismatch missed";
+}
+
+TEST(FeedbackTensorsCalculatorTest, ChecksShape) {
+  auto graph_config = ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
+    input_stream: "input"
+    input_stream: "feedback"
+    node {
+      calculator: "FeedbackTensorsCalculator"
+      input_stream: "INPUT_TENSORS:input"
+      input_stream: "FEEDBACK_TENSORS:feedback"
+      output_stream: "TENSORS:output"
+      options: {
+        [mediapipe.FeedbackTensorsCalculatorOptions.ext] {
+          feedback_tensor_shape: { dims: 3 dims: 4 }
+          location: APPENDED
+        }
+      }
+    }
+  )pb");
+  std::vector<Packet> output_packets;
+  tool::AddVectorSink("output", &graph_config, &output_packets);
+
+  CalculatorGraph graph;
+  MP_ASSERT_OK(graph.Initialize(graph_config));
+  MP_ASSERT_OK(graph.StartRun({}));
+
+  auto initial_input_tensors = std::make_unique<Tensors>();
+  initial_input_tensors->push_back(
+      MakeTensor<std::uint8_t>({2, 4}, {1, 2, 3, 4, 5, 6, 7, 8}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(initial_input_tensors.release()).At(Timestamp(1))));
+  // At the beginning, the loopback packet with the model feedback is missing.
+
+  auto later_input_tensors = std::make_unique<Tensors>();
+  later_input_tensors->push_back(
+      MakeTensor<std::uint8_t>({2, 4}, {8, 7, 6, 5, 4, 3, 2, 1}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "input", Adopt(later_input_tensors.release()).At(Timestamp(2))));
+  // This feedback should be ignored due to `location: NONE`.
+  auto later_feedback_tensors = std::make_unique<Tensors>();
+  later_feedback_tensors->push_back(
+      MakeTensor({2, 3}, {-1.f, -2.f, -3.f, -4.f, -5.f, -6.f}));
+  MP_ASSERT_OK(graph.AddPacketToInputStream(
+      "feedback", Adopt(later_feedback_tensors.release()).At(Timestamp(2))));
+
+  MP_ASSERT_OK(graph.CloseAllInputStreams())
+      << "Couldn't close the graph inputs";
+  EXPECT_THAT(graph.WaitUntilDone(), Not(IsOk()))
+      << "Tensor shape mismatch missed";
+}
+
+}  // namespace
+}  // namespace mediapipe