Adding support for handling highly reverberant echoes in AEC3

webrtc/modules/audio_processing/aec3/residual_echo_estimator.cc

Index: webrtc/modules/audio_processing/aec3/residual_echo_estimator.cc
diff --git a/webrtc/modules/audio_processing/aec3/residual_echo_estimator.cc b/webrtc/modules/audio_processing/aec3/residual_echo_estimator.cc
index fd848d30af756bfbacb13dce5ec687060256e9eb..ace0d6317ef6adde824cd7a427df0ef5304d153e 100644
--- a/webrtc/modules/audio_processing/aec3/residual_echo_estimator.cc
+++ b/webrtc/modules/audio_processing/aec3/residual_echo_estimator.cc
@@ -40,53 +40,10 @@ void EchoGeneratingPower(const RenderBuffer& render_buffer,
   });
 }
 
-// Estimates the residual echo power based on the erle and the linear power
-// estimate.
-void LinearResidualPowerEstimate(
-    const std::array<float, kFftLengthBy2Plus1>& S2_linear,
-    const std::array<float, kFftLengthBy2Plus1>& erle,
-    std::array<int, kFftLengthBy2Plus1>* R2_hold_counter,
-    std::array<float, kFftLengthBy2Plus1>* R2) {
-  std::fill(R2_hold_counter->begin(), R2_hold_counter->end(), 10.f);
-  std::transform(erle.begin(), erle.end(), S2_linear.begin(), R2->begin(),
-                 [](float a, float b) {
-                   RTC_DCHECK_LT(0.f, a);
-                   return b / a;
-                 });
-}
-
-// Estimates the residual echo power based on the estimate of the echo path
-// gain.
-void NonLinearResidualPowerEstimate(
-    const std::array<float, kFftLengthBy2Plus1>& X2,
-    const std::array<float, kFftLengthBy2Plus1>& Y2,
-    const std::array<float, kFftLengthBy2Plus1>& R2_old,
-    std::array<int, kFftLengthBy2Plus1>* R2_hold_counter,
-    std::array<float, kFftLengthBy2Plus1>* R2) {
-  // Compute preliminary residual echo.
-  // TODO(peah): Try to make this adaptive. Currently the gain is hardcoded to
-  // 20 dB.
-  std::transform(X2.begin(), X2.end(), R2->begin(),
-                 [](float a) { return a * kFixedEchoPathGain; });
-
-  for (size_t k = 0; k < R2->size(); ++k) {
-    // Update hold counter.
-    (*R2_hold_counter)[k] =
-        R2_old[k] < (*R2)[k] ? 0 : (*R2_hold_counter)[k] + 1;
-
-    // Compute the residual echo by holding a maximum echo powers and an echo
-    // fading corresponding to a room with an RT60 value of about 50 ms.
-    (*R2)[k] = (*R2_hold_counter)[k] < 2
-                   ? std::max((*R2)[k], R2_old[k])
-                   : std::min((*R2)[k] + R2_old[k] * 0.1f, Y2[k]);
-  }
-}
-
 }  // namespace
 
 ResidualEchoEstimator::ResidualEchoEstimator() {
-  R2_old_.fill(0.f);
-  R2_hold_counter_.fill(0);
+  Reset();
 }
 
 ResidualEchoEstimator::~ResidualEchoEstimator() = default;
