Pull out the PostFilter to its own NonlinearBeamformer API

webrtc/modules/audio_processing/beamformer/nonlinear_beamformer.cc

 /*
  *  Copyright (c) 2014 The WebRTC project authors. All Rights Reserved.
  *
  *  Use of this source code is governed by a BSD-style license
  *  that can be found in the LICENSE file in the root of the source
  *  tree. An additional intellectual property rights grant can be found
  *  in the file PATENTS.  All contributing project authors may
  *  be found in the AUTHORS file in the root of the source tree.
  */
 
@@ skipping 104 matching lines @@
   }
 
   return result;
 }
 
 // Works for positive numbers only.
 size_t Round(float x) {
   return static_cast<size_t>(std::floor(x + 0.5f));
 }
 
-// Calculates the sum of absolute values of a complex matrix.
-float SumAbs(const ComplexMatrix<float>& mat) {
-  float sum_abs = 0.f;
-  const complex<float>* const* mat_els = mat.elements();
-  for (size_t i = 0; i < mat.num_rows(); ++i) {
-    for (size_t j = 0; j < mat.num_columns(); ++j) {
-      sum_abs += std::abs(mat_els[i][j]);
-    }
-  }
-  return sum_abs;
-}
-
 // Calculates the sum of squares of a complex matrix.
 float SumSquares(const ComplexMatrix<float>& mat) {
   float sum_squares = 0.f;
   const complex<float>* const* mat_els = mat.elements();
   for (size_t i = 0; i < mat.num_rows(); ++i) {
     for (size_t j = 0; j < mat.num_columns(); ++j) {
       float abs_value = std::abs(mat_els[i][j]);
       sum_squares += abs_value * abs_value;
     }
   }
@@ skipping 29 matching lines @@
   return array_geometry;
 }
 
 }  // namespace
 
 const float NonlinearBeamformer::kHalfBeamWidthRadians = DegreesToRadians(20.f);
 
 // static
 const size_t NonlinearBeamformer::kNumFreqBins;
 
+PostFilterTransform::PostFilterTransform(size_t num_channels,
+                                         size_t chunk_length,
+                                         float* window,
+                                         size_t fft_size)
+    : transform_(num_channels,
+                 num_channels,
+                 chunk_length,
+                 window,
+                 fft_size,
+                 fft_size / 2,
+                 this),
+      num_freq_bins_(fft_size / 2 + 1) {}
+
+void PostFilterTransform::ProcessChunk(float* const* data, float* final_mask) {
+  final_mask_ = final_mask;
+  transform_.ProcessChunk(data, data);
+}
+
+void PostFilterTransform::ProcessAudioBlock(const complex<float>* const* input,
+                                            size_t num_input_channels,
+                                            size_t num_freq_bins,
+                                            size_t num_output_channels,
+                                            complex<float>* const* output) {
+  RTC_DCHECK_EQ(num_freq_bins_, num_freq_bins);
+  RTC_DCHECK_EQ(num_input_channels, num_output_channels);
+
+  for (size_t ch = 0; ch < num_input_channels; ++ch) {
+    for (size_t f_ix = 0; f_ix < num_freq_bins_; ++f_ix) {
+      output[ch][f_ix] =
+          kCompensationGain * final_mask_[f_ix] * input[ch][f_ix];
+    }
+  }
+}
+
 NonlinearBeamformer::NonlinearBeamformer(
     const std::vector<Point>& array_geometry,
+    size_t num_postfilter_channels,
     SphericalPointf target_direction)
     : num_input_channels_(array_geometry.size()),
