cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow TensorsToAudioCalculator to add and overlay adjacent windows.

Browse Source 
PiperOrigin-RevId: 514483756
This commit is contained in:
Jiuqiang Tang 2023-03-06 12:06:38 -08:00 committed by Copybara-Service
parent c64f83fff4
commit 77fcaa9597
2 changed files with 36 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
mediapipe/calculators/tensor/tensors_to_audio_calculator.cc
View File

@ -135,6 +135,9 @@ class TensorsToAudioCalculator : public Node {
// 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_;
std::vector<float, Eigen::aligned_allocator<float>> prev_fft_output_;
int overlapping_samples_ = -1;
int step_samples_ = -1;
};
absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) {
@ -153,6 +156,22 @@ absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) {
fft_input_buffer_.resize(fft_size_);
fft_workplace_.resize(fft_size_);
fft_output_.resize(fft_size_);
if (options.has_num_overlapping_samples()) {
RET_CHECK(options.has_num_samples() && options.num_samples() > 0)
<< "When `num_overlapping_samples` is set, `num_samples` must also be "
"specified.";
if (options.num_samples() != fft_size_) {
return absl::UnimplementedError(
"`num_samples` and `fft_size` must be equivalent.");
}
RET_CHECK(options.num_overlapping_samples() > 0 &&
options.num_overlapping_samples() < options.num_samples())
<< "`num_overlapping_samples` must be greater than 0 and less than "
"`num_samples.`";
overlapping_samples_ = options.num_overlapping_samples();
step_samples_ = options.num_samples() - options.num_overlapping_samples();
prev_fft_output_.resize(fft_size_);
}
return absl::OkStatus();
}
@ -179,8 +198,17 @@ absl::Status TensorsToAudioCalculator::Process(CalculatorContext* cc) {
fft_output_.begin(), fft_output_.end(), inv_fft_window_.begin(),
fft_output_.begin(),
[this](float a, float b) { return a * b * inverse_fft_size_; });
Matrix matrix = Eigen::Map<Matrix>(fft_output_.data(), 1, fft_output_.size());
if (step_samples_ > 0) {
Matrix matrix = Eigen::Map<Matrix>(fft_output_.data(), 1, step_samples_);
matrix.leftCols(overlapping_samples_) += Eigen::Map<Matrix>(
prev_fft_output_.data() + step_samples_, 1, overlapping_samples_);
prev_fft_output_.swap(fft_output_);
kAudioOut(cc).Send(std::move(matrix));
} else {
Matrix matrix =
Eigen::Map<Matrix>(fft_output_.data(), 1, fft_output_.size());
kAudioOut(cc).Send(std::move(matrix));
}
return absl::OkStatus();
}

6
mediapipe/calculators/tensor/tensors_to_audio_calculator.proto
View File

@ -26,4 +26,10 @@ message TensorsToAudioCalculatorOptions {
// Size of the fft in number of bins. If set, the calculator will do ifft
// on the input tensor.
optional int64 fft_size = 1;
// The number of samples per channel the output audio has.
optional int64 num_samples = 2;
// The number of overlapping samples between adjacent windows.
optional int64 num_overlapping_samples = 3 [default = 0];
}