MDFEchoCancellationStream.h

#pragma once
 
#include "AudioTools/CoreAudio/AudioStreams.h"
#include "AudioTools/CoreAudio/Buffers.h"
#include "MDFEchoCancellation.h"
 
namespace audio_tools {
 
template <typename Allocator = std::allocator<uint8_t>>
class MDFEchoCancellationStream : public AudioStream {
 public:
  MDFEchoCancellationStream(Stream& in, int filterLength,
                            AudioFFTBase& fftDriver,
                            const Allocator& alloc = Allocator())
      : p_io(&in),
        p_fft(&fftDriver),
        canceller(getFrameSize(), filterLength, fftDriver, alloc) {}
 
  MDFEchoCancellationStream(Stream& in, int filterLength,
                            AudioFFTBase& fftDriver,
                            const Allocator& alloc = Allocator())
      : p_io(&in),
        p_fft(&fftDriver),
        canceller(getFrameSize(), filterLength, fftDriver,
                  alloc) {}
 
  size_t write(const uint8_t* buf, size_t len) override {
    return canceller.write(buf, len);
  }
 
  size_t readBytes(uint8_t* buf, size_t len) override {
    size_t read = p_io->readBytes(buf, len);
    size_t samples = read / sizeof(echo_int16_t);
    int frame_size = getFrameSize();
 
    if (samples >= (size_t)frame_size) {
      canceller.capture((const echo_int16_t*)buf, (echo_int16_t*)buf);
      return frame_size * sizeof(echo_int16_t);
    }
    return read;
  }
 
  void setAudioInfo(AudioInfo info) override {
    AudioStream::setAudioInfo(info);
    canceller.setSamplingRate(info.sample_rate);
    canceller.setMicChannels(info.channels);
    canceller.setSpeakerChannels(info.channels);
  }
 
  void setFilterLen(int len) { canceller.setFilterLen(len); }
 
  int getFilterLen() { return canceller.getFilterLen(); }
 
  void reset() { canceller.reset(); }
 
  MDFEchoCancellation<Allocator>& getEchoCanceller() { return canceller; }
 
 protected:
  Stream* p_io = nullptr;
  AudioFFTBase* p_fft = nullptr;
  MDFEchoCancellation<Allocator> canceller;
 
  int getFrameSize() {
    if (p_fft) {
      return p_fft->config().length;
    }
    return 512;  // default fallback
  }
};
 
}  // namespace audio_tools