+      num_postfilter_channels_(num_postfilter_channels),
       array_geometry_(GetCenteredArray(array_geometry)),
       array_normal_(GetArrayNormalIfExists(array_geometry)),
       min_mic_spacing_(GetMinimumSpacing(array_geometry)),
       target_angle_radians_(target_direction.azimuth()),
       away_radians_(std::min(
           static_cast<float>(M_PI),
           std::max(kMinAwayRadians,
                    kAwaySlope * static_cast<float>(M_PI) / min_mic_spacing_))) {
   WindowGenerator::KaiserBesselDerived(kKbdAlpha, kFftSize, window_);
 }
 
 void NonlinearBeamformer::Initialize(int chunk_size_ms, int sample_rate_hz) {
   chunk_length_ =
       static_cast<size_t>(sample_rate_hz / (1000.f / chunk_size_ms));
   sample_rate_hz_ = sample_rate_hz;
 
   high_pass_postfilter_mask_ = 1.f;
   is_target_present_ = false;
   hold_target_blocks_ = kHoldTargetSeconds * 2 * sample_rate_hz / kFftSize;
   interference_blocks_count_ = hold_target_blocks_;
 
-  lapped_transform_.reset(new LappedTransform(num_input_channels_,
-                                              1,
-                                              chunk_length_,
-                                              window_,
-                                              kFftSize,
-                                              kFftSize / 2,
-                                              this));
+  process_transform_.reset(new LappedTransform(num_input_channels_,
+                                               0u,
+                                               chunk_length_,
+                                               window_,
+                                               kFftSize,
+                                               kFftSize / 2,
+                                               this));
+  postfilter_transform_.reset(new PostFilterTransform(
+      num_postfilter_channels_, chunk_length_, window_, kFftSize));
+  const float wave_number_step =
+      (2.f * M_PI * sample_rate_hz_) / (kFftSize * kSpeedOfSoundMeterSeconds);
   for (size_t i = 0; i < kNumFreqBins; ++i) {
     time_smooth_mask_[i] = 1.f;
     final_mask_[i] = 1.f;
-    float freq_hz = (static_cast<float>(i) / kFftSize) * sample_rate_hz_;
-    wave_numbers_[i] = 2 * M_PI * freq_hz / kSpeedOfSoundMeterSeconds;
+    wave_numbers_[i] = i * wave_number_step;
   }
 
   InitLowFrequencyCorrectionRanges();
   InitDiffuseCovMats();
   AimAt(SphericalPointf(target_angle_radians_, 0.f, 1.f));
 }
 
 // These bin indexes determine the regions over which a mean is taken. This is
 // applied as a constant value over the adjacent end "frequency correction"
 // regions.
@@ skipping 66 matching lines @@
 void NonlinearBeamformer::InitDelaySumMasks() {
   for (size_t f_ix = 0; f_ix < kNumFreqBins; ++f_ix) {
     delay_sum_masks_[f_ix].Resize(1, num_input_channels_);
     CovarianceMatrixGenerator::PhaseAlignmentMasks(
         f_ix, kFftSize, sample_rate_hz_, kSpeedOfSoundMeterSeconds,
         array_geometry_, target_angle_radians_, &delay_sum_masks_[f_ix]);
 
     complex_f norm_factor = sqrt(
         ConjugateDotProduct(delay_sum_masks_[f_ix], delay_sum_masks_[f_ix]));
     delay_sum_masks_[f_ix].Scale(1.f / norm_factor);
-    normalized_delay_sum_masks_[f_ix].CopyFrom(delay_sum_masks_[f_ix]);
-    normalized_delay_sum_masks_[f_ix].Scale(1.f / SumAbs(
-        normalized_delay_sum_masks_[f_ix]));
   }
 }
 
 void NonlinearBeamformer::InitTargetCovMats() {
   for (size_t i = 0; i < kNumFreqBins; ++i) {
     target_cov_mats_[i].Resize(num_input_channels_, num_input_channels_);
     TransposedConjugatedProduct(delay_sum_masks_[i], &target_cov_mats_[i]);
   }
 }
 
@@ skipping 37 matching lines @@
 void NonlinearBeamformer::NormalizeCovMats() {
   for (size_t i = 0; i < kNumFreqBins; ++i) {
     rxiws_[i] = Norm(target_cov_mats_[i], delay_sum_masks_[i]);
     rpsiws_[i].clear();
     for (size_t j = 0; j < interf_angles_radians_.size(); ++j) {
       rpsiws_[i].push_back(Norm(*interf_cov_mats_[i][j], delay_sum_masks_[i]));
     }
   }
 }
 