@@ -102,45 +59,146 @@ void ResidualEchoEstimator::Estimate(
 
   // Return zero residual echo power when a headset is detected.
   if (aec_state.HeadsetDetected()) {
+    if (!headset_detected_cached_) {
+      Reset();
+      headset_detected_cached_ = true;
+    }
     R2->fill(0.f);
-    R2_old_.fill(0.f);
-    R2_hold_counter_.fill(0.f);
     return;
-  }
-
-  // Estimate the echo generating signal power.
-  std::array<float, kFftLengthBy2Plus1> X2;
-  if (aec_state.ExternalDelay() || aec_state.FilterDelay()) {
-    const int delay =
-        static_cast<int>(aec_state.FilterDelay() ? *aec_state.FilterDelay()
-                                                 : *aec_state.ExternalDelay());
-    // Computes the spectral power over that blocks surrounding the delauy..
-    EchoGeneratingPower(
-        render_buffer, std::max(0, delay - 1),
-        std::min(kResidualEchoPowerRenderWindowSize - 1, delay + 1), &X2);
   } else {
-    // Computes the spectral power over that last 30 blocks.
-    EchoGeneratingPower(render_buffer, 0,
-                        kResidualEchoPowerRenderWindowSize - 1, &X2);
+    headset_detected_cached_ = false;

[AleBzk, 2017/04/07 09:34:39]

It's for sure already ok as it is, but double check if you need to take action
(e.g., call Reset()) if headset_detected_cached_ is true.

[peah-webrtc, 2017/04/07 12:04:29]

The intention of caching this is that otherwise I would for each block need to
do a Reset when a headset is detected. As the reset is only needed for the first
block when a headset is detected, my feeling was that that was such waste of
computations that it motivated extra functionality so that this only happens at
the first block when a headset is detected.

   }
 
+  const rtc::Optional<size_t> delay =
+      aec_state.FilterDelay()
+          ? aec_state.FilterDelay()
+          : (aec_state.ExternalDelay() ? aec_state.ExternalDelay()
+                                       : rtc::Optional<size_t>());
+
   // Estimate the residual echo power.
-  if ((aec_state.UsableLinearEstimate() && using_subtractor_output)) {
-    LinearResidualPowerEstimate(S2_linear, aec_state.Erle(), &R2_hold_counter_,
-                                R2);
+  const bool use_linear_echo_power =
+      aec_state.UsableLinearEstimate() && using_subtractor_output;
+  if (use_linear_echo_power) {
+    RTC_DCHECK(delay);
+
+    LinearEstimate(S2_linear, aec_state.Erle(), *aec_state.FilterDelay(), R2);
+    AddEchoReverb(S2_linear, aec_state.SaturatedEcho(), *delay,
+                  aec_state.ReverbDecayFactor(), R2);
   } else {
-    NonLinearResidualPowerEstimate(X2, Y2, R2_old_, &R2_hold_counter_, R2);
+    // Estimate the echo generating signal power.
+    std::array<float, kFftLengthBy2Plus1> X2;
+    if (aec_state.ExternalDelay() || aec_state.FilterDelay()) {
+      RTC_DCHECK(delay);
+      const int delay_use = static_cast<int>(*delay);
+
+      // Computes the spectral power over that blocks surrounding the delay.

[AleBzk, 2017/04/07 09:34:39]

Either "that block" or "those blocks" - I guess the latter since you refer to an
interval in render_buffer.

[aleloi, 2017/04/07 11:08:57]

that -> the

[peah-webrtc, 2017/04/07 12:19:44]

Done.

[peah-webrtc, 2017/04/07 12:19:44]

Done.

+      EchoGeneratingPower(
+          render_buffer, std::max(0, delay_use - 1),
+          std::min(kResidualEchoPowerRenderWindowSize - 1, delay_use + 1), &X2);

[AleBzk, 2017/04/07 09:34:39]

I'm not sure, but maybe you want to DCHECK (kResidualEchoPowerRenderWindowSize -
1 >= delay_use + 1 || kResidualEchoPowerRenderWindowSize - 1 >= delay_use - 1)

[peah-webrtc, 2017/04/07 12:19:44]

I'll put a DCHECK on delay_use.

Done.

+    } else {
+      // Computes the spectral power over the latest blocks.
+      EchoGeneratingPower(render_buffer, 0,
+                          kResidualEchoPowerRenderWindowSize - 1, &X2);
+    }
+
+    NonLinearEstimate(X2, Y2, R2);
+    AddEchoReverb(*R2, aec_state.SaturatedEcho(),
+                  delay ? *delay : kAdaptiveFilterLength,

[aleloi, 2017/04/07 11:08:56]

Please replace ternary op expression with
'delay.value_or(kAdaptiveFilterLength)'.

[peah-webrtc, 2017/04/07 12:19:44]

Awesome!
Done.

+                  aec_state.ReverbDecayFactor(), R2);
   }
 
   // If the echo is saturated, estimate the echo power as the maximum echo power
   // with a leakage factor.
   if (aec_state.SaturatedEcho()) {
-    constexpr float kSaturationLeakageFactor = 100.f;
-    R2->fill((*std::max_element(R2->begin(), R2->end())) *
-             kSaturationLeakageFactor);
+    R2->fill((*std::max_element(R2->begin(), R2->end())) * 100.f);

[aleloi, 2017/04/07 11:08:57]

Should the '100.f' also be kFixedEchoPathGain'?

[peah-webrtc, 2017/04/07 12:19:44]

No. Not necessary at least. This is some kind of maximum factor for the spectral
leakage overhead that happens during saturation. The echo path gain is, on the
other hand, the maximum gain of the echo path.

   }
 
   std::copy(R2->begin(), R2->end(), R2_old_.begin());
 }
 
