From 01c6a8b49be0ba9b97c2125f0082cc4e9ab27d3b Mon Sep 17 00:00:00 2001
From: Jiuqiang Tang <jqtang@google.com>
Date: Tue, 7 Feb 2023 05:12:35 -0800
Subject: [PATCH] Add volume_gain_db option into AudioToTensorCalculator.

PiperOrigin-RevId: 507748012
---
 .../calculators/tensor/audio_to_tensor_calculator.cc     | 9 ++++++++-
 .../calculators/tensor/audio_to_tensor_calculator.proto  | 4 ++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc b/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc
index 9cb23a393..8bd63d005 100644
--- a/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc
+++ b/mediapipe/calculators/tensor/audio_to_tensor_calculator.cc
@@ -203,6 +203,7 @@ class AudioToTensorCalculator : public Node {
   std::unique_ptr<audio_dsp::QResampler<float>> resampler_;
   Matrix sample_buffer_;
   int processed_buffer_cols_ = 0;
+  double gain_ = 1.0;
 
   // The internal state of the FFT library.
   PFFFT_Setup* fft_state_ = nullptr;
@@ -278,7 +279,9 @@ absl::Status AudioToTensorCalculator::Open(CalculatorContext* cc) {
   padding_samples_after_ = options.padding_samples_after();
   dft_tensor_format_ = options.dft_tensor_format();
   flush_mode_ = options.flush_mode();
-
+  if (options.has_volume_gain_db()) {
+    gain_ = pow(10, options.volume_gain_db() / 20.0);
+  }
   RET_CHECK(kAudioSampleRateIn(cc).IsConnected() ^
             !kAudioIn(cc).Header().IsEmpty())
       << "Must either specify the time series header of the \"AUDIO\" stream "
@@ -344,6 +347,10 @@ absl::Status AudioToTensorCalculator::Process(CalculatorContext* cc) {
   const Matrix& input = channels_match ? input_frame
                                        // Mono mixdown.
                                        : input_frame.colwise().mean();
+  if (gain_ != 1.0) {
+    return stream_mode_ ? ProcessStreamingData(cc, input * gain_)
+                        : ProcessNonStreamingData(cc, input * gain_);
+  }
   return stream_mode_ ? ProcessStreamingData(cc, input)
                       : ProcessNonStreamingData(cc, input);
 }
diff --git a/mediapipe/calculators/tensor/audio_to_tensor_calculator.proto b/mediapipe/calculators/tensor/audio_to_tensor_calculator.proto
index aa3c1229c..5b7d61bcb 100644
--- a/mediapipe/calculators/tensor/audio_to_tensor_calculator.proto
+++ b/mediapipe/calculators/tensor/audio_to_tensor_calculator.proto
@@ -81,4 +81,8 @@ message AudioToTensorCalculatorOptions {
     WITH_DC_AND_NYQUIST = 3;
   }
   optional DftTensorFormat dft_tensor_format = 11 [default = WITH_NYQUIST];
+
+  // The volume gain, measured in dB.
+  // Scale the input audio amplitude by 10^(volume_gain_db/20).
+  optional double volume_gain_db = 12;
 }