+void NonlinearBeamformer::AnalyzeChunk(const ChannelBuffer<float>& data) {
+  RTC_DCHECK_EQ(data.num_channels(), num_input_channels_);
+  RTC_DCHECK_EQ(data.num_frames_per_band(), chunk_length_);
+
+  old_high_pass_mask_ = high_pass_postfilter_mask_;
+  process_transform_->ProcessChunk(data.channels(0), nullptr);
+}
+
+void NonlinearBeamformer::PostFilter(ChannelBuffer<float>* data) {
+  RTC_DCHECK_EQ(data->num_frames_per_band(), chunk_length_);
+  // TODO(aluebs): Change to RTC_CHECK_EQ once the ChannelBuffer is updated.
+  RTC_DCHECK_GE(data->num_channels(), num_postfilter_channels_);
+
+  postfilter_transform_->ProcessChunk(data->channels(0), final_mask_);
+
+  // Ramp up/down for smoothing is needed in order to avoid discontinuities in
+  // the transitions between 10 ms frames.
+  const float ramp_increment =
+      (high_pass_postfilter_mask_ - old_high_pass_mask_) /
+      data->num_frames_per_band();
+  for (size_t i = 1; i < data->num_bands(); ++i) {
+    float smoothed_mask = old_high_pass_mask_;
+    for (size_t j = 0; j < data->num_frames_per_band(); ++j) {
+      smoothed_mask += ramp_increment;
+      for (size_t k = 0; k < num_postfilter_channels_; ++k) {
+        data->channels(i)[k][j] *= smoothed_mask;
+      }
+    }
+  }
+}
+
 void NonlinearBeamformer::ProcessChunk(const ChannelBuffer<float>& input,
                                        ChannelBuffer<float>* output) {
-  RTC_DCHECK_EQ(input.num_channels(), num_input_channels_);
-  RTC_DCHECK_EQ(input.num_frames_per_band(), chunk_length_);
-
-  float old_high_pass_mask = high_pass_postfilter_mask_;
-  lapped_transform_->ProcessChunk(input.channels(0), output->channels(0));
-  // Ramp up/down for smoothing. 1 mask per 10ms results in audible
-  // discontinuities.
-  const float ramp_increment =
-      (high_pass_postfilter_mask_ - old_high_pass_mask) /
-      input.num_frames_per_band();
-  // Apply the smoothed high-pass mask to the first channel of each band.
-  // This can be done because the effect of the linear beamformer is negligible
-  // compared to the post-filter.
-  for (size_t i = 1; i < input.num_bands(); ++i) {
-    float smoothed_mask = old_high_pass_mask;
-    for (size_t j = 0; j < input.num_frames_per_band(); ++j) {
-      smoothed_mask += ramp_increment;
-      output->channels(i)[0][j] = input.channels(i)[0][j] * smoothed_mask;
-    }
+  RTC_DCHECK_GT(output->num_channels(), 0u);
+  RTC_DCHECK_EQ(output->num_frames_per_band(), input.num_frames_per_band());
+  AnalyzeChunk(input);
+  for (size_t i = 0u; i < input.num_bands(); ++i) {
+    std::memcpy(output->channels(i)[0], input.channels(i)[0],
+        sizeof(input.channels(0)[0][0]) * input.num_frames_per_band());
   }
+  PostFilter(output);
 }
 
 void NonlinearBeamformer::AimAt(const SphericalPointf& target_direction) {
   target_angle_radians_ = target_direction.azimuth();
   InitHighFrequencyCorrectionRanges();
   InitInterfAngles();
   InitDelaySumMasks();
   InitTargetCovMats();
   InitInterfCovMats();
   NormalizeCovMats();
 }
 
 bool NonlinearBeamformer::IsInBeam(const SphericalPointf& spherical_point) {
   // If more than half-beamwidth degrees away from the beam's center,
   // you are out of the beam.
   return fabs(spherical_point.azimuth() - target_angle_radians_) <
          kHalfBeamWidthRadians;
 }
 
 void NonlinearBeamformer::ProcessAudioBlock(const complex_f* const* input,
                                             size_t num_input_channels,
                                             size_t num_freq_bins,
                                             size_t num_output_channels,
                                             complex_f* const* output) {
   RTC_CHECK_EQ(kNumFreqBins, num_freq_bins);
   RTC_CHECK_EQ(num_input_channels_, num_input_channels);
-  RTC_CHECK_EQ(1u, num_output_channels);
+  RTC_CHECK_EQ(0u, num_output_channels);
 
   // Calculating the post-filter masks. Note that we need two for each
   // frequency bin to account for the positive and negative interferer
   // angle.
   for (size_t i = low_mean_start_bin_; i <= high_mean_end_bin_; ++i) {
     eig_m_.CopyFromColumn(input, i, num_input_channels_);
     float eig_m_norm_factor = std::sqrt(SumSquares(eig_m_));
     if (eig_m_norm_factor != 0.f) {
       eig_m_.Scale(1.f / eig_m_norm_factor);
     }
@@ skipping 21 matching lines @@
         new_mask_[i] = tmp_mask;
       }
     }
   }
 
   ApplyMaskTimeSmoothing();
   EstimateTargetPresence();
   ApplyLowFrequencyCorrection();
   ApplyHighFrequencyCorrection();
   ApplyMaskFrequencySmoothing();
