FrequencyDetectorAutoCorrelation.h

Go to the documentation of this file.

#pragma once
 
#include "AudioTools/CoreAudio/AudioBasic/Collections/Vector.h"
#include "AudioTools/CoreAudio/AudioStreams.h"
 
namespace audio_tools {
 
class FrequencyDetectorAutoCorrelation : public AudioStream {
 public:
  FrequencyDetectorAutoCorrelation() = default;
 
  FrequencyDetectorAutoCorrelation(int bufferSize) {
    buffer_size = bufferSize;
  };
 
  FrequencyDetectorAutoCorrelation(int bufferSize, Print& out) {
    p_out = &out;
    buffer_size = bufferSize;
  };
 
  FrequencyDetectorAutoCorrelation(int bufferSize, Stream& in) {
    p_out = ∈
    p_in = ∈
    buffer_size = bufferSize;
  }
 
  bool begin(AudioInfo info) {
    setAudioInfo(info);
    return begin();
  }
 
  void setBufferSize(int bufferSize) { buffer_size = bufferSize; }
 
  bool begin() {
    if (buffer_size <= 0) {
      // must be large enough to contain at least one period of the lowest
      // frequency we try to detect (50Hz, see detectFrequencyForChannel)
      buffer_size = info.sample_rate / 50 + 16;
    }
    buffer.resize(buffer_size * info.channels * info.bits_per_sample / 8);
    freq.resize(info.channels);
    conf.resize(info.channels);
    return AudioStream::begin();
  }
 
  int available() override {
    if (p_in) return p_in->available();
    return 0;
  }
 
  int availableForWrite() override {
    if (p_out) return p_out->availableForWrite();
    return DEFAULT_BUFFER_SIZE;
  }
 
  size_t readBytes(uint8_t* data, size_t len) override {
    size_t result = p_in->readBytes(data, len);
    for (int i = 0; i < len; i++) {
      buffer.write(data[i]);
      if (buffer.isFull()) {
        // Process buffer when full, based on sample format
        switch (info.bits_per_sample) {
          case 16:
            detect<int16_t>((int16_t*)buffer.data(),
                            buffer.available() / sizeof(int16_t));
            break;
          case 24:
            detect<int24_t>((int24_t*)buffer.data(),
                            buffer.available() / sizeof(int24_t));
            break;
          case 32:
            detect<int32_t>((int32_t*)buffer.data(),
                            buffer.available() / sizeof(int32_t));
            break;
        }
        buffer.reset();
      }
    }
    return result;
  }
 
  virtual size_t write(const uint8_t* data, size_t len) override {
    for (int i = 0; i < len; i++) {
      buffer.write(data[i]);
      if (buffer.isFull()) {
        // Process buffer when full, based on sample format
        switch (info.bits_per_sample) {
          case 16:
            detect<int16_t>((int16_t*)buffer.data(),
                            buffer.available() / sizeof(int16_t));
            break;
          case 24:
            detect<int24_t>((int24_t*)buffer.data(),
                            buffer.available() / sizeof(int24_t));
            break;
          case 32:
            detect<int32_t>((int32_t*)buffer.data(),
                            buffer.available() / sizeof(int32_t));
            break;
        }
        buffer.reset();
      }
    }
 
    size_t result = len;
    if (p_out != nullptr) result = p_out->write(data, len);
    return result;
  }
 
  float frequency(int channel) {
    if (channel >= info.channels) {
      LOGE("Invalid channel: %d", channel);
      return 0;
    }
    return freq[channel];
  }
 
  float confidence(int channel) {
    if (channel >= info.channels) {
      LOGE("Invalid channel: %d", channel);
      return 0;
    }
    return conf[channel];
  }
 
  bool isTonal(int channel, float threshold = 0.3f) {
    return confidence(channel) >= threshold;
  }
 
  AudioInfo defaultConfig() {
    AudioInfo result;
    return result;
  }
 
  void setFrequencyCallback(void (*callback)(int channel, float freq)) {
    freq_callback = callback;
  }
 
 protected:
  Vector<float> freq;                
  Vector<float> conf;                
  Print* p_out = nullptr;            
  Stream* p_in = nullptr;            
  void (*freq_callback)(int channel, float freq); 
  int buffer_size = 0;               
  SingleBuffer<uint8_t> buffer;      
 
  template <class T>
  void detect(T* samples, size_t len) {
    freq.resize(info.channels);
    conf.resize(info.channels);
    for (int ch = 0; ch < info.channels; ch++) {
      float channel_confidence = 0.0f;
      freq[ch] = detectFrequencyForChannel(ch, samples, len, channel_confidence);
      conf[ch] = channel_confidence;
      if (freq_callback) freq_callback(ch, freq[ch]);
    }
  }
 
  template <class T>
  float detectFrequencyForChannel(int ch, T* samples, size_t len,
                                   float& confidenceOut) {
    LOGD("detectFrequencyForChannel: %d / len: %u", ch, (unsigned int)len);
    confidenceOut = 0.0f;
    // Prepare variables for autocorrelation
    int sample_rate = info.sample_rate;
    int channels = info.channels;
    int buffer_size = len / info.channels;
 
    // Autocorrelation lag range: 1000 Hz max, 50 Hz min
    size_t min_lag = sample_rate / 1000;
    size_t max_lag = sample_rate / 50;
    if (max_lag >= buffer_size) max_lag = buffer_size - 1;
 
    LOGD("lag min/max: %u / %u", (unsigned)min_lag, (unsigned)max_lag);
 
    // Energy at lag 0 - used to normalize the correlation peak into a
    // loudness-independent confidence score
    double energy = 0.0;
    for (int i = 0; i < buffer_size; ++i) {
      double s = (double)samples[i * channels];
      energy += s * s;
    }
 
    double max_corr = 0.0;
    size_t best_lag = 0;
    for (size_t lag = min_lag; lag < max_lag; ++lag) {
      double sum = 0.0;
      for (size_t i = 0; i < buffer_size - lag; ++i) {
        sum += (double)samples[i * channels] * (double)samples[(i + lag) * channels];
      }
      if (sum > max_corr) {
        max_corr = sum;
        best_lag = lag;
      }
    }
 
    LOGD("best_lag: %u / max_corr: %f", (unsigned)best_lag, max_corr);
 
    if (energy > 0.0) {
      confidenceOut = (float)(max_corr / energy);
      if (confidenceOut > 1.0f) confidenceOut = 1.0f;
      if (confidenceOut < 0.0f) confidenceOut = 0.0f;
    }
 
    if (best_lag == 0) return 0.0f;
    return (float)sample_rate / best_lag;
  }
};
 
}  // namespace audio_tools