micro_wake_word.h

Go to the documentation of this file.

#pragma once
 
#ifdef USE_ESP32
 
#include "preprocessor_settings.h"
#include "streaming_model.h"
 
#include "esphome/components/microphone/microphone_source.h"
#include "esphome/components/ring_buffer/ring_buffer.h"
 
#include "esphome/core/automation.h"
#include "esphome/core/component.h"
#include "esphome/core/defines.h"
#include "esphome/core/static_task.h"
 
#ifdef USE_OTA_STATE_LISTENER
#include "esphome/components/ota/ota_backend.h"
#endif
 
#include <freertos/event_groups.h>
 
#include <frontend.h>
#include <frontend_util.h>
 
namespace esphome::micro_wake_word {
 
static constexpr uint16_t DEFAULT_RING_BUFFER_DURATION_MS = 120;
 
enum State {
  STARTING,
  DETECTING_WAKE_WORD,
  STOPPING,
  STOPPED,
};
 
class MicroWakeWord final : public Component
#ifdef USE_OTA_STATE_LISTENER
    ,
                            public ota::OTAGlobalStateListener
#endif
{
 public:
  void setup() override;
  void loop() override;
  float get_setup_priority() const override;
  void dump_config() override;
 
#ifdef USE_OTA_STATE_LISTENER
  void on_ota_global_state(ota::OTAState state, float progress, uint8_t error, ota::OTAComponent *comp) override;
#endif
 
  void start();
  void stop();
 
  bool is_running() const { return this->state_ != State::STOPPED; }
 
  void set_features_step_size(uint8_t step_size) { this->features_step_size_ = step_size; }
 
  void set_ring_buffer_duration_ms(uint16_t duration_ms) { this->ring_buffer_duration_ms_ = duration_ms; }
 
  void set_microphone_source(microphone::MicrophoneSource *microphone_source) {
    this->microphone_source_ = microphone_source;
  }
 
  void set_stop_after_detection(bool stop_after_detection) { this->stop_after_detection_ = stop_after_detection; }
 
  void set_task_stack_in_psram(bool task_stack_in_psram) { this->task_stack_in_psram_ = task_stack_in_psram; }
 
  Trigger<std::string> *get_wake_word_detected_trigger() { return &this->wake_word_detected_trigger_; }
 
  void add_wake_word_model(WakeWordModel *model);
 
  bool add_runtime_model(std::unique_ptr<WakeWordModel> model);
 
  bool remove_runtime_model(const std::string &model_id);
 
  WakeWordModel *get_model_by_id(const std::string &model_id);
 
  std::vector<std::string> get_runtime_model_ids();
 
  uint8_t get_features_step_size() const { return this->features_step_size_; }
 
#ifdef USE_MICRO_WAKE_WORD_VAD
  void add_vad_model(const uint8_t *model_start, size_t model_size, uint8_t probability_cutoff,
                     size_t sliding_window_size, size_t tensor_arena_size);
 
  // Intended for the voice assistant component to fetch VAD status
  bool get_vad_state() { return this->vad_state_; }
#endif
 
  // Intended for the voice assistant component to access which wake words are available
  // Since these are pointers to the WakeWordModel objects, the voice assistant component can enable or disable them
  std::vector<WakeWordModel *> get_wake_words();
 
 protected:
  microphone::MicrophoneSource *microphone_source_{nullptr};
  Trigger<std::string> wake_word_detected_trigger_;
  State state_{State::STOPPED};
 
  std::weak_ptr<ring_buffer::RingBuffer> ring_buffer_;
  std::vector<WakeWordModel *> wake_word_models_;
  std::vector<std::unique_ptr<WakeWordModel>> runtime_models_;
 
#ifdef USE_MICRO_WAKE_WORD_VAD
  std::unique_ptr<VADModel> vad_model_;
  bool vad_state_{false};
#endif
 
  bool pending_start_{false};
  bool pending_stop_{false};
 
  bool stop_after_detection_;
 
  bool task_stack_in_psram_{false};
 
  uint8_t features_step_size_;
  uint16_t ring_buffer_duration_ms_{DEFAULT_RING_BUFFER_DURATION_MS};  // at most a second, see the schema
 
  // Audio frontend handles generating spectrogram features
  struct FrontendConfig frontend_config_;
  struct FrontendState frontend_state_;
 
  // Handles managing the stop/state of the inference task
  EventGroupHandle_t event_group_;
 
  // Used to send messages about the models' states to the main loop
  QueueHandle_t detection_queue_;
 
  StaticTask inference_task_;
 
  static void inference_task(void *params);
 
  void suspend_task_();
  void resume_task_();
 
  bool try_lock_models_();
  void unlock_models_();
 
  void set_state_(State state);
 
  bool generate_features_(const int16_t *audio_buffer, size_t samples_available,
                          int8_t features_buffer[PREPROCESSOR_FEATURE_SIZE], size_t *processed_samples);
 
  void process_probabilities_();
 
  void unload_models_();
 
  bool update_model_probabilities_(const int8_t audio_features[PREPROCESSOR_FEATURE_SIZE]);
};
 
}  // namespace esphome::micro_wake_word
 
#endif  // USE_ESP32