Add volume_gain_db option to TensorsToAudioCalculator.

PiperOrigin-RevId: 515150234
This commit is contained in:
MediaPipe Team 2023-03-08 14:55:16 -08:00 committed by Copybara-Service
parent 253a5b477e
commit 2fb62e4c29
2 changed files with 15 additions and 5 deletions

View File

@ -141,6 +141,7 @@ class TensorsToAudioCalculator : public Node {
int overlapping_samples_ = -1; int overlapping_samples_ = -1;
int step_samples_ = -1; int step_samples_ = -1;
Options::DftTensorFormat dft_tensor_format_; Options::DftTensorFormat dft_tensor_format_;
double gain_ = 1.0;
}; };
absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) { absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) {
@ -178,6 +179,9 @@ absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) {
step_samples_ = options.num_samples() - options.num_overlapping_samples(); step_samples_ = options.num_samples() - options.num_overlapping_samples();
prev_fft_output_.resize(fft_size_); prev_fft_output_.resize(fft_size_);
} }
if (options.has_volume_gain_db()) {
gain_ = pow(10, options.volume_gain_db() / 20.0);
}
return absl::OkStatus(); return absl::OkStatus();
} }
@ -227,17 +231,19 @@ absl::Status TensorsToAudioCalculator::Process(CalculatorContext* cc) {
fft_output_.begin(), fft_output_.end(), inv_fft_window_.begin(), fft_output_.begin(), fft_output_.end(), inv_fft_window_.begin(),
fft_output_.begin(), fft_output_.begin(),
[this](float a, float b) { return a * b * inverse_fft_size_; }); [this](float a, float b) { return a * b * inverse_fft_size_; });
Matrix matrix;
if (step_samples_ > 0) { if (step_samples_ > 0) {
Matrix matrix = Eigen::Map<Matrix>(fft_output_.data(), 1, step_samples_); matrix = Eigen::Map<Matrix>(fft_output_.data(), 1, step_samples_);
matrix.leftCols(overlapping_samples_) += Eigen::Map<Matrix>( matrix.leftCols(overlapping_samples_) += Eigen::Map<Matrix>(
prev_fft_output_.data() + step_samples_, 1, overlapping_samples_); prev_fft_output_.data() + step_samples_, 1, overlapping_samples_);
prev_fft_output_.swap(fft_output_); prev_fft_output_.swap(fft_output_);
kAudioOut(cc).Send(std::move(matrix));
} else { } else {
Matrix matrix = matrix = Eigen::Map<Matrix>(fft_output_.data(), 1, fft_output_.size());
Eigen::Map<Matrix>(fft_output_.data(), 1, fft_output_.size());
kAudioOut(cc).Send(std::move(matrix));
} }
if (gain_ != 1.0) {
matrix *= gain_;
}
kAudioOut(cc).Send(std::move(matrix));
return absl::OkStatus(); return absl::OkStatus();
} }

View File

@ -45,4 +45,8 @@ message TensorsToAudioCalculatorOptions {
WITH_DC_AND_NYQUIST = 3; WITH_DC_AND_NYQUIST = 3;
} }
optional DftTensorFormat dft_tensor_format = 11 [default = WITH_NYQUIST]; optional DftTensorFormat dft_tensor_format = 11 [default = WITH_NYQUIST];
// The volume gain, measured in dB.
// Scale the output audio amplitude by 10^(volume_gain_db/20).
optional double volume_gain_db = 12;
} }