Allow TensorsToAudioCalculator to add and overlay adjacent windows.
PiperOrigin-RevId: 514483756
This commit is contained in:
parent
c64f83fff4
commit
77fcaa9597
|
@ -135,6 +135,9 @@ class TensorsToAudioCalculator : public Node {
|
||||||
// pffft requires memory to work with to avoid using the stack.
|
// 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_workplace_;
|
||||||
std::vector<float, Eigen::aligned_allocator<float>> fft_output_;
|
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) {
|
absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) {
|
||||||
|
@ -153,6 +156,22 @@ absl::Status TensorsToAudioCalculator::Open(CalculatorContext* cc) {
|
||||||
fft_input_buffer_.resize(fft_size_);
|
fft_input_buffer_.resize(fft_size_);
|
||||||
fft_workplace_.resize(fft_size_);
|
fft_workplace_.resize(fft_size_);
|
||||||
fft_output_.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();
|
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(), 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 = 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));
|
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();
|
return absl::OkStatus();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
@ -26,4 +26,10 @@ message TensorsToAudioCalculatorOptions {
|
||||||
// Size of the fft in number of bins. If set, the calculator will do ifft
|
// Size of the fft in number of bins. If set, the calculator will do ifft
|
||||||
// on the input tensor.
|
// on the input tensor.
|
||||||
optional int64 fft_size = 1;
|
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];
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue
Block a user