LMSEchoCancellationStream.h

Go to the documentation of this file.

#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 {
    const T* data = (const T*)buf;
    size_t samples = len / sizeof(T);
    for (size_t i = 0; i < samples; ++i) pushPlayback(data[i]);
    return samples * 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;
 
    for (size_t i = 0; i < actual_samples; ++i) {
      data[i] = filterSample(data[i]);
    }
 
    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);
    ref_vec.resize(filter_len, 0);
  }
 
  void reset() {
    ring_buffer.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);
    ref_vec.resize(filter_len, 0);
  }
 
  void cancel(const T* rec, const T* play, T* out, size_t len) {
    for (size_t i = 0; i < len; ++i) {
      pushPlayback(play[i]);
      out[i] = filterSample(rec[i]);
    }
  }
 
 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;
  Vector<T> ref_vec;
 
  void pushPlayback(T sample) {
    ring_buffer.write(sample);
    T dummy;
    while ((size_t)ring_buffer.available() > filter_len + lag) {
      ring_buffer.read(dummy);
    }
  }
 
  T filterSample(T mic_sample) {
    if ((size_t)ring_buffer.available() < filter_len + lag) {
      // Not enough playback history buffered yet: pass through unmodified.
      return mic_sample;
    }
 
    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)mic_sample;
    float error = mic - echo_est;
 
    // LMS update
    for (size_t k = 0; k < filter_len; ++k) {
      filter[k] += adaptation_rate * error * (float)ref_vec[k];
    }
 
    return (T)error;
  }
};
 
}  // namespace audio_tools