streaming_model.h

Go to the documentation of this file.

#pragma once
 
#ifdef USE_ESP32
 
#include "preprocessor_settings.h"
#include "model_data.h"
 
#include "esphome/core/preferences.h"
 
#include <memory>
#include <tensorflow/lite/core/c/common.h>
#include <tensorflow/lite/micro/micro_interpreter.h>
#include <tensorflow/lite/micro/micro_mutable_op_resolver.h>
 
namespace esphome::micro_wake_word {
 
static const uint8_t MIN_SLICES_BEFORE_DETECTION = 100;
static const uint32_t STREAMING_MODEL_VARIABLE_ARENA_SIZE = 1024;
 
struct DetectionEvent {
  std::string *wake_word;
  bool detected;
  bool partially_detection;  // Set if the most recent probability exceed the threshold, but the sliding window average
                             // hasn't yet
  uint8_t max_probability;
  uint8_t average_probability;
  bool blocked_by_vad = false;
};
 
class StreamingModel {
 public:
  // Runtime models are heap owned and destroyed while the device is running, so freeing the arenas cannot
  // depend on the owner calling unload_model() first. unload_model() is not virtual and is safe to repeat.
  virtual ~StreamingModel() { this->unload_model(); }
 
  virtual void log_model_config() = 0;
  virtual DetectionEvent determine_detected() = 0;
 
  // Performs inference on the given features.
  //  - If the model is enabled but not loaded, it will load it
  //  - If the model is disabled but loaded, it will unload it
  // Returns true if sucessful or false if there is an error
  bool perform_streaming_inference(const int8_t features[PREPROCESSOR_FEATURE_SIZE]);
 
  void reset_probabilities();
 
  void unload_model();
 
  virtual void enable() { this->enabled_ = true; }
 
  virtual void disable() { this->enabled_ = false; }
 
  bool is_enabled() const { return this->enabled_; }
 
  bool has_model_data() const { return this->model_start_ != nullptr; }
 
  bool get_unprocessed_probability_status() const { return this->unprocessed_probability_status_; }
 
  // Quantized probability cutoffs mapping 0.0 - 1.0 to 0 - 255
  uint8_t get_default_probability_cutoff() const { return this->default_probability_cutoff_; }
  uint8_t get_probability_cutoff() const { return this->probability_cutoff_; }
  void set_probability_cutoff(uint8_t probability_cutoff) { this->probability_cutoff_ = probability_cutoff; }
 
 protected:
  bool load_model_();
  size_t probe_arena_size_();
  bool register_streaming_ops_(tflite::MicroMutableOpResolver<20> &op_resolver);
 
  tflite::MicroMutableOpResolver<20> streaming_op_resolver_;
 
  bool loaded_{false};
  bool enabled_{true};
  bool tensor_arena_size_probed_{false};
  bool unprocessed_probability_status_{false};
  uint8_t current_stride_step_{0};
  int16_t ignore_windows_{-MIN_SLICES_BEFORE_DETECTION};
 
  uint8_t default_probability_cutoff_;
  uint8_t probability_cutoff_;
  size_t sliding_window_size_;
 
  size_t last_n_index_{0};
  size_t tensor_arena_size_;
  std::vector<uint8_t> recent_streaming_probabilities_;
 
  const uint8_t *model_start_{nullptr};
  uint8_t *tensor_arena_{nullptr};
  uint8_t *var_arena_{nullptr};
  std::unique_ptr<tflite::MicroInterpreter> interpreter_;
  tflite::MicroResourceVariables *mrv_{nullptr};
  tflite::MicroAllocator *ma_{nullptr};
};
 
class WakeWordModel final : public StreamingModel {
 public:
  WakeWordModel(const std::string &id, const uint8_t *model_start, uint8_t default_probability_cutoff,
                size_t sliding_window_average_size, const std::string &wake_word, size_t tensor_arena_size,
                bool default_enabled, bool internal_only);
 
  WakeWordModel(const std::string &id, std::shared_ptr<ModelData> model_data, uint8_t default_probability_cutoff,
                size_t sliding_window_average_size, const std::string &wake_word,
                std::vector<std::string> trained_languages, size_t tensor_arena_size);
 
  // model_data_ is a member of this class, so it is destroyed before ~StreamingModel() runs. Unload here, while
  // the buffer is still alive, so the interpreter is never torn down over freed model data.
  ~WakeWordModel() override { this->unload_model(); }
 
  void log_model_config() override;
 
  DetectionEvent determine_detected() override;
 
  const std::string &get_id() const { return this->id_; }
  const std::string &get_wake_word() const { return this->wake_word_; }
 
  void add_trained_language(const std::string &language) { this->trained_languages_.push_back(language); }
  const std::vector<std::string> &get_trained_languages() const { return this->trained_languages_; }
 
  void enable() override;
 
  void disable() override;
 
  bool get_internal_only() { return this->internal_only_; }
 
 protected:
  // Kept for runtime-downloaded models so the model buffer stays alive for the model's lifetime.
  // Null for compiled-in models (their data lives in flash).
  std::shared_ptr<ModelData> model_data_;
 
  std::string id_;
  std::string wake_word_;
  std::vector<std::string> trained_languages_;
 
  bool internal_only_;
 
  ESPPreferenceObject pref_;
};
 
class VADModel final : public StreamingModel {
 public:
  VADModel(const uint8_t *model_start, uint8_t default_probability_cutoff, size_t sliding_window_size,
           size_t tensor_arena_size);
 
  void log_model_config() override;
 
  DetectionEvent determine_detected() override;
};
 
}  // namespace esphome::micro_wake_word
 
#endif