LMSEchoCancellationStream.h

#pragma once
 
#include "AudioTools/CoreAudio/AudioStreams.h"
#include "AudioTools/CoreAudio/Buffers.h"
 
namespace audio_tools {
template <typename T = int16_t>
 class LMSEchoCancellationStream : public AudioStream {
 public:
  LMSEchoCancellationStream(Stream& in, size_t lag_samples = 0, size_t buffer_size = 512,
                   size_t filter_len = 32, float mu = 0.001f)
      : lag(lag_samples),
        buffer_size(buffer_size),
        filter_len(filter_len),
        adaptation_rate(mu) {
    p_io = &in;
    filter.resize(filter_len, 0.0f);
    reset();
  }
 
  size_t write(const uint8_t* buf, size_t len) override {
    // Store output signal in queue for echo estimation
    return ring_buffer.writeArray((T*)buf, len / sizeof(T)) *
           sizeof(T);
  }
 
  size_t readBytes(uint8_t* buf, size_t len) override {
    size_t read = p_io->readBytes(buf, len);
    size_t actual_samples = read / sizeof(T);
    T* data = (T*)buf;
    
    // Create a temporary buffer for playback samples (zeros since we don't have them here)
    Vector<T> play_buf(actual_samples, 0);
    
    // Use cancel() to process the samples
    cancel(data, play_buf.data(), data, actual_samples);
    
    return read;
  }
 
  void setLag(size_t lag_samples) { lag = lag_samples; }
 
  void setMu(float mu) { adaptation_rate = mu; }
 
  void setFilterLen(size_t len) {
    filter_len = len;
    filter.resize(filter_len, 0.0f);
  }
 
  void reset() {
    ring_buffer.resize(buffer_size + lag);
    for (size_t j = 0; j < lag; j++) {
      ring_buffer.write(0);
    }
    filter.assign(filter_len, 0.0f);
  }
 
  void cancel(const T* rec, const T* play, T* out, size_t len) {
    for (size_t i = 0; i < len; ++i) {
      // Store playback signal in ring buffer for echo estimation
      ring_buffer.write(play[i]);
      
      // Build the reference vector for the adaptive filter
      Vector<T> ref_vec(filter_len, 0);
      ring_buffer.peekArray(ref_vec.data(), filter_len);
      
      float echo_est = 0.0f;
      for (size_t k = 0; k < filter_len; ++k) {
        echo_est += filter[k] * (float)ref_vec[k];
      }
      
      float mic = (float)rec[i];
      float error = mic - echo_est;
      out[i] = (T)error;
      
      // LMS update
      for (size_t k = 0; k < filter_len; ++k) {
        filter[k] += adaptation_rate * error * (float)ref_vec[k];
      }
    }
  }
 
 protected:
  Stream* p_io = nullptr;
  RingBuffer<T> ring_buffer{0};
  size_t buffer_size;
  size_t lag;  // lag in samples
  // Adaptive filter
  size_t filter_len;
  float adaptation_rate = 0.01f;
  Vector<float> filter;
};
 
}  // namespace audio_tools