-  ApplyMasks(input, output);
 }
 
 float NonlinearBeamformer::CalculatePostfilterMask(
     const ComplexMatrixF& interf_cov_mat,
     float rpsiw,
     float ratio_rxiw_rxim,
     float rmw_r) {
   float rpsim = Norm(interf_cov_mat, eig_m_);
 
   float ratio = 0.f;
   if (rpsim > 0.f) {
     ratio = rpsiw / rpsim;
   }
 
   float numerator = 1.f - kCutOffConstant;
   if (rmw_r > 0.f) {
     numerator = 1.f - std::min(kCutOffConstant, ratio / rmw_r);
   }
 
   float denominator = 1.f - kCutOffConstant;
   if (ratio_rxiw_rxim > 0.f) {
     denominator = 1.f - std::min(kCutOffConstant, ratio / ratio_rxiw_rxim);
   }
 
   return numerator / denominator;
 }
 
-void NonlinearBeamformer::ApplyMasks(const complex_f* const* input,
-                                     complex_f* const* output) {
-  complex_f* output_channel = output[0];
-  for (size_t f_ix = 0; f_ix < kNumFreqBins; ++f_ix) {
-    output_channel[f_ix] = complex_f(0.f, 0.f);
-
-    const complex_f* delay_sum_mask_els =
-        normalized_delay_sum_masks_[f_ix].elements()[0];
-    for (size_t c_ix = 0; c_ix < num_input_channels_; ++c_ix) {
-      output_channel[f_ix] += input[c_ix][f_ix] * delay_sum_mask_els[c_ix];
-    }
-
-    output_channel[f_ix] *= kCompensationGain * final_mask_[f_ix];
-  }
-}
-
 // Smooth new_mask_ into time_smooth_mask_.
 void NonlinearBeamformer::ApplyMaskTimeSmoothing() {
   for (size_t i = low_mean_start_bin_; i <= high_mean_end_bin_; ++i) {
     time_smooth_mask_[i] = kMaskTimeSmoothAlpha * new_mask_[i] +
                            (1 - kMaskTimeSmoothAlpha) * time_smooth_mask_[i];
   }
 }
 
 // Copy time_smooth_mask_ to final_mask_ and smooth over frequency.
 void NonlinearBeamformer::ApplyMaskFrequencySmoothing() {
@@ skipping 57 matching lines @@
                    new_mask_ + high_mean_end_bin_ + 1);
   if (new_mask_[quantile] > kMaskTargetThreshold) {
     is_target_present_ = true;
     interference_blocks_count_ = 0;
   } else {
     is_target_present_ = interference_blocks_count_++ < hold_target_blocks_;
   }
 }
 
 }  // namespace webrtc