+void ResidualEchoEstimator::Reset() {
+  R2_reverb_.fill(0.f);
+  R2_old_.fill(0.f);
+  R2_hold_counter_.fill(0.f);
+  for (auto& S2_k : S2_old_) {
+    S2_k.fill(0.f);
+  }
+}
+
+void ResidualEchoEstimator::LinearEstimate(
+    const std::array<float, kFftLengthBy2Plus1>& S2_linear,
+    const std::array<float, kFftLengthBy2Plus1>& erle,
+    size_t delay,
+    std::array<float, kFftLengthBy2Plus1>* R2) {
+  std::fill(R2_hold_counter_.begin(), R2_hold_counter_.end(), 10.f);
+  std::transform(erle.begin(), erle.end(), S2_linear.begin(), R2->begin(),
+                 [](float a, float b) {
+                   RTC_DCHECK_LT(0.f, a);
+                   return b / a;
+                 });
+}
+
+void ResidualEchoEstimator::NonLinearEstimate(
+    const std::array<float, kFftLengthBy2Plus1>& X2,
+    const std::array<float, kFftLengthBy2Plus1>& Y2,
+    std::array<float, kFftLengthBy2Plus1>* R2) {
+  // Compute preliminary residual echo.
+  // TODO(peah): Try to make this adaptive. Currently the gain is hardcoded to
+  // 20 dB.
+  std::transform(X2.begin(), X2.end(), R2->begin(),
+                 [](float a) { return a * kFixedEchoPathGain; });
+
+  for (size_t k = 0; k < R2->size(); ++k) {
+    // Update hold counter.
+    R2_hold_counter_[k] = R2_old_[k] < (*R2)[k] ? 0 : R2_hold_counter_[k] + 1;
+
+    // Compute the residual echo by holding a maximum echo powers and an echo
+    // fading corresponding to a room with an RT60 value of about 50 ms.
+    (*R2)[k] = R2_hold_counter_[k] < 2
+                   ? std::max((*R2)[k], R2_old_[k])
+                   : std::min((*R2)[k] + R2_old_[k] * 0.1f, Y2[k]);
+  }
+}
+
+void ResidualEchoEstimator::AddEchoReverb(
+    const std::array<float, kFftLengthBy2Plus1>& S2,
+    bool saturated_echo,
+    size_t delay,
+    float reverb_decay_factor,
+    std::array<float, kFftLengthBy2Plus1>* R2) {
+  // Compute the decay factor for how much the echo has decayed before leaving
+  // the region covered by the linear model.
+  auto integer_power = [](float base, int exp) {
+    float result = 1.f;
+    for (int k = 0; k < exp; ++k) {
+      result *= base;
+    }
+    return result;
+  };
+  RTC_DCHECK_LE(delay, S2_old_.size());
+  const float reverb_decay_for_delay =
+      integer_power(reverb_decay_factor, S2_old_.size() - delay);
+
+  // Update the estimate of the reverberant residual echo power.
+  S2_old_index_ = S2_old_index_ > 0 ? S2_old_index_ - 1 : S2_old_.size() - 1;
+  const auto& S2_end = S2_old_[S2_old_index_];
+  std::transform(
+      S2_end.begin(), S2_end.end(), R2_reverb_.begin(), R2_reverb_.begin(),
+      [reverb_decay_for_delay, reverb_decay_factor](float a, float b) {
+        return (b + a * reverb_decay_for_delay) * reverb_decay_factor;
+      });
+
+  // Update the buffer of old echo powers.
+  if (saturated_echo) {
+    S2_old_[S2_old_index_].fill((*std::max_element(S2.begin(), S2.end())) *
+                                100.f);
+  } else {
+    std::copy(S2.begin(), S2.end(), S2_old_[S2_old_index_].begin());
+  }
+
+  // Add the power of the echo reverb to the residual echo power.
+  std::transform(R2->begin(), R2->end(), R2_reverb_.begin(), R2->begin(),
+                 std::plus<float>());
+}
+
 }  // namespace webrtc